Technical Summary: TaskSense

Problem Statement

World models for visual control typically rely on reconstructing full observations to learn compact latent states. While effective in standardized benchmarks with static backgrounds, this approach implicitly encourages latent representations to preserve information across the entire visual input, regardless of its relevance to the control task. In environments where task-relevant content occupies only a small fraction of the observation or where background clutter and distractors are present, this reconstruction objective consumes valuable representational capacity on task-irrelevant visual content. This mismatch between visual reconstruction and control objectives biases the latent dynamics, diluting learning signals for control-relevant features and severely degrading downstream performance under visual distractions. Conversely, relying solely on task supervision (e.g., reward prediction) often provides too weak a signal to learn rich visual representations. Existing solutions that eliminate reconstruction entirely sacrifice the generative capabilities needed for imagination-based planning, while methods that apply attention or masking after feature extraction fail to control which information enters the latent dynamics model in the first place.

Methodology

The authors propose TaskSense, a task-centric world modeling framework that enforces task relevance before latent encoding. TaskSense augments the DreamerV3 architecture with a differentiable stochastic spatial attention mechanism conditioned on the previous latent belief state.

Core Components

1. Latent-Conditioned Stochastic Spatial Attention:

   • Instead of predicting attention for every pixel, the model defines attention over a coarse lattice $\Lambda$ covering the observation.
   • A lightweight MLP predicts attention logits based on the previous latent state $s_{t-1}$ (specifically the hidden state $h_{t-1}$ and stochastic latent $z_{t-1}$), allowing the model to decide where to focus based on its evolving understanding of the task rather than instantaneous visual appearance.
   • The attention field is modeled as a collection of independent stochastic variables using a Binary Concrete (Relaxed Bernoulli) distribution. This allows multiple regions to be attended to or ignored simultaneously, avoiding the competitive constraints of softmax attention.
   • During training, differentiable masks are sampled; during inference, the deterministic mean is used.

2. Attention-Aware Encoding:

   • The sampled attention mask is upsampled and used to modulate the observation via element-wise multiplication, attenuating low-attention regions.
   • To resolve ambiguity between low-intensity pixels and suppressed high-intensity pixels, the upsampled attention map is concatenated as an additional input channel to the encoder.

3. Auxiliary Inverse-Dynamics Objective:

   • To prevent the sparsity regularizer from collapsing the attention to a trivial solution (ignoring all observations), an auxiliary objective is introduced.
   • The model predicts the action $a_{t-1}$ executed between two consecutive time steps using the temporal difference of encoded features ($f_t - f_{t-1}$).
   • This objective grounds the attention mechanism in control-relevant regions by encouraging the retention of visual features that are predictive of the agent's behavior.

4. Attention-Conditioned Reconstruction:

   • Rather than reconstructing the full observation, the decoder reconstructs only the attended regions ($o^{att}_t$).
   • The decoder is conditioned on the sampled attention map (via stop-gradient) to ensure it knows which regions were retained, preventing the latent state from being penalized for missing information that was intentionally filtered out.

Training Objective

The total loss function combines the standard DreamerV3 latent dynamics losses ($L_{dyn}$) with the new components:
$$L_{WM} = L_{dyn} + L_{rec} + \lambda_{sparse}L_{sparse} + \lambda_{inv}L_{inv}$$
Where $L_{rec}$ is the reconstruction loss over filtered observations, $L_{sparse}$ encourages compact attention maps, and $L_{inv}$ is the inverse-dynamics loss.

Key Contributions

1. TaskSense Framework: A reconstruction-based world modeling framework that enforces task relevance before latent encoding through latent-conditioned stochastic spatial attention.
2. Self-Supervised Learning Mechanism: A novel approach that grounds stochastic attention in control-relevant regions by combining inverse-dynamics supervision with reconstruction over filtered observations.
3. Empirical Validation: Demonstration that TaskSense maintains competitive performance on standard benchmarks while significantly outperforming baselines under visual distractions.

Experimental Results

The authors evaluated TaskSense on the DeepMind Control Suite (DMC) and the Distracting Control Suite (DCS), which features dynamic natural-image backgrounds.

• Standard Performance: On the standard DMC (static backgrounds), TaskSense matches the performance of the DreamerV3 baseline across most tasks, demonstrating that reconstructing only a subset of the observation is sufficient for effective latent planning.
• Robustness to Distractions: On the Distracting Control Suite, TaskSense consistently outperforms DreamerV3. For example, on the Walker Run task, TaskSense achieved a mean return of 531.2 compared to DreamerV3's 286.2. This confirms that filtering irrelevant visual content before encoding improves robustness.
• Qualitative Analysis: Visualizations of the learned attention maps on the Cheetah Run task show that TaskSense consistently focuses on control-relevant regions (e.g., the cheetah's legs) while suppressing the torso and distracting backgrounds.
• Ablation Studies:
  • Removing the inverse-dynamics objective leads to diffuse attention that fails to localize control-relevant regions, causing significant performance degradation.
  • Replacing stochastic attention with deterministic attention results in a significant performance drop, highlighting the importance of exposing the model to diverse partial observations during training.
  • Removing attention-conditioned reconstruction causes a smaller but noticeable degradation, indicating the necessity of aligning the reconstruction target with the filtered input.

Significance and Claims

The paper claims that TaskSense addresses a fundamental mismatch in current world models: the tendency to preserve all visual information rather than focusing on what matters for control. By shifting the focus to observation-level filtering rather than post-hoc representation refinement, TaskSense allows reconstruction-based world models to allocate representational capacity efficiently toward control-relevant information.

The authors assert that this approach enables world models to maintain the generative benefits of reconstruction (essential for imagination-based planning) while achieving robustness against visual distractions that typically degrade standard models. The work suggests that enforcing task relevance at the observation interface is a promising direction for improving visual world models and representation learning in model-based reinforcement learning.

The paper acknowledges a limitation: it assumes that visual features predictive of immediate actions (captured by inverse dynamics) are synonymous with those relevant for long-term task success. In scenarios where these diverge, the current supervision signal may focus attention on sub-optimal features, a direction identified for future work.