Scalable Causal Imitation Learning

Paper: Scalable Causal Imitation Learning
Authors: Eylam Tagor, Mingxuan Li, Elias Bareinboim
Link: arXiv:2607.17003v1 [cs.LG], July 2026
Code: ✅ Open-source at github.com/CausalAILab/Scalable-Causal-Imitation-Learning
Keywords: Causality, Imitation Learning, Inverse Reinforcement Learning, Unobserved Confounders

1. Abstract

Imitation learning (IL) learns policies from expert demonstrations in unknown environments, but struggles when the imitator's and expert's observations are mismatched and unobserved confounders are present. Causal imitation learning (CIL) approximates the expert's policy by identifying appropriate adjustment sets via the sequential π-backdoor criterion. However, existing CIL methods — Causal Behavioral Cloning (Causal BC) and Causal Generative Adversarial Imitation Learning (Causal GAIL) — are designed for short-horizon, low-dimensional settings. In continuous control with long horizons and high-dimensional state-action spaces: Causal BC suffers from compounding errors, Causal GAIL is unstable and sample-inefficient, and sequential π-backdoor adjustment becomes impractical. This paper introduces Causal SQIL and Causal IQ-Learn — two off-policy causal imitation learning algorithms combining the causal adjustment framework with state-of-the-art inverse RL objectives. Both operate on causally-adjusted state representations produced by an efficient approximation of the sequential π-backdoor criterion, exploiting the causal structure of continuous control environments to reduce full-horizon adjustment to a fixed-size sliding window. On a suite of confounded environments, both algorithms substantially outperform prior CIL algorithms on long-horizon tasks, sometimes surpassing the expert.

2. Background and Motivation

Traditional IL assumes the expert and imitator have identical sensory capabilities (No Unobserved Confounders, NUC), under which the expert policy is identifiable from observational data. In practice, experts often have access to additional sensors or context unavailable to the imitator, and latent conditions (wind, friction, payload) may unpredictably change their distribution. These unobserved confounders jointly affect state transitions, expert actions, and rewards.

Standard IL methods ignoring confounding overfit to spurious correlations during training, failing when these correlations shift at runtime. In safety-critical tasks like autonomous driving and robotic manipulation, such failures are only revealed during deployment — an unaffordable risk.

Existing CIL methods leverage structural causal knowledge but are limited to low-dimensional, short-horizon tasks:

• Causal BC fails in long-horizon tasks: As a supervised learner, small prediction errors compound. AntMaze-Medium (H≈300) achieves 88.9% expert performance, but drops to 71.2% on AntMaze-Large (H≈700) and 30.8% on HumanoidMaze-Medium (H>1000).
• Causal GAIL fails in high-dimensional domains: Relies on on-policy rollouts (PPO) and a discriminator; long-horizon tasks suffer from insufficient exploration causing credit-assignment failure. AntMaze-Medium reaches 84.6% success, AntMaze-Large drops to 13.1%, HumanoidMaze-Medium collapses to 0%.

Figure 1: Performance of various IL algorithms on Confounded AntMaze Large. (a) Causal IQ-Learn navigates successfully; (b) Causal GAIL struggles to scale; (c) Non-causal methods fail by overfitting spurious correlations; (d) Success rates across all algorithms.

3. Core Method

3.1 Causal Model and Sequential π-Backdoor Criterion

The expert-environment system is modeled as an SCM $\mathcal{M}=\langle\mathbf{U},\mathbf{V},\mathscr{F},P(\mathbf{u})\rangle$, with $\mathbf{U}$ exogenous and $\mathbf{V}$ endogenous variables. Endogenous variables are partitioned into $\mathbf{V}^O$ (observable to imitator) and $\mathbf{V}^L=\mathbf{V}\setminus\mathbf{V}^O$ (latent to imitator).

Expert demonstrations reflect the observational distribution $P(\mathbf{V}^O)$, while the imitator induces $P(\mathbf{V}\mid do(\pi))$. With latent variables, these may differ substantially. CIL's goal is to determine sufficient subsets $\mathbf{Z}_t\subseteq\mathbf{V}^O$ for unbiased approximation $\pi_t(x_t\mid\mathbf{Z}_t)\approx P(x_t\mid\mathbf{Z}_t)$.

Definition (Sequential π-Backdoor): Sets $\{\mathbf{Z}_t\}_{t=0}^H$ satisfy the sequential π-backdoor iff for every $t$, $X_t\perp\!\!\!\perp Y\mid\mathbf{Z}_t$ or $X_t\notin\mathrm{An}(Y)$. Conditioning on $\mathbf{Z}_t$ removes all confounding from latent parents between $X_t$ and $Y$.

Figure 2: Generic causal diagram for sequential imitation learning with unobserved confounders. Causal adjustment excludes confounded variables (e.g., compass W), conditioning only on variables that causally determine expert actions.

3.2 Causal SQIL and Causal IQ-Learn

Given a π-backdoor admissible scope $\mathcal{S}=\{\langle X_t,\mathbf{Z}_t\rangle\}_{t=0}^{H-1}$, the causally-adjusted state representation is $\mathbf{z}_t$ (concatenation of $\mathbf{Z}_t$ values). Both algorithms operate on $(\mathbf{z}_t, x_t)$ pairs instead of standard $(s_t, a_t)$.

Causal SQIL: Assigns fixed reward $r=1$ to expert transitions, $r=0$ to policy transitions, training an SAC agent on the combined replay buffer. The critic minimizes the soft Bellman residual:

$$y=r+\gamma\Big(\min_{j=1,2}Q_{\bar{\theta}_j}(\mathbf{z}',a')-\alpha\log\pi_\phi(a'\mid\mathbf{z}')\Big)$$

$$\mathcal{L}_Q=\mathbb{E}\Big[\big(Q_\theta(\mathbf{z},x)-y\big)^2\Big]$$

The actor maximizes the entropy-regularized objective:

$$\mathcal{L}_\pi=\mathbb{E}_\mathbf{z}\Big[\alpha\log\pi_\phi(x\mid\mathbf{z})-\min_{j=1,2}Q_{\theta_j}(\mathbf{z},x)\Big]$$

Causal IQ-Learn: Learns a Q-function whose implicit reward is consistent with expert behavior. The critic enforces the soft Bellman equation on expert data:

$$\mathcal{L}_{\mathrm{expert}}=\mathbb{E}\Big[\big(Q_\theta(\mathbf{z},x)-\gamma V_{\bar\theta}(\mathbf{z}')\big)^2\Big]$$

and applies a policy-consistency regularizer on policy data:

$$\mathcal{L}_{\mathrm{reg}}=\mathbb{E}\Big[\big(\log\pi_\phi(x\mid\mathbf{z})-Q_\theta(\mathbf{z},x)+V_\theta(\mathbf{z})\big)^2\Big]$$

The combined loss $\mathcal{L}_Q=\mathcal{L}_{\mathrm{expert}}+\lambda\mathcal{L}_{\mathrm{reg}}$ grounds the implicit reward $r(\mathbf{z},x)=Q(\mathbf{z},x)-\gamma V(\mathbf{z}')$ in deconfounded state-action associations.

3.3 Scalable Causal Adjustment: Sliding Window Approximation

With $H\geq1000$, the FindOX algorithm operates over thousands of nodes, $|\mathbf{Z}_t|$ grows linearly, making full-horizon adjustment intractable.

Assumption 1 (Time-Homogeneous Local Causal Structure): There exists window size $k$ such that: (i) all causal influence on $X_t$ is captured within $k$ steps; (ii) the local causal structure is isomorphic at any timestep. This holds naturally in physics-based continuous control like MuJoCo — dynamics depend on immediate state, external forces have temporally localized effects.

Algorithm 1 (Windowed Sequential π-Backdoor Adjustment) has two stages:

• Stage 1 (Lines 1-10): Solve exact sequential π-backdoor on short-horizon proxy graph $\mathcal{G}_k$ (horizon $k+1$). FindOX identifies the maximal admissible set $\mathbf{O}^X$; per-action adjustment sets constructed from Markov boundary MB and boundary actions BA.
• Stage 2 (Lines 11-15): Transfer proxy-graph adjustment sets to full horizon using time-homogeneity. Clip to $k$-step window, reducing dimensionality from $O(H)$ to $O(k)$. Output fixed-dimensional sliding window $S$.

Theorem 1 (Correctness of Windowed Adjustment): Under Assumption 1, adjustment sets $\{\mathbf{Z}_t^k\}$ satisfying the sequential π-backdoor on proxy graph $\mathcal{G}_k$, when transferred to full horizon as $\{\mathbf{Z}_t^H\}$, satisfy the sequential π-backdoor for $(\mathcal{G}_H, \mathbf{X}, Y)$.

Algorithm 1: Windowed Sequential π-Backdoor Adjustment pseudocode. Two stages: proxy-graph solving + full-horizon windowed transfer.

flowchart TD
    A[Causal Graph G + horizon H] --> B{Assumption 1: Time-Homogeneous?}
    B -->|Yes| C[Build Proxy Graph G_k
horizon=k+1]
    C --> D[FindOX: Maximal Admissible Set O_X]
    D --> E{X subseteq O_X?}
    E -->|No| F[Not Imitable]
    E -->|Yes| G[Markov Boundary MB + Boundary Actions BA]
    G --> H[Per-Action Adjustment Sets Z_t^k]
    H --> I[Windowed Transfer
Dim O(H) to O(k)]
    I --> J[Sliding Window S
Fixed Dimension]
    J --> K[Causal Encoding z_t]
    K --> L[Causal SQIL / IQ-Learn
Off-Policy Training]
    L --> M[Deconfounded Policy π]

4. Key Experiments

4.1 Confounded Environment Suite

Confounded continuous-control environments based on OGBench:

• Confounded AntMaze (H=1000): 8-DoF ant navigates maze under latent wind. Imitator observes wind-affected compass $\mathbf{W}$ instead of orientation $\mathbf{O}$.
• Confounded HumanoidMaze (H=2000): 21-DoF humanoid navigates under latent seismic tremors. Imitator observes tremor-affected vibration sensor $\mathbf{W}$ instead of hidden center-of-mass velocity $\mathbf{C}$.

Causal methods exclude $\mathbf{W}$ from adjustment sets; non-causal methods condition on it, failing under distributional shift. 8 algorithms compared (4 causal + 4 non-causal).

4.2 Main Results

Table 1: Evaluation results on confounded tasks. Causal SQIL and Causal IQ-Learn are best in most tasks, sometimes surpassing the expert.

Key Findings:

• All non-causal methods fail catastrophically: 0% or near-0% success across all environments. The failure is not about learning paradigm but about conditioning on all observed variables.
• Causal adjustment necessary but not sufficient: Causal adjustment significantly improves performance, but Causal BC and Causal GAIL degrade with increasing horizon/dimensionality. Causal SQIL and Causal IQ-Learn scale more gracefully.
• Confounding cannot be revealed by in-distribution evaluation: During training, causal and non-causal variants achieve comparable returns; the gap only emerges at runtime when $P(\mathbf{U})$ shifts.
• Windowed approximation is necessary: Moderate $k\in[1,10]$ is ideal; $k=0$ (pure Markov) and $k=100$ see significant drops.
• Causality gap dominates algorithmic gap: Every algorithm's causal variant outperforms its non-causal counterpart — causal reasoning is the primary determinant of success in confounded environments.

Figure 3: Causal BC on Confounded AntMaze. Medium (left) reaches goal; Large (right) follows expert path initially but drifts off course, cannot recover.

Figure 4: Training and runtime evaluation returns. During training, causal/non-causal variants are comparable; at runtime, distributional shift causes non-causal methods to collapse.

5. Limitations and Future Work

• Known causal graph: The method assumes the causal graph is known; in real scenarios it may need to be learned from data or provided by domain experts.
• Bounded temporal influence of confounders: Assumption 1 requires confounder effects to dissipate within $k$ steps, which may not hold for all environments.
• No exploitation of reward structure: The method doesn't leverage latent reward structure, learning only implicitly through expert behavior.
• Sensitivity to expert optimality: Expert policy suboptimality may affect causal adjustment correctness.
• Future: Causal graph learning, relaxing time-homogeneity, exploiting reward structure, handling non-optimal experts.

6. Conclusion

This paper bridges a critical gap in the imitation learning literature: scalable methods can imitate long-horizon expert policies but are fragile when NUC is violated; causal methods are robust under confounding but struggle to scale. By combining the causal adjustment framework with off-policy soft Q-learning methods (SQIL and IQ-Learn), and exploiting continuous control structural properties to reduce full-horizon adjustment to a fixed-size sliding window, two new algorithms substantially outperform prior CIL methods on long-horizon confounded tasks, achieving 90% of expert success rate on average and even surpassing the expert on some tasks. All non-causal methods fail catastrophically under confounding, confirming that causal adjustment is a necessary condition for success. The windowed approximation makes the method scalable to arbitrarily long horizons, and the causal adjustment layer is algorithm-agnostic — any future IL algorithm built on soft Q-learning can be made causal by the same procedure.