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 / 500	82.1% (mean)	0.212 (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).

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

Very early snapshot (epoch 5 of ~120 needed for convergence) — newer-epoch checkpoints will be added as training progresses.

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_5 was 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; it is not compatible with the original Param-of-list layout.