用于策略学习的任意点轨迹建模

24年1月来自伯克利分校、清华、斯坦福、上海AI实验室、上海姚期智研究院和香港中文大学的论文“”Any-point Trajectory Modeling for Policy Learning“。

从演示中学习是教授机器人新技能的有效方法，拥有更多演示数据通常可以提高策略学习。然而，收集演示数据的高成本是一个重大瓶颈。视频作为丰富的数据源，包含行为、物理和语义知识，但由于缺乏动作标签，从中提取特定于控制的信息具有挑战性。这项工作引入一个框架，即任意点轨迹建模 (ATM)，它利用视频演示通过预训练轨迹模型来预测视频帧内任意点的未来轨迹。经过训练后，这些轨迹可提供详细的控制指导，从而能够以最少的动作标记数据学习强大的视觉运动策略。在模拟和现实世界中评估的 130 多个语言条件任务中，ATM 的表现平均比强大的视频预训练基线高出 80%。

近年来，计算机视觉和自然语言理解取得了重大进展 [22, 7]，其中大型数据集的可用性起着至关重要的作用。同样，在机器人技术中，扩大人类演示数据一直是学习新技能的关键 [6, 34, 14]，随着数据集的增大，性能有明显的提升趋势 [29, 6]。然而，人类演示（通常是通过遥控设备 如 GELLO 和VR Vive 收集的动作标记轨迹  [56, 53]）的收集既费时又费力。例如，在 RT-1 [6] 中收集 13 万条轨迹需要 17 个月，这使得数据收集成为机器人学习的主要瓶颈。

视频包含有关行为、物理和语义的知识，可提供另一种数据源。然而，缺乏动作标签使得在策略学习中使用视频数据变得困难。先前的研究已经通过使用自监督目标进行视频预训练来解决这个问题，以学习策略学习的观察特征表示 [44, 33, 42]。然而，特征表示只描述了当前时间步的状态，很大程度上忽略了预测未来状态的转换动态。为了明确地对转换动态进行建模，先前的研究已经开发了视频预测模型，可以根据当前图像帧预测未来的图像帧来指导策略学习 [12, 55, 13]。然而，学习用于控制的视频预测模型带来了两个挑战。首先，视频预测任务通过对每个像素的变化进行建模来避免任何抽象，将物理运动与纹理、光照等视觉外观结合起来。这种耦合使建模变得困难，常常导致幻觉和不切实际的未来预测 [12]。其次，这些模型在训练和推理方面都需要大量计算。在计算资源有限的情况下，性能会显著下降。此外，高推理成本迫使这些模型采用开环执行 [12, 5]，这往往会导致策略的鲁棒性降低。

本文提出ATM，如图所示：给定任务指令和图像帧中任意一组点的初始位置，任意点轨迹模型 (ATM) 可以根据任务预测这些点的未来轨迹。在无动作视频数据集上训练模型后，预测的轨迹，对一组语言为条件的操作任务，可作为学习视觉运动策略的有效指导。

有了足够的视频预训练，将能够从有限的动作标记轨迹中学习可推广的策略。先前的研究 [12、23、5] 主要依赖像素级未来帧预测作为视频预训练。虽然视频预测在训练和推理阶段都是资源密集型的，但它专注于重建像素级细节，而这通常与策略学习无关，可能会对后续策略学习的效率产生不利影响。

如图所示，任意点轨迹建模 (ATM) 是一个两阶段框架：首先使用大规模无动作视频学习预测视频帧中的未来点轨迹作为预训练，然后使用预测的轨迹通过少量带有动作标记的演示来指导策略学习。

在训练轨迹 Transformer 根据观察结果预测未来轨迹后，便可以学习由这些预测轨迹指导的策略。

a) 任意点跟踪：在轨迹 Transformer 预训练期间，可以过滤没有大移动的轨迹。但是，使用这种启发式方法需要知道每个点的未来位置，这在策略推理期间计算起来可能很昂贵。相反，只需在网格上使用一组固定的 32 个点作为策略就足够了。这种采样方法避免了学习关键点或寻找要跟踪点的潜在复杂性 [48]，并且在实践中效果很好。ATM 对输入点集具有置换不变性，还对点的分布具有鲁棒性，这样能够使用与训练不同的点采样方案进行策略学习。

b) 轨迹引导策略学习：如上图（b）所示，轨迹引导策略 π(at|ot, pt:t+H) 输入当前观察 ot 和预测轨迹 pt:t+H，并预测动作。如下图显示了策略架构的简化说明。Transformer 策略架构遵循先前研究 [27, 21]。尽管仅预测轨迹就已经提供了预测动作的丰富信息，但仍将上下文图像观察纳入策略中，以免丢失任何信息，如先前研究 [26]方法。

将融合前后的轨迹tokens与图像tokens（早期融合和晚期融合）结合起来，以确保可以有效整合来自轨迹的引导信息。令人惊讶的是，由于轨迹已经提供了细粒度的子目标，该策略在此阶段不再需要语言指导作为任务规范。本质上，提供的轨迹将困难的策略学习问题简化为更容易的子目标跟踪问题，将策略简化为