Generalizable Visual Reinforcement Learning with Segment Anything Model

📄 arXiv: 2312.17116v1 📥 PDF

作者: Ziyu Wang, Yanjie Ze, Yifei Sun, Zhecheng Yuan, Huazhe Xu

分类: cs.LG, cs.CV, cs.RO

发布日期: 2023-12-28

备注: Project page and code: https://yanjieze.com/SAM-G/

💡 一句话要点

提出SAM-G框架，利用SAM提升视觉强化学习在未知环境中的泛化能力

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 视觉强化学习 泛化能力 Segment Anything Model DINOv2 零样本学习 机器人控制 图像分割

📋 核心要点

  1. 现有视觉强化学习方法在泛化性方面存在不足，未能充分利用现代视觉基础模型。
  2. SAM-G框架利用SAM的可提示分割能力，为RL智能体提供高质量的掩码图像，提升泛化性。
  3. 实验结果表明，SAM-G在DMControl和Adroit任务上显著提升了视觉泛化能力，优于现有方法。

📝 摘要（中文）

本文提出了一种名为SAM-G的框架，旨在提升视觉强化学习（RL）智能体在未见环境中的泛化能力。现有方法主要集中于通过辅助监督、预训练或数据增强来获取鲁棒的视觉表征，而忽略了现代视觉基础模型的潜力。SAM-G框架利用Segment Anything Model (SAM)的可提示分割能力，通过DINOv2和SAM提取的图像特征寻找对应关系作为SAM的点提示，进而生成高质量的掩码图像供智能体直接使用。在8个DMControl任务和3个Adroit任务上的评估表明，SAM-G显著提高了视觉泛化能力，且无需改变RL智能体的架构，仅需改变其观测。值得注意的是，在DMControl和Adroit上具有挑战性的视频hard设置中，SAM-G分别实现了44%和29%的相对改进，优于最先进的方法。

🔬 方法详解

问题定义：视觉强化学习中的一个核心挑战是训练能够泛化到未见环境的策略。现有的方法，如通过辅助监督、预训练或数据增强来学习鲁棒的视觉表征，虽然取得了一定的进展，但仍然存在泛化能力不足的问题，并且没有充分利用近年来涌现的强大的视觉基础模型，例如Segment Anything Model (SAM)。

核心思路：SAM-G的核心思路是利用SAM强大的零样本分割能力，将原始的视觉输入转化为更具语义信息的掩码图像，从而帮助RL智能体学习到更鲁棒和泛化的策略。通过将图像分割成不同的对象，智能体可以更好地理解环境，并做出更明智的决策。

技术框架：SAM-G框架主要包含以下几个阶段：1) 使用DINOv2提取图像特征；2) 利用DINOv2和SAM的特征寻找对应关系，生成SAM的点提示；3) 使用SAM根据点提示生成高质量的掩码图像；4) 将掩码图像作为RL智能体的输入，训练或执行策略。整个框架无需修改RL智能体的架构，只需要改变其观测方式。

关键创新：SAM-G的关键创新在于将SAM引入到视觉强化学习中，并利用其可提示分割能力来提升泛化性。与以往的方法不同，SAM-G不依赖于特定的任务或环境进行训练，而是利用SAM的零样本能力来处理未见环境。此外，SAM-G通过DINOv2特征来引导SAM的分割，提高了分割的准确性和效率。

关键设计：SAM-G的关键设计包括：1) 使用DINOv2作为视觉特征提取器，提供丰富的语义信息；2) 设计了一种基于特征对应关系的提示生成方法，能够有效地引导SAM进行分割；3) 将SAM生成的掩码图像直接作为RL智能体的输入，简化了训练流程。具体的参数设置和网络结构细节在论文中有详细描述，例如DINOv2的特定层选择，以及SAM的提示策略。

📊 实验亮点

SAM-G在DMControl和Adroit任务上取得了显著的性能提升。在DMControl的video hard设置中，SAM-G实现了44%的相对改进，在Adroit的video hard设置中，实现了29%的相对改进，均优于当前最先进的方法。这些结果表明，SAM-G能够有效地提升视觉强化学习的泛化能力，使其在未见环境中表现更好。

🎯 应用场景

SAM-G框架具有广泛的应用前景，可以应用于机器人控制、自动驾驶、游戏AI等领域。通过提升视觉强化学习的泛化能力，SAM-G可以帮助智能体更好地适应复杂和动态的环境，从而实现更智能和自主的决策。未来，SAM-G可以进一步扩展到其他类型的视觉任务，例如目标检测、图像分类等，为人工智能的发展做出更大的贡献。

📄 摘要（原文）

Learning policies that can generalize to unseen environments is a fundamental challenge in visual reinforcement learning (RL). While most current methods focus on acquiring robust visual representations through auxiliary supervision, pre-training, or data augmentation, the potential of modern vision foundation models remains underleveraged. In this work, we introduce Segment Anything Model for Generalizable visual RL (SAM-G), a novel framework that leverages the promptable segmentation ability of Segment Anything Model (SAM) to enhance the generalization capabilities of visual RL agents. We utilize image features from DINOv2 and SAM to find correspondence as point prompts to SAM, and then SAM produces high-quality masked images for agents directly. Evaluated across 8 DMControl tasks and 3 Adroit tasks, SAM-G significantly improves the visual generalization ability without altering the RL agents' architecture but merely their observations. Notably, SAM-G achieves 44% and 29% relative improvements on the challenging video hard setting on DMControl and Adroit respectively, compared to state-of-the-art methods. Video and code: https://yanjieze.com/SAM-G/