Technical Summary: Planning-Aligned Token Compression for Long-Context Autonomous Driving

Problem Statement

Monolithic Vision-Action (VA) and Vision-Language-Action (VLA) policies represent a paradigm shift in autonomous driving, mapping multi-modal inputs directly to vehicle trajectories via unified transformer backbones. However, scaling these models to handle complex interactions requires extended temporal context, which rapidly inflates the visual token sequence length beyond real-time computational budgets.

Existing solutions face a trade-off:

• Linear Transformers/External Memory: Often require backbone modifications or introduce architectural complexity.
• Token Compression: The most practical approach as it requires no backbone changes. However, current methods rely on rule-based heuristics (e.g., temporal decay), which retain recent frames and discard older ones indiscriminately.
• The Core Limitation: Rule-based compression is decoupled from the planning objective. It fails to distinguish between redundant information and decision-critical historical cues (e.g., the arrival order of vehicles at a four-way stop). As illustrated in the paper, discarding distant observations (e.g., at $t=-3s$) can lead to incorrect right-of-way decisions, causing safety violations like rolling through stops.

Methodology: COMPACT-VA

The authors propose COMPACT-VA (Compression via Planning-Aligned Context Tokens), a working memory framework designed to compress extended context into bounded representations while explicitly aligning compression with driving intent.

1. Hierarchical Temporal Context Buffering

Instead of a flat context window, the observation history ($T$ timesteps) is organized into a hierarchical FIFO memory bank with $K$ compression layers.

• Temporal Decay Inductive Bias: Recent frames retain high token resolution (no compression), intermediate frames undergo moderate compression, and distant frames undergo aggressive compression.
• Q-former Compression: At each layer, learnable query tokens are concatenated with raw observation tokens, timestep embeddings, and camera embeddings. A Q-former module (based on MMDiT) processes this via self-attention to selectively aggregate task-relevant visual features.
• Positional Consistency: Rotary Positional Embeddings (RoPE) are scaled by the compression ratio to maintain alignment with the uncompressed sequence.

2. Planning-Aligned Variational Coupling

The core innovation is a Conditional Variational Auto-Encoder (cVAE) with Vector Quantization (VQ) that ties compression quality directly to trajectory prediction.

• Driving Intent Latent ($z$): A compact discrete latent variable is distilled from future trajectories during training.
• Posterior Encoder (Q-net): Trains only during the training phase. It takes the compressed observations ($o_{compressed}$) and the ground-truth future trajectory ($\tau_{future}$) to extract the driving intent latent $z_q$.
• Prior Encoder (P-net): Operates during both training and inference. It predicts the intent latent $z_p$ using only the compressed observations, without access to future information.
• Vector Quantization: Both $z_q$ and $z_p$ are mapped to a shared discrete codebook. This discretization forces the model to learn a finite set of driving "skills" or intents.
• End-to-End Optimization: The system is optimized via a composite loss:
  $$\mathcal{L} = \mathcal{L}_{traj} + \lambda_{KL} \cdot D_{KL}(q_\phi(z|o, \tau_{future}) \parallel p_\theta(z|o_{compressed})) + \lambda_{commit} \cdot \mathcal{L}_{commit}$$
  The KL divergence term ensures the prior (inference-time) distribution matches the posterior (training-time) distribution. If the Q-former discards critical history, the prior cannot predict the intent, resulting in high KL divergence and degraded trajectory prediction. This creates a closed-loop where the model learns what to keep.

3. Inference

At test time, only the prior pathway is active. The model compresses observations, predicts the latent $z_p$, quantizes it, and appends it as a special token to the compressed memory. This combined sequence feeds the unified transformer backbone for autoregressive trajectory generation.

Evaluation Framework

The authors identify high-signal dynamic scenarios where extended context (5–10 seconds) is critical for discrete decision-making, distinct from continuous trajectory optimization:

1. Four-Way Stops: Determining right-of-way based on arrival order.
2. Stop/Yield with Dynamic Occlusion: Maintaining state of occluded vehicles.
3. Unprotected Turns: Assessing oncoming traffic gaps over time.

Metrics: Traditional metrics like minADE are deemed misaligned. The paper introduces behavioral metrics:

• Go Success Rate (Go SR): Ability to proceed after stopping.
• Stop Success Rate (Stop SR): Ability to stop completely when required.
• Roll-Through Rate: Frequency of failing to stop completely (safety critical).
• Stop Position/Duration Errors: Spatial and temporal deviations.

Key Results

Evaluated on the Alpamayo dataset (specifically curated memory-dependent scenarios) and Alpasim simulator:

• Performance Gains: Under comparable token budgets (~1,424 tokens), COMPACT-VA achieves a 68.3% Go Success Rate, outperforming the standard 1-second baseline (63.8%) and a dense 5-second uncompressed baseline (61.9%).
• Safety Improvements: It reduces the Roll-Through Rate by 22% (7.0% vs. 9.0% for baselines) and improves Stop Success Rate to 89.2%.
• Ablation Insights:
  • Removing planning alignment drops Go SR to 65.6%, confirming the necessity of coupling compression with intent.
  • Naïve compression (without history conditioning) yields 63.5%, showing the value of trajectory-conditioned compression.
  • Hierarchical compression is essential; recent frames must retain full resolution.
• Efficiency: Compared to uncompressed 5-second processing (6,400 tokens), COMPACT-VA achieves a 3.3× inference speedup and 2.7× memory reduction (3.95 GB vs. 10.51 GB peak GPU memory) while maintaining general driving competence in closed-loop simulations.

Significance and Claims

The paper claims that task-aware working memory is key to reasoning in diverse autonomous driving scenarios. By moving from rule-based heuristics to planning-aligned compression, the model implicitly discovers and retains decision-critical historical information without hand-crafted retention rules.

The authors position this work as a step toward effective memory for monolithic VA/VLA policies, demonstrating that bounded context windows, when coupled with variational objectives, can resolve the tension between computational constraints and the need for long-horizon reasoning in safety-critical driving decisions. They note that while gains on general driving are modest, the approach is particularly impactful in high-signal scenarios where discrete decisions (stop/go) depend on behavioral cues from the past 5–10 seconds.