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

Folder	Type	Gripper(s)	Epoch	SR	NJD	Hardware
se-panda_5000_170	single-embodiment	Panda (2 DOF)	170 / 500	97.8%	0.293	RTX 6000 Ada
se-shadow_5000_40	single-embodiment	Shadow Hand (22 DOF)	40 / 500	75.9%	0.232	RTX 6000 Ada
me-full_25000_5	multi-embodiment	all 5 + z0	5 / 120	82.1% (mean)	0.212 (mean)	MI300X (ROCm)
me-full_25000_30	multi-embodiment	all 5 + z0	30 / 120	84.3% (mean)	0.203 (mean)	MI300X (ROCm)
me-full_25000_90	multi-embodiment	all 5 + z0	90 / 120	86.4% (mean)	0.212 (mean)	MI300X (ROCm)
me-full_25000_120	multi-embodiment	all 5 + z0	120 / 120	86.1% (mean)	0.209 (mean)	MI300X (ROCm)

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)

Gripper	DOF	SR	NJD
Panda	2	94.9%	0.273
VX300	2	93.3%	0.157
DexEE	12	65.7%	0.109
Allegro	16	81.3%	0.287
Shadow Hand	22	75.3%	0.232
mean		82.1%	0.212

me-full_25000_30 per-gripper eval (epoch 30)

Gripper	DOF	SR	NJD
Panda	2	96.0%	0.252
VX300	2	93.3%	0.159
DexEE	12	76.0%	0.095
Allegro	16	84.4%	0.283
Shadow Hand	22	71.6%	0.224
mean		84.3%	0.203

me-full_25000_90 per-gripper eval (epoch 90)

Gripper	DOF	SR	NJD
Panda	2	96.8%	0.265
VX300	2	96.7%	0.155
DexEE	12	73.1%	0.109
Allegro	16	86.4%	0.289
Shadow Hand	22	79.2%	0.243
mean		86.4%	0.212

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.

Gripper	DOF	SR	NJD
Panda	2	97.5%	0.257
VX300	2	97.4%	0.134
DexEE	12	73.0%	0.114
Allegro	16	85.4%	0.288
Shadow Hand	22	77.3%	0.250
mean		86.1%	0.209

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

Folder	Fine-tune	Epoch	Best for	seeded targeted
me-seed-anneal_25000_4	annealed LR	4	best overall · Panda · Allegro	0.4582 (Panda)
me-seed-anneal_25000_6	annealed LR	6	VX300	0.4339
me-seed-anneal_25000_8	annealed LR	8	DexEE	0.1599
me-seed_25000_2	constant LR	2	Shadow Hand	0.3502

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).

Gripper	DoF	unseeded	seeded	+ collision projection	total gain
Panda	2	0.2067	0.4611	0.5190	+151%
VX300	2	0.2027	0.4410	0.4752	+134%
DexEE	12	0.0914	0.1578	0.1597	+75%
Allegro	16	0.1261	0.2601	0.3220	+155%
Shadow Hand	22	0.1510	0.3060	0.3813	+152%
mean		0.1556	0.3252	0.3714	+139%

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.