TaskSense: Focusing on What Matters in World Models
SM Mazharul Islam, Manfred Huber
cs.AI, cs.CV, cs.LG
Submitted: 2026-08-06
Updated: 2026-08-10
License: http://creativecommons.org/licenses/by/4.0/
The gist: World models for visual control typically learn compact latent states by reconstructing observations, implicitly encouraging representations to preserve information across the entire visual input.
Terminology
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.
Sources
- Learning Transformer-based World Models with Contrastive Predictive Coding
- Provable RL with Exogenous Distractors via Multistep Inverse Dynamics
- World Models
- Dream to Control: Learning Behaviors by Latent Imagination
- Mastering Atari with Discrete World Models
- Temporal Difference Learning for Model Predictive Control
- Categorical Reparameterization with Gumbel-Softmax
- Guaranteed Discovery of Control-Endogenous Latent States with Multi-Step Inverse Models
- The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables
- Transformers are Sample-Efficient World Models
- R2-Dreamer: Redundancy-Reduced World Models without Decoders or Augmentation
- SOLD: Slot Object-Centric Latent Dynamics Models for Relational Manipulation Learning from Pixels
- The 2017 DAVIS Challenge on Video Object Segmentation
- Loss is its own Reward: Self-Supervision for Reinforcement Learning
- The Distracting Control Suite -- A Challenging Benchmark for Reinforcement Learning from Pixels
- Learning Latent Dynamic Robust Representations for World Models
- DeepMind Control Suite
- Rethinking Full Connectivity in Recurrent Neural Networks
- Learning Invariant Representations for Reinforcement Learning without Reconstruction
Related papers
- MAVEN-T: Reinforced Heterogeneous Distillation for Real-Time Multi-Agent Trajectory Prediction
- Model Discovery Agent: LLM-assisted Bayesian experiment design for data-efficient discovery of mechanistic world models
- Agents' Last Exam
- The Clinician's Veto: Navigating Trust, Liability, and Uncertainty in Autonomous AI Prescribing
- MindHelper: Closed-Loop Embodied Mental-State Reasoning for Precision Intervention
- Incumbent Advantage: Brand Bias and Cognitive Manipulation Dynamics in LLM Recommendation Systems