Causal Imitation Learning Under Measurement Error and Distribution Shift

1. Problem Statement

The paper addresses offline imitation learning (IL) in settings where two critical challenges coexist:

1. Measurement Error: The true decision-relevant state is not directly observed. Instead, the learner observes noisy proxies (measurements) of a latent state component.
2. Distribution Shift: The joint distribution of the latent state and the measurement process may change between the training domain (where expert demonstrations are collected) and the deployment domain.

In such regimes, standard Behavioral Cloning (BC) fails. BC typically learns a policy $\pi(a|o)$ by minimizing the observational conditional probability of actions given observations. The authors argue that under distribution shift, these observational correlations become spurious. Specifically:

• Conditioning on raw measurements ($\pi_{BC2}$) makes the policy brittle to changes in the measurement mechanism (e.g., sensor calibration shifts).
• Ignoring measurements and conditioning only on observed states ($\pi_{BC1}$) makes the policy brittle to changes in the latent dynamics or state transition kernels.
• Crucially, this bias is systematic and does not vanish even with infinite data or longer trajectories.

2. Methodology

A. Causal Framework and Target Policy

The authors propose a new target policy, CausIL, defined not by observational conditioning but by counterfactual intervention.

• Model: The system consists of an observed state $S_t$, a latent state $U_t$ (unobserved), a noisy measurement $W_t$ (proxy for $U_t$), and an expert action $A_t$. The expert acts based on $S_t$ and a delayed latent state $U_{t-1}$ (i.e., $A_t \sim \pi_E(S_t, U_{t-1})$).
• Target Definition: Instead of imitating $p(A_t | S_t, W_{t-1})$, the goal is to learn the causal-optimal policy $\pi_{opt}(s)$, which maximizes the probability of the expert's action under the intervention $do(S_t = s)$:
  $$\pi_{opt}(s) := \arg\max_{a \in \mathcal{A}} p(A_t^{(s)} = a)$$
  where $A_t^{(s)}$ is the potential outcome of the action if the state were set to $s$.
• Robustness: This target depends only on the marginal distribution of the latent state $U_{t-1}$ and the invariant expert mechanism $\pi_E$. It is robust to shifts in the measurement channel $p(W_{t-1}|U_{t-1})$ and, under certain conditions, robust to shifts in state transition kernels that preserve the marginal distribution of $U_{t-1}$.

B. Identification

Since $U_{t-1}$ is unobserved, the authors derive conditions under which $\pi_{opt}$ is identifiable from observational data alone. They adapt the Proximal Causal Inference (PCI) framework, treating the lagged observed state $Z_t = S_{t-1}$ as a "treatment-inducing proxy" and the noisy measurement $W_{t-1}$ as an "outcome-inducing proxy" for the latent confounder $U_{t-1}$.

• Discrete Case: Identification is achieved via matrix inversion. If the proxies $Z_t$ and $W_{t-1}$ are sufficiently informative (satisfying a rank condition), the interventional distribution is identified as:
  $$p(A_t^{(s)}) = P_{A|Z',s} (P_{W'|Z',s})^{-1} P_{W'}$$
  where $P$ matrices represent conditional and marginal probabilities of the proxies.
• Continuous Case: Identification relies on completeness and confounding bridge assumptions. The authors show that there exists a "bridge function" $h_a(w, s)$ satisfying a Fredholm integral equation of the first kind, linking observed data to the counterfactual distribution:
  $$p(A_t^{(s)} = a) = \int h_a(w, s) p(w) dw$$

C. Estimation

The paper proposes estimators for both discrete and continuous settings:

1. Discrete Estimator: A plug-in estimator using empirical frequencies to construct the probability matrices and performing the matrix inversion derived in the identification theorem.
2. Continuous Estimator: An adversarial estimation procedure over Reproducing Kernel Hilbert Spaces (RKHS). The bridge function $h_a$ is learned by solving a regularized minimax problem (saddle-point optimization) that enforces the conditional moment restriction $E[Y - h(W, S) | Z, S] = 0$.

3. Key Contributions

• Theoretical Formulation: The paper formally defines the regime where measurement error and distribution shift coexist and demonstrates that standard BC targets are fundamentally biased in this setting.
• Causal Target: It introduces $\pi_{opt}$, a policy target that is robust to changes in both the measurement process and latent dynamics, provided the expert mechanism is invariant.
• Identification Results: It provides novel identification conditions for offline IL with unobserved confounders and measurement error, leveraging proximal causal inference ideas (specifically the use of lagged states and noisy measurements as proxies).
• Algorithms: It develops practical estimation strategies for both discrete and continuous state spaces, including a closed-form solution for the discrete case and an RKHS-based adversarial estimator for the continuous case.

4. Experimental Results

The authors evaluate CausIL against standard BC baselines ($\pi_{BC1}$ and $\pi_{BC2}$) in two settings:

1. Semi-Simulated Longitudinal Data (PhysioNet 2019):

   • Setup: Using ICU vital signs as observed states ($S_t$) and lactate levels as a delayed latent context ($U_t$). Expert actions (hemodynamic support intensity) are simulated based on $S_t$ and $U_{t-1}$.
   • Shifts: The study introduces "population shifts" (different trajectory compositions) and "measurement shifts" (changes in noise/bias of the lactate proxy).
   • Result: CausIL consistently achieved lower Mean Squared Error (MSE) and remained stable across shifts. In contrast, $\pi_{BC2}$ degraded sharply under measurement shifts, and $\pi_{BC1}$ showed increased error under population shifts.

2. Simulation Study:

   • Setup: Synthetic data generated from the model in Section 2.
   • Result: The experiments confirmed the theoretical robustness claims. $\pi_{opt}$ remained invariant when the measurement channel changed, whereas $\pi_{BC2}$ failed. Similarly, $\pi_{opt}$ was more stable than $\pi_{BC1}$ when state transition kernels shifted but the marginal latent distribution remained fixed.

5. Significance and Claims

The paper claims that CausIL offers a principled solution to the "causal delusion" problem in offline imitation learning when data is noisy and environments shift.

• Robustness: The primary significance is the demonstration that targeting interventional quantities (via proximal identification) yields policies that generalize better to new domains than observational targets, even without reward signals or interactive expert queries.
• Practicality: The method works in fully offline settings, which is critical for high-stakes domains like healthcare (e.g., ICU decision support) and robotics where interactive exploration is impossible or dangerous.
• Limitations: The approach relies on specific identification assumptions (e.g., the existence of valid proxies and the rank/completeness conditions). The authors note that these conditions are non-trivial but are satisfied in many practical scenarios involving delayed latent states and noisy sensors.

The authors conclude that in realistic, shifting environments with measurement error, learning policies from purely observational demonstrations requires moving beyond behavioral cloning to causal frameworks that explicitly model the relationship between latent states, proxies, and actions.