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