SimWAM: A Simple World Action Model for End-to-End Autonomous Driving

Zongchuang Zhao¹, Xin Zhou¹, Tianyang Xu¹, Zhengyang Sun¹, Kaixuan Zhou², Honglin Li², Dingkang Liang^1†, Xiang Bai¹

¹ Huazhong University of Science & Technology
² Dongfeng Research & Development Institute
^† Project leader.

This repository provides the official implementation of SimWAM for the paper A Simple World Action Model for End-to-End Autonomous Driving, including supervised training and action-only reinforcement learning on NAVSIM.

• 2026.08.07: Released the SimWAM code and weight.

World-Action Models (WAMs) improve end-to-end autonomous driving by transferring video dynamics priors to action prediction, but existing methods incur costly test-time future imagination. We present SimWAM, a simple yet effective WAM that leverages future-video prediction as a training-time supervision signal. It co-trains a pretrained video expert and a lightweight action expert with joint flow matching. An isolated attention mask keeps action prediction independent of future frames, allowing trajectory prediction without explicit future-frame generation at inference. Since the two experts share no parameters and interact only through a unified attention interface, the video backbone could be replaced and the action expert scaled independently without modifying the learning objective or inference pipeline. We further apply reinforcement learning to optimize a compositional driving reward beyond trajectory imitation. Our SimWAM achieves 91.5 PDMS on NAVSIM, surpasses state-of-the-art WAM-based planners with substantially lower latency, and transfers zero-shot to nuScenes. These results position SimWAM as a simple yet solid baseline that could readily benefit from advances in video generation for efficient autonomous driving.

_{Overview of SimWAM. During training, the video and action DiTs are jointly optimized for future-frame generation and trajectory prediction via shared attention, while the isolated mask prevents the action tokens from accessing future-frame tokens. During inference and reinforcement learning, the model directly predicts trajectories without explicitly predicting future frames.}

• Joint flow-matching co-training. The video expert — a video Diffusion Transformer initialized from Wan2.2-5B together with its video VAE and T5 text encoder — and a lightweight action DiT (hidden size 1024) are co-trained with joint flow matching over future-frame latents and trajectories. Future-video prediction serves as a training-time supervision signal that transfers a traffic-aware dynamics prior into the shared observation representation used for planning.
• Isolated attention mask. Future-frame tokens and action tokens both attend to the current observation latents while remaining mutually invisible, keeping action prediction independent of future frames. This mask is the only structural modification required to isolate the action tokens from future-frame information.
• Direct trajectory prediction at inference. Because the action expert never depends on future-frame tokens, explicit future-frame generation is omitted at deployment: the standalone action DiT directly predicts trajectories without auxiliary motion modules, substantially reducing inference latency.
• Reinforcement learning. The deterministic flow ODE is reformulated as a marginal-preserving SDE, and a group of candidate trajectories per scenario is optimized with FlowGRPO against the compositional NAVSIM PDM reward. RL focuses on the hard navtrain scenarios with the lowest PDMS after imitation learning and updates only the LoRA adapters of the action expert.
• Flexibility. The two experts share no weights and interact only through the unified attention interface: the video backbone is replaceable (e.g., LTX-Video, Wan2.1-1.3B, Cosmos2.5, Wan2.2-5B) and the action expert is independently scalable (0.21B–1.02B) without modifying the learning objective or inference pipeline.

_{SimWAM achieves the best PDMS and the lowest inference latency among recent world-model-based planners on NAVSIM.}

Using only a single front camera at 384×672, SimWAM establishes a new state of the art on NAVSIM navtest, surpassing the strongest VLM-based planner (SGDrive) by 0.4 PDMS and recent imagine-then-act WAMs (DriveLaW, DriveWAM) by 2.4 and 1.4 points, respectively.

Video co-training and reinforcement learning contribute complementary gains, improving PDMS by 4.9 points while preserving the simplicity and efficient inference of the standalone action expert.

_{Qualitative comparison on two navtest scenarios. After reinforcement learning, the ego commits further along the route while staying collision-free within the drivable area.}

Clone the repository:

git clone https://github.com/H-EmbodVis/SimWAM.git
cd SimWAM

Create a Python 3.10 environment and install the pinned runtime dependencies:

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

python -m pip install -r requirements.txt
python -m pip install -e navsim --no-deps
python -m pip install -e . --no-deps

For NAVSIM, nuPlan, maps, and dataset preparation, follow the official NAVSIM v1.1 repository and the nuPlan devkit. The required release subset is included under navsim/.

All commands are intended to run from the repository root and use relative paths by default.

bash scripts/model_prepare.sh

bash scripts/precomput_text_embed.sh

Use +overwrite=false to keep existing embeddings.

NNODES=4 \
NPROC_PER_NODE=8 \
bash scripts/train_navsim_zero1_torchrun.sh \
  task=navsim_uncond_front_384x672_1e-4 \
  num_workers=8

NPROC_PER_NODE=8 \
bash scripts/train_navsim_grpo_zero1_torchrun.sh \
  task=navsim_grpo_action_pdm_384x672_flowgrpo_lora \
  num_workers=8 \
  model.checkpoint_path=./runs/navsim_uncond_front_384x672_1e-4/<run-id>/checkpoints/weights/step_XXXXXX.pt

Alternatively, set SIMWAM_IL_CHECKPOINT to the supervised checkpoint.

CKPT=./runs/navsim_grpo_action_pdm_384x672_flowgrpo_lora/<run-id>/checkpoints/weights/step_XXXXXX.pt \
TASK=navsim_grpo_action_pdm_384x672_flowgrpo_lora \
NPROC_PER_NODE=8 \
bash experiments/navsim/run_eval_navsim.sh

To verify the released supervised checkpoint with a one-sample smoke test:

CKPT=./weights/SimWAM.pt \
TASK=navsim_uncond_front_384x672_1e-4 \
NPROC_PER_NODE=1 \
bash experiments/navsim/run_eval_navsim.sh \
  EVALUATION.max_samples=1 \
  EVALUATION.num_inference_steps=2 \
  EVALUATION.save_videos=false

Training outputs are written to runs/; evaluation outputs are written to evaluate_results/navsim/.

The supplied launchers use DeepSpeed ZeRO-1 through scripts/ds_configs/ds_zero1_config.json. The matching Accelerate configuration is scripts/accelerate_configs/accelerate_zero1_ds.yaml.

SimWAM builds upon the following projects and resources:

• NAVSIM for the planning benchmark and evaluation tooling.
• nuPlan / OpenScene for the driving datasets.
• Wan2.2 for the pretrained video generation backbone.

If SimWAM is useful in your research, please consider citing the paper:

@article{zhao2026simwam,
  title={SimWAM: A Simple World Action Model for End-to-End Autonomous Driving}, 
  author={Zongchuang Zhao and Xin Zhou and Tianyang Xu and Zhengyang Sun and Kaixuan Zhou and Honglin Li and Dingkang Liang and Xiang Bai},
  journal={arXiv preprint arXiv:2608.07468},
  year = {2026}
}

See LICENSE. NAVSIM, Wan2.2, nuPlan, and OpenScene retain their own licenses and distribution terms. Refer to the official NAVSIM license, Wan2.2 repository, and nuPlan license for upstream terms.