异步机器人推理：解耦动作预测和执行

摘要： 机器人策略变得越来越庞大，并且倾向于预测未来一系列的动作，而不是单一的下一个动作。这导致机器人在等待新动作执行时处于空闲状态，在执行时引入了明显的延迟，并且缺乏响应性。异步推理通过将动作预测与动作执行解耦，收紧了控制循环，消除了运行时的延迟，并实现了更具适应性的控制。在这篇博文中，我们介绍了异步推理背后的基础知识，以及如何利用它来提高机器人策略在现实世界中的性能。

目录

• 快速入门
• 异步推理：深度解析
• 1. 为什么顺序推理不够用
• 2. 异步推理，简而言之
• 3. 系统架构
  • 机器人客户端
  • 策略服务器
• 4. 分析异步推理
• 5. 在您的设置中使用异步推理
• 结论

开始入门

按照我们的教程开始使用异步推理。

顺序推理（上）与异步推理（下）。通过允许重新规划和更紧凑的控制循环，异步推理实现了 (1) 尝试恢复，以及 (2) 任务完成速度提升约 2 倍。顺序推理在抓取物体失败后仍会继续执行当前的动作序列，而异步推理可以重新规划并执行新的动作序列。两种设置使用相同的策略！

异步推理：深度解析

通过异步推理，我们将动作执行与动作预测解耦。考虑到当前流行的模型如 [ACT]、[OpenVLA]、[PI0] 和 [SmolVLA] 的趋势，这一点尤其重要，因为它们倾向于根据观察 $o_t$ 输出一系列动作 $a_{t:t+H}$，而不是单个动作 $a_t$。您可以通过使用 LeRobot 运行所有这些模型来亲自验证这一点。

顺序使用动作序列会导致 (1) 运行时延迟，影响任务执行时间，以及 (2) 由于在很大程度上是开环行动而缺乏响应性。异步推理通过将动作预测与动作执行解耦来缓解这两个限制。我们在 SmolVLA 中引入了异步推理，发现它可以在任务成功率相当的情况下，将任务完成时间缩短约 2 倍。

具体来说，我们设计了一个包含两个组件的系统，其中策略推理和动作执行在两个不同的进程中执行，甚至可以在通过网络连接的两台不同机器上进行。

• 一个 PolicyServer（策略服务器），托管在加速硬件上，能够使用比真实机器人上分配的更多计算资源来运行推理。
• 一个 RobotClient（机器人客户端）将接收到的动作入队并执行它们，同时计算下一个动作序列。

PolicyServer 和 RobotClient 之间的通信依赖于 gRPC，它保证了比同类 REST API 快约 5 倍的性能。所有这些的结果是一个从不等待推理的机器人。

异步推理，突出显示：(1) 客户端发送第一个观察进行推理，不久后接收到第一个动作序列；(2) 客户端在尚未用尽当前动作序列时发送另一个观察进行处理；(3) 客户端接收到一个更新的动作序列，并将其与之前正在执行的剩余部分进行聚合。

---

1. 为什么顺序推理不够用

假设一个策略 $\pi$ 将当前观察 $o_t$ 映射到未来 $H$ 个动作的序列。形式上，$\pi : \mathcal{O} \;\mapsto\; \mathcal{A}, \mathbf{A}_t = \begin{pmatrix} a_{t}, & a_{t+1}, & \dots & a_{t+H} \end{pmatrix} = \pi(o_t)$。

因此，传统的控制循环将包括以下步骤：

1. 捕获 $o_t$。
2. 运行 $\pi(o_t)$ 以获得 $\mathbf{A}_t = \pi(o_t)$。
3. 将 $\mathbf{A_t}$ 入队，并开始从队列中弹出动作来执行。
4. 如果队列为空，等待 $\mathbf{A}_{t+H}$，否则重复步骤 3。

在步骤 2 中，机器人处于空闲状态。延迟随着模型大小的增加而增长（并且模型会随着时间的推移而变得越来越庞大），并可能迅速占据交互时间（通常约为 1/fps），如下面的视频所示（来自我们的 Discord 社区 🤗）

这直接导致 (1) 任务完成时间性能降低——机器人需要等待下一个动作序列被计算出来——以及 (2) 响应性降低，原因在于 (2.1) 在有动作可用时进行大范围的开环行动，以及 (2.2) 在等待下一个动作序列时完全空闲。

（左）顺序推理，突出显示了空闲时段。(右)选择动作的时间，显示了当本地队列耗尽时触发推理导致的峰值（在 2021 款 MacBook Pro 上使用 ACT 模型，推理延迟约为 100 毫秒，即 30fps 下约 3 帧）。

---

2. 异步推理，简而言之

我们的系统通过重叠计算和执行来消除空闲期

1. RobotClient 将最新的观察流式传输到 PolicyServer。
2. 当服务器执行推理时，客户端执行当前队列中的动作。
3. 新动作到达后，会被合并到队列中，循环继续。

关键思想是，机器人已经知道接下来几个时间步该做什么，因此它可以在服务器计算新动作时继续移动。

异步推理通过解耦这两个过程，在时间上将当前动作序列的执行与下一个动作序列的计算重叠起来，这两个过程甚至可以在通过网络连接的完全不同的机器上运行。

这带来了更紧凑的控制循环，和一个从不等待推理的机器人。反过来，这使得任务完成时间缩短了约 2 倍，任务成功率相当，并且由于更紧凑的循环，控制更具适应性（见下方视频）。

---

3. 系统架构

因为 gRPC 基于 HTTP/2 并使用协议缓冲区，它能实现低延迟的二进制消息传递和开箱即用的双向流，这反过来帮助我们维持更紧凑的控制循环和低于 100ms 的往返延迟（在我们的本地网络上，并将 SmolVLA 托管在 NVIDIA RTX 4090 上）。

RobotClient 在机器人上运行，并通过 gRPC 将观察数据流式传输到 PolicyServer。PolicyServer 准备接收到的观察数据以进行推理，并将一个动作序列发回给 RobotClient。

机器人客户端

从客户端的角度看，观察数据根据本地队列状态流式传输到服务器。传入的动作序列在重叠部分与当前可用的动作队列进行聚合。

RobotClient 维护一个本地动作队列，并遵循一个简单而有效的策略：当队列长度低于可配置的阈值时发送新的观察数据（在 SmolVLA 论文中为 \(g\)，在代码中为 chunk_size_threshold）。这个阈值以最大序列大小的比例表示，作为一个触发条件，平衡了计算负载和响应性。

客户端根据本地队列状态将观察数据流式传输到服务器。

从客户端的角度来看，过程如下：

1. 队列监控：客户端持续监控其动作队列长度与序列大小阈值参数。当队列低于此阈值时，它会发出信号，表示应发送新的观察数据进行处理。

2. 观察数据流式传输：一旦满足阈值条件，客户端捕获当前观察数据并通过 gRPC 将其流式传输到 PolicyServer。至关重要的是，观察数据是流式传输的，而不是通过一元 RPC 发送的，因为它们通常会超过 4MB 的最大消息大小（高分辨率的多个摄像头捕获会导致这种情况）。

3. 动作序列聚合：当从服务器收到新的动作序列时，客户端会将其与当前队列中任何剩余的动作在重叠部分进行合并。这就是自定义聚合器发挥作用的地方，它们以不同方式处理当前和传入序列之间的重叠部分。目前，我们支持通过指定一个自定义的 aggregate_fn(chunk1: torch.Tensor, chunk2: torch.Tensor) -> torch.Tensor 函数来灵活地聚合序列，该函数对每个重叠的时间步调用，并且可以由用户提供。重叠部分（图中以浅蓝色显示）需要仔细处理。我们可以设计不同的聚合策略：

   • 替换：简单地用新的预测替换重叠的动作
   • 加权混合：使用时间权重组合重叠动作（较近的动作获得更高的权重）

这个系统是高度可配置的，因为可以根据网络延迟、模型推理时间和期望的响应性来调整序列大小阈值。较低的阈值意味着更频繁的更新（和更高的计算成本），而较高的阈值则以可能导致队列枯竭为代价减少了通信开销。最后，我们通常在一个线程中从 PolicyServer 接收动作，并在另一个线程中执行它们。这使得客户端在一个单独的线程中监听传入的序列，而不会阻塞执行，并始终消耗当前序列，直到新的序列完全可用。

策略服务器

在收到来自 RobotClient 的观察数据后，PolicyServer 会对接收到的观察数据进行必要的清理，使其准备好进行推理。这个过程如下图所示：

在服务器上运行的观察清理流程，突出显示了与 (1) 键匹配、(2) 预处理和 (3) 推理准备相关的三个主要步骤。

一旦观察数据准备好，就会与上一次用于推理的观察数据进行比较。这样可以避免陷入一个循环，即处理非常相似的观察数据，从而触发不必要的推理和执行相似的动作（这反过来又会导致处理非常相似的观察数据）。我们根据关节空间的相似性来比较观察数据，这为我们提供了一种近似且快速的方法来测量机器人的变化。显然，这个指标不适应环境的动态变化（物体改变位置或施加干扰），但我们发现它在大多数情况下是一个很好的权衡，并且在避免不必要的推理和状态崩溃方面非常有效。至关重要的是，RobotClient 保留了对给定观察数据是否必须处理的控制权，以避免死锁。客户端发送并标记为 must_go=True 的观察数据无论相似性度量如何都会被处理。

策略工作流，其中传入的观察数据与上一次用于推理的数据进行比较，只有在足够不同或标记为 `must_go` 时才进行处理。

最后，为确保 PolicyServer 始终处理最新的可用观察数据，我们会阻塞传入的观察数据，直到前一个观察数据被成功处理。为此，我们在 PolicyServer 上利用队列来确保传入的观察数据在服务器准备好处理它们之前不会入队（见下文）。

客户端每 1/fps 秒 ping 一次服务器，但在前一个观察数据成功处理之前，观察数据不会入队进行处理。

---

4. 分析异步推理

就所有实际目的而言，在异步推理中有两个重要的时间尺度：

• 环境步长 $\texttt{environment\_dt} = 1/\texttt{fps}$，描述机器人执行一个动作的速度。
• 推理延迟 $\texttt{inference\_time}$：前向传播 + 网络往返时间。我们可以假设网络往返时间相对于策略推理时间可以忽略不计，尽管这可能不适用于所有设置。

重要的是，比率 $c = \frac{\texttt{environment\_dt}}{\texttt{inference\_time}}$ 会导致不同的行为

• $c \ll 1$：环境演化速度快于推理速度。在这种情况下，队列会很快清空，我们会退化为顺序控制。
• $c \ge 1$：服务器能跟上。队列总是（几乎）满的。

关键是，$c$ 影响在任何给定时间队列中可用动作的数量。为了避免上述顺序极限控制，可以：

1. 为策略服务器使用更多计算资源，将服务器托管在 GPU 上，通过分配更多计算资源来减少 $\texttt{inference\_time}$。
2. 更频繁地向服务器发送观察数据，当队列长度 $k$ 低于其最大大小的一部分 $g = k/H$ 时发送新的观察数据。
   • $g=0$ 再现了顺序推理（队列为空，等待）。
   • $g=1$ 每个时间步都发送一个观察数据（最大计算量，最小延迟）。

实验（见下图）表明，当发送的观察数据未被过滤掉时（它们都是必须处理的），$g\approx0.7$ 提供了一个很好的权衡。我们建议设置 $g=0.5$ 并遵循我们的文档来根据您的需求调整此参数。

在任何给定时间队列中可用动作的数量，作为 g 的函数。g 值越大，更新越频繁，计算成本也越高。g 值接近 0 则再现了顺序推理（队列为空，等待）。在我们的实验中，我们发现 g~0.7 是一个很好的权衡。

---

5. 在您的设置中使用异步推理

异步推理是提高机器人策略性能的一种简单而有效的方法。在我们使用 SmolVLA 的实验中，异步推理在任务成功率相当的情况下，任务完成时间提速约 2 倍，并且由于更紧凑的循环，控制更具适应性。

要使用异步推理运行您的策略，您只需按照我们的教程，并使用您自己的自定义参数（例如，策略路径或动作序列大小阈值）。异步推理支持支持动作序列分块的策略！

---

结论

我们介绍了异步推理，这是一种简单而有效的方法，可以提高机器人策略的性能。在我们使用 SmolVLA 的实验中，异步推理在任务成功率相当的情况下，任务完成时间提速约 2 倍，并且由于更紧凑的循环，控制更具适应性。

我们很高兴能与社区分享这项工作，并期待看到它如何被用来提高机器人策略的性能。我们欢迎在 huggingface/lerobot 上提交 PR 来改进和扩展异步推理框架，并乐于在我们的 Discord 社区 🤗 中进一步讨论。