World Model Robustness via Surprise Recognition

TL;DR

This paper presents a novel approach that uses surprise recognition in world models to improve the robustness of reinforcement learning agents against out-of-distribution noise and sensor faults, maintaining performance across diverse environments.

Contribution

The authors introduce surprise-based rejection sampling techniques for world models, enhancing robustness to sensor faults and noise in reinforcement learning agents.

Findings

Techniques preserve performance under various noise conditions.

Methods improve stability of different world model architectures.

Approach tested successfully in self-driving simulation environments.

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 .