Pre-trained Text-to-Image Diffusion Models Are Versatile Representation Learners for Control

📄 arXiv: 2405.05852v1 📥 PDF

作者: Gunshi Gupta, Karmesh Yadav, Yarin Gal, Dhruv Batra, Zsolt Kira, Cong Lu, Tim G. J. Rudner

分类: cs.CV, cs.AI, cs.CL, cs.LG, cs.RO, stat.ML

发布日期: 2024-05-09

💡 一句话要点

提出基于预训练文本到图像扩散模型的稳定控制表征，提升具身智能体的控制能力。

🎯 匹配领域: 支柱一：机器人控制 (Robot Control) 支柱二：RL算法与架构 (RL & Architecture) 支柱三：空间感知与语义 (Perception & Semantics) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 具身智能 控制策略 扩散模型 表征学习 视觉-语言模型

📋 核心要点

  1. 现有对比学习视觉-语言模型在具身智能体细粒度场景理解方面存在不足，限制了控制策略的泛化能力。
  2. 利用预训练文本到图像扩散模型，提取包含细粒度视觉空间信息的稳定控制表征，提升控制策略的性能。
  3. 实验证明，该方法在模拟控制任务和开放词汇导航基准上均取得了优异的性能，优于现有方法。

📝 摘要（中文）

具身智能体需要对物理世界进行细粒度的理解，这依赖于视觉和语言输入。仅从特定任务数据中学习这些能力非常困难。因此，预训练的视觉-语言模型被广泛应用于将从互联网规模数据中学到的表征迁移到下游任务和新领域。然而，常用的对比学习表征（如CLIP）在使具身智能体获得足够细粒度的场景理解方面表现不佳，而这对于控制至关重要。为了解决这个问题，我们考虑来自预训练文本到图像扩散模型的表征，这些模型经过显式优化以从文本提示生成图像，因此包含反映高度细粒度视觉空间信息的文本条件表征。我们利用预训练的文本到图像扩散模型构建了稳定控制表征，从而可以学习推广到复杂、开放环境的下游控制策略。实验表明，使用稳定控制表征学习的策略在各种模拟控制设置（包括具有挑战性的操作和导航任务）中，与最先进的表征学习方法相比具有竞争力。最值得注意的是，稳定控制表征能够学习在OVMM（一个困难的开放词汇导航基准）上表现出最先进性能的策略。

🔬 方法详解

问题定义：现有基于对比学习的视觉-语言模型（如CLIP）在具身智能体控制任务中，无法提供足够细粒度的场景理解能力，导致控制策略难以泛化到复杂、开放的环境中。这些模型学习到的表征可能不够精细，无法捕捉到对于控制至关重要的视觉空间信息。

核心思路：利用预训练的文本到图像扩散模型，因为这些模型被显式地训练成从文本提示生成图像，所以其内部的表征包含了高度细粒度的视觉空间信息。通过提取这些表征，可以为具身智能体提供更丰富的场景理解，从而提升控制策略的性能和泛化能力。

技术框架：该方法的核心是构建“稳定控制表征”。首先，使用预训练的文本到图像扩散模型（例如Stable Diffusion）提取文本条件下的图像表征。然后，将这些表征作为下游控制策略的输入，训练控制策略以执行特定的任务。整体流程包括：文本提示输入 -> 扩散模型表征提取 -> 控制策略训练 -> 任务执行。

关键创新：关键创新在于利用了文本到图像扩散模型中蕴含的细粒度视觉空间信息，并将其作为控制策略的表征。与传统的对比学习方法相比，扩散模型通过生成图像的过程，学习到了更丰富的场景结构和细节，从而能够提供更有效的控制信号。

关键设计：具体实现上，使用了Stable Diffusion模型作为表征提取器。控制策略可以使用各种强化学习算法进行训练，例如PPO。关键参数包括扩散模型的选择、表征提取的位置（例如，扩散模型的中间层）、以及控制策略的网络结构和训练参数。损失函数通常是标准的强化学习损失函数，例如PPO的clip loss。

📊 实验亮点

实验结果表明，使用稳定控制表征学习的策略在OVMM开放词汇导航基准上取得了最先进的性能。此外，在模拟控制任务中，该方法也优于或与最先进的表征学习方法具有竞争力，证明了其有效性和泛化能力。

🎯 应用场景

该研究成果可应用于机器人导航、物体操作、自动驾驶等领域。通过提升具身智能体的场景理解能力，可以实现更智能、更灵活的机器人系统，从而在复杂环境中完成各种任务，例如家庭服务、工业自动化和医疗辅助等。

📄 摘要（原文）

Embodied AI agents require a fine-grained understanding of the physical world mediated through visual and language inputs. Such capabilities are difficult to learn solely from task-specific data. This has led to the emergence of pre-trained vision-language models as a tool for transferring representations learned from internet-scale data to downstream tasks and new domains. However, commonly used contrastively trained representations such as in CLIP have been shown to fail at enabling embodied agents to gain a sufficiently fine-grained scene understanding -- a capability vital for control. To address this shortcoming, we consider representations from pre-trained text-to-image diffusion models, which are explicitly optimized to generate images from text prompts and as such, contain text-conditioned representations that reflect highly fine-grained visuo-spatial information. Using pre-trained text-to-image diffusion models, we construct Stable Control Representations which allow learning downstream control policies that generalize to complex, open-ended environments. We show that policies learned using Stable Control Representations are competitive with state-of-the-art representation learning approaches across a broad range of simulated control settings, encompassing challenging manipulation and navigation tasks. Most notably, we show that Stable Control Representations enable learning policies that exhibit state-of-the-art performance on OVMM, a difficult open-vocabulary navigation benchmark.