Training ManiFlow Policies on Automatically Generated Data#

ManiFlow policies are trained with the maniflow-tb submodule (allenai/maniflow-tb, Python package hvla) and its local launcher, scripts/train_maniflow.py. The script turns its flags into a single training command,

accelerate launch --config_file hvla/config/deepseeds/deepspeed_zero2.yaml ... hvla/training/train_hvla.py \
    --config_yaml ./examples/Franka/train_maniflow_multimodal.yaml --framework.name ManiFlowMultiModal ...

and runs it locally: accelerate launch on all visible GPUs by default, or torchrun across several nodes. It trains on a merged LeRobot dataset, the output of the data generation pipeline. Every flag is listed in the train_maniflow.py Reference.

Setup#

These docs describe maniflow-tb commit c04eab7, the commit the maniflow-tb submodule is pinned to. The maniflow extra installs it in editable mode, together with molmo2 (see Installation and Setup):

git submodule update --init maniflow-tb molmo2
uv pip install -e ".[dev,isaac,torch-cu128,maniflow]"

The dataloader decodes the RGB videos with TorchCodec, which needs the FFmpeg shared libraries (FFmpeg 4 to 7) on the machine. python -c "import torchcodec" should succeed before you train. maniflow-tb’s own maniflow-tb/INSTALL.md describes a standalone conda environment instead, including FFmpeg and the optional muon-optimizer package for --optimizer muon.

Training downloads pretrained weights from the Hugging Face Hub on first use (the timm RGB backbone and t5-small); set HF_HOME to choose the cache location. Metrics go to W&B when credentials are available (WANDB_API_KEY or a wandb login); otherwise W&B runs offline and its logs stay in the run directory.

Training on a generated dataset#

python maniflow-tb/scripts/train_maniflow.py  --data_root_dir data/datasets/merged/train/my_dataset

Relative --data_root_dir and --run_root_dir paths are resolved against the directory you run the script from; the training command itself runs inside maniflow-tb/.

  • --data_root_dir: the merged LeRobot dataset (by default the pipeline writes it to
    <output-root>/datasets/merged/train/<experiment-id>). Pass several directories to co-train
    on them; --dataset_ratio sets their sampling ratio.

Everything else keeps its default, which gives:

  • absolute 9-D joint actions (7 joints, gripper, Play mode), predicting chunks of 50 actions
    and executing 25;

  • a 9-D state (joint positions, gripper, Play mode); The play mode is used my Maniflow to indicate that it has completed the task and request the next instruction.

  • front and gripper RGB through a fine-tuned CLIP ViT-B/16. Depth is supported but off by default.

  • the target heatmap as cross-attention tokens through the shared RGB encoder (dropped with probability 0.1, sigma 0.8 patches);

  • T5-small language conditioning on the action string only (action_only);

  • one node with all visible GPUs, AdamW with DeepSpeed ZeRO-2, a batch of 128 per GPU, 200k steps, a checkpoint every 10k steps and a 5% validation split.

The script prints the run ID, the run directory and the full command before it starts; --dry_run prints them and exits. The run ID is built from the dataset name and the flags, plus a timestamp. --user_prefix puts a prefix in front of it, and --run_id replaces it.

Choosing GPUs and nodes#

--num_gpus_per_node defaults to the number of visible GPUs, so CUDA_VISIBLE_DEVICES picks them. A short debug run on one GPU:

CUDA_VISIBLE_DEVICES=0 WANDB_MODE=offline python maniflow-tb/scripts/train_maniflow.py \
    --run_root_dir /tmp/maniflow_debug --duration 100 --save_interval 100 \
    --logging_frequency 10 --dbs 4 --num_workers 0 --max_episodes 500 \
    --data_root_dir data/datasets/merged/train/my_dataset

For several nodes, run the script on every node with the same flags, plus --num_nodes N --node_rank R --master_addr <address of node 0> --run_id <name>. It then launches torchrun instead of accelerate launch; every node must use the same --run_id (or --resume) so that all of them write to the same run directory, which therefore has to be on a filesystem the nodes share.

The last batch of an epoch#

The train dataloader keeps the last, partial batch of every epoch, and ManiFlow splits every batch into a flow part (int(0.75 * B) samples) and a consistency part (int(0.25 * B) samples). A last batch of 1 to 3 samples leaves the consistency part empty, and training stops at that step with

IndexError: list index out of range

raised from the T5 tokenizer (encode_text_input_T5), because it is handed an empty list of instructions. The train: <episodes> episodes, <N> samples line that [FrankaDataset] prints before training gives N; the last batch holds N mod --dbs samples (per GPU). Pick a --dbs for which that is 0 or at least 4. For example, 242 samples crash at --dbs 16 (last batch: 2) and train at --dbs 32 (last batch: 18). The debug run above, at --dbs 4, crashes whenever N is not a multiple of 4.

N is a property of the dataset, not of the training flags, and it is not fixed for a given datagen command. The pipeline seeds task generation (--task-gen-seed, default 6303), reset state sampling (a fixed 0 with --vecgen-problems-per-env 1, --seed plus the group offset otherwise) and execution (--seed, default 42), so the problems and start states repeat. The executed trajectories are not guaranteed to: vecgen steps GPU physics and builds motion plans on worker threads, and a skill that is retried from the live state (a flip that misses keeps its failed attempt in the episode) changes the episode length. Rerunning the same command can therefore change N, and a --dbs that worked before can stop working. Check the log line after each regeneration.

Resuming and fine-tuning#

  • --resume <run_id> continues a run in --run_root_dir from its latest complete checkpoint, and fails if there is none. The architecture (actions, state, cameras, conditioning, depth) comes from the run’s saved config.yaml, so the flags that define it are not emitted; only the training-loop flags still apply. --data_root_dir is not needed. The run must resume on the same number of GPUs it was saved with, since the optimizer state is sharded per GPU.

  • --pretrained_checkpoint <file> starts a new run from a checkpoint’s weights. The optimizer, scheduler, EMA and W&B state start fresh.

Outputs#

Each run writes to <run_root_dir>/<run_id>/ (--run_root_dir defaults to ./runs):

Path

Contents

config.yaml

The full training config, which the policy server reads.

stats.json

Normalization statistics of the (merged) training data.

checkpoints/steps_<N>_model.safetensors

Weights every --save_interval steps.

checkpoints/steps_<N>_training_state.pt, checkpoints/state_<N>/, checkpoints/steps_<N>.done

Scheduler and EMA state, the per-GPU optimizer state (newest only) and the completion marker --resume looks for.

final_model/

model.safetensors, config.yaml and config.json at the end of training.

summary.jsonl

One line per saved checkpoint.

wandb/, wandb_run_id.txt

W&B logs (also the offline ones) and the W&B run ID.

Metrics go to the W&B project maniflow-tb (--wandb_project, or $WANDB_PROJECT), in your default entity unless --wandb_entity or $WANDB_ENTITY names another.

Evaluating the checkpoint#

Pass a checkpoint to launch_eval.py. It finds config.yaml in the run directory, and the policy server reads the training data’s normalization stats from it:

python scripts/launch_eval.py --policy maniflow \
    --model-path runs/<run_id>/checkpoints/steps_200000_model.safetensors \
    --dataset-path /path/to/lerobot_dataset

The policy server imports hvla from the maniflow-tb submodule, and the eval pipeline refuses to start if that checkout is not the commit ToyBlocks pins (git submodule update --init maniflow-tb fixes it). A checkpoint trained with this code is on Hugging Face, allenai/maniflow-tb, with the same layout (hf download allenai/maniflow-tb --local-dir <dir>). See Policy Evaluation for the rest of the eval workflow, and Setting up a new policy server to evaluate a policy that is not ManiFlow.