Training and Evaluating Diffusion Policies with Long Context Lengths

2026-06-15

AI Summary

This work addresses the limitations of existing imitation learning methods, which often rely on short-horizon observations and struggle with tasks requiring long-term memory, frequently leading to repeated failures. Challenging the conventional wisdom that blindly extending context length inevitably degrades performance, the study systematically investigates the impact of context length on imitation learning efficacy. The authors propose a multi-context-length joint training algorithm combined with a denoising diffusion backbone network integrating UNet and cross-attention mechanisms, enabling efficient policy learning from RGB observations alone. Experiments demonstrate that, under standard data conditions, even a straightforward extension of context length significantly improves success rates across most tasks for single-task policies, thereby validating the feasibility and effectiveness of long-context imitation learning while reducing its sample complexity.

Abstract

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.

Problem

• imitation learning
• context length
• robotic manipulation
• memory
• long-context

Innovation

• context length
• imitation learning
• diffusion policies
• long-context learning
• cross-attention