arXiv:2608.06544cs.AIcs.CV2026-08

让世界模型专注关键视觉信息,提升复杂场景下的控制能力。

TaskSense: Focusing on What Matters in World Models

论文配图:TaskSense: Focusing on What Matters in World Models
图 1 · 摘自论文原文
  • 用注意力机制提前筛选任务相关区域,避免无关内容干扰
  • 在干扰环境中性能优于基线模型,尤其在复杂视觉背景下表现更稳
  • 适合需要抗干扰的强化学习与机器人视觉控制任务

视觉控制的世界模型通常通过重建观测来学习紧凑的隐状态,但这一过程会保留整个视觉输入的信息,导致大量表征容量被背景杂讯和干扰物占用。任务相关信息往往仅占观测的一小部分,这种不匹配使隐状态偏向建模无关视觉内容,削弱了控制相关特征的学习信号。为此,我们提出 TaskSense,一种以任务为中心的世界建模框架:通过基于前一隐状态的可微分随机空间注意力机制,在编码前就聚焦任务相关区域,并引入辅助逆动力学目标引导注意力。世界模型仅重构注意力区域,迫使隐状态保留任务关键信息并丢弃无关内容。解码器还受采样注意力图条件控制,确保重建一致性。相比 DreamerV3 基线,TaskSense 在 DeepMind Control Suite 上保持竞争力,且在 Distracting Control Suite 中持续领先,显著提升了对视觉干扰的鲁棒性。定性分析表明,注意力机制在逆动力学监督下能稳定定位控制相关区域,抑制无关内容。

原文摘要 · Abstract (English)

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.

世界模型注意力机制强化学习视觉控制

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。