model card

README.md

@@ -1,23 +1,37 @@
 ---
 license: agpl-3.0
 tags:
-  - grasping
-  - jax
-  - equivariance
-  - multi-embodiment
+- robotics
+- grasping
+- flow-matching
+- jax
+- flax
 ---
 
-# Kinematics Flow checkpoints
+# kinematics-flow — multi-embodiment grasp generation (early checkpoint)
 
-Single-embodiment checkpoints for [Kinematics Flow](https://github.com/boschresearch/kinematics-flow), from ["Towards a Multi-Embodied Grasping Agent"](https://arxiv.org/abs/2510.27420). These are **mid-training checkpoints**, not final converged models.
+Checkpoint of the full multi-embodiment model from
+[Towards a Multi-Embodied Grasping Agent](https://arxiv.org/abs/2510.27420)
+(code: [boschresearch/kinematics-flow](https://github.com/boschresearch/kinematics-flow)),
+trained on AMD MI300X (ROCm) with the paper's full configuration:
+all five grippers (Panda, VX300, DexEE, Allegro, ShadowHand) + z0, 25,000 scenes,
+fp32, batch 5 scenes × 128 grasps, warmup-cosine LR (peak 3e-4).
 
-## Checkpoints
+## Contents
 
-| Folder | Gripper | DOF | Epoch | Test success rate | Test NJD |
-|---|---|---|---|---|---|
-| `se-panda_5000_170` | Panda | 2 | 170 / 500 | 97.8% | 0.293 |
-| `se-shadow_5000_40` | Shadow Hand | 22 | 40 / 500 | 75.9% | 0.232 |
+- `me-full_25000_5/` — Orbax (OCDBT) checkpoint at **epoch 5** (~25,000 train steps).
+  **Early checkpoint** — the paper trains ~120 epochs for full quality; this is a
+  work-in-progress snapshot, expect rough grasps.
 
-Both trained with `num_scenes=5000` on a single RTX 6000 Ada (48GB). Eval numbers are simulation-based grasp success rate and normalized joint diversity (NJD), computed via `kin_flow.cli.bench` on 10 held-out test scenes (100 sampled grasps/scene each).
+## Loading
 
-Format: orbax/OCDBT checkpoint directories, loadable via `kin_flow.ctrl.trainer.Trainer.get_model_from_checkpoint`.
+```python
+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: trained with flax 0.11 and a per-path `nnx.Param` layout in
+`TPWithWeightsAndBiases` (`kin_flow/net/module/fctp.py`) — restore with a matching
+code state; checkpoints are not compatible with the original Param-of-list layout.