World Model Robustness via Surprise Recognition
This addresses robustness issues for AI agents in safety-critical domains like autonomous driving, though it is incremental as it builds on existing world model techniques.
The paper tackles the problem of AI systems being destabilized by out-of-distribution noise in real-world deployments by developing an algorithm that uses a world model's surprise measure to reduce noise impact in reinforcement learning agents, showing that it preserves performance relative to baselines under varying noise types and levels in self-driving simulations like CARLA and Safety Gymnasium.
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 .