Diffusion Policy: Visuomotor Policy Learning via Action Diffusion

Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, Shuran Song

★ Must read March 2023

Real Push-T on a UR5 (4x): push the T into the target, then move to the end zone.

TL;DR Represent a visuomotor policy as a denoising diffusion process that generates a short sequence of future actions, executed with receding-horizon control. It beats prior behavior-cloning methods by 46.9% on average across simulation benchmarks and reaches 95% success on a real Push-T task.

What they do

Demonstration data is multimodal, temporally correlated and needs precision, which strains standard behavior cloning. Diffusion Policy models the policy as a conditional denoising diffusion process over robot actions: starting from Gaussian noise, a noise-prediction network refines a sequence of future actions over K steps, conditioned on recent observations. Because it learns the score gradient instead of an energy, it avoids the negative sampling that makes implicit policies such as IBC unstable to train.

How it works

Receding-horizon action sequences. From the last To observation steps the policy predicts Tp actions and executes Ta of them before replanning; 8 executed steps worked best for most tasks.
Observations as conditioning. Images are encoded once by a ResNet-18 trained end-to-end, and only actions are denoised. With DDIM (10 inference steps) latency is 0.1 s on an Nvidia 3080.
Two denoisers. A 1D temporal CNN with FiLM conditioning works out of the box on most tasks; a time-series transformer does better when actions change quickly but needs more tuning.
Position control. Predicting positions consistently beat velocity commands for Diffusion Policy, while the baselines worked best with velocity.

Results

Simulation (Robomimic, Push-T, Block Pushing, Franka Kitchen): beats LSTM-GMM, IBC and BET on every task, by 46.9% on average; +213% on Kitchen’s p4 metric and +32% on Block Push’s p2.
Real Push-T on a UR5: 95% success and 0.80 IoU (human demos 0.84), against 20% for LSTM-GMM and 0% for IBC.
Mug flipping succeeds 90% of the time (LSTM-GMM failed to grasp in every trial). Sauce coverage is 0.74 for pouring and 0.77 for spreading, against 0.79 for humans.
Bimanual Franka tasks, 20 trials each: egg beater 55%, mat unrolling 75%, shirt folding 75%.

Limitations

It inherits behavior cloning’s dependence on adequate demonstrations. Inference costs more and is slower than simpler policies such as LSTM-GMM; predicting action sequences partly offsets this, but may not suffice for tasks that need high-rate control.

Videos

Mug flipping (4x): set a mug lip-down with the handle pointing left.

Sauce pouring onto pizza dough (4x), with the dough moved by hand mid-task.

Figures

From Figure 1: Policy representations. a) Explicit policy with different types of action representations. b) Implicit... — From Figure 1: Policy representations. a) Explicit policy with different types of action representations. b) Implicit policy learns an energy function over action and observation and picks actions that minimize it. c) Diffusion policy refines noise into actions via a learned gradient field.

From Figure 2: Diffusion Policy overview. a) At each time step the policy takes the latest To steps of observations a... — From Figure 2: Diffusion Policy overview. a) At each time step the policy takes the latest To steps of observations and outputs Ta steps of actions. b) CNN version: observation features condition every convolution layer through FiLM; starting from Gaussian noise, the predicted noise is subtracted K times to give the denoised action sequence. c) Transformer version: the observation embedding enters each decoder block through cross-attention, and action tokens use causal attention.

From Figure 3: Multimodal behavior in Push-T. The end-effector (blue) can go left or right to push the block. Diffusi... — From Figure 3: Multimodal behavior in Push-T. The end-effector (blue) can go left or right to push the block. Diffusion Policy learns both modes and commits to one per rollout; LSTM-GMM and IBC are biased toward one mode, and BET fails to commit to a single mode.

From Figure 8: Robustness tests on real Push-T. Left: a hand waved in front of the camera for 3 seconds causes slight... — From Figure 8: Robustness tests on real Push-T. Left: a hand waved in front of the camera for 3 seconds causes slight jitter, but the actions still work. Middle: the policy immediately corrects a block shifted during pushing. Right: when the block is shifted on the way to the end zone, the policy turns back and returns it to the goal, a behavior never demonstrated.

Abstract

This paper introduces Diffusion Policy, a new way of generating robot behavior by representing a robot's visuomotor policy as a conditional denoising diffusion process. We benchmark Diffusion Policy across 12 different tasks from 4 different robot manipulation benchmarks and find that it consistently outperforms existing state-of-the-art robot learning methods with an average improvement of 46.9%. Diffusion Policy learns the gradient of the action-distribution score function and iteratively optimizes with respect to this gradient field during inference via a series of stochastic Langevin dynamics steps. We find that the diffusion formulation yields powerful advantages when used for robot policies, including gracefully handling multimodal action distributions, being suitable for high-dimensional action spaces, and exhibiting impressive training stability. To fully unlock the potential of diffusion models for visuomotor policy learning on physical robots, this paper presents a set of key technical contributions including the incorporation of receding horizon control, visual conditioning, and the time-series diffusion transformer. We hope this work will help motivate a new generation of policy learning techniques that are able to leverage the powerful generative modeling capabilities of diffusion models. Code, data, and training details is publicly available diffusion-policy.cs.columbia.edu

Summaries are written automatically from each paper's text. Figures and videos come from the paper and its project page; all credit goes to the authors.

Matt Strong