CCIL: Continuity-based Data Augmentation for Corrective Imitation Learning

TL;DR

CCIL introduces continuity-based corrective labels to robustify imitation learning without extra data beyond expert demonstrations. By learning a locally Lipschitz dynamics model from expert data and generating corrective state-action pairs near the expert support, CCIL augments behavioral cloning with high-quality synthetic data and bounded-error labels. The method combines a Lipschitz-constrained dynamics learning objective with two corrective-label generation techniques (BackTrack and DisturbedAction) solved via a backward Euler root finder and filtered by rejection sampling. Empirical results across classic control, drone, driving with LiDAR, and locomotion/manipulation domains demonstrate improved robustness to disturbances and discontinuities, often outperforming strong baselines like MILO, MOREL, and BC. The work provides theoretical bounds on label quality and presents a practical, data-efficient approach to stabilize imitation learning in diverse robotic scenarios.

Abstract

We present a new technique to enhance the robustness of imitation learning methods by generating corrective data to account for compounding errors and disturbances. While existing methods rely on interactive expert labeling, additional offline datasets, or domain-specific invariances, our approach requires minimal additional assumptions beyond access to expert data. The key insight is to leverage local continuity in the environment dynamics to generate corrective labels. Our method first constructs a dynamics model from the expert demonstration, encouraging local Lipschitz continuity in the learned model. In locally continuous regions, this model allows us to generate corrective labels within the neighborhood of the demonstrations but beyond the actual set of states and actions in the dataset. Training on this augmented data enhances the agent's ability to recover from perturbations and deal with compounding errors. We demonstrate the effectiveness of our generated labels through experiments in a variety of robotics domains in simulation that have distinct forms of continuity and discontinuity, including classic control problems, drone flying, navigation with high-dimensional sensor observations, legged locomotion, and tabletop manipulation.

Figures (6)

• Figure 1: Our proposed framework, CCIL. To enhance the robustness of imitation learning, we propose to augment the dataset with synthetic corrective labels. We leverage the local continuity in the dynamics, learn a regularized dynamics function and generate corrective labels near the expert data support. We provide theoretical guarantees on the quality of the generated labels. We present empirical evaluations CCIL over 14 robotic tasks in 4 domains to showcase CCIL improving imitation learning agents' robustness to disturbances.
• Figure 2: Generating corrective labels from demonstration. Given expert demonstration state $s^*$ and action $a^*$ that arrives at $s^*_{next}$. BackTrack Labels: search for an alternative starting state $s^\mathcal{G}$ that can arrive at expert state $s^*$ if it executes expert action $a^*$. DisturbedAction Labels: search for an alternative starting state $s^\mathcal{G}$ that would arrive at $s^*_{next}$ if it executes a slightly disturbed expert action $a^\mathcal{G}=a^*_t+\Delta$, where $\Delta$ is a sampled noise.
• Figure 3: Evaluation on the Pendulum and Discontinuous Pendulum Task.
• Figure 4: F1tenth
• Figure 5: LiDar POV

Theorems & Definitions (3)

• Definition 4.1
• Theorem 4.2
• Theorem 4.3