EquiBot: SIM(3)-Equivariant Diffusion Policy for Generalizable and Data Efficient Learning
Authors: Jingyun Yang
Ziang Cao
Congyue Deng
Rika Antonova
Shuran Song
Jeannette Bohg
Venue: CoRL 2024 (2024)
Keywords: Imitation Learning
Equivariance
Data Efficiency
TL;DR: We propose a visuomotor policy learning method that by combining equivariance with diffusion policies is capable of generalizable and data-efficient policy learning in a wide range of robot 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.

EquiBot: SIM Diffusion Policy for Generalizable and Data Efficient Learning
Jingyun Yang, Zi-ang Cao", Congyue Deng, Rika Antonova, Shuran Song, Jeannette Bohg
Stanford University {j [jingyuny, ziangcao, congyue, rika. antonova, shuran, bohg edu
Generalizable
and Sample Efficient
Visuomotor Policy for Everyday Tasks
SUICESSS
EquiBot
10 demos
Diffusion
stucki
Policy
Figure I:
We propose method for learning generalizable and sample-efficient visuomotor policies that can be applied to everyday 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 long-standing problem in robot learning. We propose EquiBot, robust, data- efficient, and generalizable approach for robot manipulation task learning. Our approach combines -qquvvaaan neural network architectures with diffu- sion models. This ensures that our learned policies are invariant to changes in scale, rotation, and translation, enhancing their applicability to unseen environ- ments while retaining the benefits of diffusion-basee policy learning such as multi- modality and robustness. We show on suite of simulation tasks that our pro- posed method reduces the data requirements and improves generalization 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 af- ter learning from just minutes of human demonstrations in each task. Website: https://equi- bott gith....n
Keywords: Imitation Learning, Equivariance, Data Efficiency
Introduction
The unpredictability and variability real-world environments have posed significant challenges the development of fully autonomous robotic agents. Existing visuomotor policies learned via imitation learning are effective in controlled settings 2, 4, 5, 6, 7J, but they require substantial data and generalize poorly to unseen scenarios [8, 9, 10, 11, 12], limiting their practical deployment.
In this work, we introduce EquiBot, an equivariant policy learning architecture based on diffusion models [13]. The equivariant neural network guarantees that outputs scale, translate, and rotate with inputs, even if not fully trained [14], enabling generalization unseen scenarios with changes in
These authors contributed equally.
8th Conference Robot Learning (CoRL. 2024), Munich, Germany.

---

object visuals, sizes, and placements. This architecture also has data efficiency benefits, as it can infer actions for object placements and poses that are in distribution but insufficiently demonstrated, outperforming non-equivariant policies. Using a diffusion model-based policy architecture also of- fers robust learning performance, including support for idle actions and multi-modal behaviors [8]. Leveraging this design, EquiBot enables wide range of houschold manipulation tasks using only handful of single-view human demonstration videos. At test time, our method takes single-view, object-centric point clouds and robot proprioception as input and outputs sequences of robot end- effector actions that consist of 6D velocity commands and gripper open-close signals.
We evaluate our method in both simulation and real-world experiments. In simulation, we test on a suite of six tasks adapted from prior benchmarks [8, 15, 16], showing that our method is more data-efficient and generalizes better to unseen scenarios than prior works. In real robot experiments, quantitatively test on 6 real-world mobile manipulation tasks in home settings, including pushing a chair towards a desk, closing a laundry machine door, folding towels, making bed, and closing a suitcase. Our method successfully performs these tasks with unseen objects and scenes from just 5 minutes of human demonstrations, outperforming baselines that rely on augmentations for generalization or do not use a diffusion process to predict actions.
Related Work
Data-efficient imitation learning, Recent imitation learning approaches for imitation learning as- sume large quantities of data to be available for learning manipulation policies [4, 8]. To reduce the amount of data that policy training requires per task, some work [2, 17, 18, 19] formulate a one-shot or few-shot imitation learning setup where the policy is trained on multi-task demonstration data more demonstrations as well as the
and can then output actions in novel task after seeing current state. However, this approach requires the availability of multi-task data in a task domain. Some other works [20] achieve imitation learning from small demonstration datasets with sampling- based optimization methods like Bayesian Optimization, but these methods are often limited to small action spaces and open-loop settings. In contrast to prior works, we show that embedding equivari- ance into the policy architecture can effectively improve the data efficiency of the imitation learning algorithm, allowing robust policies be learned with just handful of demonstrations.
Equivariance in robot manipulation. To help learned robot policies generalize to unseen envi- ronments and object placements, prior works [21, 22, 23] use data augmentation to improve cross- domain transfer of the learned policy. However, these methods increase training time significantly and do not guarantee generalization to visual appearances, object scales, and poses that are unseen in the training data distribution even after the augmentation. In contrast, utilizing equivariant represen- tations in policy learning allows the learned policies to generalize objects and initial conditions not previously seen training time. Prior works have explored the use of equivariance in robot manipulation in several different setups [24, 25, 26, 27, 28, 29], but most of them either focus on only simple nick-and-plas place like tasks, do not support closed-loop policies, or do not support scale equivariance, In our prior work [16], we developed a SIM(3)-equivarian) visuomotor policy learning method that can go beyond pick-and-pe and-place tasks with deformable and articulated objects. However, EquivAct cannot handle multi-modal training data because of deterministic architecture. Equiv- Act also lacks the necessary architectural design that prevent it from predicting multiple actions into the future, which is crucial for producing temporally consistent actions [8]. Compared to prior works, our method trains stably, handles multi-modal data with diffuusion, and generalizes to unseen object appearance, initial states, scales, and poses with
Equivariant diffuuion architectures. Prior works have integrated equivariance in diffusion models in various non-robotics domains [30, 31, 32, 33]. Some works have attempted to integrate equivari- architectures in diffusion models for robotics [34, 35, 36]. Diffusion-EDDFs [34] produces target end-effector poses and can only solve pick-and-place tasks. EquiDiff [35] can only handle SO(2)- cquivariance with simple 2D trajectories. EDGI [36] assumes ground-truth scene states as input
2

---

and cannot handle complex visual observations. Compared to prior works, our work proposes an equivariant diffusion policy for closed-loop 3D manipulation tasks with point cloud observations.
3
Method
In this section, we describe the design of EquiBot. Our method builds upon recent advances that use a diffusion process to represent visuomotor policies [8]. Starting from that, we provide key insights for injecting equivariant architectures. By building an equivariant noise prediction network, we enforce each diffusion step to be equivariant by construction (see Figure 2), and because of the self- symmetry of the initial Gaussian noise, the overall framework outputs equivariant distribution under stochasticity. Section 3.2 provides our formal argument for this. The equivariant update in the diffusion process plays the role of letting the network "see" different input variations under transfor- mations, which replaces the need for data augmentations and makes our framework lata-efficient. Below, we introduce concepts related to equivariance and diffusion policy, then describe our method.
3.1 Preliminaries
Point cloud & rbbott poie
Training
Problem setup. We assume an imitation learn- ing setup, where our method receives a demon- Jr=) which con-
stration dataset sists of N demonstration trajectories r" Each demonstration trajectory consists of sequences of observation-action pairs Az). The goal of the policy is to learn a mapping æ from past observations Oz-T,at to either the next action A, or the next set of actions A1:+T,, where T. and T the observation and prediction horizons. At evaluation time, the policy re- ceives state O, and predicts the next one or more actions to be executed in the environ- ment. In this work, we assume the observation O, = (X,, S,) is composed the scene point cloud X, and robot proprioception S..
stepii-1
Model inputs and outputs translate, scale construction
difiusion step
ditlusion step
pasa
Deployment: Zero-shot Generalization
Figure 2: Method overview. Given input scene point cloud & robot pose, our method performs series of diffusion steps to obtain denoised actions with -qquvvaaiur, i.e. when the inputs translate, rotate, and scale, the outputs are guar- anteed to translate, rotate, and scale accordingly.
See())-qquvvariant network architectures. Let f be function that takes a point cloud X E IR X33 as input. This function is considered SiM(3)-equivarrian if f(TX) = TX(X) for any rigid 3D transformation T = (R, t, s) E SIM(3), where R, t, and s denote rotation, translation, and scale respectively. In this work, we use the same -equivaaau encoder architecture and network layers [16].
Diffusion process as policy representation. Our method uses Denoising Diffusion Probabilistic Models (DDPMs) to model the conditional distribution p(A, O2)) similar to [8]. Starting from Gaussian noise AK where K is the number of diffusion steps. DDPM performs K iterations of denoising to predict actions with decreasing levels noise, A, .. This process follows
AQ.
A A-1
(AA ThEE (O, A k) (0,
(1)
where is a denoising network, A/.0. I) is Gaussian noise, and O::;, 7k, are functions of k set by noise scheduler. The policy outputs A₁ as its inference output. In this work, we use the CNN-based Diffusion Policy variant specified [8] as the starting point of our architecture design. The original CNN-based Diffusion Policy architecture uses a noise prediction network that takes observation O diffusion iteration k. and noisy action A, as input, and predicts the gradient VE(A, for denoising A,. The network first uses an encoder encode the visual observations. The encoded visual features and positional embeddings of the diffusion iteration parameter passed into FiLM layers [37] so that the encoded visual inputs are integrated into the network. Then, the policy network uses

---

convolutional U-net [38] to process the imput noisy actions At, the conditioned observations, and diffusion iteration l to predict the output denoising gradients.
Observation and action spaces. We use point clouds as input observations, since these contain the necessary 3D information ensure policies can be structured as equivariant translation, rotation,
and scaling. We represent robot proprioception information with S. (S s s) with 3D positions in S normalized directions in s and scalars in S(() Robot proprioception can be converted into such format that uses positions, velocities, offsets, and scalars in most cases, e.g. end-effector positions go S end-effector velocities can be converted to position targets and go to s end-effector orientations can be converted into rotation matrices and placed in S gripper
open-close states go S(()) Similarly, our representation for actions A, (Al (0) A (4) A(")
consists of 3D offsets or velocities A, normalized directions A, and scalars S Similar to proprioception information, most existing action spaces can also be converted into this format.
3.2 Equivariant Distributions and Diffusion
We make a diffusion process equivariant by making the architecture for performing each diffusion step equivariant. In this section, we show that with equivariance in the per-step diffusion update, the final output action distribution is also equivariant under stochasticity. We mainly discuss equiv- ariance to rotatooo in the diffusion process. Equivariance translations and scaling (to get -eqiiaaaqnvaariant architectures) achieved via canonicalization before the diffusion process.
Proposition 1. Let p(x cc) be an SQl3)-equivariaat density function conditioned on c, i.e. VR. E 5O(3), p(*" c) = p(R.. Re). If the Markov transitions p(**=1|<** c) are SO(3)- equivariant for all k i.e. p(a*-110* c) = p(R=*~ Re), then the density p(x0lc ,)) also S0(3)-equivvaiimn,
Please view the complete proof in Section of the supplementary materials. In our case, obser- vation O, conditions the diffusion process of the actions A. The prior distribution p( A O,) is standard Gaussian distribution equivariant to SO(3) transformations. The transition probabili- ties p( p(A, 1-1 A O₂) are predictions by an equivariant network, and thus are equivariant to SO(3)- rotations, Therefore, the final output action p( O()))O)) is also S0(3)-cquivariant
3.3 Equiva-aqn Diffusion Policy
To design SIM(3)-equivariann model architecture, our approach is to modify each part of the CNN- based Diffusion Policy architecture [8] to make them individually equivariant architectures. First, we design our point cloud encoder to be equiiaaequvaaiian and additionally output a centroid vector Oe as well as a scalar O, quantifying object scale. O, and O are then used to scale the inputs to subsequent layers of the network so that they are invariant to positions and scales. We then modify the FiLM layers, the convolutional U-net architecture, and other connecting layers to be SO(3)-equivariant Finally, before producing the output actions, we scale relevant parts of the action back using +, and θ, the output is Siu(a)-equivariian to the input observation. Please see supplementary materials (Figure 13) for detailed method figure.
Encoder. We use a PointNet-based [39] encoder in this work. For SIM(3) -quivvaiiie we reuse the encoder & introduced in [40] and [16]. This encoder takes a point cloud X as input and outputs latent code θ (X), comprised of four components: O (O Qinv O.., O.), where On is rotation equivariant latent representation, Oiy is invariant latent representation, scalar O is the computed object scale, and vector O denotes the object centroid. For more details on this encoder, we refer to [16]. While [16] pre-trains the encoder using generated simulation data, we do not perform pre-training on the encoder and learn it from scratch. This eliminates the need to build task-specific simulation environments and collect custom pre-training data in these environments.
Routing input observations and actions into conditional U-net. The conditional U-net takes two inputs: action representation Z, and conditioning information Z.. To construct these inputs

---

from point cloud encoding O proprioception St, and noisy action A,, we need to first translate and scale S, and A, using θ and θ so the resulting values are invariant to scale and position, and then merge relevant inputs. More concretely, we define the action representation as
Jtuse (LA (0) ,0,, Al") (,A)
Zu
where fruse a Vector Neuron layer that takes vector information as its first and scalar information as its second input argument. We define conditioning information as
2. = (Z) Nectorr) (scalur) - (IOmms (s e,) [s]]] pos emb()mmb)))))
(3)
where Z and Zaaalar are vector and scalar conditioning used as input to the FiLM layers [37] in the conditional U-net, and pos_emb(k) is the positional embedding of the diffusion iteration k.
-equiaaquvaarian conditional U-net. A conditional U-net is composed of ID convolution layers, upsampling layers, and FiLM layers. We make this network -03)--qqivvriiat by converting every layer of this network to -033veequivariant layer.
To make ID convolution layers SO(3) -0uu)-qquivariant, we treat vector channels of the layer inputs as batch dimensions and perform the original convolution operations. This simple change makes the convolution layer SO(3)-equivariint. We do not make any modifications to the upsampling layer, as it is naturally -Q(3)-equivariant To make the FiLM layer -quivvarau,, we substitute vanilla linear layers with vectorized linear layers introduced in [14]. More formally, FiLM layer is for- mulated as FiLM((F n, B) nfF + B, where 7 f(x) and B h(x) are parameters predicted from learned functions used to modulate a neural network layer's activations F, and x is this neural network layer's 's input. We replace non-equivariant layers f and h with 'vector neuron' layers [14], achieving rotation equivariance.
Output. The conditional U-net with SOuuivequivarian layers processes Za and Ze and outputs translation and scale invariant actions Amyy- To process this value into the final output of the policy, we assemble the final output action as A. = (A)) iny O,, A inv
inv and Á (s) are the position, direction, and scalar components the predicted invariant action.
(d) A (d)), where A (==) A (d)
Please refer to Section E in supplementary materials for implementation details.
4 Experiments
Through experiments, want answer the following questions: (1) does method general- ize to unseen scenarios better than imitation learning methods that do not leverage equivariance; (2) does our method demonstrate more robust performance than prior methods for equivariant visuomo- tor policy learning that do not leverage diffusion models; (3) does our method achieve better data efficiency when there is limited training data; (4) how do different components of the method con- tribute to the final performance? We perform quantitative experiments in both simulations (Section 4.1) and the real world (Section 4.2) to answer these questions. Please view ablation experiments in Section B of the supplementary materials.
4.1 Simulation Experiments
4.1.1 Comparisons to Vanilla Diffusion Policies and Other Equivariant Policy Architectures
Below, we evaluate our method on nut-of-distribution generalization and compare to prior methods.
Comparisons. We compare our method to three baselines. (1) Diffusion Policy (DP) [8]: Vanilla dif- fusion policy using a point cloud as input; we substitute the imaged-base basdd encoder to a PointNet++ encoder [39] similar to what EquiBot using. (2) Diffusion Policy with Augmentations (DP+Aug): This baseline uses the same architecture as the vanilla diffusion policy baseline, but trains with syn- thetically generated data augmentation. (3) EquivAct [16]: A re-implementation of [16] that drops the pre-training phase that requires task-specific simulated data.
5

---

Cloth Folding
Object Covering
Box Closing
Push
Can
Square
Figure 3: Visualizations of simulation environments. The three mobile manipulation tasks feature varied rigid, deformable, and articulated objects. The Push T task features multi modal demonstra- tion data that challenge the learning algorithms. The Can and Square tasks from the Robomimic benchmark require precise position and orientation movements successfully complete the tasks.
Augmentations. The APgAug baseline requires augmentations in the training phase. In all four environments, we augment the training data to (1) rotate the observation around the z-axis, (2) uniformly scale the observation within the range 0.5 -1.5x, and (3) apply a random Gaussian offset to the observation with standard deviation equal to 0.1 times the approximate workspace size.
Tasks. We use four simulated tasks: Cloth Folding, Object Covering, Box Closing, and Push T (see Figure 3). The first three tasks in- volve two mobile robots manipulating various deformable and articulated objects. In these tasks, simulated depth camera records point clouds of relevant objects in the scene from third-person viewpoint. The policy takes these point clouds as input and commands the end- effector position, rotation, and gripper open- close actions of both robots. We train policies in these tasks with 50 synthetically generated demonstrations. The Push T' benchmark task is simulated 2D T-shape pushing game devel- oped in [8] to showcase learning from multi- modal demonstrations. To make this task com- patible with our setup, we assume the agent and object are placed on the ground plane (= = 0) in 3D space. In this task, the policy receives the eight comers of the T-shape as input and outputs the velocity command to the pushing operator. We use the same 200 demonstrations as [8] to train policies in this task.
Fodding
Object Covering
Crghe
000
Organ
18:5497
D2-4 Aug
DhAt:: EquiBet
Figure 4: Results of out-of-distribatio gener- alization experiments. We show that our method achieves more robust out-of-distribution general- ization performance than methods that do diffusion processes to model policies and ones that do utilize equivariance. Error bars show the mean and standard deviation over check- points and seeds.
Training and evaluation. We train all methods
for 2,000 epochs 3 random seeds. For every training run, we a checkpoint every 50 epochs and evaluate the last checkpoints saved at the end of training. For each evaluation, run the policy in randomly initialized environment for 10 episodes and record the mean final reward the policy achieves. This means that each bar in the resulting plot reports the mean and standard deviation of evaluation results over 3 seeds x5 last checkpoints x 10 episodes = 150 trials of evaluation.
Evaluation setups. To gain insight into the generalizability of competing methods, we design four different evaluation setups test policies. The Original setup evaluates the policy at the same initial poses and goals as in the demos; the OOD (R) setup randomizes initial object rotation around the z-axis; the OOD (S) setup scales the scene up 1x to 2x with up to a 1.33 aspect ratio change; the OOD (P) setup adds dramatic position randomization to the scene; the OOD (R+S+P) combines previous OOD setups by randomizing rotation, scale, and translation of the scene.
Results. We show the results of this experiment in Figure The DP baseline performs very well in the Original setup, but its performance drops significantly when it comes to any of the O0D setups. The DP +P+ baseline performs especially well in the 0OD (R) setup because the augmentations

---

data this baseline learns from are in-distribution with respect to the evaluation scenarios. But in other scenarios such as OOD (P) and OOD (R+S+P), suffers from significant performance drops. The EquivAct baseline performs very well in the Cloth Folding task but displays subpar performance in Object Covering and Box Closing tasks. also cannot perform well in the Push task because the deterministic network architecture with behavior cloning loss cannot handle multi-modal training data well. We also show in supplementary materials Section B that EquivAcr displays unstable training performance. Our method performs stably in all four tasks and suffers from the least amount of performance drop compared to all baselines. This shows that our method indeed outperforms prior methods in out-of-distribution generalization.
Con
Square
4.1.2 Data Efficiency Experiments
In this experiment, we aim to test if our method outperforms prior methods in a low- data regime, even if it is evaluated in dis- tribution. Because only care about in- distribution performance in this experiment, we only compare our method with the DP baseline. We adopt two Robomimic environments [15] for this experiment: Can and Square.
Denos
Denos
Demos Derros Derros
Figure 5: Results of data efficiency experi- ments. Our method achieves better data efficiency than the Diffusion Policy when evaluated in distri- bution on two benchmark tasks.
Setup. In each task, we train all methods 2.000 epochs in three setups: learning from 100 demos, 50 demos, and 25 demos. Following the same eval- uation standard for mobile dual-robot and Push 7 environments, we report the mean and standard deviation over 150 trials of evaluation.
Results. The results of this experiment can be found Figure In both tasks, the performance of DP drops dramatically when the number of demos decreases from 100 25, while our method retains relatively higher performance when the amount of data drops. This is because when the training data size is too small to cover the whole distribution of initial poses, the DP baseline can- not naturally generalize to unseen initial poses during evaluation. Our method, on the other hand, leverages the equivariance nature to cope with novel initial object poses at test time.
4.2 Real Robot Experiments
We show series of real robot experiments where train mobile robots perform everyday manipulation tasks from 5 minutes of single-view human demonstration videos. We select suite of 6 tasks that involve diverse everyday objects, including rigid, articulated, and deformable objects (see Figure 6): (1) Push Chair: A robot pushes chair towards a desk; (2) Luggage Packing: A robot picks up a pack clothes and places it in an open suitcase; (3) Luggage Closing: A robot closes an open suitcase on the floor; (4) Laundry Door Closing: robot pushes the door of laundry machine to close it: (5) Bimanual Folding: Two robots collaboratively fold a piece of cloth on couch; (6) Bimanual Make Bed: Two robots unfold a comforter make it cover the bed completely.
Data collection and robot setup. We collect 15 human demonstration videos for each real robot task. We use a ZED stereo camera to record the movement of a human operator using their fingers to manipulate the objects of interest at Hz. After data collection, we use an off-the-shelf hand detection model [41], an object segmentation model [42], and a proprietary learned stereo-to -eeeo-to-depth model to parse out the human hand poses and object point clouds in each frame of the collected demos. We then subsample this data to 3 Hz and convert it into a format supported by our policy training algorithm. In all real robot experiments, we use holonomic mobile bases [43] with Kinova Gen3 7 DoF arms mounted on top. Similar to human demonstration processing, we use a ZED2 camera and an object segmentation model to obtain the processed segmented point cloud as part of the input of our policy. See Section of supplementary materials for more details.
Training and evaluation. We train all methods for 1,000 epochs. After training, evaluate each method for 10 episodes and record the success rate of the method. We vary the evaluation
7

---

Chah
Luggage Packing
Luggage Closing
Laundry Door Closing
Binsrusllolding
Binonual Red
Figure 6: Real robot evaluation setups. Each block represents one task. In each block, we show sample training demos collected by human on the left and evaluation scenarios the right.
Push Chair
Laundry Door Closing
Luggage Closing
Unseen Prsees +
Object Variations
Long Desk
Round Table
Lange Luggage
DP DP+Aug
0/10
0/10
3/10
0/10
1/10
2/10
2/10
2/10
Luggage Closing
Bimanual Folding
Bimunual Make
Unseen Poses
Translate Rotate
Translate Rotate
Object Variations +
Tishirt
Towel Roll
Cap
Shorts
Long Bath Towel
Comforter
DP
0/10
0/10
0/10
0/10
Ours
7/10
3/10
8/10
8/10
6/10
8/10
Table 1: Results of real robot experiments.
suite of 6 mobile manipulation tasks, we show that our method can learn from just 5 minutes of human demonstration, outperforming the Diffusion Policy and the Diffusion Policy with Augmentation baselines by a large margin.
scenarios from the training scenarios differently in each task. In Laundry Door Closing, perform evaluations in-distribution. In Push Chair, Luggage Closing, and Bimanual Make Bed, we evaluate with out-of-distributinn objects. In Luggage Packing and Bimanual Folding, we not only switch to novel objects but also translate and rotate the layout of the scene.
Results. The results of real robot experiments are shown in Table 1. The evaluation shows that method can generalize diverse unseen objects, outperforming the DP baseline in both in- distribution and out-of-distribution scenarios with novel objects and unseen object poses. See sup- plementary materials Section C for more results and detailed analysis.
5 Conclusions, Limitations and Future work
We proposed EquiBot, visuomotor policy learning method for generalizable and data-efficient policy learning in wide range of robot manipulation tasks. In a suite of 6 simulated and 6 real robot tasks, we showed that our proposed method outperforms vanilla diffusion policies and prior imitation learning methods using equivariant architectures. We demonstrated that our method can learn from just 5 minutes of human demonstrations and generalize to unseen scenarios that are dramatically different from training scenarios.
While method generalizes scenes with unseen object positions, scales, and orientations, does not handle nonlinear changes in object shapes or dynamics by construction. Our method also does not handle variations in the relative positioning of objects when multiple objects are present. Resolving these limitations might involve explicitly modeling scene dynamics and individual ob- jects. Our method might also fail when the scene is partially occluded or the camera angle changes dramatically. This issue can be solved by learning gerrroe-purpose 3D representations robust to incomplete point clouds novel viewing angles. Please see more discussions on limitations in Section H of supplementary materials. Solving the above challenges in ization and extending to multi-task setups are interesting directions for future work.
general-

---

Acknowledgments
This work was supported in part by the Toyota Research Institute, Intrinsic and the Human-Centered Al Institute. The Stanford Robotics Center provided the space for our experiments.
We thank Jimmy Wu for helping with the hardware setup. We thank Xingze Dai and Yumeng Lu for helping with data collection. We thank Cheng Chi and Yilun Du for sharing insights about Diffusion Policy. We thank Yihuai Gao and Haoyu Xiong for helping with the fin ray gripper setup.
References
[1] Ross, G. Gordon, and D. Bagnell. reduction of imitation learning and structured predic- tion to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Pro- ceedings, 2011.
[2] Y. Duan, M. Andrychowiez, Stadie, O. Jonathan Ho, Schneider, I. Sutskever, P. Abbeel, and W. Zaremba. One-shot imitation learning. Advances in neural information processing systems, 30, 2017.
[3] Florence, C. Lynch, A. Zeng, O. A. Ramirez, A. Wahid, L. Downs, A. Wong, J. Lee, I. Mor- datch, and Tompson. Implicit behavioral cloning. In Conference Robot Learning, pages 158-168. PMLR, 2022.
[4] A. Brohan, N. Brown, J. Carbajal, Y. Chebotar, J. Dabis, C. Finn, K. Gopalakrishnan, K. Haus- man, A. Herzog, J. Hsu, et al. Rt-1: Robotics transformer for real-world control at scale. arXiv preprint 2rXiv:22/2.06811 2022.
[5] T. Z. Zhao, V. Kumar, S. Levine, and C. Finn. Learning fine mne-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:2304.13705, 2023.
[6] L. X. Shi, A. Sharma, T. Z. Zhao, and C. Finn. Waypoint-based imitation learning for robotic manipulation. In Conference on Robot Learning, pages 1199-2209. PMLR, 2023.
[7] A. Brohan, N. Brown, Carbajal, Chebotar, X. Chen, K. Choromanski, T. Ding, D. Driess,
A. Dubey, C. Finn, al. Rt-2: Vision-language-action models transfer web knowledge to robotic control. arXiv preprint 1rXiv:2307.15818, 2023.
[8] C. Chi, Feng, Y. Du. Z. Xu, E. Cousineau, Burchfiel, and S. Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
[9] A. Padalkar, A. Pooley, A. Jain, A. Bewley, A. Herzog, A. Irpan, A. Khazatsky, A. Rai, A. Singh, A. Brohan, al. Open x-embodiment: Robotic learning datasets and rt-x mod- els. arXiv preprint arXiv:2310.08864 2023.
[10] O. M. Team, D. Ghosh, H. Walke, K. Pertsch, K. Black, O. Mees, S. Dasari, Hejna, T. Kreiman, C. Xu, al. Octo: An open-source generalist robot policy. arXiv preprint drXiv:2405.12213 2024.
[11] A. Khazatsky, K. Pertsch, S. Nair, A. Balakrishna, S. Dasari, S. Karamcheti, S. Nasiriany, M. K. Srirama, L. Y. Chen, K. Ellis, al. Droid: A large-scale in-the-wild robot manipulation dataset. arXiv preprini arXiv:2403. 10000.11244 12945, 2024.
[12] M. Kim, K. Pertsch, S. Karamcheti, T. Xiao, A. Balakrishna, S. Nair, R. Rafailov, E. Foster, G. Lam, P. Sanketi, et al. Openvla: An open-source vision-language-action model. arXiv
preprint arXiv:2406.099 09246, 2024.
[13] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 13:514068811, 2020. 2020.

---

[14] C. Deng, O. Litany, Y. Duan, A. Poulenard, A. Tagliasacchi, and L. J. Guibas. Vector neu- rons: A general framework for so (3)-qqivvrriant networks. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 12200-12209,2000
[15] A. Mandlekar, D. Xu, J. Wong, S. Nasiriany, C. Wang. R. Kulkarni, L. Fei-Fei, S. Savarese, Y. Zhu, and R. Martin-Martin, What matters learning from offline human demonstrations for robot manipulation. In arXiv preprint htth::20080329966 2021.
[16] Yang, C. Deng, J. Wu, R. Antonova, L. Guibas, and Bohg. Equivact: Sim(3)-equivariant visuomotor policies beyond rigid object manipulation, 2023.
[17] C. Finn, T. Yu, T. Zhang, Abbeel, and S. Levine. One-shot visual imitation learning via meta-learning. In Conference robot learning, pages 357-368. PMLR, 2017.
[18] S. James, M. Bloesch, and A. J. Davison. Task-embedded control networks for few-shot imi- tation learning. In Conference on robot learning, pages 783-795, PMLR, 2018.
[19] Z. Mandi, F. Liu, K. Lee, and P. Abbeel. Towards more generalizable one-shot visual imitation learning. In 2022 International Conference on Robotics and Automation (ICRA), pages 2434- 2444. IEEE, 2022.
[20] Yang, J. Zhang, C. Settle, A. Rai, R. Antonova, and Bohg. Learning periodic tasks from human demonstrations. IEEE International Conference on Robotics and Automation (ICRA), 2022.
[21] S. James, P. Wohlhart, M. Kalakrishnan, Kalashnikov, Irpan, J. Ibarz, S. Levine, R. Hadsell, and K. Bousmalis. Sim-to-real via sim-to-sim: Data-efficient robotic grasping via cetommized-t-ccaoniie adaptation networks. In Proceedings the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 12627-12637, 2019.
[22] B. Mehta, M. Diaz, F. Golemo, C. J. Pal, and L. Paull. Active domain randomization. In Conference on Robot Learning, pages 1 1162-1176. PMLR, 2020.
[23] T. Yu, T. Xiao, A. Stone, J. Tompson, A. Brohan, S. Wang, J. Singh, C. Tan, J. Peralta, B. Ichter, et al. Scaling robot learning with semantically imagined experience. arXiv preprint arXiv:2302. 11550, 1550, 2023.
[24] A. Simeonov, Y. Du, A. Tagliasacchi, J. B. Tenenbaum, Rodriguez, P. Agrawal, and Sitz- mann. Neural descriptor fields: -3)equivari object representations for manipulation. In 2022 International Conference Robotics and Automation (ICRA), pages 6394-6400. IEEE, 2022.
[25] A. Simeonov, Y. Du, Y.-C. Lin, A. R. Garcia, L. P. Kaclbling, T. Lozano- Pérez, and P. Agrawal. Se (3)-cquivariant relational rearrangement with neural descriptor fields. In Conference Robot Learning, pages 835-846 PMLR, 2023.
[26] Z. Xue, Z. Yuan, J. Wang, X. Wang, Y. Gao. and H. Xu. Useek: Unsupervised se (3)- equivariant 3d keypoints for generalizable manipulation. 2023 IEEE International Con- ference on Robotics and Automation (ICRA), pages 1715-1722. IEEE, 2023.
[27] H. Ryu, H.-i. Lee, J. Lee, and J. Choi. Equivariant descriptor fields: (3)-equivariant energy-based models for end-to-end visual robotie manipulation learning. arXiv preprint wwXiv:226608832/ 2022.
[28] Т. Weng, D. Held, F. Meier, and M. Mukadam. Neural grasp distance fields for robot manip- ulation. In 2023 IEEE International Conference Robotics and Automation (ICRA), pages 1814-1821. IEEE, 2023.
10

---

[29] J. Seo, N. P. S. Prakash, X. Zhang, C. Wang, J. Choi, M. Tomizuka, and R. Horowitz. Robot manipulation task learning by leveraging se (3) group invariance and equivariance. arXiv preprint arXiv:2308. 14984, 2023.
[30] E. Hoogeboom, V. G. Satorras, C. Vignac, and M. Welling. Equivariant diffusion for molecule generation in 3d. In International conference on machine learning, pages 8867-8887. PMLR, 2022.
[31] A. Schneuing. Y. Du, C. Harris, A. Jamasb, Igashov, W. Du, Blundell, Lió, C. Gomes, M. Welling, et al. Structure-based drug design with equivariant diffusion models. arXiv preprint arXiv:2210. 13695, 2022.
[32] J. Guan, W. W. Qian, X. Peng, Y. Su, J. Peng, and J. Ma. 3d equivariant diffusion for target- aware molecule generation and affinity prediction. arXiv preprint arXiv:2303.0333 0303003343, 2023.
[33] I. Igashov, H. Stärk, C. Vignac, A. Schneuing, G. Satorras, P. Frossard, M. Welling, M. Bron- stein, and B. Correia. Equivariant 3d-conditional diffusion model for molecular linker design. Nature Machine Intelligence, pages 1-11, 2024.
[34] H. Ryu, J. Kim, J. Chang. H. S. Ahn, Seo, T. Kim, J. Choi, and Horowitz. Diffusion-edfs:: Bi-equivariant denoising generative modeling (3) for visual robotic manipulation. arXiv preprint arXiv:2309.02668 02685, 2023.
[35] K. Chen, X. Chen, Yu, M. Zhu, and Yang. Equidiff: A conditional equivariant diffusion model for trajectory prediction. In 2023 IEEE 26th International Conference on Intelligent Transportation Systems (ITSC), pages 746-751, IEEE, 2023.
[36] J. Brehmer, J. Bose, P. De Haan, and T. S. Cohen. Edgi: Equivariant diffusion for planning with embodied agents. Advances in Neural Information Processing Systems, 36, 2024.
[37] E. Perez, F. Strub, H. De Vries, V. Dumoulin, and A. Courville. Film: Visual reasoning with general conditioning layer. In Proceedings of the AAAI conference artificial intelligence, volume 32, 2018.
[38] O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical im- age segmentation. In Medical image computing and computer-assisted Enervention-MICAA 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234.241.5 Springer, 2015.
[39] C. R. Qi, L. Yi, H. Su, and L. J. Guibas. Pointnet++: Deep hicrarchical feature learning on point sets metric space. Advances neural information processing systems, 30, 2017.
[40] J. Lei, C. Deng, K. Schmeckpeper, L. Guibas, and K. Daniilidis. Efem: Equivariant neural field expectation maximization for 3d object segmentation without scene supervision. In Pro- ceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1902-4912. 2023.
[41] G. Pavlakos, D. Shan, I. Radosavovic, A. Kanazawa, D. Fouhey, and J. Malik. Reconstructing hands in 3D with transformers. In CVPR, 2024.
[42] A. Kirillov, E. Mintun, N. Ravi, Mao, Rolland, L. Gustafson, Xiao, S. Whitehead, A. C. Berg, W.-Y. Lo, P. Dollár, and R. Girshick. Segment anything. arXiv:2304.06643 2023.
[43] R. Holmberg and O. Khatib. Development and control of a holonomic mobile robot for mobile manipulation tasks. The International Journal of Robotics Research, 1961171006-1074, 2000.
[44] Y. Ze, G. Zhang, K. Zhang, C. Hu, M. Wang, and H. Xu. 3d diffusion policy: Generalizable visuomotor policy learning via simple 3d representations, In Proceedings of Robotics: Science and Systems (RSS), 2024.
11

---

[45] H. Huang, D. Wang. R. Walters, and R. Platt. Equivariant transporter network. arXiv preprint arXiv2200.009400, 2022.
[46] H. Huang, O. Howell, X. Zhu, D. Wang, R. Walters, and R. Platt. Fourier transporter: Bi- equivariant robotic manipulation in 3d. arXiv preprint arXiv:240.112046, 2024.
[47] D. Wang, R. Walters, and Platt. So())-equivarian reinforcement learning. arXiv preprint arXiv:2203.044: 2022.
[48] M. Jia, D. Wang, G. Su, D. Klee, X. Zhu, R. Walters, and R. Platt. Seil: Simulation-ann Simulation-augmented equivariant imitation learning. In 2023 IEEE International Conference on Robotics and Au- tomation (ICRA), pages 1845-1851. IEEE, 2023.
[49] J. Köhler, L. Klein, and F. Noé. Equivariant flows: exact likelihood generative learning for symmetric densities. In International conference on machine learming, pages 5361-5370. PMLR, 2020.
[50] M. Xu, L. Yu, Y. Song, C. Shi, S. Ermon, and J. Tang. Geodiff: A geometric diffusion model for molecular conformation generation. arXiv preprint arXiv:2203.02923. 2022.
[51] Song, C. Meng, and Ermon. Denoising diffusion implicit models. arXiv preprint arXiv:2010.02502, 2020.
[52] H. K. Cheng, S. W. Oh, B. Price, A. Schwing, and J.-Y. Lee. Tracking anything with decoupled video segmentation. In ICCV, 2023.
[53] T. Chen, S. Korbllith, M. Norouzi, and G. Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020.
12

---

A Further Discussion of Related Works
Imitation learning from 3D inputs. While many imitation learning works assume RGB images visual observations [15, 4, 20, 8], some works, similar the design choice of work, assume 3D point clouds or depth inputs their methods. Recently, Ze et al. [44] proposed 3D Diffusion Policy, depth-only variant of diffusion policy for visuomotor policy learning. Their method is very similar to our DP baseline, with two differences: (1) they use a simpler DP3 encoder in their work; (2) they use a 2-layer MLP to encode the robot proprioceptive states before concatenating with the point cloud representation.
In Figure 7, show quantitative comparison results between our method and baselines re- lated to [44]. As in our main paper results, we run 3 seeds of training runs, each for 2,000 epochs, for all experiments. We evaluate the last 5 checkpoints for each training run and col-
0.7
Fioal Rowarrd
lect the final task reward for 10 episodes in each evaluation. Comparisons show that [44] displays similar performance as the DP base- line in the main paper, performing very well in-distribution setups and poorly in out-of- distribution setups.
43
Orignal
OOD (K4S6)
COD 0145m!
OOD [R-SS
Figure 7: Comparisons with DP3-related architec- tures in the Cloth Folding task. We compare method with two DP3-related baselines and one varia- tion of our method that uses DP3 encoder: (1) DP3 a variation of the DP baseline with the PointNet- based encoder replaced by the DP3 encoder proposed in [44], with the code of DP3 encoder copied verbatim variant of are DP3 baseline trained with augmentations the same as the DP+ baseline in the main paper; (3) DP3 w/ Equivariance is the integration of DP3 into our method.
We also show that [44] can be easily integrated with method by switching our PointNet- based encoder to the DP3 encoder. In Figure 7, we show a comparison between our method and a variation of our method with a modifi- cation of the DP3 encoder to make it SO(3)- equivariant. Results show that the DP3 variant has slightly lower but comparable performance method in the Cloth Folding task.
from the public codebase; (2) DP3+Aug
Additional works in equivariant architectures for robot manip- ulation. Aside from prior equivariant architectures for robot ma- nipulation introduced in the main paper, there are also robotics works that attempt to utilize equivariance various setups. Some works [45, 46] attempt
Object Covering
0.9
equivariance in pick-and-place setup, while others [47, 48] propose SO(()-cquivarian robot poli- cies for tabletop manipulation tasks. In comparison to prior works, our proposed architecture is equivariant to position, orientation, and uniform scaling. In addition, our method can be applied in various 3D manipulation tasks that involve rigid, deformable, and articu- lated objects.
Rewerd
0.5
0.3
0.1
Original
000 (R+S+P)
EquivAct + Diffusior (T,=
EquivAct Diffattr (To=2,7,
B Ablations and Analysis
IT,
Equivict Diflusion (T7= -
EquieAct Diffiezer nTo7
EquiBot
Ablations on architecture designs. We perform more ablation experiments to understand how various aspects of our proposed method contribute to the final performance of the method. The first question we want to answer is how much the architecture of our method matters compared that of the prior work EquivAcr [16]. To show this, we provide a comparison of EquiBot with naive ex- tension of EquivAcr with diffusion head. As seen in Figure 8, EquiBot significantly outperforms this architecture for all hyperparameter variations tried in the OOD setups. When naively adding
Ablations model architecture. EquiBot outperforms naive extensions of EquivAct [16] for a variety of hyperparameter choices.
Figure
13

---

a diffusion head to EquivAct with action horizon 1, the policy's 's poor performance can be explained by its inability to generate multiple actions sequence to ensure temporal action consistency [8]. When we increase prediction horizons, the EquivAct architecture, which not designed to deal with large prediction horizons, suffers from dramatic drops in performance. Our method performs the best across all comparisons we performed in Figure 8.
Ablations on equivariance. To show how different forms of equivariance play role in the final performance of our method, we per- form ablation experiments the Object Cov- ering task by subtracting rotation, translation, and scaling equívariance implementations from method. The result of this ablation ex- periment is shown in Figure 9. As shown in the results, when we subtract out any form of equivariance from the implementation, the per- formance of this ablated method suffers from performance drops when the evaluation setup out-of-distribution the equivariance this method supports.
0.3-
01
Original
000
OODD
OOD
000
Scale Equlvariance
Figure 9: Ablations on equivariance object covering. We show that translation, rotation, and scale equivariance each contribute to our method.
Training stability. In the real world, it is impractical to test multiple training checkpoints on the robot. There- fore, high training stability can make checkpoint selection and getting robust test-time performance easy. Here, we specifically compare the training stability of our method with prior work EquivAct [16]. In Figure 10, plot the in-distribution performance of EquivAct against our method in the Push T task during the training process. We plot average reward over 40 evaluation episodes for every 50 epochs of training. Our method achieves more stable training performance across checkpoints than EquivAct.
Average Rewarrd
0.8
0.6
0.4
0.0
500
Esulwict
Ours
Figure 10: Training stability.
C
Further Evaluations and Performance Analysis
C.1 Failure Analysis
Although our method outperforms vanilla Diffusion Policy [8] and prior equivariant visuomotor policy architectures, our method still presents various failure cases. Below, we focus our analysis on execution failure our method the real robot experiments. Figure 11, we show the failure breakdown of all real robot executions have performed with our method.
In most packing tasks (packing t-shirt, towel, and cap), the main failure cases are the end-effector opening too early. We believe this is because the out-of-distribution scenarios resulted in the agent thinking that it has moved to the dropping location for the object and opened the end-effector. For the packing shorts task variation, half of the failures come from end-effector opening too late, and half of the failures come from the shorts not slipping off from the end-effector due the friction of the end-effector. This problem can potentially be solved by designing end-effectors that handle deformable objects better performing online adaptation after training, which is out of the scope of our work.
In the Push Chair, Laundry Door Closing, and Bimanual Folding tasks, the majority of failures come from the end-effector not performing the full motion or not performing gripper open close actions the right time. This most likely happens because the errors in predicted actions accumulated and the observation became too out-of-diitribution scenarios for the policy behave correctly. The Bimanual Make Bed task appears to be more difficult for our object segmentation method than other
14

---

Push Chair Long Desk
Push Chair to Circular Desk
Packing Tishirt
Packing Towel
Packing Cap
1
10
8
Packing Shorts
Laundry Clesing
Clesing
Bimanual Fold
Bimanual Make Bed
4
Success
End-effector opened too late
Deformable object is over-stretched Incorrect movernent based point cloud
Did perform the motion fully End-effeccoor opened early
Object did not fall off after end-effector
End-effector missed object during interactions
Figure 11: Failure cases our method during real robot executions. The pie charts show failure breakdown in every real robot task variation. The navy color denotes success, while other colors denote different types of failures. Each pic chart shoes total of 10 trials since we run 10 episodes evaluation.
Ours
DP
Number of Demos +
10
25
50
10
25
50
Success Rate
8/10
9/10
10/10
3/10
5/10
7/10
Close with Click Sound
2/10
2/10
5/10
2/10
1/10
5/10
Total Missing Angle
17.46
7.18%
6.3"
140.42
33.85
21.55
Collision Safety Issue
0/10
0/10
0/10
2/10
0/10
0/10
Table 2: Detailed performance of the laundry door closing task.
tasks, causing failures to segment the full comforter in some scenarios, since the folded comforter
looks like two pieces of cloth, one laid on top of another.
C.2 Additional Real Robot Results
Detailed performance analysis of the laundry door closing task. To understand the performance of the Diffusion Policy [8] and our method better, perform more detailed analysis in the laundry door closing task. In Table tion setup. Success Rate measures the percentage of evaluations that end within the success criteria we set; Close with Click Sound measures the percentage of episodes that end with the laundry door closed completely after making clicking sound; Total Accumulated Missing Angle measures the sum of the opening angles the laundry door the end of the 10 evaluation episodes; Collision or Safery Issue measures the percentage of evaluation runs that are terminated because of undesired collisions critical safety issues, such the robot arm getting stuck at the laundry door. From the evaluation results, we see that the DP policy not only has lower success rate but also has much larger accumulated missing angle. The baseline also suffers from many safety issues requir- ing episodes to be manually terminated by the robot operator. Our method has much fewer safety issues when executes.
report four different metrics of policy performance in each evalua-
Qualitative results. In Figure 12, we show qualitative rollout samples for all evaluation scenarios mentioned in the paper, plus one bonus task where two robots lift woven basket onto coffee table.
15

---

Figure 12: Qualitative rollout samples for real robot evaluation scenarios. From top to bottom, we have: (1) pushing a chair towards a long standing desk; (2) pushing a chair towards circular table; (3) packing t-shirts; (4) packing towel roll; (5) packing cap; (6) packing shorts: (7)
closing check-in luggage; (8) closing laundry door; (9) bimanual folding: (10) bimanual make bed;
(11) bonus task where two robots lift woven basket onto coffee table.
D Equivariant Distributions and Diffusion
In this section, we prove Proposition that introduced in Section 3.2 of the main paper.
Proposition 1. Let (c) be a equivaraua density function conditioned on e, i.e.
VR SO(3), p(x (c) p(Ra Rcc. If Markov transitions p1a* c) are S0(3)-
equivariant for all k. i.e. pa c) p(R===1 R Rc). then the density p(x%lc) p(xTT ( ( (,.,)))))))) also SO (qu3)--quivariant.
16

---

Proof. We say that distribution p(y) invariant to the SO(3) (0(3)-root) rotation group actions if:
p(y) p (Ry), VR SO(3).
(4)
We say that a conditional distribution p(y|=) is equivariant to SO(3) rotations if:
p(ylx) p(yy Ræ), vR E SO(3).
(5)
[49] shows that invariant distribution composed with equivariant invertible function results in an invariant distribution. Moreover, given Markov chain o0:K [50] shows that if the initial
distribution xx p(=") is invariant to group and the transition probabilities -²~1 ~ p(x*=11 [xx))
are equivariant at each time step to the same group, then the marginal distribution of == is also
invariant to the group actions at each time step. Specifically, p(z") is invariant:
/ pl (lak)ppe 1p(x 0.K |/0 "dxx" 1:K
p(x")
(6)
f pla" II pla* )d
K
= (Ro* II
p(Rx R
k-1 R.. )dx I:K
(8)
Rm
Now consider additional condition and equivariant initial distribution p(x with transitions p(a*-11 b e) follows:
1cx* PR. Rc), p(a"~110" pl p(R. R Rc).
(10)
The following shows that the marginal distribution p(z0lc) ()) is also equivariant:
p(z)lc) = / pl (c)p((00KK-1100 c)dx
(11)
= -f p(x p(a*|c) (c) II p(x p(a |.*t ccdd 1:K
K
(12)
K
/ "(Ra Re)
II p Ra Ro Rejdx 1:
(13)
k=1
p(Ro
(14)
E Method Architectures and Implementation Detail
In this section, we describe in detail the architecture of our method. We visualize the architecture of our model in Figure 13.
Observation and action spaces. In all simulated and real robot tasks except for Push T, use 13- dimensional proprioception information and 7-dimensional action space for each robot. The pro- prioception data for each robot consists of the following information: a 3-dimensional end-effector position, a 6-dimensional vector denoting end-effector orientation (represented by two columns of the end-effector rotation matrix), a 3-dimensional vector indicating the direction of gravity, and a scalar that represents the degree to which the gripper opened. The action space for each robot consists of the following information: 3-dimensional vector for the end-effector position veloc- ity, 3-dimensional vector for the end-effector angular velocity in axis-angle format, and scalar denoting the gripper action.
17

---

A A-1
VecLinear
inv
Point Cloud
ax+b
Z
(vector)
Z (calar))
s
(d)
VecConv1D
VecLinear
Proprioception
ax+b
pos_ emb(k)
Timestep
(Positional Embedding)
X
A
VecConv1D
Noisy Action
Z
Figure 13: Architecture of EquiBot. Given input point cloud, robot proprioception, noisy actions, and the diffusion timestamp, our architecture processes position, direction, and scalar in- formation independently, uses the encoder outputs scale position information into position and scale invariant values, and then routes them into an SOuui-equivariant conditional U-net to predict denoised actions. In the figure, omit scaling for the ease of viewing. VecLinear and VecConv1D refer to -euui-aqua version of linear and convolution ID layers.
In the Push T task, the robot proprioception is 3-dimensional and consists of the agent's 3D position in the scene, while the action space is 3-dimensional and denotes the absolute position target of the agent.
In all simulated and real robot tasks, our policy uses an observation horizon steps, a prediction horizon of 16 steps, and an action horizon of 8 steps. This is identical to the setup used in the diffusion policy paper [8].
Encoder architecture. In all tasks except for Push T, use SIM(3)-quivvriian version of PointNet++ with layers and hidden dimensionality 128. In the Push 7 task, we decrease the number of layers to since the number of points in the point cloud observation is much smaller in this task.
Noise prediction network. Our noise prediction network inherits hyperparameters from the original diffusion policy paper [8]. In all simulation experiments, perform 100 denoising steps during inference. In real robot experiments, optimize for inference speed, we use the DDIM scheduler [51] with 8 denoising steps.
the DDPM scheduler [13] and
Point cloud size. Picking the number of points to sample in the point cloud observation is a key hyperparameter to consider when designing an architecture that takes point cloud inputs. our experiments, we found that using 512 1024 points is sufficient for all tasks. In particular, for all real robot experiments and simulated mobile manipulation tasks, clouds. In Can and Square tasks, we use 256 and 512 points respectively since there is relatively more training data in these tasks, and decreasing the number of points in the point cloud makes training faster without hurting performance.
1024-point point
Normalization. Data normalization can be important the performance of diffusion models. Vanilla diffusion policy normalizes the observations and actions separately. We instead normal- ize all 3D-ved vector inputs (including observation and action) together due SlM()--qqivvariance
18

---

assumptions. Implementation-wisee we take subset of training data and compute the mean point cloud scale and mean action scale as Spe and 8sc- Then, all position- and velocity-related information divided by Spe/Sw at the start of the network forward pass and multiplied back before the output returned. This normalization factor ensures that the diffusion process always works with actions with values within the -1 to range. We normalize scalar information in the same way as the vanilla diffusion policy.
Real Robot Setup and Experiments - Further Details
F.1 Human Demo Parsing Infrastructure
We use a single ZED camera to record human natural motion in real-time, which is way more flexible and time-efficient than the expert demonstration from human teleoperation of a robot. With that, we create a set of human demonstrations D {T"} In=1- Each human demonstration consists of series of RGB-D image frames 7 ( ((T") ) where T is the episode horizon.
The human demonstration processing module has three parts:
an off-the-shelf object detection and tracking model X; 17 that takes the current and previous demonstration frames Ie and 1 input, and outputs parsed point cloud of objects of interest X
2. an off-the-shelf hand detection model Hü = ((T") that takes the current demonstration frame as input and outputs the keypoints on each human hand the input frame HY::
an alignment module !(X}, , that takes the outputs of the previous two steps as input, aligns the 3D coordinates of the outputs, and outputs the aligned human finger pose y? in the same coordinate system of X
We use Grounded Segment Anything Model with DEVA [52] as the object detection and tracking model v and HaMeR [41] as the hand detection model 4. In the alignment module, we find a set of matching points between the point cloud X and HF, and then fit a rotation transformation R, and an offset 4 to transform points H to the coordinate frame of the point cloud. Then, we extract the thumb and index finger positions the transformed keypoint set to predict the human "end- effector" pose y?. We found that the alignment module Q is crucial, as the hand detection model produces hand poses in a different coordinate frame from the point cloud. Without this module, the resulting hand poses will not align with the object point cloud.
F.2
Mobile Robot Control Infrastructure
Our robot control setup consists of a centralized workstation and two mobile robots. The workstation reads and parses visual observations, performs policy inference, and communicates with the robots. The mobile robots take the output actions of the policy and execute them. On the workstation side, build multi-proce infrastructure to handle observation parsing, policy inference, and action execution.
In the observation parsing process, we obtain visual observations from a single ZED 2 camera di- rectly connected to the workstation via cable. We use the Grounded Segment Anything Model with DEVA [52] to obtain segmented point clouds that contain only relevant objects in the scene. We then downsample this segmented point cloud to 1024 points. The downsampled point cloud and the robot proprioception information are sent to the policy inference process. The policy inference process then outputs sequence of 16 predicted actions.
In the action execution process, we first reset all robots to their initial poses. Then, for each step in an evaluation episode, we read out the latest policy inference results from the policy inference process. To ensure accurate execution, we compute the elapsed time between the policy input time and action execution time. If this elapsed time exceeds threshold, we skip the first few predicted actions during action execution. After skipping the first few predicted actions to account for latency,
19

---

we select the 8 actions that immediately follow the skipped actions to execute. This means that no matter how many actions are skipped, we always execute actions at a time.
On the mobile robot, we also build a multi multi-process control infrastructure to control the Kinova arm and mobile base. We utilize motion capture system obtain mobile base position. Then, we (1) transfer each control signal from the global world frame to the local frame of the Kinova arm, (2) convert the target pose at the gripper fingertip to an expected pose of the Kinova end-effector, and (3) convert it into velocity command for the arm. We use position control for the mobile base and only move it when the robot end-effector is too close to or too far from the base.
F.3 Task Details
Push chair. In this task, the human demonstrations are collected using a standing desk (48 x 30 inches). The policies are evaluated on two different tables: longer rectangular desk (58 x 23 inches) and circular table (diameter of 36 inches). An episode is considered successful if the center of the chair goes beneath the desk.
Luggage packing. In the human demonstrations, a human picks up pack of white t-shirts and places them into white carry-on luggage. At evaluation time, test four different packing items: white t-shirts (same training object), gray towel roll, blue cap, and navy shorts. An episode is considered successful if at least half of the packed object ends up within the luggage.
Luggage closing, In this task, human demonstrations are collected on a small carry-on luggage (55 40 23 cm), while the policies are evaluated on a large check-in luggage (76 x 48 x 25 cm). An episode is considered successful if the luggage ends up in closed state.
Laundry door closing, In this task, the human demonstrations and the robot work with the same laundry machine (front-loader). The goal is to close the door of the laundry machine that is open at the start of the episode. An episode is considered successful if the door ends up with an opening of at most 5cm.
Bimanual folding. In this task, the human demonstrations are collected by using two hands to fold small piece of cloth (34 x 38 cm). At evaluation time, the robot is asked to fold a large gray towel (140 x 75 cm). After each evaluation episode, we measure the mean distance between each grasped corner to their corresponding target cloth corners and mark the episode as successful if this mean distance is less than 0.2 times the length of the folding side of the cloth.
Bimanual make bed. In this task, the human demonstrations are collected by using two hands to unfold towel (34 38 em). At evaluation time, the robot is asked to make the bed by unfolding a much larger comforter on top of the bed. After each evaluation episode, we measure the mean distance between each grasped corner to the bed headboard and mark the episode as successful if this mean distance is less than 0.2 times the length of the bed.
G Simulation Tasks - Further Details
Cloth folding, In this task, the demonstrations show two robots folding a piece of cloth (27.5 x 27.5 cm). During an evaluation, we compute the task reward as 1.0 - (d1 + d2)/(0.275 x 2), where dj and d2 denote the distance from the two grasped cloth corners to the target cloth corners.
Object covering. In this task, the demonstrations show two robots moving a piece of cloth (27.5 x 27.5 cm) onto rigid box (10 x 7 x 5 cm). During an evaluation, the task reward is computed as Vintessea// Veeavex hull, where Vinerseca is the volume intersection between the box and the convex hull of the cloth and Veonve hull is the volume of the convex hull of the cloth.
Box closing. In this task, the demonstrations show two robots closing box (14.5 12 x 11.5 cm) with three flaps. Success in this task is evaluated (aj a2 + a3)/(3 x 180), where 0₁ to ag denote the angle in degrees at which each flap of the box is closed.
20

---

Push T. In this task, 2D anchor pushes a T-shaped object on a plane of dimension 512 x 512 pixels. The task reward is computed as the percentage of the shape that overlaps with the target T pose.
Robomimic tasks. We use the same object and reward specifications in these tasks as the original benchmark. Please check out the Robomimic [15] paper for more details.
H Limitations and Future Work Further Details
Although only enforced equivariance to tranfrrrrassooomattion by construction, in practice still observed generalization to geometric variations beyond SIM(3), such as non-uniform scaling. While this capability is not enabled by construction, our hypothesis is that ensuring equivariance in all policy layers is conducive to learning a more general feature representation. This may relate to the observations in prior feature-learning works, e.g. [53], which noted that by training the network to be (in their case) invariant to data transformations, they gained better features for downstream per- ception tasks. In future work, would be interesting to study the intermediate features of equivariant networks with systematic probing and evaluation techniques.
21


Reviews:
[{"review_id": "uACqTtBSZy", "reviewer": "Reviewer_8ciL", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 2, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "In this paper, the authors present a new SIM(3)-equivariant diffusion policy that by definition is invariant w.r.t. environment changes in scale, rotation and translation. The authors evaluate the method's efficiency across simulated and real robot experiments.", "summary_of_recommendation": "Overall I think the paper tackles an interesting and important problem. However, in its current form it lacks clarity in multiple locations, the authors need to better highlight the difference w.r.t. the very related work [17], better stress the contributions of this work, and better explain what they mean with training stability.", "technical_quality": 3, "clarity_of_presentation": 2, "potential_impact": 2, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "- Please explain in more detail how this work is different from [17]\n- What exactly do you mean in unstable training performance of [17] w.r.t. to this work?\n- How exactly are the actions encoded in the representation and how is the delta in the actions received as output and used to update the actions?\n- Why are end effector positions integrated and translated into a set of points instead of using the velocity representation?\n- For the iterative denoising process, does the computation time scale linear with the number of denoising iterations? Also, how does this compare with the computational requirements with [17]?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "vbOso2cCo5", "reviewer": "Reviewer_U9p1", "summary": "", "questions": "", "limitations": "", "rating": 2, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 2, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This paper propses EquiBot, an extension of EquiAct to close-loop visuomotor control tasks. By combining SO(3) equivariant models with diffusion the authors show they can solve several robot manipulation tasks in simulation and in the real world.", "summary_of_recommendation": "This paper proposes a extension of a previous work (EquiAct) to close-loop control and is therefor somewhat incremental. The real world experiments are very impressive but the baseline results in simulation seem worse than previously reported.", "technical_quality": 3, "clarity_of_presentation": 3, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "1. Can the authors provide more detail on the simulation experiments to explain the poor performance of diffusion when compared to that in the original work?\n2. Additional evaluation episodes for simulation experiments would strengthen the claims. \n3. How many demonstrations were used in the simulation experiments?\n4. Is our understanding of this method being an extension of EquiAct with a diffusion policy head correct?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "58ZVxa2lL4", "reviewer": "Reviewer_vrP9", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": 2, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "This work proposes a SIM(3) (translation, rotation, scale) equivariant diffusion policy. The translation and scale equivariance are achieved by canonicalization. The rotation equivariance is based on type1 irreps (vector). First the VN point net++ extract type1 irreps, then a modified FiLM conditioned temporal Unet achieves SO(3) equivariance, due to 3 copies of standard Unet are used to process the XYZ- components of the type1 irreps independently. Processing XYZ- components of type1 irreps is also a form of canonicalization.  The method achieves strong generalization to novel SIM(3) transformed scenes as well as sample efficiency for a specific class of tasks. However, due to the limitation of canonicalization, the method do not generalize to broader tasks that involving muti-objects.", "summary_of_recommendation": "This work leverages SIM(3) equivariance in the diffusion policy by using cannonicalization. However, the method is limited to simple tasks, i.e., when there is only one target object in the scene.", "technical_quality": 3, "clarity_of_presentation": 4, "potential_impact": 3, "robotics_focus": 4, "quality_of_limitations": 3, "questions_for_rebuttal": "Does the proposed method outperform the baseline (Diffusion Policy) on broader class of tasks? I.e., coffee preparation in the MimicGen that includes multiple objects and long horizon reasoning?\n\nCould author provide ablation experiments that show amount the SIM(3) (rotation, translation, and scale) equivariance, which equivariance is most critical?\n\nCould author provide experiments similar to “out-of-distribution generalization experiments” that test the model on rotation, translation, and scale out-of-distribution individually?\n\nSection. Equivariant diffusion architectures in the Appendix seems the most related works, could you put it in the main text?", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: The paper introduces a SIM(3)-equivariant diffusion policy, which aims to maintain invariance to changes in scale, rotation, and translation in robotic environments. The reviewers acknowledge the importance and potential impact of equivariant policies for enhancing data efficiency in imitation learning. However, concerns are raised about the clarity of the paper, particularly in relation to prior work (EquiAct) and the detailed methodology, which requires further explanation and citations. Reviewers also noted the necessity for additional experiments, including ablation studies and broader task evaluations, to fully validate the approach. Despite these issues, the real-world experiments and the inclusion of scale invariance are appreciated.
During the rebuttal period, although there was not much discussion between the authors and the reviewers, there was a discussion amongst the reviewers. All agreed that the rebuttal had lead to a better draft of the paper, and most concerns were addressed. There were however concerns regarding the performance gap of diffusion policy in pushT between this paper and the original diffusion policy paper.
Decision: Accept