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:

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 uses the Apache-2.0 license.

Downloads last month

-

Downloads are not tracked for this model. How to track
Video Preview
loading

Model tree for shitijgovil/FloMo

Adapter
(9)
this model