| --- |
| license: mit |
| library_name: pytorch |
| pipeline_tag: robotics |
| tags: |
| - robotics |
| - robot-learning |
| - imitation-learning |
| - embodied-ai |
| - world-model |
| - vision-language-action |
| - robotwin |
| --- |
| |
| # Flex-π — RoboTwin 2.0 (3-camera, 384×320) |
|
|
| Flex-π checkpoint for **RoboTwin 2.0**, trained jointly on all 50 tasks. |
|
|
| - Paper: [arXiv:2608.10860](https://arxiv.org/abs/2608.10860) |
| - Project page: <https://flex-pi.github.io/> |
| - Code: <https://github.com/geyan21/flex-pi> |
|
|
| ## Results |
|
|
| RoboTwin 2.0, success rate (%) over 50 tasks, as reported in the paper (Table 1). |
| *Clean* and *Randomized* are background conditions; the two rows are **inference |
| regimes served by these same weights**. |
|
|
| | Inference regime | Clean | Randomized | Avg. | |
| | --- | --- | --- | --- | |
| | Action-only | 94.5 | 94.6 | 94.6 | |
| | Full joint | 94.3 | 94.8 | 94.6 | |
|
|
| `K = 4` Euler denoising steps throughout. Full joint additionally denoises the |
| future-video, DINO, and pointmap streams; action-only skips them, trading them |
| for cheaper inference at the same average success rate. Selecting a regime is an |
| inference-time flag — no retraining, no separate weights. |
|
|
| ## Architecture |
|
|
| A Mixture-of-Transformers pairing a video DiT with an action DiT, coupled by |
| HBridge. Alongside actions the model can denoise three auxiliary streams: |
| future video, DINO features, and pointmaps. |
|
|
| | | | |
| | --- | --- | |
| | Video expert | Wan2.2-TI2V-5B, 5.00 B params | |
| | Action expert | ActionDiT, 1.02 B params | |
| | Layers | 30 (HBridge: 7 bottom / 16 middle / 7 top) | |
| | Semantic encoder | DINOv3 `vit_base_patch16_dinov3.lvd1689m`, 768-d, frozen | |
| | Cameras | `cam_high`, `cam_left_wrist`, `cam_right_wrist` @ 240×320 | |
| | Composite video | 384×320, 33 frames | |
| | Action / proprio | 14-d each (bimanual ALOHA-AgileX), `ConcatLeftAlign` | |
| | Action : video rate | 4:1 | |
|
|
| ## Training |
|
|
| | | | |
| | --- | --- | |
| | Data | 2,500 clean + 25,000 randomized demos, all 50 tasks | |
| | Epochs | 6 | |
| | Learning rate | 1e-4 | |
| | Precision | bf16 | |
|
|
| Trained with flex-joint sampling at `p = 0.5` on every present and joint flag, |
| with cross-modal prediction enabled for all three streams. That is what lets one |
| set of weights serve any regime in the results table above. |
|
|
| ## Files |
|
|
| ```text |
| config.yaml # architecture + processor; autoloaded by the eval |
| dataset_stats.json # action/state normalization statistics |
| checkpoints/weights/step_048060.pt # 12 GB |
| ``` |
|
|
| Keep this directory layout. The eval locates `config.yaml` and |
| `dataset_stats.json` by walking up from the checkpoint path. |
|
|
| ## Usage |
|
|
| This repository holds the **policy weights only**. The Wan2.2 base components |
| and the ActionDiT backbone are separate downloads, resolved through |
| `DIFFSYNTH_MODEL_BASE_PATH` — see [`docs/INSTALL.md`](https://github.com/geyan21/flex-pi/blob/main/docs/INSTALL.md) |
| and [`docs/ROBOTWIN.md`](https://github.com/geyan21/flex-pi/blob/main/docs/ROBOTWIN.md). |
|
|
| ```bash |
| hf download flex-pi/flexpi-robotwin --local-dir runs/flexpi-robotwin |
| |
| export DIFFSYNTH_MODEL_BASE_PATH="$(pwd)/checkpoints" # Wan2.2 weights |
| ``` |
|
|
| Then set the checkpoint at the top of `scripts/eval_flexpi_robotwin.sh`: |
|
|
| ```bash |
| CKPT="./runs/flexpi-robotwin/checkpoints/weights/step_048060.pt" |
| DATASET_STATS="./runs/flexpi-robotwin/dataset_stats.json" |
| ``` |
|
|
| and run: |
|
|
| ```bash |
| bash scripts/eval_flexpi_robotwin.sh |
| ``` |
|
|
| The launcher defaults reproduce the **full joint** row: `NUM_INFERENCE_STEPS=4`, |
| `INSTRUCTION_TYPE=unseen`, `EVAL_NUM_EPISODES=100`, and all six regime flags |
| `true`. For the **action-only** row, set the three `INFER_JOINT_*` flags to |
| `false`. `PHASES=clean,random` covers both background conditions. |
|
|
| ## License |
|
|
| MIT — see [LICENSE](https://github.com/geyan21/flex-pi/blob/main/LICENSE). |
|
|
| ## Citation |
|
|
| ```bibtex |
| @article{yan2026flexpi, |
| title = {Flex-$\pi$: A Multi-Stream World-Action Model with Compute Flexibility}, |
| author = {Yan, Ge and Liu, Jinghao and Fan, Yuzhi and Cai, Lei and Liao, Minwen |
| and Zhang, Jesse and Fox, Dieter}, |
| journal = {arXiv preprint arXiv:2608.10860}, |
| year = {2026}, |
| url = {https://arxiv.org/abs/2608.10860} |
| } |
| ``` |
|
|