Vector Quantization for Trajectory Prediction

IROS 2026

Carter Fang, Thomas Gilles, Sergio Casas, and Raquel Urtasun

by Waabi

We present VQTraj, a novel trajectory decoding framework that transforms motion forecasting into a classification problem by learning discrete trajectory latents with Vector Quantization (VQ). Unlike previous classification-based methods constrained by heuristics, we learn priors directly from trajectories. Our novel formulation yields an interpretable latent space where discrete codes directly correspond to semantic driving behaviors. Designed as a modular, drop-in replacement for standard regression-based decoders, we demonstrate VQTraj's effectiveness by upgrading state-of-the-art models. Crucially, VQTraj improves the accuracy of the single most-likely prediction on Argoverse 2 and improves zero-shot generalization performance on the Waymo Open Motion Dataset.

Why Vector Quantization?

From recurrent neural networks to transformers, architectures for scene encoding in motion forecasting networks have evolved significantly over time. Despite this, trajectory decoders have remained relatively simplistic with modern networks often opting for linear regressors that predict a fixed number of modes and are supervised with Winner-Takes-All. Alternatives to regression like classification have been explored in the past, but they relied on handcrafted methods for generating discrete trajectory labels. VQTraj revisits the classification-based paradigm but leverages discrete latent trajectory priors learned with Vector Quantization.

Learning Discrete Trajectory Codebooks

In the first stage of the framework, VQTraj acts as an autoencoder to learn a codebook of discrete latent priors by reconstructing future trajectories. The quantized autoencoder has three core modules: an encoder that compresses a future trajectory into a continuous latent, a vector quantizer that discretizes the latent, and a decoder that maps quantized latents back to future trajectories.

Trajectory Prediction as Code Classification

In the second stage of the framework, we train a standard motion forecasting network with our pretrained VQTraj network occupying the place of the conventional decoder. Using the learned codebook, trajectory decoding is framed as a classification task over discrete latent priors. In this stage, all pretrained VQTraj components are frozen. Supervision combines cross-entropy classification loss with an auxiliary loss in trajectory-space based on Gumbel-top-k.

Results on Public Benchmarks

We demonstrate the effectiveness and generality of VQTraj by applying it to two state-of-the-art forecasting architectures: RealMotion and ForecastMAE on the Argoverse 2 dataset. In both settings, VQTraj achieves significant improvements in single-mode accuracy as well as ranked multi-modal predictions. RealMotion-I with VQTraj achieves state-of-the-art results for primary mode accuracy (minFDE1 and minADE1) while improving multi-modal ranking metrics (bminFDE6) by 4.3% over baseline models. We also highlight that VQTraj flexibly allows us to set mode count at runtime for each actor, in contrast to regression-based networks that predict a fixed number of modes for every actor.

We also evaluate VQTraj on the Waymo Option Motion Dataset (WOMD). Without fine-tuning, models trained on Argoverse 2 equipped with VQTraj were evaluated zero-shot on the WOMD validation split. VQTraj demonstrated remarkable generalization, improving bminFDE6 by 10.1%, minADE1 by 14.3%, and minFDE1 by 6.9% over baseline architectures.

Model Introspection

While many vector quantization schemes are possible, quantizing trajectories with FSQ facilitates interpretability of our learned latent codes. We can inspect how the different FSQ dimensions affect decoded trajectories by perturbing the code of an actor along each dimension. For example, in the figure, level 0 captures macro rightward lane, level 1 captures macro left-turn behaviors and lateral maneuvers, and level 2 controls fine-grained longitudinal endpoint variation and speed profile adjustments.

@inproceedings{fang2026vqtraj,
  title     = {Vector Quantization for Trajectory Prediction},
  author    = {Fang, Carter and Gilles, Thomas and Casas, Sergio and Urtasun, Raquel},
  booktitle = {IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)},
  year      = {2026}
}