EasyWAM README

The project's README from github.com/OpenMOSS/EasyWAM. As one markdown file: README.md.

EasyWAM Logo

🚀 EasyWAM: A Unified and Efficient Framework for Training and Evaluating World Action Models

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

English README 中文 README EasyWAM Website Hugging Face Models WeChat

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

✨ Overview and Key Features

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.

EasyWAM Overview

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

📰 News

🤖 Supported Models and Benchmarks

🧠 Models

Model Full-Parameter Training LoRA Fine-Tuning
EasyWAM-Unified ✅ ✅
EasyWAM-MoT ✅ ✅
EasyWAM-MoT-Joint ✅ ✅
EasyWAM-MoT-IDM ✅ ✅
EasyWAM-Hidden ✅ ✅

🏗️ Backbone

Backbone Supported
Wan2.2-TI2V-5B ✅
Cosmos-Predict2.5-2B ✅
FLUX.2 Klein-4B (ImageWAM-compatible) ✅

See FLUX.2 / ImageWAM backbone integration for checkpoint migration, configuration, and evaluation details.

🧪 Benchmarks

Benchmark Training Evaluation
LIBERO Full-parameter and LoRA Standard evaluation
LIBERO-Plus Uses LIBERO checkpoints Robustness evaluation
RoboTwin Full-parameter and LoRA Clean and randomized evaluation

🏆 Benchmark Results

All benchmark results reported below use Wan2.2-TI2V-5B as the backbone.

LIBERO

Full-Parameter

Model Spatial Object Goal Long Avg.
EasyWAM-Unified 99.0 99.4 99.2 98.2 99.0
EasyWAM-MoT 97.8 98.4 97.6 95.6 97.4
EasyWAM-Hidden 99.4 100.0 97.0 97.8 98.6

LoRA (Rank 128)

Model Spatial Object Goal Long Avg.
EasyWAM-Unified 84.0 97.8 92.0 81.2 88.8
EasyWAM-MoT 96.8 98.8 94.4 90.4 95.1
EasyWAM-Hidden 96.8 99.4 92.6 86.8 93.9
LIBERO-Plus
Model Orig (LIBERO) Background Camera Language Layout Light Noise Robot Avg.
EasyWAM-Unified 99.0 55.8 33.7 93.7 80.6 92.2 50.2 71.4 67.5
EasyWAM-MoT 97.4 52.8 20.6 80.4 65.2 85.1 51.5 49.7 56.8
EasyWAM-Hidden 98.6 56.8 49.2 95.3 81.0 90.4 58.2 77.4 72.4

For an analysis of these benchmark results, see our blog: What WAM Architecture Do We Need? (中文).

⚡ Efficiency Results

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 121.9 samples/s, a 2.37× throughput improvement over FastWAM's 51.5 samples/s. It also reduces data, forward, and backward time per step by 66.0%, 58.4%, and 54.7%, respectively.

EasyWAM and FastWAM training efficiency comparison

Framework Throughput (samples/s) ↑ Data Time (s) ↓ Forward Time (s) ↓ Backward Time (s) ↓
EasyWAM-MoT 121.9 0.0258 0.5162 0.4842
FastWAM 51.5 0.0759 1.2412 1.0680

On LIBERO, training for 20,000 steps takes approximately 6 hours with EasyWAM, compared with approximately 14 hours using the original FastWAM codebase—a 57.1% 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.

🚀 Quick Start

🛠️ Installation

conda create -n easywam python=3.10 -y
conda activate easywam
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 .

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

📦 Prepare Models

Released EasyWAM checkpoints are available in the OpenMOSS-Team/EasyWAM collection on Hugging Face. For checkpoint-specific download commands and usage requirements, see the corresponding model cards.

Run the following commands from the project root. The paths match the values in configs/model/backbone/wan22.yaml and configs/model/backbone/cosmos25.yaml.

mkdir -p checkpoints

# Wan2.2 video DiT, VAE, UMT5 encoder, and tokenizer
huggingface-cli download Wan-AI/Wan2.2-TI2V-5B --local-dir checkpoints/Wan2.2-TI2V-5B

# Cosmos post-trained 2B video DiT and its Wan2.1 video tokenizer
huggingface-cli download nvidia/Cosmos-Predict2.5-2B \
  --include "base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt" "tokenizer.pth" \
  --local-dir checkpoints/Cosmos-Predict2.5-2B

# Reason1 text encoder and tokenizer used to build the Cosmos text cache
huggingface-cli download nvidia/Cosmos-Reason1-7B --local-dir checkpoints/Cosmos-Reason1-7B

After downloading and generating ActionDiT initialization weights, the files used by the default configs are:

checkpoints/
├── Wan2.2-TI2V-5B/
├── Cosmos-Predict2.5-2B/
│   ├── base/post-trained/81edfebe-bd6a-4039-8c1d-737df1a790bf_ema_bf16.pt
│   └── tokenizer.pth
├── Cosmos-Reason1-7B/
├── ActionDiT_Wan22_5B_alphascale_1024hdim.pt
└── ActionDiT_CosmosPredict25_2B_alphascale_1024hdim.pt

EasyWAM-MoT and EasyWAM-Hidden also use an interpolated ActionDiT initialization:

# Writes the exact path referenced by configs/model/backbone/wan22.yaml
python scripts/preprocess_action_dit_backbone.py \
  --model-config configs/model/easywam_mot_wan22.yaml \
  --backbone wan22 \
  --output checkpoints/ActionDiT_Wan22_5B_alphascale_1024hdim.pt \
  --device cuda \
  --dtype bfloat16

# Writes the exact path referenced by configs/model/backbone/cosmos25.yaml
python scripts/preprocess_action_dit_backbone.py \
  --model-config configs/model/easywam_mot_cosmos25.yaml \
  --backbone cosmos25 \
  --output checkpoints/ActionDiT_CosmosPredict25_2B_alphascale_1024hdim.pt \
  --device cuda \
  --dtype bfloat16

EasyWAM-Unified does not require this ActionDiT checkpoint.

📝 Precompute Text Embeddings

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
python scripts/precompute_text_embeds.py task=libero_easywam_mot_cosmos25

🏋️ Train

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
NPROC_PER_NODE=8 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=robotwin_easywam_unified_wan22_lora

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

📊 Evaluate

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

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

The managers default to 8 GPUs. Override MULTIRUN.num_gpus and MULTIRUN.max_tasks_per_gpu to match your machine. See the benchmark guides for installation, data layout, checkpoint examples, filtering, and resume behavior.

📚 Documentation

Documentation English 中文
Documentation index Index 文档索引
Configuration Guide 配置指南
LIBERO Guide 使用指南
LIBERO-Plus Guide 使用指南
RoboTwin Guide 使用指南

🗂️ Repository Layout

EasyWAM/
├── configs/          # Data, model, task, training, and evaluation configs
├── docs/             # Repository documents
│   ├── README.md      # Documentation index
│   ├── blogs/        # Architecture blogs
│   ├── configurations/ # Configuration guides
│   └── instructions/   # Integration and benchmark 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

🤝 Contributing

EasyWAM is built with the community. You can help by fixing bugs, improving documentation, adding model or benchmark support, or sharing ideas that make World Action Model research more accessible.

We are happy to welcome contributions in any form. If you would like to contribute or help build EasyWAM together, please email [email protected]. See docs/CONTRIBUTING.md for the contribution guidelines.

🙏 Acknowledgements

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

📝 Citation

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

Star History

Star History Chart

EasyWAM (OpenMOSS/EasyWAM)

Searching Database…

这是什么?EasyWAM 是一个统一的研究代码库,用于训练、微调和评估世界动作模型(World Action Model,WAM)——这类模型为具身智能和机器人研究联合预测未来的视频与机器人动作。
这是开箱即用的机器人控制器,还是用来构建 WAM 的研究框架?它是一个研究框架,不是即插即用的机器人控制器——它提供模型实现、训练方案和评估管线,让研究者能够构建、微调并对自己的世界动作模型进行基准测试。
我如何用它训练或评估模型?你要从 Hugging Face 上的 OpenMOSS-Team/EasyWAM 集合下载预训练骨干和检查点权重,为基准数据集预先计算文本嵌入,然后用 Hydra 覆盖项和 NPROCPERNODE 这类环境变量启动训练或评估脚本,进行分布式运行。
运行它需要特定的 GPU 配置吗?你需要带 PyTorch 的 Python 以及支持 CUDA 的 GPU;报告中的基准测试使用的是 8×NVIDIA H100 GPU。FlashAttention 2/3/4 是可选的,用于加速注意力计算,如果没有安装则回退到 PyTorch SDPA。
训练实际需要多少算力?即使有 EasyWAM 的优化,在 8 块 H100 GPU 上训练 LIBERO 模型 20,000 步仍需约 6 小时——比原 FastWAM 代码库的 14 小时要快,但仍然是相当重的多 GPU 工作负载,不是能在笔记本上跑的东西。
这到底是为谁准备的?它是为想要训练、微调或对世界动作模型进行基准测试的机器人和具身智能研究者构建的,而不是为想要成品可部署产品的开发者准备的。
EasyWAM 支持哪些骨干?它目前支持 Wan2.2-TI2V-5B 和 Cosmos-Predict2.5-2B 作为视频扩散骨干,另外还有 FLUX.2 Klein-4B 用于 ImageWAM 兼容的变体。
我可以在哪些基准上评估?EasyWAM 提供 LIBERO、LIBERO-Plus(一个侧重鲁棒性的变体)和 RoboTwin 的训练与评估方案。
里面包含哪些模型架构?EasyWAM-Unified 用单一的 Video DiT 骨干联合预测视频和动作,而 EasyWAM-MoT、EasyWAM-MoT-Joint、EasyWAM-MoT-IDM 和 EasyWAM-Hidden 则使用各自独立的视频与动作 DiT 专家,并采用不同的交互设计。
EasyWAM 比原 FastWAM 代码快多少?在 8×H100 GPU 上,EasyWAM-MoT 达到 121.9 samples/s,而 FastWAM 是 51.5 samples/s,吞吐量提升 2.37 倍,并把 LIBERO 的总训练时间缩短约 57%。
Creator
OpenMOSS
Size
3.3 MB
Stars
401
Updated
CommitOct 09
Age
1mo
Commits · 90 Days