PaperlayerView exact revision on arXiv

Source-linked AI summary

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, Hangjun Ye, Yong Li

arXiv:2610.35044v1cs.CVcs.RO

TL;DR

Driving world models often represent appearance, metric geometry, and ego motion separately rather than in one decodable latent. DriveUniTok unifies these modalities, and DriveUniWAM predicts future unified latents for scene prediction and planning. On NAVSIM, the unified design improves prediction and planning at lower downstream computational cost than separate-stream alternatives, with 91.54 PDMS after reinforcement learning post-training.

  • Problem

    Existing approaches often model appearance, projected-LiDAR geometry, and ego motion partially or independently, raising whether they should be represented within one unified tokenization framework.

    Source: §1 · p. 1

  • Method

    DriveUniTok encodes RGB appearance, projected LiDAR geometry, and clip-relative ego pose in one jointly decodable continuous latent, which frozen DriveUniTok and diffusion-based DriveUniWAM use to predict future latents for joint scene prediction and planning.

    Source: §1 · p. 2

  • Results

    91.54 PDMS after reinforcement learning post-training gives DriveUniWAM state-of-the-art performance among evaluated WAM-based planners, while the unified design improves prediction and planning at lower downstream computational cost than separate-stream controls.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The unified latent provides a shared representation in which appearance and motion correspondence is established for world-action modeling.

    Source: §2.2 · p. 3, §3.3 · p. 6

Abstract

from arXiv · show

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.

1 INTRODUCTION

Existing driving world models often separate appearance, metric geometry, and ego motion, leaving a central question of whether they should share one decodable representation. DriveUniTok unifies them for future prediction and planning, and NAVSIM results show stronger performance at lower downstream cost than separate-stream controls.

1 INTRODUCTION

  • The paper frames the latent as the space where a driving world model must capture information needed for both future prediction and planning.It describes a shift from world models as scene simulators and data engines toward generative future observations integrated into policy backbones.Source: §1 · p. 1
  • Existing approaches often represent appearance, metric geometry, and ego pose partially or independently, prompting the question of whether all three should share one decodable tokenization.RGB video latents primarily capture appearance; some approaches add geometry separately, while driving BEV and cross-modal tokenizers commonly leave or separate ego pose.Source: §1 · p. 1
  • DriveUniTok uses shared learnable queries and a causal temporal compressor to encode RGB appearance, projected LiDAR geometry, and clip-relative ego pose in one jointly decodable continuous latent.Modality-specific readouts reconstruct all three modalities without a privileged anchor; a training-only semantic decoder encourages preservation of high-level scene semantics.Source: §1 · p. 2
  • DriveUniWAM adapts a pretrained diffusion model to predict future unified latents, which decode into future observations and ego trajectories for joint scene prediction and planning.It is conditioned on encoded history, clip-relative ego poses, and a navigation signal; GRPO post-training uses the PDM reward while regularizing toward the supervised reference model.Source: §1 · p. 2
  • 91.54 PDMS after reinforcement learning post-training gives DriveUniWAM state-of-the-art performance among evaluated WAM-based planners, with stronger prediction and planning at lower downstream cost than separate-stream controls.The comparison is reported on NAVSIM.Source: §1 · p. 2

2 RELATED WORK

Prior driving tokenizers and planners represent scene information in different ways, but generally do not jointly encode appearance, metric geometry, and ego motion in one decodable latent. DriveUniTok makes that unification at tokenization for world-action modeling.

2.1 GENERAL AND DRIVING TOKENIZERS

  • General-purpose tokenizers target visual generation or cross-modal alignment, while driving-specific tokenizers represent camera or occupancy scenes; the described approaches do not jointly encode appearance, metric geometry, and ego pose.BEVWorld renders fused camera and LiDAR modalities but its latent does not encode ego pose.Source: §2.1 · p. 3

2.2 DRIVING PLANNERS AND WORLD-ACTION MODELS

  • Driving planners use future information through rollouts, BEV features, or planning latents, while vision-language-action variants add semantic priors and reasoning; these approaches generally separate future-scene prediction from trajectory generation.The cited approaches include feature fusion and waypoint regression, diffusion-refined trajectories, and autoregressive decoding.Source: §2.2 · p. 3
  • World-action models couple future prediction and action generation, but their described designs do not generally provide one decodable latent for appearance, metric geometry, and ego motion.DriveUniTok unifies those modalities at tokenization, and DriveUniWAM forecasts that latent for joint scene prediction and planning.Source: §2.2 · p. 3

3 METHODOLOGY

DriveUniTok maps each driving sequence to a compact latent from which RGB, projected LiDAR depth, and ego pose can be decoded. Its shared-query encoder, temporal compressor, and shared decoder implement that representation.

3.1 PROBLEM SETUP

  • Each observation sequence has five consecutive frames sampled at 2 Hz, with each frame containing a front-camera image, projected sparse LiDAR depth, and ego pose relative to the first frame.The stated image resolution is 288 × 512.Source: §3.1 · pp. 3–4

3.2 DRIVEUNITOK: QUERY ALIGNED MULTIMODAL TOKENIZER

  • The encoder fuses RGB, projected depth, and ego pose into shared queries, compresses their per-frame states into one latent, and the decoder expands it for target-specific readouts.The architecture includes a training-only semantic branch; snowflakes in the figure mark frozen components.Source: Figure 2 · p. 4, §3.2 · p. 4

3.1 PROBLEM SETUP

  • DriveUniTok maps a sequence to a 584-token latent of width 16, from which RGB, depth, and pose are decoded; DriveUniWAM models latent evolution across sequences.The tokenizer is trained first and frozen, and DriveUniWAM is post-trained against the NAVSIM PDM score.Source: §3.1 · p. 4

3.2 DRIVEUNITOK: QUERY ALIGNED MULTIMODAL TOKENIZER

  • Frozen DINOv3 features process RGB and projected depth in a shared feature hierarchy, while learned queries jointly gather evidence from both streams and ego tokens.Depth receives a single-channel patch embedding, and relative poses are mapped by a trainable ego state encoder; the 576 image queries and eight ego tokens are updated jointly.Source: §3.2 · pp. 4–5

Q + DA(Q, F rgb)

DriveUniTok combines appearance, projected depth, and ego-pose information in shared queries, compresses them into a unified latent, and decodes all three modalities from shared memory. DriveUniWAM predicts future unified latents with a shared denoiser and can be post-trained using NAVSIM rewards while the tokenizer remains frozen.

Q + DA(Q, F rgb)

  • Shared queries attend to RGB and depth features, while self-attention lets image- and ego-initialized queries exchange information before temporal compression.No modality is a privileged reference, and fusion is not deferred until after independent encoders finish.Source: Q + DA(Q, F rgb) · p. 5
  • Causal attention and learned pooling compress each query’s T frame states into one Gaussian-bottleneck code, yielding S tokens regardless of modality count.The resulting sequence is roughly half the length of separate per-modality representations.Source: Q + DA(Q, F rgb) · p. 5
  • A shared Transformer reads frame-level memory through modality-specific mask tokens, producing RGB and depth reconstructions and regressed ego poses.The RGB, depth, and ego branches share memory; their mask tokens, positional embeddings, and output heads differ.Source: Q + DA(Q, F rgb) · p. 5
  • The tokenizer combines RGB, valid-pixel depth, ego-pose, semantic, and KL losses; shared decoder memory keeps the queries from degenerating into unrelated modality-specific subspaces.The depth loss uses confidence-weighted regression on valid LiDAR pixels plus a multiscale gradient term; the semantic target is training-only.Source: Q + DA(Q, F rgb) · p. 5

3.3 DRIVEUNIWAM: JOINT WORLD-ACTION MODELING

  • DriveUniWAM uses a frozen DriveUniTok and a pretrained Wan-initialized denoiser to predict future unified latents conditioned on encoded history and navigation.History queries attend only to history tokens, while future queries attend to the complete joint sequence.Source: §3.3 · p. 6
  • For NAVSIM reward optimization, sampled joint futures are decoded into trajectories and scored, with relative advantages driving a clipped policy update while the tokenizer and decoder remain frozen.Only denoiser LoRA adapters are trained; the update is regularized toward the supervised reference model.Source: §3.3 · pp. 6–7, Figure 3 · p. 6

4 EXPERIMENTS

The NAVSIM experiments assess reconstruction, joint prediction and planning, and backbone generalization. DriveUniTok improves reported prediction and planning results against separate-stream controls, with DriveUniWAM reaching 91.54 PDMS after reward post-training.

4.1 EXPERIMENTAL SETUP

  • On NAVSIM, evaluation covers multimodal reconstruction, joint future prediction, and planning, using front-camera RGB, projected LiDAR depth, and clip-relative ego poses encoded into 584 tokens.NAVSIM is a non-reactive planning benchmark derived from the Open-Scene subset of nuPlan.Source: §4.1 · p. 7

4.2 TOKENIZER RECONSTRUCTION

  • DriveUniTok attains the lowest LPIPSDINO (0.1063) and rFVDVJ (3.4546) among benchmark-scale methods, improving over their respective second-best results by 18.7% and 39.7%.Large pretrained VAEs retain an advantage on pixel-oriented metrics but do not encode geometry or ego motion.Source: §4.2 · pp. 7–8

4.3 BENEFITS OF UNIFIED TOKENIZATION

  • At equal representation and DiT budgets, DriveUniTok reduces gFVDVJ by about 48%, future depth AbsRel by about 43%, and trajectory ADE by nearly 9%, while raising PDMS by 3.5 points to 91.0.Against the natural-budget control, it uses roughly half as many tokens and latent scalars, halves sequence length, cuts per-step DiT computation by about 52%, and improves PDMS from 90.1 to 91.0.Source: §4.3 · p. 8, Table 2 · p. 8

4.4 PLANNING PERFORMANCE

  • 91.54 PDMS after reward post-training; DriveUniWAM reaches 91.0 PDMS with supervised training, while the tokenizer stays frozen during both stages.The planning evaluation also reports 90.5 PDMS with camera-only input and 91.0 when projected LiDAR is added.Source: §4.4 · p. 8

4.5 BACKBONE AND TOKENIZER GENERALIZATION

  • Across five video diffusion backbones, the unified latent achieves 88.9–91.0 PDMS; replacing the native VAE improves PDMS by 2.5 points for Wan2.1-1.3B and 3.1 points for Wan2.2-5B.Table 4 compares native VAE versus DriveUniTok with backbones fixed, then evaluates DriveUniWAM across backbones using DriveUniTok.Source: §4.5 · p. 8, Table 4 · p. 9

5 CONCLUSION

The conclusion presents DriveUniTok and DriveUniWAM as a unified-latent approach to joint future-scene and ego-trajectory prediction. On NAVSIM, it reports improved prediction and planning with lower downstream computational cost than the natural-budget separate-stream control.

5 CONCLUSION

  • On NAVSIM, DriveUniWAM reaches 91.0 PDMS with supervised training and 91.54 after reinforcement learning post-training, while the unified design improves prediction and planning at lower downstream computational cost than the natural-budget separate-stream control.DriveUniTok and DriveUniWAM jointly predict future scenes and ego trajectories through a unified multimodal latent.Source: §5 · p. 9

AI USE STATEMENT

The authors used AI tools for code development and debugging, manuscript editing, and literature search and review, while retaining responsibility for the paper's content.

AI USE STATEMENT

  • AI tools assisted with code development and debugging, language editing, and literature search and review; the authors checked technical claims and take full responsibility for the paper.The authors describe the tools as research assistance rather than a substitute for author verification.Source: AI USE STATEMENT · p. 10

A TOKENIZER IMPLEMENTATION DETAILS

DriveUniTok aggregates RGB, depth, and ego information into shared per-frame tokens, compresses them into one latent, and decodes all three modalities from shared features. The implementation comparison matches representation and downstream computation budgets while retaining differences in tokenizer parameter counts and cross-modal connectivity.

A TOKENIZER IMPLEMENTATION DETAILS

  • The encoder initializes 576 spatial-reference queries and 8 learned-reference queries, then jointly attends over them and RGB/depth feature hierarchies to form 584 tokens.The frozen DINOv3-L feature backbone supplies the features; the two query groups differ in initialization, not modality.Source: A TOKENIZER IMPLEMENTATION DETAILS · p. 14
  • Temporal pooling reduces each unified query's T states to one sequence-level state, and a Variational Latent Bottleneck compresses it to width 16.The decoder maps the latent back to width 1024 and uses shared ViT blocks and modality-specific readouts.Source: A TOKENIZER IMPLEMENTATION DETAILS · p. 14
  • The equal-budget comparison matches 584 tokens, 9,344 latent scalars, downstream sequence length, nominal attention budget, sensors, data, and optimization recipe.It does not match trainable parameter counts: the separate baseline has independent components, while DriveUniTok uses a shared aggregation pathway.Source: A TOKENIZER IMPLEMENTATION DETAILS · pp. 14–15
  • Both comparison arms use frozen DINOv3-L features and a frozen semantic teacher, with matched reconstruction objectives, semantic supervision, depth preprocessing, and mask-token reconstruction.Neither arm uses a separate masking stage; the same training procedure includes mask-token reconstruction.Source: A TOKENIZER IMPLEMENTATION DETAILS · p. 14
  • DriveUniTok uses one shared query sequence for cross-modal exchange before compression; the separate baseline uses independent RGB, depth, and ego streams without tokenizer cross-modal interaction.The separate streams have their own query sets, encoder stacks, temporal compressors, and Variational Latent Bottlenecks.Source: A TOKENIZER IMPLEMENTATION DETAILS · p. 14
  • Tokenizer loss weights are 1 for RGB, depth, ego, and semantic terms, and 0.1 for KL; depth terms are skipped for batches with fewer than 100 valid points.Depth is clipped to [1, 150] m and normalized by an average scale of 60; the semantic target comes from a frozen DINOv3-L teacher.Source: A TOKENIZER IMPLEMENTATION DETAILS · p. 15, Table 5 · p. 15

B QUALITATIVE TOKENIZER RECONSTRUCTIONS

The qualitative reconstructions show that one unified latent preserves both RGB scene appearance and the spatial structure of sparse LiDAR depth. The supervised flow-matching setup trains a shared denoiser to predict velocity over the unified latent.

B QUALITATIVE TOKENIZER RECONSTRUCTIONS

  • Across nine scenes and five frames per scene, reconstructions preserve salient RGB details and sparse LiDAR depth structure within one compact representation.The RGB images retain road geometry and vehicles, while depth retains spatial structure and distance patterns.Source: B QUALITATIVE TOKENIZER RECONSTRUCTIONS · p. 15
  • The reconstruction comparisons pair RGB ground truth with reconstructed RGB, and projected-depth ground truth with depth reconstructed at valid LiDAR pixels.The same four-column layout is used across reconstruction groups 1–3.Source: Figure 4 · p. 16, Figure 5 · p. 17, Figure 6 · p. 18

C.1 SUPERVISED FLOW MATCHING

  • The supervised model uses one history clip and two future clips, with five frames per clip sampled at 2 Hz.Clip-relative poses are transformed into successive endpoint-anchored coordinate frames so shared boundary frames form a continuous trajectory.Source: §C.1 · p. 19
  • A single shared denoiser predicts velocity over the complete unified latent while the history remains clean and the future clips are independently noised.The two future clips share the same flow time.Source: §C.1 · p. 19
  • Training optimizes only the Wan DiT parameters, leaving DriveUniTok and its decoder frozen.The reported supervised training runs for 130,000 optimizer updates.Source: §C.1 · p. 19

C.2 REINFORCEMENT LEARNING POST-TRAINING

DriveUniWAM is post-trained with Flow-GRPO against the NAVSIM PDM reward using groups of jointly sampled futures. The update is regularized toward a frozen reference model, while the tokenizer and decoder remain fixed.

C.2 REINFORCEMENT LEARNING POST-TRAINING

  • For each scene, Flow-GRPO samples eight joint futures over ten stochastic denoising steps and scores their decoded trajectories with NAVSIM.Group rewards provide relative advantages for a clipped policy update.Source: §C.2 · p. 19
  • The clipped update computes likelihood ratios over transitions of the complete unified token sequence.The transition ratio compares the current and old policies; transitions are weighted by their contribution to the final sample.Source: §C.2 · p. 19
  • A frozen reference model supplies a behavior-cloning regularizer, and the tokenizer, decoder, and reference policy remain frozen during post-training.Only rank-16 LoRA adapters on denoiser attention projections are trained.Source: §C.2 · p. 19

D TRAINING OBJECTIVE ABLATION

The objective ablation finds that the full training recipe produces the best reported RGB, depth, and ego metrics. The comparison removes masked reconstruction and semantic supervision while holding the unified architecture and evaluation protocol fixed.

D TRAINING OBJECTIVE ABLATION

  • All variants share the unified architecture, 584 × 16 latent, complete camera, depth, and ego inputs at evaluation, and the same NAVSIM training and evaluation protocol.The comparison is intended to isolate training-objective effects from representation differences.Source: Table 6 · p. 20, D TRAINING OBJECTIVE ABLATION · p. 20
  • The ablation removes masked reconstruction and semantic supervision while retaining the RGB DINO perceptual objective when semantic supervision is removed.Intra-modal masking reconstructs hidden patches or tokens from the same modality; cross-modal masking reconstructs a hidden modality from the others.Source: D TRAINING OBJECTIVE ABLATION · pp. 19–20
  • The full recipe reduces LPIPSDINO from 0.1540 to 0.1063, rFVDVJ from 8.4000 to 3.4546, depth AbsRel from 0.0695 to 0.0461, and trajectory ADE from 0.1450 to 0.0728.It achieves the best results across all reported RGB, depth, and ego metrics.Source: D TRAINING OBJECTIVE ABLATION · p. 20

E NAVSIM-V2 PLANNING RESULTS

On NAVSIM-v2, DriveUniWAM reaches 90.3 EPDMS while maintaining high scores on individual metrics. Modality sensitivity results favor using all three inputs, and qualitative comparisons show trajectory improvements after reinforcement-learning post-training.

E NAVSIM-V2 PLANNING RESULTS

F MODALITY SENSITIVITY

  • Using camera, LiDAR/depth, and ego-trajectory inputs together yields the strongest overall reconstruction and future-generation results.The modality variants mask inputs at evaluation time, so they are not treated as a complete sensor-failure robustness study.Source: F MODALITY SENSITIVITY · p. 21, Table 8 · p. 21

G QUALITATIVE NAVSIM PLANNING VISUALIZATIONS

  • Qualitative comparisons show improved trajectory shape, turning curvature, and endpoint placement for DriveUniWAM, with further gains after reinforcement-learning post-training.The figures compare front-view observations with VLA methods, a world-action baseline, and the two DriveUniWAM variants.Source: G QUALITATIVE NAVSIM PLANNING… · p. 21, Figure 7 · p. 22
  • Figures 7–12 present qualitative NAVSIM navtest visualizations for scene groups 1–6; Figure 7 identifies columns for the front-view observation, AutoVLA, ReCogDrive, PWM, DriveUniWAM, and its reinforcement-learning post-training variant.Source: Figure 7 · p. 22, Figure 8 · p. 23, Figure 9 · p. 24

Related papers

Loading 2610.35044v1…