RVT-2

提出了 RVT-2，一种多任务三维操作模型，其训练速度比前代 RVT 快 6 倍，推理速度快 2 倍。在 RLBench [24] 基准测试中，RVT-2 的成功率从 65% 提升至 82%。

架构改进：

• 引入了多阶段推理管道，使网络能够聚焦于感兴趣区域并预测更精确的末端执行器姿态；
• 采用了凸上采样技术，节省 GPU 内存并提高训练速度；
• 改进了末端执行器旋转预测，使用位置条件特征代替 RVT 中仅使用全局特征的方法。

系统级优化：

• 开发了一个定制的虚拟图像渲染器，替代 RVT 中使用的通用渲染器（PyTorch3D [35]）。这一加速渲染库显著提高了训练和推理速度，同时减少了内存使用；
• 采用了最新的 Transformer 模型训练技术，包括快速优化器和混合精度训练。

---

Architectural Changes

Multi-stage Design

RVT 使用围绕机器人的固定视图预测夹爪姿态。然而，当感兴趣的物体非常小且夹爪姿态需要非常精确时（例如插销插入任务），这些固定视图可能不足。

因此，RVT-2 采用了多阶段设计。在第一阶段（粗阶段），它使用固定视图预测感兴趣区域。然后，RVT-2 对感兴趣区域进行放大并在其周围重新渲染图像。使用 $4$ 倍放大因子，这些放大的图像被用于更精确的夹爪姿态预测。

Convex Upsampling

RVT-2 移除了特征上采样，直接从 token 分辨率的特征预测形状为 $h×w$ 的热力图。具体来说，它使用了 [46] 提出的凸上采样层。凸上采样层使用粗网格中特征的学习凸组合在更高分辨率上进行预测。

Location Conditioned Rotation

RVT 和 PerAct 使用全局视觉特征预测末端执行器的旋转。当存在多个有效的末端执行器位置且旋转取决于位置时，这可能会出现问题。例如，在堆叠积木的任务中，场景中可能有两个相似的积木，但它们的方向不同。在这种情况下，选择任意一个积木都是有效的步骤。然而，由于积木的方向不同，末端执行器的旋转将取决于所选的位置。

为了解决这一问题，RVT-2 使用从特征图中末端执行器位置处提取的局部特征进行旋转预测。

System Related Changes

Point-Renderer

在 CUDA 中实现了一个自定义的基于投影的点云渲染器，通过以下三个步骤渲染点云：

• 投影：将每个三维点投影到图像平面，计算其深度和像素坐标；
• 深度排序：为每个像素找到深度最小的点，并将其 RGB 值和深度分配给该像素；
• 屏幕空间扩散：在屏幕空间中对点进行扩散，以减少噪声。

这一自定义渲染器相比起 PyTorch3D 显著加速了训练和推理过程。