LeapLabTHU / Simple-WAM

# Simple-WAM
Renping Zhou*, Zanlin Ni*, Zihao Fan, Guohao Fu, Zeyu Liu, Hao Shi, Jie Zhang, Chi Bene Chen, Yang Yue, Xueyang Fu, Gao Huang† \
(* indicates equal contribution, † means Co-corresponding author) 
This repository contains the official implementation of the paper "What Makes World Action Models Generalize?
An Empirical Study of Test-Time Future Modeling".
## Environment Setup
### 1. Create the Environment
 [code] 
### 2. Install Dependencies
Core model and data dependencies are pinned in
requirements.txt; benchmark-specific environments are not included.
 [code] 
## Dataset Download
These trajectory datasets are needed for training, not for
simulation evaluation with the released checkpoints.
### LIBERO
Download the dataset from
yuanty/LIBERO-fastwam
into `data/libero_mujoco3.3.2/`, then extract each archive:
 [code] 
The existing data configs expect this layout:
 [code] 
### RoboTwin
Download all `robotwin2.0.tar.gz.part-*` files from
yuanty/robotwin2.0-fastwam
into `data/robotwin2.0/`. Concatenate the parts in filename order and extract:
 [code] 
Expected layout:
 [code] 
## Evaluation
### 1. Download Checkpoints and Simulator Assets
Download the released checkpoints and simulator assets from
rpzhou/Simple-WAM:
 [code] 
If the files are already downloaded, skip the download command. Expected layout:
 [code] 
Wan VAE, text encoder, and tokenizer files are downloaded automatically on first
use if missing. For offline evaluation, prepare those files in the same
`checkpoints/` model-cache directory beforehand. The trained VideoDiT and
ActionDiT weights come from the released `.pt` file.
### 2. LIBERO
Install the simulator dependencies into the existing environment.
 [code] 
The archive restores the omitted `third_party/LIBERO/libero/` directory, including
assets, BDDL files, and initial states. Its paths are resolved through the bundled
`third_party/LIBERO/.libero/config.yaml`.
Evaluate the full-data checkpoint on all four suites, with 50 trials per task:
 [code] 
Use the same command with a different `config_name` to evaluate another released
checkpoint. For task generalization, replace `MULTIRUN.task_suite_names` with the
held-out suite:
| Checkpoint/config suffix | Evaluation suites |
| --- | --- |
| `_fewshot5`, `_fewshot10` | All four suites |
| `_taskgen_spatial`, `_taskgen_video_spatial` | `[libero_spatial]` |
| `_taskgen_object`, `_taskgen_video_object` | `[libero_object]` |
| `_taskgen_goal`, `_taskgen_video_goal` | `[libero_goal]` |
| `_taskgen_long`, `_taskgen_video_long` | `[libero_10]` |
All suffixes above follow `libero_joint_2cam224_1e-4`. Results are saved under
`evaluate_results/libero/ /`.
### 3. LIBERO-plus
Install the simulator dependencies listed in **LIBERO** above, then add the
image-corruption dependencies and extract the LIBERO-plus archive:
 [code] 
The printed path must be under `third_party/LIBERO-plus`. Both benchmarks use
the same Python package name, `libero`; installing LIBERO-plus replaces the
editable LI