Starwm improves visual world models by focusing on relevant details
StarWM: Self-Supervised Trained Attention Routing for Robust World Models
Computer Vision and Pattern RecognitionMachine Learning
Summary
World models try to understand and predict environments by looking at what they see. The authors found that methods that recreate the whole scene often waste effort on unimportant parts, while those that avoid reconstruction risk losing useful info. They developed StarWM, which learns to focus reconstruction only on important moving parts of the scene, ignoring distractions. This helps the model better predict what really matters over time, even with changing backgrounds.
What this means in practice
- •For robotics engineers: Build robots that can better understand and predict their environment despite visual distractions using focused attention in their world models.
- •For video game developers: Improve AI agents that learn and plan in complex games by ignoring irrelevant visual details during training.
Authors
Zeqiang Zhang, Fabian Wurzberger, Maximilian Otte, Daniel Schmid, Sebastian Gottwald, Arne Peter Raulf, Daniel Alexander Braun
Abstract
A robust world model must strike the balance between faithfully capturing environmental dynamics and abstracting away from irrelevant content. While reconstruction-based world models ensure faithful supervision, they misallocate representational capacity by pixel area rather than dynamics relevance for visual tasks, which can cause task-irrelevant content to dominate the learned representation. Alternatively, reconstruction-free methods avoid this bias but risk discarding possibly relevant information. We propose StarWM, which uses a cross-attention module trained on self-supervised dynamics to decide where reconstruction applies. A dual-stream decoder then restricts reconstruction to the attended regions, with stop-gradient barriers preventing interference between the two objectives. These components allows reconstruction to supervise the visual content of attended regions without contaminating the latent with non-predictive information. On DeepMind Control with dynamic video backgrounds, default (reward-free) StarWM achieves the strongest performance under random-frame distractors and substantially outperforms reconstruction-based baselines under sequential video. In addition, its reward-augmented variant matches or exceeds reconstruction-free methods on sequential video, achieving the highest overall return across all distractor regimes. Mechanistic probing confirms StarWM preserves state attributes with near-perfect fidelity through long-horizon imagination while systematically discarding distractors.