Dimios45/kinematics-flow

Kinematics Flow checkpoints

Checkpoints for Kinematics Flow, from "Towards a Multi-Embodied Grasping Agent". These are mid-training checkpoints, not final converged models.

All eval numbers are simulation-based grasp success rate (SR) and normalized joint diversity (NJD), computed via kin_flow.cli.bench on 10 held-out test scenes with 100 sampled grasps per scene.

Models

Single-embodiment models: num_scenes=5000. Multi-embodiment model: the paper's full configuration — num_scenes=25000, fp32, batch 5 scenes × 128 grasps, warmup-cosine LR (peak 3e-4). Training moved to 3x MI300X partway through (after epoch 30), with per-device batch unchanged (5 scenes each) and 3x the scenes processed per optimizer step; epoch numbers here are the corrected true dataset-pass count.

me-full_25000_5 per-gripper eval (epoch 5)

me-full_25000_30 per-gripper eval (epoch 30)

me-full_25000_90 per-gripper eval (epoch 90)

me-full_25000_120 — final model (epoch 120, training complete)

Training finished 2026-07-15 at exactly 120.0 true dataset passes (gradient step 302,000 = 3,000,000 scene-visits / 25,000 scenes), the convergence point recommended by the original repo's README. Saved from the step-302,000 train-state checkpoint, so it also contains optimizer state for further fine-tuning.

Mean SR is flat vs. epoch 90 (86.1% vs 86.4%, within eval noise at 10 scenes × 100 grasps), i.e. the model has converged — consistent with the original repo's guidance that ~120 epochs give reasonable convergence.

Loading

Format: orbax/OCDBT checkpoint directories.

from kin_flow.ctrl.trainer import Trainer
from kin_flow.net.kinematics_flow import KinematicsFlow, KinematicsFlowConfiguration
# build `model` from the repo's train.yaml config, then:
model = Trainer.get_model_from_checkpoint(model, "<path>/me-full_25000_5")

Note: me-full_25000_* checkpoints were trained with flax 0.11 using a per-path nnx.Param layout in TPWithWeightsAndBiases (kin_flow/net/module/fctp.py) — restore with a matching code state; they are not compatible with the original Param-of-list layout.

Seed-conditioned models

Extension of the multi-embodiment model with an equivariant seed point — a 3D point marking which object to grasp — after SeededGrasp (arXiv 2607.20207), expressed in irreps so SE(3) equivariance is preserved by construction.

These are fine-tuned from me-full_25000_120. The seed enters the scene encoder as one virtual graph edge from every node to the seed point, contributing an ℓ=1 direction vector and ℓ=0 distance — the two quantities the encoder already consumes. Only ~10k new parameters (unet.seed_mlp) on an 18.8M-parameter model.

Models

me-seed-anneal_25000_4 has the highest mean across all five grippers (0.3289) and is the recommended default. The -anneal runs use a learning rate annealed 8e-5 → 0; the original schedule held a constant 8e-5 (decay_steps was 10,000 against a 200k-step run) which degraded results over training.

Evaluation

targeted = yield × hit × SR|target — successful grasps on the intended object per raw sample drawn. Measured on 119 held-out cluttered scenes (4 objects each), 100 grasps per scene, both conditions scored on identical grasps (same RNG, same seed point, differing only in whether the seed mask is on).

Seed conditioning raises the fraction of grasps aimed at the intended object from chance (~0.28 for ~3.5 objects/scene) to 0.61–0.79. Collision projection is a separate, training-free post-processing step that recovers the collision-free yield seeding costs.

Supporting numbers — 46 evaluations across 15 checkpoints and 5 grippers. Seeded and unseeded results are identical to 4 decimal places on the pre-fine-tune checkpoint, where the seed pathway is zero-initialised, confirming the conditioning is inert at initialisation.

Caveats

• —Test clouds are analytically sampled from MuJoCo geometry, not rendered scans (the renderer segfaults on our hardware). Validated to 2 mm geometric agreement but carrying a measured −2.7 pp systematic offset — valid for A/B comparison, not directly comparable to the me-full_25000_* numbers above.
• —Performance saturates within ~4,000 steps; epochs 2–8 are statistically indistinguishable and longer training does not help.
• —DexEE is an outlier (+75%, SR|target 0.26 vs 0.55–0.77 elsewhere): it aims correctly but still fails to hold objects.

Loading

Same as the me-full_* checkpoints, but the model must be built with seed conditioning so unet.seed_mlp exists:

# model config needs seed_zero_init; then
model = Trainer.get_model_from_checkpoint(model, "<path>/me-seed-anneal_25000_4",
                                          partial=True)

# inference: seed is (xyz, mask) in world metres; mask=0 => unconditional
se3, dof = inference("Flow", model, sample, num_samples, cfg,
                     seed=(seed_xyz, np.float32(1.0)))

partial=True restores only the parameters present in a checkpoint, which is also how these were created from the unconditional model.