Learning Long-Context Diffusion Policies via Past-Token Prediction

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.

Overview of Learning Long-Context Diffusion Policies via Past-Token Prediction

The paper explores the challenge of learning long-context robotic policies through imitation learning. It acknowledges that reasoning over long sequences of observations and actions is vital for complex robotic tasks but identifies key issues: increased training complexity due to memory demands, and the degradation of policy performance through spurious correlations in the data. Traditionally, these issues are often circumvented by limiting the context length, disregarding potentially crucial historical data. This work proposes a novel approach that emphasizes retaining historical information through Past-Token Prediction (PTP).

Key Contributions

1. Challenge Identification: The authors discuss the contrasts between conventional imitation learning policies and more recent diffusion-based models. Specifically, while classical methods exhibit a copycat problem, over-relying on past actions, diffusion policies often fail to demonstrate essential dependency across temporal sequences.
2. Past-Token Prediction (PTP): PTP is introduced as an auxiliary task where policies predict past action tokens alongside future ones. This technique encourages the model to leverage historical data effectively, improving overall temporal modeling within the policy. The benefits are primarily realized in the policy head rather than the visual encoder, thereby enhancing task performance significantly.
3. Multistage Training Strategy: The authors develop a multistage training strategy to optimize model performance while minimizing computational demands. By pre-training visual encoders on short context lengths and subsequently fine-tuning the policy head using cached embeddings from longer contexts, they preserve the advantages of PTP and reduce memory overhead substantially.
4. Self-Verification Mechanism: PTP is extended into a self-verification mechanism to be used during test time. This enables the policy to score and select action candidates that align with past actions, thereby fostering more robust inference.

Numerical and Empirical Results

The paper demonstrates that the proposed method improves the efficiency of policy training by tenfold and enhances performance by threefold across ten diverse robotic tasks, both simulated and real-world. In history-critical tasks, the approach achieved success rates where baseline techniques faltered.

Implications and Future Work

This research opens avenues for leveraging historical data robustly in the field of robotic imitation learning. The approach has profound implications, setting the stage for further advancements in diffusion-based policy learning and potentially extending these benefits to other models reliant on past context. Future exploration could involve integration into tokenization-based policies or adaptation to various real-world scenarios where historical context influences decision-making processes.

In conclusion, the paper advances the understanding of long-context policies in robotics, proposing methods that enhance training efficiency and policy effectiveness without the sacrifice of critical historical information. This aligns well with the evolving demand for more nuanced, temporally-aware robots capable of handling complex tasks with greater contextual understanding.