OpenMOSS/EasyWAM

A unified framework for training, fine-tuning, and evaluating World Action Models

EasyWAM – 世界行动模型的统一框架

是什么 – EasyWAM 是一个开源研究代码库，可让您训练和评估 世界行动模型（WAMs）。WAMs 是神经网络，能够从多模态输入中预测未来的视频帧和机器人动作，使具身AI代理能够在模拟环境中进行规划和行动。

为何重要 – 训练此类模型计算量巨大，且需要大量工程工作（数据加载、视频解码、分布式训练、检查点处理、在多个机器人模拟器上进行评估）。EasyWAM 将所有这些组件整合到一个统一、模块化的框架中，使研究人员可以专注于模型构思，而非工程开销。

核心特性（如 README 所述）

统一且模块化的设计 – 模型、数据、训练和评估共享通用接口，便于插入新的 WAM 架构。
性能优先 – 内置对 FlashAttention 2/3/4、BF16、梯度检查点、DeepSpeed ZeRO-1/2 和 LoRA 的支持。作者报告称，在 8 × H100 GPU 上，相比原始 FastWAM 代码，速度提升了 2.7 倍。
端到端流水线 – 稀疏视频解码、索引文本缓存、持久化工作进程、提示缓存和可恢复评估，减少重复工作。
基于 Hydra 的配置 – 所有超参数均以易于覆盖的 YAML 文件表达；启动脚本可直接接受 Hydra 覆盖参数。
多基准支持 – 提供 LIBERO、LIBERO-Plus、RoboTwin、RoboDojo 和 RoboCasa365 的即用型配方，每项均支持全参数和 LoRA 微调模式。
后端灵活性 – 支持三种大型视觉-语言后端：Wan2.2-TI2V-5B、Cosmos-Predict2.5-2B 和 FLUX.2 Klein-4B。

支持的模型（快速概览）

模型	架构	全参数训练	LoRA 微调
EasyWAM-Unified	单后端视频 DiT（联合视频 + 动作）	✅	✅
EasyWAM-MoT	双后端混合自注意力（仅动作）	✅	✅
EasyWAM-MoT-Joint	双后端，联合视频 + 动作去噪	✅	✅
EasyWAM-MoT-IDM	双后端，教师强制视频用于动作预测	✅	✅
EasyWAM-Hidden	视频 DiT 特征条件化独立的 Action DiT	✅	✅

快速上手（来自 README）

创建 conda 环境

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

安装 PyTorch（CUDA 12.8）和该包

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 以获得最快的注意力内核。README 中列出了 FA-2、FA-3 和 FA-4 的安装命令。
从 README 中徽章链接的 Hugging Face 集合下载检查点。

为将要使用的基准预计算文本嵌入，例如：

python scripts/precompute_text_embeds.py task=libero_easywam_mot_wan22

训练 – 选择一个启动脚本和 Hydra 任务配置。例如在 8 个 GPU 上进行全参数训练：

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

在 4 个 GPU 上进行 LoRA 训练：

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

评估 – 运行适当的管理脚本并指向检查点：
```
python experiments/libero/run_libero_manager.py \
    task=libero_easywam_mot_wan22 \
    ckpt=<path/to/checkpoint.pt>
```
LIBERO-Plus、RoboTwin、RoboDojo 和 RoboCasa365 也存在类似的命令。

接下来可以查看

文档 – docs/ 文件夹包含后端准备、数据布局、基准设置和配置参考的详细指南（中英文）。
基准结果 – 查看 docs/results/result.md 获取完整表格；README 已展示 LIBERO 和 LIBERO-Plus 的数据。
贡献 – 项目遵循典型的开源流程：发现 bug 请开 issue，提交新模型/基准的 PR，中英文文档同步更新。

TL;DR

EasyWAM 是一个研究级、高度优化的框架，用于在机器人模拟基准上构建和测试世界行动模型。它集成了现代高效技术（FlashAttention、DeepSpeed、LoRA）和多个前沿后端的即用型配置，让您从“我有一个新的 WAM 想法”到“在 LIBERO/RoboDojo 上评估的训练模型”只需几条命令即可完成。