Diffusion-Reward Adversarial Imitation Learning

TL;DR

Diffusion-Reward Adversarial Imitation Learning (DRAIL), which integrates a diffusion model into GAIL, aiming to yield more robust and smoother rewards for policy learning, is proposed.

Abstract

Imitation learning aims to learn a policy from observing expert demonstrations without access to reward signals from environments. Generative adversarial imitation learning (GAIL) formulates imitation learning as adversarial learning, employing a generator policy learning to imitate expert behaviors and discriminator learning to distinguish the expert demonstrations from agent trajectories. Despite its encouraging results, GAIL training is often brittle and unstable. Inspired by the recent dominance of diffusion models in generative modeling, we propose Diffusion-Reward Adversarial Imitation Learning (DRAIL), which integrates a diffusion model into GAIL, aiming to yield more robust and smoother rewards for policy learning. Specifically, we propose a diffusion discriminative classifier to construct an enhanced discriminator, and design diffusion rewards based on the classifier's output for policy learning. Extensive experiments are conducted in navigation, manipulation, and locomotion, verifying DRAIL's effectiveness compared to prior imitation learning methods. Moreover, additional experimental results demonstrate the generalizability and data efficiency of DRAIL. Visualized learned reward functions of GAIL and DRAIL suggest that DRAIL can produce more robust and smoother rewards. Project page: https://nturobotlearninglab.github.io/DRAIL/

Figures (10)

• Figure 1: Denoising diffusion probabilistic model. Latent variables $x_1, ..., x_N$ are produced from the data point $x_0$ via a forward diffusion process, i.e., gradually adding noises to the latent variables. A 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-Reward Adversarial Imitation Learning. Our proposed framework DRAIL incorporates a diffusion model into GAIL. (a) Our proposed diffusion discriminative classifier $D_\phi$ learns to distinguish expert data $(\mathbf{s}_E, \mathbf{a}_E) \sim \tau_E$ from agent data $(\mathbf{s}_\pi, \mathbf{a}_\pi) \sim {\tau_i}$ using a diffusion model. $D_\phi$ is trained to predict a value closer to $1$ when the input state-action pairs are sampled from expert demonstration and predict a value closer to $0$ otherwise. (b) The policy $\pi_\theta$ learns to maximize the diffusion reward $r_\phi$ computed based on the output of $D_\phi$ that takes the state-action pairs from the policy as input. The closer the policy resembles expert behaviors, the higher the rewards it can obtain.
• Figure 3: Environments & tasks.(a) Maze: A point-mass agent (green) within a 2D maze is trained to move from its initial position to reach the goal (red). (b) FetchPush: The manipulation task is implemented with a 7-DoF Fetch robotics arm. FetchPush requires picking up or pushing an object to a target location (red). (c) HandRotate: For this dexterous manipulation task, a Shadow Dexterous Hand is employed to in-hand rotate a block to achieve a target orientation. (d) AntReach: This task trains a quadruped ant to reach a goal randomly positioned along the perimeter of a half-circle with a radius of 5 m. (e) Walker: This locomotion task requires training a bipedal walker policy to achieve the highest possible walking speed while maintaining balance. (f) CarRacing This image-based racing game task requires driving a car to navigate a track as quickly as possible.
• Figure 4: Learning efficiency. We report success rates (Maze, FetchPush, HandRotate, AntReach) and return (Walker, CarRacing), evaluated over five random seeds. Our method DRAIL learns more stably, faster, and achieves higher or competitive performance compared to the best-performing baseline in all the tasks.
• Figure 5: Generalization experiments in FetchPush. We present the performance of our proposed DRAIL and baselines in the FetchPush task, under varying levels of noise in initial states and goal locations. The evaluation spans three random seeds, and the training curve illustrates the success rate dynamics.