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 미세조정 모드를 모두 지원합니다.
백본 유연성 – 세 가지 대규모 비전-언어 백본과 호환: 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용 명령어가 나와 있습니다.
배지에 링크된 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의 수치가 표시되어 있습니다.
기여 – 프로젝트는 일반적인 오픈소스 흐름을 따릅니다: 버그는 이슈를 열고, 새로운 모델/벤치마크는 PR로 제출하며, 양국어 문서를 함께 업데이트하세요.

TL;DR

EasyWAM은 로봇 시뮬레이션 벤치마크에서 세계 행동 모델을 구축하고 테스트하기 위한 연구용 고성능 프레임워크입니다. FlashAttention, DeepSpeed, LoRA와 같은 최신 효율 기술을 통합하고, 여러 최첨단 백본용 즉시 실행 구성 파일을 제공하여, '새로운 WAM 아이디어'에서 'LIBERO/RoboDojo에서 평가된 모델'까지 몇 가지 명령어로 빠르게 전환할 수 있습니다.