通过预测历史动作提升长序列决策能力,显著改善机器人策略性能。
Learning Long-Context Diffusion Policies via Past-Token Prediction
- 引入过去动作预测任务,强制模型捕捉历史与未来的依赖关系。
- 在真实和仿真任务中,长上下文策略性能提升3倍,训练速度加快10倍以上。
- 适合需要长期记忆的机器人决策、自动驾驶等复杂控制场景。
对长序列观测与动作进行推理是许多机器人任务的核心需求。然而,从示范中学习有效的长上下文策略仍具挑战性:随着上下文长度增加,训练因内存需求上升而变得昂贵,且策略性能常因虚假相关性下降。现有方法通常通过截断上下文来规避问题,但会丢弃关键的历史信息。本文提出一种新方法,显式正则化过去信息的保留。我们重新审视模仿学习中的复制问题,发现当前扩散策略存在相反缺陷:过度依赖先验动作,却忽视历史与未来动作间的必要关联。为此,我们引入过去动作预测(PTP)作为辅助任务,让策略同时预测过去和未来的动作标记。该正则化显著提升了策略头的时序建模能力,且对视觉表征依赖极少。基于此,我们进一步设计多阶段训练策略:先用短上下文预训练视觉编码器,再用缓存的长上下文嵌入微调策略头。该策略在保留PTP优势的同时,大幅降低内存与计算开销。最后,我们将PTP扩展为测试时的自验证机制,使策略在推理阶段评估并选择与历史动作一致的候选动作。在四个真实世界和六个模拟任务上的实验表明,本方法使长上下文扩散策略性能提升3倍,训练速度加速超过10倍。
原文摘要 · Abstract (English)
Reasoning over long sequences of observations and actions is essential for many robotic tasks. Yet, learning effective long-context policies from demonstrations remains challenging. As context length increases, training becomes increasingly expensive due to rising memory demands, and policy performance often degrades as a result of spurious correlations. Recent methods typically sidestep these issues by truncating context length, discarding historical information that may be critical for subsequent decisions. In this paper, we propose an alternative approach that explicitly regularizes the retention of past information. We first revisit the copycat problem in imitation learning and identify an opposite challenge in recent diffusion policies: rather than over-relying on prior actions, they often fail to capture essential dependencies between past and future actions. To address this, we introduce Past-Token Prediction (PTP), an auxiliary task in which the policy learns to predict past action tokens alongside future ones. This regularization significantly improves temporal modeling in the policy head, with minimal reliance on visual representations. Building on this observation, we further introduce a multistage training strategy: pre-train the visual encoder with short contexts, and fine-tune the policy head using cached long-context embeddings. This strategy preserves the benefits of PTP while greatly reducing memory and computational overhead. Finally, we extend PTP into a self-verification mechanism at test time, enabling the policy to score and select candidates consistent with past actions during inference. Experiments across four real-world and six simulated tasks demonstrate that our proposed method improves the performance of long-context diffusion policies by 3x and accelerates policy training by more than 10x.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。