Training and Evaluating Diffusion Policies with Long Context Lengths

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.

• The paper demonstrates through nearly 200 policies across simulation and hardware tasks that naive context scaling can remain effective up to 92 observations when data and architecture are appropriate.
• The study finds that long-context failures primarily reflect manipulation-skill difficulty and covariate shift, while UNet with cross-attention outperforms DiT and reduces sample complexity.
• The paper introduces variable history training, with progressive scheduling improving limited-data performance by 1.25–2× while preserving the memory benefits of longer observations.

Overview

This paper from MIT presents a systematic empirical study of context length scaling in Diffusion Policy-based imitation learning, together with a new training algorithm for the limited-data regime (2606.16447). The central claim is contrarian: prior work has characterized naive scaling of observation history as brittle and prone to catastrophic failure, motivating auxiliary losses, VLM-based frame selection, and language-guided retrieval mechanisms. The authors argue this conclusion is premature, and support it with close to 200 trained policies spanning five simulation tasks (plus one hardware task), three data scales per task, multiple context lengths up to $T_o = 80$ (and $T_o = 92$ on hardware), and three denoising/conditioning architectures.

Benchmark design

The evaluation spans two categories of tasks. Tasks solvable with short context include push-T, robomimic square, and robomimic lift, evaluated at $T_o \in \{1, 2, 5, 10, 12, 16, 20\}$. Tasks requiring long context are two newly introduced benchmarks: push-and-return (push a block to center, retreat to neutral, push it back to its original spawn location) and grasp-and-return (same structure via prehension). These tasks deliberately pair memory-dependent decisions with manipulation primitives of differing local stability — planar pushing is locally unstable, whereas grasping is locally stable. Success is decomposed into task success, manipulation completion (memory requirement relaxed), and contextual success (fraction of manipulation-completion rollouts that return the object to the correct location). A hardware marshmallows task, requiring exactly two scoops before pressing a button, provides an out-of-distribution validation at $T_o = 92$. Each simulated success rate is computed over 100–200 trials with Wilson score confidence intervals.

Naive scaling is not catastrophically brittle

The headline finding contradicts prior claims directly. With a UNet backbone using cross-attention conditioning, single-task policies achieve high success rates on many tasks even under naive history scaling:

• On lift, performance shows no drop at any data scale as context length grows.
• On square and push-T, performance drops beyond short contexts only in the low ($N/2$) data regime; the gap shrinks with more data and becomes statistically insignificant for square at high scale.
• On grasp-and-return, long-context policies perform well even with $N/2$ trajectories; on push-and-return, acceptable performance requires $N$ trajectories.
• On hardware, naive scaling to $T_o = 92$ achieves high success on marshmallows with only 100 demonstrations, where prior work reported that VLM-based frame filtering was necessary for success using 250 demonstrations. The authors concede that differences in embodiment and environment could explain this discrepancy.

A key diagnostic result concerns why long-context policies fail when they do. Contextual success remains high for long-context policies across all data scales, while manipulation completion tracks task success closely. This implies failures stem primarily from difficulty learning the manipulation skill itself, not from failure to track history-dependent requirements. The authors further attribute the sample complexity of long-context learning to local stability of the manipulation primitive: grasp-and-return policies succeed despite mean gripper contact offsets of roughly 1.2–1.4 cm from the brick center, because the prehensile primitive tolerates error. Conversely, short-context policies fail on grasp-and-return regardless of data scale due to unobservability of the expert's latent phase — grasping looks visually similar in both directions of motion — indicating fundamental limits of short-context imitation for some tasks.

The authors identify three methodological gaps behind prior criticism of naive scaling: absence of dataset-size benchmarking, use of the DiT architecture (which they show fails even with enlarged parameter counts, removal of causal attention within action chunks, and additional sampling steps), and lack of comparison against short-context-solvable tasks.

Architecture matters: cross-attention conditioning

Comparing UNet+FiLM, UNet+Cross-Attention, and DiT, the paper finds that DiT exhibits catastrophic failure in several cases — including behavior matching the trends cited by past-token prediction work, suggesting that criticism of naive scaling may partly reflect a poor baseline architecture. UNet+Cross-Attention matches FiLM at short context lengths but outperforms it at longer contexts, particularly in the limited-data regime. Because cross-attention keeps model size constant (~151M parameters) as context length grows, whereas FiLM's conditioning vector scales with $T_o$, the authors also ran a parameter-matched comparison and found no change in conclusions in the $N/2$ regime. As data scales, the gap between the two UNet variants shrinks, supporting the interpretation that cross-attention provides an inductive bias reducing the sample complexity of long-context learning.

Variable history training

Motivated by loss–success diagnostics showing that long-context checkpoints achieve lower training and validation loss yet worse closed-loop success — consistent with overfitting compounded by covariate shift — the authors propose variable history training: a single policy trained over a curriculum of context lengths sampled from a step-varying distribution $T_o = 92$0, with short-context datagrams regularizing against overfitting. Two schedules are studied: random sprinkle (80% probability of full context) and progressive (a curriculum expanding from minimum to maximum context). In the $T_o = 92$1 regime, progressive scheduling with a short past-action prediction horizon yields 1.25x–2x improvement over the strongest naive-scaling baseline, matching short-context performance while retaining memory benefits. When data suffices for naive scaling, random sprinkle with full past prediction maintains or improves performance. The practical recommendation is therefore regime-dependent, which the authors state plainly rather than claiming a universally superior configuration.

Revisiting past-token prediction

The paper re-examines "Learning Long-Context Diffusion Policies via Past-Token Prediction" (Torne et al., 14 May 2025), which attributes gains to predicting past actions as an auxiliary loss and treats freezing the observation encoder as a mere compute optimization. The reproduction finds that encoder freezing contributes materially to success, independently or jointly with past-action prediction — a claim contradicting the original attribution. Past-action prediction alone shows unclear benefit across data regimes. Push-and-return is an exception where the frozen-encoder effect does not hold, and the proposed variable history training matches or exceeds the PTP-style configurations. Notably, these comparisons replace the original DiT backbone with UNet+Cross-Attention, since the authors found DiT unsuitable even after hyperparameter corrections. The inference-time chunk verification technique from PTP was excluded after being found statistically insignificant (e.g., 176 vs. 175 successes out of 200 on square).

Limitations and open questions

The paper concedes several limitations. All reported results come from single training runs without accounting for training-run variance, a choice justified by computational cost but acknowledged as noisy — particularly for the loss–success diagnostic analysis. The push-and-return benchmark uses only two return modes; scaling combinatorial complexity while holding local complexity fixed remains untested. Inference-time overhead of processing many image frames per step is not addressed, and adapting distillation methods to long-context policies is left open. Finally, whether variable-length context at inference time can be actively regulated, analogous to agentic LLM context management, is posed as an open algorithmic question.

Conclusion

This paper demonstrates, through an unusually broad controlled study, that naive context length scaling in Diffusion Policy is substantially less brittle than reported in prior literature, provided the conditioning architecture (UNet+Cross-Attention) and sufficient data are used. It establishes that long-context failures are driven primarily by manipulation-skill learning difficulty and covariate shift rather than memory-tracking failure, introduces variable history training to reduce sample complexity in the low-data regime, and corrects the attribution of past-token prediction gains to encoder freezing. Its main lasting contribution may be methodological: any claim about long-context imitation learning should be conditioned on data scale, architecture, and manipulation-primitive stability.