OSWAM

An Effective World Action Model with One-Step Action Generation

Guowen Zhang^* · Liyi Chen^* · Xu Hu · Pengfei Wang
Jingfan Chen · Chenhang He · Lei Zhang^†

The Hong Kong Polytechnic University
^* Equal contribution. ^† Corresponding author.

Structured imagination. One-step action.

Overview · Installation · Checkpoints · NAVSIM · RoboTwin · Citation

Official implementation of OSWAM for autonomous driving on NAVSIM and embodied control on RoboTwin 2.0, including privileged teacher training, student distillation, and evaluation.

🔍 Overview

OSWAM builds an observation-conditioned action prior with compact Action Imagination Tokens (AIT) and Scene Imagination Tokens (SIT) in a pretrained Video DiT. An Action DiT refines this prior in one denoising step. During training, a privileged teacher transfers future-conditioned features to the student; inference uses the current observation without a future-video rollout.

🔧 Installation

Run the commands below from the repository root. OSWAM requires Linux and Python 3.10 or later. First install a compatible PyTorch / torchvision pair for your accelerator, then install the project:

conda create -n oswam python=3.10 -y
conda activate oswam

# Install a compatible PyTorch/torchvision build for your hardware first.
pip install -r requirements.txt

For NAVSIM, install the nuPlan and geospatial dependencies following the NAVSIM v1.1 setup and nuPlan devkit, then register the bundled v1 devkit:

pip install -e ./navsim --no-deps

For RoboTwin, install the additional data-processing dependencies:

pip install -r requirements-robotwin.txt

For distributed ZeRO-1 training, install DeepSpeed:

pip install 'deepspeed>=0.14'

📦 Pretrained Models and Checkpoints

Wan backbone

OSWAM uses the Wan2.2-TI2V-5B Video DiT, Wan2.2 VAE, and T5 text components. Set the model-cache directory before preprocessing, training, or evaluation:

export DIFFSYNTH_MODEL_BASE_PATH="$PWD/checkpoints"

• Wan-AI/Wan2.2-TI2V-5B: Video DiT.
• DiffSynth-Studio/Wan-Series-Converted-Safetensors: VAE and T5 weights.
• Wan-AI/Wan2.1-T2V-1.3B: tokenizer assets.

OSWAM weights

Download checkpoints from VCLab-PolyU/OSWAM on Hugging Face and save them under checkpoints/.

Benchmark	Model	Checkpoint	Hugging Face
NAVSIM	OSWAM student	navsim_oswam.pt	Download
NAVSIM	Privileged teacher	navsim_privileged_teacher.pt	Download
RoboTwin 2.0	OSWAM student	robotwin_oswam.pt	Download
RoboTwin 2.0	Privileged teacher	robotwin_privileged_teacher.pt	Download

🚗 NAVSIM

1. Prepare the data

Follow the NAVSIM dataset instructions to obtain the logs, camera blobs, and nuPlan maps. Organize the data as follows, or point the environment variables to equivalent locations:

data/navsim/
├── navsim_logs/
│   ├── trainval/
│   └── test/
├── sensor_blobs/
│   ├── trainval/
│   └── test/
└── maps/

export OPENSCENE_DATA_ROOT="$PWD/data/navsim"
export NUPLAN_MAPS_ROOT="$OPENSCENE_DATA_ROOT/maps"
export NAVSIM_TEXT_EMBED_CACHE="$PWD/data/text_embeds_cache/navsim"
source scripts/navsim_data_env.sh

2. Prepare the Action DiT and text embeddings

Before training, initialize the Action DiT by interpolating the Video DiT weights and cache the task's text embeddings. Skip either command if its matching assets already exist:

bash scripts/model_prepare.sh
bash scripts/precomput_text_embed.sh

3. Train the privileged teacher and OSWAM

# Stage 1: privileged teacher with full-video context.
python scripts/train.py --config-name train_teacher

# Stage 2: OSWAM with frozen-teacher distillation.
# Use the supplied teacher, or point this variable to your trained teacher.
export OSWAM_TEACHER_CHECKPOINT="$PWD/checkpoints/navsim_privileged_teacher.pt"
python scripts/train.py --config-name train_oswam

For multi-GPU training, use the ZeRO-1 launcher. Adjust the process count and batch size to your hardware:

NPROC_PER_NODE=8 bash scripts/train_wam_zero1_torchrun.sh \
  --config-name train_teacher

NPROC_PER_NODE=8 bash scripts/train_wam_zero1_torchrun.sh \
  --config-name train_oswam

4. Evaluate on NAVSIM v1

For PDMS, build the v1 metric cache, then run the combined prediction-and-scoring launcher:

bash navsim/scripts/evaluation/run_metric_caching.sh

bash experiments/navsim/run_eval_sft_pdms.sh \
  --checkpoint "$PWD/checkpoints/navsim_oswam.pt" \
  --num-inference-steps 2 \
  --nproc-per-node 1 \
  --metric-cache-path "$PWD/data/metric_cache_navtest" \
  model.skip_dit_load_from_pretrain=true

Increase --nproc-per-node to distribute prediction across GPUs. The metric cache and predictions must cover the same navtest scenes. The scoring launcher writes to evaluate_results/navsim_pdms/ by default.

5. Evaluate on NAVSIM v2

The v2 wrapper selects the bundled navsim_v2/ devkit and test data automatically. Prepare its separate cache, generate predictions, and score them:

export CKPT="$PWD/checkpoints/navsim_oswam.pt"
export EXP_NAME=oswam_navsim
export NAVSIM_EXP_ROOT="$PWD/evaluate_results_v2"

bash experiments/navsim/run_eval_oswam_v2.sh check
bash experiments/navsim/run_eval_oswam_v2.sh cache
bash experiments/navsim/run_eval_oswam_v2.sh predict \
  model.skip_dit_load_from_pretrain=true
bash experiments/navsim/run_eval_oswam_v2.sh score

The defaults are data/metric_cache_navtest_v2/ for the cache and evaluate_results_v2/ for outputs. Override NAVTEST_CACHE and NAVSIM_EXP_ROOT as needed. See NAVSIM v2 details.

NAVSIM v1 uses PDMS, while this v2 wrapper uses reactive, one-stage EPDMS; their scores are not directly comparable, their caches are incompatible, and the v2 wrapper differs from the official two-stage v2 protocol.

🤖 RoboTwin 2.0

1. Prepare the dataset and simulator

The data pipeline uses the preprocessed LeRobot-format RoboTwin dataset, following FastWAM's data preparation. The upstream dataset is available at yuanty/robotwin2.0-fastwam. Download all split archives into data/robotwin2.0/, then extract them from the repository root:

mkdir -p data/robotwin2.0
cat data/robotwin2.0/robotwin2.0.tar.gz.part-* | tar -xz -C data/robotwin2.0

Expected layout:

data/robotwin2.0/
├── robotwin2.0/
│   ├── data/
│   ├── meta/
│   └── videos/
└── dataset_stats.json

Supply the normalization statistics matching the checkpoint as dataset_stats.json; these are required in addition to the dataset. Install the simulator and download its assets following the official RoboTwin instructions. Set the paths before training or evaluation:

export OSWAM_ROBOTWIN_DATASET="$PWD/data/robotwin2.0/robotwin2.0"
export OSWAM_ROBOTWIN_STATS="$PWD/data/robotwin2.0/dataset_stats.json"
export OSWAM_ROBOTWIN_TEXT_CACHE="$PWD/data/text_embeds_cache/robotwin"
export OSWAM_ROBOTWIN_ROOT=/path/to/RoboTwin
source scripts/robotwin_env.sh

2. Prepare the Action DiT and text embeddings

Prepare the Action DiT initialization:

python scripts/preprocess_robotwin_action_dit_backbone.py \
  --model-config configs/model/oswam_robotwin.yaml \
  --output checkpoints/ActionDiT_linear_interp_Wan22_alphascale_1024hdim.pt \
  --device cuda --dtype bfloat16

Cache the dataset's instruction embeddings before training:

python scripts/precompute_robotwin_text_embeds.py

3. Train the privileged teacher and OSWAM

# Stage 1: privileged teacher.
bash scripts/train_robotwin.sh teacher

# Stage 2: OSWAM with its paired frozen teacher.
export OSWAM_ROBOTWIN_TEACHER_CHECKPOINT="$PWD/checkpoints/robotwin_privileged_teacher.pt"
bash scripts/train_robotwin.sh student

For distributed ZeRO-1 training:

NPROC_PER_NODE=8 USE_DEEPSPEED=1 bash scripts/train_robotwin.sh teacher
NPROC_PER_NODE=8 USE_DEEPSPEED=1 bash scripts/train_robotwin.sh student

4. Evaluate in the simulator

Run a single task, using a task name present in your RoboTwin installation:

bash scripts/eval_robotwin.sh \
  ckpt="$PWD/checkpoints/robotwin_oswam.pt" \
  EVALUATION.task_name=click_alarmclock \
  EVALUATION.task_config=demo_randomized

Use EVALUATION.task_config=demo_clean for the Clean setting.

For evaluation across tasks, run the task manager after sourcing scripts/robotwin_env.sh:

python experiments/robotwin/run_robotwin_manager.py \
  ckpt="$PWD/checkpoints/robotwin_oswam.pt" \
  MULTIRUN.num_gpus=1 MULTIRUN.max_tasks_per_gpu=1

Evaluation outputs are stored under evaluate_results/robotwin/. See the RoboTwin guide for more details.

📁 Repository Layout

OSWAM/
├── assets/                    # Paper figures
├── configs/                   # Data, model, training, and evaluation presets
├── src/oswam/
│   ├── models/wan22/          # NAVSIM student and privileged teacher
│   └── robotwin/              # RoboTwin models, data pipeline, and trainers
├── scripts/                   # Preprocessing, training, checkpoint tools
├── experiments/
│   ├── navsim/                # NAVSIM prediction and scoring
│   └── robotwin/              # RoboTwin evaluation and policy adapter
├── navsim/                    # NAVSIM v1 devkit
├── navsim_v2/                 # NAVSIM v2 devkit
├── checkpoints/               # Local model weights and pretrained assets
├── data/                      # Datasets, statistics, and caches
├── docs/                      # Detailed setup and validation notes
└── tests/                     # Regression checks

📖 Citation

If you find OSWAM useful in your research, please cite:

@article{zhang2026oswam,
  title={OSWAM: An Effective World Action Model with One-Step Action Generation},
  author={Zhang, Guowen and Chen, Liyi and Hu, Xu and Wang, Pengfei and Chen, Jingfan and He, Chenhang and Zhang, Lei},
  journal={arXiv preprint arXiv:2610.35850},
  year={2026}
}

🙏 Acknowledgements

Our implementation builds on SimWAM, FastWAM, and Wan2.2. We thank NAVSIM, nuPlan, and RoboTwin for their datasets and evaluation tools.

See LICENSE and the preserved FastWAM license. Third-party code, datasets, simulator assets, and pretrained models retain their respective licenses.

For questions, please contact Guowen Zhang.