SE(3)-DiffusionFields: Learning smooth cost functions for joint grasp and motion optimization through diffusion

In plain words

A robot arm that has to pick up a mug needs to know which hand poses will hold it well, and how to move the arm there without bumping into anything. Usually these two problems are solved one after the other: first pick a grasp, then plan an arm motion to reach it. If the grasp turns out to be unreachable, the planner fails and must start again.

This paper teaches the robot a "goodness score" for hand poses, learned from examples of successful grasps, and makes that score smooth everywhere, so that it tells the optimiser which way to move. The score is a diffusion model: it learns to push random hand poses toward the places where good grasps live. Because the score is smooth, it can be added to the other costs of a motion plan, such as staying smooth and avoiding collisions, and the whole plan is improved by gradient steps at once. The grasp and the motion are then chosen together.

TL;DR

• Learns a 6DoF (SE(3)) grasp energy with a diffusion model, trained by denoising score matching on the Lie group §III-A §III-B.
• Uses that energy as one cost term in a trajectory optimiser that improves many random trajectories at once with Langevin-style steps §III-C.
• Author-reported: joint optimisation with 25 particles matches a decoupled baseline with 800 particles in simulated clutter picking fig. 5.
• Code is MIT-licensed PyTorch for grasp training and sampling; the trajectory optimiser is not in that repository.

Problem

Grasp selection and motion planning are often decoupled: a grasp generator proposes poses, then trajectory optimisation tries to reach one. A generated grasp can be infeasible for the arm, clutter or joint limits, so the second stage fails §I §IV-B.

Learned cost functions could remove this split, since a scalar cost can be summed with other costs. But cost functions learned with cross-entropy or contrastive divergence have flat, zero-slope regions far from the data, so gradient descent cannot move samples toward good grasps §I §IV-A. The authors assume the object pose is known, and pose estimation and perception are outside the scope of the experiments §III-B §IV.

Method

Diffusion on SE(3). The model is a scalar energy E(H, k) over a grasp pose H in SE(3) and a noise level k. Its score is the negative gradient of the energy with respect to the pose, and that gradient is what the sampler follows §III-A. Noise is applied on the group: a pose is perturbed as H times the exponential of a 6-vector drawn from a Gaussian with standard deviation sigma_k eq. 4. The training loss matches the energy gradient to the score of that perturbed Gaussian on the Lie group eq. 5. The score uses the inverse left Jacobian of SE(3) and the adjoint of H^{-1} app. A. The authors compute the gradient of the Lie-group log density with automatic differentiation (Theseus with PyTorch) §III-A.

Sampling. Samples start from a random pose and take steps of the form in eq. 6: an exponential map of a drift term (half the squared step size times the score, which is the negative energy gradient) plus Gaussian noise, applied on the left. Each step uses the noise schedule of the current level. The sampler returns the final poses alg. 3 app. B.

Architecture. For an object with a learned shape code z_m, a set of N points is placed around the gripper and transformed by the grasp pose and by the object pose into the object frame §III-B fig. 3. A feature encoder F maps each point and the noise level k, conditioned on z_m, to an SDF value and additional features. The features of all points are flattened and a decoder D outputs the scalar energy e §III-B. The SDF head is trained jointly with a supervised SDF loss, which the authors state improves grasp generation §III-B alg. 2. Shape codes are learned per object, as in auto-decoder models app. D.

Training. Each iteration samples a minibatch of objects, perturbs successful grasps, and minimises the sum of the denoising loss and the SDF loss alg. 1 alg. 2. Only successful grasps are used app. B. The object pose is set to identity for training alg. 2.

Grasp and motion optimisation. Trajectories are sequences of T joint configurations. The total cost is a weighted sum of the learned energy evaluated at the end-effector pose, plus heuristic terms such as smoothness, table and box collisions, initial-configuration fixing, pregrasp approach and grasp-place similarity §III-C app. C.3. A set of initial trajectories is evolved with the inverse diffusion update in eq. 7, and the lowest-cost particle at the final noise level is executed alg. 4 §III-C. Initial trajectories are straight lines in joint space towards random configurations app. C.3.

# Grasp energy training, denoising score matching  [alg. 1] [alg. 2]
# Input:  grasps (B, 4, 4): successful grasp poses of one object  [§III-B]
#         shape_code (latent): learned per-object code  [§III-B]
# Output: scalar loss

def training_step(grasps, shape_code):
    # 1. Draw a noise level and perturb each grasp on SE(3)  [§III-A] [eq. 4]
    noise_level = sample_noise_level()        # sigma_k
    noise = randn(B, 6) * noise_level         # (B, 6)
    noisy = grasps @ expmap(noise)            # (B, 4, 4)

    # 2. Map N gripper points into the object frame  [§III-B]
    points = noisy @ gripper_points           # (B, N, 3)
    points = object_pose_inverse @ points     # (B, N, 3)

    # 3. Encode points, conditioned on shape code and noise level  [§III-B]
    sdf, features = encoder(points, shape_code, noise_level)
    energy = decoder(flatten(features))       # (B,)

    # 4. Denoising score matching against the Lie-group Gaussian score  [eq. 5]
    target = lie_score(noisy, grasps, noise_level)  # (B, 6)
    loss_dsm = mean(abs(grad_energy(energy, noisy) + target))

    # 5. Add the SDF regression loss on the same object  [alg. 2]
    loss_sdf = mse(sdf, sdf_labels)

    return loss_dsm + loss_sdf

# SE(3) grasp sampling by annealed Langevin steps  [eq. 6] [alg. 3]
# Input:  energy E(H, k); noise levels sigma_1..sigma_L; step rate eps
# Output: n_s grasp poses

def sample_grasps(energy, n_s, sigmas, eps):
    # 1. Start from random SE(3) poses  [app. C.1]
    poses = expmap(randn(n_s, 6) * sigmas[L])

    # 2. Walk noise levels from k = L down to 1  [alg. 3]
    for k in reversed(range(1, L + 1)):
        # 3. Step size grows with the noise level  [alg. 3] (eps not stated)
        alpha = eps * sigmas[k] / sigmas[L]
        # 4. Gradient step on the energy, plus Gaussian noise  [eq. 6]
        grad = grad_energy(energy, poses, k)  # (n_s, 6)
        noise = randn(n_s, 6)
        poses = expmap(-0.5 * alpha**2 * grad + alpha * noise) @ poses

    return poses

# Joint grasp and motion optimisation by inverse diffusion  [eq. 7] [alg. 4]
# Input:  cost J(tau, k) = energy at final pose + heuristic costs  [app. C.3]
#         n_s initial trajectories, each (Q, T); steps; step rate eps
# Output: lowest-cost trajectory

def optimise_trajectories(cost, n_s, T, sigmas, eps):
    # 1. Straight-line initial trajectories to random configurations  [app. C.3]
    trajs = init_straight_lines(n_s, T)       # (n_s, Q, T)

    # 2. Walk noise levels from k = L down to 1  [alg. 4]
    for k in reversed(range(1, L + 1)):
        # 3. Step size scaled by the noise level  [alg. 4]
        alpha = eps * sigmas[k] / sigmas[L]
        # 4. Gradient step on the total cost, plus noise  [eq. 7]
        grad = grad_cost(cost, trajs, k)      # (n_s, Q, T)
        drift = -0.5 * alpha**2 * grad        # (n_s, Q, T)
        trajs = trajs + drift + alpha * randn_like(trajs)

    # 5. Return the particle with the lowest cost at k = 1  [alg. 4]
    return trajs[argmin(cost(trajs, 1))]

Hyperparameters

Values not listed in the paper are marked. Where a value comes from the code, the code path is cited and the paper's text should be read alongside it.

Results

Author-reported. Setup: Acronym mugs in Isaac Gym, with a VAE, a classifier and a VAE-plus-refinement baseline, and a 6DoF grasp classifier in trajectory optimisation §IV-A §IV-B.

The paper reports the grasp-generation success and EMD values only in fig. 4; the text gives their ranking, not numbers. The real-robot runs use a different mug from the one the model was trained on, and the mug pose comes from an external motion-capture system §IV-C. The authors attribute the higher real-world than simulated success to simpler real scenes; their simulation includes flying obstacles that the real scene lacks §IV-C.

Paper vs code

Repository: TheCamusean/grasp_diffusion at code_commit 59b633c. The paper links it in the caption of fig. 3. Its LICENSE is MIT, and the copyright line names Julen Urain, the first author. The code covers grasp-energy training, grasp sampling and the point-cloud variants. It does not contain the trajectory optimiser or the robot experiments of §IV-B and §IV-C; those are described only in the paper.

• Denoising loss scale. Paper writes a norm in the DSM loss with no constant eq. 5. Code uses mean L1 divided by 20 denoising_loss.py.
• Noise levels. Paper uses a discrete set k in {0..L} with sigma_k §III-A alg. 1. Code samples a continuous time t in [1e-5, 1] and uses sigma = 0.5 in the closed-form marginal standard deviation denoising_loss.py denoising_loss.py.
• SDF loss. Paper's algorithm uses mean squared error alg. 1 alg. 2. Code uses L1 on SDF values clipped to [-10, 0.6] sdf_loss.py sdf_loss.py.
• Sampler step size. Paper sets alpha_k = eps sigma_k / sigma_L alg. 3. Code uses a fixed base of 1e-3, a ratio of squared standard deviations and noise scaled by 0.5 grasp_samplers.py grasp_samplers.py. The grasp sampler also performs an extra noise-free fitting phase grasp_samplers.py.
• Sampler defaults. Grasp generation script uses T=30, T_fit=20 and k_steps=2 generate_6d_grasp_poses.py. The point-cloud script uses T=70, T_fit=50, k_steps=2 generate_pointcloud_6d_grasp_poses.py.
• Learning rate. Training passes lr=1e-4 to the trainer, but the optimiser is built with per-module rates 5e-4, 1e-3, 1e-3 train_6d_grasp_diffusion.py. The trainer only uses its own lr when no optimiser is given trainer.py.
• Schedule never decays. Learning-rate schedules are evaluated only at epoch 0, so the step decay in the config is not applied train_6d_grasp_diffusion.py.
• Energy network. Paper gives no widths; code uses 8 layers of width 512 with weight norm, dropout 0.2 and a 128-dimensional shape code params.json.
• Points. Paper uses an unspecified N points; code uses 30 points with a fixed location and scale params.json.
• Time conditioning. The energy takes the noise level as input to the feature encoder grasp_dif.py.
• Trajectory experiments. Not in the repository, so the cost weights and trajectory settings in tab. II to tab. V cannot be checked against code.

No discrepancies were checked for the trajectory optimiser, since no code was found for it.

Pitfalls

• Six-vector order and convention. Paper writes Expmap of a 6-vector without stating the order. Code puts translation first and rotation second denoising_loss.py. Mixing the order gives wrong gradients silently.
• Mixed units in the Gaussian. An isotropic sigma mixes metres and radians. The paper notes a factor of 10 on the position error of a desired-grasp term to account for metres versus radians app. C.2.
• Left Jacobian and adjoint. The paper derives the score with the left Jacobian and says the right Jacobian could also be used; it does not use it app. A.
• Sign of the score. The model's gradient is the negative of the energy gradient, so the loss compares the energy gradient to the negative of the target score eq. 5 denoising_loss.py.
• Classifier fails for a reason. Grasp classifiers trained with cross-entropy have flat regions far from data, which stalls gradient steps. The authors present the diffusion energy as avoiding this §I §IV-A.
• Step sizes are unspecified. The paper gives neither eps nor L, and the code's sampler uses a different step rule (see Paper vs code).
• Cost balance. The authors list the difficulty of weighting many cost terms as a limitation §IV-C. All weights in tab. II to tab. V are hand-set.
• Compute. Run time is not reported in the paper. Each sampler step needs one energy gradient per particle alg. 3 alg. 4.
• Perception. The object pose and shape are assumed known §III-B §IV. Real-world runs use an external motion-capture system for the mug pose §IV-C.

Lineage

• Builds on: score-based generative modelling and annealed Langevin sampling §II. The Lie-group Gaussian follows Chirikjian §II app. A.
• Compared against: 6DoF-GraspNet style VAEs with refinement §IV-A app. D, a grasp classifier used as an energy §IV-B, and the OMG-Planner and CHOMP-based decoupled pipelines §IV-B.
• Related: learned cost functions for grasp and motion §V.