# Training and Evaluating Diffusion Policies with Long Context Lengths

arXiv: https://arxiv.org/abs/2606.16447 · HTML: https://mlsanity.com/paper/2606.16447/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Diffusion and flow policies for robotics, one of 14 directions in Diffusion models.
- Builds on 26 of its 36 references that are in the map.
- First posted on arXiv 15 Jun 2026.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)
- Offline RL and imitation learning from data / [Robot manipulation imitation learning](https://mlsanity.com/offlinerl/robot-manipulation-imitation-learning/)

## Builds on (most cited first, 26 of 26)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [U-Net: Convolutional Networks for Biomedical Image Segmentation](https://mlsanity.com/paper/1505.04597/) (2015)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [FiLM: Visual Reasoning with a General Conditioning Layer](https://mlsanity.com/paper/1709.07871/) (2017)
- [π0.5: a Vision-Language-Action Model with Open-World Generalization](https://mlsanity.com/paper/2504.16054/) (2025)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [Sim-to-Real Transfer of Robotic Control with Dynamics Randomization](https://mlsanity.com/paper/1710.06537/) (2017)
- [RMA: Rapid Motor Adaptation for Legged Robots](https://mlsanity.com/paper/2107.04034/) (2021)
- [TraceVLA: Visual Trace Prompting Enhances Spatial-Temporal Awareness for Generalist Robotic Policies](https://mlsanity.com/paper/2412.10345/) (2024)
- [Causal Confusion in Imitation Learning](https://mlsanity.com/paper/1905.11979/) (2019)
- [MemoryVLA: Perceptual-Cognitive Memory in Vision-Language-Action Models for Robotic Manipulation](https://mlsanity.com/paper/2508.19236/) (2025)
- [A Careful Examination of Large Behavior Models for Multitask Dexterous Manipulation](https://mlsanity.com/paper/2507.05331/) (2025)
- [Interactive Supercomputing on 40,000 Cores for Machine Learning and Data Analysis](https://mlsanity.com/paper/1807.07814/) (2018)
- [Reinforcement Learning for Versatile, Dynamic, and Robust Bipedal Locomotion Control](https://mlsanity.com/paper/2401.16889/) (2024)
- [Instruction-driven history-aware policies for robotic manipulations](https://mlsanity.com/paper/2209.04899/) (2022)
- [Fighting Copycat Agents in Behavioral Cloning from Observation Histories](https://mlsanity.com/paper/2010.14876/) (2020)
- [RoboMME: Benchmarking and Understanding Memory for Robotic Generalist Policies](https://mlsanity.com/paper/2603.04639/) (2026)
- [Learning Long-Context Diffusion Policies via Past-Token Prediction](https://mlsanity.com/paper/2505.09561/) (2025)
- [The NLP Task Effectiveness of Long-Range Transformers](https://mlsanity.com/paper/2202.07856/) (2022)
- [Learning Bipedal Walking for Humanoids with Current Feedback](https://mlsanity.com/paper/2303.03724/) (2023)
- [BPP: Long-Context Robot Imitation Learning by Focusing on Key History Frames](https://mlsanity.com/paper/2602.15010/) (2026)
- [Empirical Analysis of Sim-and-Real Cotraining of Diffusion Policies for Planar Pushing from Pixels](https://mlsanity.com/paper/2503.22634/) (2025)

## Built on by (most cited first, 2 of 2)

- [WorldToken: Time-First Sequence Modeling for Robotic Imitation Learning](https://mlsanity.com/paper/2608.22591/) (2026)
- [Minimal Recurrent Behavioral Memory for Imitation under Partial Observability](https://mlsanity.com/paper/2609.25757/) (2026)

## Abstract (arXiv)

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 for Diffusion Policies 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.
