[Paper Review] 3D Diffuser Actor: Policy Diffusion with 3D Scene Representations

Tsung-Wei Ke, Nikolaos Gkanatsios+1 more|arXiv (Cornell University)|Feb 16, 2024

Reinforcement Learning in Robotics4 citations

TL;DR

3D Diffuser Actor introduces a diffusion-based robotic policy that leverages 3D scene representations and 3D relative attention to denoise end-effector pose predictions iteratively, achieving state-of-the-art performance on RLBench (18.1% absolute gain in multi-view, 13.1% in single-view) and CALVIN (7% relative improvement in zero-shot generalization), while generalizing effectively from few real-world demonstrations.

ABSTRACT

Diffusion policies are conditional diffusion models that learn robot action distributions conditioned on the robot and environment state. They have recently shown to outperform both deterministic and alternative action distribution learning formulations. 3D robot policies use 3D scene feature representations aggregated from a single or multiple camera views using sensed depth. They have shown to generalize better than their 2D counterparts across camera viewpoints. We unify these two lines of work and present 3D Diffuser Actor, a neural policy equipped with a novel 3D denoising transformer that fuses information from the 3D visual scene, a language instruction and proprioception to predict the noise in noised 3D robot pose trajectories. 3D Diffuser Actor sets a new state-of-the-art on RLBench with an absolute performance gain of 18.1% over the current SOTA on a multi-view setup and an absolute gain of 13.1% on a single-view setup. On the CALVIN benchmark, it improves over the current SOTA by a 9% relative increase. It also learns to control a robot manipulator in the real world from a handful of demonstrations. Through thorough comparisons with the current SOTA policies and ablations of our model, we show 3D Diffuser Actor's design choices dramatically outperform 2D representations, regression and classification objectives, absolute attentions, and holistic non-tokenized 3D scene embeddings.

Motivation & Objective

• To improve generalization and performance in multi-task 3D robot manipulation by combining diffusion policies with 3D scene representations.
• To address the limitations of 2D-based policies and deterministic action heads by modeling action distributions with multimodal, conditional diffusion.
• To enable zero-shot generalization across camera viewpoints and unseen environments using 3D scene feature learning and 3D relative attention.
• To demonstrate strong real-world transfer from a small number of human demonstrations using a 3D-aware, iterative denoising policy.
• To validate the architectural components—3D tokenization and 3D relative attention—through ablation studies on generalization and performance.

Proposed method

• The model encodes multi-view RGB-D observations into 2D image feature tokens using a vision encoder.
• 2D feature tokens are lifted into 3D space using depth information to form 3D scene tokens representing the environment and robot end-effector.
• A 3D relative position attention Transformer processes 3D scene tokens, language instruction embeddings, and noisy pose estimates to predict 3D translation and rotation errors.
• The policy uses iterative denoising: starting from a noisy pose, it iteratively refines the end-effector trajectory using noise prediction in a diffusion process.
• The model conditions on language instructions and maintains a short history of end-effector poses to support long-horizon, multi-step task execution.
• Pose predictions are sparsely generated (keyposes), enabling efficient inference and compatibility with motion planners like BiRRT.

Experimental results

Research questions

• RQ1 Can combining 3D scene representations with diffusion-based action modeling improve zero-shot generalization in multi-task 3D robot manipulation?
• RQ2 How does 3D relative attention improve spatial reasoning and equivariance in visuomotor policy learning compared to 2D or non-relative attention?
• RQ3 To what extent does 3D scene featurization enhance generalization across camera viewpoints and novel environments?
• RQ4 Can a diffusion policy conditioned on 3D scene representations generalize effectively from a small number of real-world demonstrations?
• RQ5 What architectural components—such as 3D tokenization or iterative denoising—most significantly contribute to performance and robustness?

Key findings

• 3D Diffuser Actor achieves a 18.1% absolute performance gain over the prior SOTA on RLBench in a multi-view setup, and a 13.1% gain in a single-view setup.
• On the CALVIN benchmark, it improves zero-shot unseen scene generalization by successfully completing 0.2 more tasks on average, a 7% relative improvement over the prior SOTA.
• The model generalizes effectively to real-world tasks, achieving high success rates on 12 diverse manipulation tasks using only 15 demonstrations per task.
• Ablation studies confirm that 3D scene featurization and 3D relative attention are critical for performance, with both components significantly improving generalization.
• The model runs in 600 ms per inference on a single GPU, demonstrating efficient inference due to sparse keypose prediction, and supports long-horizon tasks with an average of 3.27 successfully completed tasks in a chain.
• Stronger conditioning on language instructions significantly improves performance, especially in long-horizon multi-task settings, as shown in ablation studies.