3D Diffusion Policy: Generalizable Visuomotor Policy Learning via Simple 3D Representations

Yanjie Ze, Gu Zhang, Kangning Zhang, Chenyuan Hu, Muhan Wang, Huazhe Xu

Seminal March 2024

An Allegro hand autonomously wraps plasticine into a roll-up.

TL;DR DP3 conditions a diffusion policy on a 64-dim feature from a small MLP encoder over a sparse, colorless point cloud from a single camera. Across 72 simulation tasks it averages 74.4% success versus 59.8% for image-based Diffusion Policy, and it reaches 85% on 4 real tasks from 40 demonstrations each.

What they do

Image-based Diffusion Policy needs 100-200 demonstrations per real task. DP3 keeps its action backbone but replaces image input with a sparse point cloud from a single camera, encoded into a compact feature. It is tested on 72 simulated tasks from 7 domains and 4 real tasks, three with an Allegro hand.

How it works

Perception. An 84x84 depth image from one camera is converted to a point cloud, cropped to a bounding box to remove table and ground, and downsampled to 512 or 1024 points with farthest point sampling. Color is dropped.
DP3 Encoder. A three-layer MLP with LayerNorm, max-pooling and a projection head produces a 64-dim feature. It beats PointNet, PointNet++, PointNeXt and Point Transformer, including pre-trained versions.
Decision. A convolutional diffusion policy, conditioned on the point feature and robot pose, denoises action sequences with DDIM and sample prediction (100 steps in training, 10 at inference).

Results

Simulation (72 tasks, mostly 10 demos each): 74.4% average success vs 59.8% for Diffusion Policy, a 24.2% relative gain. DP3 typically converges within about 500 epochs.
On 6 ablation tasks, point clouds (78.3) beat image (40.7), depth (32.0), RGB-D (34.7) and voxel (32.3) inputs.
Real robot (40 demos, 10 trials per task): 85% average, vs 35% for image and 20% for depth Diffusion Policy. DP3 had a 0% safety-violation rate, vs 32.5% and 25.0% for the baselines.
Single-trial generalization tests: 4/5 unseen bowl positions, 5/5 cube colors, 5/5 new objects and 3/3 camera views; both baselines failed every position, color and view.

Limitations

The best 3D representation for control is still open, and very long-horizon tasks are untested. View generalization covers only slight camera changes and needs a manual point cloud transform. The safety assessment is qualitative, and generalization tests use one trial per condition.

Videos

A Franka gripper picks up a bowl and pours meat floss onto the plasticine.

The Allegro hand wraps plasticine, then pinches it into dumpling pleats.

Figures

From Figure 2: Overview of DP3. Perception and decision are trained end-to-end on expert demonstrations. A single-vie... — From Figure 2: Overview of DP3. Perception and decision are trained end-to-end on expert demonstrations. A single-view point cloud is encoded by a lightweight MLP into a compact 3D representation, and a diffusion backbone generates actions conditioned on it and the robot state.

From Figure 3: MetaWorld Reach with 5 demonstrations (red), evaluated 1000 times; blue points are successful targets.... — From Figure 3: MetaWorld Reach with 5 demonstrations (red), evaluated 1000 times; blue points are successful targets. DP3 succeeds across 3D space, Diffusion Policy and IBC only in part of it, and BC-RNN fails. Successes from left to right: 0 / 285 / 327 / 415.

From Figure 10: The four multi-stage real tasks. The Allegro hand makes a Dumpling, uses a Drill to touch a cube and ... — From Figure 10: The four multi-stage real tasks. The Allegro hand makes a Dumpling, uses a Drill to touch a cube and makes a Roll-Up; the gripper does Pour. Each row shows point clouds of a collected trajectory, stage by stage, ending in the end state.

Abstract

Imitation learning provides an efficient way to teach robots dexterous skills; however, learning complex skills robustly and generalizablely usually consumes large amounts of human demonstrations. To tackle this challenging problem, we present 3D Diffusion Policy (DP3), a novel visual imitation learning approach that incorporates the power of 3D visual representations into diffusion policies, a class of conditional action generative models. The core design of DP3 is the utilization of a compact 3D visual representation, extracted from sparse point clouds with an efficient point encoder. In our experiments involving 72 simulation tasks, DP3 successfully handles most tasks with just 10 demonstrations and surpasses baselines with a 24.2% relative improvement. In 4 real robot tasks, DP3 demonstrates precise control with a high success rate of 85%, given only 40 demonstrations of each task, and shows excellent generalization abilities in diverse aspects, including space, viewpoint, appearance, and instance. Interestingly, in real robot experiments, DP3 rarely violates safety requirements, in contrast to baseline methods which frequently do, necessitating human intervention. Our extensive evaluation highlights the critical importance of 3D representations in real-world robot learning. Videos, code, and data are available on https://3d-diffusion-policy.github.io .

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