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

TL;DR

The paper tackles offline imitation learning under covariate shift by introducing MILO, a model-based framework that learns a dynamics model from offline data and applies pessimism via a calibrated uncertainty penalty to constrain policy optimization. By formulating the imitation objective as an IPM-based min-max problem and leveraging partial coverage, MILO achieves horizon-efficient performance and provable bounds that depend on the relative coverage between the expert and offline data. Theoretical results cover discrete and continuous MDPs, including non-parametric models like KNRs and Gaussian processes, and extend to a practical offline RL variant. Empirically, MILO outperforms Behavior Cloning and ValueDICE on MuJoCo tasks, demonstrating robustness to limited expert data and varying offline data coverage, with pessimism proving crucial for stability and performance.

Abstract

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.

Figures (4)

• Figure 1: (Left) Frames at timesteps 200, 400, 600, 800, and 1000 for Humanoid-v2 from policies trained with BC on 100 state-action pairs from the expert (blue), BC on 1M offline samples plus 100 expert samples (yellow), and our algorithm MILO (red). The expert has a performance of 3248 and the behavior policy used to collect the offline dataset has performance of $1505\pm 473$ ($\approx 46\%$ of the expert's). (Right) Expert performance normalized scores averaged across 5 seeds.
• Figure 2: Learning curves across five seeds for MILO plotted against the best performance of BC after 1000 epochs of training on the expert/offline+expert data and the best performance of ValueDICE after 10 thousand iterations. The bottom right bar graph shows the expert performance normalized scores where we plot the performance at the last iteration for MILO.
• Figure 3: (Left 2) Learning curves for Hopper and Walker2d with (red) and without (blue) pessimism. MILO generally performs worse without pessimism. (Right 2) Learning curve for Walker2d and Humanoid with more expert samples.
• Figure 4: Performance of MILO with one expert trajectory. Note MILO performance just as well with trajectory inputs as with state-action pair sample inputs.

Theorems & Definitions (42)

• Example 1: Discrete MDPs
• Example 2: KNRs
• Example 3: GPs
• Theorem 3: Bound of MILO
• Remark 4: Implications on offline RL
• Theorem 5
• Theorem 6: Error for KNRs
• Remark 7
• Definition 9: Effective dimension
• Theorem 10: Error for GPs