Deep SE(3)-Equivariant Geometric Reasoning for Precise Placement Tasks

TL;DR

The paper addresses precise relative placement in robotic manipulation by enforcing $SE(3)$-equivariance through a two-part design: an $SE(3)$-invariant RelDist representation of cross-object relationships and differentiable geometric reasoning layers (multilateration and Procrustes) that recover the cross-pose. The approach enables end-to-end training from few demonstrations and generalizes across object variations, outperforming baselines in high-precision placement tasks on RLBench, NDF, and real-world mug-hanging scenarios. Key contributions include the RelDist invariant representation, the differentiable multilateration layer $\texttt{MUL}$, and the differentiable Procrustes layer $\texttt{PRO}$, all guaranteeing $SE(3)$-equivariance by construction. The results demonstrate substantially improved placement precision and robust real-world applicability, with limitations noted for symmetric objects and the need for object segmentation.

Abstract

Many robot manipulation tasks can be framed as geometric reasoning tasks, where an agent must be able to precisely manipulate an object into a position that satisfies the task from a set of initial conditions. Often, task success is defined based on the relationship between two objects - for instance, hanging a mug on a rack. In such cases, the solution should be equivariant to the initial position of the objects as well as the agent, and invariant to the pose of the camera. This poses a challenge for learning systems which attempt to solve this task by learning directly from high-dimensional demonstrations: the agent must learn to be both equivariant as well as precise, which can be challenging without any inductive biases about the problem. In this work, we propose a method for precise relative pose prediction which is provably SE(3)-equivariant, can be learned from only a few demonstrations, and can generalize across variations in a class of objects. We accomplish this by factoring the problem into learning an SE(3) invariant task-specific representation of the scene and then interpreting this representation with novel geometric reasoning layers which are provably SE(3) equivariant. We demonstrate that our method can yield substantially more precise placement predictions in simulated placement tasks than previous methods trained with the same amount of data, and can accurately represent relative placement relationships data collected from real-world demonstrations. Supplementary information and videos can be found at https://sites.google.com/view/reldist-iclr-2023.

Figures (11)

• Figure 1: Invariance of relative placement tasks under transformations. In this case, a ring on peg maintains the same relative position under a rigid transformation $\mathbf{T}$.
• Figure 2: Method overview. First, the point clouds $\mathbf{P}_{{\mathcal{A}}}, \mathbf{P}_{{\mathcal{B}}}$ are each encoded with a dense $SE(3)$-equivariant encoder, after which cross-attention is applied to yield task-specific dense representations. Then, the kernel matrix $\mathbf{R}_{{\mathcal{A}}{\mathcal{B}}}$ is constructed through the learned kernel $\mathcal{K}_\psi$. This matrix is then passed into $\texttt{MUL}$ to infer the desired final point clouds, and then passed into $\texttt{PRO}$ to extract a final transform which moves object ${\mathcal{A}}$ into its goal position.
• Figure 3: Reasoning with multilateration (a) in a 2D environment with block and gripper. (b) For each sampled point on the gripper, (c) we estimate the desired distances between the gripper point and points on the block. We then use multilateration (d) to extract a least-squares solution to compute the desired gripper point location. Doing this for every point on the gripper, (e) we can reconstruct the desired position for each gripper point. (f) These corresponding points can be used to infer a rigid transform which brings the gripper to the final goal position.
• Figure 4: RLBench (James2020-mj) relative placement tasks. Top: the initial state of a demonstration. Bottom: the final state of a demonstration, where a successful placement has been achieved.
• Figure 5: A real-world demonstration of mug-hanging.

Theorems & Definitions (2)

• Theorem 1
• proof