Objects matter: object-centric world models improve reinforcement learning in visually complex environments
作者: Weipu Zhang, Adam Jelley, Trevor McInroe, Amos Storkey
分类: cs.LG, cs.CV
发布日期: 2025-01-27
💡 一句话要点
提出OC-STORM，利用对象中心世界模型提升视觉复杂环境中强化学习的样本效率。
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 强化学习 基于模型的强化学习 世界模型 对象中心表示 视觉复杂环境
📋 核心要点
- 传统基于像素的MBRL在视觉复杂环境中难以捕捉关键决策细节，导致样本效率低下。
- OC-STORM通过整合对象分割和预训练视觉模型，使智能体聚焦于与奖励和目标相关的关键对象。
- 实验表明，OC-STORM在Atari游戏和Hollow Knight游戏中优于传统MBRL方法，提升了性能。
📝 摘要（中文）
深度强化学习在从像素中学习控制策略方面取得了显著成功，但其应用仍受到样本效率低的限制，需要比人类更多的环境交互才能达到相当的性能。基于模型的强化学习（MBRL）通过利用学习到的世界模型来生成模拟经验，从而提高样本效率。然而，在视觉复杂的环境中，小而动态的元素对于决策至关重要。传统的基于像素的MBRL方法通常依赖于带有L2损失的自编码器，这种方法容易被大的区域所主导，并且常常无法捕捉到与决策相关的细节。为了解决这些限制，我们提出了一个对象中心的MBRL流程，它整合了计算机视觉的最新进展，使智能体能够专注于关键的决策相关元素。我们的方法包括四个主要步骤：（1）用分割掩码标注与奖励和目标相关的关键对象，（2）使用预训练的、冻结的基础视觉模型提取对象特征，（3）将这些对象特征与原始观测相结合来预测环境动态，以及（4）使用由该对象中心世界模型生成的想象轨迹来训练策略。在高效的MBRL算法STORM的基础上，我们称这个流程为OC-STORM。我们证明了OC-STORM在克服传统MBRL方法在Atari游戏和视觉复杂游戏Hollow Knight中的局限性方面的实际价值。
🔬 方法详解
问题定义：论文旨在解决在视觉复杂环境中，基于模型的强化学习（MBRL）方法样本效率低下的问题。现有方法，特别是基于像素的MBRL，通常使用自编码器和L2损失来学习世界模型，这使得模型容易被图像中大的、静态的区域所主导，而忽略了对于决策至关重要的小而动态的元素。
核心思路：论文的核心思路是利用对象中心（object-centric）的表示来改进世界模型的学习。通过让智能体关注场景中的关键对象，而不是整个像素空间，可以更有效地学习环境动态，并提高样本效率。这种方法借鉴了计算机视觉领域的最新进展，特别是对象分割和预训练的视觉模型。
技术框架：OC-STORM流程包含四个主要步骤：(1) 对象标注：使用分割掩码标注与奖励和目标相关的关键对象。(2) 特征提取：使用预训练的、冻结的基础视觉模型提取对象特征。(3) 动态预测：将提取的对象特征与原始观测相结合，用于预测环境动态。(4) 策略训练：使用由对象中心世界模型生成的想象轨迹来训练策略。该流程建立在高效的MBRL算法STORM之上。
关键创新：最重要的技术创新在于将对象中心的表示引入到MBRL中。与传统的基于像素的MBRL方法不同，OC-STORM显式地关注场景中的关键对象，从而能够更有效地学习环境动态。此外，使用预训练的视觉模型可以减少从头开始学习视觉表示的需求，进一步提高样本效率。
关键设计：论文使用预训练的、冻结的视觉模型（具体模型未知）来提取对象特征。对象分割掩码的获取方式（人工标注或自动分割）未知，但标注质量对最终性能至关重要。动态预测模块的具体网络结构未知，但需要能够有效地融合对象特征和原始观测。策略训练使用STORM算法，该算法的具体细节需要参考原始STORM论文。
🖼️ 关键图片
📊 实验亮点
OC-STORM在Atari游戏和视觉复杂游戏Hollow Knight中进行了评估，实验结果表明，OC-STORM能够显著优于传统的MBRL方法。具体的性能提升数据未知，但论文强调OC-STORM能够克服传统方法在视觉复杂环境中的局限性，证明了对象中心表示的有效性。
🎯 应用场景
该研究成果可应用于各种视觉复杂的强化学习任务中，例如游戏AI、机器人导航、自动驾驶等。通过关注关键对象，智能体可以更有效地学习环境动态，从而提高决策能力和样本效率。该方法在资源受限或环境交互成本高的场景下具有重要价值，例如在真实机器人上进行训练。
📄 摘要（原文）
Deep reinforcement learning has achieved remarkable success in learning control policies from pixels across a wide range of tasks, yet its application remains hindered by low sample efficiency, requiring significantly more environment interactions than humans to reach comparable performance. Model-based reinforcement learning (MBRL) offers a solution by leveraging learnt world models to generate simulated experience, thereby improving sample efficiency. However, in visually complex environments, small or dynamic elements can be critical for decision-making. Yet, traditional MBRL methods in pixel-based environments typically rely on auto-encoding with an $L_2$ loss, which is dominated by large areas and often fails to capture decision-relevant details. To address these limitations, we propose an object-centric MBRL pipeline, which integrates recent advances in computer vision to allow agents to focus on key decision-related elements. Our approach consists of four main steps: (1) annotating key objects related to rewards and goals with segmentation masks, (2) extracting object features using a pre-trained, frozen foundation vision model, (3) incorporating these object features with the raw observations to predict environmental dynamics, and (4) training the policy using imagined trajectories generated by this object-centric world model. Building on the efficient MBRL algorithm STORM, we call this pipeline OC-STORM. We demonstrate OC-STORM's practical value in overcoming the limitations of conventional MBRL approaches on both Atari games and the visually complex game Hollow Knight.