Training and Evaluating Diffusion Policies with Long Context Lengths

• URL: http://arxiv.org/abs/2606.16447v1
• Date: Mon, 15 Jun 2026 09:19:34 GMT
• Title: Training and Evaluating Diffusion Policies with Long Context Lengths
• Abstract summary: We first benchmark policy performance as context length is incrementally increased from short to long.<n>We propose a training algorithm to jointly train policies at multiple context lengths.<n>We re-evaluate some previously proposed solutions to long-context imitation learning.
• Score: 19.39755600541638
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• 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.

Related papers

• Improving Robotic Imitation Learning via Trajectory Standardization [7.969825103657315]
  A common preprocessing strategy is time-uniform downsampling to shorten sequences, but it cannot effectively remove speed-induced non-uniformity or redundant pauses.<n>We propose Information-Standardized Trajectory Resampling (ISR), an offline preprocessing method for effective imitation learning.<n>ISR improves task success rates by about 25%, remains robust across datasets collected from different operators, and reduces both dataset size and training cost.
  arXiv  Detail & Related papers  (2026-06-22T06:46:26Z)
• Beyond the Majority: Long-tail Imitation Learning for Robotic Manipulation [121.31598641762251]
  Long-tail distribution of training demonstrations hinders policy learning.<n>Approaching-Phase Augmentation (APA) transfers knowledge from data-rich head tasks to data-scarce tail tasks without requiring demonstrations.
  arXiv  Detail & Related papers  (2026-02-06T09:10:02Z)
• Let's (not) just put things in Context: Test-Time Training for Long-Context LLMs [39.99645732873852]
  We show that inference-time strategies show rapidly diminishing returns and fail at long context.<n>We propose a simple method that overcomes limitations of static self-attention.<n>Our method leads to large 12.6 and 14.1 percentage point improvements for Qwen3-4B on average across subsets of LongBench-v2 and ZeroScrolls benchmarks.
  arXiv  Detail & Related papers  (2025-12-15T21:01:37Z)
• Test-time Offline Reinforcement Learning on Goal-related Experience [50.94457794664909]
  Research in foundation models has shown that performance can be substantially improved through test-time training.<n>We propose a novel self-supervised data selection criterion, which selects transitions from an offline dataset according to their relevance to the current state.<n>Our goal-conditioned test-time training (GC-TTT) algorithm applies this routine in a receding-horizon fashion during evaluation, adapting the policy to the current trajectory as it is being rolled out.
  arXiv  Detail & Related papers  (2025-07-24T21:11:39Z)
• Learning Long-Context Diffusion Policies via Past-Token Prediction [48.86967836229684]
  We propose an alternative approach that explicitly regularizes the retention of past information.<n>We introduce Past-Token Prediction, an auxiliary task in which the policy learns to predict past action tokens alongside future ones.<n> 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.
  arXiv  Detail & Related papers  (2025-05-14T17:00:47Z)
• A Controlled Study on Long Context Extension and Generalization in LLMs [85.4758128256142]
  Broad textual understanding and in-context learning require language models that utilize full document contexts. Due to the implementation challenges associated with directly training long-context models, many methods have been proposed for extending models to handle long contexts. We implement a controlled protocol for extension methods with a standardized evaluation, utilizing consistent base models and extension data.
  arXiv  Detail & Related papers  (2024-09-18T17:53:17Z)
• Effective Long-Context Scaling of Foundation Models [90.57254298730923]
  We present a series of long-context LLMs that support effective context windows of up to 32,768 tokens. Our models achieve consistent improvements on most regular tasks and significant improvements on long-context tasks over Llama 2.
  arXiv  Detail & Related papers  (2023-09-27T21:41:49Z)
• A Character-Level Length-Control Algorithm for Non-Autoregressive Sentence Summarization [23.495225374478295]
  Sentence summarization aims at compressing a long sentence into a short one that keeps the main gist, and has extensive real-world applications such as headline generation. In our work, we address a new problem of explicit character-level length control for summarization, and propose a dynamic programming algorithm based on the Connectionist Temporal Classification (CTC) model.
  arXiv  Detail & Related papers  (2022-05-28T21:09:53Z)
• DDPG++: Striving for Simplicity in Continuous-control Off-Policy Reinforcement Learning [95.60782037764928]
  We show that simple Deterministic Policy Gradient works remarkably well as long as the overestimation bias is controlled. Second, we pinpoint training instabilities, typical of off-policy algorithms, to the greedy policy update step. Third, we show that ideas in the propensity estimation literature can be used to importance-sample transitions from replay buffer and update policy to prevent deterioration of performance.
  arXiv  Detail & Related papers  (2020-06-26T20:21:12Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.