| --- |
| license: agpl-3.0 |
| tags: |
| - grasping |
| - jax |
| - equivariance |
| - multi-embodiment |
| --- |
| |
| # Kinematics Flow checkpoints |
|
|
| 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. |
|
|
| 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. |
|
|
| ```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: `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. |
|
|
