BPP: Long-Context Robot Imitation Learning by Focusing on Key History Frames

1. Problem Statement

The paper addresses a critical limitation in robot imitation learning: the inability of current policies to effectively utilize long-term history for non-Markovian tasks (tasks where the current observation is insufficient to determine the correct action).

• The Challenge: While many robot tasks (e.g., searching for an object, counting ingredients, multi-stage assembly) require remembering past events, standard policies often condition only on the current observation or a short temporal window.
• Naïve History Conditioning Failure: Simply feeding a sequence of past observations into a policy often leads to catastrophic failure. The authors identify that this is not merely an architectural issue but a coverage problem. The space of possible observation histories grows exponentially with the horizon. Since training data (demonstrations) covers only a tiny fraction of this space, policies learn spurious correlations (e.g., latching onto incidental background features or specific temporal patterns) that exist in the training data but do not generalize to out-of-distribution (OOD) trajectories encountered during deployment.
• Limitations of Existing Solutions: Previous attempts to fix this using regularization (e.g., auxiliary loss functions to predict ground-truth states) or architectural constraints have shown inconsistent results. The paper argues these methods fail because they do not fundamentally solve the lack of coverage over the history space; they merely modify the loss function without changing the input distribution.

2. Methodology: Big Picture Policies (BPP)

The authors propose Big Picture Policies (BPP), a method that conditions the robot policy on a minimal set of semantic keyframes rather than raw observation histories.

• Core Concept: Instead of processing a long, dense sequence of frames, BPP compresses the history into a small set of behaviorally salient events (keyframes) that are critical for task progress (e.g., "object grasped," "drawer opened," "button pressed").
• Keyframe Detection via VLM:
  • BPP utilizes an off-the-shelf Vision-Language Model (VLM) (specifically Gemini 3 Pro) as a binary detector.
  • The VLM is prompted with task-specific criteria to identify if a specific frame represents a "key event."
  • To avoid redundancy, the system detects the "rising edge" of an event (the first frame where the event is detected after a period of non-detection).
• Latency-Aware Training:
  • Since VLM inference introduces latency (3–5 seconds), BPP incorporates a latency mask during training. The policy is trained to condition only on keyframes detected up to time $t - \Delta$ (where $\Delta$ is the latency), simulating the delay encountered at deployment.
• Policy Architecture:
  • The underlying policy is a standard Diffusion Transformer.
  • Inputs: The current observation + the set of detected keyframes $\{o_k\}$.
  • This approach reduces the effective input space size, increasing the overlap between training and test-time distributions, thereby mitigating spurious correlations.

3. Key Contributions

1. Diagnosis of the Coverage Problem: The paper provides empirical evidence that the failure of history-conditioned policies stems from insufficient coverage of the history space, leading to spurious correlations. It demonstrates that even "golden" auxiliary regularizers (predicting ground-truth state) can hurt OOD performance if coverage is lacking.
2. BPP Framework: Introduction of a simple yet effective method that leverages VLMs to filter observation histories into semantic keyframes, bypassing the need for complex architectural changes or massive data collection.
3. State-of-the-Art Performance: BPP achieves significant improvements over both memoryless baselines and strong history-conditioned baselines (including those with auxiliary losses) across diverse real-world and simulation tasks.
4. Analysis of Action Chunking: The authors show that while action chunking (predicting a sequence of future actions) helps mitigate spurious correlations by forcing robust feature learning, it is insufficient on its own to close the gap to an oracle policy without the keyframe abstraction.

4. Experimental Results

The authors evaluated BPP on four real-world bimanual manipulation tasks (using an ALOHA 2 platform) and three simulation tasks.

Real-World Tasks:

• Mug Replacement, Marshmallows (counting), Drawer Search, Stacking Puzzle.
• Results: BPP achieved an average success rate of 53.6%, significantly outperforming the best prior approach (Past-Token Prediction, PTP) at 31.8% and the Naïve History baseline at 12.8%.
• Improvement: BPP showed a ~70% relative improvement over the best comparison method on average.
• Behavioral Analysis:
  • Memoryless policies failed due to repetitive loops (e.g., opening the same drawer repeatedly).
  • Naïve history policies often overfit to training trajectories, leading to failures like dropping objects or retrying successful grasps unnecessarily.
  • BPP demonstrated robust progress tracking, systematically searching drawers without revisiting them and correctly counting ingredients.

Simulation & Ablation Studies:

• Oracle Comparison: In some tasks (e.g., Variable Password), BPP even outperformed an "Oracle" policy that had access to privileged ground-truth state information. This suggests that the representation provided by keyframes is more learnable than raw state encodings.
• Data Efficiency: BPP achieved high performance with significantly fewer demonstrations (e.g., 200 demos) compared to naïve history conditioning, which required much larger datasets to converge.
• Robustness to VLM Errors: Even with imperfect VLM detection (simulated oracle keyframes vs. VLM keyframes), the performance drop was modest (70% vs. 60% success rate), indicating the method is robust to detection noise.

5. Significance and Future Directions

• Paradigm Shift: BPP shifts the focus from "how to process more data" to "how to process the right data." It suggests that for long-horizon tasks, abstracting continuous control into semantic milestones is more effective than raw temporal conditioning.
• Scalability: By using VLMs to identify keyframes, the approach offers a path to automatically labeling diverse, non-expert datasets for training generalist robot policies without manual annotation.
• Limitations: The method currently relies on VLM inference latency, which can be a bottleneck for highly dynamic, real-time tasks. Future work aims to distill keyframe detection into low-latency models.
• Broader Impact: The concept of "keyframes" extends beyond memory tasks to in-context learning and adaptation, where identifying relevant past segments (rather than raw frames) can improve learning efficiency.

In summary, BPP solves the long-context imitation learning problem by using VLMs to distill history into a compact, semantically meaningful representation, effectively reducing distribution shift and enabling robust performance in complex, memory-dependent robotic tasks.