Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers

全文总结

研究背景

1. 背景介绍: 这篇文章的研究背景是机器人学习的异质性问题。传统的机器人学习方法通常针对特定任务和特定机器人本体收集数据，这种方法既昂贵又容易过拟合。为了克服这一挑战，研究者们提出了通过在大规模、异构的机器人数据上进行预训练来学习任务和本体无关的共享表示。
2. 研究内容: 该问题的研究内容包括：提出一种名为Heterogeneous Pre-trained Transformers (HPT)的架构，通过异构预训练来学习任务和本体无关的共享表示，并在多个模拟器和真实世界设置中进行实验验证。
3. 文献综述: 该问题的相关工作有：自然语言处理和计算机视觉领域的基础模型通过在大规模、高质量和多样化的数据上进行预训练，实现了任务无关的通用模型。机器人领域的相关工作包括在人类视频上进行视觉部分的预训练以及在统一模型和数据格式下进行全策略预训练，但这些方法通常假设没有本体感知信息，并在迁移学习中添加。

研究方法

这篇论文提出了Heterogeneous Pre-trained Transformers (HPT)架构。具体来说：

• Stem: HPT的stem由本体感知tokenizer和视觉tokenizer组成，将不同本体的异构输入映射到固定数量的tokens。本体感知tokenizer使用MLP将本体感知信息映射到一个特征空间，然后应用正弦位置编码和注意力机制将其映射到tokens。视觉tokenizer使用预训练的冻结特征网络（如ResNet）将相机图像映射到tokens。
• Trunk: trunk是一个共享的可扩展Transformer架构，负责将拼接的tokens映射到共享表示空间。trunk的参数在预训练阶段固定，不随本体和任务变化。
• Head: head将trunk的输出映射到各个下游任务的动作空间。每个本体和任务对应一个独立的head，head可以在迁移学习时重新初始化。

实验设计

• 默认设置: 使用27个机器人远程操作数据集进行预训练，包括Open-X Embodiment数据集的一个子集。使用一个相机视角的场景和预训练的冻结ResNet18图像编码器计算视觉特征。总共使用16k轨迹进行训练。
• 扩展设置: 使用52个数据集和200k轨迹进行预训练，包括模拟数据、部署机器人数据和人类视频数据。实验使用了更大的模型和更高的批量大小。

结果与分析

• 数据扩展: 在不同异构本体上的验证损失表现出稳定且可扩展的趋势。随着数据量和模型规模的增加，验证损失逐渐接近收敛。
• 模型扩展: 随着模型规模的增加（从1M到1B参数），验证损失逐渐降低，直到达到一个平台期。实验发现，深度和宽度扩展的效果相似。
• 预训练在合成数据和互联网人类视频上的效果: 预训练在模拟和人类视频数据上也是可行的，尽管这些数据与真实机器人存在较大的本体差异。这些数据提供了补充的本体数据，展示了HPT框架处理更大异质性的能力。
• 迁移学习: 在多个模拟器基准测试和真实世界任务中进行迁移学习实验，结果表明预训练模型在新本体、任务和环境中的表现优于基线模型。例如，在Sweep Leftover任务中，HPT-B Fine-tuned模型的平均成功率为70.0%，而从头开始训练的模型为43.3%。

优点与创新

1. 异构预训练Transformer（HPT）：提出了HPT架构，通过异构预训练学习任务无关和实体无关的共享表示。该架构将不同实体的 proprioception 和 vision 输入对齐到一个短的令牌序列，并处理这些令牌以映射到不同任务的控制机器人。
2. 大规模多实体数据集：利用了大规模的多实体真实世界机器人数据集以及仿真、部署的机器人和人类视频数据集，研究了跨异构性进行预训练的策略。
3. 实验验证：在超过52个数据集上进行了实验，验证了HPT在不同模拟器基准和现实世界设置中的扩展性和性能提升。
4. 模块化设计：HPT架构模块化为实体特定的stem、共享的trunk和任务特定的head，使得模型可以灵活地适应新的实体和任务。
5. 公开代码和权重：为了促进未来的研究，HPT的代码和权重已开源。

不足与反思

1. 数据集划分简单：平衡数据集中的实体划分相对简单，未来工作应探索更复杂的数据集划分和过滤方法以确保数据质量。
2. 监督学习目标：当前工作专注于监督学习作为预训练目标，未来可以探索超越监督学习的训练目标。
3. 数据规模有限：尽管模型架构和训练过程是模块化的且独立于实体设置，但异构预训练可能需要较长时间才能收敛。
4. 评估限制：模拟和现实世界的评估任务仅限于短视操作任务，固定实体，这可能限制了使用更高容量模型的好处。
5. 任务成功率不高：测试的任务上学习到的策略仍然没有提供非常高的可靠性，通常低于90%。
6. 未来研究方向：包括探索不同的算法、网络架构、训练目标，以及高质量多样化数据集的收集和理解。此外，还需要大规模的统一仿真基准来一致比较不同模型，并扩展到更长视线的精细操作任务。

---

摘要

当今训练通用机器人模型的一个障碍是异质性。以前的机器人学习方法通常收集数据以针对一个特定实现进行训练，这既昂贵又容易过拟合。本研究研究了通过跨不同实现和任务规模在机器人数据上进行异质预训练来学习策略表示的问题。我们提出了异质预训练变换器（HPT），它预训练一个大型、可共享的策略神经网络主干，以学习任务和实现无关的共享表示。这种通用架构将来自不同实现的特定本体感知和视觉输入映射到短序列的令牌，然后处理这些令牌以将它们映射到控制不同任务的机器人。利用最近的大规模多实现现实世界机器人数据集以及仿真、部署的机器人和人类视频数据集，我们调查了跨异质性的预训练策略。我们进行了实验，以研究训练目标的扩展行为，涉及52个数据集。HPT在多个模拟基准测试和真实世界设置中，对未见过的任务提高了超过20%的微调策略性能。

1 引言

当今构建机器人策略是困难的：它通常需要为每个机器人、任务和环境收集特定数据，而学习到的策略并不能在这些特定设置之外泛化。一个历史教训已经彻底改变了机器学习，即在大量高质量和多样化的数据上进行预训练可以带来通常比特定模型表现更好的通用模型。开源大规模数据收集的最新进展使得这条路径成为可能，但大规模机器人数据中存在的异质性（如不同的机器人硬件和不同环境）提出了一个重大挑战。该领域目前的一个核心问题是