Risk-Aware World Model Predictive Control for Generalizable End-to-End Autonomous Driving

This paper proposes Risk-aware World Model Predictive Control (RaWMPC), a unified framework that enhances the generalization and safety of end-to-end autonomous driving by leveraging a risk-aware world model and self-evaluation distillation to make reliable decisions in unseen scenarios without relying on expert demonstrations.

The Gist

Imagine you are teaching a child how to drive a car.

The Old Way (Imitation Learning):
Most current self-driving cars are taught like a student who only watches a master driver. The computer says, "Copy exactly what the expert does." If the expert turns left when it's sunny, the car learns to turn left.

• The Problem: What happens when it starts raining, or a deer jumps out, or the road looks nothing like the videos the car studied? The car panics. It has never seen a "deer" or "rain" in its training data, so it doesn't know how to react. It's like a student who memorized the answers to a math test but fails when the teacher changes the numbers.

The New Way (RaWMPC):
This paper introduces a new system called RaWMPC. Instead of just copying a teacher, this system learns by imagining the future.

Think of RaWMPC as a cautious chess player or a daydreaming driver. Before it actually moves the car, it runs a "mental simulation" in its head.

How It Works (The 3-Step Magic)

1. The "Crystal Ball" (The World Model)
The car builds a mental model of the world. It's like having a crystal ball that can show you what happens if you do different things.

• Scenario: You are approaching a red light.
• The Simulation: The car imagines three futures:
  • Future A: "If I speed up, I might crash into the car ahead." (The crystal ball shows a crash).
  • Future B: "If I swerve left, I might hit a pedestrian." (The crystal ball shows a collision).
  • Future C: "If I slow down and stop, everything is safe." (The crystal ball shows a smooth stop).

2. The "Risk Hunter" (Risk-Aware Interaction)
Here is the genius part. Most cars are afraid to make mistakes during training. They only practice safe driving.
RaWMPC is different. It deliberately practices making mistakes in a safe, virtual environment (a video game simulator).

• It intentionally tries to drive off the road, hit imaginary walls, or run red lights in the simulation.
• Why? Because it needs to learn what a "crash" looks like so it can recognize it in real life. It's like a firefighter practicing with a fire hose so they aren't scared when a real fire starts. By "touching the hot stove" in the simulation, it learns to avoid the real stove later.

3. The "Self-Teacher" (Self-Evaluation Distillation)
Once the "Crystal Ball" is smart enough to predict crashes, the system teaches a smaller, faster version of itself how to make good choices quickly.

• The big brain says, "Don't do that, it's dangerous. Do this instead."
• The small brain learns to pick the safe option without needing to run the full simulation every single second. It's like a student who, after studying hard, can instantly answer the question without needing to re-derive the whole formula.

Why Is This Better?

• No "Expert" Needed: You don't need a human to drive perfectly for the car to learn. The car learns by exploring and seeing what happens when it fails.
• Handles the Unknown: Because it understands risk (what causes a crash) rather than just memorizing moves, it can handle weird situations it has never seen before. If a cow walks onto the road, it doesn't freeze; it calculates the risk and slows down.
• Explainable: You can ask the car, "Why did you stop?" and it can say, "Because I imagined that if I kept going, I would hit that pedestrian." It's not a black box; it's a cautious planner.

The Bottom Line

Current self-driving cars are like parrots (they repeat what they heard).
RaWMPC is like a wise old driver (it thinks ahead, remembers what bad outcomes look like, and chooses the safest path).

This paper proves that by letting the car "dream" about crashes and learn from them, we can build self-driving cars that are safer, smarter, and don't need a human teacher to show them every single possible scenario.

Technical Summary

1. Problem Statement

Current End-to-End Autonomous Driving (E2E-AD) systems primarily rely on Imitation Learning (IL), where models learn to replicate expert driving behaviors from large-scale datasets. While effective in common scenarios, this approach suffers from a fundamental generalization dilemma:

• Limited Coverage: Expert demonstrations cannot cover all "long-tail" or rare scenarios (e.g., extreme weather, complex interactions).
• Safety Risks: When encountering unseen situations, IL models often produce unsafe decisions because they lack prior experience with hazardous outcomes.
• Lack of Risk Awareness: Most existing methods, including Model-Based Reinforcement Learning (MBRL), focus on maximizing expected rewards or mimicking experts, lacking explicit mechanisms to systematically discover, model, and avoid rare but catastrophic outcomes.

The paper asks: Can an E2E-AD system make reliable decisions without any expert action supervision?

2. Methodology: RaWMPC

The authors propose RaWMPC (Risk-aware World Model Predictive Control), a unified framework that replaces expert supervision with risk-aware predictive control. The system operates without requiring expert action labels for policy learning.

Core Components:

1. Risk-Aware World Model:

   • Instead of predicting a single trajectory, the world model predicts the near-future states ($\hat{s}_{t+1:t+H}$) for a set of candidate action sequences.
   • It utilizes a Semantic-Guided Decoding mechanism. A segmentation decoder predicts future semantic maps (e.g., pedestrians, vehicles), and this semantic attention is fused into an Event Decoder to predict specific traffic events (collisions, off-road, traffic sign violations).
   • The model recursively rolls out states over a planning horizon $H$.

2. Predictive Control & Action Selection:

   • For each candidate action, the system computes a cost function $C(\cdot)$ that balances progress (moving toward the target) and risk (probability of collisions or rule violations).
   • The system selects the action sequence that minimizes this cost: $A^* = \arg\min_n C(\hat{s}^n_{t+1:t+H})$.
   • This provides interpretability, as the decision is based on explicitly comparing the predicted consequences of multiple options.

3. Risk-Aware Interactive Training:

   • To teach the world model to recognize danger without expert labels, the authors introduce a two-stage training strategy:
     • Offline Warm-up: A small amount of logged data initializes the world model's ability to predict state transitions.
     • Online Simulator Interaction: The model explores the environment using a risk-aware sampling strategy. It deliberately samples:
       • Good modes: Low-cost (safe) actions.
       • Bad modes: High-cost (risky/hazardous) actions.
       • Rand modes: Random exploration.
   • By exposing the model to "bad" rollouts (hazards), it learns to predict catastrophic outcomes, making them avoidable.

4. Self-Evaluation Distillation:

   • To ensure efficient inference at test time (avoiding expensive online optimization for every step), the authors distill the world model's evaluation capability into a lightweight Generative Action Proposal Network.
   • Using Contrastive Learning (InfoNCE), the proposal network is trained to generate action sequences that the world model rates as "positive" (low risk) while pushing away "negative" (high risk) sequences.
   • This allows the system to generate high-quality candidate actions without needing expert demonstrations.

3. Key Contributions

1. Zero-Expert Framework: RaWMPC is the first E2E-AD framework that achieves state-of-the-art performance without relying on expert action supervision for policy learning. It learns risk awareness purely through interaction.
2. Risk-Aware Interaction Strategy: A novel training method that intentionally exposes the world model to hazardous behaviors, enabling it to predict and avoid rare, catastrophic events that are typically absent in standard datasets.
3. Self-Evaluation Distillation: A mechanism to transfer the world model's risk-evaluation capabilities into a fast, generative policy, enabling real-time deployment while maintaining high safety standards.
4. Enhanced Generalization: By focusing on minimizing predicted risk rather than mimicking experts, the system demonstrates superior robustness in out-of-distribution (OOD) scenarios.

4. Experimental Results

The method was evaluated on two major benchmarks: Bench2Drive (CARLA simulator) and NAVSIM (real-world data).

• State-of-the-Art Performance:
  • Bench2Drive: RaWMPC achieved a Driving Score (DS) of 88.31 and Success Rate (SR) of 70.48%, outperforming all existing IL and RL methods (e.g., HiP-AD, SimLingo).
  • NAVSIM: Achieved a PDMS of 91.3, the highest among learning-based methods.
• Robustness to Domain Shift:
  • In a "Sunny-only training $\to$ Rainy testing" scenario, RaWMPC significantly outperformed IL baselines. While IL methods suffered severe performance drops due to perception failures in rain, RaWMPC maintained high safety margins by explicitly evaluating risk under uncertainty.
• Ablation Studies:
  • Removing the risk-aware sampling (using only random or $\epsilon$-greedy) caused significant performance degradation, proving the necessity of learning from "bad" rollouts.
  • Removing predictive control (directly outputting actions) collapsed performance, confirming that evaluating multiple candidates is crucial.
  • Self-evaluation distillation outperformed training the proposal network directly on expert actions, suggesting that risk-minimization is a better objective than imitation.

5. Significance

• Paradigm Shift: RaWMPC challenges the dominant "imitation is king" paradigm in E2E-AD. It demonstrates that a system can learn to drive safely and generally by understanding the consequences of actions rather than copying human behavior.
• Safety & Reliability: By explicitly modeling and avoiding high-risk outcomes, the system is better suited for the "long-tail" problems of real-world driving where expert data is scarce or non-existent.
• Interpretability: Unlike black-box neural networks that output a single steering angle, RaWMPC's decision process is transparent: it selects the action with the lowest predicted risk based on a forward-looking simulation.
• Reduced Data Dependency: The ability to train effectively with minimal or no expert demonstrations reduces the cost and complexity of deploying autonomous driving systems, as it does not require massive, perfectly labeled datasets of expert driving.