Learning Long-Context Diffusion Policies via Past-Token Prediction

TL;DR

This work tackles the problem of learning long-context diffusion policies for robotics by identifying a gap in temporal action dependencies between expert demonstrations and learned policies. It introduces Past-Token Prediction (PTP), an auxiliary objective that jointly predicts past and future actions to reinforce temporal modeling, paired with a multi-stage training regimen that caches embeddings to reduce memory and compute. A test-time self-verification mechanism further selects action sequences consistent with past actions, improving reliability. Across ten tasks in simulation and real robots, the approach yields about 3× performance gains and over 10× faster training, highlighting its practical impact for efficient, history-aware robotic control.

Abstract

Reasoning over long sequences of observations and actions is essential for many robotic tasks. Yet, learning effective long-context policies from demonstrations remains challenging. As context length increases, training becomes increasingly expensive due to rising memory demands, and policy performance often degrades as a result of spurious correlations. Recent methods typically sidestep these issues by truncating context length, discarding historical information that may be critical for subsequent decisions. In this paper, we propose an alternative approach that explicitly regularizes the retention of past information. We first revisit the copycat problem in imitation learning and identify an opposite challenge in recent diffusion policies: rather than over-relying on prior actions, they often fail to capture essential dependencies between past and future actions. To address this, we introduce Past-Token Prediction (PTP), an auxiliary task in which the policy learns to predict past action tokens alongside future ones. This regularization significantly improves temporal modeling in the policy head, with minimal reliance on visual representations. Building on this observation, we further introduce a multistage training strategy: pre-train the visual encoder with short contexts, and fine-tune the policy head using cached long-context embeddings. This strategy preserves the benefits of PTP while greatly reducing memory and computational overhead. Finally, we extend PTP into a self-verification mechanism at test time, enabling the policy to score and select candidates consistent with past actions during inference. 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.

Figures (14)

• Figure 1: We propose a simple framework for learning long-context diffusion policies from human demonstrations. Our method leads to 3× gains in performance while reducing the training expense by more than 10×.
• Figure 2: Comparison of regression-based and diffusion-based policies in temporal action dependency, normalized by that in demonstrations.
• Figure 3: Illustration of past-token prediction. The policy head is trained to jointly predict both past and future action tokens, encouraging the model to capture the temporal dependencies that are otherwise lost between past and future actions.
• Figure 4: Overview of multistage training with embedding caching. As PTP acts on the decoder, caching embeddings substantially improves inference speed without sacrificing performance. We use a visual encoder from a short-range policy with low validation loss to compute the embeddings of the images in the buffer and cache them in the buffer. With the cached embeddings we can train the long-horizon policy much faster. At test time we take the original encoder.
• Figure 5: Test-time verification. Multiple action sequences are sampled from the same observation, and the policy selects the sequence that is most consistent compared to ground-truth previous actions.