Training and Evaluating Diffusion Policies with Long Context Lengths

📄 arXiv: 2606.16447v1 📥 PDF

作者: Abhinav Agarwal, Adam Wei, Taylan Kargin, Michael Zeng, Cole Becker, Arif Kerem Dayi, Pablo Parrilo, Asuman Ozdaglar, Russ Tedrake

分类: cs.RO, cs.AI

发布日期: 2026-06-15


💡 一句话要点

提出长上下文长度的扩散策略以解决模仿学习中的记忆问题

🎯 匹配领域: 支柱一:机器人控制 (Robot Control) 支柱二:RL算法与架构 (RL & Architecture)

关键词: 模仿学习 长上下文 策略训练 去噪网络 机器人操作

📋 核心要点

  1. 现有模仿学习方法在处理需要记忆的任务时表现不佳,容易陷入重复失败的动作。
  2. 论文提出了一种新的训练算法,能够在多个上下文长度下联合训练策略,降低样本复杂度。
  3. 实验结果表明,使用UNet+Cross-Attention的策略在多种任务中实现了高成功率,验证了长上下文的有效性。

📝 摘要(中文)

模仿学习使得机器人能够从RGB观察中实现高灵活性的操作。然而,现有策略通常仅基于短期观察历史进行决策,导致无法处理需要记忆的任务。本文首次详细研究了模仿学习中上下文长度的影响,发现简单地增加上下文长度并不如文献中所述那样脆弱。通过适当的条件方法和去噪骨干(UNet+Cross-Attention),单任务策略在多种任务中实现了高成功率。我们还提出了一种联合训练算法,进一步降低了长上下文学习的样本复杂度,并重新评估了一些先前提出的长上下文模仿学习解决方案。

🔬 方法详解

问题定义:本文旨在解决模仿学习中策略对短期观察历史的依赖,导致无法有效处理需要长时间记忆的任务。现有方法在面对复杂任务时,容易陷入重复执行失败动作的困境。

核心思路:论文的核心思路是通过增加上下文长度来提升策略的记忆能力,并提出了一种联合训练算法,以便在多个上下文长度下共同优化策略,从而降低样本复杂度。

技术框架:整体架构包括数据收集、策略训练和评估三个主要阶段。在训练阶段,采用UNet与Cross-Attention的组合作为去噪骨干,确保策略能够有效利用长上下文信息。

关键创新:最重要的技术创新在于首次系统性地研究了上下文长度对模仿学习策略性能的影响,提出的联合训练算法显著提高了长上下文学习的效率。

关键设计:在参数设置上,采用了适当的条件方法来处理上下文信息,并设计了特定的损失函数以优化策略性能。网络结构上,结合了UNet和Cross-Attention,以增强模型对长上下文的处理能力。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示,使用UNet+Cross-Attention的单任务策略在多种任务中实现了高成功率,成功率显著高于传统方法,验证了长上下文长度的有效性和必要性。

🎯 应用场景

该研究的潜在应用领域包括高灵活性机器人操作、复杂任务的自动化执行以及人机协作等。通过提升机器人在长上下文任务中的表现,能够显著提高其在实际应用中的适应性和效率,推动智能机器人技术的发展。

📄 摘要(原文)

Imitation learning has enabled highly-dexterous robotic manipulation from RGB observations. Policies trained with these methods, however, typically condition robot actions on only a short history of observations. These policies cannot solve tasks that require memory and can get stuck repeatedly executing the same failing motions. In this work, we first benchmark policy performance as context length is incrementally increased from short to long, across a spectrum of tasks with varying local stability and memory requirements, and in multiple data regimes. To our knowledge, this is the first study to investigate context length in imitation learning at this level of detail. Our results challenge prior claims: naively scaling context length is not as brittle as advertised in literature. With an appropriate conditioning method and denoising backbone (UNet+Cross-Attention), single-task policies achieve high success rates on many tasks in the usual data regime even with naive scaling. Next, we propose a training algorithm to jointly train policies at multiple context lengths, further reducing the sample complexity of long-context learning. Finally, we apply our findings to re-evaluate some previously proposed solutions to long-context imitation learning.