Risk-Aware World Model Predictive Control for Generalizable End-to-End Autonomous Driving
| Authors | Jiangxin Sun et al. |
| Year | 2026 |
| Field | Computer Vision |
| arXiv | 2602.23259 |
| Download | |
| Categories | cs.CV, cs.AI, cs.RO |
Abstract
With advances in imitation learning (IL) and large-scale driving datasets, end-to-end autonomous driving (E2E-AD) has made great progress recently. Currently, IL-based methods have become a mainstream paradigm: models rely on standard driving behaviors given by experts, and learn to minimize the discrepancy between their actions and expert actions. However, this objective of "only driving like the expert" suffers from limited generalization: when encountering rare or unseen long-tail scenarios outside the distribution of expert demonstrations, models tend to produce unsafe decisions in the absence of prior experience. This raises a fundamental question: Can an E2E-AD system make reliable decisions without any expert action supervision? Motivated by this, we propose a unified framework named Risk-aware World Model Predictive Control (RaWMPC) to address this generalization dilemma through robust control, without reliance on expert demonstrations. Practically, RaWMPC leverages a world model to predict the consequences of multiple candidate actions and selects low-risk actions through explicit risk evaluation. To endow the world model with the ability to predict the outcomes of risky driving behaviors, we design a risk-aware interaction strategy that systematically exposes the world model to hazardous behaviors, making catastrophic outcomes predictable and thus avoidable. Furthermore, to generate low-risk candidate actions at test time, we introduce a self-evaluation distillation method to distill riskavoidance capabilities from the well-trained world model into a generative action proposal network without any expert demonstration. Extensive experiments show that RaWMPC outperforms state-of-the-art methods in both in-distribution and out-of-distribution scenarios, while providing superior decision interpretability.
Engineering Breakdown
Plain English
This paper addresses a critical limitation in end-to-end autonomous driving systems: imitation learning models that train only on expert demonstrations fail catastrophically when encountering rare, out-of-distribution scenarios not present in training data. The authors propose a risk-aware world model predictive control framework that enables autonomous driving without relying on expert action supervision, instead using a learned world model to predict future states and make decisions that account for uncertainty and risk. By decoupling perception from action through world models and adding explicit risk awareness, the approach achieves better generalization to long-tail driving scenarios that would cause standard imitation-learning baselines to produce unsafe actions.
Core Technical Contribution
The core novelty is replacing traditional imitation learning's expert action supervision with a two-stage approach: (1) learning a world model that predicts future observations given planned actions, and (2) using model predictive control (MPC) that explicitly optimizes for both task success and risk minimization rather than imitating expert behavior. This eliminates the distribution mismatch problem where models overfit to common expert behaviors and fail on rare scenarios. The key insight is that by planning in a learned latent space and evaluating plan quality through uncertainty estimates, the system can handle out-of-distribution scenarios by naturally preferring conservative actions when prediction confidence is low, without needing explicit expert demonstrations for those cases.
How It Works
The system operates in three integrated stages. First, a world model encoder-decoder network learns to predict future image frames and state variables given current observations and a sequence of planned action commands—this is trained on unlabeled driving video data through reconstruction loss. Second, at inference time, an MPC controller generates candidate action sequences and uses the world model to rollout predictions for 4-8 timesteps forward. Third, a risk-aware cost function evaluates each predicted trajectory by combining task rewards (staying on road, reaching goal) with epistemic uncertainty penalties—trajectories passing through high-uncertainty regions are downweighted, naturally avoiding out-of-distribution areas. The controller selects the action sequence with the lowest combined cost, effectively trading off goal progress against decision reliability based on what the model is confident about.
Production Impact
Adopting this approach would fundamentally change your autonomous driving pipeline from a supervised behavior cloning architecture to an unsupervised planning-based system, eliminating the need to collect and label expert demonstrations for every scenario you want to handle. In practice, you'd deploy a learned world model (likely 50-150M parameters) that runs inference 10-20 times per decision cycle for MPC planning, adding 100-300ms of latency per decision compared to direct imitation networks—acceptable for highway driving but challenging for urban environments requiring <50ms decisions. The major production benefit is graceful degradation on edge cases: instead of producing random unsafe actions on out-of-distribution scenarios, the system conservatively slows down or stops when encountering uncertainty, improving safety metrics. However, you'd need to carefully validate that the world model's uncertainty estimates are actually well-calibrated (not overconfident on truly novel scenarios), as poorly calibrated uncertainty could create new failure modes where the system appears confident in wrong predictions.
Limitations and When Not to Use This
The paper's core assumption—that a learned world model's uncertainty estimates reliably indicate out-of-distribution scenarios—requires strong validation in real deployment, as world models can suffer from overconfidence on systematic distribution shifts (e.g., new weather conditions or camera angles not in training data). The approach requires sufficient compute for real-time MPC planning with multiple rollouts, making it impractical for edge devices or extremely latency-constrained applications, and the paper doesn't clearly detail computational requirements or how to scale this beyond highway driving scenarios. The framework still needs high-quality unlabeled driving video data at scale to learn good world models, and the paper doesn't address how to handle multimodal futures (when multiple valid next actions exist) or long-horizon planning beyond 8-10 steps where world model errors compound. Finally, evaluation appears limited to simulation or controlled datasets—real-world deployment results on truly novel long-tail scenarios are needed to validate whether this matches or exceeds human-driven safety performance.
Research Context
This work builds on the recent shift away from pure behavioral cloning toward world-model-based planning in imitation learning (following work by Dreamer, PlaNet, and Latent World Models), but applies these techniques specifically to autonomous driving where safety-critical decisions demand uncertainty awareness. It addresses known failure modes of IL-based driving systems documented in benchmarks like CARLA and nuScenes, where standard supervised approaches achieve <60% success on rare scenarios. The paper opens a research direction on integrating epistemic uncertainty estimation with planning in safety-critical domains, suggesting that learned uncertainty can replace hand-engineered safety margins in production systems. This also connects to broader trends in robotics toward world models and model predictive control as alternatives to end-to-end deep learning for generalizable behavior.
:::tip Subscribe Get weekly breakdowns of papers like this in AI Letters - the newsletter for engineers building production AI systems. :::