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 相机。

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

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

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 位颜色来渲染彩色点云。将其扩展到具有任意数量通道的图像，通过打包点索引来代替。