| --- |
| license: mit |
| library_name: pytorch |
| base_model: Wan-AI/Wan2.2-TI2V-5B-Diffusers |
| tags: |
| - robotics |
| - forward-dynamics |
| - video |
| - scene-flow |
| - lora |
| --- |
| |
| # FloMo midtrain base checkpoints |
|
|
| FloMo does any-to-any forward-dynamics cotraining for robot learning from video. |
| This repo has the three midtrain base checkpoints from the paper. Each one is a LoRA |
| on the Wan2.2-TI2V-5B DiT. Code: https://github.com/pairlab/flomo |
|
|
| | Folder | Outputs | Train steps | Size | |
| |---|---|---|---| |
| | `ma/flomo_ma.pt` | motion tokens + action | 13,750 | 646 MB | |
| | `va/flomo_va.pt` | video + action | 14,000 | 646 MB | |
| | `vma/flomo_vma.pt` | video + motion tokens + action | 13,500 | 646 MB | |
|
|
| Motion tokens are dense 3D scene flow encoded by the frozen Wan VAE. All three |
| checkpoints come from epoch 28 of training on the paper mixture: EgoVerse 0.65, |
| MolmoAct2 0.15, RealYAM 0.15, and RealYAM human video 0.05. Action supervision uses |
| the MolmoAct2 and RealYAM robot data only. |
|
|
| ## Contents |
|
|
| Each `.pt` file is a `torch.save` dict with these keys: |
|
|
| - `model`: trainable parameters only (LoRA adapters and FloMo heads). The frozen Wan |
| VAE and umT5 text encoder are loaded from the Wan weights. |
| - `lora_enabled`: `True`. |
| - `normalizer_stats`: action normalizers for `realyam` and `molmoact2`. |
| - `config`: the resolved training config. Absolute paths are set to `???`, and the |
| W&B ids are removed. |
|
|
| The files do not contain optimizer, scheduler, or scaler state. |
|
|
| ## Usage |
|
|
| Install FloMo and the Wan weights as the GitHub README describes. Then: |
|
|
| ```bash |
| hf download shitijgovil/FloMo vma/flomo_vma.pt --local-dir ./flomo_ckpts |
| |
| # Inference: (T, 14) YPR action chunk for a bimanual robot |
| python infer_example.py --checkpoint ./flomo_ckpts/vma/flomo_vma.pt --wan-ckpt-dir ./wan_checkpoints \ |
| --image primary.png --image left.png --image right.png \ |
| --task "pick up the croissant and place it in the oven" \ |
| --ee-pose="x,y,z,yaw,pitch,roll,grip,x,y,z,yaw,pitch,roll,grip" |
| |
| # Fine-tune (weights only), e.g. one RoboTwin 2.0 task |
| python train.py --config-name train_robotwin data_root=<RoboTwin>/data \ |
| robotwin_task=open_laptop checkpoint=./flomo_ckpts/vma/flomo_vma.pt resume=false run_name=robotwin_open_laptop |
| ``` |
|
|
| When you fine-tune VA or VMA on a config with no video output, the loader warns |
| `Unexpected keys: modalities.video_primary.{mask_token,view_pe}`. This is expected. |
|
|
| ## SHA-256 |
|
|
| ``` |
| 8d68b7f66b702c53f826e4be155865df4b1bd002ad4196b577c0754bb617f354 ma/flomo_ma.pt |
| c50dd53732951f338ed28d34f57aa5f2dabd5da98fe495fc3194b5270655ceaa va/flomo_va.pt |
| 196b9faf8d8a58de9ca3559faa77c8144b8a6facbb87e7896803522a31f84e90 vma/flomo_vma.pt |
| ``` |
|
|
| ## License |
|
|
| The FloMo code and these weights use the MIT license. The base model |
| [Wan2.2-TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) uses the Apache-2.0 license. |
|
|
