⚡ electrical engineering

Amortizing Trajectory Diffusion with Keyed Drift Fields

该论文提出了一种名为“键控漂移策略”（KDP）的单步轨迹生成方法，通过引入条件感知的键空间距离和漂移场目标，成功将扩散模型的迭代推理过程摊销至训练阶段，从而在保持多模态规划多样性的同时，显著降低了强化学习中的推理延迟。

原作者： Gokul Puthumanaillam, Melkior Ornik

发布于 2026-03-17

📖 1 分钟阅读☕ 轻松阅读

原作者： Gokul Puthumanaillam, Melkior Ornik

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇文章介绍了一种名为 KDP (Keyed Drifting Policies，带键的漂移策略) 的新方法，旨在解决机器人“思考”太慢的问题。

为了让你轻松理解，我们可以把机器人做决策的过程想象成**“在迷宫里找路”**。

1. 以前的困境：慢吞吞的“试错大师”

想象一下，你让一个机器人去拿桌上的苹果。

旧方法（扩散模型 Diffusion）：就像是一个极其谨慎但动作缓慢的画家。
- 它先画出一团乱糟糟的墨迹（随机噪声）。
- 然后，它需要反复擦拭、修改、再擦拭、再修改（比如 20 次），每次只把画面修正一点点，直到最后画出一条完美的拿苹果路线。
- 问题：虽然画出来的路线很完美，但每次画都要花很长时间。如果机器人每走一步都要重新画一次（因为环境在变），它可能还没画完，苹果就被别人拿走了，或者它自己撞墙了。这就叫**“推理延迟太高”**。

2. 新方法的灵感：聪明的“一步到位”

作者问：能不能让机器人只画一笔，就直接画出完美的路线？

新目标：像经验丰富的老司机一样，看一眼路况，瞬间就能决定怎么打方向盘，不需要反复修改。
挑战：如果你直接训练一个模型让它“一笔成型”，它往往会犯傻。因为它会盯着那些还没发生的、不确定的未来（比如“如果我不小心撞墙了怎么办”），导致它画出的路线是所有可能路线的平均值——也就是**“原地不动”或者“像僵尸一样僵硬”**，因为它不敢冒险去走任何一条具体的路。

3. KDP 的核心魔法：带“钥匙”的导航员

为了解决这个问题，作者发明了一个叫 KDP 的方法。它的核心思想可以用一个**“带钥匙的图书馆”**来比喻：

比喻 A：错误的找书方式（旧方法）

假设你想找一本关于“在雨天开车”的书。

旧方法：你走进图书馆，不管书的内容，只看书的厚度和颜色（就像旧算法只看整个轨迹的长度和形状）。结果你发现，很多书虽然颜色厚度一样，但内容可能是“在晴天开车”或者“在沙漠开车”。你被这些不相关的书搞晕了，最后决定“不看了，就坐在地上吧”（这就是动作坍缩，机器人不动了）。

比喻 B：KDP 的“钥匙”机制（新方法）

KDP 给每本书配了一把**“钥匙”**（Key）。

钥匙是什么？ 就是当前的状态（比如：现在的天气是“雨天”，现在的车速是"60 码”）。
怎么找书？ 当你需要找“雨天开车”的路线时，KDP 不会看整本书，而是先用钥匙去开锁。它只把那些钥匙匹配（确实是雨天）的书挑出来。
漂移（Drift）： 在这些匹配的书里，它不是简单地取平均值，而是像磁铁一样：
1. 吸引（Attract）： 把生成的路线往那些“雨天”的真实成功案例上拉。
2. 排斥（Repel）： 把生成的路线从那些“太相似、太拥挤”的路线上推开，防止大家都走同一条死胡同（保持多样性）。

4. 训练过程：把“反复修改”变成“肌肉记忆”

这是 KDP 最厉害的地方。

传统做法：每次遇到新情况，都要现场反复修改（推理时慢）。
KDP 的做法：在训练阶段（机器人还在学校学习时），它模拟了“反复修改”的过程。
- 它先让机器人画一笔（生成一个粗糙路线）。
- 然后利用上面的“钥匙 + 磁铁”机制，告诉机器人：“嘿，往左偏一点，离那个成功的雨天案例近一点，离那个失败的案例远一点。”
- 关键点：它把这个“修正过程”直接**压缩（Amortize）**进了机器人的大脑（神经网络参数）里。
- 结果：等到真正上路（推理）时，机器人不需要再反复修改了。它看一眼钥匙（当前状态），大脑里直接输出修正好的完美路线。这就好比把“反复擦黑板”的过程，变成了“一次性写对”的肌肉记忆。

5. 实际效果：快如闪电，稳如泰山

作者在机器人上做了实验，效果惊人：

速度：以前画一张图要 20 步（耗时几百毫秒），现在只要1 步（耗时几毫秒）。机器人的反应速度提升了几十倍。
质量：在复杂的任务中（比如用机械臂穿针引线，或者无人机在拥挤房间里飞行），KDP 不仅快，而且比慢吞吞的旧方法更聪明、更灵活。
实时性：因为算得快，机器人可以每秒重新规划几十次。如果前面突然冒出一个人，它能瞬间调整路线；而旧方法可能还在算第一步，人已经撞上了。

总结

KDP 就像是一个把“反复思考”内化为“直觉”的超级机器人。

它不再需要在每次行动前都停下来慢慢“想”（迭代去噪），而是通过一种聪明的**“钥匙匹配”**机制，在训练时就学会了如何根据当前情况，直接生成一条既符合目标又多样化的完美路线。这让机器人从“慢吞吞的画家”变成了“反应神速的赛车手”，真正实现了在现实世界中的实时智能控制。

1. 研究背景与问题 (Problem)

背景：
基于扩散模型（Diffusion Models）的轨迹规划器在离线强化学习（Offline RL）中表现出色，能够生成丰富、多模态的动作序列。它们通常用于滚动时域控制（Receding-Horizon Control），即根据当前状态采样短期轨迹并执行第一个动作。

核心痛点：

推理延迟高： 传统的扩散规划器需要迭代去噪（Iterative Denoising），通常需要 $T$ 个步骤。在紧密的计算预算下，这导致闭环规划速度极慢，难以满足实时控制需求。
单步生成的困境： 现有的尝试减少去噪步数或蒸馏的方法，往往仍依赖迭代循环。如果直接训练一个“单步生成器”（One-step Generator）来匹配数据分布，在高维状态 - 动作序列中会遭遇模式坍塌（Mode Collapse）。
- 原因： 在无条件或 naive 的距离度量下，未受约束的未来维度（Future Dimensions）主导了距离计算，导致生成的轨迹倾向于“平均轨迹”，动作多样性丧失，控制器表现近乎静态。

研究目标：
如何在保持条件可控性（Conditioning）和样本多样性的同时，实现单步推理（One-step Inference）的扩散式轨迹规划，从而大幅降低延迟。

2. 核心方法论：键控漂移策略 (Keyed Drifting Policies, KDP)

作者提出了一种名为 KDP 的单步轨迹生成器，其核心思想是将迭代优化的过程“摊销”（Amortize）到训练阶段，并在推理时直接输出完整轨迹。

2.1 关键洞察：条件感知的邻域 (Conditioning-Aware Neighborhood)

问题： 在高维空间中，如果直接用整个轨迹窗口的距离来寻找“邻居”，未受约束的未来部分会掩盖当前状态（条件）的匹配度。
解决方案： 引入键空间（Key Space）。
- 键（Key）： 定义为轨迹的初始状态部分（即当前观测值 $c = s_0$ ）。
- 邻域构建： 在键空间中计算距离（ $d_\kappa$ ）来确定哪些数据轨迹是“邻居”（即与当前条件匹配）。
- 更新执行： 虽然邻居的选择基于键空间，但实际的轨迹更新（漂移）是在完整轨迹空间中进行的。

2.2 训练机制：键控漂移场 (Keyed Drift Field)

KDP 通过构建一个漂移场来指导生成器，该场包含两个作用：

吸引（Attraction）： 将生成的轨迹拉向与当前条件匹配的数据集窗口。
排斥（Repulsion）： 将生成的轨迹推开，远离当前模型生成的其他样本，以维持多样性。

具体步骤：

硬约束（Hard Conditioning）： 使用钳制算子（Clamping Operator）强制生成轨迹的初始状态等于当前观测值 $c$ 。
漂移计算：
- 计算生成样本 $x_i$ 与数据集样本 $y^+_j$ 和当前批次生成样本 $y^-_j$ 在键空间的距离。
- 基于键距离计算高斯核权重（Softmax），生成加权平均的“吸引原型” $\mu^+$ 和“排斥原型” $\mu^-$ 。
- 漂移向量 $V = \mu^+ - \mu^-$ 。
停止梯度目标（Stop-Gradient Target）：
- 构建目标轨迹 $\tilde{x}_i = C_c(\text{sg}(x_i + V))$ 。这里 $\text{sg}$ 表示停止梯度，意味着模型学习的是“如果进行一步迭代优化后的结果”，而不是直接优化 $x_i$ 本身。
- 通过最小化生成输出 $x_i$ 与目标 $\tilde{x}_i$ 之间的加权均方误差来训练生成器。
动作多样性保护： 在损失函数中，对动作维度（Action dimensions）赋予更高的权重，防止动作坍塌。

2.3 推理过程

在推理时，KDP 接收当前状态 $c$ ，并行采样 $K$ 个噪声向量。
通过单次前向传播生成 $K$ 个完整轨迹窗口。
使用评分器（Scorer）选择最佳轨迹，执行第一个动作。
计算优势： 将扩散模型的 $T$ 次串行计算转化为 $K$ 次并行计算（ $T$ 被摊销到训练中），极大地降低了延迟。

3. 主要贡献 (Key Contributions)

条件感知的邻域定义： 提出了一种新的训练信号构建方式，通过在键空间（条件空间）而非全轨迹空间计算相似度，解决了高维条件生成中的模式坍塌问题。
单步规划配方（One-step Planning Recipe）： 设计了一种基于“停止梯度漂移目标”的训练方法，将迭代 refinement 摊销到模型参数中，实现了无需迭代去噪的单步轨迹生成，同时保留了扩散规划器的接口。
端到端评估： 在标准离线 RL 基准（D4RL）和真实硬件（四旋翼无人机、机械臂）上进行了验证，证明了其在严格延迟预算下的实时闭环控制能力。

4. 实验结果 (Results)

4.1 仿真基准 (D4RL)

短视距控制（Locomotion）： KDP 在 HalfCheetah、Hopper 等任务上，性能与最先进的低 NFE（神经函数评估次数）生成方法相当，且显著优于需要 20 步去噪的 Diffuser。
长视距与稀疏奖励（Maze2D, AntMaze）：
- 在 Maze2D 中，KDP 在保持实时规划成本（0.03s/步）的同时，性能优于 Diffuser（3.8s/步）。
- 在 AntMaze 中，虽然任务极难，KDP 仍保持了与扩散模型相当的竞争力，且延迟极低。
灵巧操作（Adroit）： 在 Pen, Door, Hammer 等高精度任务中，KDP 不仅速度极快（毫秒级），而且得分显著高于 Diffuser。这表明单步漂移目标比在线迭代去噪更适合处理接触动力学。

4.2 硬件部署 (Real-World)

平台： Crazyflie 微型四旋翼（导航）和 SO-100 机械臂（抓取/堆叠）。
结果：
- 重规划频率（Replan Rate）： KDP 将重规划频率从 Diffuser 的 ~4Hz 提升至 38Hz（导航）和 17Hz（操作）。
- 延迟： 规划延迟从数百毫秒降低至 <10ms。
- 成功率： 在保持甚至略微提高任务成功率的同时，实现了实时的闭环控制。
- 意义： 证明了 KDP 生成的轨迹具有足够的局部一致性，能够应对现实世界中传感器和执行器的延迟与噪声。

4.3 消融实验 (Ablation)

无键空间（Full-window distances）： 性能崩溃，动作多样性丧失，验证了“条件感知邻域”的核心作用。
无排斥项（Attraction-only）： 导致模式坍塌，动作方差爆炸。
包含自负样本（Self-negatives）： 性能略有下降，验证了排除自样本掩码的必要性。

5. 意义与结论 (Significance)

打破延迟瓶颈： KDP 成功打破了扩散模型在实时控制中的延迟瓶颈，证明了通过训练时的“摊销优化”（Amortized Refinement），可以在推理时实现单步生成，同时保留扩散模型的多模态和条件控制优势。
实时闭环控制的新范式： 该方法使得基于生成式模型的规划器能够以高频（>30Hz）运行，这对于需要快速反应的真实机器人系统（如无人机避障、灵巧手操作）至关重要。
理论贡献： 揭示了在高维条件生成中，距离度量的选择（键空间 vs 全空间）对防止模式坍塌的决定性作用，为未来的条件生成模型设计提供了新的视角。

局限性： 在极度稀疏奖励和长视距的目标导向任务中，虽然 KDP 表现良好，但相比扩散模型仍有提升空间，可能需要结合更好的排序器或分层结构。

总结： 这篇论文提出了一种高效、实时的轨迹规划方法，通过巧妙的“键控漂移”机制，将扩散模型的强大建模能力与单步推理的低延迟完美结合，为机器人实时决策提供了强有力的工具。