# Offline Imitation Learning with Model-based Reverse Augmentation

arXiv: https://arxiv.org/abs/2406.12550 · HTML: https://mlsanity.com/paper/2406.12550/

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Placed in Imitation learning methods, one of 13 directions in Offline RL and imitation learning from data.
- Builds on 25 of its 49 references that are in the map.
- First posted on arXiv 18 Jun 2024.

## Placed in

- Offline RL and imitation learning from data / [Imitation learning methods](https://mlsanity.com/offlinerl/imitation-learning-methods/)

## Builds on (most cited first, 25 of 25)

- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Offline Reinforcement Learning with Implicit Q-Learning](https://mlsanity.com/paper/2110.06169/) (2021)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [A Minimalist Approach to Offline Reinforcement Learning](https://mlsanity.com/paper/2106.06860/) (2021)
- [AWAC: Accelerating Online Reinforcement Learning with Offline Datasets](https://mlsanity.com/paper/2006.09359/) (2020)
- [Reinforcement Learning in Healthcare: A Survey](https://mlsanity.com/paper/1908.08796/) (2019)
- [Uncertainty-Based Offline Reinforcement Learning with Diversified Q-Ensemble](https://mlsanity.com/paper/2110.01548/) (2021)
- [A Survey of Inverse Reinforcement Learning: Challenges, Methods and Progress](https://mlsanity.com/paper/1806.06877/) (2018)
- [MimicPlay: Long-Horizon Imitation Learning by Watching Human Play](https://mlsanity.com/paper/2302.12422/) (2023)
- [Reward (Mis)design for Autonomous Driving](https://mlsanity.com/paper/2104.13906/) (2021)
- [Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2005.06800/) (2020)
- [Toward the Fundamental Limits of Imitation Learning](https://mlsanity.com/paper/2009.05990/) (2020)
- [Error Bounds of Imitating Policies and Environments](https://mlsanity.com/paper/2010.11876/) (2020)
- [Discriminator-Weighted Offline Imitation Learning from Suboptimal Demonstrations](https://mlsanity.com/paper/2207.10050/) (2022)
- [S4RL: Surprisingly Simple Self-Supervision for Offline Reinforcement Learning](https://mlsanity.com/paper/2103.06326/) (2021)
- [Adaptivity and Non-stationarity: Problem-dependent Dynamic Regret for Online Convex Optimization](https://mlsanity.com/paper/2112.14368/) (2021)
- [Offline Reinforcement Learning with Reverse Model-based Imagination](https://mlsanity.com/paper/2110.00188/) (2021)
- [Optimal Transport for Offline Imitation Learning](https://mlsanity.com/paper/2303.13971/) (2023)
- [Bidirectional Model-based Policy Optimization](https://mlsanity.com/paper/2007.01995/) (2020)
- [Recall Traces: Backtracking Models for Efficient Reinforcement Learning](https://mlsanity.com/paper/1804.00379/) (2018)
- [Inducing Structure in Reward Learning by Learning Features](https://mlsanity.com/paper/2201.07082/) (2022)
- [Offline Inverse Reinforcement Learning](https://mlsanity.com/paper/2106.05068/) (2021)
- [How to Leverage Diverse Demonstrations in Offline Imitation Learning](https://mlsanity.com/paper/2405.17476/) (2024)

## Built on by (most cited first, 8 of 8)

- [Robust Semi-Supervised Learning in Open Environments](https://mlsanity.com/paper/2412.18256/) (2024)
- [Learning for Long-Horizon Planning via Neuro-Symbolic Abductive Imitation](https://mlsanity.com/paper/2411.18201/) (2024)
- [Robust Offline Imitation Learning Through State-level Trajectory Stitching](https://mlsanity.com/paper/2503.22524/) (2025)
- [Model-Based Offline Reinforcement Learning with Reliability-Guaranteed Sequence Modeling](https://mlsanity.com/paper/2502.06491/) (2025)
- [A Context Engineering Framework for Improving Enterprise AI Agents based on Digital-Twin MDP](https://mlsanity.com/paper/2603.22083/) (2026)
- [Mitigating State Aliasing in Vision-Language-Action Models via Inverse Dynamics Learning](https://mlsanity.com/paper/2605.29577/) (2026)
- [Affordance-Conditioned Decision Making: Bridging the Semantic-Spatial Gap in Zero-Shot Cross-Floor Vision-and-Language Navigation](https://mlsanity.com/paper/2609.32292/) (2026)
- [Beyond Conservatism: Recoverability-Conditioned Exploration for Model-Based Imitation Learning](https://mlsanity.com/paper/2609.33336/) (2026)

## Abstract (arXiv)

In offline Imitation Learning (IL), one of the main challenges is the textitcovariate shift between the expert observations and the actual distribution encountered by the agent, because it is difficult to determine what action an agent should take when outside the state distribution of the expert demonstrations. Recently, the model-free solutions introduce the supplementary data and identify the latent expert-similar samples to augment the reliable samples during learning. Model-based solutions build forward dynamic models with conservatism quantification and then generate additional trajectories in the neighborhood of expert demonstrations. However, without reward supervision, these methods are often over-conservative in the out-of-expert-support regions, because only in states close to expert-observed states can there be a preferred action enabling policy optimization. To encourage more exploration on expert-unobserved states, we propose a novel model-based framework, called offline Imitation Learning with Self-paced Reverse Augmentation (SRA). Specifically, we build a reverse dynamic model from the offline demonstrations, which can efficiently generate trajectories leading to the expert-observed states in a self-paced style. Then, we use the subsequent reinforcement learning method to learn from the augmented trajectories and transit from expert-unobserved states to expert-observed states. This framework not only explores the expert-unobserved states but also guides maximizing long-term returns on these states, ultimately enabling generalization beyond the expert data. Empirical results show that our proposal could effectively mitigate the covariate shift and achieve the state-of-the-art performance on the offline imitation learning benchmarks. Project website: urlhttps://www.lamda.nju.edu.cn/shaojj/KDD24SRA/.