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

Authors: jingyun  yang
zi-ang  cao
congyue  deng
Venue: robot-learning_org_CoRL_2024_Conference
Published: 2024-06-07 06:52:23

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.

Main paper: 
 introduction: 
 The unpredictability and variability of real-world environments have posed significant challenges to the development of fully autonomous robotic agents. Existing visuomotor policies learned via imitation learning are effective in controlled settings [1, 2, 3, 4, 5, 6, 7] , 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 to unseen scenarios with changes in 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 offers robust learning performance, including support for idle actions and multi-modal behaviors [8] . Leveraging this design, EquiBot enables a wide range of household manipulation tasks using only a 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 endeffector 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, we 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 a 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 assume 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 and can then output actions in a novel task after seeing one or more demonstrations as well as the 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 samplingbased 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 equivariance into the policy architecture can effectively improve the data efficiency of the imitation learning algorithm, allowing robust policies to be learned with just a handful of demonstrations.Equivariance in robot manipulation. To help learned robot policies generalize to unseen environments and object placements, prior works [21, 22, 23] use data augmentation to improve crossdomain 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 representations in policy learning allows the learned policies to generalize to objects and initial conditions not previously seen at 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 pick-and-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)-equivariant visuomotor policy learning method that can go beyond pick-and-place tasks with deformable and articulated objects. However, EquivAct cannot handle multi-modal training data because of its 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 diffusion, and generalizes to unseen object appearance, initial states, scales, and poses with SIM(3)-equivariance.Prior works have integrated equivariance in diffusion models in various non-robotics domains [30, 31, 32, 33] . Some works have attempted to integrate equivariant architectures in diffusion models for robotics [34, 35, 36] . Diffusion-EDFs [34] produces target end-effector poses and can only solve pick-and-place tasks. EquiDiff [35] can only handle SO(2)equivariance with simple 2D trajectories. EDGI [36] assumes ground-truth scene states as input 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. 
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 selfsymmetry of the initial Gaussian noise, the overall framework outputs an 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 transformations, which replaces the need for data augmentations and makes our framework data-efficient. Below, we introduce concepts related to equivariance and diffusion policy, then describe our method.Model inputs and outputs translate, rotate, and scale by construction Problem setup. We assume an imitation learning setup, where our method receives a demonstration dataset D = {τ n } N n=1 , which consists of N demonstration trajectories τ n . Each demonstration trajectory consists of sequences of observation-action pairs (O t , A t ). The goal of the policy is to learn a mapping π from past observations O t-To:t to either the next action A t or the next set of actions A t:t+Tp , where T o and T p are the observation and prediction horizons. At evaluation time, the policy receives state O t and predicts the next one or more actions to be executed in the environment. In this work, we assume the observation O t = (X t , S t ) is composed of the scene point cloud X t and robot proprioception S t .Let f be a function that takes a point cloud X ∈ R N ×3 as input. This function is considered SIM(3)-equivariant if f (TX) = Tf (X) for any rigid 3D transformation T := (R, t, s) ∈ SIM(3), where R, t, and s denote rotation, translation, and scale respectively. In this work, we use the same SIM(3)-equivariant encoder architecture and network layers as [16] .Diffusion process as policy representation. Our method uses Denoising Diffusion Probabilistic Models (DDPMs) to model the conditional distribution p(A t |O t ) similar to [8] . Starting from Gaussian noise A K t , where K is the number of diffusion steps, DDPM performs K iterations of denoising to predict actions with decreasing levels of noise, A K-1 t , . . . , A 0 t . This process followsEQUATIONwhere ϵ θ is a denoising network, N (0, I) is Gaussian noise, and α k , γ k , σ k are functions of k set by a noise scheduler. The policy outputs A 0 t as its inference output. In this work, we use the CNN-based Diffusion Policy variant specified in [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 t , diffusion iteration k, and noisy action A t as input, and predicts the gradient ∇E(A t ) for denoising A t . The network first uses an encoder to encode the visual observations. The encoded visual features and positional embeddings of the diffusion iteration parameter are passed into FiLM layers [37] so that the encoded visual inputs are integrated into the network. Then, the policy network uses a convolutional U-net [38] to process the input noisy actions A t , the conditioned observations, and diffusion iteration k to predict the output denoising gradients.Observation and action spaces. We use point clouds as input observations, since these contain the necessary 3D information to ensure policies can be structured as equivariant to translation, rotation, and scaling. We represent robot proprioception information with S t = (S t . Similar to proprioception information, most existing action spaces can also be converted into this format.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 equivariance to SO(3)-rotations in the diffusion process. Equivariance to translations and scaling (to get SIM(3)-equivariant architectures) is achieved via canonicalization before the diffusion process.Proposition 1. Let p(x K |c) be an SO(3)-equivariant density function conditioned on c, i.e. ∀R ∈ SO(3), p(x K |c) = p(Rx K |Rc). If the Markov transitions p(x k-1 |x k , c) are SO(3)- equivariant for all k, i.e. p(x k-1 |x k , c) = p(Rx k-1 |Rx k , Rc), then the density p(x 0 |c) = p(x T |c)Π K k=1 p(x k-1 |x k , c) is also SO(3)-equivariant.Please view the complete proof in Section D of the supplementary materials. In our case, observation O t conditions the diffusion process of the actions A 0:K t . The prior distribution p(A K t |O t ) is a standard Gaussian distribution equivariant to SO(3) transformations. The transition probabilities p(A k-1 t |A k t , O t ) are predictions by an equivariant network, and thus are equivariant to SO(3)rotations. Therefore, the final output action p(A 0 t |O t ) is also SO(3)-equivariant.To design a SIM(3)-equivariant model architecture, our approach is to modify each part of the CNNbased Diffusion Policy architecture [8] to make them individually equivariant architectures. First, we design our point cloud encoder to be SIM(3)-equivariant and additionally output a centroid vector Θ c as well as a scalar Θ s quantifying object scale. Θ c and Θ s 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 Θ c and Θ s so the output is SIM(3)-equivariant to the input observation. Please see supplementary materials (Figure 13 ) for a detailed method figure.Encoder. We use a PointNet-based [39] encoder in this work. For SIM(3)-equivariance, we reuse the encoder Φ introduced in [40] and [16] . This encoder takes a point cloud X as input and outputs a latent code Θ = Φ(X), comprised of four components:Θ := (Θ R , Θ inv , Θ c , Θ s ), where Θ R is a rotation equivariant latent representation, Θ inv is an invariant latent representation, scalar Θ s is the computed object scale, and vector Θ c 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 a conditional U-net. The conditional U-net takes two inputs: action representation Z a and conditioning information Z c . To construct these inputs from point cloud encoding Θ, proprioception S t , and noisy action A t , we need to first translate and scale S t and A t using Θ c and Θ s so the resulting values are invariant to scale and position, and then merge relevant inputs. More concretely, we define the action representation asEQUATIONwhere f fuse is a Vector Neuron layer that takes vector information as its first and scalar information as its second input argument. We define conditioning information asEQUATIONt , pos emb(k) ,where Z vector c and Z scalar c 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.A conditional U-net is composed of 1D convolution layers, upsampling layers, and FiLM layers. We make this network SO(3)-equivariant by converting every layer of this network to an SO(3)-equivariant layer.To make 1D convolution layers SO(3)-equivariant, 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)-equivariant. We do not make any modifications to the upsampling layer, as it is naturally SO(3)-equivariant. To make the FiLM layer SO(3)-equivariant, we substitute vanilla linear layers with vectorized linear layers introduced in [14] . More formally, a FiLM layer is formulated as FiLM(F|η, β) = ηF + β, where η = f (x) and β = 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 input. We replace non-equivariant layers f and h with 'vector neuron' layers [14] , achieving rotation equivariance.Output. The conditional U-net with SO(3)-equivariant layers processes Z a and Z c and outputs translation and scale invariant actions Âinv . To process this value into the final output of the policy, we assemble the final output action as Ât = ( Please refer to Section E in supplementary materials for implementation details.EQUATION 
experiments: 
 Through our experiments, we want to answer the following questions: (1) does our method generalize 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 visuomotor 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 contribute 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.Below, we evaluate our method on out-of-distribution generalization and compare to prior methods.Comparisons. We compare our method to three baselines. (1) Diffusion Policy (DP) [8] : Vanilla diffusion policy using a point cloud as input; we substitute the imaged-based encoder to a PointNet++ encoder [39] similar to what EquiBot is using. (2) Diffusion Policy with Augmentations (DP+Aug):This baseline uses the same architecture as the vanilla diffusion policy baseline, but trains with synthetically generated data augmentation. (3) EquivAct [16] : A re-implementation of [16] that drops the pre-training phase that requires task-specific simulated data. Tasks. We use four simulated tasks: Cloth Folding, Object Covering, Box Closing, and Push T (see Figure 3 ). The first three tasks involve two mobile robots manipulating various deformable and articulated objects. In these tasks, a simulated depth camera records point clouds of relevant objects in the scene from a third-person viewpoint. The policy takes these point clouds as input and commands the endeffector position, rotation, and gripper openclose actions of both robots. We train policies in these tasks with 50 synthetically generated demonstrations. The Push T benchmark task is a simulated 2D T-shape pushing game developed in [8] to showcase learning from multimodal demonstrations. To make this task compatible with our setup, we assume the agent and object are placed on the ground plane (z = 0) in 3D space. In this task, the policy receives the eight corners 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.Training and evaluation. We train all methods for 2,000 epochs on 3 random seeds. For every training run, we save a checkpoint every 50 epochs and evaluate the last 5 checkpoints saved at the end of training. For each evaluation, we run the policy in a 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 ×5 last checkpoints ×10 episodes = 150 trials of evaluation.Evaluation setups. To gain insight into the generalizability of competing methods, we design four different evaluation setups to test our 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 1× to 2× 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 4 . The DP baseline performs very well in the Original setup, but its performance drops significantly when it comes to any of the OOD setups. The DP+Aug baseline performs especially well in the OOD (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), it 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. It also cannot perform well in the Push T 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 EquivAct 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.Figure 5 : Results of data efficiency experiments. Our method achieves better data efficiency than the Diffusion Policy when evaluated in distribution on two benchmark tasks.In this experiment, we aim to test if our method outperforms prior methods in a lowdata regime, even if it is evaluated in distribution. Because we only care about indistribution 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.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 evaluation standard for mobile dual-robot and Push T environments, we report the mean and standard deviation over 150 trials of evaluation.Results. The results of this experiment can be found in Figure 5 . In both tasks, the performance of DP drops dramatically when the number of demos decreases from 100 to 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 cannot 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.We show a series of real robot experiments where we train mobile robots to perform everyday manipulation tasks from 5 minutes of single-view human demonstration videos. We select a suite of 6 tasks that involve diverse everyday objects, including rigid, articulated, and deformable objects (see Figure 6 ): Data collection and robot setup. We collect 15 human demonstration videos for each real robot task. We use a ZED 2 stereo camera to record the movement of a human operator using their fingers to manipulate the objects of interest at 15 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-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 F of supplementary materials for more details.Training and evaluation. We train all methods for 1,000 epochs. After training, we evaluate each method for 10 episodes and record the success rate of the method. We vary the evaluation scenarios from the training scenarios differently in each task. In Laundry Door Closing, we perform evaluations in-distribution. In Push Chair, Luggage Closing, and Bimanual Make Bed, we evaluate with out-of-distribution 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 our method can generalize to diverse unseen objects, outperforming the DP baseline in both indistribution and out-of-distribution scenarios with novel objects and unseen object poses. See supplementary materials Section C for more results and detailed analysis. 
conclusions, limitations and future work: 
 We proposed EquiBot, a visuomotor policy learning method for generalizable and data-efficient policy learning in a 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 our method generalizes to scenes with unseen object positions, scales, and orientations, it 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 objects. Our method might also fail when the scene is partially occluded or the camera angle changes dramatically. This issue can be solved by learning general-purpose 3D representations robust to incomplete point clouds or novel viewing angles. Please see more discussions on limitations in Section H of supplementary materials. Solving the above challenges in out-of-distribution generalization and extending to multi-task setups are interesting directions for future work. 
 
 Appendix: Imitation learning from 3D inputs. While many imitation learning works assume RGB images as visual observations [15, 4, 20, 8] , some works, similar to the design choice of our work, assume 3D point clouds or depth inputs to their methods. Recently, Ze et al. [44] proposed 3D Diffusion Policy, a 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 , we show quantitative comparison results between our method and baselines related 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 collect the final task reward for 10 episodes in each evaluation. Comparisons show that [44] displays similar performance as the DP baseline in the main paper, performing very well in in-distribution setups and poorly in out-ofdistribution setups.We also show that [44] can be easily integrated with our method by switching our PointNetbased encoder to the DP3 encoder. In Figure 7 , we show a comparison between our method and a variation of our method with a modification of the DP3 encoder to make it SO(3)equivariant. Results show that the DP3 variant has slightly lower but comparable performance as our method in the Cloth Folding task.Figure 8 : Ablations on model architecture. EquiBot outperforms naive extensions of EquivAct [16] for a variety of hyperparameter choices.Additional works in equivariant architectures for robot manipulation. Aside from prior equivariant architectures for robot manipulation introduced in the main paper, there are also robotics works that attempt to utilize equivariance in various setups. Some works [45, 46] attempt to use equivariance in a pick-and-place setup, while others [47, 48] propose SO(2)-equivariant robot policies 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 articulated objects.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 to that of the prior work EquivAct [16] .To show this, we provide a comparison of EquiBot with a naive extension of EquivAct with a diffusion head. As seen in Figure 8 , EquiBot significantly outperforms this architecture for all hyperparameter variations we tried in the OOD setups. When naively adding a diffusion head to EquivAct with action horizon 1, the policy's poor performance can be explained by its inability to generate multiple actions in a sequence to ensure temporal action consistency [8] . When we increase prediction horizons, the EquivAct architecture, which is 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 .Figure 9 : Ablations on equivariance in object covering. We show that translation, rotation, and scale equivariance each contribute to our method.Ablations on equivariance. To show how different forms of equivariance play a role in the final performance of our method, we perform ablation experiments in the Object Covering task by subtracting rotation, translation, and scaling equivariance implementations from our method. The result of this ablation experiment is shown in Figure 9 . As shown in the results, when we subtract out any form of equivariance from the implementation, the performance of this ablated method suffers from performance drops when the evaluation setup is out-of-distribution to the equivariance this method supports. Training stability. In the real world, it is impractical to test multiple training checkpoints on the robot. Therefore, 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 , we 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.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 of our method in the real robot experiments. In Figure 11 , we show the failure breakdown of all real robot executions we 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 to the friction of the end-effector. This problem can potentially be solved by designing end-effectors that can handle deformable objects better or 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 at the right time. This most likely happens because the errors in predicted actions accumulated and the observation became too out-of-distribution scenarios for the policy to behave correctly. The Bimanual Make Bed task appears to be more difficult for our object segmentation method than other 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.Detailed performance analysis of the laundry door closing task. To understand the performance of the Diffusion Policy [8] and our method better, we perform a more detailed analysis in the laundry door closing task. In Table 2 , we report four different metrics of policy performance in each evaluation 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 a clicking sound; Total Accumulated Missing Angle measures the sum of the opening angles of the laundry door at the end of the 10 evaluation episodes; Collision or Safety Issue measures the percentage of evaluation runs that are terminated because of undesired collisions or critical safety issues, such as the robot arm getting stuck at the laundry door. From the evaluation results, we see that the DP policy not only has a lower success rate but also has a much larger accumulated missing angle. The baseline also suffers from many safety issues requiring episodes to be manually terminated by the robot operator. Our method has much fewer safety issues when it executes.Qualitative results. In Figure 12 , we show qualitative rollout samples for all evaluation scenarios we mentioned in the paper, plus one bonus task where two robots lift a woven basket onto a coffee table. In this section, we prove Proposition 1 that we introduced in Section 3.2 of the main paper. Proof. We say that a distribution p(y) invariant to the SO(3)-rotation group actions if:EQUATIONWe say that a conditional distribution p(y|x) is equivariant to SO(3) rotations if: p(y|x) = p(Ry|Rx), ∀R ∈ SO(3).(5) [49] shows that an invariant distribution composed with an equivariant invertible function results in an invariant distribution. Moreover, given a Markov chain x 0:K , [50] shows that if the initial distribution x K ∼ p(x K ) is invariant to a group and the transition probabilities x k-1 ∼ p(x k-1 |x k ) are equivariant at each time step to the same group, then the marginal distribution of x k-1 is also invariant to the group actions at each time step. Specifically, p(x 0 ) is invariant:EQUATIONNow consider an additional condition c, and equivariant initial distribution p(x K |c) with transitions p(x k-1 |x k , c) as follows:EQUATIONThe following shows that the marginal distribution p(x 0 |c) is also equivariant:EQUATIONIn 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, we use a 13dimensional proprioception information and a 7-dimensional action space for each robot. The proprioception 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 is opened. The action space for each robot consists of the following information: a 3-dimensional vector for the end-effector position velocity, a 3-dimensional vector for the end-effector angular velocity in axis-angle format, and a scalar denoting the gripper action. Figure 13 : Architecture of EquiBot. Given input scene point cloud, robot proprioception, noisy actions, and the diffusion timestamp, our architecture processes position, direction, and scalar information independently, uses the encoder outputs to scale position information into position and scale invariant values, and then routes them into an SO(3)-equivariant conditional U-net to predict denoised actions. In the figure, we omit scaling for the ease of viewing. VecLinear and VecConv1D refer to a SO(3)-equivariant version of linear and convolution 1D 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 of 2 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, we use a SIM(3)-equivariant version of PointNet++ with 4 layers and hidden dimensionality 128. In the Push T task, we decrease the number of layers to 2 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, we use the DDPM scheduler [13] and perform 100 denoising steps during inference. In real robot experiments, to optimize for inference speed, we use the DDIM scheduler [51] with 8 denoising steps.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. In our experiments, we found out that using 512 or 1024 points is sufficient for all tasks. In particular, for all real robot experiments and simulated mobile manipulation tasks, we use 1024-point point 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.Normalization. Data normalization can be important to the performance of diffusion models. Vanilla diffusion policy normalizes the observations and actions separately. We instead normalize all 3D-vector inputs (including observation and action) together due to our SIM(3)-equivariance assumptions. Implementation-wise, we take a subset of training data and compute the mean point cloud scale and mean action scale as s pc and s ac . Then, all position-and velocity-related information is divided by s pc /s ac at the start of the network forward pass and multiplied back before the output is returned. This normalization factor ensures that the diffusion process always works with actions with values within the -1 to 1 range. We normalize scalar information in the same way as the vanilla diffusion policy. We use a single ZED 2 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 = {τ n } N n=1 . Each human demonstration consists of a series of RGB-D image frames τ n = {I n t } T t=1 , where T is the episode horizon. The human demonstration processing module has three parts:1. an off-the-shelf object detection and tracking model X n t = Ψ(I n t , I n t-1 ) that takes the current and previous demonstration frames I n t and I n t-1 as input, and outputs a parsed point cloud of objects of interest X n t ; 2. an off-the-shelf hand detection model H n t = Φ(I n t ) that takes the current demonstration frame as input and outputs the keypoints on each human hand in the input frame H n t ; 3. an alignment module Ω(X n t , H n t , I n t ) 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 n t in the same coordinate system of X n t . We use Grounded Segment Anything Model with DEVA [52] as the object detection and tracking model Ψ and HaMeR [41] as the hand detection model Φ. In the alignment module, we find a set of matching points between the point cloud X n t and H n t , and then fit a rotation transformation R h and an offset t h to transform all points H n t to the coordinate frame of the point cloud. Then, we extract the thumb and index finger positions on the transformed keypoint set to predict the human "endeffector" pose y n t . We found that the alignment module Ω 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.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, we build a multi-process 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 directly 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 a 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 a threshold, we skip the first few predicted actions during action execution. After skipping the first few predicted actions to account for latency, 


Decision: 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.

Reviews: review: title: title: Initial Review
paper_summary: paper_summary: 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.
main_review: main_review: Overall the problem setting investigated in this paper is very important as equivariant / invariant policies have the potential to significantly improve data-efficiency which is crucial in imitation learning. However, for me, the relation and contributions w.r.t. the prior work of Equivact is unclear and needs further clarification. Moreover, I found the paper rather difficult to read and follow.

In the following, I provide a more detailed review.

1) Introduction. Overall, I am missing more citations in the introduction. While the introduction in its current form nicely motivates the direction of the work I feel that one single citation for an entire intriduction section is simply not sufficient given the vast amount of work in imitation learning, imitation learning for manipulation, and also in policy representation. The paper definitely needs more evidence (in the form of citations) to underline the authors claims. Moreover (but this is just a minor comment), I felt that the sentence in line 26-27 a bit too much of an overstatement, i.e., ".... that depart drastically from the scenarios in which it is trained". While I fully agree with the authors that equivariance is a strong property and can circumvent already quite some variations that would otherwise cause problems. Nevertheless, there exist way more environment variations where equivariance alone is not sufficient. Thus I suggest to tone down the statement.

2) Related work and in particular the relation with the other work [17] is not entirely clear and should be further clarified. Is it right that this work here essentially shares the same encoder as [17] and only uses a different policy backbone / policy parameterisation? If yes, this should be written and explained more clearly. Moreover, it would be essential that the authors provide more explanation and also evidence w.r.t. the fact that "... we show that this method displays unstable training performance" (line 72). As far as I see, at the moment the paper and supplementary only shows that empirically during evaluation the performance is different, however, I would argue that different performance is not necessarily related with "unstable training performance".

3) "Assumptions to observation and action spaces" at the end of Section 3.1. Since the representation of the observation and action spaces is essential, additional information here would significantly improve the paper. In particular, what is the magnitude of the end effector orientation that are then expressed as velocity vectors? Why are end effector positions integrated and translated into a set of points instead of using the velocity representation? Also, how exactly is the action representation retrieved / converted to desired end effector poses that the robot should go through? In particular, the authors mention the executed actions consist of 3D offsets or velocities, normalized directions and scalars. In my opinion more details are required here. As I see, there are also no details provided on this in the Appendix which makes it difficult to fully understand the method. This also continues in the later part (Section 3.2.) of the paper, where the sentence in lines 140-142 of how the "relevant part of the action" is scaled back was not understandable for me. I therefore suggest to define how actions are encoded, processed in the architecture, updated, etc. more strictly, in the form of equations.
On top of that I wonder, when the iterative denoising process is run, do you have to go multiple times through the Point-net based encoder, and if yes, what is the computational complexity of the method, i.e., what is its runtime?

4) Moreover (again rather minorily), I feel the statement "... other forms of visual information do not have the necessary 3D information to make it equivariant to translation, rotation, and scaling" (lines 122-124) a bit too strong. Combined with proprioception / pose information of a camera I think it might also be possible to integrate such observations into equivariant policies. I do agree that this might be less straightforward but still possible. Thus I would suggest to tone down this statement.

5) In the simulation experiments to improve readability, I propose to put the augmentations before the tasks, I feel that this would improve readability.

6) In the experiment results, I did not understand what is meant by the fact that the Equivact policy cannot handle multi-modal training data well (line 251). I assume this is related with the behavior cloning loss employed in EquivAct w.r.t. the DiffusionPolicy employed in this work. However, it would help to write this explicitly here in the paper.

7) To potentially save space, I propose to merge the Data collection paragraph and the Mobile robot setup as many parts are repeated and it is also inconsistent as in the Mobile robot setup there are the citations that are missing in the Data collection paragraph.

In light of the thorough rebuttal provided by the authors, I raise my score to weak accept. However, I want to point out that I agree with the other reviewers in that the paper would benefit significantly if also results were shown on more complex tasks (including more than one object) or long horizon tasks.
questions: questions: - Please explain in more detail how this work is different from [17]
- What exactly do you mean in unstable training performance of [17] w.r.t. to this work?
- 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?
- Why are end effector positions integrated and translated into a set of points instead of using the velocity representation?
- 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]?
limitations: limitations: 3
review_summary: review_summary: 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.
title: title: EquiBot demonstrates impressive real-world results but questionable simulation results.
paper_summary: paper_summary: 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.
main_review: main_review: **Strengths:**
- The real world experiments are fairly impressive with complex realistic tasks being solved using the proposed method.
- Paper is well written and easy to understand/follow.
- Scale invariance is typically overlooked in equivariant models and is included in this work. 
- This work is generally valuable. Data efficiency is a big challenge in robotics especially in these complex mobile robot manipulation domains where expert demonstrations are expensive. Equivariant models has been shown in both this and other works to be effective at addressing this.

**Weaknesses:**
- The primary weakness of this work is the performance of the baselines. Namely the performance of the diffusion policy baseline is lower than reported in the original Diffusion Policy work. While some of this could be explained by the change from RGB-D observations to point clouds, the authors state that they use the key point observations for the Push-T task and report significantly worse performance for the diffusion baseline than in the original work. This makes it difficult to trust the baseline performance for the other tasks as it implies some issue w/the baseline. This might be explained by the amount of expert data used for training but this is not stated anywhere as far as I see.
- Additionally, when benchmarking in simulation more than 10 episodes should be used to evaluate the methods.
- Limitations section is a bit bare-bones. Would be interesting to discuss the effect of incomplete point clouds on the method as this is a common failing of SO(3) equivaraint models.
- As this builds on top of EquiAct, the paper assumes a decent amount of pre-knowledge about the workings of EquiAct encoders.
- Novelty is a bit limited as this is primarily an extension to EquivAct where a diffusion head is placed onto of the EquiAct backbone.
questions: questions: 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?
2. Additional evaluation episodes for simulation experiments would strengthen the claims. 
3. How many demonstrations were used in the simulation experiments?
4. Is our understanding of this method being an extension of EquiAct with a diffusion policy head correct?
limitations: limitations: 3
review_summary: review_summary: 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.
title: title: The presented policy achieves SIM(3) equivariance, but only to a limited degree.
paper_summary: paper_summary: 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.
main_review: main_review: quality

Overall the paper achieves SIM(3) equivariance, but the method is constrained by canonicalization. Experiments showed the efficacy of the method, but ablation experiment is missing.

clarity

The presentation is clear, except that the Figure.3 in Appendix does not agree with Equation.1, 2 in the main text.

originality

This work has marginal novelty, because SE(3) equivariant diffusion planner has been explored previous work [1]. The contribution could be more significant if the method is applied to a broader manipulation tasks.

significance of this work

-strengths

Experiments demonstrate the generalization and sample efficiency by leveraging SIM(3) equivariance.

-weaknesses

Achieving SIM(3) equivariance by canonicalization limited the expressiveness of the method. Moreover, the method is based on type1 irreps in the feature, which potentially limits the rotational accuracy of the policy. In contrast to canonicalization, group convolution or higher types of irreducible representation of SO(3) could be utilized to increase expressiveness.

Lastly, there is no ablation experiment.

[1] 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.
questions: questions: 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?

Could author provide ablation experiments that show amount the SIM(3) (rotation, translation, and scale) equivariance, which equivariance is most critical?

Could author provide experiments similar to “out-of-distribution generalization experiments” that test the model on rotation, translation, and scale out-of-distribution individually?

Section. Equivariant diffusion architectures in the Appendix seems the most related works, could you put it in the main text?
limitations: limitations: 3
review_summary: review_summary: 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.