💻 computer science

TPS-Drive: Task-Guided Representation Purification for VLM-based Autonomous Driving

TPS-Drive introduces a novel framework for VLM-based autonomous driving that employs task-guided representation purification via an agent-centric tokenizer to eliminate spatial redundancy and hallucinations, enabling a decoupled reasoning pipeline that achieves state-of-the-art safety and forecasting performance on both open-loop and closed-loop benchmarks.

Original authors: Jiaxiang Li, Yumao Liu, Ke Ma

Published 2026-05-27

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jiaxiang Li, Yumao Liu, Ke Ma

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are teaching a very smart, well-read robot to drive a car. This robot is a Vision-Language Model (VLM): it's great at reading maps, understanding traffic signs, and talking about what it sees. However, there's a big problem: it's terrible at understanding the exact 3D geometry of the road. It's like a brilliant philosopher who can describe a storm in beautiful poetry but can't tell you exactly where a raindrop will hit the windshield.

The paper introduces TPS-Drive, a new system designed to fix this by teaching the robot to "think" in a cleaner, more focused way.

Here is how it works, broken down with simple analogies:

The Problem: Two Bad Ways to Describe the Road

The authors say current methods for teaching robots to drive fall into two traps:

The "Text Translator" Trap: Some systems try to turn the 3D world into simple words or numbers (e.g., "car ahead," "box at x,y,z").
- The Analogy: Imagine trying to describe a complex sculpture by only using a list of words like "round," "tall," "red." You lose the shape, the distance, and the smooth curves. The robot gets confused and starts "hallucinating" (imagining things that aren't there or placing them in the wrong spots).
The "Overloaded Camera" Trap: Other systems feed the robot raw, high-definition video or dense grids of the whole scene.
- The Analogy: Imagine giving the robot a 4K video feed of a busy street, but the video is 90% static background (like a brick wall, the sky, or a parked car that hasn't moved in years). The robot's brain gets overwhelmed by all this "noise." It spends all its energy processing the boring wall and misses the pedestrian stepping off the curb. This is called "representation interference."

The Solution: The "Task-Guided Purification" Filter

TPS-Drive solves this by introducing a special filter called the Agent-Centric Tokenizer.

The Metaphor: Think of the robot's brain as a library. Usually, the library is flooded with books about everything: the weather, the color of the pavement, the trees, and the cars. The robot can't find the important books.
The Fix: TPS-Drive installs a "Librarian" (a frozen 3D detection head) who knows exactly what the robot needs to know right now. This Librarian scans the scene and throws away 99% of the books (the static background, the sky, the textures).
The Result: The robot is left with a "Purified Space" containing only the critical, moving actors: the cars, pedestrians, and cyclists. It can now "think" clearly because it's only looking at what matters.

How the Robot Drives (The Three-Step Process)

Once the robot has this clean, purified view of the world, it drives using a three-step reasoning pipeline:

Scene Understanding: The robot looks at the purified scene and writes a structured summary. It doesn't just say "I see a car"; it understands the physics: "There is a car 10 meters ahead, moving at 5 mph, and I need to brake."
Future Forecasting: The robot predicts what will happen next. It asks, "If I keep going, where will that car be in 2 seconds?" Because it's only looking at the moving agents (thanks to the purification filter), this prediction is much more accurate.
Action Generation: Finally, the robot decides what to do. It uses a "diffusion planner" (a tool that generates smooth, safe paths) to draw a line on the road showing exactly where the car should go to avoid a crash.

The Training: From Student to Expert

The authors didn't just train the robot once; they used a three-stage training process:

Pretraining: Teaching the "Librarian" (the tokenizer) how to filter out the noise.
Supervised Fine-Tuning: Teaching the robot to read the filtered scene and predict the future, just like a student studying for a test.
Reward-Driven Refinement: This is the most important step. The robot practices driving in a simulator. If it drives safely and follows the rules, it gets a "reward." If it crashes or drives recklessly, it gets a penalty. The robot learns to prioritize safety over just copying human drivers.

The Results: Safer Driving

The paper claims that TPS-Drive works significantly better than previous methods:

Fewer Crashes: In open-loop tests (where the robot plans a path but doesn't actually drive), it had the lowest collision rates compared to other top models.
Better Predictions: It is much better at guessing where other cars and people will be in the future.
Safety Records: In closed-loop tests (where the robot actually drives in a simulated environment), it set new safety records, avoiding crashes and following traffic rules (like stopping at red lights) better than its competitors.

The Bottom Line

TPS-Drive is like giving a smart robot a pair of "smart glasses." Instead of seeing a blurry, noisy mess of the whole world, these glasses automatically blur out the boring background and highlight only the moving cars and people. This allows the robot to focus its brainpower on the things that actually matter, leading to safer and more accurate driving decisions.

Technical Summary: TPS-Drive

Problem Statement

Vision-Language Models (VLMs) offer strong semantic reasoning capabilities for autonomous driving but struggle to bridge the gap between high-level reasoning and precise 3D spatial forecasting. Existing representation strategies face a critical trade-off:

Text-Aligned Methods: These flatten continuous spatial states into discrete symbols (e.g., natural language descriptions or flattened coordinates). While compatible with standard tokenizers, this discretization compromises the underlying geometric structure of physical space, inducing "spatial hallucinations" and limiting the model's ability to capture continuous agent states.
Dense Visual Methods: These preserve spatial topology (e.g., dense grids or occupancy maps) but overwhelm standard tokenizers with redundant background textures and static scenes. Driven by reconstruction-guided Vector Quantization (VQ), these methods disproportionately allocate limited codebook capacity to static backgrounds rather than critical dynamic agents, causing "representation interference" that hinders semantic reasoning.

The core challenge is designing a representation that bridges semantic reasoning with precise forecasting without succumbing to geometric distortion or information redundancy.

Methodology: TPS-Drive

TPS-Drive introduces a framework centered on Task-Guided Representation Purification, enabling VLMs to "Think in Purified Space." The architecture decouples the driving process into three sequential stages: Scene Understanding, Future Forecasting, and Action Generation.

1. Agent-Centric Tokenizer (Representation Purification)

The core innovation is an Agent-Centric Tokenizer that constructs a purified Bird's-Eye-View (BEV) vocabulary.

Task-Guided Vector Quantization: Unlike standard VQ-VAEs that optimize for global reconstruction, this tokenizer is supervised by a frozen 3D detection head (CenterPoint). This supervision explicitly reallocates limited codebook capacity from pervasive static backgrounds to critical dynamic agents.
Primary and Residual Layers: The tokenizer employs a primary codebook ( $V_1$ ) to capture core structural patterns of agents, supervised by detection losses (heatmap and 3D bounding box). To recover precise positional nuances lost in quantization, $L-1$ residual codebooks are added in a hierarchical manner to encode residual errors.
Outcome: This process isolates spatial redundancy, yielding a distilled token sequence focused exclusively on dynamic agents, effectively preventing representation interference.

2. Decoupled Reasoning Pipeline

The framework processes inputs through three distinct stages:

Scene Understanding: A vision-language backbone extracts a unified structured representation ( $s_t$ ) encompassing scene descriptions and physical rules (e.g., braking capabilities). This is generated via a predefined vocabulary and supervised by annotations from a large language model (Qwen3.5-27B).
Future Forecasting: Conditioned on the current context and $s_t$ , the model predicts future world tokens ( $z_{t+\Delta t}$ ) in the discrete purified space. It uses an autoregressive approach for primary tokens and parallel classification heads for residual tokens.
Action Generation: A conditional diffusion planner generates continuous ego trajectories ( $\tau$ ) based on the predicted future states and structured scene context.

3. Progressive Three-Stage Training

The model is optimized via a specific training paradigm:

Tokenizer Pretraining: The primary encoder, codebook, and decoder are trained with task-guided supervision, followed by training of residual codebooks.
Supervised Fine-Tuning (SFT): The vision-language backbone is fine-tuned to jointly master scene understanding, future forecasting, and action generation.
Reward-Driven Refinement: To surpass pure imitation learning, the policy undergoes end-to-end grouped relative optimization. Stochastic rollouts are scored by an offline reward based on geometric errors and motion smoothness, refining the world modeling and diffusion planner to align with safety-oriented objectives.

Key Contributions

Agent-Centric Tokenizer: A novel tokenizer supervised by a frozen 3D detection head that achieves task-guided representation purification, explicitly isolating spatial redundancy to empower VLMs to reason in a purified space.
Decoupled Framework (TPS-Drive): A system that separates scene understanding, future forecasting, and action generation, optimized via a three-stage training paradigm featuring reward-driven refinement.
Performance Records: The framework establishes new safety records on rigorous benchmarks while maintaining high spatial forecasting accuracy.

Experimental Results

Extensive experiments were conducted on nuScenes, NAVSIMv1, and NAVSIMv2.

Open-Loop Planning (nuScenes): TPS-Drive achieved a record-low Collision Rate (CR) of 0.10% (with ego status) under the ST-P3 protocol, significantly outperforming leading VLM-based methods like OmniDrive (0.94%) and OccWorld (0.32%). It also demonstrated superior agent-centric spatial forecasting, achieving 34.60% NDS and 24.03% mAP, surpassing world models like WoTE and BEVWorld.
Closed-Loop Planning (NAVSIMv1 & v2):
- On NAVSIMv1, TPS-Drive achieved a Predictive Driver Model Score (PDMS) of 89.7, outperforming Recogdrive (89.6) and AutoVLA (89.1).
- On the more rigorous NAVSIMv2, it attained an Extended PDMS (EPDMS) of 86.7, setting a new safety record and outperforming expert-level planners like HydraMDP++ (85.6) and concurrent VLA methods.
Ablation Studies: The results confirmed that task-guided tokenization provided a +12.72% absolute increase in forecasting accuracy (NDS) compared to standard reconstruction-guided tokenization. Reward-driven refinement further improved the EPDMS from 84.1 to 86.7 and reduced the open-loop collision rate to 0.19%.

Significance and Claims

The paper claims that TPS-Drive successfully addresses the fundamental limitations of existing VLM-based driving systems by decoupling semantic reasoning from continuous control through representation purification. By shifting from global reconstruction to task-guided purification, the framework effectively bridges the gap between high-level semantic reasoning and safe, precise trajectory generation.

The authors emphasize that isolating safety-critical semantics from raw visual data is foundational for effective reasoning in Vision-Language-Action (VLA) models. The results validate that purified geometric representations can establish new safety standards in both open-loop and closed-loop interactive environments.

Limitations: The authors acknowledge that relying on a frozen detection head bounds the representation potential and that the decoupled multi-stage architecture currently restricts real-time inference speed. Future work is proposed to explore end-to-end differentiable spatial purification and hardware optimization.