Restoring Noisy Demonstration for Imitation Learning With Diffusion Models

TL;DR

This work tackles the challenge of learning from noisy offline demonstrations in imitation learning. It introduces Diffusion-Model-Based Demonstration Restoration (DMDR), a two-stage framework that first filters clean samples using autoencoders plus Local Outlier Factor, then trains conditional diffusion models with noise-timestep predictors to restore the remaining corrupted data. By aggregating clean and restored demonstrations, DMDR enables IL methods (e.g., BC, IBC, diffusion-policy) to achieve substantial performance gains across robot arm manipulation, dexterous manipulation, and locomotion tasks, and proves robust to diverse noise types and levels. The approach offers a practical path to leveraging imperfect real-world demonstrations for reliable policy learning, with ablations confirming the value of each component and analysis highlighting limitations and future directions for data efficiency and suboptimal-but-uncontaminated data handling.

Abstract

Imitation learning (IL) aims to learn a policy from expert demonstrations and has been applied to various applications. By learning from the expert policy, IL methods do not require environmental interactions or reward signals. However, most existing imitation learning algorithms assume perfect expert demonstrations, but expert demonstrations often contain imperfections caused by errors from human experts or sensor/control system inaccuracies. To address the above problems, this work proposes a filter-and-restore framework to best leverage expert demonstrations with inherent noise. Our proposed method first filters clean samples from the demonstrations and then learns conditional diffusion models to recover the noisy ones. We evaluate our proposed framework and existing methods in various domains, including robot arm manipulation, dexterous manipulation, and locomotion. The experiment results show that our proposed framework consistently outperforms existing methods across all the tasks. Ablation studies further validate the effectiveness of each component and demonstrate the framework's robustness to different noise types and levels. These results confirm the practical applicability of our framework to noisy offline demonstration data.

Figures (8)

• Figure 1: Illustration of Diffusion-Model-Based Demonstration Restoration (DMDR): We propose a two-stage learning framework that first identifies and filters clean samples from the noisy demonstrations. Then, by learning diffusion models using the clean samples, we restore the remaining noisy samples to provide more reliable demonstrations.
• Figure 2: Demonstration filtering.(a) Training. We train the state autoencoder ${\phi}_s$ using the reconstruction loss and apply the Local Outlier Factor (LOF) on the feature space, i.e., $z_s$. We only show the state autoencoder ${\phi}_s$ in the figure for illustration, while the action autoencoder ${\phi}_a$ follows the identical architecture. (b) Inference. We use autoencoders and LOF to individually identify outliers for states and actions. With the predictions of outliers, we filter the noisy demonstrations into four subsets: $D_{(\hat{s}, \hat{a})}$, $D_{(s', a')}$, $D_{(\hat{s}, a')}$, and $D_{(s', \hat{a})}$, which contains clean state-action pairs, noisy state-action pairs, clean states with noisy actions, and noisy states with clean actions, respectively.
• Figure 3: Demonstration restoration.(a) Training of conditional Diffusion Models. We train the state diffusion model $\theta_s$ using the condition of clean action $\hat{a}$ to restore the noise-injected state $s_t$. We employ an identical architecture for the action diffusion model $\theta_a$, which restores the noise-injected action $a_t$ using the clean state $\hat{s}$. (b) Training of Noise timestep predictors. To estimate the appropriate noise timestep during inference, we train the state noise predictor $\psi_{s}$ given a clean action $\hat{a}$. The action noise predictor $\psi_{a}$ follows the identical architecture.
• Figure 4: Inference of demonstration restoration. Given a noisy state $s'$ and a clean action $\hat{a}$ from $D_{(s', \hat{a})}$, we first predict the noise timestep $t^{*}$ for the noisy state. If $t^{*}$ is less than a predefined threshold $t_{\text{thres}}$, then we append it to the clean subset $D_{(\hat{s}, \hat{a})}$ directly. Otherwise, we denoise the noisy state using the state diffusion model $\theta_{s}$, conditioned on the clean action and the predicted noise timestep. Similarly, noisy actions in $D_{(\hat{s}, a')}$ can be restored using the action noise predictor $\psi_{a}$ and the action diffusion model $\theta_{a}$.
• Figure 5: Environments & Tasks.(a)-(b) FetchPick and FetchPush: The robot arm manipulation tasks employ a 7-DoF Fetch robotics arm to pick up/push an object from the table and move it to a target location. (c) HandRotate: This dexterous manipulation task requires a Shadow Dexterous Hand to in-hand rotate a block to a target orientation. (d) Walker: These locomotion tasks require learning agents to walk as fast as possible while maintaining their balance.