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

Paper: 2602.15010 Authors: Max Sobol Mark, Jacky Liang, Maria Attarian, Chuyuan Fu, Debidatta Dwibedi, Dhruv Shah, Aviral Kumar Categories: cs.RO, cs.LG

Abstract

Many robotic manipulation tasks inherently require reasoning over historical observations—such as remembering which drawers have been searched when looking for an object. While current state-of-the-art robot policies excel at reactive behaviors conditioned on current observations, they struggle with tasks requiring memory of past states. This paper identifies that naively conditioning on observation history leads to spurious correlations: policies overfit to incidental features in training trajectories that don’t generalize at deployment. The root cause is exponentially growing history space that cannot be adequately covered during training. The authors propose Big Picture Policies (BPP), which uses vision-language models to extract a minimal set of task-relevant keyframes from observation histories. By compressing diverse trajectories into compact event representations, BPP achieves 70% higher success rates than baselines on real-world manipulation tasks.

Key Contributions

• Problem Analysis: Identifies that spurious correlations in history-conditioned policies stem from insufficient coverage of the exponentially large history space during training
• BPP Framework: Introduces a method that uses vision-language models (VLMs) to detect meaningful keyframes, projecting diverse rollouts onto a compact set of task-relevant events
• Distribution Shift Reduction: Demonstrates that keyframe-based conditioning substantially reduces the gap between training and deployment distributions without sacrificing policy expressivity
• Empirical Validation: Achieves 70% improvement over best baselines across four real-world manipulation tasks and three simulation environments

Methodology

The BPP approach consists of three main components:

Keyframe Detection: A vision-language model analyzes the observation history to identify frames corresponding to significant task events. Rather than conditioning on all past observations, BPP selects only those frames where meaningful state changes occur. This is formulated as detecting frames $\{o_{t_1}, o_{t_2}, ..., o_{t_k}\}$ where $k \ll T$ and each $o_{t_i}$ represents a semantically important moment.

Policy Architecture: The policy $\pi(a_t | o_t, \{o_{t_i}\}_{i=1}^k)$ conditions on both the current observation and the detected keyframes. This architecture maintains expressivity for complex tasks while dramatically reducing the effective history space that must be covered during training.

Training Procedure: BPP is trained using standard imitation learning on demonstration data, but the keyframe detector preprocesses trajectories to extract relevant historical context. This preprocessing step is task-agnostic and leverages the semantic understanding of pre-trained VLMs.

The key insight is that by projecting high-dimensional observation histories onto a low-dimensional space of task-relevant events, BPP achieves better generalization. The exponential growth of possible histories is tamed by focusing only on semantically meaningful transitions.

Experimental Results

The authors evaluate BPP on seven tasks total: four real-world manipulation tasks and three simulation environments. All tasks require history conditioning to succeed.

Real-World Tasks:

• Drawer search and retrieval
• Multi-step assembly requiring memory of completed sub-tasks
• Object rearrangement with occlusions
• Sequential manipulation with state dependencies

Performance Metrics: BPP achieves an average 70% improvement in success rate over the best baseline methods. Notably, naive history conditioning (using all past frames) performs worse than current-observation-only policies on several tasks, confirming the spurious correlation problem.

Ablation Studies: The paper demonstrates that:

• Keyframe selection is crucial—random frame sampling performs poorly
• VLM-based detection outperforms hand-crafted heuristics
• The number of keyframes can be task-adaptive without manual tuning

Generalization: BPP shows strong out-of-distribution generalization, succeeding on trajectories with novel orderings of sub-tasks and environmental variations not seen during training.

Technical Insights

The paper provides several important theoretical and practical insights:

Coverage Problem: The authors formalize why history-conditioned policies fail. Given a horizon $T$, the space of possible observation histories grows as $|\mathcal{O}|^T$. Training data provides limited coverage, leading policies to memorize spurious patterns. Traditional regularization (dropout, weight decay) doesn’t address this fundamental coverage issue.

Keyframe Compression: By reducing history to $k$ keyframes where $k \ll T$, the effective history space becomes $|\mathcal{O}|^k$, which is exponentially smaller and more tractable for training data to cover.

VLM Advantages: Pre-trained vision-language models bring semantic understanding that enables zero-shot keyframe detection. The VLM identifies “important” frames based on task-relevant visual changes, without requiring task-specific training.

Comparison to Alternatives: The paper compares against several baselines including recurrent policies (LSTMs, Transformers), attention mechanisms, and various regularization schemes. BPP’s advantage comes from explicitly addressing the coverage problem rather than just adding model capacity or regularization.

Takeaways

1. History-conditioned robot policies fail primarily due to insufficient coverage of the exponentially large history space, not model capacity limitations
2. Vision-language models can effectively identify task-relevant keyframes in observation histories without task-specific training
3. Compressing histories to minimal keyframe sets reduces distribution shift between training and deployment by 70% in success rate
4. The approach is general across diverse manipulation tasks and doesn’t require manual engineering of what constitutes a “keyframe”
5. BPP demonstrates that semantic compression of temporal information is more effective than naive temporal modeling for robot learning