Addressing Corner Cases in Autonomous Driving: A World Model-based Approach with Mixture of Experts and LLMs

TL;DR

This work tackles the brittleness of motion forecasting in autonomous driving when facing corner cases by introducing WM-MoE, a world-model–based framework that unifies perception, memory, and decision making. It fuses a lightweight language model via a temporal tokenizer with a Mixture-of-Experts decoder to allocate specialized capacity to rare, high-risk interactions while preserving performance on common scenarios. The approach is validated across four real-world datasets and a new corner-case benchmark, nuScenes-corner, showing state-of-the-art accuracy and robustness to data imbalance and missing inputs. The results suggest that integrating world-model reasoning with structured expert specialization and language-informed priors yields safer, more generalizable motion forecasting suitable for fully autonomous driving systems.

Abstract

Accurate and reliable motion forecasting is essential for the safe deployment of autonomous vehicles (AVs), particularly in rare but safety-critical scenarios known as corner cases. Existing models often underperform in these situations due to an over-representation of common scenes in training data and limited generalization capabilities. To address this limitation, we present WM-MoE, the first world model-based motion forecasting framework that unifies perception, temporal memory, and decision making to address the challenges of high-risk corner-case scenarios. The model constructs a compact scene representation that explains current observations, anticipates future dynamics, and evaluates the outcomes of potential actions. To enhance long-horizon reasoning, we leverage large language models (LLMs) and introduce a lightweight temporal tokenizer that maps agent trajectories and contextual cues into the LLM's feature space without additional training, enriching temporal context and commonsense priors. Furthermore, a mixture-of-experts (MoE) is introduced to decompose complex corner cases into subproblems and allocate capacity across scenario types, and a router assigns scenes to specialized experts that infer agent intent and perform counterfactual rollouts. In addition, we introduce nuScenes-corner, a new benchmark that comprises four real-world corner-case scenarios for rigorous evaluation. Extensive experiments on four benchmark datasets (nuScenes, NGSIM, HighD, and MoCAD) showcase that WM-MoE consistently outperforms state-of-the-art (SOTA) baselines and remains robust under corner-case and data-missing conditions, indicating the promise of world model-based architectures for robust and generalizable motion forecasting in fully AVs.

Figures (10)

• Figure 1: Skewed time to collision (TTC) distributions and key highlights of our proposed model. (a) and (b) illustrate the skewed distributions of TTC in the nuScenes and Argoverse 1 datasets, with a small percentage of high-risk scenarios (TTC $<$ 2s). (c) highlights examples of safety-critical situations in the nuScenes dataset, like turns and roundabouts, while other agents are around. (d) presents the key highlight of our work: the integration of a MoE model and LLMs with the proposed temporal tokenizer to enhance performance in challenging scenarios while maintaining accuracy in simpler ones.
• Figure 2: Illustration of the world model guiding driving decisions. Left: The world model maintains a compact latent scene state via three interacting modules: Perception, Memory, and Decision. Right: Imagine a corner-case scenario involving a partially occluded crosswalk with pedestrians entering. The human brain's internal world model navigates this unseen scene through a sequence of operations: (a) First, the Perception module abstracts the complex external world into a simplified, task-relevant that localizes the target agent, nearby pedestrians, and road geometry; (b) Next, the Memory module contextualizes this sketch, integrating recent history with commonsense physical and social rules (e.g., "yield to pedestrians"), enabling counterfactual reasoning that contrasts a safe yield with an unsafe acceleration. (c) Finally, the Decision module evaluates these simulated futures and selects the optimal action, resulting in a safe and rational maneuver.
• Figure 3: Overall pipeline of WM-MoE. WM-MoE follows a world model-based pipeline with three key parts (a): Perception Module, Memory Module, and Decision Module. The Perception Module encodes traffic scenarios by capturing both spatial and temporal dynamics, resulting in a scene encoding $\bm{S}_{\text{enc}} \in \mathbb{R}^{T \times B \times D_{\text{emb}}}$. This encoding is fed into the Memory Module (b), which includes the Intention-Aware Encoder and the Language-Enhanced Encoder, producing multimodal queries $\bm{q}_{\text{mode}}$ and language-informed features $\bm{t}_{\text{llm}}$, respectively. In the Decision Module (c), a Cross-Modal Fusion employs cross-attention to integrate $\bm{q}_{\text{mode}}$, $\bm{t}_{\text{llm}}$, and the BEV feature $\bm{v}_{\text{enc}}$ (extracted via a CNN backbone), yielding the enhanced cross-modal feature $\bm{f}_{\text{c}}$. Finally, the MoE Decoder utilizes the cross-modal feature to generate the candidate multimodal future trajectories.
• Figure 4: Overall architecture of the MoE decoder. It comprises ($N$) stacked MoE blocks, each with a gating layer that routes inputs to ($K$) experts. The decoder aggregates contributions from multiple experts rather than relying on a single expert for any scenario.
• Figure 5: Joint spatial distributions of corner-case scenes in the nuScenes-corner dataset. The plots visualize the probability density (z-axis) of (a) Turning, (b) Congested, (c) Braking, and (d) Acceleration events over the map's x-y coordinates. The high-density clusters (z-axis) reveal that these critical events are not uniformly distributed but are highly concentrated in specific geographic areas, suggesting a strong correlation between road topology and the occurrence of corner cases.