RVT-2：通过少量演示学习精确操作

24年6月来自Nvidia的论文“RVT-2: Learning Precise Manipulation from Few Demonstrations”。

本工作研究如何构建一个机器人系统，该系统可以根据语言指令解决多个 3D 操作任务。为了在工业和家庭领域发挥作用，这样的系统应该能够通过少量演示学习新任务并精确解决它们。先前的研究，如 PerAct [40] 和 RVT [17]，已经研究过这个问题，但它们在执行需要高精度的任务时往往举步维艰。本文研究如何使它们更有效、更精确、更快速。通过结合架构和系统级改进，提出的 RVT-2，是一个多任务 3D 操作模型，与前身 RVT 相比，它的训练速度提高 6 倍，推理速度提高 2 倍。RVT-2 在 RLBench [24] 上取得新的最高水平，将成功率从 65% 提高到了 82%。RVT-2 在现实世界中也很有效，它只需 10 次演示就可以学习需要高精度的任务，例如拿起和插入插头。

机器人学习的终极目标之一是构建通用机器人系统，该系统可以解决多项任务并推广到未知的环境配置。要发挥作用，这样的系统应该能够进行精确操作，并且只需要对新任务进行几次演示。例如，在工业制造环境中，期望一个人向机器人演示几次高精度任务，比如将钉子插入机器人，之后机器人应该开始独立执行该任务。在家庭和零售等其他领域也可以找到类似的例子。这项工作研究构建一个只需几次演示就能精确解决各种任务的操作系统的问题。该系统应具有三个关键特征：（1）处理多项任务，（2）只需要几次演示，（3）高精度地解决任务。

先前的工作已朝着这一目标取得重大进展。从研究平面拾取和放置任务的 Transporter Networks [52] 和 IFOR [16] 等开始，最近的工作已经超越 2D 平面，并通过一些示例研究 3D 中的操作 [25]。一些值得注意的方法是 PerAct [40] 和 RVT [17]。给定语言指令，PerAct [40] 通过预测下一个关键帧姿势，采用了多任务 Transformer 模型进行 3D 操作。尽管 PerAct 取得令人印象深刻的性能，但它使用基于体素的场景表示，限制其可扩展性。RVT [17] 通过提出一种用于编码场景的新型多视图表示来解决 PerAct 的局限性。多视图表示具有多种优势，包括更快的训练速度、更快的推理速度和更好的任务性能。与 PerAct 相比，RVT 的训练速度提高 36 倍，并在 RLBench [24] 的 18 个任务上将性能从 48% 提高到 63%。

有一个问题：是什么阻碍了 RVT 实现更高的性能？经过仔细分析，RVT 在需要高精度的任务上表现不佳，例如拧灯泡或插入钉子。在分析过程中，还发现几个进一步提高系统训练和推理速度的机会。通过架构和系统级改进，能够提高 RVT 的速度和效率。因此，本文推出 RVT-2，它将 RVT 的训练速度提高 6 倍（从每天 240 万个样本提高到每天 1600 万个样本），推理速度提高 2 倍（从 11.6 fps 提高到 20.6 fps），任务成功率提高 15 个百分点（从 62.9 提高到 77.6），在 RLBench 基准上取得最先进的结果。

如图所示RVT-2 执行高精度任务。给定一个语言指令，单个 RVT-2 模型可以执行多个 3D 操作任务，包括需要毫米级精度的任务，例如将钉子插入孔中和将插头插入插座。RVT-2 每项任务使用 ∼10 次演示进行训练，并且仅使用单个第三人称 RGB-D 相机。

与在 2D 自上而下设置中操作 [52, 16, 39] 相比，由于动作空间的自由度更高以及 3D 空间推理的复杂性 [14]，推断机器人在完整 3D 空间中的运动和交互要困难得多。为了解决 3D 空间中的操作，最近的研究利用各种感知表示。相机图像已广泛用于基于视觉的操作，例如在 RT-1 [1]、RT-2 [2] 和 ALOHA [53] 等模型中。为了更有效的 3D 空间推理，通常需要深度信息，其中 RGB-D 图像被假定为操作策略的输入 [44]。PolarNet [5] 和 M2T2 [51] 直接使用从 RGB-D 图像重建的点云，并使用编码器和Transformer对其进行处理以预测动作。 C2F-ARM [25]、PerAct [40] 和 FourTran [22] 对点云进行体素化，并使用 3D 卷积网络作为动作推理的骨干。Act3D [13] 和 ChainedDiffuser [48] 将场景表示为多尺度 3D 特征云。为了提高时间效率和任务功效，RVT [17] 建议使用多视角虚拟图像作为场景表示。尽管如此，这些先前模型中的大多数，仅适用于不需要高精度动作的现实任务。

RVT-2 允许进行三维和精确操作。单个 RVT-2 模型经过训练，可以根据语言指令解决多个任务，并且每个任务只需要几次演示。RVT-2 以 RVT [17] 为基础，RVT 是一种最先进的 3D 目标操作模型。与 RVT 类似，RVT-2 基于关键帧操作范例。但它通过一系列改进实现更好的任务性能、精度和速度。这些改进可以分为两类：与神经网络变化相关的架构改进，以及与软件优化相关的系统改进。

如图所示RVT-2的架构图：

RVT 使用机器人周围的一组固定视图来预测夹持器姿势。当感兴趣的物体非常小且夹持器姿势需要非常精确（例如将钉子插入孔中）时，这些固定视图可能不足以完成任务。因此，如图所示RVT-2 采用了多阶段设计，其中在第一阶段或粗略阶段，它使用一组固定视图来预测感兴趣的区域。然后，RVT-2 放大感兴趣的区域并重渲染其周围的图像。用 4 倍的放大倍数，这意味着放大的相机覆盖的区域大小是粗略相机的 1/4。它使用这些特写图像来精确预测夹持器姿势。这种自适应渲染是可能的，因为 RVT 中提出的虚拟渲染提供了灵活性。

RVT 基于 ViT（视觉Transformer）[12]，它将图像划分为 t1×t2 个patches。每个图像patch都作为维度为 d 的单个token进行处理。为了根据图像token预测热图，RVT 首先将图像token排列在相应的patch位置。这会产生形状为 t1 × t2 × d 的 3D 特征。然后，RVT 使用转置卷积将这些特征上采样到 h×w 的图像分辨率，从而创建形状为 h×w×d 的特征。最后，这些形状为 h×w×d 的特征用于预测热图。此操作序列非常有效，但由于形状为 h×w×d 的中间特征较大，因此占用大量内存。

为了解决这个问题，RVT-2 删除了特征上采样，并直接从 token 分辨率的特征中预测形状为 h×w 的热图。具体来说，它使用 [46] 提出的凸上采样层。凸上采样层（convex upsampling）使用粗网格中学习到的特征凸组合来以更高的分辨率进行预测。[46] 展示了它如何在更高的分辨率下实现更清晰的预测。经验发现，凸上采样可以节省内存而不会牺牲性能。凸上采样层不需要任何特殊实现，可以使用 PyTorch 中的 fold 函数表示。

RVT 中的网络参数（如虚拟图像大小（220）和patch大小（11））可能不是 GPU 的最佳参数，因为它们不能被 2 的幂（如 16）整除。RVT-2 合理化了这些参数，使神经网络更加 GPU 友好，从而提高了其速度。 RVT-2 采用与 ViT [12] 类似的参数，即图像大小为 224，patch大小为 14。除了更加 GPU 友好之外，这些参数还减少了多视图Transformer内的总tokens数，该tokens等于 (image size/patch size)^2，从而进一步提高了速度。这些选择使 RVT-2 在训练和测试期间更快，而不会影响性能。

RVT 和 PerAct 使用全局视觉特征（例如整个图像的最大池化）来预测末端执行器的旋转。当有多个有效的末端执行器位置并且末端执行器的旋转取决于位置时，这可能会出现问题。例如，考虑堆叠块的任务，其中场景有两个相似的块但方向不同。在这里，选择两个块中的任何一个都是有效的步骤。但是，由于块具有不同的方向，因此末端执行器的旋转将取决于所选的末端执行器位置。由于 RVT 仅使用全局视觉特征来预测旋转，因此它无法处理这种情况。为了解决这个问题，RVT-2 使用从末端执行器位置的特征图中池化的局部特征进行旋转预测。这允许 RVT-2 进行位置相关的旋转预测。

RVT 使用放置在正交位置（即后、前、上、左和右）的五个虚拟摄像机渲染场景点云。这种选择基于他们的观察，即较少的摄像机视图会降低性能。然而，在多阶段 RVT-2 模型中，仅使用三个视图（即正面、顶部和右侧）就足够了，并且不会牺牲性能。这可能是因为 RVT-2 使用放大视图进行最终预测。更少的虚拟视图减少了渲染器要渲染的图像数量和多视图转换器要处理的tokens数量。因此，这提高了训练和推理速度。

RVT 使用 PyTorch3D [35] 来渲染虚拟 RGB-D 图像。PyTorch3D 是一个有吸引力的选择，因为它的界面易于使用。然而，它是一个功能齐全的可微分渲染器，在点云渲染中会产生大量的时间和内存开销。为了避免这种情况，在 CUDA 中实现一个基于投影的自定义点云渲染器。该渲染器执行 3 个步骤将具有 N 个点的点云渲染到大小为 (h, w) 的 RGB 图像和深度图像：

a) 投影：对于索引 n ∈ {0,1…N} 和 RGB 值 fn 的每个 3D 点，它使用相机内外参计算深度 dn 和图像像素坐标 (xn,yn)。从 2D 像素坐标 (xn,yn)，它计算线性像素索引 in = xn · w + yn 。使用 GPU 矩阵乘法可以轻松加速投影操作。

b) Z 排序：对于图像中线性索引为 j 的每个像素，它会在投影到像素 {n | in = j} 的点集中找到深度为 dn 最小的点索引。它将点的 RGB 值 fn 分配给 RGB 图像的像素 j，将深度 dn 分配给深度图像的像素 j。

为了加速 Z 排序，将每个点的深度和索引打包成一个 64 位整数，这样最高有效 32 位编码深度，而最低有效位编码点索引。然后，可以使用两个 CUDA 内核实现 Z 排序。首先，对点云点进行并行循环，尝试使用 atomicMin 操作将每个打包的深度索引存储到像素 j 处的深度索引图像中。只有每个像素处最小深度点存储的深度索引才会保留。第二个内核在对像素进行循环时，解包深度-索引并查找点特征，创建深度和特征图像。这个技巧是由 Schu ̈tz [38] 提出的，用于通过打包 32 位颜色来渲染彩色点云。将其扩展到具有任意数量通道的图像，通过打包点索引来代替。

c) 屏幕-空间splatting：前两个步骤足以生成渲染图像。但是，这些点被视为无限小的光源，这会在屏幕空间点云分辨率不高于图像分辨率的区域产生噪声。解决这个问题的常用方法是 3D splatting，其中每个点都由某种有限大小的几何建模。将每个点表示为面向相机的半径为 r 圆盘。可以在投影和 z 排序之后在屏幕空间中计算此splatting，从而减少投影和 z 排序所需的计算。对于图像中的每个像素 j，在邻域中搜索深度最低的另一个像素 k。如果像素 k 的深度 dk < d j ，并且比 r · focal length/dk 更近，则用像素 k 的特征和深度替换像素 j 的特征和深度。

通过采用训练Transformer的最新发展来优化 RVT 的训练流水线。分析各种技术，并采用那些在不影响性能的情况下提高速度的技术。具体来说，用混合精度训练、8 位 LAMB 优化器 [10] 和基于 xFormers 注意层的快速 GPU 实现 [31]。

仿真环境实验。在 RLBench [24] 中开发的标准多任务操作基准上进行实验，该基准已被以前的研究 [40, 13, 17] 采用。基准包含 18 项任务，包括非抓握任务（如按钮）、常见的拾取和放置任务（如放置葡萄酒）以及需要高精度的钉入孔任务（如插入钉）。每个任务都由语言描述指定，包含 2 到 60 个变型，例如处理不同颜色或位置的物体。带有平行钳口夹持器的 Franka Panda 机器人被命令完成任务。该任务和机器人通过 CoppelaSim [36] 进行模拟。输入的 RGB-D 图像分辨率为 128×128，由安装在机器人前部、左肩、右肩和手腕处的四个无噪声摄像头拍摄。用与 PerAct 和 RVT 相同的数据集训练和测试 RVT-2，每个任务有 100 个演示用于训练，25 个未见过的演示用于测试。

真实世界实验。在与 RVT 类似的真实世界操控设置上将 RVT-2 与 RVT 进行比较。该设置由静态安装的 Franka Emika Panda 手臂和静态第三人称视角 Azure Kinect RGB-D 摄像头组成。没有遵循 RVT 中的摄像头位置，而是将摄像头移近机器人的工作空间，以确保点云的质量以实现高精度操控。将这个摄像头位置用于所有任务。除了 RVT 中使用的相同五个任务（堆叠积木、按压消毒剂、将记号笔放入杯子/碗中、将物体放入抽屉、将物体放入架子）之外，还评估了 IndustRealKit [45] 中的三个高精度任务：拾取并插入 16 毫米钉子、拾取并插入 8 毫米钉子、拾取并插入插头。这两个钉子任务包括从一个孔中拾起钉子并将其插入另一个孔中。在评估过程中，拾取和放置位置都是在工作表面上随机分配的。插头任务包括从托盘中拿起一个 2 脚插头并将其插入垂直安装的插座。此任务超越了 2D 拾取和放置，需要在 3D 空间中进行精确操作。同样，在评估过程中，插头托盘和插座的位置在其工作表面上是随机的。对于 RVT [17] 中的任务，收集与它们报告中相同数量的演示（∼ 10）。对于新的高精度任务，每个任务收集 10 个演示。

如图是真实世界设置（顶部）和来自[45]的钉子拾取和插入任务：

如图所示：来自 [45] 的两脚插头拾取和插入任务；以及考虑的物体位置变化。