OpenMOSS/EasyWAM

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

EasyWAM – 世界行動モデルの統合フレームワーク

何であるか – EasyWAM は、世界行動モデル（WAM）の学習と評価を可能にするオープンソースの研究用コードベースです。WAM は、マルチモーダル入力から将来の動画フレームとロボットの行動を予測するニューラルネットワークであり、エムベデッド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微調整モードが備わっています。
バックボーンの柔軟性 – 3つの大規模なビジョン・ランゲージバックボーンに対応：Wan2.2-TI2V-5B、Cosmos-Predict2.5-2B、FLUX.2 Klein-4B。

対応モデル（概要）

モデル	アーキテクチャ	フルパラメータ学習	LoRA微調整
EasyWAM-Unified	単一バックボーンのVideo DiT（動画＋行動の統合）	✅	✅
EasyWAM-MoT	二重バックボーン＋混合自己注意（行動専用）	✅	✅
EasyWAM-MoT-Joint	二重バックボーン、統合動画＋行動ノイズ除去	✅	✅
EasyWAM-MoT-IDM	二重バックボーン、教師強制型動画による行動予測	✅	✅
EasyWAM-Hidden	Video 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用のコマンドが記載されています。
バッジにリンクされたHugging Faceコレクションからチェックポイントをダウンロード。
使用するベンチマーク用のテキスト埋め込みを事前計算。例：
```
python scripts/precompute_text_embeds.py task=libero_easywam_mot_wan22
```

学習 – 起動スクリプトとHydraタスク構成を選択。8GPUでのフルパラメータ学習の例：

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

4GPUでの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の数値が表示されています。
貢献 – プロジェクトは典型的なオープンソースフローに従います：バグはIssueを開き、新しいモデルやベンチマークはPRで提出し、両言語のドキュメントを更新してください。

TL;DR

EasyWAM は、ロボットシミュレーションベンチマーク上で世界行動モデルを構築・テストするための研究用高効率フレームワークです。FlashAttention、DeepSpeed、LoRAといった最新の効率化テクニックと、複数の最先端バックボーン用の即時実行構成を統合しており、『新しいWAMのアイデア』から『LIBERO/RoboDojoで評価されたモデル』まで、数回のコマンドで実現できます。