File size: 2,829 Bytes
57c45e4 1b1574a 57c45e4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | ---
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.
|