[Paper Review] EquiBot: SIM(3)-Equivariant Diffusion Policy for Generalizable and Data Efficient Learning

Jingyun Yang, Ziang Cao+4 more|arXiv (Cornell University)|Jul 1, 2024

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

EquiBot proposes a SIM(3)-equivariant diffusion policy that achieves state-of-the-art data efficiency and generalization in robot manipulation by enforcing invariance to scale, rotation, and translation through a geometrically principled neural architecture. Trained on just 5 minutes of human demonstrations, it generalizes robustly to unseen objects, poses, and scenes in both simulation and real-world mobile manipulation tasks.

ABSTRACT

Building effective imitation learning methods that enable robots to learn from limited data and still generalize across diverse real-world environments is a long-standing problem in robot learning. We propose Equibot, a robust, data-efficient, and generalizable approach for robot manipulation task learning. Our approach combines SIM(3)-equivariant neural network architectures with diffusion models. This ensures that our learned policies are invariant to changes in scale, rotation, and translation, enhancing their applicability to unseen environments while retaining the benefits of diffusion-based policy learning such as multi-modality and robustness. We show on a suite of 6 simulation tasks that our proposed method reduces the data requirements and improves generalization to novel scenarios. In the real world, with 10 variations of 6 mobile manipulation tasks, we show that our method can easily generalize to novel objects and scenes after learning from just 5 minutes of human demonstrations in each task.

Motivation & Objective

• Address the challenge of data inefficiency and poor generalization in imitation learning for real-world robot manipulation.
• Overcome limitations of prior methods that rely on data augmentation or fail to generalize to unseen object scales, rotations, and translations.
• Develop a visuomotor policy that maintains robust performance across diverse, out-of-distribution scenarios with minimal demonstration data.
• Integrate equivariance with diffusion policy learning to enable multi-modal, closed-loop action generation with improved sample efficiency.

Proposed method

• Design a SIM(3)-equivariant neural network architecture that ensures model outputs transform consistently with input transformations (scale, rotation, translation).
• Integrate the equivariant architecture into a diffusion policy framework to enable multi-modal, robust action sequence prediction.
• Use single-view scene point clouds and robot proprioception as input, with object segmentation via Grounded SAM for scene parsing.
• Train the policy end-to-end using human demonstration videos, with data augmentation via geometric transformations to improve robustness.
• Apply a motion prediction head that outputs 8-step action sequences at 3 Hz, enabling real-time execution on mobile manipulators.
• Leverage equivariance to implicitly generalize to unseen object scales, orientations, and positions without explicit data augmentation.

Experimental results

Research questions

• RQ1 Can SIM(3)-equivariance in a diffusion policy architecture significantly improve data efficiency and zero-shot generalization in robot manipulation?
• RQ2 To what extent does equivariance enable generalization to unseen object scales, rotations, and translations beyond the training distribution?
• RQ3 How does the combination of equivariance and diffusion modeling compare to baseline imitation learning and non-equivariant diffusion policies in real-world deployment?
• RQ4 Can a single policy generalize across diverse household tasks with minimal demonstrations, including those involving deformable and articulated objects?
• RQ5 Does enforcing equivariance lead to better feature representations that generalize beyond the inductive bias of the symmetry group?

Key findings

• EquiBot achieves state-of-the-art performance in both simulation and real-world settings, outperforming baseline diffusion policies and prior equivariant imitation learning methods.
• In real-world experiments, EquiBot generalizes successfully to novel objects and scene layouts after only 5 minutes of human demonstration data.
• The method generalizes to out-of-distribution scenarios, including unseen object poses, scales, and scene configurations, with higher success rates than baselines.
• The model demonstrates robustness to geometric variations beyond SIM(3), such as non-uniform scaling and changes in object dynamics, despite not being explicitly trained for them.
• The combination of SIM(3)-equivariance and diffusion modeling enables multi-modal behavior generation and idle actions, improving policy robustness.
• Quantitative evaluation shows consistent success across 6 real-world mobile manipulation tasks, including pushing a chair, closing a suitcase, and folding towels, under diverse test conditions.