SE(3)-Equivariant Relational Rearrangement with Neural Descriptor Fields
Abstract: We present a framework for specifying tasks involving spatial relations between objects using only ∼ 5-10 demonstrations and then executing such tasks given point cloud observations of a novel pair of objects in arbitrary initial poses.
Our method uses Neural Descriptor Fields (NDFs) to achieve this by assigning a consistent local coordinate frame to the task-relevant parts of objects in demonstrations and localizing the corresponding coordinate frame on unseen object instances.
We propose an optimization method that uses multiple NDFs and a single annotated 3D keypoint in one of the demonstrations, to directly assign a set of consistent coordinate frames to the task-relevant object parts.
We also propose an energy-based learning scheme to model the joint configuration of the objects that satisfies a desired relational task.
We validate our pipeline on three multi-object rearrangement tasks in simulation and on a real robot.
Results demonstrate that our method can infer relative transformations that satisfy the desired relation between novel objects in unseen initial poses using just a few demonstrations.
Introduction
Many tasks we want robots to perform – e g., stacking bowls and plates to declutter a table, putting objects together to build an assembly, and hanging mugs on a rack with hooks – involve rearranging objects relative to one another.
Such tasks can be described in terms of spatial relations between part features of a set of objects, where a local coordinate frame is attached to the task-relevant part of the object and the relation is achieved by transforming the objects to bring these coordinate frames into a specified alignment.
For example, hanging a mug on a rack is a relation between the mug’s handle and the rack’s hook, while stacking a bowl on a mug involves aligning the bottom of the bowl with the top of the mug (see Fig. 1).
By representing tasks via relative configurations between local parts, a system can generalize across objects with diverse sizes and shapes [1–3].
Specifying task goals in this way requires the ability to (i) assign a consistent local coordinate frame to the task-relevant part of each object, and (ii) detect the corresponding coordinate frames on new object instances.
Prior works have demonstrated these capabilities using techniques such as supervised keypoint detection [1, 2], but the use of large task-specific datasets labeled by humans limits easy deployment for a wide diversity of tasks.
Neural Descriptor Fields (NDFs) [3] instead use task-agnostic self-supervised pretraining to achieve this ability using just a small set ( ∼ 5-10) of task demonstrations.
If consistent coordinate frames are assigned to the relevant local parts of the objects in the demonstrations, then NDFs can localize a corresponding coordinate frame on a novel object.
To avoid requiring the demonstrator to specify these consistent coordinate frames by hand, NDFs exploit an assumption that the task involves an interaction between an unknown object and a known secondary object.
This allows the user to position a single coordinate frame near the task-relevant part of the secondary object, and then use the demonstrations to associate this coordinate frame with the part of the unknown object that is used in the interaction.
While the results demonstrate generalization
[Figure]
to unseen objects in diverse poses [3], assuming a known secondary object is limiting.
For the “mug hanging on a rack” example, the system generalizes to scenarios with unseen mugs, but fails when given both an unseen mug and a rack with a new shape.
In this paper, we address this limitation.
In particular, we present Relational Neural Descriptor Fields (R-NDFs), a framework, using ∼ 5demonstrations, that takes as input 3D point clouds of a pair of unseen objects in arbitrary initial poses and outputs a relative transformation between them that satisfies a relational task objective.
The central difficulty in applying NDFs to scenarios with changing pairs of objects is to assign a set of consistent local coordinate frames to the task-relevant parts of the objects in the demonstrations, which may be both unaligned and differently shaped .
We propose an optimization method that uses two NDFs (one per object) and a single 3D keypoint label in just one of the demonstrations, to assign a set of local coordinate frames that are consistently posed relative to the task-relevant parts of the objects.
We then apply NDFs to localize the corresponding coordinate frames for unseen pairs of objects presented in arbitrary initial poses, and solve for the relative transformation between them that satisfies the desired relation.
However, errors can accumulate when inferring a relative transformation based on a pair of coordinate frames that have been independently localized.
To mitigate this effect, we also propose a learning approach that directly models the joint configuration of the pair of objects and helps refine the transformation for satisfying the relation.
We validate R-NDFs on three relational rearrangement tasks in both simulation and the real world.
Our simulation results show that R-NDFs outperform a set of baseline approaches, and our proposed optimization and learning-based refinement schemes benefit overall task success.
Finally, our real world results exhibit the effectiveness of R-NDFs on pairs of diverse real world objects in tabletop pick-and-place, and highlight the potential for applying our approach to multi-step tasks.
Background: Neural Descriptor Fields
A Neural Descriptor Field (NDF) [3] represents an object using a function f that maps a 3D coordinate x ∈ R 3 and an object point cloud P ∈ R 3 × N to a spatial descriptor in R d :
[Equation]
The function f is parameterized as a neural network constructed to be SE (3) -equivariant, such that if an object is subject to a rigid body transform T ∈ SE (3) its spatial descriptors transform accordingly * :
[Equation]
This enables NDFs to behave consistently for the same object, regardless of the underlying pose.
NDFs are also trained to learn correspondence over objects in the same category, so that points near similar geometric features of different instances (e.g., a point near the handle of two different mugs) are mapped to similar descriptor values.
The equivariance property is obtained by using
SO(3)-equivariant neural network layers [4] and mean-centered point clouds, while the category-level correspondence is obtained by training f on a category-level 3D reconstruction task [3, 5].
NDFs can also be redefined to model a field over full SE (3) poses, rather than individual points.
This is achieved by concatenating the descriptors of the individual points in a rigid set of query points
X ∈ R 3 × N q , i.e., a set of three or more non-collinear points x
i , i = 1 ...N q , that are constrained to transform together rigidly.
This construction allows NDFs to represent an SE (3) pose T via its action on X , i.e., via the coordinates of the transformed query point cloud T X :
[Equation]
Thus, F maps a point cloud P and an SE (3) pose T to a category-level pose descriptor Z ∈ R d × N q , where F inherits the same SE (3) -equivariance from f .
General Problem Setup and Preliminaries
Our high level goal is to enable a user to specify a task involving a geometric relationship between a pair of rigid objects, and enable a robot to perform this task on unseen object instances presented in arbitrary initial poses.
Examples of relations we consider include “ mug hanging on a rack ”, “ bowl stacked upright on a mug ”, and “ bottle placed upright on a tray ”.
Concretely, our goal is to build a system that takes as input two (nearly complete) 3D point clouds
P A and P B (each segmented out from the overall scene) of objects O A and O B , and outputs an
SE (3) transformation T B for transforming O B into a configuration that satisfies a desired relation between O A and O B .
We represent the relation as an alignment between a pair of local coordinate frames attached to task-relevant geometric features of the objects, and break down the problem of obtaining T B into (i) assigning a set of consistent coordinate frames to the task-relevant local object parts and (ii) localizing these coordinate frames on the relevant parts of the new objects.
Furthermore, we assume a user specifies the relational task by providing a small handful of K task demonstrations {D i } Ki =1 , such that it’s intuitive and efficient to specify a wide diversity of tasks with minimal engineering effort.
A demonstration D consists of point clouds ˆ P A and ˆ P B (of objects ˆ O A
and ˆ O B ) and relation-satisfying transformation ˆ T B .
NDFs for Encoding Single Unknown Object Relations .
Prior work on NDFs may be applied to a simplified version of this task, where the geometry and state of O A is known .
Given that O A is known, we can initialize a set of query points X A near the task-relevant part of O A and use the query points to encode the relative pose ˆ T B via Equation (3).
Thus, a demonstration D is mapped to a target pose descriptor ˆ Z = F (ˆ T − 1 B | ˆ P B ) representing the (inverse of the) final pose of ˆ O B relative to O A .
In practice, pose descriptors from multiple demonstrations {D i } Ki =1 are averaged to obtain an overall descriptor ˆ Z = 1 K (cid:80)
Ki
=
ˆ Z i for the whole set, which has important implications in the version of the task with two unknown objects (see Section 4.1 for further discussion).
Given a novel object instance represented by point cloud P B , we can compute a transformation T B such that transforming O B by T B satisfies the demonstrated relation between O A and O B .
This is achieved by minimizing the L1 distance to the target pose descriptor ˆ Z :
[Equation]
Intuitively, Equation (4) performs well across different objects due to the fact that NDFs are pretrained to enable reconstruction across a large dataset of 3D shapes.
As a result, shared descriptors are discovered across different instances in a shape category.
In contrast, training a model directly on the few demonstrations (e.g., for regressing pose T B ) would be more susceptible to overfitting.
Method
We now describe how we apply NDFs to infer relations between pairs of unknown objects.
In
Section 4.1, we propose an iterative optimization method for assigning consistent task-relevant coordinate frames to multiple objects.
In Section 4.2, we discuss how we train a neural network on top of NDF features to model the joint object configuration and refine an inferred transformation.
The system inputs consist of pretrained NDFs f A and f B for each object category, demonstrations {D i } Ki =
= { (ˆ P A , ˆ P B , ˆ T B ) i } Ki =1 , and a single labeled 3D coordinate x AB for one of the demonstrations, indicating approximately where the respective demonstration objects interact.
[Figure]
4.1 Multiple NDFs for Inferring Pairs of Task-Relevant Local Coordinate Frames
Consider a scenario where O A and O B have unknown underlying shapes and configurations.
We now show how NDFs can be used for inferring a pair of task-relevant local coordinate frames on both objects and recovering a transformation T B that satisfies the relation.
The key idea of our approach is to formulate this problem as a bi-level optimization (illustrated in Figure 2), where we first optimize to find a task-relevant portion of O A , and subsequently optimize a relative transform of a local part of O B with respect to the local region of O A .
We begin with two pretrained NDFs, f A and f B , and query points X A in a canonical pose at the world frame origin.
We obtain X A by sampling N q points from a zero-mean Gaussian and scaling such that X A has scale similar to the salient object parts.
We then use the keypoint x AB to transform X A near the task-relevant features in the demonstration associated with x AB .
Denote this transformation as ˆ T X A .
Finally, we encode world-frame pose ˆ T X A into a descriptor conditioned on ˆ P A , as ˆ Z A
=
F A (ˆ T X A | ˆ P A ) , and relative pose ˆ T − 1 B as ˆ Z B =
F B (ˆ T − 1 B ˆ T X
A
|
ˆ P B ) ,
conditioned on ˆ P B .
At test-time, we optimize both the world-frame pose of the query points T X A and the (inverse of) pose T B relative to the initial pose found in the first step:
[Equation]
Figure 2 shows an example of this pipeline, where the resulting T B is applied to the point cloud P B of O B to satisfy the “hanging” relation.
[Figure]
Minimizing Descriptor Variance .
In practice, solving Equations (5) and (6) works better if pose descriptors { ˆ Z i } Ki =1 from multiple demonstrations are averaged together to obtain an overall target descriptor ˆ Z = K (cid:80) Ki
=
ˆ Z i (see Sec.
6.1 and [3]).
The reason is that a single demonstration underspecifies which object parts are relevant for the task, allowing ˆ Z to be sensitive to object features which are not relevant to the desired relation.
Instead, a set of demonstrations using slightly different objects (e.g., with different scales) reveals regions near local interactions that are shared across the demonstrations, which helps disambiguate between parts that are critical vs. irrelevant for the specified relation.
However, to avoid the pitfalls of averaging across a potentially multimodal or disjoint set, we want descriptors in the set { ˆ Z i } Ki =1 to be sensitive to nearby local geometry in a way that is consistent (i.e., unimodal ) across the demos.
This only occurs if the query points used to obtain the descriptors are themselves consistently aligned relative to each respective object (see Figure 3).
Therefore, we need to find a transformation ˆ T X A ,i for each demonstration D i that transforms the canonical query points
X A into a configuration that leads the descriptors { ˆ Z i = F A (ˆ T X A ,i | ˆ P A,i ) }
Ki =1 to be consistent with each other.
We address this by finding the set of transformations { ˆ T X A ,i } Ki =1 that minimizes the variance across the descriptor set { ˆ Z i =
F A (ˆ T X A ,
i | ˆ P A,i ) }
Ki =1 :
[Equation]
X A,i } [-Ki =
 Var ( {-][Equation]
[Equation]
subject to
[Equation]
where Var ( · ) denotes the sum of the per-element variance across a set of vectors.
We perform this minimization by applying NDFs in an alternating optimization procedure.
Starting with an initial reference pose (constructed using x AB ) placing X A near the task-relevant object parts in one of the demonstrations, we iteratively apply Equation (5) to obtain a descriptor for each demonstration that matches the reference.
At the outer level, we refit the reference descriptor using the mean of the most recently obtained individual descriptors, and repeat.
More details can be found in the Appendix.
4.2 Capturing Joint Descriptor Alignment through Learned Energy Functions
The method in Section 4.1 proposes to infer a desired relation by sequentially localizing independent coordinate frames for each object.
While this approach is generally effective, small errors can accumulate and cause slight misalignments that lead to failure in the execution.
We thus propose to learn a neural network which directly captures the joint configuration of O A and O B that satisfies the desired relation, and use this model to refine predictions made by the method in Section 4.1.
Pairwise Energy Functions .
We train an Energy-Based Model (EBM) E θ ( · ) [6] to parameterize a learned energy landscape over NDF encodings of relative poses between O A and O B (i.e,.
E θ acts as a learned analogue for the L1 distance in Section 4.1).
The energy function E θ ( · ) is trained so that the ground truth transform of O B with respect to O A is recovered given NDFs f A and f B (note that f corresponds to descriptor evaluation at single coordinate x while F is defined over sets of coordinates).
Explicitly, our energy function is trained so that:
[Equation]
Since each NDF is a continuous field, it is difficult to input them directly into our energy function
E θ ( · ) .
We represent the energy function as the sum of the point-wise evaluation of each NDF on a set of different query points X E sampled from transformed pointcloud TP B .
[Equation]
At test-time, we use Equation (8) to refine the transformation obtained using Equations (5) and (6).
4.3 Learning
NDF training .
We represent NDFs f A and f B as two neural networks with identical architecture and separate weights.
Following [3], the architecture consists of a PointNet [7] point cloud encoder with
SO(3) equivariant Vector Neuron [4] layers, and a multi-layer perceptron (MLP) decoder.
The NDF is represented as a function mapping a 3D coordinate and a point cloud to the vector of concatenated activations of the MLP.
The models are trained end-to-end to reconstruct 3D shapes given object point clouds.
We use a dataset of ground truth 3D shapes and generate a corresponding set of 3D point clouds in simulation.
More architecture and training data details can be found in the Appendix.
Energy-Based Model Training .
We supervise the EBM E θ so that optimization over the learned energy landscape recovers the relative transform between O A and O B .
In particular, we follow the training objective in [8] and train arg min T [ E θ ( · )] to match a target pose using the following procedure.
We first apply a small delta perturbation T ∆ to ˆ T B ˆ P B (i.e., the point cloud of O B in its final configuration) to obtain ˆ P B, ∆ =
T ∆ ˆ T B ˆ P B .
We then train E θ to iteratively refine an initial random pose T 0 with translation t 0 and rotation R 0 to undo the perturbation pose T ∆ .
We run n steps of optimization on t 0 and R 0 , where an individual step is given by t k = t k − 1 − λ ∇ t E θ ( f A ( ·| ˆ P A ) , f B ( ·| T ˆ P B, ∆ )) and R k
= R k − 1 − λ ∇ R E θ ( f c ( ·| ˆ P A ) , f B ( ·| T ˆ P B, ∆ )) .
We may train the energy function so that T n corresponds to the inverse of the perturbation pose T ∆ using L trans =
∥ t n − t − 1∆ ∥ and L rot =
∥ R n − R − 1∆ ∥ .
However, with symmetric objects, there are multiple different rotations R n which may satisfy the desired relation (e.g., a bowl is still “on” a mug, regardless of the angle about its radial axis).
To account for these symmetries, we implicitly enforce consistency between an optimized transform T n and T − 1∆ by enforcing that its application on ˆ P B, ∆ leads to a similar point cloud to ˆ T B ˆ P B .
We achieve this by minimizing the Chamfer loss [9] between the optimized transformed point cloud T n ˆ P B, ∆ and the demonstration point cloud ˆ T B ˆ P B .
Application to Tabletop Manipulation
Robot and Environment Setup .
We apply the method in Section 4 to the problem of tabletop object rearrangement using a Franka Panda robotic arm with a Robotiq 2F140 parallel jaw gripper.
The arm is used to collect the demonstrations and to execute the inferred transformation at test-time.
Our environment consists of the arm on a table with four calibrated depth cameras.
Providing and Encoding Demonstrations .
When collecting a demonstration, initial object point clouds ˆ P A and ˆ P B of objects ˆ O A
and ˆ O B are obtained by fusing a set of back projected depth images.
The demonstrator moves the gripper to a pose ˆ T grasp , grasps ˆ O B , and finally moves the gripper to a pose ˆ T place that satisfies the desired relation between ˆ O A and ˆ O B .
ˆ T B is obtained as ˆ T place ˆ T − 1 grasp .
In one of the demonstrations, a 3D keypoint x AB is labeled near the parts of the objects that interact with each other by moving the gripper to this region and recording its position.
Test-time Task Setup and Inference .
At test time, we are given point clouds P A and P B of new objects O A and O B .
Equations (5 ), ( 6 ), and ( 8 ) are applied in sequence to obtain T B .
T B is applied to O B by transforming an initial grasp pose T grasp (obtained using a separate grasp generation pipeline) by T B to obtain a placing pose T place =
T B T grasp , and off the shelf inverse kinematics and motion planning is used to reach T grasp and T place .
Experiments and Results
Our experiments are designed to evaluate R-NDFs in executing relational rearrangement tasks with unseen objects using only a few demonstrations.
We seek to answer three questions: (1) How well do
R-NDFs predict transformations that satisfy a relational task?
(2) How important is each component in R-NDFs?
(3) Can R-NDFs be used to perform multi-object pick-and-place tasks in the real world?
Baselines .
As existing rearrangement methods are not directly applicable with so few demonstrations, we compare with two constructed baselines.
The first is to train an MLP to directly regress the relative transformation between objects (“Pose Regression”).
The MLP takes as input the point cloud encodings obtained from the same PointNet [7] encoder with Vector Neuron [4] layers used in
NDFs, and is trained directly on the demonstrations.
The second method is based on 3D point cloud registration (“Patch Match”).
We use a state of the art registration method [10] to align the test-time shapes to the demonstration shapes and then compute the resulting relative transformation.
Task Setup and Evaluation Metrics .
We consider three relational rearrangement tasks for evaluation: (1) Hanging a mug on the hook of a rack, (2) Stacking a bowl upright on top of a mug, and (3) Placing a bottle upright inside of a box-shaped container.
We provide 10 demonstrations of each task and evaluate if each method, using the demonstrations, can infer a transformation that satisfies the desired relation for unseen pairs of object instances with randomly sampled poses.
Experiments are conducted in both the real-world and in simulation using PyBullet [11].
In simulation, the transformation obtained by each method is directly applied by resetting the simulator to the transformed object states.
To quantify performance, we report the success rate over 100 trials, where we use the ground truth simulator state to compute success (objects must be in contact, have the correct relative orientation, and not interpenetrate).
6.1 Simulation Results
We begin by evaluating how well R-NDFs can infer the desired transformations in simulation.
We consider two settings of varying difficulty.
First, the pair of unseen objects are positioned randomly on the table with a randomly sampled “upright” orientation (similar to those used in the demonstrations).
Second, the orientation of O B is randomly sampled from the full space of 3D rotations.
Results in Table 4a compare the performance of our approach to the baselines.
We find that training a model to directly regress the relative pose using the point cloud embeddings leads to substantial
Bowl on Mug
Mug on Rack
[Figure] [Figure]
overfitting and much lower success rates.
On the other hand, the registration-based method can sometimes find transformations that correctly align the unseen shapes to the demonstration objects, and thus achieves higher success rates than pose regression.
However, 3D registration is susceptible to locally optimal results that align the task irrelevant parts of the objects.
Common failure modes of using 3D registration in the tasks we consider include aligning the body of the mug but ignoring the handle, or aligning the racks to be upside down.
Figure 4b illustrates the final simulator state after applying some of the representative predictions of each method.
In contrast, R-NDFs more accurately localize the task-relevant object parts and assign coordinate frames to these parts that are consistent with the demonstrations, leading to the highest success rates.
Consistent with [3], the performance gap between the “upright” and “arbitrary” pose settings is small, which can be attributed to the built-in equivariance of the features used in R-NDF.
6.2 Ablations
Next, we analyze the importance of the individual components of R-NDFs.
We investigate ablations on the simulated “mug on rack” task, again considering both “upright” and “arbitrary” pose settings.
The top row of Table 5a illustrates that R-NDF performs worse with a single demonstration.
Since there are multiple possible explanations for the alignment between two objects when given one example of the desired relation, pose descriptors obtained from a single demonstration are more sensitive to task irrelevant object features.
The second row of Table 5a investigates the effect of averaging descriptors across the set of demonstrations without first aligning the query points relative to the objects in each demo.
We modified the demonstrations to provide keypoints { x AB,i } Ki =1 near the relevant region in each demonstration, and then transform the query points to this region without aligning their orientations.
Removing the query point alignment reduces the performance.
The third row of Table 5a shows that removing the EBM refinement also decreases the success rate.
We further examine the importance of accurately specifying the 3D keypoint x AB near the taskrelevant region on one of the demonstrations.
We run the trials multiple times with Gaussian distributed noise added to the labeled point.
Figure 5b shows a plot of the success rate vs. the noise magnitude normalized by the approximate size of the object.
The plot indicates that with limited noise perturbation, the success rate does not suffer significantly, though we observe a steep decline with more substantial perturbations.
These larger perturbations shift the query points to regions near geometric features that are less relevant to the desired relation.
[Figure]
6.3 Real Results
Finally, we validate that R-NDFs can be used to perform pick-and-place on pairs of unseen objects in the real world.
Figure 6 shows the execution on our three tasks.
Our method successfully infers a tranformation between the objects that satisfies the relations, despite the objects being presented in a challenging array of initial configurations.
Figure 1 shows a multi-step rearrangement application of R-NDFs for the “bowl on mug” task.
First, a relation between the mug and the table is specified and inferred for placing the mug upright.
Then, the system executes the “stacking” relation between the bowl and the upright mug .
This highlights how R-NDFs can enable executing sequential chains of relations to satisfy task objectives involving more than two objects.
Please see our attached supplemental video for additional real world results.
Related Work
Novel Object Rearrangement .
Several methods exist for novel object rearrangement [1, 12–29], many of which don’t consider multiple varying objects that interact.
CatBC [30] uses dense correspondence models to achieve impressive pick-and-place policy generalization from a single demonstration, but assumes a known receptacle for placing.
Neural shape mating [31], OmniHang [32], and kPAM 2.0 [2] generalize to pairs of unseen objects, but these approaches train on large task-specific datasets.
TransporterNets [33, 34] enables rearrangement with varying pick and place locations from a few demonstrations, but focuses on top-down manipulation and struggles with out-of-plane reorientation.
In contrast, we focus on executing relations involving large 3D reorientations.
Neural Fields in Robotics .
Neural fields use neural networks to parameterize functions over continuous spatial or temporal coordinates [35].
They have been applied to model various signals and scene properties, such as images [36], geometry [5, 37, 38], appearance [39, 40], tactile imprints [41], and sound [42], with high fidelity and memory efficiency.
Neural fields have been applied to represent objects for manipulation [3, 43–46] and environment states for dynamics and policy learning [47– 49].
They have also been used for pose estimation [50, 51], SLAM [52, 53], and representing object geometry without depth cameras [54, 55].
Limitations and Conclusion
Limitations .
R-NDFs require a pretrained NDF for each category used in the task, which can be nontrivial to obtain for novel object categories without existing 3D model datasets.
Our approach also requires an annotated keypoint to localize task-relevant object parts.
Future work could explore automated discovery of task-relevant regions directly from a set of demonstrations.
Our system uses depth cameras, which often struggle with noise and objects with thin and transparent features.
An
RGB-only approach offering a similar level of generalization would be interesting to investigate.
Finally, we require segmented object point clouds.
While object instance segmentation is quite mature, pretrained segmentation models regularly struggle when objects are in diverse orientations.
Conclusion .
This work presents an approach for learning from a limited number of demonstrations to rearrange novel objects into configurations satisfying a relational task objective.
We develop methods that build upon prior applications of neural fields for representing objects and increase the scope of tasks they can achieve.
Our results illustrate the general applicability of our framework across a diverse range of relational tasks involving pairs of novel objects in arbitrary initial poses.