World Model Robustness via Surprise Recognition

TL;DR

This work tackles the fragility of world-model-based reinforcement learning under real-world sensor noise and OOD conditions. It proposes Surprise-Guided Rejection Sampling, using the world model's Bayesian surprise to dynamically downweight or drop unreliable sensor inputs, with both multi-sensor and single-sensor variants and a representation-dropout training regime. Empirical results in CARLA and Safety Gymnasium show improved robustness and stable performance across diverse perturbations, and the method also enhances stability for both DreamerV3 and Cosmos world models, with an efficient $O(n\log n)$ inference mechanism. The approach offers architecture-agnostic, practical robustness to sensor failures and provides publicly available code for replication and deployment.

Abstract

AI systems deployed in the real world must contend with distractions and out-of-distribution (OOD) noise that can destabilize their policies and lead to unsafe behavior. While robust training can reduce sensitivity to some forms of noise, it is infeasible to anticipate all possible OOD conditions. To mitigate this issue, we develop an algorithm that leverages a world model's inherent measure of surprise to reduce the impact of noise in world model--based reinforcement learning agents. We introduce both multi-representation and single-representation rejection sampling, enabling robustness to settings with multiple faulty sensors or a single faulty sensor. While the introduction of noise typically degrades agent performance, we show that our techniques preserve performance relative to baselines under varying types and levels of noise across multiple environments within self-driving simulation domains (CARLA and Safety Gymnasium). Furthermore, we demonstrate that our methods enhance the stability of two state-of-the-art world models with markedly different underlying architectures: Cosmos and DreamerV3. Together, these results highlight the robustness of our approach across world modeling domains. We release our code at https://github.com/Bluefin-Tuna/WISER .

Figures (19)

• Figure 1: The change in the world model's measure of surprise as noise increases in all tested CARLA Dosovitskiy17 and Safety-Gymnasium Dosovitskiy17 environments over 15000 sampled steps per intensity. We leverage this insight to understand the degree to which surprise signal can help identify noise given the absence of ground truth.
• Figure 2: Multi-Sensor Rejection Sampling. In practice, we utilize Algorithm \ref{['alg:nonadaptive_masking']} to simulate this behavior in the agent.
• Figure 3: Performance (top) and Score-Cost ratio (bottom) across three Safety Gymnasium tasks. We display the Score-Cost ratio to measure how safely the task can be accomplished with respect to the score. Each column corresponds to a noise type, with clean being the nominal setting.
• Figure 4: Agent performance as the number of sensor failures (from 0-5 sensors) during the stop sign task. We observe that as the sensors begin to become affected, we are able to reduce the effect that noise has on the main policy. For all 72 tested settings over the Stop Sign, Right Turn, and Four Lane Driving Tasks, see Appendix \ref{['fig:All_Coarse_Carla']}.
• Figure 5: Rejection sampling process for noise classification and context state machine for the world model's latent state.