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.

---

📣 News

• 2026.08.07: Released the SimWAM code.

---

📄 Abstract

World-Action Models (WAMs) improve end-to-end autonomous driving by transferring video dynamics priors to action prediction, but existing methods require costly future generation at inference. We present SimWAM, a simple yet effective WAM that uses video generation purely as a training 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 the video branch to be discarded after training and leaving a self-contained planner that directly predicts trajectories. 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 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

_{Overview of SimWAM. During joint training, the video and action DiTs learn future dynamics and trajectory generation through MoT, while the attention mask restricts the action tokens to the current observation. Only the action DiT is retained for inference and reinforcement learning.}

• Joint flow-matching co-training. A pretrained video expert (Wan2.2-5B, with its VAE and T5 encoder) and a lightweight action DiT are co-trained with joint flow matching, so future-frame prediction transfers a traffic-aware motion prior into the shared observation representation.
• Isolated attention mask. Future-frame tokens and action tokens both attend to the current observation but remain mutually invisible, keeping action prediction independent of future frames. This is the only structural modification needed to separate the two experts.
• Self-contained planner at inference. Because the action expert never depends on future video tokens, the entire video branch is discarded after training — no future-scene generation and no auxiliary motion modules at deployment.
• 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, updating only LoRA adapters of the action expert.
• Flexibility. The experts share no weights and interact only through the attention stream: the video backbone is replaceable (e.g., LTX-Video, Wan2.1-2B, Cosmos2.5, Wan2.2-5B) and the action expert is independently scalable (0.21B–1.02B) without changing the objective or inference pipeline.

---

📈 Performance

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

NAVSIM navtest

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.

Method	Reference	Sensors	NC↑	DAC↑	EP↑	TTC↑	C↑	PDMS↑
Human Agent	-	-	100.0	100.0	87.5	100.0	99.9	94.8
Traditional E2E planners
UniAD	CVPR'23	6×C	97.8	91.9	78.8	92.9	100.0	83.4
TransFuser	TPAMI'22	3×C+L	97.7	92.8	79.2	92.8	100.0	84.0
Hydra-MDP++	arXiv'25	3×C+L	97.6	96.0	80.4	93.1	100.0	86.6
ARTEMIS	RA-L'26	3×C+L	98.3	95.1	81.4	94.3	100.0	87.0
WorldRFT	AAAI'26	3×C	97.5	96.0	80.9	94.0	100.0	87.0
DiffusionDrive	CVPR'25	3×C+L	98.2	96.2	82.2	94.7	100.0	88.1
WoTE	ICCV'25	3×C+L	98.5	96.8	81.9	94.9	99.9	88.3
SeerDrive	NeurIPS'25	3×C+L	98.4	97.0	83.2	94.9	99.9	88.9
VLM-based planners
UniWorldVLA	arXiv'26	1×C	98.7	96.7	83.2	96.1	100.0	89.4
DriveDreamer-Policy	arXiv'26	3×C	98.4	97.1	83.5	95.1	100.0	89.2
Vega	arXiv'26	1×C	98.9	95.3	81.6	96.1	100.0	87.9
ImagiDrive	ICRA'26	1×C	98.6	96.2	80.5	94.5	100.0	87.4
AutoVLA	NeurIPS'25	3×C	98.4	95.6	81.9	98.0	99.9	89.1
ReCogDrive	ICLR'26	1×C	97.9	97.3	87.3	94.9	100.0	90.8
ExploreVLA	ECCV'26	1×C	98.8	98.4	83.5	96.5	99.9	90.4
DriveVLA-W0	ICLR'26	1×C	98.7	99.1	83.3	95.3	99.3	90.2
SGDrive	CVPR'26	1×C	98.6	97.8	85.8	96.2	100.0	91.1
World-model-based planners
Epona	ICCV'25	1×C	97.9	95.1	80.4	93.8	99.9	86.2
PWM	NeurIPS'25	1×C	98.6	95.9	81.8	95.4	100.0	88.1
DriveLaW	CVPR'26	1×C	99.0	97.1	81.3	96.7	100.0	89.1
DriveWAM	arXiv'26	1×C	98.3	98.1	84.3	95.2	100.0	90.1
SimWAM (Ours)	-	1×C	98.4	98.7	86.4	95.5	100.0	91.5

Component analysis

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.

Configuration	NC	DAC	EP	TTC	PDMS
Action-only	97.6	95.7	81.7	92.6	86.6
+ Video	98.7	98.0	83.9	95.9	90.3
+ RL	98.4	98.7	86.4	95.5	91.5

Qualitative results

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

---

⚙️ Installation

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.

---

📦 Preparation

ActionDiT initialization

bash scripts/model_prepare.sh

Text embeddings

bash scripts/precomput_text_embed.sh

Use +overwrite=false to keep existing embeddings.

---

🏋️ Training and Evaluation

Supervised training

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

FlowGRPO LoRA fine-tuning

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.

NAVSIM evaluation

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.

---

👍 Acknowledgement

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.

---

📖 Citation

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}
}

---

License

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.