使用异构预训练 Transformer 扩展本体感受-视觉的学习

24年9月来自 MIT 何凯明团队和 Meta 的论文“Scaling Proprioceptive-Visual Learning with Heterogeneous Pre-trained Transformers”。

当今训练通用机器人模型的障碍之一是异构性。以前的机器人学习方法通​​常收集数据，用一个特定的具身来训练一个任务，这种方法成本高昂，而且容易过拟合。这项工作研究通过对不同具身和任务的机器人数据进行大规模异构预训练来学习策略表示的问题。提出异构预训练 Transformer (HPT)，它预训练一个大型、可共享的策略神经网络主干，以学习与任务和具身无关的共享表示。这种通用架构将来自不同具身的特定本体感觉和视觉输入对齐到一小段token序列，然后处理这些token以映射到控制机器人以执行不同的任务。利用最近的大规模多具身真实世界机器人数据集以及模拟数据集、部署的机器人数据集和人类视频数据集，研究跨异构的预训练策略。其实验研究训练目标的扩展行为，范围达到 52 个数据集。

机器人领域的大规模策略学习分别利用来自真实机器人 [6, 72]、人类视频 [53, 48] 和模拟域 [32, 62, 82, 79] 的各种数据。此外，还有多任务学习 [64, 65, 84, 22]、元学习 [78, 54, 18]、小样本学习 [83] 和成队学习 [81] 方面的研究。最近，RT-X、Octo、OpenVLA [6, 14, 55, 35] 在来自各种机器人实例的数据集上训练通用的视觉-语言-动作机器人策略。

所谓异构预训练概念如图所示。它将不同的具身方案（每个都有自己的本体感觉和视觉传感器）通过针对具身方案的token化器（“stem”）映射到共享的潜空间上。这将来自不同具身方案的异构数据对齐到联合表示空间中。这样就能够在所有异构数据集的联合上训练共享的 Transformer 主干。预训练的 Transformer 可以迁移到新的具身方案，并在迁移时学习一个小型的新token化器。

异构预训练 Transformers (HPT)，是一类架构，旨在可扩展地从异构具身数据中学习。HPT 模块化通用策略网络架构（如图所示），并使用监督学习预训练潜 Transformer 的策略表示。受到从多模态数据中学习的启发 [1, 73, 19, 30]，用针对具身的token化器（称为“stem”）来对齐各种传感器输入，例如摄像机视图和本体感觉输入。“trunk”在数据集之间共享和预训练，并在适应预训练期间未知的新具身和任务时进行传输。此外，使用针对任务的操作解码器（称为“头”）来产生操作输出。至关重要的是，在“对每个具身进行token化”之后，HPT 在一小段潜token 序列的共享空间上运行。这种层次结构的灵感来源于人类如何在脊髓神经回路层面处理特定运动反应和感知刺激之间的反馈回路 [68]。

在具有跨具身的异构机器人学习中，数据来自不同的领域（例如模拟和真实机器人），跨越 RGB 图像、语言指令、深度图、3D 点云和触觉图像等传感模态。每个机器人都是一个独特的硬件具身，具有不同的自由度、末端执行器、传感器配置、控制器和动作空间以及特定于应用程序的物理设置。

网络架构模块化为具身特定的stem、共享chunk和任务特定的头部。直观地讲，如图所示的stem，是神经网络的早期层，可将来自异构具身和模态的传感输入对齐到共享表示空间中。网络的共享中间部分称为trunk，它将传感表征处理为可用于多项任务的潜表征。最后，网络的最后一部分是头部，它将潜表示映射到感兴趣的单个任务动作空间。训练过程称为异构预训练，根据采样的方案和任务数据分配和对齐特定的stem/头部对，同时仍然享有在共享trunk中进行联合训练的好处。这可以被认为是使用神经网络对每个实施方案进行token化，并减轻在标准训练过程中将具身方案统一为同质数据形式的需要。

HPT 中的stem（上图所示）由本体感受token化器和视觉token化器组成。这些token化器将来自不同具身的异构输入映射到固定数量、固定维度的tokens，这使得stem能够以相同的方式处理它们，尽管异构性很大，并且能够享受固定上下文长度的扩展和推理优势。关键思想是利用注意机制 [77, 30, 9] 将固定数量的可学习token关注到观察的特征上。虽然主要关注本体感受和视觉，但处理触觉、3D 和动作输入中的其他类型传感器异构，可以在stem中灵活扩展。

作为预训练的核心组件，trunk架构遵循一个transformer，参数化为维度为 d 潜空间参数。输出token序列长度 L 与输入token序列长度相同。输出token序列只是池化作为观察的最终组合特征。trunk在不同的具身和任务之间共享，以捕获复杂的输入输出关系（即trunk参数的数量是固定的，与具身和任务的数量无关）。

策略头获取stem transformer的输出并将其映射到每个数据集中的动作空间 A。对于每个具身和任务，策略头可以是任意架构（例如 MLP），它将trunk的池化特征作为输入并输出规范化的动作轨迹。策略头重新初始化以迁移到新的具身。

其训练过程有两个数据扩展轴：一个数据集的数量，和数据集的总数。在预训练阶段，每次迭代只更新trunk参数，而每个异构具身和任务的stem和头部则基于训练批次采样进行更新。

使用 AdamW [46] 优化器训练 HPT，权重衰减率为 0.05，基本学习率为 0.0002，采用带有预热和丢弃的余弦学习率调度。根据批次大小应用基本学习率的比例缩放。为了在迁移学习期间支持各种范围，在训练期间沿每个批次的时间维度应用随机掩码。由于动作轨迹在预测范围内可能会出现不平衡的损失，使用 Huber 损失。预训练阶段在学习率调度和优化器中的各种超参以及验证数据集的选择方面都是稳定的。代码是开源的，预发布的模型可以从 Huggingface 轻松下载。

在实践中，由于训练损失可能因不同的数据集而异，并且目标是在所有具身和任务上表现良好，因此应用加权采样程序进行数据加载。对于每次训练迭代，都会以数据集大小的指数作为温度，逆概率方式对数据集进行采样。具体来说，计算每个数据集大小的平方根，并对这些大小求和以计算归一化常数。然后，对于每个批次，以相应的概率从这些数据集中采样。这可以防止大型数据集主导整个训练周期，这是多任务学习中的常见做法。

注：每个具身的stem和头部的更新频率与trunk不同，类似于混合专家 [71] 训练程序。特别是在分布式训练设置下，每个stem和头部都使用来自特定具身和任务的数据进行训练，trunk将从训练工作器的所有批次中积累梯度。这些预训练实验的计算资源范围从 8 V-100 到 128 V-100，训练时间从半天到 1 个月不等。总数据集磁盘大小约为 10Tb，RAM 内存要求低于 50Gb。

对于模拟基准，用已发布的数据集作为专家演示 [89, 81, 13]。总之，Metaworld [89] 使用腕部摄像头视图，Robomimic [49] 以及 Simpler [42] 使用第三人称视图，并通过数据集使用自己的本体感受定义。Fleet-Tools [81] 使用两个视图作为输入，并使用末端执行器姿势作为本体感受输入。用预训练的冻结 ResNet 特征对图像进行编码，并在将本体感受输入传递到 stem 之前对其进行规范化。为除 Metaworld 之外的所有这些模拟基准训练单任务策略。

对于 Simpler[42] 基准，专注于 Close-Drawer、Move Near 和 Pick Coke Can 任务以及具有视觉匹配设置的 Google EDR 实例。测试 9 种不同的初始化，总共 218 个episodes。注：模拟任务侧重于语言调节，不暴露本体感受输入，这不是最合适的 HPT 测试平台。为了解决这些问题，在具有 79 条轨迹的 RTX 监督数据集上微调 HPT 作为其他模拟基准。用 HPT-base 作为本次实验的主干（backbone）。

默认情况下，用 20000 次迭代进行训练，批量大小为 512，学习率较小，为 1e 5。图像和状态主干是隐维度为 128 的单层 MLP，头部是双层 MLP。只使用长度为 1 的观察窗口和 MLP 作为策略头。每个训练数据集每个任务使用 10-100 条轨迹，每个测试涵盖具有不同初始条件的 50 个episodes。模拟中的每个轨迹在场景初始化方面略有不同。为了减少方差，进行 5 次独立训练，并取每个基线的平均值。

在真实世界中，在两种不同的机器人设置中试验机器人工具使用的任务 Sweep Leftover, Fill Water, Scoop Food 和 Switch Insertion。虽然在这两种设置中，都使用 Franka Panda 作为机器人，但传感器位置以及动作空间截然不同。为每个任务收集大约 100 个演示，并对每个任务进行 15 次试验评估，以测量平均成功率。

在评估过程中，一个人始终监督机器人。评估 episode 可以因安全问题、机器人故障、超时等原因终止。如果episode 完成了任务，则认为它成功了。在 Fill Water 任务中，成功分数为 1 表示将一些水倒入碗中。在 Sweep Leftover 任务中，成功分数为 1 表示将所有推入盘子中，成功分数为 0.5 表示将一些推入盘子中。在Scoop Food 任务中，成功分数为 1 表示舀出一些狗粮并全部倒入碗中，分数为 0.5 表示只舀出一些。在Switch Insertion 任务 [90] 中，成功分数为 1 表示开关精确插入 PCB 板上的三个引脚。机器人在尝试插入之前会移动到预定义的姿势。选择这些具有挑战性的任务，是因为它们需要与工具和颗粒体进行丰富的接触交互，并且需要高精度和密集的接触。确保机器人的初始条件相同。由于这些任务的复杂性和人为错误，目标设置的初始条件并不完全相同。

迁移学习。对于现实世界实验中的策略头，尝试 MLP 和扩散策略 [13]。与其他方法中通常用于单个任务的更大模型（例如 100M）相比，微调的活动参数不超过 3Mb。用长度为 2 的观察历史窗和较小的学习率 2e 5。在单个 NVIDIA RTX 2080Ti GPU 上，以批量大小 256 进行 20000 次迭代（大约 4 小时时间）。