Learning Long-Context Diffusion Policies via Past-Token Prediction

1 Introduction

• Historical observations and actions can inform choices in non-Markovian tasks, including when scenes are occluded, task stages shape strategy, or prior actions set an execution style.The cited examples include speed, curvature, and strategy.Source: §1 · p. 1
• Longer histories can expose policies to spurious features and increase memory and computation costs, while truncation or compact summaries risk discarding information critical to later decisions.The introduction notes that deployment performance may degrade when policies rely on spurious historical features.Source: §1 · pp. 1–2
• Past-token prediction (PTP) trains policies to predict past actions alongside future ones, strengthening temporal dependencies primarily in the policy head rather than the visual encoder.The paper identifies weaker temporal dependencies in learned action sequences than in expert data.Source: §1 · p. 2
• A multistage recipe pre-trains the visual encoder with short contexts, then fine-tunes a long-context decoder on cached image embeddings to reduce memory and computational overhead.At inference, the policy selects among candidate actions by consistency with previously executed actions.Source: §1 · p. 2
• Across six simulation and four real-world tasks, the method increases long-context policy success rate by 3× while reducing training overhead by over 10 times.On history-critical tasks, it enables 80% success where existing methods fail entirely.Source: §1 · p. 2, Figure 1 · p. 2

2 Related Work

• Imitation learning shifted from supervised observation-to-action prediction toward modeling demonstration distributions [15] [19] [46] [47], an approach recently successful toward generalist robot policies [4].Imitation learning remains susceptible to covariate shift [32], with compounding errors when learned policies diverge from demonstrations.Source: §2 · pp. 2–3
• Recent image-conditioned policies have been found to degrade with history, prompting some works to exclude history altogether [4] [6] [37] [38] [47].The paper presents its training recipe as counteracting this degradation.Source: §2 · p. 3
• Long-context methods reduce or summarize history, but assume much of it is irrelevant; autoregressive action-token approaches instead face open challenges in preserving long-range temporal structure [39].The paper's approach is orthogonal: it regularizes diffusion policies to retain past-action information.Source: §2 · p. 3
• Test-time scaling methods use extra inference compute for sample reranking or self-verification, and this paper applies past-action accuracy to self-verify robotic actions [35].The paper reports that more test-time compute can bring shorter-training or shorter-history checkpoints closer to optimal-checkpoint performance.Source: §2 · p. 3

3 Preliminaries

• Each observation is assumed to reveal only part of the underlying state, while the complete state can be inferred from observation history.The setting includes latent demonstration strategies, task stage, and visual occlusions.Source: §3 · p. 3
• The target policy takes the current observation and a k-step observation history as input, then predicts the current and next l actions.The paper notes that longer context can provide richer history but may reduce policy performance.Source: §3 · p. 3
• Longer histories can add spurious features correlated with but not causally influencing expert actions, and copycat behavior occurs when a policy predicts future actions by mimicking previous ones.Such spurious reliance can yield high training-distribution accuracy yet poor deployment generalization.Source: §3 · pp. 3–4
• Past-token prediction expands the policy’s prediction window in both temporal directions, jointly predicting past and future action tokens to encourage retention of past-action information.Figure 3 illustrates the joint prediction objective.Source: Figure 3 · p. 4
• The action predictability ratio divides expert prior-action prediction error by policy prior-action prediction error; values above 1 indicate over-reliance, while values below 1 indicate weaker-than-expert dependency.The policy error is measured using a two-layer MLP trained to predict current actions from previous actions alone.Source: §3 · p. 4
• Across three RoboMimic simulation tasks, regression policies’ action predictability exceeded expert levels, whereas diffusion policies had ratios below 1, indicating weaker-than-expert dependence on past actions.Figure 2 compares temporal action dependency normalized by that in demonstrations.Source: §3 · p. 4, Figure 2 · p. 4

4.1 Past-Token Prediction

• PTP trains the policy to predict past action tokens alongside future ones, expanding the prediction window in both temporal directions to preserve information about past actions.Future-chunk prediction alone remains insufficient for modeling dependencies between past and future decisions.Source: §4.1 · pp. 4–5

4.2 Memory-Efficient Training with PTP

• The three-stage recipe trains a short-context encoder, freezes it to cache frame embeddings, then trains a long-context policy head at near short-context computational footprint.The original encoder is used at test time.Source: §4.2 · p. 5, Figure 4 · p. 5

4.3 Test-Time Verification with PTP

• Longer histories generally improve performance, with gains saturating beyond four steps on simple Square but substantial boosts on more complex tasks.The cited examples include Transport, Long-Horizon Square, and Long-Horizon Aloha.Source: Figure 7 · p. 6
• Increasing past-token supervision aligns the learner more closely with expert action dependencies and results in higher success rates.Source: Figure 6 · p. 6
• Embedding caching speeds training by over 5× without hurting performance; complex Tool Hang policies fail to perform after two days without caching.Source: Figure 8 · p. 6
• At deployment, the policy samples candidate action sequences and selects the one whose reconstructed past best matches executed actions, rather than trusting potentially suboptimal prior predictions.The selection procedure is GPU-parallelizable with minimal computational overhead.Source: §4.3 · pp. 5–6

5.1 Simulation Experiments

• Over 30% improvement over no-history diffusion policies and over 60% over no-PTP diffusion policies across six simulation tasks.Gains are especially pronounced on long-horizon ALOHA and square tasks.Source: Figure 9 · p. 7
• Full PTP brings temporal action dependencies close to expert demonstrations, while the non-PTP baseline has 10× to 100× weaker action predictability ratios.The tested variants predict one, eight, or sixteen past action tokens.Source: §5.1 · p. 7
• Longer histories generally improve PTP-trained policies, with gains saturating beyond four steps on square but continuing on more complex tasks.Long-horizon square, long-horizon ALOHA, and transport benefit substantially from longer contexts.Source: §5.1 · p. 7
• Cached-embedding training matches performance in 20% of the training time and surpasses it within 40% of the compute budget.The comparison uses two days of training, with checkpoints evaluated every 50 epochs.Source: §5.1 · p. 8
• Increasing PTP-guided candidates from one to five improves success rate by approximately 5% across three challenging simulation tasks.The tasks are Tool Hang, Transport, and Long Square.Source: §5.1 · p. 8

5.2 Real-World Experiments

• 70% average success rate across real-world tasks, with 80% on Tape Replacement, where both baselines fail.The no-history baseline averages 15% success, while the no-PTP baseline is near zero on three of four tasks.Source: §5.2 · p. 8

6 Conclusion

• PTP strengthens temporal action dependencies, while multistage training and self-verification support its use during training and inference.Experiments cover six simulated and four real-world tasks.Source: §6 · p. 9

7 Limitations and Discussion.

• The method is studied specifically for extending context length in diffusion policies; its effectiveness on other policy classes remains a possibility, not an established result.The authors point to tokenization-based policies as a potential extension.Source: §7 · p. 9
• Inference overhead remains a practical bottleneck in closed-loop operation, and the history-downsampling and longer-action-chunk adjustments compromise policy reactivity.The authors identify further inference acceleration as future work.Source: §7 · p. 9

A.1 Which component of the policy benefits most from PTP?

• Decoder PTP performs on par with full PTP, whereas Encoder PTP performs significantly worse.The comparison is reported for Push-T.Source: §A.1 · p. 14, Figure 12 · p. 14
• The ablation attributes PTP's primary benefit to temporal modeling in the action decoder rather than changes to the visual encoder.This motivates decoupling encoder pretraining from long-context policy learning in the multistage recipe.Source: §A.1 · p. 14

A.2 Does our method reduce reliance on action chunking?

• Long-context, single-step policies trained by this method substantially outperform short-context counterparts and recover most long-context chunk-8 performance, reducing reliance on open-loop chunking.The comparison uses 16 past frames for long-context policies and 2 for short-context policies.Source: Figure 13 · p. 14, §A.2 · p. 14

A.3 Is PTP still critical when conditioning on past actions?

• Even when given the previous 16 actions, vanilla diffusion policies perform poorly without PTP, while PTP consistently yields substantially better results.Source: Figure 14 · p. 15, §A.3 · p. 15

B.1 Real-World Tasks

• Real-world evaluation uses Franka Panda and ALOHA, with 50–200 demonstrations per Franka task and 150 for ALOHA tape replacement.Franka observations include RGB images and end-effector pose; ALOHA uses dual-arm RGB views and proprioception.Source: §B.1 · p. 15

B.2 Simulation Tasks

• The six simulation tasks comprise Push-T, three Robomimic benchmarks, and two new long-horizon tasks requiring memory of an initial state or block position.Table 1 lists hyperparameters for simulation and real-world experiments.Source: §B.2 · p. 15, Table 1 · p. 16

B.3.1 Policy Architecture

• Long-context policies use a frozen encoder pretrained on 2-frame inputs; cached visual embeddings avoid repeated encoder passes and speed training.The paper reports that freezing the encoder does not impair performance, based on its Fig. 12 analysis.Source: §B.3.1 · p. 16

B.3.2 Subsampling Rate

• Temporal subsampling feeds every Kth observation, reducing effective input size while retaining broad temporal coverage.The sequence length T is a multiple of K.Source: §B.3.2 · p. 16

B.3.3 Context Length

• As observation history grows, the prediction horizon for past and future tokens scales proportionally; future-token prediction length has only a minor effect on task performance.Table 2 provides context-length configurations.Source: §B.3.3 · p. 16, Table 2 · p. 16

B.3.4 Action Dependency Metric

• Action predictability measures how accurately the current action can be predicted from 15 past actions; higher values indicate stronger learned temporal dependencies.The metric is compared between policy rollouts and expert demonstrations.Source: §B.3.4 · p. 16

• The appendix reports per-task success rates across training, history, and chunking conditions; Table 3 compares diffusion and regression baselines, while Table 4 covers closed-loop chunk-size-1 settings.Source: C Additional Results · p. 16, Table 3 · p. 17, Table 4 · p. 17