Diffusion Policies

A diffusion policy is a robot control policy that generates actions the way a diffusion model generates images: it starts from random noise and removes the noise step by step, conditioned on what the robot currently observes, until a sequence of motor commands remains. Trained on human demonstrations, it is a form of behavioral cloning that can represent several distinct ways of doing a task instead of averaging them into one. Cheng Chi, Shuran Song, and colleagues at Columbia University, the Toyota Research Institute, and MIT introduced it at RSS 2023 in Diffusion Policy: Visuomotor Policy Learning via Action Diffusion, outperforming earlier imitation learning methods by an average of 46.9% across 12 tasks from 4 robotics manipulation benchmarks.

Why Regression Fails on Multimodal Demonstrations#

Human demonstrations often contain more than one valid action for the same observation. A person pushing a block may approach it from the left in one demonstration and from the right in the next, and a driver facing an obstacle may pass it on either side. A policy trained by regression with a mean squared error loss function predicts the average of those choices, which can be an action no demonstrator took, such as driving straight into the obstacle. Earlier fixes discretized actions into bins, fit Gaussian mixtures, or used energy-based models such as Implicit Behavioral Cloning, which needs negative samples during training and was often unstable. A diffusion model can represent any distribution of actions, and its noise-prediction loss trains as stably as ordinary regression.

This PyTorch example trains a tiny diffusion policy on demonstrations that steer left or right around an obstacle in equal numbers:

import torch
from torch import nn

torch.manual_seed(0)
# Demonstrations at one observation: facing an obstacle, experts steer left (-1) or right (+1) equally often
actions = torch.cat([torch.randn(500, 1) * 0.05 - 1, torch.randn(500, 1) * 0.05 + 1])
print(f"Regression policy (mean squared error) outputs: {actions.mean():.2f}")

# Diffusion policy: train a network to predict the noise mixed into an action at denoising step k
K, betas = 50, torch.linspace(1e-4, 0.2, 50)
alpha_bar = torch.cumprod(1 - betas, 0)
net = nn.Sequential(nn.Linear(2, 64), nn.SiLU(), nn.Linear(64, 64), nn.SiLU(), nn.Linear(64, 1))
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)
for _ in range(3000):
    k, noise = torch.randint(0, K, (1000, 1)), torch.randn(1000, 1)
    noisy = alpha_bar[k].sqrt() * actions + (1 - alpha_bar[k]).sqrt() * noise
    loss = (net(torch.cat([noisy, k / K], 1)) - noise).pow(2).mean()
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# Sampling: start from pure noise and remove the predicted noise one step at a time
a = torch.randn(1000, 1)
with torch.no_grad():
    for k in reversed(range(K)):
        eps = net(torch.cat([a, torch.full_like(a, k / K)], 1))
        a = (a - betas[k] / (1 - alpha_bar[k]).sqrt() * eps) / (1 - betas[k]).sqrt()
        a += betas[k].sqrt() * torch.randn_like(a) if k > 0 else 0
print(
    f"Diffusion policy samples: {(a < -0.5).float().mean():.0%} steer left, {(a > 0.5).float().mean():.0%} steer right"
)

The regression policy outputs 0.00, straight into the obstacle. The diffusion policy's samples split into 44% steering left and 55% steering right, with only about 1% between the two modes, so each sample commits to one valid maneuver.

How Diffusion Policies Work#

Training takes an action sequence from a demonstration, mixes in Gaussian noise at a randomly chosen step out of 100, and trains a neural network to predict that noise from the noisy actions and the observation. At run time, the robot samples a sequence of pure noise and denoises it while conditioning on its latest observation. The original paper used 10 DDIM denoising steps on its real robots, giving an inference latency of 0.1 seconds on an NVIDIA RTX 3080 GPU.

The policy plans in chunks and replans in a receding-horizon loop. In most of the paper's tasks, it observes the last 2 time steps, predicts the next 16 actions, executes the first 8, and then denoises a new sequence from fresh observations, which keeps motion smooth while still reacting to change:

The authors built two noise-prediction networks: a 1D temporal convolutional neural network that injects observation features through FiLM conditioning, which they recommend as the starting point, and a transformer for tasks whose actions change quickly. Camera images pass through a ResNet-18 trained from scratch, with spatial softmax pooling in place of global average pooling and group normalization in place of batch normalization. Predicting target positions instead of velocities consistently worked better and made the policy more robust to latency.

Policies can also observe a low-dimensional state instead of pixels. In the Push-T benchmark, where a robot pushes a T-shaped block onto a target, the state-based variant observes 9 2D keypoints of the block plus the end effector's position. An Ultralytics YOLO26 pose estimation model trained with a custom keypoint layout, set by kpt_shape in the dataset YAML, returns such keypoints from camera frames through result.keypoints.xy.

Faster and Larger Successors#

Denoising takes several network passes per action chunk, so much follow-up work targets real-time inference. Consistency Policy, by Aaditya Prasad and colleagues in 2024, distills a diffusion policy into a consistency model that needs one or a few steps, speeding up inference by an order of magnitude. 3D Diffusion Policy (DP3), by Yanjie Ze and colleagues in 2024, conditions on compact point clouds instead of images, handling most of its 72 simulated tasks with only 10 demonstrations and reaching 85% success on 4 real robot tasks with 40 demonstrations each.

The same recipe now drives robot foundation models. Physical Intelligence's π0 (2024) attaches an action expert trained with flow matching, a close relative of diffusion, to a 3-billion-parameter vision-language model and generates chunks of 50 actions to control robots at up to 50 Hz. Hugging Face's open-source LeRobot library implements Diffusion Policy alongside other imitation learning policies for embodied AI research.

Diffusion Policies vs. Related Terms#

• Action chunking: Predicting a sequence of future actions instead of one. Diffusion policies generate their chunks by denoising, and the Action Chunking with Transformers (ACT) policy generates them with a conditional variational autoencoder.
• Diffusion forcing: Gives every token of a sequence its own noise level so one causal model can roll out step by step and plan toward goals. On a robot task that required memory, it succeeded 80% of the time where the memoryless Diffusion Policy failed.
• Diffusion models: The general family of generative models that learn to reverse a gradual noising process, used for image, video, and audio generation. A diffusion policy is a diffusion model over actions, conditioned on observations.