Official codebase for Fast-WAM: Do World Action Models Need Test-time Future Imagination?

This repository contains the training and evaluation code for FastWAM on LIBERO / RoboTwin.

What's New

FastWAM is now faster, better suited to large-scale datasets, and more flexible for research. This update brings substantially faster training and inference, native LeRobot v3.0 support, and a new model that can switch between acting with and without future imagination.

⚡ Approximately 2x faster inference

End-to-end FastWAM inference is now approximately 2x faster, including text encoding and VAE encoding:

• NVIDIA H20: 470 ms → 210 ms
• NVIDIA RTX 4090: 190 ms → 110 ms

The accelerated path is enabled by default for LIBERO with EVALUATION.compile_action_infer=true. We gratefully acknowledge PR #43 for proposing the optimization ideas that inspired this work. Existing checkpoints remain fully compatible with the accelerated inference path.

🚀 Approximately 10% faster training

FastWAM training is approximately 10% faster on NVIDIA H20 GPUs. The new training path combines a compiled denoising core with batched VAE encoding and a lightweight CUDA Graph backend. Enable denoising compilation with:

bash scripts/train_zero1.sh 8 task=libero_uncond_2cam224_1e-4 \
  model.compile_training_denoise=true

Both cached text embeddings and on-the-fly T5 encoding are supported. The latter skips text-cache preprocessing and is more convenient, at the cost of approximately 10% lower training throughput.

📦 Native LeRobot 2.1 and 3.0 support

FastWAM now supports both LeRobot 2.1 and LeRobot 3.0 datasets. LeRobot 3.0's chunked parquet and video layout scales better to large datasets, with faster data loading and dataset-statistics computation as the dataset grows.

Download the released LeRobot 3.0 LIBERO dataset from Hugging Face and select the v3.0 data config:

huggingface-cli download yuanty/LIBERO-fastwam \
  --repo-type dataset \
  --include "lerobot_v30/**" \
  --local-dir ./data
python scripts/train.py task=libero_uncond_2cam224_1e-4 \
  data=libero_2cam_lerobot_v30

For another LeRobot 3.0 dataset, copy configs/data/libero_2cam_lerobot_v30.yaml, update train.dataset_dirs, and select the new config with data=<config_name>. Existing LeRobot 2.1 configs continue to work unchanged.

🧠 Optional IDM: one model, two thinking modes

Optional IDM is a new FastWAM variant that supports two inference modes in a single model:

• IDM mode: imagine the future video first, then predict actions.
• First-frame mode (Fast-WAM): skip test-time future imagination and predict actions directly from the current observation.

Download the released Optional IDM checkpoint from Hugging Face:

huggingface-cli download yuanty/fastwam \
  libero_optional_idm_2cam224.pt \
  libero_optional_idm_2cam224_dataset_stats.json \
  --local-dir ./checkpoints/fastwam_release

Train the optional-IDM variant once:

bash scripts/train_zero1.sh 8 task=libero_optional_idm_2cam224_1e-4

Then choose either inference mode at evaluation time without retraining, making it easy to study when future imagination helps:

python experiments/libero/run_libero_manager.py \
  task=libero_optional_idm_2cam224_1e-4 \
  ckpt=./checkpoints/fastwam_release/libero_optional_idm_2cam224.pt \
  EVALUATION.dataset_stats_path=./checkpoints/fastwam_release/libero_optional_idm_2cam224_dataset_stats.json \
  EVALUATION.sigma_shift=1.0 \
  +EVALUATION.action_infer_mode=idm \
  MULTIRUN.num_gpus=8

Replace idm with first_frame to use the Fast-WAM inference mode.

The released checkpoint, trained with action scheduler shift 1.0, achieves the following success rates on the full LIBERO benchmark (40 tasks, 50 episodes per task):

Other improvements

• The action scheduler shift now defaults to 1.0 for both training and evaluation; shifts from 1.0 to 3.0 perform similarly in our experiments. When evaluating the original released checkpoints, set EVALUATION.sigma_shift=5.0 to reproduce the original setting.
• Upgraded LIBERO evaluation with persistent model workers, dynamic task scheduling, bad-GPU quarantine, failure recovery, and resumable results.
• Optimized action-only inference for IDM and Optional IDM: infer_action returns actions and video latents directly, while VAE decoding is performed only by infer_joint when video output is requested, eliminating redundant computation in action-only deployments.

Index

• File Structure
• Environment Setup
• Model Preparation
• Dataset Download
• Inference with Released Checkpoints
• Training
• Inference with Your Trained Checkpoints
• Acknowledgements
• BibTeX

File Structure

FastWAM/
├── configs/
│   ├── data/                 # Dataset configs (LIBERO, RoboTwin, etc.)
│   ├── model/                # Model architecture and component configs
│   └── task/                 # Task-level configs (training task names)
├── scripts/
│   ├── train.py
│   ├── train_zero1.sh        # Deepspeed zero1 training entrypoint
│   ├── preprocess_action_dit_backbone.py  # Preprocess ActionDiT backbone before training
│   └── precompute_text_embeds.py  # Precompute T5 text embedding cache before training
├── experiments/
│   ├── libero/
│   │   └── run_libero_manager.py
│   └── robotwin/
│       └── run_robotwin_manager.py
├── src/fastwam/              # Core code
├── runs/                     # Training outputs (ckpt, logs)
├── checkpoints/              # Pretrained or external checkpoints
├── data/                     # Data directory
└── evaluate_results/         # Inference / evaluation results

Environment Setup

conda create -n fastwam python=3.10 -y
conda activate fastwam
pip install -U pip
pip install torch==2.7.1+cu128 torchvision==0.22.1+cu128 --extra-index-url https://download.pytorch.org/whl/cu128
pip install -e .

Model Preparation

This step is required before both training and inference.

Step 1: set the Wan model directory first (opional, default ./checkpoints):

mkdir -p checkpoints
export DIFFSYNTH_MODEL_BASE_PATH="$(pwd)/checkpoints"

Step 2: pre-generate the ActionDiT backbone (interpolated from Wan22 DiT):

# uncond (fastwam)
python scripts/preprocess_action_dit_backbone.py \
  --model-config configs/model/fastwam.yaml \
  --output checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt \
  --device cuda \
  --dtype bfloat16

Dataset Download

LIBERO

The preprocessed LIBERO dataset used by Fast-WAM is available at:

• https://huggingface.co/datasets/yuanty/LIBERO-fastwam

Download all compressed files first, then extract them all:

mkdir -p data/libero_mujoco3.3.2
cd data/libero_mujoco3.3.2
# Run after downloading all 4 tar.gz files
for f in *.tar.gz; do
  tar -xzf "$f"
done

The extracted directory structure should be:

data/libero_mujoco3.3.2/
├── libero_10_no_noops_lerobot/
├── libero_goal_no_noops_lerobot/
├── libero_object_no_noops_lerobot/
└── libero_spatial_no_noops_lerobot/

RoboTwin

The preprocessed RoboTwin dataset used by Fast-WAM is available at:

• https://huggingface.co/datasets/yuanty/robotwin2.0-fastwam

Download all split archive files first, then concatenate and extract:

mkdir -p data/robotwin2.0
cd data/robotwin2.0
# Run after downloading all robotwin2.0.tar.gz.part-* files
cat robotwin2.0.tar.gz.part-* | tar -xzf -

The extracted directory structure should be:

data/robotwin2.0/
└── robotwin2.0/
    ├── data/
    ├── meta/
    └── videos/

If you also keep:

data/robotwin2.0/dataset_stats.json

in the root directory, it can be used directly as the statistics file for the current configs in this repo. You can also recompute it.

Inference with Released Checkpoints

The released checkpoints and their corresponding dataset stats are available on Hugging Face.

Optional: download released checkpoints and dataset stats from Hugging Face:

pip install -U huggingface_hub
huggingface-cli download yuanty/fastwam \
  libero_uncond_2cam224.pt \
  libero_uncond_2cam224_dataset_stats.json \
  libero_optional_idm_2cam224.pt \
  libero_optional_idm_2cam224_dataset_stats.json \
  robotwin_uncond_3cam_384.pt \
  robotwin_uncond_3cam_384_dataset_stats.json \
  --local-dir ./checkpoints/fastwam_release

After downloading, the local directory is expected to contain:

checkpoints/fastwam_release/
├── libero_uncond_2cam224.pt
├── libero_uncond_2cam224_dataset_stats.json
├── libero_optional_idm_2cam224.pt
├── libero_optional_idm_2cam224_dataset_stats.json
├── robotwin_uncond_3cam_384.pt
└── robotwin_uncond_3cam_384_dataset_stats.json

Before running the LIBERO benchmark, install the official LIBERO environment first from the LIBERO repository. Then run this final step:

pip install mujoco==3.3.2

The mujoco environment should ideally stay consistent with the LIBERO data version.

We have already copied the RoboTwin evaluation-related code into third_party/RoboTwin. You still need to follow the official RoboTwin instructions from the RoboTwin repository to finish environment installation and download the required assets, then create the policy symlink:

ln -sfn "$(pwd)/experiments/robotwin/fastwam_policy" "$(pwd)/third_party/RoboTwin/policy/fastwam_policy"

Optional: evaluate released LIBERO checkpoint:

The released LIBERO / RoboTwin evaluation managers default to 8 GPUs (MULTIRUN.num_gpus=8 in configs/sim_libero.yaml and configs/sim_robotwin.yaml). If you want to evaluate with fewer GPUs, pass a smaller value such as MULTIRUN.num_gpus=4.

python experiments/libero/run_libero_manager.py \
  task=libero_uncond_2cam224_1e-4 \
  ckpt=./checkpoints/fastwam_release/libero_uncond_2cam224.pt \
  EVALUATION.dataset_stats_path=./checkpoints/fastwam_release/libero_uncond_2cam224_dataset_stats.json \
  EVALUATION.sigma_shift=5.0 \
  MULTIRUN.num_gpus=8

Optional: evaluate released RoboTwin checkpoint:

python experiments/robotwin/run_robotwin_manager.py \
  task=robotwin_uncond_3cam_384_1e-4 \
  ckpt=./checkpoints/fastwam_release/robotwin_uncond_3cam_384.pt \
  EVALUATION.dataset_stats_path=./checkpoints/fastwam_release/robotwin_uncond_3cam_384_dataset_stats.json \
  EVALUATION.sigma_shift=5.0 \
  MULTIRUN.num_gpus=8

For faster RoboTwin evaluation, we have enabled EVALUATION.skip_get_obs_within_replan=true in configs/sim_robotwin.yaml. This skips RGB rendering while consecutively executing an action chunk within one replan window, which speeds up evaluation but makes the saved video look very low-FPS. Set it to false if you want to save a fully rendered video.

Note: We evaluate with unseen instructions, following Motus. Lingbot-VA uses seen instructions instead. You can try EVALUATION.instruction_type=seen to use seen instructions, which should theoretically improve performance by one or two points.

Training

1) Precompute T5 embedding cache before training

Use scripts/precompute_text_embeds.py to precompute embeddings for each training task:

# LIBERO
python scripts/precompute_text_embeds.py task=libero_uncond_2cam224_1e-4
# RoboTwin
python scripts/precompute_text_embeds.py task=robotwin_uncond_3cam_384_1e-4

For multi-GPU:

torchrun --standalone --nproc_per_node=8 scripts/precompute_text_embeds.py task=libero_uncond_2cam224_1e-4

2) Training (using fastwam as an example)

When running a new task for the first time, set pretrained_norm_stats in the corresponding configs/data/*.yaml to null first. After one training run, a dataset_stats.json file will be generated in the current run directory (for example, runs/{task_name}/{run_id}/dataset_stats.json). You can then update pretrained_norm_stats to that file path for subsequent runs.

# LIBERO
bash scripts/train_zero1.sh 8 task=libero_uncond_2cam224_1e-4
# RoboTwin
bash scripts/train_zero1.sh 8 task=robotwin_uncond_3cam_384_1e-4

For LIBERO, we train on a single node with 8 GPUs. For RoboTwin, we use 64 GPUs to accelerate training. You can try reducing the GPU count or training epochs.

Inference with Your Trained Checkpoints

The mujoco environment should ideally stay consistent with the LIBERO data version. Then run LIBERO evaluation:

# LIBERO
python experiments/libero/run_libero_manager.py task={task_name} ckpt={ckpt_path}

We have already copied the RoboTwin evaluation-related code into third_party/RoboTwin. You still need to follow the official RoboTwin instructions from the RoboTwin repository. Finish installation and download the required assets, then create the policy symlink:

ln -sfn "$(pwd)/experiments/robotwin/fastwam_policy" "$(pwd)/third_party/RoboTwin/policy/fastwam_policy"

Then run RoboTwin evaluation:

python experiments/robotwin/run_robotwin_manager.py task={task_name} ckpt={ckpt_path}

Common task_name examples:

libero_uncond_2cam224_1e-4
robotwin_uncond_3cam_384_1e-4

Acknowledgements

The RoboTwin evaluation code in this repository is adapted from the official RoboTwin repository. We thank the RoboTwin team for releasing their codebase and assets.

BibTeX

If you find our work helpful, please consider citing:

@article{yuan2026fastwam,
  title={Fast-WAM: Do World Action Models Need Test-time Future Imagination?},
  author={Tianyuan Yuan and Zibin Dong and Yicheng Liu and Hang Zhao},
  journal={arXiv preprint arXiv:2603.16666},
  year={2026},
  url={https://arxiv.org/abs/2603.16666}
}