Planning-aligned Token Compression for Long-Context Autonomous Driving
arXiv:2606.07464v1 [cs.RO] 5 Jun 2026
Zhixuan Liang∗1,2 , Yuxiao Chen1 , Yurong You1 , Peter Karkus1 , Wenhao Ding1 , Boyi Li1 , Alexander Popov1 , Yan Wang1 , Maximilian Igl1 , Yiming Li1 , Danfei Xu1 , Nikolai Smolyanskiy1 , Boris Ivanovic1 , Ping Luo†2 , Marco Pavone†1
Abstract— Monolithic vision-action models represent an emerging paradigm in autonomous driving. However, this architecture produces token sequences that quickly exceed realtime computational budgets when encoding extended temporal context for complex interactions. While approaches like linear transformers and external memory try to make the context lightweight, token compression is most compatible with the architecture as it requires no backbone modifications. Yet existing compression adopts rule-based heuristics like temporal decay, decoupled from planning, risking loss of decision-critical information. We propose COMPACT-VA, a planning-aligned working memory framework built on conditional VQ-VAE, compressing extended context into bounded representations. Compression is conditioned on both historical trajectory and a learned planning intent that the posterior encoder distills from future trajectories during training, while the prior encoder learns to predict it from compressed observations. The compressed memory, concatenated with the predicted latent, feeds the policy for end-to-end optimization, planning with retained decisioncritical information. We evaluate on high-signal dynamic scenarios where historical context is most critical for behavior correctness (e.g., stop, yield, or proceed), and accordingly design behavioral metrics. Under comparable token budgets, we achieve >6% improvement (68.3%) on success rates with consistent gains across metrics. Ablations validate planning-aligned coupling effectiveness. Closed-loop evaluation confirms that COMPACTVA maintained general driving performance with 3.3× speedup and 2.7× memory reduction over uncompressed processing.
I. INTRODUCTION Vision-action (VA) and vision-language-action (VLA) policies [1] represent the latest paradigm in autonomous driving, directly mapping all modality inputs to vehicle trajectories through a unified transformer backbone. Unlike modular pipelines with separate perception–prediction–planning stages [2], [3] that maintain explicit state representations (e.g., bounding boxes) across modules, VLA policies instead encode the entire history directly within the observation token sequence, enabling fully end-to-end learning. Scaling these models to handle complex driving scenarios introduces a fundamental challenge that a longer history context quickly increases the visual sequence length beyond 1 Z. Liang, Y. Chen, Y. You, P. Karkus, W. Ding, B. Li, A. Popov, Y. Wang, M. Igl, Y. Li, D. Xu, N. Smolyanskiy, B. Ivanovic, and M. Pavone are with NVIDIA Research. {yuxiaoc, bivanovic,
mpavone}@nvidia.com 2 Z. Liang and P. Luo are with School of Computing and Data Science, The University of Hong Kong, {zxliang, pluo}@cs.hku.hk *This work was done during Zhixuan’s internship at NVIDIA. † Ping Luo and Marco Pavone are corresponding authors.
How to save decision-critical history under token budgets?
…
-5s
-4s
-3s
-2s
-1s
0s
Maximum context length
At t=0s, ego arrives at intersection and observes the cross-traffic vehicle. However, memory from t=-3s reveals this vehicle arrived first and has right-of-way → Ego must yield and stop.
Previous: No Context Extension / Rule-based Heuristics What to keep? Temporal Decay Heuristics
Planning-aligned Memory (Ours) Coupled with Driving Intent Driving Intent Latent
Keep recent, discard distant Context Window
Learn what to keep from future intent Memory Bank
…
-5s
-4s
-3s
-2s
-1s
0s …
Unknown Arrival Order Predict to rollthrough
No Context
Unable to determine go or wait
-5s
-4s
-3s
-2s
-1s
0s
Compressed Memory is Kept
Preserve decision-relevant history
Stop
Fig. 1: From limited context to planning-aligned memory. At an all-way sign intersection, determining right-of-way requires observing traffic over extended periods. Rule-based temporal decay with a 2s window (t=-2s to t=0s) discards distant observations, losing information about which vehicle arrived first (observed at t=-3s) and causing incorrect go decisions (62% success). Our planning-aligned compression retains decision-critical cues through coupling with driving intent via conditional VAE, preserving arrival order information for correct behavior (68.3% success).
the real-time computational budgets. Several approaches have been proposed to address this, including linear transformers [4] and external memory modules [5]. Among them, token compression has emerged as the most practical solution, it requires no backbone modifications and naturally supports the short to medium term memory horizon (on the order of tens of seconds) that driving decisions demand [6]. Despite this promise, existing compression methods rely on rule-based heuristics such as temporal decay, which retains recent frames while discarding older ones. These strategies are decoupled from the planning objective and cannot distinguish genuinely critical historical cues from other redundant information. As illustrated in Fig. 1, rule-based compression with temporal decay discards decision-critical historical information, while
our planning-aligned approach learns to retain it by coupling compression with driving intent prediction. To address this limitation, we propose a COMpression via Planning-Aligned Context Tokens framework (abbr. COMPACT-VA), a working memory mechanism that learns what historical information to retain through closedloop optimization of driving performance. Our approach is built on a conditional variational auto-encoder (cVAE) with vector quantization (VQ), introducing a compact discrete latent representation bridging compressed observations and trajectory prediction. We instantiate this framework with a Qformer-based compression module [7]. During compression, learnable query tokens are concatenated with raw observation tokens, historical trajectory information, and the learned latent, then processed through self-attention. This design enables compression conditioned on both past trajectory context and future driving intent, constituting planning-aligned working memory. Compressed tokens are organized into a hierarchical FIFO buffer retaining more tokens for recent frames and fewer for distant ones. The compressed memory feeds the policy backbone for end-to-end trajectory prediction (detailed in Sec. III). Through joint optimization, the model learns decision-relevant memory without hand-crafted rules. To evaluate our approach, we focus on high-signal scenarios where extended historical context plays a prominent role, such as four-way stops, dynamic occlusion and unprotected turns. These scenarios share a fundamental characteristic that good behaviors depend on making right discrete decisions (i.e., whether the vehicle stops when required and proceeds when appropriate) rather than trajectory optimization over a continuous spectrum. We design behavioral metrics including stop/go success rates, roll-through rates, and stop position/duration errors tailored to these decision-critical scenarios. Under comparable token budgets, COMPACT-VA achieves 68.3% go success rate versus ∼62.0% for baselines (+6.3%) and reduces safety-critical roll-throughs by 22%. Compared to compression without planning alignment (65.6%), COMPACTVA demonstrates +2.7% improvement, confirming that coupling compression with planning forces the model to retrain decision-critical information. Ablations validate that each of the sub-components, such as hierarchical compression, history conditioning, and planning coupling, contribute to improvements in high-signal scenarios. Closed-loop evaluation on 910 diverse scenarios confirms COMPACT-VA maintains general driving competence while achieving 3.3× speedup and 2.7× memory reduction over uncompressed scheme. In summary, our contributions are threefold: (1) We identify high-signal dynamic scenarios, including four-way stops, dynamic occlusion and unprotected turns as critical testbeds for memory-dependent driving, and introduce behavioral metrics assessing decision correctness beyond trajectory displacement. (2) We propose a planning-aligned working memory for VLA driving policies built on a conditional VAE, where compression quality is explicitly tied to trajectory prediction through a variational objective, enabling end-to-end learning of task-relevant compression. (3) Experiments demonstrate sizable improvements under comparable token budgets, and
achieve consistent gains across all metrics. Ablations validate each component’s contribution to performance. II. RELATED WORK A. Long-term Memory for Physical AI Long-term memory is essential for physical AI systems under partial observability, including robot manipulation [8], [9], [10], [11], [12], [13], navigation [14], [15], and autonomous driving [16], [17]. Classical robotics maintain persistent world state through explicit estimation (e.g., SLAM), while learningbased approaches embed memory via recurrence or beliefstate updates [18]. Latent world models have also been widely used to encode history and enable long-horizon reasoning [19], [20]. Recent VLA policies revisit external memory modules [5], [21] begin to adopt similar architectures. In contrast, autonomous driving and other high-rate embodied systems require memory mechanisms that are both computationally bounded and causally grounded, retaining only task-relevant history while discarding redundant observations. Our approach aligns with this view of memory as a compressed, persistent interface between past and present [22]. B. Long-Context Transformers with Recurrence A prominent direction for long-context modeling is to augment attention with explicit recurrence or external memory. Transformer-XL introduces segment-level recurrence [23], and the Compressive Transformer maintains dual memories to compress old activations to retain salient information [22]. A separate but related line explores memory as a learnable component updated at inference. TTT layers [24] make hidden states an expressive model updated via self-supervised learning, while Titans memorize during inference to complement local attention [25]. These works reinforce long-term memory to be an explicit, persistent bank rather than an ever-growing context window. Our compression draws inspiration from this compressive-memory view and adapts it to driving policy with multimodal inputs. Moreover, we introduce a variational objective tying compression directly to planning performance, ensuring decision-relevant information retained. C. Token Compression and Efficient Seq Architectures Orthogonal to recurrence, token-level compression reduces computational cost by merging or pruning tokens. ToMe [26] merges similar tokens based on feature similarity, while StreamingLLM [27] retains initial tokens as attention sinks to stabilize long-sequence decoding. Sparse attention, pruning, and pooling strategies have been widely explored in vision and video transformers. To address quadratic attention cost, sparseattention transformers (e.g., Longformer [28], BigBird [29]) limit patterns to local windows plus global tokens, while statespace models (e.g., Mamba [4]) achieve linear scaling through selective mechanisms. Hybrid designs (e.g., Hymba [30]) combine attention with SSM heads for efficient context summarization. While these methods accelerate inference or extend context, they apply rule-based or architecture-driven compression agnostic to downstream tasks. Hint-AD [31] shows that explicitly aligning intermediate representations
with planning semantics improves both interpretability and task performance in end-to-end driving. Our approach extends this alignment philosophy to token compression, learning what historical information to retain by coupling compression directly with the planning objective, discovering task-relevant patterns that may be overlooked by rule-based methods.
Ncompressed =
III. METHOD
K X
nk · Ncam · ⌊Nimg /rk ⌋.
(2)
k=1
A. Unified Vision-Action Model Backbone Our approach builds upon the unified vision-action (VA) policy variant of Alpamayo [1], which consists of three core components: a vision encoder processing multi-camera observations into visual tokens, a transformer backbone performing temporal reasoning without text instruction, and a trajectory decoder generating future vehicle motion. Multi-camera images from the current and past T timestamps are encoded into visual tokens through a pre-trained vision encoder (e.g., DINOv2 [32]). Each image produces Nimg tokens. With Ncam cameras per timestep for multi-view observation, the raw vision token count grows as: Nraw = T × Ncam × Nimg .
then progressively compressed through the layers, with layer Lk producing ⌊Nimg /rk ⌋ tokens per camera per frame, where rk is the cumulative compression ratio (i.e., the product of compression factors from layer 1 to layer k relative to the original Nimg ). The total compressed token count becomes:
(1)
These vision tokens, along with temporal positional embeddings and camera-specific embeddings, are concatenated with encoded historical trajectory information, and then fed into a causal transformer backbone. For historical trajectories, we apply sinusoidal positional embeddings with MLP compression to produce a single continuous token representing the ego history. For future trajectory, we adopt an FSQ-based tokenizer [33], compressing future waypoints into discrete tokens via finite scalar quantization, enabling autoregressive generation while maintaining reconstruction quality (c.f. Alpamayo [1]). While this unified architecture eliminates explicit intermediate modules, the sequence length grows linearly with the context length. For complex driving scenarios requiring extended temporal context, the token count grows substantially, easily surpassing typical VLM context windows. Without effective compression, the quadratic transformer attention 2 cost O(Nraw ) becomes intractable for deployment. B. Learned Hierarchical Temporal Context Buffering To manage token sequence growth while preserving temporal information, we compress raw observations through learned query-based aggregation organized into a hierarchical memory bank. The overall architecture of COMPACT-VA is shown in Fig. 2. We first describe the hierarchical compression module in this section, then introduce the planning-aligned variational coupling in Sec. III-C, which determines what information to retain. Hierarchical Buffer Structure. The observation history spanning T timesteps is organized into K compression layers {L1 , L2 , . . . , LK }, each applying different compression ratios to balance token efficiency and information preservation. Each PK layer Lk contains nk consecutive frames ( k=1 nk = T ). Compression is applied hierarchically in a cascading manner where frames are first encoded into Nimg tokens per camera,
We employ a multi-layer hierarchy following a temporaldecay heuristic where the most recent layer retains full token resolution (no compression), the intermediate layer applies moderate compression, and the distant layer uses aggressive compression. The layer durations are proportioned such that recent history occupies a smaller temporal window but retains higher token density, while distant history spans a longer period with sparse representation. This design achieves substantial compression while preserving fine-grained information where behavioral cues are most critical. Specific configuration details are provided in Sec. V-A. Learned Compression via Q-former. The compression at each level is realized through a Q-former module that takes the buffered observation tokens {obuffer } as input, along with their t associated timestep embeddings etime and camera embeddings ecam . For frames at level Lk , we concatenate ⌊Nimg /rk ⌋ learnable query tokens with the raw observation tokens and other embeddings, then process this combined sequence through self-attention in an MMDiT manner [34]. The query tokens selectively aggregate task-relevant visual features through bidirectional attention, producing the compressed representation for each frame. After compression, the tokens from all levels are reordered chronologically, from early to late frames. To accommodate the hierarchical compression structure, we employ the RoPE positional embeddings [35] from Alpamayo such that, for tokens at level Lk with compression ratio rk , the RoPE frequency step is scaled by rk . This keeps positional encodings consistent across compression levels and aligned with the uncompressed sequence. Rather than relying on hand-crafted rules, the compression is learned end-to-end. While the hierarchical buffer structure provides an inductive bias for temporal decay, the query tokens adaptively determine which visual features to retain. The resulting compressed tokens form the trajectory-conditioned memory. Next, we introduce the variational framework that enables planning-aligned compression. C. Planning-aligned Variational Token Compression While the hierarchical compression (Sec. III-B) reduces token count, it does not explicitly couple compression with the planning objective to determine what information to retain based on downstream planning needs. We introduce a conditional variational auto-encoder (cVAE) framework with vector quantization (VQ) that addresses this by coupling compression quality with trajectory prediction. The key idea is to distill driving intent from future trajectories into a compact discrete latent z ∈ Rd , then train compressed observations to
Timestep and Camera Embedding
Multi-view Long-term History Observations …
Conditional VAE with VQ KL Loss
Prior Encoder (P-net)
…
Posterior Encoder (Q-net)
…
𝑡−𝑇
𝑡
𝑡−1
Learnable Queries Historical Driving Intent Trajectory
Vector Quantize
z
Vision Encoder Q-former Compression
MLP
Raw Observation Tokens …
… Hierarchical FIFO Memory Bank
Distant Frames
Skill Codebook
Middle Frames
Timestep and Camera Embedding
Recent Frames
…
Historical Trajectory
…
Future Traj (Train only) Planning Loss
Predicted Trajectory
𝑡−1
Policy Transformer 𝑡
(Unified Backbone)
Fig. 2: Overall architecture of COMPACT-VA. Multi-view raw observation tokens from current and past are compressed via Q-former conditioning on camera embeddings, historical trajectory and a learned driving intent latent. The posterior encoder (Q-net) distills intent from future trajectories during training, while the prior encoder (P-net) predicts this latent from compressed observations alone. The hierarchical FIFO memory bank, concatenated with the predicted latent as a special token, feeds the policy transformer, with end-to-end optimization coupling compression quality to planning performance.
be sufficient for predicting this latent, ensuring retention of decision-critical historical cues (Fig. 2). Variational Encoder Architecture. We employ two encoders with distinct roles during training and inference. Both encoders use lightweight architectures to remain computationally efficient relative to the policy backbone. 1) Posterior encoder qϕ (z | o, τfuture ) (training only): Extracts driving intent from future trajectories. Following Alpamayo [1], ground-truth trajectories are first converted to unicycle control sequences (a, κ) for acceleration and curvature via least-squares optimization with Tikhonov regularization [36] to attenuate high-frequency noise, then uniformly quantized into discrete tokens. We find that this approach better captures trajectory dynamics compared to FSQ tokenization. These tokens are then compressed via an MLP into Nagg tokens including one global token and Nlocal local tokens from uniformly divided segments. A small transformer outputs mean µϕ and log-variance log σϕ2 parameterizing a Gaussian distribution N (µϕ , σϕ2 ), from which latent zq is sampled. 2) Prior encoder pθ (z | ocompressed ) (training and inference): This encoder predicts the driving intent latent using only the compressed observations ocompressed from the Q-former (Sec. III-B), without access to future information. It processes the compressed tokens through attention pooling followed by MLPs to produce zp . During training, it learns to match the posterior distribution; at inference, it operates independently to predict driving intent from historical context alone. 3) Vector Quantization: Both zq and zp are mapped to a shared discrete codebook via argmink ∥z − ck ∥, yielding quantized embedding zskill = ci∗ . Following [37], gradient flow uses the straight-through estimator zskill = z + (ci∗ − z).detach(), with commitment loss Lcommit = ∥z − ci∗ ∥2 encouraging alignment with the codebook.
Policy Input Composition. The discrete skill embedding zskill obtained from VQ, using prior encoder’s zp at training and inference, is re-projected through a learned linear layer and appended as a special token. This token is concatenated with the trajectory-conditioned memory ocompressed , the historical trajectory token (Sec. III-A), and re-applied timestep and camera embeddings. This combined sequence is fed into the unified transformer backbone following Alpamayo [1], which autoregressively predicts the future trajectory tokens. End-to-End Training. The entire system that consists of Q-former compression, prior/posterior encoders, VQ codebook, and policy backbone, is optimized end-to-end with a composite objective, L = Ltraj + λKL · DKL qϕ (z|o, τfuture ) ∥ pθ (z|ocompressed ) + λcommit · Lcommit ,
(3)
where Ltraj is the cross-entropy loss over discrete future trajectory tokens, and the KL divergence encourages the prior distribution to match the posterior. During training, the policy conditions on the latent sampled from the prior encoder zp , instead of the posterior zq , ensuring consistency between training and inference. This design creates a closed-loop coupling between compression and planning. If the Q-former discards any decision-critical historical information, the prior encoder cannot accurately predict the latent inferred from future trajectories, resulting in both high KL divergence and degraded trajectory prediction. Through this joint optimization, the model implicitly discovers what historical information matters for downstream decisions, without requiring hand-crafted retention rules. Inference. At test time, only the prior pathway is active. The model compresses observations via Q-former, predicts the latent zp from trajectory-conditioned memory, quantizes it via VQ to retrieve the discrete skill embedding, re-projects and
The oncoming car yield or maintain speed?
Does it get right-of-way?
(a) Four-Way Stop Intersection
Obstacle
(b) Dynamic Occlusion
(c) Unprotected Turns
Fig. 3: Illustration of high-signal dynamic scenarios.
observed seconds ago may begin decelerating to yield, or maintain speed requiring the ego vehicle to stop and wait for a safe gap. Without tracking oncoming vehicles’ trajectories over several seconds, the model risks either turning into oncoming traffic or waiting unnecessarily. These three scenario classes encompass the core decisionmaking challenges in driving. According to [6], fundamental skills tested here, i.e., errors in gap acceptance, right-of-way negotiation, and stopping behavior, account for approximately 40% of intersection crashes. B. Behavioral Evaluation Metrics
appends it as a special token, then autoregressively generates trajectory tokens. This preserves full compatibility with the unified VA architecture while enabling effective long-horizon planning under strict token budgets. IV. MEMORY-DEPENDENT DRIVING SCENARIOS AND EVALUATION FRAMEWORK We focus on scenarios where extended historical context plays a particularly important role in determining correct behavior. Prior studies [6] identify that critical driving decisions rely on behavioral cues captured within a 510 second temporal window, constituting extended context relative to standard driving policies that typically process only 1-2 seconds, and distinct from navigation tasks requiring long-term spatial memory over entire routes. We identify highsignal dynamic scenarios where extended context determines behavioral correctness and design metrics for these decisioncritical outcomes beyond trajectory displacement. A. Stop-Controlled Intersections as Memory Testbeds We identify three scenario classes (Fig. 3) where extended historical context plays a prominent role: (1) four-way stops requiring right-of-way negotiation based on arrival order, (2) stop / yield signs requiring assessment of dynamic crosstraffic, and (3) unprotected turns requiring gap acceptance decisions. These share a fundamental characteristic that correct behavior depends on discrete decision correctness, i.e., whether the vehicle stops when required and proceeds when appropriate, instead of trajectory smoothness. • Four-Way Stop Intersections. Multiple vehicles arrive at an intersection with all-way signs on all approaches. Right-of-way follows arrival order [6], demanding tracking which vehicles arrived earlier over several seconds. The ego vehicle needs to determine its position in this temporal queue. The model must maintain information about arrival times, 5-10 seconds ago, to correctly infer yielding order and proceed promptly when gaining right-of-way. • Stop / Yield with Dynamic Occlusion. As ego vehicle approaches the intersection, previously visible participants may become occluded or exit the field of view. The model must persist their state observed seconds earlier rather than relying solely on currently visible objects to avoid incorrectly assessing an occluded intersection as clear. • Unprotected Turns. The ego vehicle turns across oncoming traffic without a protected signal. An oncoming vehicle
Traditional metrics like minADE [38] are misaligned with these decision-critical scenarios. A rolling stop may achieve low minADE while constituting a safety violation [39], whereas stopping correctly but slightly late incurs higher minADE despite being safer. We propose behavioral metrics directly assessing decision correctness across all three scenario classes: • Stop Success Rate (Stop SR) measures whether the vehicle achieves complete stop (velocity < vstop ) when required. For sustained stopping, we check whether the predicted trajectory maintains stopped state during groundtruth stopped periods within a temporal tolerance window. • Go Success Rate (Go SR) evaluates whether the vehicle proceeds after stopping rather than remaining indefinitely stopped. This includes assessing prompt departure after gaining right-of-way, ensuring the model does not create traffic flow disruptions. • Roll-Through Rate quantifies the percentage where the vehicle fails to achieve complete stop, instead performing a rolling stop (minimum velocity ≥ vstop ). Rolling through is both illegal and dangerous [39], [6]. • Stop Position Error measures spatial deviation between stop location and the designated stop line. • Stop Duration Error quantifies the deviation from the ground-truth stop duration and that observed in human driving behavior. V. EXPERIMENTS A. Evaluation Settings Dataset and Scenarios. We evaluate on the Alpamayo physical AI dataset [40] with two complementary protocols: open-loop evaluation on curated memory-dependent scenarios, and closed-loop evaluation on general driving to verify nominal performance (Sec. V-C). To rigorously evaluate the memory performance, we curate a subset from the dataset containing the above scenarios where ground-truth trajectories exhibit: (1) deceleration to below 1 m/s within 6.4 s; (2) a stopped state (velocity < 0.5 m/s) maintained for at least 0.5 s; and (3) subsequent acceleration, indicating successful gap acceptance and departure. This yields approximately 16% of the dataset, where behavioral correctness is unambiguous and memory-dependent reasoning essential. For open-loop evaluation, from this subset, we extract a validation set of 20,000 clips (20s each, 200
TABLE I: Overall performance on stop-controlled intersection scenarios. Under comparable token budgets, COMPACT-VA achieves consistent improvements across all behavioral metrics while enabling 3.3× inference speedup and 2.7× memory reduction compared to uncompressed long-context processing (see Tab. III). Method
Settings
Obs Tokens
Go SR (↑)
Stop SR (↑)
Roll-through Rate (↓)
Stop Pos. Error (↓)
Stop Dur. Err (sec) (↓)
Standard Alpamayo [1] Sparse Obs w/ Long Hist. Dense Obs w/ Long Hist.
1s 8 imgs 5s 8 imgs 5s 40 imgs
1280 1280 6400
63.8% ±0.1% 62.0% ±0.1% 61.9% ±0.1%
86.8% ±0.2% 86.2% ±0.1% 85.8% ±0.2%
9.0% ±0.1% 9.3% ±0.0% 9.9% ±0.1%
1.21 ±0.02 1.25 ±0.00 1.26 ±0.01
0.50 ±0.01 0.53 ±0.00 0.54 ±0.00
Compression w/o plan-align
5s 40 imgs
1424
65.6% ±0.2%
87.5% ±0.2%
8.5% ±0.1%
1.15 ±0.01
0.49 ±0.01
1424 1424
68.2% ±0.3% 68.3% ±0.2%
89.2% ±0.1% 88.5% ±0.1%
7.0% ±0.2% 7.1% ±0.3%
1.15 ±0.03 1.10 ±0.02
0.48 ±0.01 0.48 ±0.01
COMPACT-VA (Disc.) COMPACT-VA (Cont.)
5s 40 imgs 5s 40 imgs
TABLE II: Closed-loop evaluation in Alpasim. Metrics avg dist between incidents at fault collision at fault dist traveled m min [email protected](gt) min [email protected](gt) min [email protected](gt) offroad offroad or collision at fault plan deviation progress wrong lane
Baseline 2s 18 imgs
Ours 5s 40 imgs
0.49 0.04 155.98 3.77 3.87 4.14 0.28 0.32 0.79 0.73 0.23
0.48 0.05 152 3.82 3.91 4.16 0.27 0.32 0.75 0.71 0.23
frames at 10Hz, critical decision point at frame 50) and use the remaining for training. Models trained on this subset specialize in memory-dependent decision-making. Implementation Details. For all experiments, we use T = 20 timesteps (5s at 4Hz), Ncam = 2 cameras, and Nimg = 160 tokens per image. Without compression, this yields Nraw = 6,400 vision tokens. Our hierarchical compression employs K = 3 layers: Layer 1 (n1 = 4 frames, r1 = 1), Layer 2 (n2 = 5 frames, r2 = 16), and Layer 3 (n3 = 11 frames, r3 = 80), compressing to Ncompressed = 1,424 tokens (4.5× reduction). The driving latent dimension is dz = 32. For the posterior encoder, trajectories are quantized into 128 discrete tokens (2 per waypoint for 64 waypoints), then compressed to Nagg = 5 tokens (Nlocal = 4 local + 1 global). The VQ codebook size is K = 20. For behavioral metrics, the stop velocity threshold is vstop = 0.5 m/s. Baseline Settings. We compare COMPACT-VA against baselines under varying history and token budgets: Standard Alpamayo retains only the recent 1s observation (8 frames, 1280 tokens), representing basic setting [1]; Sparse Obs w/ Long Hist extends context to 5s through sparse sampling (8 frames, 1280 tokens); Dense Obs w/ Long Hist maintains full 4Hz sampling over 5s without compression (40 frames, 6400 tokens); Compression w/o plan-align applies hierarchical compression without planning-aligned module (40 frames, 1424 tokens); and our COMPACT-VA (Disc./Cont.) add planning-aligned variational compression using discrete FSQbased or continuous latent encodings (1424 tokens). All methods are trained end-to-end on the same dataset.
TABLE III: Efficiency Performance in Alpasim Efficiency Metrics
Baseline 2s 18 imgs
Baseline (long) 5s 40 imgs
Ours 5s 40 imgs
Mean Time Std Median Time Peak GPU Memory
498.50 ms 5.59 ms 498.08 ms 5.94 GB
1253.52 ms 27.99 ms 1242.11 ms 10.51 GB
377.08 ms 12.85 ms 374.05 ms 3.95 GB
B. Overall Performances Table I presents the overall performance across all evaluated methods. We prioritize Go Success Rate as the primary metric, as it directly tests whether the model maintains effective memory to determine when to proceed, requiring long-horizon memory to assess cross-traffic patterns and right-of-way. In contrast, Stop Success Rate, while important, can often be achieved reactively by observing immediate deceleration trends without extended memory. The standard Alpamayo already achieves competitive performance (63.8% Go SR), demonstrating the effectiveness of its architecture. Sparse observation with long history (62.0%) underperforms, revealing that overly sparse sampling discards critical intermediate frames and actually hurts performance. Surprisingly, dense observation with full history (61.9%) performs worst despite accessing all 40 frames, suggesting that indiscriminate token abundance confuses the model, and quadratic attention cost over 6400 tokens hinders temporal reasoning. Learned hierarchical compression without planning alignment (65.6%) improves over Alpamayo, validating that structured compression retains decision-relevant information. Under comparable token budgets to standard baseline (1424 vs 1280 tokens), COMPACT-VA achieves sizable improvements of 68.3% Go SR (+4.5% absolute over Alpamayo, +2.7% over compression w/o plan-align), 22% relative reduction in Roll-through Rate (7.0% vs. 9.0%), higher Stop SR (89.2% vs. 86.8%), and lower Stop Duration Error (0.48s vs. 0.50s). These gains are achieved while maintaining the real-time constraint—crucially, our approach also delivers substantial computational efficiency over longcontext alternatives, i.e., 3.3× inference speedup and 2.7× memory reduction compared to uncompressed 5s 40imgs processing (Tab. III). This dual advantage (improved decisionmaking with enhanced efficiency) demonstrates that coupling compression with planning through the cVAE framework
TABLE IV: Ablation on architecture. All variants use 5s 40imgs with compression ratios 1-16-80. Progressive addition of components validates that planning-aligned compression with future latent as special token achieves optimal performance. Architecture Settings
Compression History Future Module Info Info × ✓ ✓ ✓
No compression Naı̈ve compression Compression w/o plan-aligned COMPACT-VA (Ours)
× × ✓ ✓
× × × ✓
TABLE V: Ablation on hierarchical compression rates. Format: (X×Y ) indicates X images with Y tokens per image. Compression Settings Layer1 Layer2 Layer3
Token Num
Go SR
Stop SR
(8×160)
(0×0)∗
(0×0)∗
1280
63.8%
86.8%
(8×80) (4×160)
(10×10) (10×10)
(22×2) (26×2)
784 792
64.6% 66.1%
86.9% 88.3%
(8×160) (8×160) (8×160) (8×160)
(10×10) (10×20) (20×10) (20×20)
(22×2) (22×2) (12×2) (12×2)
1424 1524 1504 1704
68.3% 67.5% 66.5% 67.1%
88.5% 88.9% 88.2% 88.3%
∗ (0×0) indicates this compression layer is not used.
TABLE VI: Ablation on history length. Longer history improves performance, with 5s 40imgs peaking at go success rates, due to the base model’s pretraining distribution. History Length
Token Num
Go SR
Stop SR
Overall
5s 20 imgs 5s 40 imgs 5s 60 imgs 5s 80 imgs
712 1424 2136 2848
65.6% 68.3% 66.6% 68.2%
88.2% 88.5% 88.7% 89.0%
76.9% 78.2% 77.7% 78.6%
enables retention of critical historical information while remaining computationally practical for deployment. C. Closed-loop Evaluation While our open-loop evaluation stress-tests memory-critical stop-controlled scenarios, closed-loop evaluation validates general driving competence. We conduct closed-loop simulation on 910 diverse scenarios from the Physical AI AV NuRec dataset1 in Alpasim [41], a state-of-the-art neuralreconstruction based end-to-end simulator. We focus closedloop evaluation on nominal driving scenarios, as there is currently a lack of sufficient reconstructions for stopcontrolled intersections in the Alpasim environments. We compare against the baseline (2s 18imgs), both trained on the same general driving data. Table II shows our method (5s 40imgs with compression) performs on par with the baseline across key safety metrics including, collision at fault, and average distance between incidents. Importantly, while maintaining this comparable driving performance with minimal token overhead, our approach delivers substantial computational efficiency gains. We measure inference time 1 https://huggingface.co:nvidia/PhysicalAI-Autonomous-Vehicles-NuRec
Go SR (↑)
Stop SR (↑)
Stop Pos. Error (↓)
Stop Dur. Err (sec) (↓)
61.9% ±0.1% 63.5% ±0.2% 65.6% ±0.2% 68.3% ±2.1%
85.8% ±0.2% 86.6% ±0.2% 87.5% ±0.2% 88.5% ±0.1%
1.26 ±0.01 1.21 ±0.01 1.15 ±0.01 1.10 ±0.02
0.54 ±0.00 0.51 ±0.01 0.49 ±0.01 0.48 ±0.01
Bird’s eye view of ego plan Clip-id: clipgt-clipgt-0b54c8d4-8119-431e-9ae8-5c70f5e64716
All-way Sign
Yield to this car
Fig. 4: Closed-loop example at an all-way sign controlled right turn. Our method correctly identifies the all-way sign (top right) and yields to the oncoming through vehicle that arrived first. The bird’s-eye view (left) shows the reference trajectory (yellow) and predicted trajectory (blue) which are both short, indicating the decision to stop and wait. Camera views (right) capture the ego vehicle decelerating as the crosstraffic vehicle proceeds, demonstrating proper right-of-way negotiation enabled by maintaining historical context.
and memory usage (averaged over 20 runs) on NVIDIA A100 and include a long-context baseline without compression (5s 40imgs) for reference. Compared to the short-context baseline, our method achieves 1.32× faster inference (377ms vs 499ms) with 33% lower peak GPU memory (3.95GB vs 5.94GB). The efficiency gains become more pronounced against the long-context baseline processing the same temporal extent. Our compression enables 3.3× speedup and 2.7× memory reduction. These results confirm COMPACT-VA not only maintains general driving competence while improving memory-dependent decision-making (Sec. V-B), but also achieves extended context with reasonable computational overhead. A qualitative example of closed-loop evaluation on all-way sign scenario is shown in Fig. 4. D. Ablation Study and Analysis Ablation on Architecture. Table IV validates each component using 5s 40imgs with compression ratios 1-16-80 (8 frames at 160 tokens, 10 at 10 tokens, 22 at 2 tokens). Without compression, Go SR is only 61.9% despite all 40 frames, confirming raw token abundance hinders performance. Naı̈ve compression improves to 63.5%, while Standard compression incorporating historical trajectory as Q-former conditioning achieves 65.6%, demonstrating trajectory-conditioned com-
pression retains decision-relevant features. Adding planningaligned variational coupling (Plan-aligned V1) reaches 66.9%, validating that the cVAE framework guides compression. Our final design (Plan-aligned V2) achieves 68.3% by appending the predicted future latent as a special token, suggesting that explicit separation of compressed observations and driving intent enables more effective conditioning. Ablation on Learned Skills. To validate effective skill learning, we analyze VQ codebook utilization (K = 20). At each training step, the prior encoder outputs a probability distribution over the K skills. To obtain a stable measure of skill utilization, we compute an exponentially weighted moving average (EWMA) of skill probabilities: batch psmooth (k) = α · psmooth (k), t t−1 (k) + (1 − α) · pt
(4)
where α = 0.99 is the decay factor. We define a skill k as active if its smoothed probability exceeds the uniform baseline psmooth (k) > 1/K = 0.05, ensuring we identify t skills genuinely preferred by the model over random selection. Across four random seeds over 50k training steps, the model consistently activates 15-17 out of 20 skills after the initial training phase (∼10k steps), achieving 80% codebook utilization. This high utilization confirms our framework learns diverse driving behaviors without mode collapse, successfully discretizing driving intent into meaningful latent representations. Ablation on Different Compression Rates. Table V examines hierarchical compression configurations under 5s 40imgs, where (X×Y ) denotes X images with Y tokens per image. Comparing layer 1, (8×160) (63.8% Go SR) substantially outperforms (8×80) (64.6%), confirming recent frames should retain full resolution. Allocating more images to layer 1 is beneficial where (8×160) (68.3%) outperforms (4×160) (66.1%). For layer 2, token density (10 vs 20 tokens/frame) yields comparable results, but increasing layer 2 from 10 to 20 frames degrades Go SR (rows 4 vs 6 and rows 5 vs 7). This suggests that under fixed token budgets, overextending distant history with sparse token allocation fails to capture useful information. This aligns with Table VI that extending history helps within a range but plateaus beyond the model’s effective modeling capacity. Ablation on History Length. Table VI evaluates different frame counts with fixed compression ratios (1-16-80) and layer proportions (4:5:11). While 5s 40imgs (1424 tokens) achieves the highest Go SR (68.3%), 5s 80imgs (2848 tokens) achieves competitive result (68.2%) and the best overall score (78.6%), showing that longer history generally benefits performance. The slight performance variation suggests that history length interacts with model capacity and pretraining distribution, i.e., our fine-tuning approach is constrained by the base model’s exposure to specific history lengths during pretraining, affecting optimal length within explored range. VI. CONCLUSION This work presents COMPACT-VA, a planning-aligned working memory framework for monolithic autonomous driving policies. By coupling compression with trajectory
prediction through conditional VQ-VAE, we address the fundamental limitation that unified VA/VLA policies either lack explicit memory mechanisms or rely on rule-based compression unable to guarantee retention of decision-relevant history. Through systematic evaluation on memory-critical stop-controlled scenarios, we demonstrate that planningaligned compression achieves sizable improvements in decision correctness while maintaining real-time token budgets. Our working memory approach with bounded context windows is well-aligned with driving, where critical decisions typically depend on behavioral cues within 5-10 seconds. We believe this work advances the field toward effective memory for VA/VLA policies, demonstrating task-aware working memory is key to reasoning in diverse scenarios. Future Work: While gains on general driving scenarios are modest, future work could explore higher-complexity scenarios with severe occlusions or multiple contenders where preliminary observations suggest larger gains, as well as recurrent memory mechanisms and alternative efficient architectures such as state-space models to extend planningaligned compression to broader embodied AI domains. R EFERENCES [1] Y. Wang, W. Luo, J. Bai, Y. Cao, T. Che, K. Chen, Y. Chen, J. Diamond, et al., “Alpamayo-r1: Bridging reasoning and action prediction for generalizable autonomous driving in the long tail,” arXiv, 2025. [2] Y. Hu, J. Yang, L. Chen, et al., “Planning-oriented autonomous driving,” in Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, 2023, pp. 17 853–17 862. [3] X. Weng, B. Ivanovic, et al., “Para-drive: Parallelized architecture for real-time autonomous driving,” in IEEE/CVF CVPR, 2024. [4] A. Gu and T. Dao, “Mamba: Linear-time sequence modeling with selective state spaces,” in Conference on language modeling, 2024. [5] H. Shi, B. Xie, Y. Liu, L. Sun, F. Liu, T. Wang, E. Zhou, H. Fan, X. Zhang, and G. Huang, “Memoryvla: Perceptual-cognitive memory in vision-language-action models for robotic manipulation,” arXiv, 2025. [6] M. Al-Sharman, L. Edes, et al., “Autonomous driving at unsignalized intersections: A review of decision-making challenges and reinforcement learning-based solutions,” arXiv, 2024. [7] J. Li, D. Li, et al., “Blip-2: Bootstrapping language-image pretraining with frozen image encoders and large language models,” in International conference on machine learning. PMLR, 2023. [8] C. Chi, Z. Xu, S. Feng, E. Cousineau, Y. Du, B. Burchfiel, R. Tedrake, and S. Song, “Diffusion policy: Visuomotor policy learning via action diffusion,” The International Journal of Robotics Research, vol. 44, no. 10-11, pp. 1684–1704, 2025. [9] K. Black, N. Brown, J. Darpinian, K. Dhabalia, D. Driess, A. Esmail, M. R. Equi, C. Finn, N. Fusai, M. Y. Galliker, et al., “π0.5 : a visionlanguage-action model with open-world generalization,” in 9th Annual Conference on Robot Learning, 2025. [10] Z. Liang, Y. Mu, M. Ding, F. Ni, M. Tomizuka, and P. Luo, “Adaptdiffuser: diffusion models as adaptive self-evolving planners,” in Proceedings of the 40th International Conference on Machine Learning, 2023, pp. 20 725–20 745. [11] Z. Liang, Y. Li, T. Yang, C. Wu, S. Mao, T. Nian, L. Pei, S. Zhou, X. Yang, J. Pang, et al., “Discrete diffusion vla: Bringing discrete diffusion to action decoding in vision-language-action policies,” in Proceedings of the 43rd International Conference on Machine Learning, 2026. [12] J. Bjorck, F. Castañeda, N. Cherniadev, X. Da, R. Ding, L. Fan, Y. Fang, D. Fox, F. Hu, S. Huang, et al., “Gr00t n1: An open foundation model for generalist humanoid robots,” arXiv preprint arXiv:2503.14734, 2025. [13] Z. Liang, Y. Mu, H. Ma, M. Tomizuka, M. Ding, and P. Luo, “Skilldiffuser: Interpretable hierarchical planning via skill abstractions in diffusion-based task execution,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2024, pp. 16 467–16 476.
[14] J. Zhang, A. Li, Y. Qi, M. Li, J. Liu, S. Wang, H. Liu, G. Zhou, Y. Wu, X. Li, et al., “Embodied navigation foundation model,” arXiv preprint arXiv:2509.12129, 2025. [15] J. Zhang, K. Wang, S. Wang, M. Li, H. Liu, S. Wei, Z. Wang, Z. Zhang, and H. Wang, “Uni-navid: A video-based vision-language-action model for unifying embodied navigation tasks,” Robotics: Science and Systems, 2025. [16] H. Chi, H.-a. Gao, Z. Liu, J. Liu, C. Liu, J. Li, K. Yang, Y. Yu, Z. Wang, W. Li, et al., “Impromptu vla: Open weights and open data for driving vision-language-action models,” Advances in Neural Information Processing Systems, vol. 38, 2026. [17] Y. Ma, Y. Cao, W. Ding, S. Zhang, Y. Wang, B. Ivanovic, M. Jiang, M. Pavone, and C. Xiao, “dvlm-ad: Enhance diffusion vision-languagemodel for driving via controllable reasoning,” in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2026, pp. 1050–1061. [18] S. Thrun, “Probabilistic robotics,” Communications of the ACM, vol. 45, no. 3, pp. 52–57, 2002. [19] D. Hafner, T. Lillicrap, J. Ba, and M. Norouzi, “Dream to control: Learning behaviors by latent imagination,” in ICLR, 2019. [20] L. P. Kaelbling, M. L. Littman, and A. R. Cassandra, “Planning and acting in partially observable stochastic domains,” Artificial intelligence, vol. 101, no. 1-2, pp. 99–134, 1998. [21] D. Driess, F. Xia, M. S. Sajjadi, et al., “Palm-e: an embodied multimodal language model,” in ICML, 2023. [22] J. W. Rae, A. Potapenko, S. M. Jayakumar, C. Hillier, and T. P. Lillicrap, “Compressive transformers for long-range sequence modelling,” in International Conference on Learning Representations, 2019. [23] Z. Dai, Z. Yang, Y. Yang, J. G. Carbonell, Q. Le, and R. Salakhutdinov, “Transformer-xl: Attentive language models beyond a fixed-length context,” in Proceedings of the 57th ACL, 2019, pp. 2978–2988. [24] Y. Sun, X. Li, K. Dalal, et al., “Learning to (learn at test time): Rnns with expressive hidden states,” in Forty-second International Conference on Machine Learning, 2024. [25] A. Behrouz, P. Zhong, and V. Mirrokni, “Titans: Learning to memorize at test time,” arXiv preprint arXiv:2501.00663, 2025. [26] D. Bolya, C.-Y. Fu, X. Dai, P. Zhang, C. Feichtenhofer, and J. Hoffman, “Token merging: Your vit but faster,” in The Eleventh International Conference on Learning Representations, 2023. [27] G. Xiao, Y. Tian, B. Chen, S. Han, and M. Lewis, “Efficient streaming language models with attention sinks,” in The Twelfth International Conference on Learning Representations, 2024. [28] I. Beltagy, M. E. Peters, and A. Cohan, “Longformer: The longdocument transformer,” arXiv preprint arXiv:2004.05150, 2020. [29] M. Zaheer, G. Guruganesh, et al., “Big bird: Transformers for longer sequences,” Advances in neural information processing systems, vol. 33, pp. 17 283–17 297, 2020. [30] X. Dong, Y. Fu, S. Diao, et al., “Hymba: A hybrid-head architecture for small language models,” in The Thirteenth International Conference on Learning Representations, 2025. [31] K. Ding, B. Chen, Y. Su, H.-a. Gao, B. Jin, C. Sima, X. Li, W. Zhang, P. Barsch, H. Li, et al., “Hint-ad: Holistically aligned interpretability in end-to-end autonomous driving,” in Conference on Robot Learning. PMLR, 2025, pp. 3742–3765. [32] M. Oquab, T. Darcet, T. Moutakanni, H. V. Vo, M. Szafraniec, V. Khalidov, P. Fernandez, D. HAZIZA, F. Massa, A. El-Nouby, et al., “Dinov2: Learning robust visual features without supervision,” Transactions on Machine Learning Research, 2024. [33] F. Mentzer, D. Minnen, E. Agustsson, and M. Tschannen, “Finite scalar quantization: Vq-vae made simple,” in The Twelfth International Conference on Learning Representations, 2024. [34] P. Esser, S. Kulal, A. Blattmann, et al., “Scaling rectified flow transformers for high-resolution image synthesis,” in ICML, 2024. [35] J. Su, M. Ahmed, Y. Lu, S. Pan, W. Bo, and Y. Liu, “Roformer: Enhanced transformer with rotary position embedding,” Neurocomputing, vol. 568, p. 127063, 2024. [36] G. H. Golub, P. C. Hansen, and D. P. O’Leary, “Tikhonov regularization and total least squares,” SIAM journal on matrix analysis and applications, vol. 21, no. 1, pp. 185–194, 1999. [37] A. Van Den Oord et al., “Neural discrete representation learning,” Advances in neural information processing systems, vol. 30, 2017. [38] S. Pellegrini, A. Ess, K. Schindler, and L. Van Gool, “You’ll never walk alone: Modeling social behavior for multi-target tracking,” in International conference on computer vision, 2009, pp. 261–268.
[39] H. Zhou, W. Cao, A. Sui, and Z. Bing, “What matters to enhance traffic rule compliance of imitation learning for end-to-end autonomous driving,” arXiv preprint arXiv:2309.07808, 2023. [40] NVIDIA, “Physicalai autonomous vehicles dataset,” https://huggingface. co/datasets/nvidia/PhysicalAI-Autonomous-Vehicles, 2025, one of the largest, geographically diverse datasets with 1,727+ hours of driving data (camera, LiDAR, radar) across 25 countries and 2500+ cities. [41] NVIDIA et al., “AlpaSim: A modular, lightweight, and data-driven research simulator for autonomous driving,” 2025. [Online]. Available: https://github.com/NVlabs/alpasim