Mitigating Covariate Shift in Behavioral Cloning via Robust Stationary Distribution Correction

AI summary

Overview

Research area: Offline imitation learning (IL), specifically behavioral cloning (BC) and distributionally robust optimization (DRO) applied to covariate shift in demonstration datasets.

Technical level: Advanced. The paper assumes familiarity with Markov decision processes, Bellman flow constraints, f-divergences, Lagrangian duality, and the DICE (stationary distribution correction ratio estimation) family of algorithms.

Scope: The paper proposes DrilDICE, a distributionally robust behavioral cloning objective that restricts its worst-case distribution set to valid stationary distributions of policies, and evaluates it on a Four Rooms gridworld and three covariate-shifted variants of D4RL -expert datasets.

What This Paper Is About

Behavioral cloning trains a policy by supervised learning on expert state-action pairs, and it is known to fail when the learned policy visits states the expert data does not cover. The paper targets a specific, practical version of this problem: the dataset contains correct expert actions, but the states were not sampled from the expert's stationary distribution (for example, because a sensor recorded some rooms more often than others, or because trajectories are incomplete). The goal is to learn a policy that stays close to the expert even when the data's state distribution is arbitrarily skewed.

Key Contributions

1. A precise problem formalization. The authors isolate a covariate shift where the dataset distribution $d_D(s)$ differs from the expert's stationary state distribution $d_E(s)$, while the action labels remain exactly the expert's, i.e. $\pi_D(s) = \pi_E(s)$ for all $s \in \mathcal{S}$. They explicitly exclude shifts caused by transition dynamics or noisy expert actions, which prior robust IL work addresses.

2. A constrained distributionally robust BC objective. Instead of taking the worst case over all state distributions (which is overly pessimistic and collapses onto a single state), they restrict the uncertainty set with two constraints: (a) the distribution must satisfy the Bellman flow constraint, so it is the stationary distribution of some policy, and (b) it must stay within an f-divergence ball of the dataset distribution $d_D$.

3. A tractable DICE-based solver (DrilDICE). Using a Lagrangian multiplier $\nu$ and a ratio $w(s,a) = d(s,a)/d_D(s,a)$, the inner maximization admits a closed-form solution (Proposition 1, attributed to Lee et al. [16]), reducing the problem to alternating minimization over $\nu$ and $\pi$ in the style of recent DICE methods.

4. The soft TV-distance. Because the TV-distance generator $f_{\mathrm{TV}}(x) = \frac{1}{2}|x-1|$ has a non-invertible derivative (a step function), the authors replace it with $f_{\text{SoftTV}}(x) = \frac{1}{2}\log(\cosh(x-1))$, whose inverse derivative is $\tanh^{-1}(2y)+1$. This yields the closed form $w^*{\pi,\nu}(s,a) = \text{ReLU}\left(\tanh^{-1}\left(\frac{2e{\pi,\nu}(s,a,s')}{\alpha}\right)+1\right)$ while retaining TV-like behavior.

5. Practical shift scenarios and an empirical comparison. They design eight Four Rooms manipulation settings (four room-marginal, four action-marginal) plus three D4RL-based scenarios: rebalanced datasets, time-dependently collected datasets, and segment/incomplete-trajectory collection.

Main Findings

• Four Rooms, all eight scenarios: DrilDICE achieved the highest normalized score in every reported setting. For Room 1 marginal manipulation ($p(u)=0.4$), it scored 95.04 ± 0.48 against 90.84 ± 0.69 for BC, 91.62 ± 0.65 for DemoDICE, 91.38 ± 0.73 for DR-BC, and 94.30 ± 0.41 for OptiDICE-BC.

• Large gains on the worst-case tail: Worst-25% performance, which is where the covariate shift bites hardest, improved substantially. In Room 3, DrilDICE reached 80.16 ± 3.44 versus 55.28 ± 4.74 for BC, 59.60 ± 4.56 for DemoDICE, and 76.80 ± 3.92 for OptiDICE-BC.

• Room 3 as the stress test: The original dataset covered Room 3 only 8.9% of the time; the manipulated version raised this to 40%, causing severe degradation in BC's worst-25% performance. DrilDICE brought worst-25% performance back to levels comparable with the other scenarios.

• D4RL rebalanced datasets: On hopper rebalanced by state with $p(D_1)=0.5$, DrilDICE scored 67.1 ± 8.2 versus 35.4 ± 4.1 for BC, 24.6 ± 1.7 for AW-BC, and 8.6 ± 5.0 for OptiDICE-BC. On walker2d rebalanced by action with $p(D_1)=0.9$, it scored 69.5 ± 8.5 versus 37.6 ± 9.0 for DR-BC.

• D4RL time-dependent collection: With timesteps subsampled from a Beta$(a,b)$ distribution, on hopper with $(a,b)=(1,1)$ DrilDICE reached 45.4 ± 5.1 versus 28.9 ± 3.8 for BC, 26.4 ± 5.8 for DemoDICE, and 22.8 ± 3.9 for OptiDICE-BC.

• Not universally best on every metric: On target 0-1 loss, DrilDICE was occasionally behind OptiDICE-BC (Room 1: 8.31 ± 0.07 versus 7.85 ± 0.05; Action RIGHT: 6.56 ± 0.05 versus 6.48 ± 0.05). On halfcheetah rebalanced by state at $p(D_1)=0.1$, AW-BC scored 53.2 ± 4.3 versus DrilDICE's 52.5 ± 3.6.

• Mechanism visualized: In the Room 3 plot, the BC policy predicted correct expert actions in 26 of 27 states inside Room 3 but performed poorly elsewhere; the learned weights $w^*_{\text{bc}}$ assigned higher values to states BC mispredicted or that were over-represented in the dataset, redirecting training toward the corrected regions.

• Results for Scenario 3 (segment-based collection) are not fully available in the supplied text; the table is truncated mid-entry, so no numerical claims are made here for that setting.

Methodology in Plain English

Start from the ideal BC objective: minimize the imitation loss over states drawn from the expert's state distribution. That distribution is unknown and the dataset does not follow it, so the authors instead minimize an upper bound — the loss under the worst state distribution in a carefully limited set.

Two rules define that set. First, the candidate distribution must correspond to an actual policy's long-run state visitation, which means it must satisfy the Bellman flow constraint (an accounting rule: the probability of visiting a state equals the starting probability plus the probability of flowing in from other states). Second, it must stay close to the observed dataset distribution, enforced with an f-divergence penalty weighted by $\alpha$; straying too far would make finite-sample estimates unreliable.

Solving this max-min problem directly is hard. The trick is to introduce a per-state multiplier $\nu$ and a density ratio $w(s,a) = d(s,a)/d_D(s,a)$. The inner maximization over $w$ then has a closed-form answer given by Proposition 1, so what remains is to alternate between fitting $\nu$ and training the policy on reweighted BC data, $\min_\pi \mathbb{E}{(s,a)\sim d_D}[w^*{\pi,\nu}(s,a) C_\pi(s)]$ — the standard DICE pattern.

For experiments, the practical loop is two-step: train an initial BC policy $\pi_0$, freeze its cost $C_{\pi_0}$, solve for the optimal weights $w^*_{\pi_0}$, then retrain the policy with weighted BC. Four Rooms uses linear function approximators with RBF features (not tabular, to avoid trivial zero-loss solutions) and KL-divergence for solver convenience; D4RL experiments use the soft TV-distance.

Baselines compared: BC, DemoDICE (distribution matching, no supplementary data), AW-BC (unconstrained adversarial weighting), DR-BC (robustness radius $\rho$ on TV distance, no Bellman flow constraint), and OptiDICE-BC (described as best-case weighting). Evaluation uses 100 episodes per policy, 50 repetitions for Four Rooms and 5 for D4RL, reporting normalized score (0 = random, 100 = expert), worst-25% score, and target 0-1 loss.

Why This Matters

Research impact: The paper sharpens the definition of covariate shift in offline IL and shows that the commonly assumed equivalence $d_D \approx d_E$ is load-bearing for distribution-matching methods. It also demonstrates that combining a Bellman flow constraint with an f-divergence ball produces a tighter upper bound than a TV-ball alone (DR-BC's approach), giving a principled answer to the question of how to choose an uncertainty set that is not vacuously pessimistic.

Real-world applications:

• Robotics from teleoperation logs, where demonstrations are recorded unevenly across workspaces or tasks because of operator time or equipment limitations.
• Autonomous driving from fleet data, where certain road types, weather, or times of day are over- or under-represented relative to what the expert driver would actually encounter.
• Medical decision support from retrospective records, where expert actions are correct but the states come from whichever patients happened to be monitored most often.
• Industrial control and sensor-driven data collection, where recording frequency varies by location or season — the exact motivating example the authors use for the time-dependent scenario.

Industry relevance: The method requires only expert demonstrations (no supplementary suboptimal data, no environment interaction), which matches the constraints of most production imitation-learning pipelines. The reweighting step adds cost over plain BC, but the robustness gain in worst-case episodes is the metric that matters for safety-adjacent deployments.

Future Directions

1. Scaling and stability of the alternating optimization. The D4RL standard errors are large (for example 73.7 ± 5.4 and 22.3 ± 2.9), suggesting variance across seeds that a more stable solver or better $\nu$ parameterization could address.

2. Sensitivity to $\alpha$ and the divergence choice. The paper only partially characterizes how the f-divergence penalty $\alpha$ trades pessimism against sample efficiency, and uses KL in Four Rooms but soft TV in D4RL; a systematic study of divergence choice across domains is a natural follow-up.

3. Extending beyond state-distribution shift. The authors deliberately exclude transition shifts and noisy expert actions, leaving open whether Bellman-flow-constrained uncertainty sets can be combined with the robust-IL methods designed for those shifts.

4. Incomplete-trajectory settings. Scenario 3 (segment-based collection) is motivated by incomplete demos but its results are not fully visible in the supplied text; a fuller treatment of how much trajectory fragmentation the method tolerates would strengthen the practical case.

Target Audience

Researchers and graduate students working on offline imitation learning, offline reinforcement learning, and distributionally robust optimization who already understand MDPs and DICE-style ratio estimation. It is also useful for practitioners building behavior-cloning systems from biased or unevenly collected demonstration data, provided they are comfortable with the underlying optimization machinery — the paper is not an introductory treatment, and readers without a background in f-divergences and Lagrangian duality will find the derivations in Sections 3.2 and 3.3 dense.

Authors’ abstract

We consider offline imitation learning (IL), which aims to train an agent to imitate from the dataset of expert demonstrations without online interaction with the environment. Behavioral Cloning (BC) has been a simple yet effective approach to offline IL, but it is also well-known to be vulnerable to the covariate shift resulting from the mismatch between the state distributions induced by the learned policy and the expert policy. Moreover, as often occurs in practice, when…