# Learning Long-Context Diffusion Policies via Past-Token Prediction

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

## Key facts

- Cited by 30 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 39 of its 46 references that are in the map.
- First posted on arXiv 14 May 2025.

## Placed in

- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

## Builds on (most cited first, 30 of 39)

- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [Let's Verify Step by Step](https://mlsanity.com/paper/2305.20050/) (2023)
- [RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control](https://mlsanity.com/paper/2307.15818/) (2023)
- [OpenVLA: An Open-Source Vision-Language-Action Model](https://mlsanity.com/paper/2406.09246/) (2024)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [π0: A Vision-Language-Action Flow Model for General Robot Control](https://mlsanity.com/paper/2410.24164/) (2024)
- [RT-1: Robotics Transformer for Real-World Control at Scale](https://mlsanity.com/paper/2212.06817/) (2022)
- [Scaling LLM Test-Time Compute Optimally can be More Effective than Scaling Model Parameters](https://mlsanity.com/paper/2408.03314/) (2024)
- [Learning Fine-Grained Bimanual Manipulation with Low-Cost Hardware](https://mlsanity.com/paper/2304.13705/) (2023)
- [Octo: An Open-Source Generalist Robot Policy](https://mlsanity.com/paper/2405.12213/) (2024)
- [What Matters in Learning from Offline Human Demonstrations for Robot Manipulation](https://mlsanity.com/paper/2108.03298/) (2021)
- [RoboCasa: Large-Scale Simulation of Everyday Tasks for Generalist Robots](https://mlsanity.com/paper/2406.02523/) (2024)
- [Large Language Models are Better Reasoners with Self-Verification](https://mlsanity.com/paper/2212.09561/) (2022)
- [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)
- [Inference-Time Scaling for Diffusion Models beyond Scaling Denoising Steps](https://mlsanity.com/paper/2501.09732/) (2025)
- [RoboAgent: Generalization and Efficiency in Robot Manipulation via Semantic Augmentations and Action Chunking](https://mlsanity.com/paper/2309.01918/) (2023)
- [Behavior Generation with Latent Actions](https://mlsanity.com/paper/2403.03181/) (2024)
- [A Survey of Imitation Learning: Algorithms, Recent Developments, and Challenges](https://mlsanity.com/paper/2309.02473/) (2023)
- [On the Self-Verification Limitations of Large Language Models on Reasoning and Planning Tasks](https://mlsanity.com/paper/2402.08115/) (2024)
- [HAMSTER: Hierarchical Action Models For Open-World Robot Manipulation](https://mlsanity.com/paper/2502.05485/) (2025)
- [3D Diffusion Policy: Generalizable Visuomotor Policy Learning via Simple 3D Representations](https://mlsanity.com/paper/2403.03954/) (2024)
- [Steering Your Generalists: Improving Robotic Foundation Models via Value Guidance](https://mlsanity.com/paper/2410.13816/) (2024)
- [What Matters in Building Vision-Language-Action Models for Generalist Robots](https://mlsanity.com/paper/2412.14058/) (2024)
- [OVM, Outcome-supervised Value Models for Planning in Mathematical Reasoning](https://mlsanity.com/paper/2311.09724/) (2023)
- [BAKU: An Efficient Transformer for Multi-Task Policy Learning](https://mlsanity.com/paper/2406.07539/) (2024)
- [Fighting Copycat Agents in Behavioral Cloning from Observation Histories](https://mlsanity.com/paper/2010.14876/) (2020)
- [In-Context Imitation Learning via Next-Token Prediction](https://mlsanity.com/paper/2408.15980/) (2024)
- [Feedback in Imitation Learning: The Three Regimes of Covariate Shift](https://mlsanity.com/paper/2102.02872/) (2021)
- [Equivariant Diffusion Policy](https://mlsanity.com/paper/2407.01812/) (2024)

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

- [MemER: Scaling Up Memory for Robot Control via Experience Retrieval](https://mlsanity.com/paper/2510.20328/) (2025)
- [MEM: Multi-Scale Embodied Memory for Vision Language Action Models](https://mlsanity.com/paper/2603.03596/) (2026)
- [Unified Diffusion VLA: Vision-Language-Action Model via Joint Discrete Denoising Diffusion Process](https://mlsanity.com/paper/2511.01718/) (2025)
- [ContextVLA: Vision-Language-Action Model with Amortized Multi-Frame Context](https://mlsanity.com/paper/2510.04246/) (2025)
- [ReMem-VLA: Empowering Vision-Language-Action Model with Memory via Dual-Level Recurrent Queries](https://mlsanity.com/paper/2603.12942/) (2026)
- [RLDX-1 Technical Report](https://mlsanity.com/paper/2605.03269/) (2026)
- [EventVLA: Event-Driven Visual Evidence Memory for Long-Horizon Vision-Language-Action Policies](https://mlsanity.com/paper/2606.20092/) (2026)
- [RoboTTT: Context Scaling for Robot Policies](https://mlsanity.com/paper/2607.15275/) (2026)
- [BPP: Long-Context Robot Imitation Learning by Focusing on Key History Frames](https://mlsanity.com/paper/2602.15010/) (2026)
- [Gated Memory Policy: In-Context Memorization and Adaptation](https://mlsanity.com/paper/2604.18933/) (2026)
- [ABPolicy: Asynchronous B-Spline Flow Policy for Real-Time and Smooth Robotic Manipulation](https://mlsanity.com/paper/2602.23901/) (2026)
- [MemoryVLA++: Temporal Modeling via Memory and Imagination in Vision-Language-Action Models](https://mlsanity.com/paper/2606.09827/) (2026)
- [VPWEM: Non-Markovian Visuomotor Policy with Working and Episodic Memory](https://mlsanity.com/paper/2603.04910/) (2026)
- [Persistent Visual Memory: Sustaining Perception for Deep Generation in LVLMs](https://mlsanity.com/paper/2605.00814/) (2026)
- [MilliVid: Hierarchical Latents for Long-Range Consistency in Video Generation](https://mlsanity.com/paper/2606.09056/) (2026)
- [Diffusion Forcing Planner: History-Annealed Planning with Time-Dependent Guidance for Autonomous Driving](https://mlsanity.com/paper/2606.11019/) (2026)
- [Training and Evaluating Diffusion Policies with Long Context Lengths](https://mlsanity.com/paper/2606.16447/) (2026)
- [Imitation Learning as Return Distribution Matching](https://mlsanity.com/paper/2509.12026/) (2025)
- [Present but Not Remembered: Auditing How Frozen VLAs Encode, Deploy, and Steer Visual History](https://mlsanity.com/paper/2607.03372/) (2026)
- [HiMe: Hierarchical Embodied Memory for Long-Horizon Vision-Language-Action Control](https://mlsanity.com/paper/2607.03449/) (2026)
- [Workspace Models: Lightweight Robotic Memory via Saliency-Driven Supervision](https://mlsanity.com/paper/2609.20820/) (2026)
- [In-Context Learning for Robots: Methods and Applications](https://mlsanity.com/paper/2609.36012/) (2026)
- [Trace-Focused Diffusion Policy for Multi-Modal Action Disambiguation in Long-Horizon Robotic Manipulation](https://mlsanity.com/paper/2602.07388/) (2026)
- [Expanding Spatial and Temporal Context for Robotic Imitation Learning With Scene Graphs](https://mlsanity.com/paper/2606.01072/) (2026)
- [UniMem: Unifying Multimodal Memory and Control for Vision-Language-Action Models](https://mlsanity.com/paper/2608.22869/) (2026)
- [Learning to Use Imagination: Progress-Conditioned Future Utilization for World Action Models](https://mlsanity.com/paper/2609.06578/) (2026)
- [Minimal Recurrent Behavioral Memory for Imitation under Partial Observability](https://mlsanity.com/paper/2609.25757/) (2026)
- [Measuring the Stability Assumption Behind Action Chunking](https://mlsanity.com/paper/2610.01626/) (2026)
- [Adaptive Mean Flow for Responsive Closed-Loop Robot Control](https://mlsanity.com/paper/2610.06089/) (2026)
- [Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning](https://mlsanity.com/paper/2610.10857/) (2026)

## Abstract (arXiv)

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.