DriveUniTok: A Unified Multimodal Tokenizer for Driving World-Action Models

Jingbo Xu; Hanchao Leng; JiaHui Huang; Hanbing Li; Jingwei Zhao; Kuiyuan Yang; Chen Gao; Guang Chen

cs.CVOct 26, 2026

DriveUniTok: A Unified Multimodal Tokenizer for Driving World-Action Models

Authors: Jingbo Xu, Hanchao Leng, JiaHui Huang, Hanbing Li, Jingwei Zhao, Kuiyuan Yang, Chen Gao, Guang Chen,

Organizations: Tsinghua University · Xiaomi EV

Abstract

World models are emerging as a fundamental component of policy learning for real-world agents, enabling agents to predict future world states and plan actions within a learned representation of the dynamic environment. Existing approaches, however, often model appearance, metric geometry represented by projected LiDAR depth, and ego motion either partially or independently, overlooking the fact that these modalities are different projections of a shared underlying world state. This separation can lead to information loss and cross-modal inconsistency, ultimately limiting downstream performance. Motivated by this unified view of the world, we propose DriveUniTok, a unified multimodal tokenizer for driving world-action models (WAMs) that encodes these modalities in a single decodable continuous latent. We freeze DriveUniTok and adapt a pretrained diffusion model, DriveUniWAM, to predict future unified latents; the tokenizer then decodes them into future appearance, projected LiDAR depth, and ego trajectories. Experiments on NAVSIM demonstrate that our unified design achieves better future prediction and planning performance at lower downstream computational cost than separate-stream alternatives. With reinforcement learning post-training, DriveUniWAM achieves a PDMS of 91.54, reaching state-of-the-art performance among the evaluated WAM-based planners.

Figures & tables

Figure 1: Comparison of latent representations for world-action modeling. (a) BEV tokenizers fuse RGB and LiDAR into a visual latent, while (b) Occupancy tokenizers represent the scene with an occupancy latent; both keep ego motion separate. (c) Modality-specific tokenizers encode RGB, LiDAR, and ego motion into separate streams. (d) DriveUniTok jointly encodes RGB, LiDAR-derived geometry, and ego motion into one unified latent.

Figure 2: DriveUniTok architecture. (a) Encoder. RGB, projected depth, and ego pose are fused into $S$ shared queries, whose $T$ per-frame states are compressed into one latent. (b) Decoder. The latent is expanded back to $T$ frame positions, and each target is read out through its own set of mask tokens, including a training-only semantic branch. Snowflakes mark frozen components, and all remaining modules are trained.

Figure 3: Training and policy optimization with DriveUniWAM. (a) Supervised training. DriveUniWAM forecasts the future latents with flow matching. (b) Reinforcement learning. Per scene, a group of sampled futures is decoded into trajectories and scored by the NAVSIM PDM reward, whose relative advantages drive the Flow-GRPO update. DriveUniTok stays frozen in both stages.

	Tokenizer configuration		NAVSIM navtest
Method	Comp.	Dim.	LPIPS ${}_{\mathrm{VGG}}\downarrow$	LPIPS ${}_{\mathrm{DINO}}\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	rFVD ${}_{\mathrm{I3D}}\downarrow$	rFVD ${}_{\mathrm{VJ}}\downarrow$
Large-Scale Pretrained Video Tokenizers
Wan2.1 VAE ( Team Wan et al., 2025 )	$4{\times}8{\times}8$	16	0.1151	0.0649	31.50	0.8819	2.06	1.0984
Wan2.2 VAE ( Team Wan et al., 2025 )	$4{\times}16{\times}16$	48	0.1081	0.0608	31.35	0.8775	2.98	0.7958
HunyuanVideo VAE ( Kong et al., 2024 )	$4{\times}8{\times}8$	16	0.1230	0.0573	32.36	0.8924	4.03	0.6278
CogVideoX VAE ( Yang et al., 2025 )	$4{\times}8{\times}8$	16	0.1401	0.0943	32.26	0.8909	5.09	0.9298

Table 1: Native-configuration comparison of video tokenizers on NAVSIM navtest. Comp. reports native temporal

\times

height

\times

width compression for grid tokenizers and tokens per clip for token-sequence methods; for DriveUniTok, it denotes the complete 584-token multimodal clip latent. The baselines use their official RGB-only preprocessing, whereas DriveUniTok additionally encodes projected depth and ego pose with a frozen DINOv3-L feature backbone. Results characterize fidelity–compression trade-offs under each tokenizer’s native configuration.

	Representation budget			Efficiency			Reconstruction			Future prediction			Planning
Representation	Tokens	Scalars	Len.	Tok. params (M)	Tok. cost	DiT/step cost	LPIPS DINO	AbsRel	ADE	gFVD VJ	AbsRel	ADE	PDMS
Separate, natural budget	1,160	18,560	3,480	957.8	4.6	34.2	0.1083	0.0723	0.0702	7.92	0.1194	0.3997	90.1
Separate, equal budget	584	9,344	1,752	956.6	4.3	16.3	0.1124	0.0766	0.0720	12.80	0.1850	0.4186	87.5
Unified (DriveUniTok)	584	9,344	1,752	847.3	4.6	16.3	0.1063	0.0461	0.0728	6.61	0.1059	0.3816	91.0

Table 2: Unified vs. separate tokenization on NAVSIM. Scalars counts latent values per clip; Len. is the DriveUniWAM sequence length at

K{=}1,F{=}2

. Costs are in TFLOPs; DiT cost is per denoising step. Lower reconstruction and prediction errors and higher PDMS are better.

Method	Reference	Sensors	NC $\uparrow$	DAC $\uparrow$	EP $\uparrow$	TTC $\uparrow$	Comfort $\uparrow$	PDMS $\uparrow$
Reference
Human Agent	–	–	100.0	100.0	87.5	100.0	99.9	94.8
Vision-Based End-to-End Methods
UniAD ( Hu et al., 2023 )	CVPR’23	6 $\times$ C	97.8	91.9	78.8	92.9	100.0	83.4
TransFuser ( Chitta et al., 2023 )	TPAMI’23	3 $\times$ C+L	97.7	92.8	79.2	92.8	100.0	84.0
ARTEMIS ( Feng et al., 2025 )	RA-L’25	3 $\times$ C+L	98.3	95.1	81.4	94.3	100.0	87.0

Table 3: NAVSIM navtest performance. C and L denote camera and LiDAR. The best and second-best learned results, including ties, are shown in bold and underlined .

Video backbone	Tokenizer	NC $\uparrow$	DAC $\uparrow$	EP $\uparrow$	TTC $\uparrow$	PDMS $\uparrow$
Part I: Visual Tokenizer Representation Ablation
Wan2.1-T2V-1.3B ( Team Wan et al., 2025 )	Native VAE	98.1	94.7	80.9	94.5	86.5
Wan2.1-T2V-1.3B ( Team Wan et al., 2025 )	DriveUniTok	98.4	97.0	82.9	95.4	89.0 (+2.5)
Wan2.2-TI2V-5B ( Team Wan et al., 2025 )	Native VAE	98.4	95.5	82.6	94.8	87.9
Wan2.2-TI2V-5B ( Team Wan et al., 2025 )	DriveUniTok	99.0	97.7	85.6	96.3	91.0 (+3.1)
Part II: Scalability across Video DiT Backbones (with DriveUniTok)

Table 4: Ablation of video backbones and latent representations on the NAVSIM navtest benchmark. Part I compares the native VAE and DriveUniTok with fixed backbones, while Part II evaluates DriveUniWAM across video DiT backbones using the same unified latent. Across both parts, best and second-best results are shown in bold and underlined ; improvements over the native VAE appear in parentheses.

Appendix figures & tables17 assets

Supplementary material from the paper’s appendix.

Appendix

Group	Term	Form	Weight
$\mathcal{L}_{\mathrm{rgb}}$	pixel $L_{1}$	$\\|\hat{I}-I\\|_{1}$	1.0
	pixel $L_{2}$	$\\|\hat{I}-I\\|_{2}^{2}$	1.0
	LPIPS	VGG feature distance ( Zhang et al., 2018 )	1.0
	perceptual	ConvNeXt feature distance	0.1
	perceptual	DINOv3 feature distance	1.0
$\mathcal{L}_{\mathrm{depth}}$	confidence	$\mathbb{E}_{\Omega}[r\kappa-\alpha\log\kappa]$ , $\alpha{=}0.2$	1.0

Appendix

Table 5: Components of the DriveUniTok training objective.

\Omega

denotes the set of pixels where LiDAR projects to a valid depth.

Figure 4: Qualitative tokenizer reconstructions for reconstruction group 1. From left to right, the columns show RGB ground truth, reconstructed RGB, LiDAR-projected depth ground truth, and reconstructed depth at valid LiDAR pixels.

Figure 5: Qualitative tokenizer reconstructions for reconstruction group 2. From left to right, the columns show RGB ground truth, reconstructed RGB, LiDAR-projected depth ground truth, and reconstructed depth at valid LiDAR pixels.

Figure 6: Qualitative tokenizer reconstructions for reconstruction group 3. From left to right, the columns show RGB ground truth, reconstructed RGB, LiDAR-projected depth ground truth, and reconstructed depth at valid LiDAR pixels.

	RGB				Depth		Ego
Variant	LPIPS ${}_{\mathrm{DINO}}\downarrow$	PSNR $\uparrow$	SSIM $\uparrow$	rFVD ${}_{\mathrm{VJ}}\downarrow$	AbsRel $\downarrow$	RMSE $\downarrow$	ADE $\downarrow$	FDE $\downarrow$
w/o masked reconstruction and semantic supervision	0.1540	24.55	0.68	8.40	0.0695	9.30	0.1450	0.2670
w/o intra-modal masking	0.1215	25.20	0.72	4.85	0.0520	8.00	0.0900	0.1650
w/o cross-modal masking	0.1298	25.02	0.72	5.60	0.0565	8.30	0.1040	0.1910
w/o semantic supervision	0.1195	25.55	0.74	4.28	0.0480	7.72	0.0775	0.1420
Full training objective	0.1063	25.62	0.74	3.45	0.0461	7.57	0.0728	0.1329

Appendix

Table 6: Training-objective ablation on NAVSIM navtest. Each row removes the named objectives from the full recipe. All variants share the architecture,

584\times 16

latent, and complete camera, depth, and ego inputs at evaluation. Best values, including ties, are bold.

Method	Reference	NC $\uparrow$	DAC $\uparrow$	DDC $\uparrow$	TLC $\uparrow$	EP $\uparrow$	TTC $\uparrow$	LK $\uparrow$	HC $\uparrow$	EC $\uparrow$	EPDMS $\uparrow$
Reference
Human Agent	–	100.0	100.0	99.8	100.0	87.4	100.0	100.0	98.1	90.1	94.5
Vision-Based End-to-End Methods
TransFuser ( Chitta et al., 2023 )	TPAMI’23	96.90	89.90	97.80	99.70	87.10	95.40	92.70	98.30	87.20	76.7
DriveSuprim ( Yao et al., 2026 )	AAAI’26	97.50	96.50	99.40	99.60	88.40	96.60	95.50	98.30	77.00	83.1
DiffusionDrive ( Liao et al., 2025a )	CVPR’25	98.20	95.90	99.40	99.80	87.50	97.30	96.80	98.30	87.70	84.5

Appendix

Table 7: Performance on the NAVSIM-v2 navtest benchmark. Human Agent is included as a reference. DriveUniWAM (Ours) denotes the imitation-trained DriveUniWAM with a frozen DriveUniTok tokenizer. TLC denotes traffic-light compliance, HC denotes history comfort, EC denotes extended comfort, and EPDMS is the NAVSIM-v2 aggregate score. The best and second-best learned results, including ties, are shown in bold and underlined .

		Tokenizer Reconstruction				Joint Future Generation
Input	Output	LPIPS ${}_{\mathrm{DINO}}\downarrow$	rFVD ${}_{\mathrm{VJ}}\downarrow$	Depth AbsRel $\downarrow$	Ego ADE $\downarrow$	LPIPS ${}_{\mathrm{DINO}}\downarrow$	gFVD ${}_{\mathrm{VJ}}\downarrow$	Depth AbsRel $\downarrow$	Traj. ADE $\downarrow$
C	C+L+T	0.1671	8.8645	0.1274	1.0512	0.1170	12.1153	0.2185	3.1991
L	L+T	–	–	0.1202	0.3350	–	–	0.2934	4.5854
T	T	–	–	–	0.2978	–	–	–	2.3534
C+L	C+L+T	0.1070	3.6430	0.0463	0.0876	0.0985	8.8587	0.1669	2.6611
C+L+T	C+L+T	0.1063	3.4546	0.0461	0.0728	0.0955	6.6076	0.1059	0.3816

Appendix

Table 8: Modality sensitivity for reconstruction and joint future generation on NAVSIM navtest. C, L, and T denote camera, LiDAR/depth, and ego-trajectory modalities. Input and reconstruction Output follow the original tokenizer masking setup; the DriveUniWAM block evaluates future generation under the same input condition. Best results are highlighted in bold .

Figure 7: Qualitative NAVSIM navtest visualizations for scene group 1. From left to right, the columns show the front-view observation, AutoVLA, ReCogDrive, PWM, DriveUniWAM (Ours), and DriveUniWAM with reinforcement learning post-training. AutoVLA and ReCogDrive are vision-language-action methods, whereas PWM is a world-action baseline.

Figure 8: Qualitative NAVSIM navtest visualizations for scene group 2. The columns follow the ordering defined in Fig. 7 .

Figure 9: Qualitative NAVSIM navtest visualizations for scene group 3. The columns follow the ordering defined in Fig. 7 .

Figure 10: Qualitative NAVSIM navtest visualizations for scene group 4. The columns follow the ordering defined in Fig. 7 .

Figure 11: Qualitative NAVSIM navtest visualizations for scene group 5. The columns follow the ordering defined in Fig. 7 .

Figure 12: Qualitative NAVSIM navtest visualizations for scene group 6. The columns follow the ordering defined in Fig. 7 .

Figure 13: Qualitative NAVSIM navtest visualizations for scene group 7. The columns follow the ordering defined in Fig. 7 .

Figure 14: Qualitative NAVSIM navtest visualizations for scene group 8. The columns follow the ordering defined in Fig. 7 .

Figure 15: Qualitative NAVSIM navtest visualizations for scene group 9. The columns follow the ordering defined in Fig. 7 .

Figure 16: Qualitative NAVSIM navtest visualizations for scene group 10. The columns follow the ordering defined in Fig. 7 .

World Models for Autonomous Driving

World models for autonomous driving predict how scenes and traffic evolve under candidate actions to support planning, control, and simulation.

See all 143 papers

DriveDreamer-Policy: A Geometry-Grounded World-Action Model for Unified Generation and Planning
Apr 2, 2026Yang Zhou, Xiaofeng Wang, Hao Shao +8World Action Models World Models for Autonomous Driving
GigaAI · University of Toronto · CUHK MMLab
Discrete-WAM: Unified Discrete Vision-Action Token Editing for World-Policy Learning
Jun 4, 2026Ziyang Yao, Haochen Liu, Yuncheng Jiang +10Hierarchical Planning World Model Learning
Xiaomi EV
Latent-WAM: Latent World Action Modeling for End-to-End Autonomous Driving
Mar 25, 2026Linbo Wang, Yupeng Zheng, Qiang Chen +13End-to-End Autonomous Driving World Models for Autonomous Driving
Institute of Automation, Chinese Academy of Sciences · Chongqing Chang’an Technology Co., Ltd · Zhongguancun Academy +2

Apr 2, 2026J/K move · Enter open · S save

DriveDreamer-Policy: A Geometry-Grounded World-Action Model for Unified Generation and Planning

Yang Zhou, Xiaofeng Wang, Hao Shao, Letian Wang+7

GigaAI · University of Toronto · CUHK MMLab

World Action Models World Models for Autonomous Driving Autonomous Driving Planning Video World Models Vision-Language-Action Models

Open paper AskPDF arXiv

Recently, world-action models (WAM) have emerged to bridge vision-language-action (VLA) models and world models, unifying their reasoning and instruction-following capabilities and spatio-temporal world modeling. However, existing WAM approaches often focus on modeling 2D appearance or latent representations, with limited geometric grounding-an essential element for embodied systems operating in the physical world. We present DriveDreamer-Policy, a unified driving world-action model that integrates depth generation, future video generation, and motion planning within a single modular architecture. The model employs a large language model to process language instructions, multi-view images, and actions, followed by three lightweight generators that produce depth, future video, and actions. By learning a geometry-aware world representation and using it to guide both future prediction and planning within a unified framework, the proposed model produces more coherent imagined futures and more informed driving actions, while maintaining modularity and controllable latency. Experiments on the Navsim v1 and v2 benchmarks demonstrate that DriveDreamer-Policy achieves strong performance on both closed-loop planning and world generation tasks. In particular, our model reaches 89.2 PDMS on Navsim v1 and 88.7 EPDMS on Navsim v2, outperforming existing world-model-based approaches while producing higher-quality future video and depth predictions. Ablation studies further show that explicit depth learning provides complementary benefits to video imagination and improves planning robustness.

Abstract

World Models for Autonomous Driving

Explore similar work

DriveDreamer-Policy: A Geometry-Grounded World-Action Model for Unified Generation and Planning

Discrete-WAM: Unified Discrete Vision-Action Token Editing for World-Policy Learning

Latent-WAM: Latent World Action Modeling for End-to-End Autonomous Driving