拥有记忆的机器人: 利用过去标记预测改进扩散策略

在机器人领域，记忆就是一切。想象一下，如果你不记得是否已经加了糖还在煮咖啡，或者试图开门却想不起刚才试了哪把钥匙。对于人类来说，这些时间依赖性——即五秒前发生的事情与接下来应该发生的事情之间的关系——是直觉性的。但对于机器人来说，这是一项巨大的计算和算法难题。

近年来, 扩散策略 (Diffusion Policies) 已成为机器人模仿学习的主流方法。它们通过学习人类演示来生成平滑、复杂的行为。然而，存在一个奇怪的悖论: 虽然我们希望机器人“记住”更多的历史信息以做出更好的决策，但向这些模型输入更长的数据序列往往会导致它们表现更差，而且训练成本也会变得极其高昂。

这篇博客文章将深入剖析一篇引人入胜的新论文: “Learning Long-Context Diffusion Policies via Past-Token Prediction” (通过过去标记预测学习长上下文扩散策略) 。 研究人员提出了一种巧妙的正则化技术和训练流程，不仅让机器人能够有效地利用长历史信息，还将训练速度提高了 10 倍以上。

展示所提框架概览，显示性能提升和训练加速。

如图 1 所示，与朴素方法相比，所提出的方法 (绿线) 在更短的训练时间内实现了显著更高的性能。让我们深入探讨为什么会存在这个问题以及他们是如何解决的。

背景: 机器人记忆的悖论

模仿学习的挑战

模仿学习 (Imitation Learning, IL) 是训练机器人模仿专家行为的过程。标准的 IL 策略接收一个观测值 (\(o_t\)) ——通常是来自摄像头的图像——并预测一个动作 (\(a_t\)) 。

然而，许多任务是非马尔可夫 (non-Markovian) 的，这意味着当前的图像并不能说明全部情况。

  • 遮挡: 如果机械臂挡住了摄像头对物体的视线，机器人需要记住物体之前在哪里。
  • 多阶段任务: 如果机器人正在打包盒子，它需要知道哪些物品已经打包好，才能决定接下来抓取什么。

为了解决这个问题，研究人员通常会将观测历史 (\(o_{t-k} \dots o_t\)) 输入模型。理论上，历史信息越多，上下文越好，性能也就越强。但在实践中，这往往导致失败。

“复读机”与“健忘症”问题

历史上，旧的基于回归的策略深受复读机问题 (Copycat Problem) 的困扰。当给定动作历史时，模型会学习到一种虚假的相关性: “对下一个动作的最佳预测就是简单地重复上一个动作。”机器人只会重复它的上一步动作，而忽略环境中视觉的变化。

然而，这篇论文的作者发现，现代扩散策略面临的问题截然相反。它们非但不过度依赖历史，反而往往完全忽略历史。

回归策略与扩散策略在时间动作依赖性方面的比较。

图 2 展示了这一现象。Y 轴代表“动作可预测性比率 (Action Predictability Ratio) ”——本质上是模型在多大程度上像人类专家那样依赖过去的动作。

  • 回归 (白色柱状图) : 这些模型过度依赖过去 (比率 > 1) ，导致复读机行为。
  • 扩散 (绿色斜线柱状图) : 这些模型过少依赖过去 (比率 < 1) 。

尽管被输入了很长的图像历史，扩散模型却未能捕捉到过去和未来动作之间的关键依赖关系。它本质上患有“健忘症”，将输入的历史信息视为噪声，而不是结构化的事件序列。

核心方法: 过去标记预测 (PTP)

为了治愈这种“健忘症”，研究人员引入了过去标记预测 (Past-Token Prediction, PTP) 。

这个想法简单却深远。标准的策略被训练用于预测动作的未来轨迹 (\(a_t, a_{t+1}, \dots\)) 。PTP 修改了损失函数，使得策略必须根据当前的上下文同时预测过去的动作 (\(a_{t-k}, \dots, a_{t-1}\)) 。

为什么这行得通？

如果一个神经网络被迫根据其内部状态准确地重建它刚刚做过的事情，那么它在数学上就被迫编码并保留有关历史的信息。如果它的损失函数依赖于正确预测过去的动作，它就无法忽略过去的观测。

过去标记预测的示意图，显示穿过策略的双向箭头。

如图 3 所示，策略头 (用 \(\pi\) 表示) 现在具有双向输出。它观察历史观测值，并生成一个从过去跨越到未来的序列。

形式上，这改变了预测目标。模型不再仅仅预测 \(\mathbf{a}_{t:t+l}\) (未来) ，而是预测:

\[ \hat { \mathbf { a } } _ { t - k : t + l } = \pi _ { \theta } \bigl ( \mathbf { o } _ { t - k : t } \bigr ) . \]

通过显式地正则化模型以保留过去的信息，PTP 确保了过去动作的“风格”和“策略”能够影响未来的动作，有效地弥合了历史上下文与未来规划之间的鸿沟。

解决计算瓶颈: 多阶段训练

在高维图像的长序列 (例如，每步 16 帧视频) 上训练模型是非常昂贵的。它会消耗大量的 GPU 显存 (VRAM) 并减慢训练速度，因为视觉编码器 (通常是 ResNet 或 Vision Transformer) 必须在每个训练步骤中处理历史中的每一帧。

作者观察到，PTP 的好处主要来自策略头 (序列建模部分) ，而不是视觉编码器 (图像处理部分) 。

基于此，他们提出了一种多阶段训练策略 , 大大减少了开销。

带嵌入缓存的多阶段训练概览。

图 4 分解了这个高效的流程:

  1. 阶段 1: 编码器训练 (短上下文) 首先，他们使用短上下文 (例如 2 帧) 训练一个标准策略。这既快又便宜。这里的目标仅仅是学习一个能够理解场景的优秀视觉编码器。

  2. 阶段 2: 特征缓存 一旦编码器训练完成，他们就将其冻结。他们让整个演示视频数据集通过这个冻结的编码器。他们不存储图像，而是将紧凑的嵌入 (特征向量) 存储到磁盘缓存中。

  3. 阶段 3: 策略训练 (长上下文) 最后，他们使用缓存的嵌入训练 PTP 策略头。由于视觉编码器不再运行，内存占用大幅下降。模型现在可以瞬间摄取预计算嵌入的长历史 (例如 16 步) 。

这种方法使他们能够比端到端方法快 10 倍地训练长上下文策略，而且不会牺牲性能。

显示特征缓存对训练速度影响的图表。

图 8 展示了这种效率。标准方法 (黑色虚线) 随着时间的推移几乎没有什么进展，因为每个 epoch 耗时太长。缓存方法 (黑色实线) 学习迅速，仅用一小部分时间就达到了很高的成功率。

智能推理: 自我验证

过去标记预测的好处不仅限于训练。在部署 (测试时间) 期间，PTP 赋予了机器人一项超能力: 自我验证 。

扩散策略是概率性的——它们从分布中采样动作。有时它们会采样到一个与机器人刚才的行为不一致的“糟糕”规划。因为 PTP 模型同时预测过去和未来的动作，我们可以使用对过去的预测作为质量检查。

逻辑如下:

  1. 机器人确切知道它在前几步采取了什么动作 (地面真值过去) 。
  2. 机器人生成 \(B\) 个不同的未来候选规划。每个规划也包含一个“重建的过去”。
  3. 机器人将每个候选者的“重建过去”与“地面真值过去”进行比较。
  4. 拥有最准确过去重建的候选者，很可能拥有最一致的未来规划。

测试时验证的图表，显示最佳动作序列的选择过程。

图 5 说明了这个选择过程。策略采样一批序列 (\(\mathcal{A}\)) :

\[ \mathcal { A } = \{ \hat { \mathbf { a } } ^ { ( 1 ) } , \ldots , \hat { \mathbf { a } } ^ { ( B ) } \} , \quad \hat { \mathbf { a } } ^ { ( i ) } \sim \pi _ { \boldsymbol { \theta } } ( \mathbf { o } _ { t - k : t } ) , \]

然后，通过最小化预测的过去动作 (\(\hat{a}_\tau\)) 与实际过去动作 (\(a_\tau\)) 之间的误差来选择最佳序列 (\(\hat{\mathbf{a}}^*\)) :

\[ \hat { \mathbf { a } } ^ { * } = \arg \operatorname* { m i n } _ { \hat { \mathbf { a } } \in \mathcal { A } } \sum _ { \tau = t - k } ^ { t - 1 } \| \hat { a } _ { \tau } - a _ { \tau } \| ^ { 2 } \]

这充当了一个运行时过滤器，在硬件执行之前丢弃幻觉或不一致的规划。

PTP 自我验证对成功率的影响。

如图 10 所示，这种推理时验证通过简单地采样更多候选者 (增加预算) 并挑选最佳者，在具有挑战性的任务上将成功率提高了约 5%。

实验与结果

强迫机器人预测过去真的能帮助它在未来表现得更好吗？结果令人信服。

仿真结果

研究人员在六个模拟任务上测试了该方法，包括“正方形 (Square) ”、“工具悬挂 (Tool Hang) ”以及特别需要记忆的复杂“长视距 (Long-Horizon) ”变体任务。

六个模拟任务中策略的比较。

图 9 比较了三种方法:

  1. 灰色: 无历史 (上下文长度 = 1 或 2) 。
  2. 红色: 无 PTP 的长历史 (标准扩散) 。
  3. 绿色: 带 PTP 的长历史 (我们的方法) 。

在标准任务 (如 Push-T) 上，PTP 的表现与最佳基线持平。然而，在长视距任务 (图表右侧) 上，差异简直是天壤之别。标准扩散策略 (红色) 几乎完全失败，因为它们无法随时间维持上下文。PTP 方法 (绿色) 则实现了接近完美的成功率。

此外，观察依赖性分析:

显示 PTP 对时间动作依赖性影响的图表。

图 6 证实，增加 PTP 监督 (在 x 轴上向右移动) 与更高的成功率直接相关。通过强迫模型遵循时间依赖性，机器人变成了一个更好的模仿者。

现实世界表现

团队还在执行诸如数勺粉末或替换物体等任务的真实机器人 (Franka Emika Panda 和 ALOHA) 上验证了该方法。如果没有记忆，这些任务是不可能完成的——如果你忘了你已经舀了几勺，你就失败了。

不同策略在现实世界任务上的比较。

图 11 展示了现实世界的结果。“无历史”和“无 PTP”基线表现非常挣扎 (15-20% 的成功率) 。PTP 方法跃升至超过 70% 的成功率 , 证明该技术足够稳健，可用于物理部署。

上下文长度的作用

最后，一旦启用了 PTP，添加更多的历史信息有帮助吗？

历史观测对 PTP 训练策略的影响。

图 7 显示，随着上下文长度的增加 (最多 16 帧) ，成功率攀升，特别是对于“长视距正方形 (Long-Horizon Square) ”任务 (橙色线) 。这证实了 PTP 成功解锁了利用长上下文的能力，将过去常常成为负担 (更多数据导致混乱) 的东西变成了资产。

结论

论文 “Learning Long-Context Diffusion Policies via Past-Token Prediction” 为诊断和修复深度学习中一个微妙的故障模式提供了一个大师级的范例。

主要收获如下:

  1. 诊断: 现代扩散策略倾向于忽略历史，无法捕捉时间一致性。
  2. 解决方案: 训练模型预测过去 , 迫使它对历史进行编码，从而改善其未来的预测。
  3. 效率: 通过缓存将视觉编码器与策略头解耦，使训练速度提高了 10 倍。
  4. 验证: 预测过去的能力使机器人能够在运行时自我验证其规划。

通过将过去视为监督信号而不仅仅是输入，这种方法使机器人能够执行以前扩散式模仿学习无法企及的复杂多阶段任务。事实证明，对于机器人来说，要想有效地前进，回头看会有所帮助。