Equivariant Diffusion Policy

Authors: Dian Wang, Stephen Hart, David Surovik, Tarik Kelestemur, Haojie Huang, Haibo Zhao, Mark Yeatman, Jiuguang Wang, Robin Walters, Robert Platt · CoRL 2024

Licences

Open Access
Paper: CC BY 4.0

Everything below draws on the full paper under its CC BY 4.0 licence; its figure is credited to the authors.

TL;DR

Equivariant Diffusion Policy builds SO(2) rotation symmetry into Diffusion Policy's denoiser, raising average simulation success by 21.9 points at 100 demos and learning real tasks from 20-60 demos.

Why it matters

Real demonstrations are expensive, and building in tabletop rotation symmetry let the policy succeed on six real tasks, including a long-horizon bagel-baking task, with 20 to 60 demonstrations where a voxel Diffusion Policy baseline mostly failed. Because the analysis covers full SE(3) gripper actions, it applies to tasks that need out-of-plane rotation, which an SE(2) variant could not solve.

Problem

Diffusion Policy captures multimodal demonstration behaviour, but it must learn a denoising function over observations, noisy actions and noise levels, which is harder than learning a direct state-to-action policy and needs more data. Earlier equivariant manipulation policies used either open-loop SE(3) or closed-loop SE(2) action spaces, and none applied symmetry to diffusion policies.

Proposed solution

The authors prove that the noise-prediction function is SO(2)-equivariant whenever the expert policy is, and derive how rotation about the gravity axis acts on full 6-DoF gripper actions under both absolute and relative pose control. The denoiser uses equivariant observation and action encoders (built with the escnn library) that produce embeddings for a discrete rotation group, a shared 1D temporal U-Net applied per group element, and an equivariant noise decoder. Inputs can be RGB images or a 64×64×64 voxel grid.

Key insights

If the expert policy is SO(2)-equivariant, so is the diffusion noise-prediction function, which makes equivariant networks a correct inductive bias for the denoiser.
Gains are largest with few demonstrations and on tasks with widely randomised initial object poses; on low-randomness Robomimic tasks the method was similar to or slightly better than Diffusion Policy.
In ablations the equivariant structure mattered more than voxel input: at 100 demos, removing equivariance cut average success by 17.6 points and removing voxels by 10.3.
Rotation data augmentation on a standard CNN helped on simple tasks such as Stack but fell well short of the equivariant network on harder ones.
A full SE(3) action space matters: an SE(2) variant scored 0.0% on Coffee Preparation versus 85.3% for SE(3), and 12.7% versus 40.0% on Threading.

Main results

Across 12 MimicGen tasks with absolute control and 100 demos, the voxel version averaged 63.9% success versus 42.0% for Diffusion Policy (DiffPo-C), a 21.9-point gain; the RGB version averaged 53.7%.
With 200 demos the voxel version (72.6%) beat every baseline trained with 1000 demos (best 71.4%).
With absolute control the voxel version outperformed the baselines in 11 of 12 environments, the exception being Hammer Cleanup D1.
On a real Franka arm, success versus a voxel Diffusion Policy baseline was: Oven Opening 95% vs 60% (20 demos); Banana in Bowl 95% vs 30%, Letter Alignment 95% vs 0%, Trash Sweeping 90% vs 5% (40 demos each); Hammer to Drawer 85% vs 5% (60 demos); Bagel Baking 80% vs 10% (58 demos).
All Bagel Baking failures came from robot joint limits, and the policy generalised zero-shot to unseen oven rotations except one pose blocked by joint limits.

Limitations

The authors note equivariance is only partly exploited because of symmetry mismatch in the vision system; the arm appearing in the voxel grid and camera noise can break symmetry.
"Incorrect equivariance" can hurt performance when the model's symmetry conflicts with the demonstrations.
The SO(2) action theory is shown only with diffusion policies, not with other pipelines such as BC RNN, which performed well with relative control.
The method is not yet extended to navigation, locomotion or mobile manipulation.
With RGB input the agent-view camera is not top-down, so image rotation does not match true state rotation; with relative control the RGB version was only marginally better than the baselines.

Field impact

Finalist for the CoRL 2024 Outstanding Paper Award and cited 107 times. The paper points to a project page at equidiff.github.io. It does not name a code repository in its text.

Abstract

Recent work has shown diffusion models are an effective approach to learning the multimodal distributions arising from demonstration data in behavior cloning. However, a drawback of this approach is the need to learn a denoising function, which is significantly more complex than learning an explicit policy. In this work, we propose Equivariant Diffusion Policy, a novel diffusion policy learning method that leverages domain symmetries to obtain better sample efficiency and generalization in the denoising function. We theoretically analyze the SO(2) symmetry of full 6-DoF control and characterize when a diffusion model is SO(2)-equivariant. We furthermore evaluate the method empirically on a set of 12 simulation tasks in MimicGen, and show that it obtains a success rate that is, on average, 21.9% higher than the baseline Diffusion Policy. We also evaluate the method on a real-world system to show that effective policies can be learned with relatively few training samples, whereas the baseline Diffusion Policy cannot.

Open Source Robotics