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