Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning

As a meticulous researcher, I have thoroughly analyzed these excerpts from Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning. The paper presents a novel framework, Keyframe Mnemonics (KM), designed to tackle the inherent challenges of Behavior Cloning (BC) in non-Markovian environments where policies must reason over long temporal contexts.

Here is a detailed and comprehensive summary synthesizing the core concepts, methodology, contributions, and limitations:

Comprehensive Summary of Keyframe Mnemonics (KM)

The Core Problem:

Behavior Cloning (BC) in non-Markovian environments is fundamentally difficult because policies must effectively reason over contextual information spanning long horizons. Existing solutions—recurrent models and attention-based mechanisms—suffer from critical weaknesses: recurrent models suffer from hidden-state collapse and gradient instability during backpropagation through time, while attention mechanisms are constrained by finite context lengths.

The Proposed Solution: Keyframe Mnemonics (KM):

The KM method proposes a novel self-supervised approach that decomposes the complex non-Markovian decision problem into two distinct subproblems:

1. Discovering Information-Critical Observations (Keyframes): Identifying a sparse set of observations from past demonstrations that contain the most relevant information for action selection.

2. Training a Conditioned BC Policy: Training the final behavior cloning policy to condition its decisions on these discovered keyframes (mnemonics).

The central innovation lies in using an objective derived from randomly sampled past observations as a reward signal to guide the discovery of these keyframes, without requiring explicit annotations of which frames are important.

Methodology: The Three-Stage KM Process:

The KM pipeline is structured into three sequential stages:

1. Proxy Policy Training (Keyframe Discovery): A preliminary proxy policy is trained to predict expert actions given a randomly sampled past observation buffer, utilizing a negative log-likelihood loss. This stage leverages the simplicity bias inherent in neural networks and combinatorial complexity arguments to bias the learned hypothesis toward keyframe-based hypotheses. The goal here is to discover which observations are most predictive of the action distribution.

2. Selector Policy Training (Keyframe Selection): A reinforcement learning (PPO) agent, acting as a selector policy (pi s), is trained in this stage to populate a memory buffer (B t) with the keyframes identified in Stage 1. The reward function for this selector is derived from the objective established in Stage 1. This process learns to prioritize observations that are crucial for future decision-making.

3. Behavior Cloning (Final Policy Training): The final BC policy is trained to model the expert action distribution, explicitly conditioned on the set of mnemonics (keyframes) selected by the selector policy from the memory buffer.

Theoretical Guarantees and Performance:

The theoretical underpinnings of KM provide strong assurances for its efficacy under specific task-structure assumptions:

• Information Retention Guarantees: Under Assumption 2 (Action disambiguation by keyframes), the expected Kolmogorov complexity gap between background-based and keyframe-based hypotheses grows linearly with the masking-induced support size (M). This leads to a bias toward invariant, keyframe-based hypotheses. Theorem 2 formally guarantees infinite-horizon keyframe retention under these assumptions.

• Horizon Invariance: The architectural decomposition enables horizon invariance—meaning the memory retention extends beyond the training horizon. This allows BC policies to generalize effectively over long time scales where traditional hybrid memory methods fail, maintaining near-perfect performance across all horizons on every task.

• Empirical Superiority: KM demonstrates significant performance gains over numerous state-of-the-art baselines, including Decision Transformer (+40.4%), Conservative Q-Learning (+42.5%), and Diffusion Policy baseline (+36.6%). Crucially, it maintains near-perfect success rates across all horizons, whereas baseline policies collapse beyond the training range.

Key Findings and Insights:

• Self-Supervision: The keyframe discovery process is fully self-supervised; no explicit temporal annotations are required for keyframe identification.

• Decision Relevance: Analysis using Grad-CAM confirms that the final BC policy attends to information cues from both the memory buffer and the current observation, demonstrating that the discovered keyframes are indeed decision-relevant and provide necessary context for disambiguating actions in partial observability (POMDPs).

• Selector Behavior: The trained selector learns a non-uniform priority policy favoring keyframes with higher mutual information with the action. This qualitative behavior aligns well with oracle keyframe definitions, suggesting the proxy objective provides a strong signal for discovery.

Limitations and Caveats:

Despite its strengths, several limitations must be acknowledged:

• Assumption 1 Restriction: The method's reliance on the selector policy being Markovian (pi s(p t o t)) restricts the assumption that keyframes must be individually identifiable from a single observation. If this assumption is violated, multiple keyframes might be required to represent the same underlying information, potentially leading to multiple information-equivalent buffer states.

• Assumption 2 Sensitivity: Assumption 2 requires strong correlation between action ambiguity and sparse keyframes. In scenarios where actions are mostly local and low-variance (e.g., VizDoom-Two-Colors), the mutual information gap might be negligible, causing the proxy policy to fail to converge because the invariant hypothesis is hard to learn from a sparse buffer.

• Hyperparameter Sensitivity: The memory budget k is a critical hyperparameter. While scaling k can improve performance on some domains (e.g., doubling it helped on some synthetic domains), halving it degrades performance across all domains, indicating that the optimal size varies by domain and task structure.

• Computational Cost: The training time for the KM pipeline is substantial, reported as approximately 1.5 hours per seed for synthetic/grid benchmarks and roughly half a day per SST task on MIKASA-Robo, with PPO training time being the primary computational bottleneck in robot domains.

Conclusion:

Keyframe Mnemonics (KM) represents a significant advancement in behavior cloning for non-Markovian environments. By intelligently decomposing the problem into keyframe discovery and policy conditioning, KM achieves remarkable horizon invariance and superior performance by leveraging self-supervision to distill essential contextual information from expert demonstrations, offering a robust solution where existing recurrent or attention-based methods fail.

1. textbf Discovery of Decision-Relevant Keyframes Without Annotation (Self-Supervised Learning): This method proposes a novel self-supervised algorithm to discover decision-relevant keyframes in nonMarkovian expert demonstrations without explicit keyframe annotations. This allows AI systems to automatically identify and store the most informative past observations from raw expert data, circumventing the need for costly manual labeling.

2. textbf Horizon-Invariant Behavior Cloning: The resulting mnemonic-conditioned BC policies achieve 100% success rates (SR) and generalize to horizons orders of magnitude beyond training without performance degradation. This enables behavior cloning policies to retain compact, decision-relevant context over extended horizons, allowing them to function reliably in real-world scenarios where memory is critical.

3. textbf Robust Context Retention via Learned Selector Policy: The system trains a selector policy πs using reinforcement learning to optimize the learned reward function and identify a set of decision-relevant keyframes, which we call Mnemonics. This selector learns to assign priorities based on the mutual information advantage of keyframes with the expert action distribution, ensuring that actions are disambiguated by selecting observations that strictly provide more information for disambiguating the action than the background frames (Assumption 2).

4. textbf Adaptive Memory Management at Runtime: The system deploys a selector to identify keyframes at runtime and provides context to disambiguate actions under partial observability. This is achieved by using a memory buffer governed by an eviction rule, such as the LIFO queue during inference, which ensures that the priority queue retains the keyframe set for the entire horizon T without eviction (Theorem 2).

5. textbf Parameter-Efficient Architecture for Memory: The proposed architecture decomposes the problem into three stages—a proxy policy, a selector policy, and a BC policy—where the selector policy is compact (under 1% of the total on the synthetic tasks and roughly 15% on the grid tasks), since it only encodes a single observation at a time. This allows for effective memory management without incurring excessive parameter overhead.

Behavior cloning (BC) in non-Markovian environments is a challenging problem because policies have to reason over contextual information over long horizons. Existing policy architectures rely on recurrent or attention-based mechanisms to capture long-term dependencies. However, recurrent models suffer from hidden-state collapse and gradient instability under backpropagation through time, while attention-based models are fundamentally limited by context length. To address these issues, we propose Keyframe Mnemonics, a novel self-supervised method that discovers a set of information-critical observations (mnemonics) by learning an objective from randomly sampled past observations and using it as a reward for keyframe selection. We then train a BC policy that conditions on the discovered keyframes to model the action distribution. Under certain task-structure assumptions, our formulation provides context retention guarantees over an infinite horizon, while maintaining a small set of decision-relevant keyframes in the policy's working memory. We evaluate our method on synthetic memory domains, where mnemonic-conditioned BC policies achieve 100 % success rates (SR) and generalize to horizons orders of magnitude beyond training without performance degradation. Additionally, we evaluate on memory-intensive robot manipulation benchmark, achieving a 13.9 % average absolute SR improvement over the strongest baseline across 23 tasks and retaining 80 % SR at 20 times longer horizons on a real robot. Code and videos are available at https://keyframe-mnemonics.github.io.