TaskSense: Focusing on What Matters in World Models

2026-08-06

AI Summary

This work addresses the limitation of existing world models that reconstruct entire images, thereby overemphasizing irrelevant background details and diluting learning signals for task-relevant features, which leads to significant performance degradation under visual distractions. To overcome this, the authors propose TaskSense, a novel framework that introduces, for the first time, a task-oriented differentiable stochastic spatial attention mechanism into world models. TaskSense dynamically selects task-relevant image regions conditioned on the previous latent state and leverages an inverse dynamics auxiliary objective to guide attention toward control-critical content, reconstructing only the attended regions to enhance representational efficiency. Experiments demonstrate that TaskSense matches DreamerV3’s performance on the DeepMind Control Suite while substantially outperforming it on the Distracting Control Suite, exhibiting superior robustness to visual distractions; qualitative results further confirm its ability to accurately localize task-relevant regions.

Abstract

World models for visual control typically learn compact latent states by reconstructing observations, implicitly encouraging representations to preserve information across the entire visual input. However, task-relevant content often occupies only a small fraction of the observation, while background clutter and distractors consume valuable representational capacity. This mismatch between visual reconstruction and control objectives biases latent representations to model task-irrelevant visual content, diluting learning signals for control-relevant features and severely degrading downstream performance under visual distractions. We introduce TaskSense, a task-centric world modeling framework that enforces task relevance before latent encoding through a differentiable stochastic spatial attention mechanism conditioned on the previous latent state. To steer attention toward control-relevant regions, we augment training with an auxiliary inverse-dynamics objective. Rather than reconstructing the full observation, the world model reconstructs only the attended regions, encouraging latent representations to preserve task-relevant information while discarding irrelevant visual content. The decoder is further conditioned on the sampled attention map, enabling consistent reconstruction despite stochastic attention. Compared with the DreamerV3 baseline, TaskSense maintains competitive performance on the DeepMind Control Suite while consistently outperforming DreamerV3 on the Distracting Control Suite, demonstrating substantially improved robustness to visual distractions. Qualitative analysis further confirms that the learned attention, guided by inverse-dynamics supervision, consistently localizes control-relevant regions while suppressing irrelevant visual content.

Problem

• world models
• visual distractions
• task-relevant representation
• latent state
• visual control

Innovation

• task-centric world models
• stochastic spatial attention
• inverse-dynamics supervision
• visual distraction robustness
• selective reconstruction