Fourier Transporter: Bi-Equivariant Robotic Manipulation in 3D

We thank the reviewer for an especially detailed and thorough review. We provide a point-by-point response to the comments and questions below:

$\textcolor{red}{ \text{ Weaknesses: }}$

While the paper demonstrates strong results on RLBench tasks, it's important to note that some tasks like stack-blocks and stack-cups primarily operate in 2D space, which may not fully reveal the strengths of the methods in 3D. It would be valuable to include additional tasks that involve more 3D rotation angles, such as put books on bookshelf.

Note that three of our benchmark tasks ( 'stack-wine', 'place-cups', 'put-plate') are fully 3D and require the ability to reason about the out-of-plane actions. Specifically, in the 'stack-wine' task, the agent must rotate the wine bottle from vertical to horizontal then place it in holder. In order to complete this task, the agent must be able to understand that objects can be rotated out of the plane. Furthermore, while we agree with the reviewer that the 'stack-blocks' and 'stack-cups' tasks do not require significant out-of-plane reasoning, our action space is still defined over all poses in $SE(3)$. These experiments are important because they indicate that our general $SO(3)$ method is competitive with purely $SO(2)$ methods. We will run an additional task which is similar to the 'put books on bookshelf' task.

$\textcolor{red}{ \text{ Questions: }}$

In section 5.3 2D Pick-Place results, the last line: 'It indicates that the $SO(2) \times SO(2)$ equivariance of FOURTRAN is more sensitive to rotations. ”. Does “sensitive” means more precise or prone to noise?

It means more precise. By 'sensitive', we mean that $SO(2) \times SO(2)$-equivariant model is better able to capture high resolution rotations. We have replaced sensitive by precise in the main text. This can be seen by noting that at $\omega=7.5$ degrees our method still outperforms the 2d baselines.

It would be interesting to conduct separate tests with high-resolution thresholds to distinguish the impact of position error and rotation error. For example, considering parameters like $\tau$ = 1cm and $\omega$ = 7.5° as well as $\tau$ = 0.5cm and $\omega$ = 15°. Additionally, a box plot of the rotation error would also provide more insight into the effect of the method.

These are great suggestions that will will incorporate into the final version of the paper. We will include these results when they are finished.

Dear Reviewer sjsh,

As we progress through the review process for ICLR 2024, I would like to remind you of the importance of the rebuttal phase. The authors have submitted their rebuttals, and it is now imperative for you to engage in this critical aspect of the review process.

Please ensure that you read the authors' responses carefully and provide a thoughtful and constructive follow-up. Your feedback is not only essential for the decision-making process but also invaluable for the authors.

Thank you,

ICLR 2024 Area Chair

Thank you very much for the useful feedback.

$\textcolor{red}{ \text{ Weaknesses: }}$

The core weakness of this work is the strength of the contribution when compared to the equivariant transporter network proposed in Huang 2022... more needs to be done to justify why the extension to 3D is non-trivial, particularly when it comes to the major claims of this work, greater angular resolution, computational benefits of Fourier implementation, and sample efficiency.

We totally disagree that this paper is not significant relative to Huang et al. 2022.

1. The generalization from the one-dimensional abelian group $SO(2)$ to the three-dimensional non-abelian group $SO(3)$ is highly non-trivial. The method proposed in (Huang 2022) cannot not be extended in three dimensions directly. (Huang 2022) requires the use of one rotation for every element of the group $G$. For subgroups of $SO(2)$, the complexity is O$(n)$, but for the case $SO(3)$, this complexity is O$(n^3)$ -- something that is totally intractable using 3D translational convolution. For example, consider that (Huang 2022) uses $C_{36}$ which has an angular resolution of $10$ degrees and requires $36$ copies of the input feature map to construct a steerable kernel. To achieve the same angular resolution in $SO(3)$, this would require copying the pick feature voxel map into $36^3/2=23228$ discrete orientations.

2. Representing the action distribution with Fourier coefficients is a significant conceptual improvement from (Huang 2022). By utilizing the Fourier Transform on groups, we can sidestep the challenges described earlier and perform memory efficient and accurate computation in three-dimensional space. It allows us to compute and represent the pick/place distribution over continuous action space. To the best of our knowledge, we are the first to represent the pick-place action distribution with Fourier coefficients.

3. The generalization into 3D is very important from a practical perspective. Many imitation learning methods reason in 3D and this paper is the first to design a Transporter-like model for this setting.

No error bars are provided in experiments (Tables 1/2), so we have no indication that the results are significant. I am sure they are, but this is important for the table 2 comparison with Huang 22

The results reported in Tables 1 and 2 are averaged over 100 trials. We have added a note to this effect in the text. We will rerun experiments and include error bars in the final version of our work.

The mapping between Figure 2 and the equations in Section 4 is incomplete, and not easily followed... the current structure of this section/group theory jargon made this difficult to follow. We want our work to be readable to a wide audience, specifically researchers in robotic manipulation who are unfamiliar with group theory. For this reason, we have added additional section on mathematical background. We have added pseudocode (see new supplementary) that describes our algorithm on an intuitive level.

$\textcolor{red}{ \text{ Questions: }}$

What level of rotational variation is present in the demonstrations and experiments? Is it possible to share some data on the typical distributions and tolerances in pick/place angles the tasks here require?

We have added Figure.5 in appendix to plot the $SO(3)$ distribution of expert actions on two tasks.

Could you explain table 3 in more detail - the ablation here appears to undermine many of the choices this work makes... why do we need bi-equivariance? Is this simply an artifact of the test scenarios not adequately evaluating 3D equivariance symmetries?

The ablation does not undermine our design, because both data augmentation and the lifting operation contribute to the bi-equivariance property of the network. In Proposition 1, we show three constraints that are needed for bi-equivariance; however, the second and third constraints can be implemented through different methods. Results in table 3 are not artifact of the test scenarios not adequately evaluating 3D equivariance; they evaluate different methods to achieve 3D equivariance.

In terms of the extension to 3D, it seems that the lifting operator introduces challenges around partial occlusions... Could you comment on the general performance/ potential limitations in this regard?

Occlusion in 3D is a challenge that all methods in this area must address. Our method does about as well as the baselines in this regard. Partial observation is common in 3D and affects the equivariance of our model to some extent. The crop centered at the pick location usually contains only part of the picked object (stack-wine, put-plate) and parts of neighboring objects. Our model lifts the dense feature map (dense descriptors) from the crop instead of the raw crop signal. Equivariance is built on the top of the 3D CNN and our model has the same robust reasoning and learning ability as CNN to deal with the partial observation.

Thank you for the quick response and clarification. Great point. We will definitely include the more explicit description of our contributions in the final draft!

We thank the reviewer for the detailed feedback.

$\textcolor{red}{ \text{ Weaknesses: }}$

Given that the Wigner D-matrix representation and corresponding 3D Fourier transform is the key insight that allows this action representation to work, it would be worth spending some more time to describe them in more detail

Thanks for pointing this out. We have added an additional section on mathematical background describing the Wigner-D matrices and non-abelian Fourier transform in more detail. A full exposition of representation theory would require a textbook, but we hope that our work can be intuitively understood by reader.

Some pseudocode/method description would be welcome Agreed. We have added pseudocode to the supplementary material.

$\textcolor{red}{ \text{ Questions: }}$

Is it possible that equivariance (2D rotational+translational) is actually more general if grasp dynamics are dependent on the object's orientation with respect to gravity?

This is a great point. We have assumed that grasp dynamics do not play a large role in our specific settings and thus $SE(3)$ symmetry is the most relevant. In the context of a good grasp and with a sufficiently large gripper closing force, the object will not shift greatly under gravity and grasping will be $SO(3)$ invariant. In cases where gravity plays a large roll in grasp dynamics, gravity may break the full $SE(3)$ symmetry to $SO(2) \ltimes R^3$ as you suggest. A simple way to address this issue is to simply add the gravity vector as an additional input to the $SE(3)$-equivariant model as done in related work [1].

[1] Fuchs, F., Worrall, D., Fischer, V., & Welling, M. (2020). Se (3)-transformers: 3d roto-translation equivariant attention networks. Advances in neural information processing systems, 33, 1970-1981.

It seems like $I_{60}$ is used in the lifting operation, but as far as I can tell there's no reason the set of rotations has to form a subgroup. Is this correct, e.g. could the granularity be increased by simply sampling more rotations (roughly evenly spaced in ) in this step?

Correct. Randomly sampling a large number of rotations is exactly what we did in our experiments. We consider the special case where the samples align with the subgroup $I_{60}$ in order to make a theoretical connection to previous work. When the lifting operation uses $I_{60}$, we prove that the resulting filters are an $I_{60}$--steerable kernel. When we sample random rotations uniformly, this produces an approximately $SO(3)$ steerable kernel. We have clarified this point in the draft.

Dear Reviewer n4JL,

As we progress through the review process for ICLR 2024, I would like to remind you of the importance of the rebuttal phase. The authors have submitted their rebuttals, and it is now imperative for you to engage in this critical aspect of the review process.

Please ensure that you read the authors' responses carefully and provide a thoughtful and constructive follow-up. Your feedback is not only essential for the decision-making process but also invaluable for the authors.

Thank you,

ICLR 2024 Area Chair

Please find our point-by-point response below:

$\textcolor{red}{ \text{ Weaknesses: }}$

The current model is limited to a single-task setting, while the baseline methods are designed for multi-task purposes. I'm concerned that the comparisons may not be fair.

This is something we specifically addressed in Table 1 where we compared the performance of our method in the single-task setting with the baselines trained both ways, i.e. with the baselines trained in the single task setting (Baseline-Single in Table 1) and in the multi-task settings (Baseline-Multi). The single task setting gives an apples-to-apples comparison in that both methods have the same training. However, we also compare to Multi-task trained baselines since, as you point out, this was the setting they were designed for. Our method outperforms the baselines either way.

It relies solely on open-loop control, disregarding path planning and collision awareness.

This is not a weakness of our method. The open loop setting which we use (also known as ``keypoint'' control) has become a standard paradigm in the literature on imitation learning for robotic manipulation. All the baseline methods, Robotic View Transformer [1], PerAct [2], and Coarse to Fine [3] are open loop control methods. Additionally all baseline methods use off-the-shelf path planning algorithms.

[1] Ankit Goyal, Jie Xu, Yijie Guo, Valts Blukis, Yu-Wei Chao, and Dieter Fox. RVT: Robotic view transformer for 3d object manipulation. In 7th Annual Conference on Robot Learning, 2023

[2] Mohit Shridhar, Lucas Manuelli, and Dieter Fox. Perceiver-actor: A multi-task transformer for robotic manipulation. In Conference on Robot Learning, pp. 785–799. PMLR, 2023.

[3] Stephen James, Kentaro Wada, Tristan Laidlow, and Andrew J Davison. Coarse-to-fine q-attention: Efficient learning for visual robotic manipulation via discretisation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2022

The paper is not well-written and some of the terms are difficult to understand. It uses a lot of notations, but many of them are not explained.

We acknowledge that the concepts could have been presented more clearly. We have tried to address this by adding a section on mathematical background and adding pseudocode to help the reader understand our method intuitively.

There are no real robot experiments.

We agree that physical robot experiments would strengthen the paper and we are currently working on that. However, we would like to note that physical experiments are not a requirement for ICLR acceptance. Several ICLR papers [1-3] have been published without physical experiments. We are nevertheless working on physical experiments and we plan to include them in the final version of this paper. As evidence that our method works on physical hardware, we have included a video clip of our method running on a Panda robot in the real-world in the supplementary material.

[1] Ryu, Hyunwoo, et al. "Equivariant Descriptor Fields: SE (3)-Equivariant Energy-Based Models for End-to-End Visual Robotic Manipulation Learning." ICLR 2023

[2] Yu, Albert, and Ray Mooney. "Using Both Demonstrations and Language Instructions to Efficiently Learn Robotic Tasks." ICLR 2023

[3] Li, Sizhe, et al. "Contact Points Discovery for Soft-Body Manipulations with Differentiable Physics." ICLR 2022

$\textcolor{red}{ \text{ Questions: }}$

What is fiber space Fourier transformation? The fiber space Fourier transformation computes the Fourier transformation (spectral decomposition) over the channels. This operation is performed separately at each voxel. The use of the fiber space Fourier transform is a key component of our method and we have added an additional explanation in the mathematical background section.

In Figure 2, how do you crop the object in the scene? What if there are multiple objects? We crop a $24^{3}$ voxel box centered at the pick location. The model does not rely on any object segmentation; it simply learns to crop around the object, since this is where the pick action occurs. Figure 2 is used to explain our idea. The crop usually contains part of the picked object and parts of the neighboring objects. [2] and [3] use similar crop approaches in the 2D case.

[1] Ethan Chun, Yilun Du, Anthony Simeonov, Tomas Lozano-Perez, and Leslie Kaelbling. Local neural descriptor fields: Locally conditioned object representations for manipulation. arXiv preprint arXiv:2302.03573, 2023

[2] Haojie Huang, Dian Wang, Robin Walters, and Robert Platt. Equivariant Transporter Network. In Proceedings of Robotics: Science and Systems, New York City, NY, USA, June 2022. doi: 10.15607/RSS.2022.XVIII.007

[3] Yen-Chen Lin, Pete Florence, Andy Zeng, Jonathan T Barron, Yilun Du, Wei-Chiu Ma, Anthony Simeonov, Alberto Rodriguez Garcia, and Phillip Isola. Mira: Mental imagery for robotic affordances. In Conference on Robot Learning, pp. 1916–1927. PMLR, 2023