Source-linked AI summary

Copilot4D: Learning Unsupervised World Models for Autonomous Driving via Discrete Diffusion

Lunjun Zhang, Yuwen Xiong, Ze Yang, Sergio Casas, Rui Hu, Raquel Urtasun

arXiv:2311.01017v4cs.CVcs.AIcs.LGcs.RO

TL;DR

Unsupervised world models have scaled less rapidly in autonomous driving than GPT-like language models, in part because observations are complex and generation must scale. Copilot4D tokenizes point-cloud observations and predicts the future with discrete diffusion; it reduces prior state-of-the-art Chamfer distance across three datasets at both 1s and 3s horizons.

  • Problem

    Unsupervised world-model scaling in autonomous driving lags GPT-like language modeling, with complex observations and scalable generation identified as bottlenecks.

    Source: §1 · p. 2
  • Method

    Copilot4D tokenizes sensor observations with VQVAE and predicts future point clouds using discrete diffusion, recasting MaskGIT for parallel token decoding.

    Source: §1 · p. 2, §4.1 · p. 4, §4.2 · p. 5
  • Results

    65%−75% lower prior-SOTA Chamfer Distance for 1s prediction and more than 50% lower for 3s prediction across NuScenes, KITTI Odometry, and Argoverse2.

    Source: §5 · p. 9
  • Takeaways & Limitations

    The results support discrete diffusion on tokenized agent experience as an approach to unsupervised world modeling for real-world autonomous-driving point clouds.

    Source: §1 · p. 2, §6 · p. 9

Abstract

from arXiv · show

Learning world models can teach an agent how the world works in an unsupervised manner. Even though it can be viewed as a special case of sequence modeling, progress for scaling world models on robotic applications such as autonomous driving has been somewhat less rapid than scaling language models with Generative Pre-trained Transformers (GPT). We identify two reasons as major bottlenecks: dealing with complex and unstructured observation space, and having a scalable generative model. Consequently, we propose Copilot4D, a novel world modeling approach that first tokenizes sensor observations with VQVAE, then predicts the future via discrete diffusion. To efficiently decode and denoise tokens in parallel, we recast Masked Generative Image Transformer as discrete diffusion and enhance it with a few simple changes, resulting in notable improvement. When applied to learning world models on point cloud observations, Copilot4D reduces prior SOTA Chamfer distance by more than 65% for 1s prediction, and more than 50% for 3s prediction, across NuScenes, KITTI Odometry, and Argoverse2 datasets. Our results demonstrate that discrete diffusion on tokenized agent experience can unlock the power of GPT-like unsupervised learning for robotics.

1 INTRODUCTION

Copilot4D addresses the challenges of unsupervised world modeling for autonomous-driving point clouds by tokenizing observations and using discrete diffusion to predict the future. Across three datasets, it substantially reduces prior state-of-the-art Chamfer distance at both 1s and 3s horizons.

  • World models predict the next observation from past observations and the current action, learning from unlabeled agent experience for learning and planning.The paper situates this approach within model-based reinforcement learning.Source: §1 · p. 1
  • Unsupervised world modeling for autonomous driving lags GPT-like language modeling, while prediction systems still rely on supervised boxes or segmentation.The paper asks what makes it difficult to predict future observations directly from unlabeled experience.Source: §1 · pp. 1–2
  • Two bottlenecks are complex, unstructured observations that are difficult to model and the need to scale generation across tens of thousands of tokens per observation.VQVAE tokenization addresses the observation-space challenge; parallel decoding addresses generation scale without assuming all tokens are conditionally independent.Source: §1 · p. 2
  • The paper recasts MaskGIT as discrete diffusion, making a few simple changes that notably improve upon it.The approach draws on MaskGIT's iterative parallel decoding of arbitrary numbers of tokens.Source: §1 · p. 2
  • Copilot4D tokenizes each observation frame with VQVAE, then autoregressively predicts future point clouds using discrete diffusion on token indices.Its tokenizer uses implicit representation for ray-depth volume rendering, and its world model interleaves spatial and temporal Transformer blocks in BEV.Source: §1 · p. 2
  • Copilot4D reduces prior state-of-the-art Chamfer distance by 65%−75% for 1s prediction and more than 50% for 3s prediction across NuScenes, KITTI Odometry, and Argoverse2.The approach also predicts multimodal future observations at a 3s horizon.Source: §1 · p. 2

2 RELATED WORK

Prior world-modeling work has combined learned observation representations with dynamics models, while discrete diffusion and 3D point-cloud representations offer relevant modeling tools. Copilot4D draws on these strands by modeling dynamics over discrete tokens and using an implicit representation in its tokenizer.

  • Ha & Schmidhuber (2018) used a VAE and RNN for world modeling, while Dreamer-v2 (Hafner et al., 2020) and IRIS (Micheli et al., 2022) improved Atari results with discrete latents and Transformer dynamics.These prior works guide the paper's use of learned representations and Transformer-based dynamics for autonomous-driving point-cloud forecasting.Source: §2 · p. 2
  • Diffusion for discrete data received less attention than continuous diffusion; MaskGIT (Chang et al., 2022) trains on aggressively masked image tokens and iteratively decodes tokens in parallel.MaskGIT has been applied to text-to-image generation, video prediction, and point-cloud generation; this work examines its connection to discrete diffusion and possible improvements.Source: §2 · p. 3
  • Point-cloud work combines BEV processing for perception with implicit neural scene representations for 3D generation, and Copilot4D's tokenizer draws on both.After tokenization, the world model operates on discrete tokens rather than the specific 3D representations.Source: §2 · p. 3

3 BACKGROUND: DIFFUSION MODELS

Diffusion models define a forward process from data to noise and learn a reverse process. For discrete data, methods differ in how they parameterize that reverse process and handle masking, while classifier-free guidance modifies conditional sampling.

  • Diffusion models define a closed-form forward process from data to noise and learn a reverse process from noise to data.Continuous diffusion typically uses Gaussian noise to simplify training objectives and speed inference.Source: §3 · p. 3
  • D3PM (Austin et al., 2021) parameterizes the discrete reverse process to directly infer x0, enabled by summing over discrete probabilities.The original diffusion loss also requires a posterior that depends on the cumulative forward transition matrix having a closed-form solution.Source: §3 · p. 3
  • Absorbing diffusion gradually masks tokens, but its posterior is not well-defined when xk differs from x0; MaskGIT instead trains masked-token prediction and samples through iterative parallel decoding.MaskGIT bases its procedures on BERT, without the diffusion loss's posterior calculation described here.Source: §3 · p. 3
  • Classifier-free guidance samples using conditional predictions amplified relative to unconditional predictions, and MaskGIT guidance modifies logits analogously.The guidance weight w controls the added difference between conditional and unconditional logits.Source: §3 · p. 3

4 METHOD: COPILOT4D

Copilot4D tokenizes LiDAR observations into Bird-Eye View (BEV) codes, then uses discrete diffusion to predict future observations. Its method combines a point-cloud tokenizer, a MaskGIT-based diffusion algorithm, mixed world-model training objectives, and a spatio-temporal Transformer.

4 METHOD: COPILOT4D

  • The world model predicts future point-cloud observations by applying discrete diffusion to tokenized observations, conditioned on past observations and actions.At inference, it starts from a fully masked next-frame token sequence, denoises it, and decodes it to render the next observation.Source: §4 · p. 4

4.1 TOKENIZE THE 3D WORLD

  • The VQVAE-like tokenizer encodes point clouds as BEV latent codes and reconstructs them through differentiable depth rendering.Its decoder also predicts voxel occupancy for spatial skipping during inference; training combines vector-quantization and rendering losses.Source: §4.1 · pp. 4–5

4.2 MASKGIT AS A DISCRETE DIFFUSION MODEL

  • Copilot4D recasts MaskGIT as absorbing-uniform discrete diffusion, adding uniform noise to unmasked tokens and applying reconstruction loss beyond masked locations.Sampling also iteratively denoises previously sampled tokens, while requiring a single trained model rather than two separate models.Source: §4.2 · p. 5

4.3 LEARNING A WORLD MODEL

  • World-model training mixes future prediction, joint past-and-future denoising, and unconditional frame modeling to support imperfect past conditioning and classifier-free guidance.The three objectives are sampled randomly during training; the joint objective is a harder pretraining task than future prediction alone.Source: Figure 4 · p. 6, §4.3 · p. 6

4.4 A SPATIO-TEMPORAL TRANSFORMER FOR WORLD MODELING

  • A spatio-temporal Transformer interleaves per-frame Swin spatial attention with GPT2 temporal attention and uses a U-Net structure with residual connections.Temporal attention is causal for conditional objectives and becomes an identity mask for unconditional frame modeling; inference decodes one frame at a time.Source: §4.4 · p. 7

5 EXPERIMENTS

Copilot4D outperforms prior unsupervised point-cloud forecasting methods across three driving datasets, while the experiments also assess evaluation scope, classifier-free guidance, and changes to MaskGIT.

Datasets and experiment setting

  • Evaluation follows prior protocols, using past point clouds and future ego poses to predict 1s or 3s; NuScenes uses different frame counts from KITTI and Argoverse2.Metrics include Chamfer distance, L1 depth, and relative L1 error; the authors report ROI-computed metrics because uncropped ground truth can inflate errors.Source: §5 · pp. 8–9

Benchmark against state-of-the-art

  • For 1s prediction, Copilot4D reduces prior-SOTA Chamfer distance by 65%–75% across three datasets; for 3s prediction, it reduces Chamfer by more than 50%.The comparisons cover NuScenes, KITTI Odometry, and Argoverse2.Source: Table 1 · p. 8, Table 2 · p. 8, §5 · p. 9

Classifier-free diffusion guidance

  • On NuScenes 3s prediction, classifier-free guidance reduces Chamfer distance by as much as 60%.The authors use past agent history as conditioning and describe guidance as amplifying conditioned information.Source: §5 · p. 9, Table 3 · p. 9

Improvement upon MaskGIT

  • On NuScenes 3s prediction, the discrete-diffusion changes reduce Chamfer distance by 29% versus MaskGIT while decoding an average of 1,600 new tokens per step.Each frame has 128 × 128 tokens and uses 10 diffusion steps; iterative denoising can revise already sampled tokens.Source: §5 · p. 9, Table 4 · p. 9

6 CONCLUSION

The conclusion presents Copilot4D as an approach to unsupervised world modeling that combines observation tokenization, discrete diffusion, and Transformers. It reports state-of-the-art point-cloud forecasting results and identifies model-based reinforcement learning as a direction for future work.

6 CONCLUSION

  • Copilot4D combines observation tokenization, discrete diffusion, and Transformers to build unsupervised world models, achieving state-of-the-art results in autonomous-driving point-cloud forecasting.The authors identify simplifying complex observations and building scalable generative models as practical bottlenecks.Source: §6 · p. 9
  • The authors propose combining this world-modeling approach with model-based reinforcement learning in future work to improve autonomous agents’ decision-making capabilities.Source: §6 · p. 9

A.1 ADDITIONAL QUANTITATIVE RESULTS

Additional results report stronger zero-shot transfer from Argoverse2 to KITTI than 4D Occupancy and describe tokenizer findings on spatial skipping and reconstruction detail.

Zero-shot transfer performance across datasets

  • Copilot4D achieves more than 4× smaller Chamfer distance than 4D Occupancy on both 1s and 3s zero-shot KITTI predictions after training on Argoverse2.The paper reports that test-set Argoverse2 results are very similar to its validation-set results.Source: §A.1 · p. 14, Table 5 · p. 14, Table 6 · p. 14

Ablation studies on the tokenizer

  • Spatial skipping improves point-cloud reconstruction, while the proposed tokenizer recovers fine geometry that voxel-presence predictions from UltraLiDAR miss.The paper also describes differentiable depth rendering with an implicit representation as crucial.Source: §A.1 · p. 14

A.2 MODEL DETAILS

Copilot4D pairs lightweight tokenizer and world-model networks with a multiscale U-Net Transformer; spatial skipping improves tokenizer reconstructions.

A.2 MODEL DETAILS

  • The tokenizer has 13 million parameters and the world model 39 million, fewer than ResNet-34 and ResNet-101, respectively.The cited ResNet models have 21.8 million and 44.5 million parameters, respectively.Source: §A.2 · p. 14
  • The world model uses a U-Net Transformer that interleaves spatial and temporal blocks across multiple feature resolutions.The tokenizer combines PointNet, BEV representation, a Neural Feature Grid, and differentiable depth rendering.Source: Figure 7 · p. 15, Figure 6 · p. 15
  • Spatial skipping improves tokenizer point-cloud reconstructions compared with training without the coarse reconstruction branch or skipping process.The ablation table reports the comparison; the accompanying figure describes the skipping process.Source: Table 7 · p. 16, Figure 8 · p. 16

A.2.1 TOKENIZER

The tokenizer maps point clouds to discrete Bird-Eye View tokens and reconstructs geometry using differentiable depth rendering over an implicit feature representation.

A.2.1 TOKENIZER

  • The VQVAE-style tokenizer encodes point clouds as Bird-Eye View tokens and reconstructs them through differentiable depth rendering.Its encoder begins with voxel-wise PointNet features, aggregates them along the z-axis into a BEV representation, and uses a Swin Transformer backbone.Source: §A.2.1 · pp. 16–17
  • The encoder outputs a 128 × 128 representation, with each grid covering a 1.25m × 1.25m BEV region, before vector quantization.The paper also notes that LayerNorm before quantization and increasing the feature dimension to 1024 are important for reconstruction.Source: §A.2.1 · p. 17
  • When more than 3% of codebook entries are unused for 256 iterations, K-Means on recent encoder outputs reinitializes the codebook after at least 200 iterations.The memory bank stores ten times as many outputs as the codebook has entries.Source: §A.2.1 · p. 17
  • The decoder's Neural Feature Grid represents continuous 3D coordinates for occupancy queries, which feed differentiable depth rendering and an L1 depth loss.A separate coarse reconstruction branch predicts occupied voxels for inference-time spatial skipping.Source: §A.2.1 · p. 18
  • Spatial skipping samples ray points only within max-pooled estimated non-empty voxels; the implicit representation and depth rendering yield finer reconstructions than UltraLiDAR's occupancy-only output.The tokenizer uses a coarse occupancy estimate to restrict queries into the 3D feature grid.Source: Figure 9 · p. 17, §A.2.1 · p. 18

A.2.2 WORLD MODEL

The world model predicts VQVAE codebook logits with a multiscale U-Net that interleaves spatial and temporal attention and conditions on observations, ego poses, and positional encodings.

A.2.2 WORLD MODEL

  • The model maps 128 × 128 token inputs to 128 × 128 × 1024 logits, using a three-level U-Net with spatial Swin and temporal GPT-2 blocks.Every two spatial blocks are followed by one temporal block, which attends at the same feature-map location across frames.Source: §A.2.2 · p. 18
  • The U-Net transitions between feature levels with Patch Merging downsampling and Level Merging upsampling, concatenation, normalization, and projection.Level Merging combines the upsampled higher-level map with the lower-level map through a residual connection.Source: §A.2.2 · pp. 18–19
  • Conditioning inputs include observation-token embeddings, ego-vehicle pose embeddings, spatial positional encodings, and temporal positional encodings.The pose matrices are flattened to 16-dimensional vectors and added to feature-map locations for their corresponding frames.Source: §A.2.2 · p. 19
  • The model uses pre-norm Transformer blocks so conditioning information added at the network input can pass through subsequent residual modules.Non-positional conditioning inputs pass through Linear, LayerNorm, and Linear layers before entering the Transformer.Source: §A.2.2 · p. 19

A.3 TRAINING TRANSFORMERS

The appendix describes Transformer training and discrete-diffusion derivations, then details absorbing-uniform diffusion training, inference, and additional qualitative outcomes and limitations.

A.3 TRAINING TRANSFORMERS

  • The Transformer uses language-model initialization and optimization strategies, with warmup followed by cosine decay and shared hyperparameters across datasets.The world model uses a 0.001 learning rate, 2,000 warmup iterations, batch size 8, and a 0.75-million-iteration cosine decay.Source: §A.3 · pp. 19–20
  • Classifier-free guidance contrasts logits conditioned on past observations and actions with unconditioned logits, while sampling uses a cosine mask schedule.The guidance can be implemented with one forward pass per diffusion step by extending the temporal sequence and using a causal mask for past frames.Source: Figure 10 · p. 20, Figure 11 · p. 20

A.4 ALTERNATIVE DERIVATION FOR DISCRETE DIFFUSION TRAINING OBJECTIVE

  • The alternative discrete-diffusion derivation presents a lower-bound objective that avoids computing the posterior and reverse transition distribution defined in the paper.The forward transition is categorical, and its cumulative transition matrix is formed by composing the stepwise matrices.Source: §A.4 · pp. 21–22

A.5 ABSORBING-UNIFORM DISCRETE DIFFUSION

  • Absorbing-uniform diffusion combines transitions that mask tokens with transitions that replace tokens by random non-mask categories.The appendix defines the probability of retaining an unmasked, unchanged token and the probability of having a mask token after k steps.Source: §A.5 · pp. 23–24

A.6 TRAINING AND INFERENCE OF ABSORBING-UNIFORM DISCRETE DIFFUSION

  • Training applies absorbing masking and uniform noise; the uniform-noise schedule makes the noise ratio linear in the step, while the absorbing schedule is cosine-based.The default maximum uniform-noise level is η = 20%.Source: §A.6 · p. 24
  • During inference, the model directly decodes predicted x0 tokens, and top-3 sampling is used instead of vanilla categorical sampling.Once a token becomes non-mask, it may be revised but cannot return to the masked state.Source: §A.6 · p. 25

A.7 ADDITIONAL QUALITATIVE RESULTS

  • Counterfactual ego trajectories produce imagined futures consistent with the actions, and the model represents other vehicles as reactive agents.In one braking scenario, the vehicle behind also brakes to avoid a collision.Source: §A.7 · p. 25, Figure 12 · p. 26
  • Three-second vehicle behavior remains less accurate than one-second prediction, and the model does not yet appear to generate new vehicles entering the region of interest.The authors note that multimodal future outcomes could partly explain the lower 3s accuracy.Source: §A.7 · p. 25
Loading 2311.01017v4…