Source-linked AI summary
Categorical Traffic Transformer: Interpretable and Diverse Behavior Prediction with Tokenized Latent
Yuxiao Chen, Sander Tonkens, Marco Pavone
TL;DR
Traffic models need accurate, diverse, interpretable behaviors that remain compatible with downstream planning and emerging LLM interfaces. CTT addresses these needs with directly supervised categorical scene modes and token-compatible architecture, achieving SOTA accuracy across multiple open datasets while enabling semantic controllability. Its evaluation scope includes a practical runtime limitation for autoregressive decoding on long-horizon WOMD.
Problem
Multimodal traffic models often use uninterpretable latent spaces that lack direct supervision, causing mode collapse and limiting semantic controllability and LLM communication.
Source: §1 · p. 1Method
CTT uses directly identifiable categorical scene modes comprising agent-to-lane and agent-to-agent modes, with consistency-trained trajectory decoding and token-compatible transformer/GNN architecture.
Source: §1 · pp. 1–2, §3 · p. 3Results
Takeaways & Limitations
Interpretable tokenized modes let CTT decouple encoder and decoder training, avoid mode collapse, and connect traffic behavior with LLM-compatible semantic representations.
Source: §1 · p. 2, §5 · p. 9On the long-horizon WOMD setting, autoregressive decoder training requires too much GPU memory, so only inference runtime was tested.
Source: §4.2 · p. 8
Abstract
from arXiv · showhide
Adept traffic models are critical to both planning and closed-loop simulation for autonomous vehicles (AV), and key design objectives include accuracy, diverse multimodal behaviors, interpretability, and downstream compatibility. Recently, with the advent of large language models (LLMs), an additional desirable feature for traffic models is LLM compatibility. We present Categorical Traffic Transformer (CTT), a traffic model that outputs both continuous trajectory predictions and tokenized categorical predictions (lane modes, homotopies, etc.). The most outstanding feature of CTT is its fully interpretable latent space, which enables direct supervision of the latent variable from the ground truth during training and avoids mode collapse completely. As a result, CTT can generate diverse behaviors conditioned on different latent modes with semantic meanings while beating SOTA on prediction accuracy. In addition, CTT's ability to input and output tokens enables integration with LLMs for common-sense reasoning and zero-shot generalization.
1 Introduction
CTT addresses mode collapse, limited interpretability, and missing token-level communication in multimodal traffic prediction by using directly supervised semantic scene modes. Its architecture combines interpretable latent modes, consistency training, and token-compatible transformer/GNN components, achieving SOTA accuracy across three public datasets.
Motivation
- CTT targets mode collapse and insufficiently diverse, interpretable multimodal predictions that limit downstream planning and simulation.Existing latent spaces lack direct supervision, causing modes to converge toward the most likely behavior; fixed discrete modes also restrict expressiveness and semantic interpretation.Source: §1 · p. 1
- Tokenized traffic predictions are designed to provide a communication channel between traffic models and LLMs for driving-domain reasoning.SOTA traffic models lack a token interface connecting multimodal predictions with semantic concepts consumed or generated by LLMs.Source: §1 · p. 1
Contribution
- CTT uses an interpretable latent space of agent-to-agent and agent-to-lane modes, directly identified from driving logs and supervised during training.The ground-truth latent mode supervises the encoder, while trajectory reconstruction is conditioned on sampled scene modes.Source: §1 · pp. 1–2, Figure 1 · p. 2
- Consistency losses train trajectories under non-ground-truth modes to conform to their semantic modes, improving diversity beyond ground-truth trajectory supervision alone.The architecture also combines an equivariant transformer using Custom Edge Embedding with GNNs for tokenized edges and language-model compatibility.Source: §1 · p. 2
Results
- CTT achieves SOTA accuracy on nuScenes, nuPlan, and WOMD while providing strong controllability and scene consistency.The introduction reports evaluation across three public datasets and highlights semantic control over generated behaviors.Source: §1 · p. 2
2 Related works
Prior traffic models span recurrent, rasterized, graph, and transformer architectures, but scene-centric interaction modeling remains difficult and latent modes are often uninterpretable. Related work explores intermediate goals, intentions, lanes, homotopies, and emerging LLM-based AV systems, each leaving expressiveness, ambiguity, scalability, or integration challenges.
Data-driven traffic models
- Traffic models use diverse architectures and are either node-centric or scene-centric, with scene-centric models better aligned to planning but harder to train.Node-centric models can achieve better marginal coverage, yet combining independent predictions can produce inconsistent joint scenes.Source: §2 · p. 2
Interpretable intermediate prediction
- Goal conditioning, anchors, intentions, and multipath representations provide intermediate structure, but they can be difficult to resolve, ambiguous, or noninterpretable.Goal-based approaches may be node-centric, intentions have ambiguous definitions, and multipath anchors are not inherently interpretable.Source: §2 · p. 3
- Lane geometry and topology provide useful intermediate information, while future agent interaction is critical for scene consistency and downstream planning.Increasing latent modes raises joint-decoding complexity, whereas fixed mode sets may be uninterpretable or insufficiently expressive; clique decomposition can lose information.Source: §2 · p. 3
Homotopy
- Homotopy methods characterize interaction modes through topological relationships, including open-end homotopy for two-dimensional agents.In road traffic, winding angle is commonly used to identify homotopy classes.Source: §2 · p. 3
LLMs for AV
- LLM-based AV research covers driving-task understanding, sub-goal setting, behavior planning, and multimodal vision-language-action systems, but suitable fast and offline integration remains open.The related-work discussion frames architecture selection for AV stacks as unresolved under frame-rate and connectivity constraints.Source: §2 · p. 3
3 Categorical Traffic Transformer
CTT uses interpretable categorical scene modes to supervise multimodal trajectory prediction directly, combining agent-to-lane and agent-to-agent modes with continuous decoding. Its training and architecture jointly support diverse mode-conditioned predictions, LLM communication, and explicit relational reasoning.
3 Categorical Traffic Transformer
- CTT defines scene modes as interpretable categorical latent variables composed of agent-to-lane and agent-to-agent modes, forming a skeleton of future scene evolution.The scene mode is identified from ground-truth future trajectories, while the encoder predicts it and the decoder reconstructs trajectories conditioned on it.Source: §3 · p. 3, §3.1 · p. 4
3.1 Scene Modes
- Pairwise a2l modes encode agent-lane relations, while unitary a2l modes identify each agent’s lane or a null label.CTT uses pairwise modes to tokenize history relationships and unitary modes for scene-mode prediction because their smaller cardinality simplifies importance sampling.Source: §3.1 · p. 3
- a2a modes classify relative agent motion into static, clockwise, and counterclockwise homotopy classes, with symmetry enforced across agent pairs.The homotopy is defined over a time window rather than independently per frame, and the scene a2a mode combines non-self agent pairs.Source: §3.1 · p. 4
3.2 Training process
- CTT predicts marginal a2l and a2a distributions, then uses importance sampling and an energy function because the joint scene-mode space grows doubly exponentially with agent count.Factors are ranked using distance, agent-to-lane distance, and marginal concentration; selected factors generate candidate joint modes whose unnormalized likelihoods are learned.Source: §3.2 · p. 4
3.1 Scene Modes
- Scene modes can communicate traffic structure to LLMs: CTT rejects GPT-4’s low-probability left-lane-change suggestion, leading to slowing down and maintaining lane.The example combines GPT’s common-sense reasoning with CTT’s traffic-specific assessment of the candidate maneuver.Source: §3.1 · p. 4, Figure 2 · p. 5
3.2 Training process
- Training combines marginal and joint mode classification, reconstruction under the ground-truth scene mode, consistency losses, and regularization while decoding diverse alternative modes.The reconstruction loss uses ground-truth future trajectories, whereas consistency losses penalize violations of selected lane and homotopy relations.Source: §3.2 · pp. 4–5
3.3 Architecture
- The encoder combines transformer attention with GNN message passing over node and edge variables to explicitly model temporal, agent, lane, a2l, and a2a relationships.The architecture shares node variables between GNN and transformer updates, with edge variables supporting mode prediction and relational message passing.Source: §3.3 · pp. 5–6, Figure 3 · p. 6
- Custom Edge Embedding supplies relative geometric information while preserving coordinate-transform equivariance, and explicit edge modeling improves scene-mode prediction over a transformer alone.The architecture uses edge-specific functions for agent-agent and agent-lane geometry; attention among lane points performed worse than the projection method.Source: §3.3 · pp. 5–6
4 Results
CTT delivers accurate, scene-consistent, and controllable multimodal predictions, while its ablations show benefits from scalable data, one-shot decoding, dynamics, and diverse lane sampling.
4.1 Performance metrics
- CTT achieves SOTA performance on WOMD and significantly outperforms benchmarks on nuScenes, despite deliberately preserving scene-level consistency.The comparison is reported for accuracy in Tables 1 and 2.Source: §4.1 · p. 7, Table 1 · p. 7, Table 2 · p. 7
- CTT produces lower collision rates than most benchmarks across the most likely and all modes, with ScePT the main exception but a significantly longer inference time.Source: §4.1 · p. 7, Table 3 · p. 7
- Scene-mode conditioning enables controllable, semantically distinct trajectories, with CTT outperforming the benchmark in ML correct rate and coverage.Figure 4 illustrates lane-mode and homotopy effects, including yielding versus not yielding to a merging vehicle.Source: §4.1 · p. 7, Figure 4 · p. 8, Table 5 · p. 8
4.2 Ablation studies
- Training on nuPlan substantially outperforms training on nuScenes even when evaluated on nuScenes, indicating data-scale benefits and cross-dataset behavior generalization.Source: §4.2 · pp. 7–8, Table 6 · p. 8
4.1 Performance metrics
- Mode accuracy and consistency are evaluated across marginal lane, agent-interaction, and joint scene modes on nuScenes and WOMD.Longer WOMD horizons affect a2a accuracy, while null lane modes may increase a2l accuracy.Source: §4.1 · p. 7, Table 4 · p. 8
4.2 Ablation studies
- One-shot decoding has negligible performance differences from autoregressive decoding while substantially reducing runtime, especially for long horizons.All reported metrics therefore use one-shot decoding; autoregressive WOMD training also exceeds available GPU memory.Source: §4.2 · p. 8, Table 7 · p. 9
- Using terminal dynamics and diverse lane sampling both improves CTT accuracy, with lane-neighbor perturbations affecting importance sampling rather than mode prediction.Source: §4.2 · p. 8, Table 8 · p. 9
5 Conclusion and discussion
CTT combines an interpretable latent space with decoupled training to avoid mode collapse and support semantically controlled predictions. Its performance and token interface also support downstream planning and LLM integration.
5 Conclusion and discussion
- CTT’s interpretable latent space decouples encoder and decoder training, avoids mode collapse, and supports semantic controllability across multiple open datasets.Source: §5 · p. 9
Custom edge functions
CTT represents agents and lane points with heading-aware coordinates, then builds edge features in local relational frames for agent-agent, agent-lane, and lane-lane interactions. These geometric features are concatenated with categorical relationship modes during encoding and decoding.
Custom edge functions
- CTT represents each agent position or lane point as [X, Y, sin(θ), cos(θ)] to avoid angle wrapping and coordinate-frame dependence.Agent auxiliary features additionally include global coordinates and static properties; lane features concatenate lane points.Source: Custom edge functions · p. 12
- For agent-agent edges, CTT returns relative position and the second agent’s static features.The edge function is defined as concatenating the relative position with the second agent’s static features.Source: Custom edge functions · p. 12
- For agent-lane edges, CTT encodes the agent’s projection onto the lane polyline together with the lane segment’s endpoints in the agent’s local frame.This representation captures the agent’s relation to the lane geometry without relying on global coordinates.Source: Custom edge functions · p. 12
- For lane-lane edges, CTT concatenates relative positions among the start and end points of both lane segments.Source: Custom edge functions · p. 12
- These edge features are concatenated with categorical a2l, a2a, and l2l modes during both encoding and decoding.The l2l modes encode relationships such as next, previous, left-adjacent, or right-adjacent lanes.Source: Custom edge functions · p. 12
GPT query example
The GPT query example describes an autonomous-driving situation in which an ambulance is merging from a ramp while nearby traffic restricts lane changes. The response recommends adapting speed or position conservatively, communicating intentions, yielding when appropriate, and avoiding abrupt maneuvers.
GPT query example
- The scenario places the vehicle in the rightmost lane at 40 mph, 10 meters from a ramp merge, with an ambulance and a passenger car nearby.Source: GPT query example · p. 12
- If a safe gap exists, the driver should signal and merge left to make room for the ambulance.Source: GPT query example · p. 12
- If changing lanes is unsafe, the driver should adjust speed to create a gap or pass the merge point, depending on traffic and ambulance speed.Source: GPT query example · p. 12
- When the ambulance signals an emergency, the response says to facilitate its passage while maintaining lane control and avoiding abrupt braking or sudden maneuvers.The response also emphasizes monitoring surrounding traffic and communicating through signals, brake lights, or a horn when needed.Source: GPT query example · pp. 12–13
- After the ambulance merges, the driver should smoothly return to an appropriate speed if road conditions and speed limits permit.Source: GPT query example · p. 13