消除推理阶段的额外开销！pi团队提出训练新思路

Tom Hardy

1103人浏览 · 2025-12-18 13:54:08

Tom Hardy · 2025-12-18 13:54:08 发布

实时动作分块（RTC）技术能够支持视觉-语言-动作模型（VLAs）生成平滑且具备响应性的机器人轨迹，其核心机制是异步预测动作块，并通过推理时补绘的方式，结合已执行的动作完成决策。然而，这种补绘方法会产生额外计算开销，进而增加推理延迟。

针对这一问题，pi团队提出了一种简便的替代方案：在训练阶段模拟推理延迟，直接基于动作前缀实现条件约束，从而完全消除推理阶段的额外开销。该方法无需修改模型架构或机器人运行时系统，仅需新增少量代码即可完成部署。在仿真实验中，研究发现当推理延迟较高时，训练阶段实时分块的性能要优于推理阶段实时分块；在基于 $\pi_{0.6}$ 视觉-语言-动作模型开展的纸箱搭建和浓缩咖啡制作两项真实场景任务实验中，本工作进一步验证：训练阶段实时分块不仅能保持与推理阶段实时分块相当的任务完成度和执行速度，还具备更低的计算成本。上述结果表明，在机器人实时控制场景中，训练阶段动作条件约束可作为推理时补绘的实用即插即用替代方案。

论文标题：Training-Time Action Conditioning for Efficient Real-Time Chunking
论文链接：https://arxiv.org/pdf/2512.05964

原文链接：消除推理阶段的额外开销！pi团队提出训练新思路

出发点与背景介绍

与聊天机器人或搜索引擎不同，具身智能体必须具备实时运行能力。智能体动作与外部环境间的反馈闭环决定了其必须拥有快速响应性——就像人类运动员一样，智能体无法在外界环境不断变化的过程中“停下来思考”。但随着前沿模型的规模持续扩大，这一要求的实现难度也越来越高。这一点在机器人学习领域表现得尤为明显：如今，参数规模达数十亿的视觉-语言-动作模型（VLAs）正被越来越多地用于高频率控制机器人，以完成各类灵巧操作任务。而当模型推理延迟达到数十至数百毫秒时，如何生成既平滑又具备响应性的动作轨迹，就成了一项不小的挑战。

实时动作分块技术整合了动作分块、流匹配和推理时补绘等技术思路，为解决上述问题提供了一套方案。在该技术框架下，动作块的预测过程是异步进行的——即在当前动作块执行的同时，就开始生成下一个动作块。为保证动作块之间的连续性，每次生成新动作块时，都会基于此前已预测的动作构成的固定前缀，并对剩余部分进行补绘。但遗憾的是，实时动作分块所采用的推理时补绘方法会引入额外计算开销，进而产生延迟，在一定程度上削弱了实时执行框架的设计初衷。此外，从实验数据来看，推理时补绘在应对高推理延迟时的能力也存在本质局限。

为此，本工作对实时动作分块技术进行了改进，提出一种在训练阶段模拟推理延迟的补绘方法，以此消除推理阶段的计算开销。该方法可直接替代推理阶段实时分块方案：既无需改动模型架构与机器人运行时系统，也仅需新增少量代码即可部署。在仿真基准测试中，研究发现高推理延迟场景下，训练阶段实时分块的性能优于推理阶段实时分块；在真实场景中，本工作证明即便是未经过动作前缀条件约束预训练的基础模型，也可通过微调实现训练阶段实时分块的部署。将该方法应用于 $\pi_{0.6}$ 视觉-语言-动作模型后，在纸箱搭建和浓缩咖啡制作这两项高复杂度任务中，模型展现出了优于推理阶段实时分块的性能。

一些预备知识

这里采用与实时分块方法相同的问题设定：首先定义一个动作分块策略，该策略可基于当前观测输出一段动作块，其中动作块包含未来多个时刻的连续动作，观测为机器人感知到的环境信息，代表控制器的时间步。将动作块包含的动作步数称为预测时域，在推理阶段，每个动作块会被执行不超过预测时域步数的动作，这一实际执行步数被称为执行时域。

为了考量模型推理带来的延迟，这里定义推理延迟为以控制器时间步为单位的延迟时长。若在时间步t启动推理，那么生成的动作块要到时间步t+d时才能就绪，因此前d个动作无法被即时执行。不过，只要推理延迟不超过预测时域与执行时域的差值，这些前d个动作就可以对应上一个动作块中尚未执行完毕的动作，可在等待当前动作块生成的过程中先行执行。将上一个动作块中与当前动作块存在重叠的这d个动作称为动作前缀。

我们采用条件流匹配方法训练策略，该方法的核心是最小化特定损失。在推理阶段，可通过将神经网络从初始时间步积分至最终时间步，来生成符合数据集分布的动作块样本。

训练时动作条件约束

推理时实时分块（RTC）会基于伪逆引导的推理时补全方法，将策略模型与动作前缀（图1中红色部分）进行关联。为了提升动作块之间的连续性，推理时实时分块还会对所有重叠动作施加约束，且对超出前缀范围的动作采用指数衰减权重（图1中黄色部分），这一机制在实时分块技术中被称为“软掩码”。尽管伪逆引导具备较强的灵活性，能够实现软掩码功能，但它要求在每一个去噪步骤中，都通过反向传播计算向量-雅可比乘积，这会产生额外的计算开销。

图1展示了两个相互重叠的动作块示意图。其中，t到t-d时间段内取自上一个动作块的d个动作即为动作前缀（红色区域）。从示意图中可以清晰看出，要实现有效的动作前缀，必须满足约束条件 $t + d \leq t - s + H$ ，即 $d \leq H - s$ 。需要注意的是，推理时实时分块会利用全部 $H - s$ 个重叠动作（红色与黄色区域）来引导当前动作块的生成，而训练时实时分块仅会采用前d个动作（红色区域）。

本工作的核心思路是，通过在训练阶段模拟推理延迟，直接让策略模型学习动作前缀的条件约束。这种方式虽然不具备推理时补全方法那样的灵活性，但能完全消除推理阶段的计算开销。具体来说，我们可以学习条件分布 $p(A_{t+d:H} | o_t, A_{t:t+d})$ ，其中 $A_{t:t+d}$ 为动作前缀（图1中红色区域）， $A_{t+d:H}$ 为动作后缀（图1中黄色与绿色区域），二者均来自同一真实动作块。对于多数标准策略模型架构，实现这一机制仅需进行3处微小修改：

图2展示了该条件约束架构在标准扩散transformer（如 $\pi_{0.6}$ 动作专家模型）上的应用方式。我们会始终输入真实、无噪声的前缀动作，同时训练模型对后缀动作进行去噪处理。不同动作标记对应的流匹配时间步存在差异，这一设计可向模型传递推理延迟的相关信息。

调整模型架构，允许不同动作时间步采用不同的流匹配时间步。对于采用adaLN-zero机制实现流匹配时间步条件约束的类扩散transformer架构，这一修改十分简便——只需让不同标记对应的缩放、偏移和门控参数独立即可，且不会增加模型可学习参数的数量。
前缀部分采用真实无噪声动作，并将其对应的流匹配时间步设为1；后缀部分则保持原有设置不变。这一操作能让模型基于真实动作前缀完成训练，同时仅对后缀动作执行去噪学习。
对损失函数进行掩码处理，仅针对后缀对应的输出计算损失。

实际应用中，由于无法提前预知精确的推理延迟，且真实场景中的推理延迟可能存在波动，因此会在训练阶段对延迟d进行随机采样。

经过上述修改后，动作生成模块可接收动作前缀 $A_{t:t+d}$ 和延迟 $d$ 作为输入，并输出动作后缀 $A_{t+d:H}$ 。由此，该模块能够与推理时实时分块的动作生成组件实现接口兼容，可作为其无缝替代方案。

实验对比分析

本研究的实验旨在对比训练时实时分块、推理时实时分块，以及基础的同步和异步基准方法的性能。仿真实验采用了与实时分块技术相关研究相同的动态Kinetix基准测试环境；真实场景实验则基于 $\pi_{0.6}$ 基础模型展开，选取了纸盒组装和意式浓缩咖啡制作这两项高精度、高难度任务，实验设置与 $\pi_{0.6}$ *的相关研究保持一致。

仿真实验结果

在动态Kinetix基准测试中，沿用相关研究的设置，训练了预测范围H=8的动作分块流策略模型，模型采用4层MLP-Mixer架构，在专家策略混合生成的数据集上训练32个轮次。每个数据点均基于2048次轨迹展开实验，测试的推理延迟范围为0（全闭环）至4（H=8时支持的最大延迟）。基础异步方法和推理时实时分块均使用未经动作前缀条件约束训练的同一份模型检查点，该模型完成了32个轮次的常规训练。

训练时实时分块的模型则从第24个轮次开始恢复训练，通过8个轮次的微调实现动作前缀条件约束的学习，以此保证所有方法的训练计算量保持一致。训练过程中，我们从集合{0,1,2,3,4}中按指数衰减权重采样延迟值，这是因为我们发现高延迟场景所需的训练监督更少。若为每个延迟值单独训练模型检查点，理论上可取得更优的实验效果。

实验结果如图3所示。发现当推理延迟达到2及以上时，训练时实时分块的性能显著优于推理时实时分块，且二者的性能差距会随延迟增大而进一步扩大。这一现象的原因在于，随着前缀动作数量增加，补全算法需要耗费更多计算资源来生成连贯的后缀动作，此时依赖模型雅可比矩阵线性化的纯推理时算法鲁棒性不足，而训练时算法的表现则更为稳定。在延迟为0和1的场景下，训练时实时分块的性能略有劣势，这可能是因为该方法无法对所有动作进行完整的训练监督，即学习生成前两个动作的训练计算量相对不足。

图3为仿真实验结果：固定执行范围 $s = ma x (d, 1)$ 时，推理延迟与任务成功率的关系。当推理延迟≥2时，训练时实时分块的表现优于推理时实时分块。每个数据点均基于2048次实验，阴影区域为95%威尔逊得分区间。

真实场景实验结果

真实场景实验采用 $\pi_{0.6}$ 基础模型，测试任务为 $\pi^*_{0.6}$ 相关研究中的意式浓缩咖啡制作和纸盒组装（任务示意图见图4）。同步基准方法与推理时实时分块使用同一份模型checkpoint，训练时实时分块则使用独立的checkpoint，两个checkpoint均基于基础模型在目标任务上微调8000个梯度步，批次大小设为512。训练阶段，在0至10的范围内均匀采样延迟值，这一设置可支持50Hz机器人实现200ms的最大推理延迟。评估阶段，模型在远端H100服务器上执行推理，去噪步骤数设为5，训练时实时分块的端到端平均延迟为108ms（对应d≈5），推理时实时分块的端到端平均延迟为135ms（对应d≈7）。