Diffusion Model-Augmented Behavioral Cloning

TL;DR

The paper addresses offline imitation learning by integrating conditional behavioral cloning with a diffusion-model-based joint-distribution signal. It introduces Diffusion Model-Augmented Behavioral Cloning (DBC), which trains a diffusion model on expert state-action pairs and jointly optimizes a BC loss $L_{BC}$ and a diffusion-model loss $L_{DM}$ to balance inference efficiency with generalization. Empirical results across navigation, manipulation, and locomotion tasks show DBC achieving state-of-the-art or competitive performance, with ablations confirming the complementary roles of BC and diffusion guidance and the importance of normalization. Overall, the work demonstrates that combining conditional and joint distribution modeling via diffusion models can yield robust, data-efficient policies for complex, multimodal tasks in settings that do not allow environment interaction.

Abstract

Imitation learning addresses the challenge of learning by observing an expert's demonstrations without access to reward signals from environments. Most existing imitation learning methods that do not require interacting with environments either model the expert distribution as the conditional probability p(a|s) (e.g., behavioral cloning, BC) or the joint probability p(s, a). Despite the simplicity of modeling the conditional probability with BC, it usually struggles with generalization. While modeling the joint probability can improve generalization performance, the inference procedure is often time-consuming, and the model can suffer from manifold overfitting. This work proposes an imitation learning framework that benefits from modeling both the conditional and joint probability of the expert distribution. Our proposed Diffusion Model-Augmented Behavioral Cloning (DBC) employs a diffusion model trained to model expert behaviors and learns a policy to optimize both the BC loss (conditional) and our proposed diffusion model loss (joint). DBC outperforms baselines in various continuous control tasks in navigation, robot arm manipulation, dexterous manipulation, and locomotion. We design additional experiments to verify the limitations of modeling either the conditional probability or the joint probability of the expert distribution, as well as compare different generative models. Ablation studies justify the effectiveness of our design choices.

Figures (13)

• Figure 1: Denoising Diffusion Probabilistic Model (DDPM). Latent variables $x_1, ..., x_N$ are produced from the data point $x_0$ via the forward diffusion process, i.e., gradually adding noises to the latent variables. The diffusion model $\phi$ learns to reverse the diffusion process by denoising the noisy data to reconstruct the original data point $x_0$.
• Figure 2: Diffusion Model-Augmented Behavioral Cloning (DBC ). Our proposed framework augments behavioral cloning (BC) by employing a diffusion model. (a) Learning a Diffusion Model: the diffusion model $\phi$ learns to model the distribution of concatenated state-action pairs sampled from the demonstration dataset $D$. It learns to reverse the diffusion process (i.e., denoise) by optimizing $\mathcal{L}_\text{diff}$ in Eq. \ref{['eq:diff_loss']}. (b) Learning a Policy with the Learned Diffusion Model: we propose a diffusion model objective $\mathcal{L}_{\text{DM}}$ for policy learning and jointly optimize it with the BC objective $\mathcal{L}_{\text{BC}}$. Specifically, $\mathcal{L}_{\text{DM}}$ is computed based on processing a sampled state-action pair $(s, a)$ and a state-action pair $(s, \hat{a})$ with the action $\hat{a}$ predicted by the policy $\pi$ with $\mathcal{L}_\text{diff}$.
• Figure 3: Environments & Tasks.(a) Maze: A point-mass agent (green) in a 2D maze learns to navigate from its start location to a goal location (red). (b) FetchPick: The robot arm manipulation tasks employ a 7-DoF Fetch robotics arm to pick up an object (yellow cube) from the table and move it to a target location (red). (c) HandRotate: This dexterous manipulation task requires a Shadow Dexterous Hand to in-hand rotate a block to a target orientation. (d)-(e) Cheetah and Walker: These locomotion tasks require learning agents to walk as fast as possible while maintaining their balance. (f) AntReach: This task combines locomotion and navigation, instructing an ant robot with four legs to reach a goal location while maintaining balance.
• Figure 4: Manifold overfitting Experiments. (a) We collect the green spiral trajectories from a script policy, whose actions are visualized as red crosses. (b) We train and evaluate $\pi_{BC}$, $\pi_{DM}$ and $\pi_{DBC}$ using the demonstrations from the script policy. The trajectories of $\pi_{BC}$ (orange) and $\pi_{DBC}$ (red) can closely follow the expert trajectories (green), while the trajectories of $\pi_{DM}$ (blue) deviates from expert's. This is because the diffusion model struggles at modeling such expert action distribution with a lower intrinsic dimension, which can be observed from incorrectly predicted actions (blue dots) produced by the diffusion model.
• Figure 5: CarRacing. This task features controlling a car navigating a track with image-based observations.