Neural Descriptor Fields: SE(3)-Equivariant Object Representations for Manipulation

In plain words

Teaching a robot to pick up a mug and hang it on a rack usually takes many demonstrations, because a new mug can sit anywhere and at any angle. This paper describes an object in a way that lets a robot find the right gripper pose from only a handful of demonstrations.

The method learns, without hand-labelled points, a function that gives every point in the space around an object a descriptor. The descriptor says where that point sits relative to the object's shape. A small set of points attached to the gripper, or to a shelf, carries these descriptors. The robot then moves that set of points until its descriptors match the ones seen in the demonstration. The descriptors are built so that they turn and move with the object, so the match still works for objects placed in poses the robot never saw in training. The authors test this in simulation and on a real robot.

TL;DR

Neural Descriptor Fields (NDFs) map each 3D point, conditioned on an object point cloud, to a descriptor built from the layer activations of an occupancy network §II-A. The network is trained by 3D reconstruction, with no keypoint labels, and is made SE(3)-equivariant with Vector Neurons §II-A. A rigid set of query points is posed by minimising the L1 distance between its pose descriptor and that of a demonstration §II-B. In simulation with 10 demos, overall pick-and-place success is 0.87-0.91 on upright objects and 0.58-0.78 on objects in arbitrary poses; the Dense Object Nets baseline reaches 0.00-0.45 tab. I. These are author-reported numbers.

Problem

Keypoint-based pipelines need hand-labelled keypoints for each task, or dense 2D correspondences that are not equivariant to 3D pose and that break down for free-space points and occlusion §I. Imitation from few demonstrations therefore struggles to transfer a grasp or a placement to a new instance in an unseen 6-DoF pose §I. The task-relevant local frame is also under-specified: which part of an object to align to depends on the task §I.

Method

The method has three parts: a point descriptor field, its extension to rigid point sets (pose descriptors), and a few-shot procedure that uses demonstration poses.

Neural point descriptor field §II-A. A function maps a 3D coordinate and an object point cloud to an n-dimensional descriptor §II-A. It is a neural network trained with an occupancy objective (3D reconstruction from a partial cloud), so no correspondence labels are needed §II-A. The descriptor is the concatenation of the activations of all layers of the occupancy network eq. 4. The authors argue that these layers form a coarse-to-fine feature set that encodes distance to the surface and therefore to salient parts such as a mug handle §II-A.

Equivariance §II-A. Subtracting the cloud's mean from both the coordinate and the cloud gives translation equivariance eq. 6. Vector Neurons give rotation equivariance of the encoder and occupancy network eq. 7. Together they give full SE(3) equivariance: transforming the cloud and the coordinate by the same rigid transform transforms the descriptor field accordingly eq. 5.

Point matching §II-A. Given a demonstration point, the test-time point is the minimiser of the descriptor distance between the two clouds eq. 8, eq. 9. The authors show this on a mug handle across poses and instances §II-A, fig. 4, fig. 5.

Pose descriptor field §II-B. A rigid body is represented by a set of query points, given in homogeneous coordinates in a canonical configuration. A pose moves the set. The pose descriptor is the concatenation of the point descriptors of the moved query points eq. 10. The authors choose the query points as uniform samples inside the bounding box of the rigid body (gripper or rack) §II-B, fig. 6. Placement of query points changes which feature the pose is aligned to §II-B, fig. 6.

Pose regression §II-B. To find the pose of a new instance, the authors start from a random pose and optimise translation and axis-angle rotation with Adam, minimising the L1 distance between the descriptor of the candidate pose on the new cloud and the demonstration descriptor eq. 11. The optimisation is shown on a grasp pose fig. 7.

Few-shot pick-and-place §II-C. Each of the K demonstrations is a tuple (cloud, grasp pose, relative pose), and the place pose is the relative pose applied to the grasp pose §II-C. Each pose is encoded as a pose descriptor conditioned on that demonstration's cloud §II-C. Pick and relative-place descriptors are averaged over the K demonstrations §II-C. For a new cloud, the method solves for the pick pose and the relative pose by matching these averages §II-C. Inverse kinematics and motion planning execute the result §II-C. The relative pose on the demonstration side is fixed to identity, so the local frame is the body frame of the demonstrated pose §II-B.

Training data and setup §III. One NDF model is trained across all three categories (mug, bowl, bottle) on partial clouds from four static depth cameras, using a dataset of 100,000 simulated objects at random tabletop poses §III.

# Neural point descriptor field  [§II-A] [eq. 4]
# Inputs:  coordinate x (3,), object cloud P (N, 3)
# Output:  descriptor z (n,)

def point_descriptor(x, P):
    # Centre the coordinate and the cloud on the cloud's mean  [§II-A] [eq. 6]
    mu = mean(P, axis=0)
    x_centred = x - mu
    P_centred = P - mu
    # Encode the cloud with a rotation-equivariant encoder  [§II-A] [eq. 7]
    latent = encoder(P_centred)
    # Keep the activations of every layer of the occupancy network and
    # concatenate them  [§II-A] [eq. 4]
    activations = occupancy_layers(x_centred, latent)
    return concatenate(activations)

# Neural pose descriptor field  [§II-B] [eq. 10]
# Inputs:  pose (4x4 rigid transform), cloud P (N, 3), query points Q (M, 3)
# Output:  pose descriptor, length M * n

def pose_descriptor(pose, P, Q):
    # Move the rigid query set by the pose, in homogeneous coordinates
    # [§II-B]
    moved = apply_pose(pose, Q)
    # Concatenate the point descriptors of all moved query points  [eq. 10]
    return concatenate([point_descriptor(q, P) for q in moved])

# Query points for one rigid body  [§II-B]
# Sampled once and reused for every pose of that body

def sample_query_points(box, n_points):
    # Uniform samples inside the bounding box of the rigid body  [§II-B]
    return uniform_in_box(box, n_points)

# Few-shot pick-and-place  [§II-C]
# demos: K tuples (cloud, pick_pose, rel_pose), where place = rel * pick

def encode_demos(demos, Q_pick, Q_place):
    # Pose descriptors of each demonstration's grasp and relative pose
    # [§II-C]
    picks = [pose_descriptor(d.pick, d.cloud, Q_pick) for d in demos]
    rels = [pose_descriptor(d.rel, d.cloud, Q_place) for d in demos]
    # Average over the K demonstrations  [§II-C]
    return mean(picks, axis=0), mean(rels, axis=0)

# Test-time pose search  [§II-B] [eq. 11]
# target: averaged demonstration descriptor; Q: query points for this body
# The number of starts and steps are not stated in the paper

def solve_pose(cloud, target, Q, n_starts, n_steps):
    # Start from random poses; the optimiser parameterises rotation as
    # axis-angle plus translation  [§II-B]
    best_pose = None
    best_loss = inf
    for pose in random_poses(n_starts):
        for step in range(n_steps):
            # L1 distance between the candidate and the target descriptors
            # [eq. 11]
            loss = l1(pose_descriptor(pose, cloud, Q), target)
            pose = adam_step(pose, loss)
        if loss < best_loss:
            best_pose = pose
            best_loss = loss
    return best_pose

Hyperparameters

Values marked "paper: not stated" come from the code only. Code citations are relative to the official repository at code_commit.

Results

Author-reported. Setup §III: PyBullet simulation with a Franka Panda; ShapeNet meshes for mug, bowl and bottle, with incompatible meshes removed §III. Ten upright demonstrations per task; tests on unseen instances with random initial poses and random uniform scaling §III. Grasp, place and overall success are measured §III. Baseline: Dense Object Nets (DON) with semantic keypoints and SVD registration, trained on 300 RGB-D views per object §III.

Table I tab. I: simulation, unseen instances (grasp / place / overall).

Reported in the text, not in a table: TransporterNets reproduced 92% rim grasping for top-down grasps, but its 6-DoF extension stayed below 10% in several attempts §III. The text attributes the arbitrary-pose drop to the PointNet encoder's sensitivity to unobserved parts of the cloud §IV-A.

Table II tab. II: descriptor ablation, upright mug (grasp / place / overall).

Table III tab. III: query-point scale, upright mug (grasp / place / overall).

Table IV tab. IV: overall success versus number of demonstrations, upright mug.

Real robot §IV-C, fig. 9. Ten upright demonstrations per category, then execution on new mugs, bowls and bottles in varied poses. The paper gives qualitative examples and a supplementary video; no quantitative real-robot numbers are reported.

Paper vs code

The official repository is anthonysimeonov/ndf_robot at code_commit (MIT). The paper's own experiment numbers come from the simulation pipeline, so the differences below matter for reproduction.

• Descriptor normalisation. The paper's descriptor is a plain concatenation eq. 4. The code L2-normalises it per query point when return_features is set, and both evaluation scripts set it vnn_occupancy_net_pointnet_dgcnn.py evaluate_ndf.py. The paper does not say whether its reported numbers used it.
• Input scale. Clouds and coordinates are multiplied by 10 before the network, which the paper does not mention vnn_occupancy_net_pointnet_dgcnn.py vnn_occupancy_net_pointnet_dgcnn.py. This changes the meaning of the query-point scale in Table III.
• Which activations. "All layers" §II-A is implemented as seven tensors: the input features (squared norm, coordinate–latent products and latent invariant), the first linear output, and five residual-block outputs vnn_occupancy_net_pointnet_dgcnn.py. The paper does not list them.
• Query points in the demo pipeline. The paper samples uniformly in the bounding box §II-B. The recorder samples 500 points uniformly in the oriented bounding box of the gripper mesh and of the rack teleop.py. A second scheme, Gaussian with σ = 0.025 around a random cloud point scaled by 1.2, appears in the alignment demo class only ndf_alignment.py ndf_alignment.py. That class is used by the demo script, not by the evaluation script ndf_demo.py.
• Loss. The paper uses an L1 distance eq. 11. The code uses PyTorch's mean-reduced L1 loss, averaged over the random starts optimizer.py optimizer.py.
• Start selection. The best start is chosen by the loss at the last iteration optimizer.py. The paper does not describe a selection rule.
• Start distribution. The paper says "random" §II-B. The code draws rotations from a healpix grid and translations from [0, 0.1) optimizer.py optimizer.py. The alignment demo class uses uniform random angles instead ndf_alignment.py.
• Noise flag is inert. The CLI default for noise is 0.05 evaluate_ndf.py, but the evaluation script does not pass it to the optimiser evaluate_ndf.py. The optimiser's own default is 0.0 optimizer.py. The effective noise is therefore 0.
• Search length. The evaluation script uses 250 iterations by default evaluate_ndf.py, while the alignment class defaults to 500 ndf_alignment.py.
• Encoder. The paper says PointNet §II-A. The code uses a Vector-Neuron ResNet-PointNet with neighbourhood features (k = 20) vnn_occupancy_net_pointnet_dgcnn.py vnn_occupancy_net_pointnet_dgcnn.py, and offers a DGCNN variant.
• Training views. The paper uses four static depth cameras §III. The dataset loader draws a random subset of 1 to 4 views dataio.py. That subset is applied only when the multiview option is on, which is not the default, so the default view set is not confirmed from this code alone.
• Training defaults. The learning-rate help text says 5e-5, but the default is 1e-4 train_vnn_occupancy_net.py.
• Side effects. The optimiser writes HTML visualisations into debug_viz/ and visualization/ in the current directory on every run optimizer.py optimizer.py.

Pitfalls

• Headline wording. The introduction says "overall success rate above 85%" §I. That holds for upright objects only (0.87-0.91). With arbitrary poses the overall success is 0.58-0.78 tab. I.
• Few-shot range. The abstract and introduction say 5-10 demonstrations §I. The main results use 10. Only the mug upright task is studied with 1, 5 or 10 demonstrations tab. IV.
• Scale unit. Table III reports scales from 0.1 to 10.0 without defining the unit, and the best tested value is 1.0 tab. III. The paper does not say whether 1.0 equals the size of the bounding box. Because the network input is multiplied by 10 in code, the effect of the scale also depends on that factor vnn_occupancy_net_pointnet_dgcnn.py.
• Non-convexity. The pose search is a non-convex optimisation with random starts. The paper does not give the number of starts, the step count, or the learning rate. Results depend on these, and the code uses 10 starts (5 with DGCNN) and 250 steps by default optimizer.py.
• Occlusion. The authors attribute the arbitrary-pose drop to the encoder's sensitivity to unobserved and newly visible surfaces §IV-A. Training uses partial clouds, so the effect is part of the trained behaviour.
• Assumptions. Clouds are segmented, and the environment is static between demonstration and test §III. Placement targets are assumed static §VI.
• Rigid objects only. The authors note the method may extend to non-rigid objects but this is untested §VI.
• Trajectories are out of scope. The method transfers poses and points, not full trajectories §VI.
• Equation and figure numbers. Equation, figure and table numbers in this card were counted from the LaTeX source order (figures 1-9, tables I-IV, equations 1-11), not from the rendered PDF. Check them against the PDF before citing externally.
• Real-robot evidence. There are no quantitative real-robot numbers §IV-C.

Lineage

• Builds on: keypoint pipelines that need hand-labelled keypoints (Manuelli et al., 2019; Gao and Tedrake, 2021) §I; dense correspondence with 2D descriptors (Florence et al., 2018) §I; neural occupancy and implicit shape models (Mescheder et al., 2019; Park et al., 2019; Chen and Zhang, 2019) §II-A; Vector Neurons for rotation equivariance (Deng et al., 2021) §II-A. None of these papers has a card in this wiki yet, so no link is given.
• Compared against: Dense Object Nets (Florence et al., 2018) tab. I–IV; TransporterNets (Zeng et al., 2020), attempted but not in the tables §III.
• Superseded by: not known from content already in this wiki.