DreamWAM: Beyond RGB Future Prediction
for World Action Models

Shanglin Yuan^1,2,* · Weiheng Zhao^1,2,* · Xin Shi² · Haoyi Jiang^1,2 · Xianda Guo³ · Liu Liu⁴ · Wenyu Liu¹ · Wei Sui^2,† · Xinggang Wang^1,‡

¹Huazhong University of Science and Technology · ²D-Robotics · ³Wuhan University · ⁴Horizon Robotics

^*Equal contribution · ^†Project Lead · ^‡Corresponding Author

This repository provides the official implementation of DreamWAM, which moves world action modeling beyond RGB by learning future appearance, motion, geometry, and semantics as complementary views of action-relevant state. It combines joint RGB-flow latent denoising with gated depth and DINO residual supervision, transferring structured future cues to action prediction through shared VideoDiT-ActionDiT attention. These beyond-RGB signals are used only during training, preserving RGB-only inference while improving robustness to visual perturbations.

Contents

• Installation
• Model Preparation
• Data Preparation
• Training
• Evaluation
• Acknowledgments
• Citation

Installation

Run all commands from the repository root. Python 3.10 and CUDA 12.8 are recommended.

git clone https://github.com/hustvl/DreamWAM.git
cd DreamWAM

conda create -n dreamwam python=3.10 -y
conda activate dreamwam
conda install -c conda-forge imagemagick -y

pip install torch==2.7.1 torchvision==0.22.1 --index-url https://download.pytorch.org/whl/cu128
pip install -e .

Model Preparation

Place the following upstream repositories and pretrained components at the repository-local targets shown below.

Component	Upstream source or weights	Local target
DreamWAM checkpoints	Hugging Face	checkpoints/
Wan2.2 TI2V-5B	Wan-AI/Wan2.2-TI2V-5B	pretrained/Wan2.2-TI2V-5B/
RAFT	source and weights	third_party/RAFT/, pretrained/raft-things.pth
DINOv2	source and ViT-B/14 register weights	third_party/dinov2/, pretrained/dinov2_vitb14_reg4_pretrain.pth
Depth Anything 3	source and DA3-BASE	third_party/Depth-Anything-3/, pretrained/da3-base/
LIBERO	benchmark	third_party/LIBERO/
LIBERO-Plus	benchmark and assets	third_party/LIBERO-Plus/; extract assets to third_party/LIBERO-Plus/libero/libero/assets/

Install the local benchmark packages after placing their source repositories:

pip install -e third_party/Depth-Anything-3
pip install -e third_party/LIBERO

Download the released checkpoints directly into the configured local directory:

hf download hustvl/DreamWAM \
  dreamwam_joint.pt dreamwam_uncond.pt \
  --local-dir checkpoints

Generate the ActionDiT initialization from Wan2.2 before training:

python scripts/prepare_action_dit.py --config configs/dreamwam_joint.yaml

The generated initialization is written to pretrained/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt.

Data Preparation

DreamWAM uses the four-suite, LeRobot v2.1 LIBERO release prepared by FastWAM: yuanty/LIBERO-fastwam. The dataset was prepared with MuJoCo 3.3.2; use the same version for benchmark consistency.

Download and extract the four suite archives under data/libero:

hf download yuanty/LIBERO-fastwam \
  --repo-type dataset \
  --include "*.tar.gz" \
  --local-dir data/libero

for archive in data/libero/*.tar.gz; do
  tar -xzf "$archive" -C data/libero
done

The resulting layout must contain:

data/libero/
|-- libero_10_no_noops_lerobot/
|-- libero_goal_no_noops_lerobot/
|-- libero_object_no_noops_lerobot/
`-- libero_spatial_no_noops_lerobot/

Precompute the RGB, optical-flow, DINO, and Depth training cache:

python scripts/precompute_cache.py --config configs/dreamwam_joint.yaml

The cache is written to cache/libero_2cam224 and is shared by both released training settings.

Training

accelerate launch --num_processes 8 scripts/train.py --config configs/dreamwam_uncond.yaml
accelerate launch --num_processes 8 scripts/train.py --config configs/dreamwam_joint.yaml

Training writes the resulting checkpoints to outputs/uncond/checkpoint.pt and outputs/joint/checkpoint.pt, respectively. Evaluation loads the checkpoint specified by paths.checkpoint in the selected YAML config; point that field to the corresponding training checkpoint above when evaluating a newly trained model.

Evaluation

LIBERO

Run one of the four official suites by selecting libero_spatial, libero_object, libero_goal, or libero_10:

python scripts/eval_libero.py \
  --config configs/dreamwam_joint.yaml \
  --suite libero_spatial

LIBERO-Plus

The default command runs the official four-suite protocol with one trial per task:

python scripts/eval_libero_plus.py --config configs/dreamwam_joint.yaml

Evaluation results are written as JSON under the corresponding outputs/<setting>/evaluation/ directory.

Acknowledgments

DreamWAM is built on the FastWAM codebase and its coupled VideoDiT-ActionDiT formulation. We thank the FastWAM authors for releasing the base model, training pipeline, and processed LIBERO data.

We also acknowledge the upstream Wan2.2, RAFT, DINOv2, Depth Anything 3, LIBERO, and LIBERO-Plus projects used by this release.

Citation

If you find this repository useful, please consider citing our paper:

@article{yuan2026dreamwam,
  title={DreamWAM: Beyond RGB Future Prediction for World Action Models},
  author={Yuan, Shanglin and Zhao, Weiheng and Shi, Xin and Jiang, Haoyi and Guo, Xianda and Liu, Liu and Liu, Wenyu and Sui, Wei and Wang, Xinggang},
  journal={arXiv preprint arXiv:2608.04996},
  year={2026}
}