让轨迹扩散模型一步生成,提速10倍以上且保持多样性。
Amortizing Trajectory Diffusion with Keyed Drift Fields
- 用条件感知的键空间距离指导轨迹更新,避免平均化
- 训练时用停止梯度目标模拟迭代优化,实现一步推理
- 适合实时控制场景,尤其算力受限的硬件部署
基于扩散的轨迹规划器可在离线强化学习中生成丰富多样的动作序列,但其迭代去噪过程带来显著推理开销,导致闭环规划在计算资源受限时速度缓慢。本文研究如何在单步推理下实现类似扩散的轨迹规划能力,同时保留生成多样化候选计划并根据当前状态进行条件规划的能力。关键观察发现:当使用相似性度量对齐生成轨迹与数据集时,若未受约束的未来维度占主导,朴素分布匹配目标会导致轨迹趋近平均值,丧失动作多样性,产生近似静止行为。核心洞察是:条件生成规划需要一种条件感知的邻域概念——轨迹更新应基于反映条件的紧凑键空间距离,同时在完整轨迹空间中应用更新。基于此,提出键控漂移策略(KDP),一种通过漂移场目标训练的一步轨迹生成器,能将生成轨迹吸引至与条件匹配的数据集窗口,并排斥附近已生成样本,利用停止梯度漂移目标将迭代优化过程摊入训练。推理时,该策略可单次前向传播生成完整轨迹窗口。在标准强化学习基准和真实硬件部署中,KDP均实现强性能,且规划延迟显著低于扩散采样。
原文摘要 · Abstract (English)
Diffusion-based trajectory planners can synthesize rich, multimodal action sequences for offline reinforcement learning, but their iterative denoising incurs substantial inference-time cost, making closed-loop planning slow under tight compute budgets. We study the problem of achieving diffusion-like trajectory planning behavior with one-step inference, while retaining the ability to sample diverse candidate plans and condition on the current state in a receding-horizon control loop. Our key observation is that conditional trajectory generation fails under naïve distribution-matching objectives when the similarity measure used to align generated trajectories with the dataset is dominated by unconstrained future dimensions. In practice, this causes attraction toward average trajectories, collapses action diversity, and yields near-static behavior. Our key insight is that conditional generative planning requires a conditioning-aware notion of neighborhood: trajectory updates should be computed using distances in a compact key space that reflects the condition, while still applying updates in the full trajectory space. Building on this, we introduce Keyed Drifting Policies (KDP), a one-step trajectory generator trained with a drift-field objective that attracts generated trajectories toward condition-matched dataset windows and repels them from nearby generated samples, using a stop-gradient drifted target to amortize iterative refinement into training. At inference, the resulting policy produces a full trajectory window in a single forward pass. Across standard RL benchmarks and real-time hardware deployments, KDP achieves strong performance with one-step inference and substantially lower planning latency than diffusion sampling. Project website, code and videos: https://keyed-drifting.github.io/
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。