TaskSense：专注于世界模型中的关键要素

# TaskSense：在世界模型中聚焦关键信息
来源：https://arxiv.org/html/2608.06544

###### 摘要

用于视觉控制的世界模型通常通过重建观测来学习紧凑的潜在状态，这隐式地鼓励表征保留整个视觉输入中的信息。然而，与任务相关的内容往往只占观测的一小部分，而背景杂乱和干扰物则消耗了宝贵的表征容量。这种视觉重建与控制目标之间的不匹配，使得潜在表征偏向于对与任务无关的视觉内容进行建模，从而稀释了与控制相关特征的学习信号，并在视觉干扰下严重降低下游性能。我们提出了 TaskSense，一个以任务为中心的世界建模框架，通过一个由前一个潜在状态条件化的可微随机空间注意力机制，在潜在编码之前强制执行任务相关性。为了将注意力引导到与控制相关的区域，我们引入了一个辅助逆动力学目标来增强训练。世界模型不重建完整观测，而仅重建被注意的区域，从而鼓励潜在表征保留与任务相关的信息，同时丢弃无关的视觉内容。解码器进一步以采样得到的注意力图为条件，从而在随机注意力的条件下仍能实现一致的重建。与 DreamerV3 基线相比，TaskSense 在 DeepMind Control Suite 上保持了具有竞争力的性能，同时在 Distracting Control Suite 上持续优于 DreamerV3，展现出对视觉干扰显著增强的鲁棒性。定性分析进一步证实，由逆动力学监督引导的所学注意力能够一致地定位与控制相关的区域，同时抑制无关的视觉内容。

## 引言

从原始视觉观测中学习紧凑的潜在动力学模型（Hafner et al. 2019b (https://arxiv.org/html/2608.06544#bib.bib32)）已成为基于模型的强化学习（MBRL）（Sutton 1991 (https://arxiv.org/html/2608.06544#bib.bib29)）的一种强大范式。通过将高维观测压缩为可预测的潜在状态，世界模型（Ha and Schmidhuber 2018 (https://arxiv.org/html/2608.06544#bib.bib30)）使得规划与策略优化可以完全在学得的潜在空间内进行。许多成功的世界模型使用强大的解码器来重建观测，使重建成为学习潜在动力学的主要训练信号（Micheli et al. 2022 (https://arxiv.org/html/2608.06544#bib.bib9)；Zhang et al. 2023 (https://arxiv.org/html/2608.06544#bib.bib10)；Hafner et al. 2023 (https://arxiv.org/html/2608.06544#bib.bib31)，2025 (https://arxiv.org/html/2608.06544#bib.bib34)）。

然而，决策通常只依赖于观测的一小部分，而常见的重建目标却奖励保留所有视觉信息。因此，优化重建完整观测会隐式地鼓励潜在表征保留整个视觉场景中的信息，无论其是否与控制相关。这种不匹配在标准化基准（Tassa et al. 2018 (https://arxiv.org/html/2608.06544#bib.bib36)；Bellemare et al. 2013 (https://arxiv.org/html/2608.06544#bib.bib37)）中常常被掩盖，因为这些基准的背景是静态的，且与任务相关的物体占据场景的主要部分。然而，当观测中充满干扰物，或与任务相关的物体只占图像的一小部分时，保留整个视觉场景变得愈发低效，导致下游性能显著下降（Stone et al. 2021 (https://arxiv.org/html/2608.06544#bib.bib38)；Morihira et al. 2026 (https://arxiv.org/html/2608.06544#bib.bib39)）。在另一个极端，仅依赖奖励预测等任务监督只能提供弱得多的学习信号，往往无法学习到足够丰富的视觉表征（Yarats et al. 2021 (https://arxiv.org/html/2608.06544#bib.bib13)）。因此，无论是重建每一个像素，还是仅从奖励中学习，都不是表征学习的理想目标。

一条研究路线完全摒弃了像素重建，转而通过对比目标（Burchi and Timofte 2025 (https://arxiv.org/html/2608.06544#bib.bib33)）、原型学习（Deng et al. 2022 (https://arxiv.org/html/2608.06544#bib.bib40)）或预测一致性（Morihira et al. 2026 (https://arxiv.org/html/2608.06544#bib.bib39)）来学习表征。尽管这些方法减少了编码无关外观的压力，但移除解码器也牺牲了基于想象规划所必需的生成式世界模型。

另一条研究路线通过以对象为中心的建模（Mosbach et al. 2024 (https://arxiv.org/html/2608.06544#bib.bib28)）、特征掩码（Seo et al. 2023 (https://arxiv.org/html/2608.06544#bib.bib23)）或基于掩码重建的注意力机制（Sun et al. 2024 (https://arxiv.org/html/2608.06544#bib.bib27)）来改进表征学习。虽然这些方法鼓励更有信息量的潜在表征，但大多数是在视觉特征已被编码器提取之后才起作用。因此，它们主要学习的是如何编码信息，而不是决定在第一时间应当让哪些信息进入潜在动力学模型。

参见标题图 1：Dreamer 与 TaskSense 的比较。虚线圆圈表示预测量。左图：Dreamer 将完整观测 o_t 编码为 z_t，并重建完整观测。右图：TaskSense 用一个随机注意力模块（黑色）增强 Dreamer（灰色组件），该模块从先前的信念状态 h_{t-1}（SG）预测注意力图 m_t。注意力图过滤 o_t 以产生 o_t^att，再由原始 Dreamer 编码器进行编码。重建仅应用于 o_t^att，同时一个辅助逆动力学目标预测先前的动作 a_{t-1}。SG 表示梯度停止。一个密切相关的方向是任务信息抽象（TIA）（Fu et al. 2021 (https://arxiv.org/html/2608.06544#bib.bib16)），它通过联合训练两个世界模型进行协作重建，同时使其中一个模型与奖励信号对抗性分离，从而鼓励与任务相关的表征。TIA 虽然有效，但显式地学习一个单独的世界模型来捕获与任务无关的信息，增加了模型复杂度和训练开销。这引出了一个基本问题：与其在编码之后学习分离相关与无关信息，世界模型能否直接控制哪些视觉信息进入其潜在动力学模型？

我们提出 TaskSense，一个以任务为中心的世界建模框架，通过潜在状态条件化的随机空间注意力，在潜在编码之前强制执行任务相关性来解决这一问题。TaskSense 从前一个潜在信念状态预测一个随机注意力图，并在编码之前使用该注意力图过滤观测，从而使模型能够基于其对任务不断发展的理解来决定关注何处，而不是依赖于瞬时的视觉外观。通过过滤输入而非修改潜在动力学模型，TaskSense 阻止了与任务无关的视觉信息消耗表征容量，同时保留了想象与规划所需的基于重建的世界模型。为了避免抑制全部观测等退化解，TaskSense 将辅助逆动力学目标与注意力条件化的重建相结合：前者鼓励注意力聚焦于与控制相关的区域，后者在无需重建无关内容的情况下保留被注意的信息。

我们的主要贡献如下：

1. 1. 我们提出了 TaskSense，一种基于重建的世界建模框架，通过潜在状态条件化的随机空间注意力，在潜在编码之前强制执行任务相关性。
2. 2. 我们开发了一种自监督学习框架，通过将逆动力学监督与对过滤后观测的重建相结合，使随机注意力扎根于与控制相关的区域。
3. 3. 我们证明 TaskSense 在标准 DeepMind Control Suite 任务上与 DreamerV3 持平，同时在视觉干扰下持续优于 DreamerV3。定性可视化和消融研究进一步验证了逆动力学监督对于学习以任务为中心的注意力至关重要。

## 相关工作

### 世界模型

学习潜在世界模型已成为从高维视觉观测进行 MBRL 的主导范式。诸如 World Models（Ha and Schmidhuber 2018 (https://arxiv.org/html/2608.06544#bib.bib30)）和 PlaNet（Hafner et al. 2019b (https://arxiv.org/html/2608.06544#bib.bib32)）等早期工作表明，紧凑的潜在动力学模型使得可以直接在学得的表征空间中进行规划。PlaNet 引入了循环状态空间模型（RSSM），后来被 Dreamer 系列（Hafner et al. 2019a (https://arxiv.org/html/2608.06544#bib.bib2)，2020 (https://arxiv.org/html/2608.06544#bib.bib8)，2023 (https://arxiv.org/html/2608.06544#bib.bib31)，2025 (https://arxiv.org/html/2608.06544#bib.bib34)）采用并扩展，用于通过潜在想象优化策略和值函数，同时保留视觉重建作为主要训练信号。最近，基于 Transformer 的世界模型，如 IRIS（Micheli et al. 2022 (https://arxiv.org/html/2608.06544#bib.bib9)）和 STORM（Zhang et al. 2023 (https://arxiv.org/html/2608.06544#bib.bib10)），进一步通过自回归序列建模改进长时程预测。

尽管架构各异，这些方法都共同依赖视觉重建来学习潜在表征，这鼓励潜在状态保留观测中的所有信息，而不管其与控制的相关性如何。虽然在标准视觉控制基准中非常有效，但这种由重建驱动的目标也可能鼓励潜在表征编码与任务无关的视觉内容，从而降低在视觉干扰下的鲁棒性（Stone et al. 2021 (https://arxiv.org/html/2608.06544#bib.bib38)）。TaskSense 保留了基于重建的世界模型的生成式优势，同时在潜在编码之前引入了任务相关性。

### 表征学习

除了对比学习和原型学习之外，多项工作还探索了改进潜在表征的替代目标。掩码重建和双模拟目标分别鼓励捕获视觉结构和行为等价性的表征（Sun et al. 2024 (https://arxiv.org/html/2608.06544#bib.bib27)；Zhang et al. 2020 (https://arxiv.org/html/2608.06544#bib.bib7)）。预测一致性和冗余减少通过鼓励稳定的潜在表征并降低对数据增强的依赖来提高鲁棒性（Morihira et al. 2026 (https://arxiv.org/html/2608.06544#bib.bib39)）。

与表征目标正交的方法，如 MuZero（Schrittwieser et al. 2020 (https://arxiv.org/html/2608.06544#bib.bib6)）和 TD-MPC（Hansen et al. 2022 (https://arxiv.org/html/2608.06544#bib.bib5)），学习针对奖励和值预测而非视觉重建进行优化的潜在动力学。

尽管这些方法提高了鲁棒性和表征质量，但许多方法减少了对像素重建的依赖，牺牲了支持潜在想象和规划的生成式观测模型（Yarats et al. 2021 (https://arxiv.org/html/2608.06544#bib.bib13)）。相比之下，TaskSense 保留重建作为主要学习信号，同时仅重建与任务相关的视觉信息。

### 学习与任务相关的视觉表征

注意力机制通过选择性地强调有信息的区域或特征来改进视觉表征学习。以对象为中心的方法将场景显式分解为对象级表征，从而能够在个体实体而非整体图像特征上进行推理（Locatello et al. 2020 (https://arxiv.org/html/2608.06544#bib.bib26)；Kipf et al. 2021 (https://arxiv.org/html/2608.06544#bib.bib22)）。最近的工作如 SOLD（Mosbach et al. 2024 (https://arxiv.org/html/2608.06544#bib.bib28)）表明，结构化的以对象为中心的表征可以进一步改进潜在世界建模和控制。然而，这些方法通常从视觉观测或学得的特征中推断重要性，而非从任务目标出发。

任务驱动的方法则引入行为监督来学习与控制相关的表征。任务信息抽象（TIA）（Fu et al. 2021 (https://arxiv.org/html/2608.06544#bib.bib16)）通过联合优化重建并同时将与任务无关的信息从奖励相关特征中分离出来，来学习与任务相关的视觉抽象。逆动力学等辅助目标也被用于鼓励表征捕获环境中可控的方面（Shelhamer et al. 2016 (https://arxiv.org/html/2608.06544#bib.bib24)；Pathak et al. 2017 (https://arxiv.org/html/2608.06544#bib.bib25)；Efroni et al. 2021 (https://arxiv.org/html/2608.06544#bib.bib20)；Lamb et al. 2022 (https://arxiv.org/html/2608.06544#bib.bib21)）。通过从状态转移中预测动作，逆动力学目标保留了有助于控制的信息，同时忽略了不可控的变化（Agrawal et al. 2015 (https://arxiv.org/html/2608.06544#bib.bib15)）。然而，这些方法主要是在视觉信息已被编码之后对表征进行正则化，而不是控制哪些观测进入潜在模型。

TaskSense 将任务驱动的监督与选择性视觉处理相结合，利用逆动力学监督使随机空间注意力扎根于与控制相关的区域。与先前从视觉特征推断重要性的注意力机制不同，TaskSense 从潜在信念状态预测观测级别的注意力，并在编码之前应用它，从而使基于重建的世界模型能够选择性地表征与控制相关的信息。

## 方法

TaskSense 用一个学得的随机空间注意力模块增强 DreamerV3（Hafner et al. 2023 (https://arxiv.org/html/2608.06544#bib.bib31)）世界模型，该模块在潜在编码之前控制信息流。注意力模块产生一个观测掩码，在视觉输入被编码器处理之前对其过滤，而解码器则以对应的注意力图为条件重建过滤后的观测。为防止退化的注意力解，我们引入了一个辅助逆动力学目标，鼓励注意力保留与控制相关的视觉信息。图 1 (https://arxiv.org/html/2608.06544#Sx1.F1) 提供了所提出架构的概览。

在描述所提出的组件之前，我们简要回顾与我们的修改相关的 DreamerV3 的方面。DreamerV3 采用循环状态空间模型（RSSM）（Hafner et al. 2019b (https://arxiv.org/html/2608.06544#bib.bib32)），其中具有块对角循环权重的 GRU（Cho et al. 2014 (https://arxiv.org/html/2608.06544#bib.bib11)）（Van Keirsbilck et al. 2019 (https://arxiv.org/html/2608.06544#bib.bib19)）维护一个确定性隐藏状态 h_t，而随机潜在变量通过学得的先验分布 ẑ_t 和后验分布 z_t 来表示。在训练期间，后验分布从当前观测 o_t 和确定性状态 h_t 中推断得到，解码……