Theia: Distilling Diverse Vision Foundation Models for Robot Learning

📄 arXiv: 2407.20179v2 📥 PDF

作者: Jinghuan Shang, Karl Schmeckpeper, Brandon B. May, Maria Vittoria Minniti, Tarik Kelestemur, David Watkins, Laura Herlant

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

发布日期: 2024-07-29 (更新: 2024-10-10)

备注: CoRL 2024

💡 一句话要点

Theia：为机器人学习提炼多样化视觉基础模型，提升泛化能力

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

关键词: 机器人学习 视觉基础模型 知识蒸馏 视觉表征 泛化能力

📋 核心要点

  1. 基于视觉的机器人策略学习需要对多样化的视觉任务有整体理解，而不仅仅是分类或分割等单任务需求。
  2. Theia通过蒸馏多个在不同视觉任务上训练的视觉基础模型，学习到包含多样化视觉知识的丰富视觉表征。
  3. 实验结果表明，Theia在机器人学习任务上，使用更少的数据和更小的模型，超越了教师模型和现有方法。

📝 摘要（中文）

本文提出Theia，一种用于机器人学习的视觉基础模型，它通过蒸馏多个在不同视觉任务上训练的现成视觉基础模型，来获得丰富的视觉表征。Theia编码了多样化的视觉知识，从而增强下游机器人学习的性能。大量实验表明，Theia使用更少的训练数据和更小的模型尺寸，优于其教师模型和先前的机器人学习模型。此外，本文还量化了预训练视觉表征的质量，并假设特征范数分布中更高的熵值能够带来更好的机器人学习性能。代码、模型和演示可在https://theia.theaiinstitute.com 获取。

🔬 方法详解

问题定义：现有的基于视觉的机器人学习方法通常依赖于在特定任务上训练的模型，泛化能力有限。它们难以捕捉到不同视觉任务之间的共性知识，导致在新的、未见过的场景中表现不佳。此外，从头开始训练这些模型需要大量的标注数据和计算资源。

核心思路：Theia的核心思路是利用知识蒸馏，将多个在不同视觉任务上预训练的视觉基础模型的知识迁移到一个更小的模型中。通过这种方式，Theia能够学习到包含多样化视觉知识的丰富视觉表征，从而提高其在机器人学习任务中的泛化能力。

技术框架：Theia的整体框架包括以下几个主要步骤：1) 选择多个在不同视觉任务上预训练的视觉基础模型作为教师模型；2) 构建一个更小的学生模型；3) 使用知识蒸馏技术，将教师模型的知识迁移到学生模型中；4) 在下游机器人学习任务上对学生模型进行微调。

关键创新：Theia的关键创新在于其利用知识蒸馏技术，将多个异构的视觉基础模型的知识融合到一个统一的模型中。这使得Theia能够学习到更全面、更鲁棒的视觉表征，从而提高其在机器人学习任务中的性能。此外，论文还提出了一个量化预训练视觉表征质量的指标，并发现特征范数分布中更高的熵值与更好的机器人学习性能相关。

关键设计：Theia使用了Transformer作为其基础架构。损失函数包括知识蒸馏损失和下游任务的损失。知识蒸馏损失用于衡量学生模型和教师模型之间的输出差异，促使学生模型学习教师模型的知识。下游任务的损失用于在特定机器人学习任务上对学生模型进行微调。论文还探索了不同的教师模型组合和蒸馏策略，以找到最佳的配置。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

Theia在多个机器人学习任务上取得了显著的性能提升。例如，在物体抓取任务中，Theia使用更少的数据和更小的模型，超越了现有的最先进方法。实验结果表明，Theia能够有效地利用从多个视觉基础模型中提取的知识，从而提高其在机器人学习任务中的泛化能力。此外，论文还发现特征范数分布中更高的熵值与更好的机器人学习性能相关，为未来的研究提供了新的思路。

🎯 应用场景

Theia具有广泛的应用前景，可应用于各种需要视觉感知的机器人任务，例如物体抓取、导航、操作等。它可以帮助机器人更好地理解周围环境，从而做出更智能的决策。此外，Theia还可以用于开发更通用、更鲁棒的机器人系统，从而降低机器人部署和维护的成本。未来，Theia有望成为机器人视觉感知领域的重要基石。

📄 摘要（原文）

Vision-based robot policy learning, which maps visual inputs to actions, necessitates a holistic understanding of diverse visual tasks beyond single-task needs like classification or segmentation. Inspired by this, we introduce Theia, a vision foundation model for robot learning that distills multiple off-the-shelf vision foundation models trained on varied vision tasks. Theia's rich visual representations encode diverse visual knowledge, enhancing downstream robot learning. Extensive experiments demonstrate that Theia outperforms its teacher models and prior robot learning models using less training data and smaller model sizes. Additionally, we quantify the quality of pre-trained visual representations and hypothesize that higher entropy in feature norm distributions leads to improved robot learning performance. Code, models, and demo are available at https://theia.theaiinstitute.com.