OpenMOSS/EasyWAM

EasyWAM – Unified Framework for World Action Models

What it is – EasyWAM is an open‑source research codebase that lets you train and evaluate World Action Models (WAMs). WAMs are neural networks that predict future video frames and robot actions from multimodal inputs, enabling embodied‑AI agents to plan and act in simulated environments.

Why it matters – Training such models is computationally heavy and requires a lot of plumbing (data loading, video decoding, distributed training, checkpoint handling, evaluation on many robot simulators). EasyWAM puts all of those pieces together in a single, modular framework, so researchers can focus on model ideas instead of engineering overhead.

Core features (as described in the README)

• Unified, modular design – model, data, training, and evaluation share a common interface, making it easy to plug in new WAM architectures.
• Performance‑first – built‑in support for FlashAttention 2/3/4, BF16, gradient checkpointing, DeepSpeed ZeRO‑1/2, and LoRA. The authors report a 2.7× speedup over the original FastWAM code on 8 × H100 GPUs.
• End‑to‑end pipeline – sparse video decoding, indexed text caches, persistent workers, prompt caching, and resumable evaluation reduce redundant work.
• Hydra‑based configuration – all hyper‑parameters are expressed in easy‑to‑override YAML files; launch scripts accept Hydra overrides directly.
• Multi‑benchmark support – ready‑to‑run recipes for LIBERO, LIBERO‑Plus, RoboTwin, RoboDojo, and RoboCasa365, each with full‑parameter and LoRA fine‑tuning modes.
• Backbone flexibility – works with three large vision‑language backbones: Wan2.2‑TI2V‑5B, Cosmos‑Predict2.5‑2B, and FLUX.2 Klein‑4B.

---

Supported models (quick glance)

Model	Architecture	Full‑param training	LoRA fine‑tuning
EasyWAM‑Unified	Single‑backbone Video DiT (joint video + action)	✅	✅
EasyWAM‑MoT	Dual‑backbone with mixed self‑attention (action‑only)	✅	✅
EasyWAM‑MoT‑Joint	Dual‑backbone, joint video + action denoising	✅	✅
EasyWAM‑MoT‑IDM	Dual‑backbone, teacher‑forced video for action prediction	✅	✅
EasyWAM‑Hidden	Video DiT features condition a separate Action DiT	✅	✅

---

Getting started (from the README)

1. Create a conda environment

   conda create -n easywam python=3.10 -y
   conda activate easywam
   

2. Install PyTorch (CUDA 12.8) and the package

   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 .
   

3. (Optional) Install FlashAttention for the fastest attention kernels. The README lists commands for FA‑2, FA‑3, and FA‑4.
4. Download a checkpoint from the Hugging Face collection linked in the badge.
5. Pre‑compute text embeddings for the benchmark you will use, e.g.:

   python scripts/precompute_text_embeds.py task=libero_easywam_mot_wan22
   

6. Train – pick a launch script and a Hydra task config. Example for full‑parameter training on 8 GPUs:

   NPROC_PER_NODE=8 bash scripts/train_zero1.sh task=libero_easywam_mot_wan22
   

   For LoRA training on 4 GPUs:

   NPROC_PER_NODE=4 bash scripts/train_zero2.sh task=libero_easywam_unified_wan22_lora
   

7. Evaluate – run the appropriate manager script, pointing to the checkpoint:

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

   Similar commands exist for LIBERO‑Plus, RoboTwin, RoboDojo, and RoboCasa365.

---

Where to look next

• Documentation – the docs/ folder contains detailed guides for backbone preparation, data layout, benchmark setup, and configuration reference (both English and Chinese).
• Benchmark results – see docs/results/result.md for full tables; the README already shows LIBERO and LIBERO‑Plus numbers.
• Contribution – the project follows a typical open‑source flow: open issues for bugs, submit PRs for new models/benchmarks, and update docs in both languages.

---

TL;DR

EasyWAM is a research‑grade, highly optimized framework for building and testing world‑action models on robot simulation benchmarks. It bundles modern efficiency tricks (FlashAttention, DeepSpeed, LoRA) with ready‑made configs for several state‑of‑the‑art backbones, letting you go from “I have a new WAM idea” to “trained model evaluated on LIBERO/RoboDojo” in a few commands.