PaperlayerView exact revision on arXiv

Source-linked AI summary

Mitigating Covariate Shift in Imitation Learning via Offline Data Without Great Coverage

Jonathan D. Chang, Masatoshi Uehara, Dhruv Sreenivas, Rahul Kidambi, Wen Sun

arXiv:2106.03207v3cs.LGstat.ML

TL;DR

Offline imitation learning must address covariate shift without additional expert or environment interactions. MILO uses static transition data from a potentially suboptimal behavior policy, requiring coverage of expert state-actions rather than global coverage. The paper reports theoretical guarantees under this condition and practical MuJoCo results where MILO imitates with few expert samples while BC fails.

  • Problem

    Offline imitation learning must handle covariate shift, which can make behavior cloning err in states not covered by expert data and compound errors.

    Source: §1 · p. 1, §2 · p. 4

  • Method

    MILO is a model-based offline imitation framework that uses static transition data with partial coverage of expert state-actions and penalizes poorly covered regions.

    Source: §1 · p. 2, §3 · p. 4

  • Results

    On benchmark MuJoCo tasks, MILO imitates with very few expert samples even with low-quality behavior policies, while BC fails.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The paper concludes that offline data with partial expert-state-action coverage can be used to mitigate covariate shift in imitation learning.

    Source: §1 · p. 2, §7 · p. 10

    MILO is only as good as the demonstrator it imitates, including with respect to the demonstrator’s biases and intentions.

    Source: §7 · p. 10

Abstract

from arXiv · show

This paper studies offline Imitation Learning (IL) where an agent learns to imitate an expert demonstrator without additional online environment interactions. Instead, the learner is presented with a static offline dataset of state-action-next state transition triples from a potentially less proficient behavior policy. We introduce Model-based IL from Offline data (MILO): an algorithmic framework that utilizes the static dataset to solve the offline IL problem efficiently both in theory and in practice. In theory, even if the behavior policy is highly sub-optimal compared to the expert, we show that as long as the data from the behavior policy provides sufficient coverage on the expert state-action traces (and with no necessity for a global coverage over the entire state-action space), MILO can provably combat the covariate shift issue in IL. Complementing our theory results, we also demonstrate that a practical implementation of our approach mitigates covariate shift on benchmark MuJoCo continuous control tasks. We demonstrate that with behavior policies whose performances are less than half of that of the expert, MILO still successfully imitates with an extremely low number of expert state-action pairs while traditional offline IL method such as behavior cloning (BC) fails completely. Source code is provided at https://github.com/jdchang1/milo.

1 Introduction

MILO addresses covariate shift in offline imitation learning using expert demonstrations and a static dataset from a potentially suboptimal behavior policy. It requires coverage of expert state-action traces rather than global state-action coverage, and the paper contrasts this setting with prior interactive and offline approaches.

1 Introduction

  • Behavior cloning can suffer covariate shift, making arbitrary mistakes in states absent from expert data and compounding errors that hurt generalization..Source: §1 · p. 1
  • Interactive methods such as DAgger [59], LOLS [15], DART, and AggreVaTe(D) [58] [66] query an expert, while model-based or interaction-based alternatives can be costly or prohibitive in real-world settings.Those alternatives can provide guarantees without a recoverable expert, but active trial-and-error may be unsafe or impossible; full expert-state coverage may also be rare.Source: §1 · pp. 1–2
  • MILO uses offline transition data with partial coverage of expert state-actions to pursue high-fidelity imitation without additional expert or environment interactions.The behavior policy may be highly suboptimal; the required coverage is of expert state-actions, not the entire state-action space.Source: §1 · p. 2
  • The analysis covers discrete MDPs and linear models, and introduces a relative-condition-number result for non-parametric models such as Gaussian Processes.The paper describes this result as new even among existing offline RL results.Source: §1 · p. 3
  • On MuJoCo continuous-control benchmarks, the practical MILO approach imitates with very few expert samples under low-quality behavior policies, while BC fails.The illustrated Humanoid-v2 behavior policy has about 46% of expert performance.Source: Figure 1 · p. 2, §1 · p. 3

1.1 Related work

  • ValueDICE [38] and broader DICE methods [46] [84] require global coverage in theory, and ValueDICE uses online interaction in practice; MILO studies static data without global coverage.Offline IL methods using only expert data still suffer covariate shift in the worst case [52].Source: §1.1 · p. 3
  • CSI requires global expert-data coverage, whereas V-MAIL [50] studies zero-shot transfer from source-task data rather than same-task data from a potentially suboptimal policy..Source: §1.1 · p. 3
  • MILO extends offline RL theory with a relative-condition-number formulation of partial coverage and distribution-dependent results for non-parametric models..Source: §1.1 · p. 3

2 Setting

The setting is episodic finite-horizon imitation learning with a small expert dataset and a static offline transition dataset from a behavior policy. The learner uses these data to seek expert-level performance without online interaction, under function-approximation assumptions.

2 Setting

  • The problem is framed as an episodic finite-horizon MDP with states, actions, transitions, horizon, costs, and an initial-state distribution.Policies map states to action distributions.Source: §2 · p. 3
  • The expert dataset contains state-action demonstrations from an expert policy that may be stochastic and need not be optimal..Source: §2 · p. 3
  • The static offline dataset contains state-action-next-state transitions sampled from a distribution induced by behavior policies that may perform much worse than the expert..Source: §2 · p. 3
  • The learner seeks a policy that performs like the expert using only expert demonstrations and offline data, without real-world or expert interaction.The objective is performance with respect to the unknown ground-truth cost.Source: §2 · p. 4
  • The analysis introduces cost, transition-model, and policy classes and assumes realizability..Source: §2 · p. 4

3 Algorithm

MILO fits a transition model from offline data, penalizes state-action pairs with weak data coverage, and optimizes a pessimistic model-based imitation objective. The framework also yields a non-parametric offline RL analysis without global coverage.

3 Algorithm

  • MILO combines an IPM distance between agent and expert with a pessimism penalty over the policy class.Its three components are model learning, penalty design for poorly covered state-actions, and offline min-max model-based policy optimization.Source: §3 · p. 4
  • The penalty is high where offline coverage is weak and low where coverage is good, steering policies away from regions where the learned model is inaccurate.For well-covered regions, the method encourages the policy to stay close to the expert.Source: §3 · pp. 4–5
  • The framework learns a calibrated transition model and uncertainty estimate from offline data, using methods including empirical transition estimates, ridge regression, Gaussian Processes, or neural-network ensembles.The penalty is designed proportional to Hσ(s, a).Source: §3 · pp. 4–5
  • The min-max objective can be solved by alternating discriminator best responses with policy updates, such as policy gradients inside the learned model..Source: §3 · p. 5

3.1 Specialization to offline RL

  • For offline RL, pessimistic policy optimization minimizes expected learned-model cost plus the penalty; the analysis also gives a non-parametric result without global coverage..Source: §3.1 · p. 5

4 Analysis

MILO’s analysis bounds imitation error under partial coverage of expert state-action traces, rather than requiring global coverage. It treats discrete MDPs, linear KNRs, and nonparametric GP models, with continuous-model results depending on the offline data distribution.

4 Analysis

  • The model estimate is assumed calibrated: its transition error is bounded by the uncertainty measure, which sets a pessimism penalty proportional to H.The specified penalty is b(s, a) = H min(σ(s, a), 2).Source: §4 · p. 5
  • Under partial coverage, MILO’s bound depends on expert-data samples and function-class complexity; asymptotically, its expert-sample error term dominates, improving horizon dependence over BC’s O(H^2√(log(|Π|)/n_e)) rate.The analysis requires only the π_e-concentrability coefficient, not global coverage over the full state-action space.Source: §4 · p. 6

4.1 Analysis: Discrete MDPs

  • For discrete MDPs, the error bound requires only finite partial coverage of expert traces and does not depend on worst-case coverage across the policy class.When partial coverage is small and offline data are sufficient, the expert-sample term dominates and the error is linear in horizon H.Source: §4.1 · p. 6

4.2 Analysis: KNRs and GPs for Continuous MDPs

  • For KNRs, when expert-trace coverage and the rank of the offline-data feature covariance are small, sufficient offline data make the expert-sample term dominate, with linear dependence on horizon H.The result is distribution-dependent and can reflect low-rank offline data rather than worst-case concentrability.Source: §4.2 · pp. 6–7
  • The continuous-model analysis replaces density-ratio coverage with a relative condition number that captures discrepancy between offline and expert data and can reflect low-rank offline data.The paper identifies its nonparametric-model result with relative condition number as new among existing offline RL results.Source: §4 · p. 6, §4.2 · p. 7
  • For GP models, the error bound depends on effective dimension and expert-trace coverage; under favorable values and enough offline data, the expert-sample term dominates.The analysis is distribution-dependent, and the relevant coverage quantity can remain finite even when the density ratio is infinite because it concerns feature subspaces.Source: §4.2 · pp. 6–7

5 Practical Algorithm

The practical MILO instantiation learns an ensemble dynamics model from offline transitions and uses model disagreement as an uncertainty bonus. It combines a discriminator-based policy update with expert behavior-cloning supervision.

5 Practical Algorithm

  • MILO uses an RBF-kernel MMD discriminator and a KL-based trust-region formulation for incremental policy updates inside the learned model.The MMD setup allows the maximum discriminator to be computed in closed form.Source: §5 · p. 7
  • An ensemble of neural network models supplies a bonus equal to the maximum pairwise model disagreement, which is used during policy updates.The algorithm also updates a discriminator and then updates the policy with NPG using the learned model and bonus.Source: §5 · p. 8
  • Expert behavior-cloning loss constrains the policy class or acts as a regularizer, while retaining the deterministic expert as a feasible policy.The paper states that the hard constraint reduces policy-class complexity without eliminating the expert policy; its practical version uses a Lagrange multiplier.Source: §5 · pp. 7–8

6 Experiments

Experiments compare MILO with offline imitation-learning baselines on MuJoCo tasks and examine pessimism, expert-sample quantity, and offline-data coverage. MILO outperforms both BC and ValueDICE across all five environments, while the ablations show that pessimism and data coverage affect performance.

6 Experiments

  • The evaluation uses five MuJoCo environments, comparing MILO with ValueDICE and BC trained on expert data or offline-plus-expert data.Offline datasets contain one million samples and come from behavior policies whose performance is often less than half the expert’s; results average five seeds.Source: §6 · pp. 8–9

6.1 Evaluation on MuJoCo Continuous Control Tasks

  • MILO outperforms both BC and ValueDICE on all five MuJoCo environments and reaches close to expert-level performance on three.Figure 2 compares MILO’s learning curves with the best BC and ValueDICE performances; its bar graph shows MILO’s final-iteration expert-normalized scores.Source: Figure 2 · p. 8, §6.1 · p. 9

6.2 Ablation

  • MILO’s final performance is stabilized and improved by pessimism in the Hopper and Walker2d comparisons.The ablation compares the imitation objective with and without its pessimism penalty.Source: §6.2 · p. 9, Figure 3 · p. 9
  • With one trajectory of expert samples, MILO achieves expert performance on Walker2d and Humanoid, where it did not with very small expert datasets.Figure 3’s right panels examine these two environments with more expert samples.Source: §6.2 · pp. 9–10, Figure 3 · p. 9
  • MILO’s performance degrades as offline-data coverage of expert state-action space decreases, yet it exceeds 20% of expert performance on three environments with random-policy data.The coverage study uses behavior-policy value as a coverage surrogate; Table 2 compares datasets from policies at approximately 50%, 25%, and random performance relative to the expert.Source: §6.2 · p. 10, Table 2 · p. 10

7 Conclusion

The conclusion presents MILO as a way to mitigate covariate shift using offline data with partial coverage, and flags demonstrator bias as a concern for real-world use. The appendix states that its penalty design is valid across tabular MDPs, KNRs, and Gaussian processes.

7 Conclusion

  • MILO uses offline environment-interaction data with partial expert state-action coverage to mitigate covariate shift; scaling to image-based control is left to future work.The paper also advises users to check demonstrations for negative bias or malicious intent before real-world applications.Source: §7 · p. 10
  • The authors warn that imitation can reproduce a demonstrator’s negative biases or intentions, so real-world demonstrations require due diligence.The caution applies to the demonstrator’s intentions as well as performance.Source: §7 · p. 10

A Penalty Designs

  • The appendix states that the penalty design is valid and the model is well-calibrated for tabular MDPs, KNRs, and GPs.The appendix organizes its penalty-design discussion into tabular, KNR, and Gaussian-process models.Source: A Penalty Designs · p. 18, §A.2 · p. 18, §A.3 · p. 19

B Proof of theorem 3

The appendix proof section develops pessimistic policy-evaluation lemmas and uses them to establish Theorem 3 and a finite-sample result for offline RL. Its discrete-MDP discussion relates the error bound to discrepancy between expert and batch-data distributions.

B Proof of theorem 3

  • The proof section uses pessimistic policy-evaluation lemmas under stated assumptions to establish Theorem 3 and an offline-RL finite-sample error bound.The RL result assumes π* ∈ Π, P ∈ P, and Assumption 2, and holds with probability 1 − 2δ.Source: B Proof of theorem 3 · pp. 19–21

C Finite sample error bound for each model

C.1 Discrete MDPs

  • For discrete MDPs, the quantity Cπe measures the distribution difference between expert and batch data and is described as smaller than common offline-RL concentrability coefficients.The text attributes the improved bound to pessimism and says the analysis extends beyond tabular models to continuous MDPs.Source: §C.1 · p. 22

C.2 KNRs

The KNR and Gaussian-process analyses give finite-sample bounds shaped by offline-data distribution and model dimension. The appendix also supplies missing proofs, including the tabular result.

C.2 KNRs

  • The KNR information-gain bound is O(rank[Σρ]^2 log(n_o)), using the offline random-design setting to account for low-rank data.The text contrasts this with the O(d log(n)) worst-case maximum-information-gain bound.Source: §C.2 · p. 22
  • For KNRs, the error bound depends on expert-data discrepancy Cπe and rank[Σρ], which can capture low-rank offline data rather than ambient feature dimension.The analysis states that when Cπe and rank[Σρ] are not too large and offline sample size is sufficient, the covariate-shift problem in BC can be avoided with horizon dependence H.Source: §C.2 · p. 23

C.3 Gaussian processes

  • For Gaussian processes, the finite-sample information-gain bound is O((d*)^2 log(n_o)), and the effective dimension extends the rank concept to infinite-dimensional models.The analysis says covariate shift can be avoided when Cπe and d* are not large and the offline sample size is sufficient.Source: §C.3 · pp. 24–26

C.4 Missing Proofs

  • The appendix says it provides missing proofs for tabular MDP, KNR, and non-parametric GP results.Its tabular-proof subsection begins the proof of Theorem 19 using Theorem 3 and Lemma 35.Source: §C.4 · p. 27, §C.4.1 · p. 27

C.4.1 Missing proofs for tabular result

  • The tabular-result proof invokes Theorem 3 and Lemma 35, then conditions on the resulting event.The passage identifies this as the proof of Theorem 19.Source: §C.4.1 · p. 27

C.4.2 Missing proofs for KNR results

The KNR proof develops bounds by changing between state-action distributions and controlling feature leverage. It uses these results to establish the theorem bounds.

C.4.2 Missing proofs for KNR results

  • The proof of Theorem 20 reuses statements (11) and (12) from the proof of Theorem 21.Source: §C.4.2 · p. 28
  • Lemma 30 compares feature-covariance quadratic forms under two distributions, enabling a change from expert visitation to offline-data distribution.The proof invokes this distribution change when moving from dπe to ρ.Source: §C.4.2 · pp. 29–30
  • The KNR derivation uses Jensen’s inequality and a variational representation as intermediate proof steps.Source: §C.4.2 · p. 30
  • The proof bounds each state-action feature’s regularized leverage by a factor involving covariance rank and confidence.The stated bound scales with rank(Σρ) + log(c2/δ).Source: §C.4.2 · p. 30

C.4.3 Missing proofs of non-parametric model

The appendix supplies proof details and auxiliary bounds for the non-parametric model, alongside experiment and environment documentation. Its non-parametric analysis uses RKHS critical-radius tools and complexity quantities tied to effective dimension.

C.4.3 Missing proofs of non-parametric model

  • The non-parametric proof uses RKHS critical-radius bounds, with d* appearing as a controlling complexity quantity.The appendix invokes a critical-radius calculation for the RKHS ball and gives a variance bound involving d*.Source: §C.4.3 · pp. 32, 35, D Auxiliary Lemmas · p. 39

D Auxiliary Lemmas

  • The auxiliary lemmas provide simulation and Gaussian-distance bounds, plus localized uniform-law and critical-radius tools for linear and RKHS function classes.Source: D Auxiliary Lemmas · pp. 35, 36, 39

E Implementation Details

E.1 Environment Details

  • The environment notes cap episodes at 500 timesteps, remove contact information from observations and reward, and add center-of-mass velocity to the state.The appendix also identifies tables for environment dimensions and reward functions.Source: §E.1 · p. 40, Table 3 · p. 40, Table 4 · p. 40

E.2 Dynamics Ensemble Architecture and Model Learning

The appendix specifies the dynamics ensemble, policy and value networks, hyperparameter selection, and discriminator update used in the experiments. It describes offline dynamics-model training and a model-discrepancy penalty in the cost.

E.2 Dynamics Ensemble Architecture and Model Learning

  • Dynamics learning uses four two-hidden-layer, 1024-unit ReLU models that predict normalized state differences and train for 300 epochs on offline data.The Adam learning rate is 5 × 10^-5.Source: §E.2 · p. 40, Table 5 · p. 41

E.3 Policy Architecture and TRPO Details

  • The policy and value networks use two hidden layers with 32 and 128 units, respectively; both use tanh, and the policy outputs a Gaussian distribution.Advantages are estimated with GAE.Source: §E.3 · p. 41, Table 6 · p. 41

E.4 Hyperparameter Selection

  • Hyperparameters were tuned on one randomly selected Hopper-v2 seed and then applied across environments, with selected TRPO, BC-regularization, and model-size values tested.The tested model hidden-layer sizes were 256, 512, and 1024.Source: §E.4 · p. 41

E.5 Discriminator Update and Cost Function Details

  • The discriminator is linear in Random Fourier Features, and its closed-form update remains valid with BC regularization in the objective.The cost adds a penalty based on the maximum discrepancy among ensemble predictions.Source: §E.5 · pp. 41–42

F Additional Experiments

With one expert trajectory, MILO still matches or nearly matches expert performance across five continuous-control tasks. The experiment uses the same hyperparameters as the main experiments.

F Additional Experiments

  • The one-trajectory experiment tests whether MILO performs well with a complete expert trajectory after the main experiments used an extremely small expert dataset.That small dataset was sampled from many expert trajectories to make BC nearly fail; the one-trajectory test kept the same hyperparameters.Source: F Additional Experiments · p. 42
  • MILO matches or nearly matches expert performance across all five continuous-control tasks using one expert trajectory.The results are averaged across five seeds; Figure 4 compares performance using trajectory inputs with state-action-pair samples.Source: F Additional Experiments · p. 42, Figure 4 · p. 42
Loading 2106.03207v3…