OpenMOSS/EasyWAM

A unified research codebase for training and evaluating World Action Models.

Start here: Quick start · Supported models · Benchmark results · Checkpoints · Documentation

EasyWAM is a unified research codebase designed to make World Action Model development efficient, reproducible, and easy to extend. It connects model implementation, data processing, distributed training, parameter-efficient fine-tuning, and large-scale evaluation through a consistent workflow.

• 🧩 Unified and modular design. Modular model components share consistent data, training, checkpoint, and evaluation interfaces, making new WAM designs easier to integrate.
• ⚡ Efficient computation. Efficiency is a first-class design goal in EasyWAM. It natively integrates FlashAttention 2/3/4 to accelerate attention workloads and provides complete LoRA support from parameter-efficient training and checkpointing to merged inference. BF16, gradient checkpointing, DeepSpeed ZeRO, and PyTorch SDPA fallback further improve speed, memory usage, and compatibility.
• 🚄 Optimized end-to-end pipeline. EasyWAM streamlines every stage of training and inference. Sparse video decoding, indexed text caches, persistent workers, prompt caching, and resumable evaluation eliminate repeated work and deliver substantial speedups across both training and inference.
• 🛠️ Easy-to-use workflows. Hydra-based configuration, standardized training and evaluation recipes, distributed launchers, automatic GPU sharding, and result summaries keep common workflows straightforward.

🌟 Hope: We hope EasyWAM will become an efficient and easy-to-use codebase for World Action Model research, enabling researchers to explore WAMs more quickly and easily. More models and benchmarks will be continuously added and supported. We warmly welcome contributions from the community to help make EasyWAM better. If you encounter any problems or have suggestions for improving EasyWAM, please open an issue. We will continue to refine and improve EasyWAM.

• [2026-09-29] EasyWAM adds official LIBERO-Pro evaluation; introduces selectable clean and random RoboTwin training modes; standardizes the data pipeline on LeRobot v3 while simplifying data preprocessing, text caching, and model runtime paths; aligns data directories and LIBERO naming; and updates evaluation concurrency defaults and benchmark result presentation.
• [2026-09-18] EasyWAM adds RoboCasa365 and RoboDojo training and evaluation, updates RoboTwin for the current upstream with independent dynamic batch serving, and scales evaluation with multiple model workers per GPU and a GPU tuning tool. Training gains configurable checkpoint retention, state tokens after actions, and generic FLUX text embedding precomputation. This update also refreshes LIBERO and LIBERO-Plus results for Unified, MoT, and Hidden on Wan2.2 and Cosmos2.5, including Wan2.2 LoRA on LIBERO, with revised architecture analysis and evaluation defaults.
• [2026-09-12] EasyWAM introduces dynamic batched evaluation, requiring only a small number of workers while significantly reducing evaluation GPU memory usage and improving GPU utilization. This update also adds task/trial progress reporting, LeRobot v3 support, FlashAttention 2/3/4 with optimized text-padding mask semantics, configurable state-token placement and causal attention, execution and caching optimizations, automatic run logging, expanded documentation, and new benchmark results.
• [2026-09-03] EasyWAM adds FLUX.2/ImageWAM backbone integration.
• [2026-09-02] EasyWAM is released with unified training and evaluation workflows for World Action Models.

• EasyWAM-Unified. A single-backbone architecture that places video and action tokens in one Video DiT to jointly predict future video and actions. The architecture is based on DreamZero.
• EasyWAM-MoT. A dual-backbone model with separate Video DiT and Action DiT experts whose tokens interact through shared mixed self-attention. It performs action-only prediction and is based on FastWAM.
• EasyWAM-MoT-Joint. A dual-backbone model that jointly denoises video and action tokens through shared mixed self-attention.
• EasyWAM-MoT-IDM. A dual-backbone model that uses teacher-forced conditional video for action prediction.
• EasyWAM-Hidden. A dual-backbone architecture that uses intermediate Video DiT features as conditional input to a separate Action DiT. The architecture is based on DiT4DiT.

All results below are reported with state_position: sequence, where state/proprio tokens are placed after the action tokens in the model sequence.

See the complete benchmark results for more results and What WAM Architecture Do We Need? (中文) for analysis.

All efficiency results reported below use Wan2.2-TI2V-5B as the backbone. Actual training time may vary with the CPU and GPU configuration; these figures are provided for reference only. Please refer to the efficiency configuration guide to choose settings appropriate for your machine and improve training and inference efficiency.

EasyWAM-MoT and FastWAM use the same model architecture, enabling an architecture-matched comparison between the EasyWAM training framework and the original FastWAM codebase. Measured on 8 × NVIDIA H100 GPUs with a per-device batch size of 16, EasyWAM-MoT achieves 138.2 samples/s, a 2.68× throughput improvement over FastWAM's 51.5 samples/s. It also reduces data, forward, and backward time per step by 85.8%, 70.5%, and 49.5%, respectively.

On LIBERO, training for 20,000 steps takes approximately 5 hours with EasyWAM, compared with approximately 14 hours using the original FastWAM codebase—a 64.3% reduction in overall training time.

🌟 Hope: Training and evaluating World Action Models often requires substantial computational resources. EasyWAM aims to lower this barrier with an efficient and lightweight codebase, enabling researchers to train models, run evaluations, and iterate quickly even with limited compute. Through continuous efficiency improvements, we hope researchers can devote more of their resources to model and algorithm innovation and that more members of the community can participate in WAM research.

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

FlashAttention is optional. When installed, EasyWAM uses the fastest compatible implementation available and otherwise falls back to PyTorch SDPA.

Install the implementation supported by your GPU. FA2 supports Ampere, Ada, and Hopper GPUs; FA3 targets Hopper GPUs; FA4 targets Hopper and Blackwell GPUs.

# FlashAttention 2
pip install flash-attn --no-build-isolation

# FlashAttention 3
git clone https://github.com/Dao-AILab/flash-attention.git
cd flash-attention/hopper
python setup.py install

# FlashAttention 4
pip install flash-attn-4

# FlashAttention 4 with CUDA 13
# pip install "flash-attn-4[cu13]"

Released EasyWAM checkpoints are available in the OpenMOSS-Team/EasyWAM collection on Hugging Face.

Prepare the selected backbone with its dedicated guide:

• Wan2.2-TI2V-5B
• Cosmos-Predict2.5-2B
• FLUX.2 Klein Base 4B

Run this once after preparing a benchmark dataset:

python scripts/precompute_text_embeds.py task=libero_easywam_mot_wan22
# Or: python scripts/precompute_text_embeds.py task=robotwin_easywam_mot_wan22
# Or: python scripts/precompute_text_embeds.py task=robocasa_easywam_mot_wan22
# Or: python scripts/precompute_text_embeds.py task=robodojo_easywam_mot_wan22
python scripts/precompute_text_embeds.py task=libero_easywam_mot_cosmos25

Existing caches are reused by default; if all prompts are cached, the text encoder is not loaded. Add +overwrite=true to regenerate embeddings.

The launchers accept Hydra overrides directly. Set the number of local processes through NPROC_PER_NODE:

# DeepSpeed ZeRO-1 on 8 local GPUs
NPROC_PER_NODE=8 bash scripts/train_zero1.sh task=libero_easywam_mot_wan22

# MoT, MoT-Joint, and MoT-IDM each have a corresponding Cosmos25 task config
bash scripts/train_zero1.sh task=libero_easywam_mot_cosmos25

# DeepSpeed ZeRO-2 LoRA training on 4 local GPUs
NPROC_PER_NODE=4 bash scripts/train_zero2.sh task=libero_easywam_unified_wan22_lora

# RoboDojo joint-only LeRobot v3 training
bash scripts/train_zero2.sh task=robodojo_easywam_mot_wan22

scripts/train_zero2_offload.sh enables ZeRO-2 CPU offload. Multi-node runs additionally use NNODES, NODE_RANK, MASTER_ADDR, and MASTER_PORT.

# LIBERO
python experiments/libero/run_libero_manager.py \
  task=libero_easywam_mot_wan22 \
  ckpt=<path/to/checkpoint.pt>

# LIBERO-Plus (uses a LIBERO checkpoint)
python experiments/libero_plus/run_libero_plus_manager.py \
  task=libero_easywam_mot_wan22 \
  ckpt=<path/to/checkpoint.pt>

# LIBERO-Pro (uses a LIBERO checkpoint)
python experiments/libero_pro/run_libero_pro_manager.py \
  task=libero_easywam_mot_wan22 \
  ckpt=<path/to/checkpoint.pt>

# RoboTwin
python experiments/robotwin/run_robotwin_manager.py \
  task=robotwin_easywam_mot_wan22 \
  ckpt=<path/to/checkpoint.pt>

# RoboDojo
python experiments/robodojo/run_robodojo_manager.py \
  task=robodojo_easywam_mot_wan22 \
  ckpt=<path/to/checkpoint.pt>

# RoboCasa365
python experiments/robocasa/run_robocasa_manager.py \
  task=robocasa_easywam_mot_wan22 \
  ckpt=<path/to/checkpoint.pt> \
  EVALUATION.dataset_stats_path=<path/to/dataset_stats.json>

Evaluation concurrency is controlled by MULTIRUN.num_gpus, MULTIRUN.gpu_ids, MULTIRUN.workers_per_gpu, and MULTIRUN.env_num_per_worker. Every model worker loads its own model copy and owns an independent dynamic inference batcher, so increasing workers_per_gpu also increases GPU memory use. MULTIRUN.inference_batch_size and MULTIRUN.inference_batch_wait_ms apply per worker. See the data and benchmark guides for dataset layout, simulator installation, checkpoint examples, filtering, and resume behavior.

EasyWAM/
├── configs/          # Data, model, task, training, and evaluation configs
├── docs/             # Repository documents
│   ├── README.md      # Documentation index
│   ├── blogs/        # Architecture blogs
│   ├── results/      # Benchmark result summaries
│   └── instructions/   # Data, benchmark, backbone, and configuration guides
├── experiments/      # Benchmark evaluators
├── scripts/          # Training, preprocessing, and caching entrypoints
├── src/              # Models, data pipeline, runtime, and trainer
├── checkpoints/      # External and trained checkpoints
├── data/             # Local datasets and text caches
└── runs/             # Training outputs

EasyWAM is built with the community, and contributions of all sizes are welcome. You can help by fixing bugs, improving documentation and examples, supporting new models or benchmarks, optimizing training and evaluation, or sharing reproducible results.

For bugs and small improvements, feel free to open an Issue or Pull Request. For substantial features or changes that may affect public interfaces, configurations, or compatibility, please open an Issue first so the scope and design can be discussed. Pull Requests should stay focused, include relevant verification, and update both English and Chinese documentation when user-facing behavior changes.

See CONTRIBUTING.md for reporting guidelines, development expectations, and the Pull Request checklist.

This project builds on code from FastWAM, and draws inspiration and references from DreamZero, DiT4DiT, and ImageWAM. Thanks to all the teams above for their valuable contributions to the open-source community.

We welcome you to cite EasyWAM's experimental results and codebase in your research. If EasyWAM is useful in your research, please cite:

@misc{easywam2026,
  title  = {EasyWAM: A Unified and Efficient Framework for Training and Evaluating World Action Models},
  author = {EasyWAM-Team},
  year   = {2026},
  url    = {https://github.com/OpenMOSS/EasyWAM}
}

If you have any questions or suggestions, feel free to join our WeChat discussion group.