unified model card: 3 checkpoints with eval results, incl. me-full epoch-5 per-gripper eval

README.md

@@ -9,30 +9,43 @@ tags:
 
 # Kinematics Flow checkpoints
 
-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.
+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.
 
-## Checkpoints
+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.
 
-| 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 |
+## Models
 
-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).
+| 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) |
 
-Format: orbax/OCDBT checkpoint directories, loadable via `kin_flow.ctrl.trainer.Trainer.get_model_from_checkpoint`.
+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).
 
-## Multi-embodiment checkpoint
+### `me-full_25000_5` per-gripper eval (epoch 5)
 
-- `me-full_25000_5/` — **epoch 5** (~25,000 train steps) of the paper's full
-  multi-embodiment configuration: all five grippers (Panda, VX300, DexEE, Allegro,
-  ShadowHand) + z0, `num_scenes=25000`, fp32, batch 5 scenes × 128 grasps,
-  warmup-cosine LR (peak 3e-4). Trained on a single AMD MI300X (ROCm).
-  **Very early snapshot** — the paper trains ~120 epochs for full quality; expect
-  rough grasps. Newer-epoch checkpoints will replace/join this one as training
-  progresses.
+| 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** |
 
-### Loading
+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.
 
 ```python
 from kin_flow.ctrl.trainer import Trainer
@@ -41,6 +54,6 @@ from kin_flow.net.kinematics_flow import KinematicsFlow, KinematicsFlowConfigura
 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.
+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.