Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning

AI summary

Overview

Research area: Imitation learning and behavior cloning for partially observable (non-Markovian) decision-making, with applications spanning synthetic memory tasks, grid worlds, robot manipulation benchmarks, and a real robot.

Technical level: Advanced. The paper combines a Kolmogorov-complexity-based theoretical analysis of neural network simplicity bias with reinforcement learning (PPO) and downstream behavior cloning, and assumes familiarity with POMDPs, mutual information, and memory-augmented policy architectures.

Scope (1 sentence): The paper proposes "Keyframe Mnemonics," a three-stage self-supervised pipeline that discovers decision-critical past observations ("mnemonics") without annotation, and trains behavior cloning policies conditioned on them so that performance does not degrade at horizons far beyond those seen in training.

What This Paper Is About

Behavior cloning policies typically condition on only a short window of recent observations, which makes them fail on tasks where a cue or piece of task-relevant information seen earlier must be remembered across a long gap before it affects the action. Existing fixes use recurrent models (which suffer hidden-state collapse and gradient instability under backpropagation through time) or attention/transformer models (which are capped by a finite context length, scale quadratically with context, and can overfit spurious visual correlations). The paper's goal is to automatically identify the sparse set of past observations that actually carry action-relevant information, keep only those in a small memory buffer, and train a policy on them so that memory retention generalizes to horizons orders of magnitude beyond the training data.

Key Contributions

1. A self-supervised keyframe discovery algorithm. The method finds decision-relevant keyframes in non-Markovian expert demonstrations without any explicit keyframe annotations, using a randomly masked sampling procedure and a learned objective (Figure 1, Algorithm 1).

2. A theoretical account under task-structure assumptions. The authors derive information retention guarantees and intuition for why masking biases the proxy policy toward keyframes, via Lemma 1, Assumption 3, Theorem 1, Proposition 1, Corollary 1, and Theorem 2, supported with empirical evidence.

3. Demonstrated horizon-invariant behavior cloning. Mnemonic-conditioned MLP policies reach 100% success rate in synthetic and grid memory domains, Transformer policies improve absolute success rate by 13.9% on the MIKASA-Robo benchmark, and a flow-matching policy retains 80% success rate at 20× longer test horizons on a real robot.

4. An architectural decomposition. The problem is split into (i) discovering information-critical observations and (ii) training a BC policy on demonstrations augmented with those observations, so that an optimal selector retains the keyframe set on an infinite horizon under the stated assumptions.

Main Findings

• Objective decreases with more keyframes. Across three synthetic tasks, the learned objective value J(ℬ) decreases as more oracle keyframes are added to the proxy policy's context buffer, supporting Corollary 1 (Figure 2a).

• Large masked-sampling support drives selection quality. Active memory regret stays high for small masked-sampling support sizes M and converges to zero as M increases, while passive memory regret grows with M as the proxy policy fits the keyframe-based hypothesis (Figure 2b, 2c). The paper notes M = 1.0× corresponds to a calibrated value chosen empirically.

• Near-perfect synthetic and grid performance. KM achieves 100% success rates on the synthetic memory and grid domains and maintains near-perfect performance across all horizons, whereas baseline policies collapse beyond the training range (Figure 3). At further horizons of 10^5 and 10^6, KM retained performance equivalent to its in-distribution behavior on all three synthetic tasks (Appendix E, Table 11).

• Gains on robot manipulation. KM improves average absolute success rate by 13.9% over the strongest published baseline, ELMUR, across the 23 MIKASA-Robo tasks. The Figure 4 caption states KM outperforms ELMUR on 17 of the 23 tasks, while the Section 3.2 text states it outperforms ELMUR on 21 of the 23 tasks; results for RememberShape3/5/9, RememberColor3/5/9 and RememberShapeColor are aggregated and the full table is in Appendix D.

• Sustained memory retention under long cue-to-choice delays. On RememberColor3-ΔT, all methods are trained at ΔT = 15, but out-of-distribution behavior diverges: ELMUR collapses to near-zero success by ΔT = 35, PTP retains roughly half its training performance through ΔT = 75 before degrading, and KM sustains its training-time success rate of ~70% across the full evaluated range up to ΔT = 250 (Figure 5a).

• No regression in the no-memory regime. KM matches the diffusion-policy baseline on three Robomimic tasks (Lift, Can, Square) where memory provides no benefit, confirming it does not regress when memory is irrelevant (Figure 5b).

• Real-robot horizon invariance. Over 20 trials per method, DP and PTP deteriorate while KM retains an average 80% success rate at 20× longer test horizons (Figure 6).

• Selector approaches the oracle. With sufficiently large calibrated support, active memory regret drops to zero while passive regret increases, giving an empirical regime of r_passive ≫ r_active ≈ 0, which the paper presents as evidence for Theorem 1.

Methodology in Plain English

The authors assume the demonstration data contains a small set of "keyframe" observations (much fewer than the trajectory length) that carry the information needed to predict the action, plus many background frames that carry little. Because they have no labels for which frames are keyframes, they build a training signal indirectly.

Stage 1 — Proxy policy training. They train a small network to predict the expert action from a buffer of past observations. Crucially, during training the buffer is assembled by randomly masking out most past frames (with the number of kept frames drawn from a Binomial distribution and capped at a memory budget k). This masking creates a combinatorially large set of input–output training pairs. The authors argue that this enormous support makes it very costly (in Kolmogorov complexity terms) for the network to learn a hypothesis that depends on background frames, so the network is pushed toward a simple hypothesis that depends only on keyframes. Once trained, the network's prediction loss acts as a scoring function over buffers: buffers containing more keyframes score better.

Stage 2 — Selector policy training. Using that frozen loss as a reward, they train a selector policy with PPO to assign priority scores to observations as they stream in. The buffer is modeled as a priority queue with a fixed budget. The reward is a normalized logarithmic function of the objective, multiplied by a "surprise" coefficient (based on action uncertainty) scaled by episode length, so that timesteps where the action strongly depends on memory get more weight. A FIFO queue is used during training (penalizing low priority for long-influential observations) and a LIFO queue at inference (preserving high-priority observations over long horizons).

Stage 3 — Behavior cloning. The trained selector produces keyframe-augmented trajectories, and a BC policy is trained conditioned on the current observation plus the buffer of selected keyframes. Because the assumptions may not hold strictly in complex domains like robot manipulation, the current observation is also always included.

Why This Matters

Impact on research. The paper reframes long-horizon memory in behavior cloning as a sparse buffer-selection problem rather than a sequence-model capacity problem. It offers a self-supervised, annotation-free way to find decision-critical observations, and it provides a theoretical argument (grounded in neural network simplicity bias and Kolmogorov complexity) for why temporal masking plus a memory budget should isolate invariant hypotheses. This is a different attack surface than recurrent, attention, or hybrid memory architectures, and the empirical claim of horizon invariance is the distinctive part: performance is maintained at 10^5 and 10^6 timesteps in synthetic tasks and at ΔT = 250 in manipulation.

Real-world applications.

• Robot manipulation with delayed cues, where the object or color the robot must act on was visible only briefly earlier in the episode.
• Industrial or warehouse robots operating in partially observable environments where a task parameter is specified once and must be remembered across a long sequence of actions.
• Assistive and service robotics where instructions or spatial context are provided once and must be carried across extended interactions.
• Any imitation-learning pipeline built from human demonstrations where the meaningful moments of a demonstration are not labeled, and where retaining the entire video history is computationally prohibitive.

Industry relevance. Practitioners who train behavior cloning policies from demonstrations face a direct tradeoff between context length and compute, and hybrid memory methods that work in-distribution can degrade outside the training horizon. A method that keeps a small, fixed set of decision-relevant frames, matches a diffusion policy on memory-free tasks (Robomimic Lift, Can, Square), and keeps 80% success rate at 20× longer horizons on real hardware, is directly relevant to deployment settings where test episodes run longer than the demonstrations used for training.

Future Directions

• Weakening the structural assumptions. The paper states that in complex domains such as robot manipulation its assumptions may not hold strongly and that actions may depend partially on background observations. Extending the guarantees to those regimes is a natural next step.

• Removing the empirically calibrated support size. The masked-sampling support size M = 1.0× is described as a calibrated value chosen empirically to be sufficient per Theorem 1. Deriving or automatically tuning this quantity rather than calibrating it would make the method easier to apply to new domains.

• Comparing learned selection against other sub-sampling schemes at scale. The paper positions KM against PTP, which sub-samples frames from the past for conditioning. Investigating when learned selection outperforms simpler sub-sampling, and on what task structures, remains open.

• Broadening evaluation and reconciling reported task counts. The paper reports an average 13.9% absolute success-rate gain over ELMUR across 23 MIKASA-Robo tasks, with the Figure 4 caption reporting outperformance on 17 of 23 tasks and the Section 3.2 text reporting 21 of 23; the benchmark itself is described as containing 32 tasks. Clarifying the exact per-task comparison and extending the evaluation to the remaining tasks would strengthen the empirical picture.

• Natural extensions to other policy classes. The authors report results with MLP policies (synthetic and grid domains), Transformer policies (MIKASA-Robo), and a flow-matching policy (real robot), leaving open how keyframe conditioning interacts with other modern policy parameterizations.

Target Audience

This paper is written for machine learning and robotics researchers working on imitation learning, behavior cloning, partially observable decision-making, and long-horizon memory in policies. It will also interest engineers building production robot policies from demonstration data who need memory retention beyond the training horizon, and theorists interested in connecting neural network simplicity bias and Kolmogorov complexity to practical representation-selection algorithms. Readers without a background in POMDPs, mutual information, or reinforcement learning optimization will find the theory sections (Lemma 1, Theorem 1, Proposition 1, Corollary 1, Theorem 2) demanding, while the three-stage training description and the empirical figures are accessible to a broader audience.

Authors’ abstract

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 $\textit{discovers}$ a set of information-critical observations ($\textit{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.