Make the Pertinent Salient: Task-Relevant Reconstruction for Visual Control with Distractions
作者: Kyungmin Kim, JB Lanier, Pierre Baldi, Charless Fowlkes, Roy Fox
分类: cs.LG, cs.AI, cs.CV, cs.RO
发布日期: 2024-10-13
💡 一句话要点
提出Segmentation Dreamer，通过任务相关重建提升视觉控制在干扰环境下的泛化性
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)
关键词: 视觉控制 强化学习 模型学习 表征学习 视觉干扰 分割掩码 任务相关重建
📋 核心要点
- MBRL在视觉控制中面临泛化性挑战，视觉干扰会引入高方差，影响表征学习。
- 提出Segmentation Dreamer (SD)，仅重建任务相关部分，降低表征学习复杂度。
- 实验表明，SD在干扰环境下显著提升了样本效率和最终性能，尤其在稀疏奖励任务中。
📝 摘要（中文）
本文针对基于模型的强化学习(MBRL)在视觉控制任务中面临的泛化性挑战，尤其是在存在视觉干扰的情况下，提出了一种简单而有效的辅助任务来促进表征学习。该方法基于DREAMER，利用图像观测的分割掩码，仅重建任务相关的图像部分。通过这种方式，显著降低了表征学习的复杂度，避免了对任务无关对象的编码。该方法被称为Segmentation Dreamer (SD)，可以使用模拟环境中易于获取的ground-truth掩码，也可以利用可能不完美的分割模型。对于后者，通过选择性地应用重建损失来避免因掩码预测误差而产生误导性的学习信号，从而进一步改进。在添加了视觉干扰的DeepMind Control suite (DMC)和Meta-World任务中，SD比现有方法实现了显著更高的样本效率和更好的最终性能。SD在稀疏奖励任务中尤其有效，无需大量的奖励工程即可训练出视觉鲁棒的智能体。
🔬 方法详解
问题定义：现有的基于模型的强化学习方法在视觉控制任务中，尤其是在存在视觉干扰的情况下，难以学习到具有良好泛化能力的表征。视觉干扰引入了大量的图像变化，使得智能体难以区分任务相关的特征和干扰信息，从而影响控制策略的学习。现有方法通常需要大量的训练数据或复杂的奖励工程才能解决这个问题。
核心思路：论文的核心思路是利用先验知识，通过分割掩码来区分图像中的任务相关和任务无关部分，并仅对任务相关的部分进行重建。这样可以减少表征学习的复杂度，使智能体更加关注任务相关的特征，从而提高在干扰环境下的泛化能力。
技术框架：Segmentation Dreamer (SD) 基于 DREAMER 框架。主要包含以下模块：1) 环境交互模块：与环境交互并收集数据。2) 表征学习模块：使用分割掩码对图像进行分割，并仅对任务相关部分进行重建。3) 世界模型模块：学习环境的动态模型。4) 策略优化模块：基于世界模型学习控制策略。整体流程是：智能体与环境交互，收集数据，使用分割掩码对图像进行分割，并仅对任务相关部分进行重建，学习环境的动态模型，最后基于世界模型学习控制策略。
关键创新：最重要的技术创新点在于引入了任务相关的重建损失，并结合分割掩码来指导表征学习。与现有方法相比，SD 不需要对整个图像进行重建，而是只关注任务相关的部分，从而降低了表征学习的复杂度，提高了在干扰环境下的泛化能力。此外，SD 还考虑了分割掩码可能不准确的情况，并提出了选择性地应用重建损失的方法，以避免因掩码预测误差而产生误导性的学习信号。
关键设计：SD 使用分割掩码来区分图像中的任务相关和任务无关部分。对于ground-truth掩码，可以直接使用。对于不完美的分割模型，SD 提出了选择性地应用重建损失的方法。具体来说，SD 会计算重建误差，并仅对重建误差小于一定阈值的部分应用重建损失。这个阈值可以通过交叉验证来确定。此外，SD 还使用了 KL 散度来约束隐变量的分布，以提高表征的稳定性和泛化能力。
🖼️ 关键图片
📊 实验亮点
SD在添加了视觉干扰的DeepMind Control suite (DMC)和Meta-World任务中，比现有方法实现了显著更高的样本效率和更好的最终性能。尤其在稀疏奖励任务中，SD能够成功训练出视觉鲁棒的智能体，而现有方法则无法解决。实验结果表明，SD能够有效地利用先验知识，降低表征学习的复杂度，提高在干扰环境下的泛化能力。
🎯 应用场景
该研究成果可应用于各种需要在复杂视觉环境下进行控制的任务，例如机器人操作、自动驾驶、游戏AI等。通过减少对大量数据的依赖和对奖励工程的需求，该方法可以降低训练成本，加速智能体的开发和部署。未来，该方法可以进一步扩展到更复杂的任务和环境，例如多任务学习和开放世界环境。
📄 摘要（原文）
Recent advancements in Model-Based Reinforcement Learning (MBRL) have made it a powerful tool for visual control tasks. Despite improved data efficiency, it remains challenging to train MBRL agents with generalizable perception. Training in the presence of visual distractions is particularly difficult due to the high variation they introduce to representation learning. Building on DREAMER, a popular MBRL method, we propose a simple yet effective auxiliary task to facilitate representation learning in distracting environments. Under the assumption that task-relevant components of image observations are straightforward to identify with prior knowledge in a given task, we use a segmentation mask on image observations to only reconstruct task-relevant components. In doing so, we greatly reduce the complexity of representation learning by removing the need to encode task-irrelevant objects in the latent representation. Our method, Segmentation Dreamer (SD), can be used either with ground-truth masks easily accessible in simulation or by leveraging potentially imperfect segmentation foundation models. The latter is further improved by selectively applying the reconstruction loss to avoid providing misleading learning signals due to mask prediction errors. In modified DeepMind Control suite (DMC) and Meta-World tasks with added visual distractions, SD achieves significantly better sample efficiency and greater final performance than prior work. We find that SD is especially helpful in sparse reward tasks otherwise unsolvable by prior work, enabling the training of visually robust agents without the need for extensive reward engineering.