基于相似性的状态注意力机制在视觉强化学习中的应用

《Neural Networks》：Similarity-Guided State Attention for Visual Reinforcement Learning

　　**摘要**

视觉强化学习（VRL）旨在从高维观测中提取有效的视觉信息，以优化决策策略。尽管现有的VRL方法通过数据增强和辅助任务在各种控制任务中取得了显著进展，但智能体仍容易受到冗余信息和不相关因素的干扰，导致过拟合和在未见环境中的泛化能力下降。为了解决这一挑战，我们提出了一种基于相似性的状态注意力（SSA）机制。在该方法中，设计了一个相似性引导模块（SGM），利用原始观测和增强观测的状态嵌入之间的相似性来指导编码器关注原始观测中的任务相关区域，从而生成相应的状态注意力图。同时，通过解码增强观测的状态嵌入来获得其状态注意力图。此外，引入了余弦相似度来衡量原始观测和增强观测的状态注意力图之间的全局相似性，并将其与二元交叉熵损失一起纳入自监督学习目标中，以促进表示空间中状态注意力图的对齐。SGM与基于余弦相似度的状态注意力图对齐的结合，有助于自监督学习获得更健壮的状态表示，从而使智能体能够学习最优策略并在未见环境中提高其泛化能力。在DeepMind控制泛化基准（DMControl-GB）上的实验结果表明，与代表性的VRL基线相比，SSA具有更优的鲁棒性和泛化性能。

**引言**

强化学习（RL）（Xu, Chen, Wen, Fu, Wang, 2025; Yang, Hou, Zhang, Hu, Zhang, 2026）是机器学习的核心领域之一。与监督学习不同，RL不依赖于真实标签，而是基于智能体与环境交互和试错过程中获得的奖励信号来优化策略。这一特性使RL在各种决策任务中展现出巨大潜力（Hassani, Hallaji, Razavi-Far, Saif, Lin, 2026; Wang, Cheng, Wang, 2025b）。然而，随着任务复杂性和环境多样性的增加，尤其是在具有高维视觉观测输入（如图像和视频）的情况下，传统RL方法面临新的挑战。智能体需要从高维视觉数据中提取有效特征，但经常受到冗余信息和不相关因素的干扰。随着深度学习的快速发展，其在特征表示和泛化性能方面的优势为视觉强化学习（VRL）的兴起奠定了基础。VRL结合了深度学习、RL和视觉信息处理技术，使智能体能够更好地理解高维视觉环境并做出决策。具体来说，深度神经网络自动提取视觉信息，将复杂的高维输入转换为低维特征，使智能体能够学习最优策略并采取有效行动。最近在视觉任务上的进展，如用于图像增强的视觉-文本融合和用于复杂视觉分类的基于图的建模，进一步表明了从高维视觉数据中学习健壮且信息丰富的表示的重要性（Li, Sun, Ji, Huang, 2025; Wang, Zhang, Xu, Li, Ren, 2026）。这使得VRL在复杂的控制任务中表现出色，例如自动驾驶（Kiran, Sobh, Talpaert, Mannion, Sallab, Yogamani, Pérez, 2022; Zhang, Wu, Cai, Ma, Leng, Sun, Shan, 2025）、机器人控制（Cai, Wang, Qi, Jin, Wu, Xu, Zhang, 2023; Wang, Man, Cui, Lu, Cai, Fan, 2024a）、导航任务（Li, Zhang, Zhao, 2020; Rao, Wu, Yang, Zhang, Lu, Lu, Zha, 2021）、视频游戏（Shao, Tang, Zhu, Li, & Zhao, Xu, Yin, Zhang, Huang, 2022）以及无人机任务（Yan et al., 2022）。然而，在提高智能体在未见环境中的鲁棒性和泛化能力的同时，如何有效过滤掉冗余信息和不相关因素仍然是当前VRL研究中的一个重大挑战。近年来，研究人员广泛探讨了VRL领域的泛化问题，相关方法大致分为两类。第一类是基于数据增强的方法（Hansen, Su, Wang, 2021; Laskin, Lee, Stooke, Pinto, Abbeel, Srinivas, 2020a; Raileanu, Goldstein, Yarats, Kostrikov, Fergus, 2021），这些方法借鉴了计算机视觉中的数据增强技术，通过对高维视觉输入进行变换或扰动来提高数据多样性，从而提高智能体的泛化能力。它们通过直接从增强观测中学习，旨在通过更丰富的输入分布提高模型的鲁棒性。然而，单纯的数据增强不能保证智能体始终关注任务相关的视觉区域，限制了其学习关键特征表示的能力，从而影响在未见环境中的泛化性能。第二类包括基于辅助任务的方法（Hansen, Sun, Abbeel, Efros, Pinto, & Wang, Wang, Wang, Cheng, Wang, 2024b; Yuan, Xue, Yuan, Wang, Wu, Gao, Xu, 2022），这些方法通常结合表示学习从高维视觉观测中提取高级特征，并使用对比损失（Laskin, Srinivas, Abbeel, 2020b; Zhu, Xia, Wu, Deng, Zhou, Qin, Liu, Li, 2023）或重建损失（Ze, Hansen, Chen, Jain, Wang, 2023; Zhang, Lv, Zhang, Yang, Zhao, Yu, Huang, Zhou, Ye, & Jiang）来优化强化学习效率。这些方法通常使用数据增强训练编码器，将原始观测和增强观测映射到低维表示空间，然后将其输入到策略网络中进行决策。同时，它们使用辅助任务和强化学习目标共同优化模型。尽管这些方法在各种控制任务中取得了优异的性能，但它们并没有明确建模原始观测和增强观测之间的关系。因此，智能体在视觉扰动下可能难以始终关注任务相关区域，导致学习到的策略容易受到任务无关因素（如背景、颜色和位置）的影响。一些研究（Bertoin, Zouitine, Zouitine, Rachelson, 2022; Grooten, Tomilin, Vasan, Taylor, Mahmood, Fang, Pechenizkiy, Mocanu, 2024）将注意力或掩码机制纳入辅助任务中，以通过强调视觉观测中的显著区域并抑制任务无关信息来提高鲁棒性和泛化性能。这些方法在指导智能体关注重要视觉特征和提高在不同环境中的泛化性能方面显示出有效性。然而，它们依赖于启发式掩码生成或间接监督信号（如奖励或显著性），而没有明确建模原始观测和增强观测之间的内在关系。因此，学习到的注意力可能无法在视觉扰动下准确捕捉一致的任务相关区域。

为了填补这一研究空白，本文提出了一种用于视觉强化学习的基于相似性的状态注意力（SSA）机制，该机制明确建模了原始观测和增强观测之间的相似性，以指导表示学习，从而在视觉扰动下实现稳健且一致的任务相关表示识别，提高智能体在未见环境中的鲁棒性和泛化能力。首先，建立了一个基于视觉注意力机制的相似性引导模块，用于学习原始观测和增强观测的状态嵌入之间的内在相似性，从而获得原始状态的注意力图，并促进状态在下游控制任务中的更有效利用。随后，引入余弦相似度来衡量原始状态和增强状态之间的全局相似性，并将其纳入自监督学习目标中，以促进表示空间中状态注意力图的对齐，从而在视觉变化下提高鲁棒性。最后，将原始观测与相应的状态注意力图融合用于Q值估计，并引入了一致性项（Bertoin et al., 2022），以鼓励价值函数优先考虑输入中的关键像素，从而促进最优策略的学习。此外，在DeepMind控制泛化基准（DMControl-GB）（Hansen and Wang, 2021）上的综合评估显示，SSA在两个级别的视觉扰动下始终优于所有VRL基线。

**总结**

主要贡献有三点：
1) 我们提出了SSA，一种用于视觉强化学习的基于相似性的状态注意力机制，旨在提高智能体在未见环境中的鲁棒性和泛化能力。
2) 相似性引导模块（SGM）利用原始观测和增强观测的状态嵌入之间的相似性来指导编码器提取任务相关表示，从而生成原始观测的状态注意力图，突出关键区域的同时抑制不相关的干扰。
3) 引入了余弦相似度来衡量原始状态和增强状态之间的全局相似性，并将其与二元交叉熵（BCE）结合形成自监督学习目标，从而细化状态注意力图，并为下游RL提供更丰富的像素级支持。

本文的后续内容组织如下：第2节回顾了VRL的相关工作。第3节提供了所提出方法的问题定义，并介绍了VRL中视觉注意力机制和自监督学习的相关先验知识。第4节详细阐述了VRL的基于相似性的状态注意力架构。第5节展示了实验结果并进行了深入分析。最后，第6节总结了研究结果并指出了未来的研究方向。

**相关工作**

在本节中，我们从两个角度回顾了VRL的相关工作：数据增强和辅助任务。

**VRL的问题定义**

VRL旨在通过智能体与环境的持续交互从高维视觉观测中学习决策策略。在本文中，我们假设环境被建模为一个无限视界的马尔可夫决策过程（MDP），由五元组???,??,??,?,???定义，其中??表示高维状态空间，??表示动作空间，??:??×??×??→[0,1]是状态转移概率函数，??(??,??):??×??→?表示奖励函数，γ表示...

**方法论**

为了更好地训练智能体在决策过程中关注重要区域，减少对无关区域的敏感性，并解决在未见环境中的鲁棒性和泛化问题，我们在本节提出了一种用于VRL的基于相似性的状态注意力架构，称为SSA。首先，我们概述了SSA的总体架构。然后，我们介绍了SSA的三个主要组成部分：基于视觉注意力机制的相似性引导模块、自监督学习中的相似性测量...

**实验**

在本节中，我们在DMControl-GB（Hansen and Wang, 2021）的五个视觉连续控制任务上进行了实验，以评估SSA的总体鲁棒性和泛化能力。DMControl-GB源自DeepMind Control Suite（DMControl）（Tassa et al., 2018），并对背景复杂性进行了修改，使其成为更具挑战性的泛化基准。

**结论**

在本文中，我们提出了一种用于VRL的基于相似性的状态注意力机制，称为SSA。在SSA中，（i）基于视觉注意力机制的SGM用于计算原始观测和增强观测的状态嵌入之间的相似性，指导编码器提取任务相关表示，并生成原始观测的状态注意力图，有效抑制冗余信息和不相关因素；（ii）我们引入余弦相似度来衡量原始状态和增强状态之间的全局相似性，并将其与二元交叉熵（BCE）结合形成自监督学习目标，从而细化状态注意力图并为下游RL提供更丰富的像素级支持。

**作者贡献声明**

Zeng Yinfeng：撰写——原始草稿、软件、方法论、调查、概念化。
Wang Xuesong：撰写——审阅与编辑、监督、资源。
Cheng Yuhu：撰写——审阅与编辑、监督、资源、概念化。

**利益冲突声明**

作者声明他们没有已知的竞争财务利益或个人关系可能影响本文报告的工作。

**致谢**

本工作得到了中国国家自然科学基金（项目编号62373364、62176259）和江苏省重点研发计划（项目编号BE2022095）的资助。