Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Causal Imitation Learning Under Measurement Error and Distribution Shift

Shi Bo, AmirEmad Ghassami

TL;DR

This paper tackles offline imitation learning when the true state is only partially observable through noisy measurements and when training and deployment distributions may differ. It introduces CausIL, a proxy-based causal-imitation framework grounded in proximal causal inference, and defines a causal imitation target $ ext{pi}_{ ext{opt}}(s)$ that depends on the interventional distribution of the expert's action under $ ext{do}(S_t=s)$. The authors provide identification results for both discrete and continuous state spaces and develop practical estimators: a matrix-inversion plug-in method for the discrete case and an RKHS-based adversarial bridge-estimation approach for the continuous case, including a regularized minimax formulation. Empirical evaluation on simulated data and semi-simulated ICU data from PhysioNet demonstrates that CausIL achieves lower imitation error and stronger robustness to measurement-channel shifts and latent-dynamics shifts compared with standard behavioral cloning baselines, highlighting the value of targeting interventional quantities in shifting, measurement-noise settings.

Abstract

We study offline imitation learning (IL) when part of the decision-relevant state is observed only through noisy measurements and the distribution may change between training and deployment. Such settings induce spurious state-action correlations, so standard behavioral cloning (BC) -- whether conditioning on raw measurements or ignoring them -- can converge to systematically biased policies under distribution shift. We propose a general framework for IL under measurement error, inspired by explicitly modeling the causal relationships among the variables, yielding a target that retains a causal interpretation and is robust to distribution shift. Building on ideas from proximal causal inference, we introduce \texttt{CausIL}, which treats noisy state observations as proxy variables, and we provide identification conditions under which the target policy is recoverable from demonstrations without rewards or interactive expert queries. We develop estimators for both discrete and continuous state spaces; for continuous settings, we use an adversarial procedure over RKHS function classes to learn the required parameters. We evaluate \texttt{CausIL} on semi-simulated longitudinal data from the PhysioNet/Computing in Cardiology Challenge 2019 cohort and demonstrate improved robustness to distribution shift compared to BC baselines.

Causal Imitation Learning Under Measurement Error and Distribution Shift

TL;DR

This paper tackles offline imitation learning when the true state is only partially observable through noisy measurements and when training and deployment distributions may differ. It introduces CausIL, a proxy-based causal-imitation framework grounded in proximal causal inference, and defines a causal imitation target that depends on the interventional distribution of the expert's action under . The authors provide identification results for both discrete and continuous state spaces and develop practical estimators: a matrix-inversion plug-in method for the discrete case and an RKHS-based adversarial bridge-estimation approach for the continuous case, including a regularized minimax formulation. Empirical evaluation on simulated data and semi-simulated ICU data from PhysioNet demonstrates that CausIL achieves lower imitation error and stronger robustness to measurement-channel shifts and latent-dynamics shifts compared with standard behavioral cloning baselines, highlighting the value of targeting interventional quantities in shifting, measurement-noise settings.

Abstract

We study offline imitation learning (IL) when part of the decision-relevant state is observed only through noisy measurements and the distribution may change between training and deployment. Such settings induce spurious state-action correlations, so standard behavioral cloning (BC) -- whether conditioning on raw measurements or ignoring them -- can converge to systematically biased policies under distribution shift. We propose a general framework for IL under measurement error, inspired by explicitly modeling the causal relationships among the variables, yielding a target that retains a causal interpretation and is robust to distribution shift. Building on ideas from proximal causal inference, we introduce \texttt{CausIL}, which treats noisy state observations as proxy variables, and we provide identification conditions under which the target policy is recoverable from demonstrations without rewards or interactive expert queries. We develop estimators for both discrete and continuous state spaces; for continuous settings, we use an adversarial procedure over RKHS function classes to learn the required parameters. We evaluate \texttt{CausIL} on semi-simulated longitudinal data from the PhysioNet/Computing in Cardiology Challenge 2019 cohort and demonstrate improved robustness to distribution shift compared to BC baselines.
Paper Structure (34 sections, 7 theorems, 118 equations, 3 figures, 2 tables, 1 algorithm)

This paper contains 34 sections, 7 theorems, 118 equations, 3 figures, 2 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Model Setting and Problem Description
  4. Problem Description
  5. Concerns regarding using behavioral-cloning.
  6. A Causal Imitation Target
  7. Comparison Between $\pi_{\mathrm{opt}}$ and BC
  8. Identification of the Optimal Policy
  9. Discrete Case: Identification via Matrix Inversion
  10. Continuous Case: Identification via a Confounding Bridge
  11. Estimation of the Optimal Policy
  12. Discrete Case: Plug-in Estimator with Coarsenings
  13. Continuous Case: RKHS-Adversarial Estimation
  14. Regularized minimax estimator.
  15. Experiments
  16. ...and 19 more sections

Key Result

Proposition 2.2

Assume $\mathcal{A}=\{1,\dots,K\}$ is finite and let $e(a)\in\mathbb{R}^K$ denote the one-hot encoding of action $a$ (i.e., $e(a)$ has a $1$ in coordinate $a$ and $0$ elsewhere). Consider the one-hot squared-loss risk over deterministic policies $\pi:\mathcal{X}\to\mathcal{A}$, Then any minimizer $\pi^\star\in\arg\min_{\pi}\mathcal{R}(\pi)$ satisfies, for almost every $x\in\mathcal{X}$, Equivale

Figures (3)

  • Figure 1: Causal graph for imitation learning under measurement error.
  • Figure 2: Mean squared error (MSE) of different policies under distributional shift with $|\mathcal{A}|=4$. From left to right: (i) no distributional shift; (ii) shift in the proxy channel $P(U_{t-1}\mid S_{t})$; (iii) simultaneous shifts in $P(W_{t-1}\mid U_{t-1})$.
  • Figure 3: Mean squared error (MSE) of different policies under distributional shift. From left to right: (i) no distributional shift; (ii) shift in ICU time and lactate level; (iii) simultaneous shifts in $P(W_{t-1}\mid U_{t-1})$. Red corresponds to $\hat{\pi}_{\mathrm{opt}}$, orange to $\hat{\pi}_{\mathrm{BC}1}$, and green to $\hat{\pi}_{\mathrm{BC2}}$.

Theorems & Definitions (9)

  • Proposition 2.2: Deterministic BC as a Bayes classifier
  • Remark 2.3
  • Proposition 3.1
  • Corollary 3.2: Population-aggregated interpretation
  • Remark 3.3
  • Proposition 3.4: Rare coincidence of BC targets and $\pi_{\mathrm{opt}}$
  • Theorem 4.5
  • Theorem 4.8
  • Proposition 5.1: Closed-form solution