Interactive-Active-AI / OneWM

# OneWM
### Can World and Action Modeling Benefit Each Other under Unified Causal Training?
Yue Su*, Bokai Lin*, Xinkai Wang*, Chubin Zhang, Xiangnan Wu, Kuan Zhang, 
Weiyang Jin, Jiwen Yu, Yiming Li, Cewu Lu, Lixin Yang†, Xihui Liu†
MMLab@HKU · MVIG@SJTU · THU · SII · UCAS · IAAI Team 
 * Equal contribution. † Corresponding authors. 
**Project Page · arXiv · PDF · Demos**
**Training and deployment code for LIBERO, RoboCasa365 and dual Franka.**
Original full-resolution PDF
## Overview
**One causal model learns to predict the world and act in it.**
OneWM is a simple framework for unified causal training of world prediction and action generation. Within a shared video–action architecture, actions serve as both **conditions for predicting visual consequences** and **targets for policy learning**. We study whether these two learning objectives can benefit each other while keeping the architecture fixed.
Our experiments on **LIBERO, LIBERO-Plus, RoboCasa365, and real robots** support reciprocal benefits: learning action-conditioned forward dynamics improves policy performance and generalization, while action learning improves multi-view world-prediction fidelity.
We plan to integrate OneWM into the OpenWAM framework soon, making it easier to train and deploy.
## How it works
OneWM connects world modeling and action generation by varying how much of the action sequence is known. A shared video–action architecture with **temporally causal attention** learns from three conditioning regimes:
- **All-clean:** specified actions condition future-video prediction.
- **Partial-clean:** a clean action prefix conditions future-video prediction and denoising of the remaining action suffix.
- **All-noisy:** the model learns to generate actions alongside future video.
At inference, the same model supports action generation from observations and instructions, or world prediction conditioned on specified actions.
## Install
Training and full-model inference require Linux, NVIDIA CUDA and a compatible
PyTorch installation. The commands below use Python 3.11 and CUDA 12.4 wheels;
adapt the PyTorch wheel to your CUDA driver when needed.
 [code] 
The Wan directory must contain its DiT safetensor shards, model configuration,
`Wan2.2_VAE.pth`, `models_t5_umt5-xxl-enc-bf16.pth`, and `google/umt5-xxl` tokenizer.
Prepare the ActionDiT transformer initialization before training:
 [code] 
This interpolates the Wan transformer weights into the action expert. It is an
initialization payload, not a trained policy. Deployment loads the complete
trained checkpoint and does not require this initialization file.
## Prepare data and train
Follow Data preparation for the dataset layout, text-cache
precomputation, RoboCasa365 downloads, and Franka split/statistics preparation.
The YAML files use portable relative paths under `data/`, `checkpoints/`, and
`outputs/`. Execute commands from the repository root.
 [code] 
The launcher calculates gradient accumulation so
`micro_batch × GPU_cou