Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning

Keyframe Mnemonics (KM) learns what a behavior cloning policy should remember: compact, decision-relevant observations discovered from demonstrations and retained across long horizons.

Paper (Coming soon) Code (Coming soon)

Method overview. A selector assigns priorities to incoming observations, stores decision-critical keyframes in a buffer, and conditions the BC policy on this compact memory along with the current observation to predict actions.

∞ horizon memory retention guarantee¹

0% annotation overhead, fully self-supervised

100% success rate (SR) on synthetic memory tasks

+13.9% avg. absolute SR over SOTA on MIKASA-Robo

¹ under certain task-structure assumptions

Definition

The horizon-invariance property.

Horizon-invariance refers to the ability of a policy to generalize at unseen horizons much longer than the training distribution without performance degradation, where horizon is defined as the time gap between a decision-relevant keyframe observation and the later decision timestep where that information is used.

The robot is shown a red or green cube for 3–5s. The cue is removed, and after a delay two cubes appear 30 cm apart. The robot must reach for the color it was shown. Demonstration data had a cue-to-choice delay of ~5s, but the test delays are 20x longer.

KM (Ours)mnemonics

PTPframe subsampling

DPdiffusion policy

Δt=3–5strain

20 / 20

16 / 20

9 / 20

Δt=100–120sOOD

16 / 20

0 / 20

2 / 20

Horizon-Invariance

Same training horizon. Much longer tests.

At the training delay $\Delta t=15$, all policies can remember the previously shown color cue and choose the corresponding color. As the cue-to-choice delay grows out of distribution (OOD), baseline memory methods fail, whereas KM retains the observed cue in its mnemonic buffer and continues to act correctly.

Use the arrows below the grid to page through additional delay sets.

KM (Ours)mnemonics

PTPframe subsampling

ELMURexternal cache

Success Ratemean +/- std

Δt=15train

Δt=15

Δt=95OOD

Δt=95

Δt=250OOD

Δt=250

Why Memory Fails

The cue disappears. The action still depends on it.

Memory-intensive imitation in partially observable environments often fails because the key observation is no longer visible when the action is needed. Rather than expanding context windows or repeatedly compressing hidden state, KM learns to preserve the few frames that disambiguate future behavior.

In non-Markovian behavior cloning, the current observation $o_t$ is not sufficient, and the policy must instead maintain a belief over the history $h_t$. Recurrent and attention-based models can help, but repeated updates can blur state information, and finite context windows eventually drop earlier evidence.

Keyframe Mnemonics turns memory into a sparse selection problem. If a small set of keyframes carries more action information than rest of the background frames, the policy can act from a fixed-size mnemonic buffer $\mathcal B_t$ instead of the full observation history.

History $h_t$ at $t=8$ all observations so far

Mnemonic buffer selected keyframes

mnemonics$\mathcal B_t$ + current$o_t$ -> action$a_t$

Visual Distribution Shift

Policy remains robust to distractors.

Objects the policy never saw during training are dropped into the workspace during the rollout.

Both the selector and BC policy remain robust to unseen distractors. The selector correctly chooses the right keyframe to form the mnemonic buffer (orange border) which is then used by the BC policy for decision-making and control.

18 / 20

MIKASA-Robo

Memory-intensive robot manipulation tasks.

Successful evaluation rollouts of mnemonic-conditioned BC policies. The top bar in each video shows the mnemonic buffer (orange border) together with the current observation (blue border) used by the policy.

Use the arrows above or below the grid to page through 23 MIKASA-Robo tasks. Bars show mean success rate with +/- 1 std.

empty

Method

Three training stages, one compact memory.

KM first learns a proxy objective from randomly sampled past observations. A selector policy then populates a priority-queue buffer with keyframes that optimize this objective. Finally, the behavior cloning policy is trained on demonstrations augmented with the discovered keyframes, or mnemonics. The discovery process is fully self-supervised and does not use any manual keyframe annotations or privileged information.

1 Keyframe discovery from randomly sampled history buffers.

Past observations are randomly sampled to build small buffers bounded by the memory budget $k$, and the proxy $f_\theta$ predicts expert actions from these buffers. We refer to this as masked-sampling training, which leverages the simplicity bias of neural networks to encourage action mappings conditioned on true keyframes while discouraging spurious shortcuts.

2 Utilize proxy loss as keyframe identification reward.

The frozen proxy loss scores buffer quality. A selector policy $\pi_{\mathrm s}$ emits priorities $p_t$ that determine whether each observation enters the mnemonic buffer. The selector is trained with PPO to minimize the expected proxy objective $J_{\theta^*}(\mathcal{B}_t)$, thereby identifying keyframes informative for the memory task.

3 Train BC policy on mnemonic-augmented demonstrations.

A trained selector identifies decision-relevant keyframes which are used to create augmented demonstrations. Finally, the BC policy learns $\pi_{\mathrm{BC}}(a_t \mid o_t,\mathcal B_t)$, using the current observation for local control and the buffer for task context.

Horizon-Invariance

Performance stays flat when the horizon grows.

In environments where the task-structure assumptions hold strongly, KM demonstrates guaranteed memory retention up to horizons five orders of magnitude beyond training, as shown across synthetic memory tasks and grid domains.

KM is plotted with the four strongest baselines for each benchmark. Shaded regions show mean +/- 1 std.

KM (Ours) DLSTM LSTM-TBPTT LSTM-BPTT DT TrXL DMamba ELMUR

Add

Final-step MSE $\downarrow$

train $h=30$

Scattered-Copy

Copy success rate $\uparrow$

train $h=20$

T-Maze

Episode success rate $\uparrow$

train $h=30$

MiniGrid-Memory

Episode success rate $\uparrow$

train $h=17$

LTMB-Hallway

Episode success rate $\uparrow$

train $h=8$

Mnemonic-conditioned BC policy rollouts in grid domains. The mnemonic buffer (orange border) and the current observation (blue border) are shown at the top center of each video.

MiniGrid-Memory

LTMB-Hallway

Under the hood

BC policy activation maps on mnemonics.

Grad-CAM activation maps show which parts of the mnemonic buffer and current observation drive the behavior cloning policy $\pi_{\mathrm{BC}}$ at decision-critical steps. Hot regions concentrate on task cues stored in discovered keyframes and on matching evidence in the current observation.

Interpretable memory. The buffer $\mathcal B_t$ contains inspectable prior observations, making it possible to see what task-relevant cue the policy is carrying through partial observability.

Decision-relevant keyframes. Activation heatmaps show that the BC policy attends to cue regions in mnemonics when disambiguating actions, rather than relying only on the latest frame.

Grid domains

MiniGrid-Memory

LTMB-Hallway

Robot domains

Remember Color 3

Shell Game Touch

Take It Back

Failure Modes

Where the policy still breaks.

Representative KM rollouts where the mnemonic-conditioned BC policy fails, highlighting directions for future work.

Spatial OOD

Δt=15

The policy struggles to generalize to out-of-distribution spatial scenarios.

Task Failure

Δt=200

Even with correct decision-making, the policy fails to execute the task successfully.

Causal Confusion

Δt=250

Spurious correlations with the mnemonic buffer lead to incorrect actions, even when the keyframes are correctly retained.

Memory Capacity Failure

Chain Of Colors 3

Policy Rollout

Mnemonic buffer + Current observation

KM struggles with memory-capacity limitations. When multiple frames contain equivalent information, the proxy objective assigns similar reward to several buffer choices, making it difficult for the selector to preserve unique keyframes and retain information reliably.

We observe this failure mode in the Chain Of Colors 3 task, where the policy must remember three colors appearing over multiple timesteps. The policy fails to complete the task successfully because, although the selector retains some keyframes in the buffer, it does not consistently preserve the full mnemonic set of three colors.