Reimagination with Test-time Observation Interventions: Distractor-Robust World Model Predictions for Visual Model Predictive Control

2025-06-19

AI Summary

World models exhibit poor robustness to out-of-distribution (OOD) visual disturbances—such as novel objects or backgrounds—in open-world scenarios, leading to distorted action-effect predictions and subsequent downstream planning failures. To address this, we propose Runtime Observation Intervention (ReOI), a test-time method that detects physically implausible disturbances, dynamically edits observations in the latent space, re-executes the world model’s forward prediction, and re-injects the original disturbance to preserve visual fidelity. ReOI establishes, for the first time, a training-free, dynamic observation correction mechanism at test time, forming a closed-loop pipeline: *detect → correct → re-predict → restore*. Evaluated on robotic manipulation tasks, ReOI improves task success rates by up to 3× over strong baselines, demonstrating substantial gains in world model robustness against OOD visual interference.

Technology Category

• Intelligent Robots: State Estimation
• Planning, Routing, and Scheduling: Replanning and Plan Repair
• Computer Vision: Diffusion Models for Vision

Application Category

• Responsible Web: Machine-in-the-loop, human agency and autonomy
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems

Abstract

World models enable robots to"imagine"future observations given current observations and planned actions, and have been increasingly adopted as generalized dynamics models to facilitate robot learning. Despite their promise, these models remain brittle when encountering novel visual distractors such as objects and background elements rarely seen during training. Specifically, novel distractors can corrupt action outcome predictions, causing downstream failures when robots rely on the world model imaginations for planning or action verification. In this work, we propose Reimagination with Observation Intervention (ReOI), a simple yet effective test-time strategy that enables world models to predict more reliable action outcomes in open-world scenarios where novel and unanticipated visual distractors are inevitable. Given the current robot observation, ReOI first detects visual distractors by identifying which elements of the scene degrade in physically implausible ways during world model prediction. Then, it modifies the current observation to remove these distractors and bring the observation closer to the training distribution. Finally, ReOI"reimagines"future outcomes with the modified observation and reintroduces the distractors post-hoc to preserve visual consistency for downstream planning and verification. We validate our approach on a suite of robotic manipulation tasks in the context of action verification, where the verifier needs to select desired action plans based on predictions from a world model. Our results show that ReOI is robust to both in-distribution and out-of-distribution visual distractors. Notably, it improves task success rates by up to 3x in the presence of novel distractors, significantly outperforming action verification that relies on world model predictions without imagination interventions.

Problem

• Improving world model predictions against novel visual distractors
• Enhancing robot action verification in open-world scenarios
• Detecting and removing distractors for reliable outcome reimagination

Innovation

• Detects distractors via implausible scene degradation
• Modifies observations to remove detected distractors
• Reimagines outcomes with cleaned observations