CCIL: Continuity-based Data Augmentation for Corrective Imitation Learning

We study imitation learning to teach robots from expert demonstrations. During robots' execution, compounding errors from hardware noise and external disturbances, coupled with incomplete data coverage, can drive the agent into unfamiliar states and cause unpredictable behavior. To address this challenge, we propose a framework, CCIL: Continuity-based data augmentation for Corrective Imitation Learning. It leverages the local continuity inherent in dynamic systems to synthesize corrective labels. CCIL learns a dynamics model from the expert data and uses it to generate labels guiding the agent back to expert states. Our approach makes minimal assumptions, requiring neither expert re-labeling nor ground truth dynamics models. By exploiting local continuity, we derive provable bounds on the errors of the synthesized labels. Through evaluations across diverse robotic domains in simulation and the real world, we demonstrate CCIL's effectiveness in improving imitation learning performance.

Our label generation algorithm consists of three steps: learning a dynamics model, generating corrective labels, and filtering out high-error labels.

We learn a dynamics model by minimizing the following loss: $$\mathbb{E}_{(s_t^*,a_t^*,s_{t+1}^*)\sim\mathcal{D}^*}\left[\hat{f}(s_t^*,a_t^*)+s_t^*-s_{t+1}^*\right]$$ Notably, a learned dynamics model can only yield reliable predictions near its data support - but not on arbitrary states and actions. CCIL decides where to query the learned dynamics models by leveraging the presence of local Lipschitz continuity in the system dynamics. CCIL encourages the learned dynamics function to exhibit local Lipschitz continuity by modifying the training objective, specifically by regularizing the continuity of the learned model with spectral normalization. Concretely, to train a dynamics model $\hat{f}$ using a neural network of $n$-layers with weight matrices $W_1,\ldots,W_n$, one can iteratively minimize the above training objective while regularizing the model by setting $$W_i\leftarrow \frac{W_i}{\max\left(\|W_i\|_2,K^{-n}\right)}\cdot K^{-n}$$ for every $W_i$, where $K$ is the Lipschitz constraint hyperparameter.

With a learned dynamics model $\hat{f}$, we can generate a corrective label $(s_t^\mathcal{G}, a_t^\mathcal{G})$ for every expert data point $(s_t^*, a_t^*)$ such that $s_t^\mathcal{G}+\hat{f}(s_t^\mathcal{G},a_t^\mathcal{G})\approx s_t^*$. One of our label generation methods is BackTrack, inspired by the backwards Euler method used in modern simulators: \begin{align*} s_t^\mathcal{G} &\leftarrow s_t^* - \hat{f}(s_t^*, a_t^*) \\ a_t^\mathcal{G} &\leftarrow a_t^* \end{align*}

By leveraging the local continuity in the environment dynamics, we can derive provable bounds on the correctness of the generated labels. Armed with this error bound, we can filter out high-error labels and only use the ones that are likely to be correct. Concretely, we set a maximum allowable error, which naturally creates a maximum allowable distance between the generated state and the expert state. This can be viewed as a trust region around each expert data point, within which we can trust the generated labels to be accurate.