Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning

Prabin Kumar Rath; Omkar Patil; Nakul Gopalan

cs.AIOct 7, 2026

Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning

Authors: Prabin Kumar Rath, Omkar Patil, Nakul Gopalan

Organizations: Arizona State University

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.

Figures & tables

Figure 1: Keyframe Mnemonics three-stage training for horizon-invariant behavior cloning. (1) Proxy policy training. A random binary mask $\Psi_{k}$ is applied to the observation history $h_{t}$ to assemble a memory buffer $\mathcal{B}_{t}$ of size at most $k$ , on which the proxy $f_{\theta}:\mathcal{B}\!\to\!\mathcal{A}$ is trained to predict the expert action $a_{t}$ under a negative log-likelihood loss $\mathcal{L}$ . Combinatorial support of training data (Lemma 1) biases $f_{\theta}$ toward keyframe-based hypotheses (Theorem 1, Proposition 1), so the frozen objective $J_{\theta^{\ast}}(\mathcal{B}_{t})=\mathcal{L}(f_{\theta^{*}}(\mathcal{B}_{t}),a_{t})$ scores keyframe-rich buffers favourably (Corollary 1; Section 2.1 ). (2) Selector policy training. The selector $\pi_{\mathrm{s}}$ is trained using PPO to populate $\mathcal{B}_{t}$ with keyframes using $J_{\theta^{\ast}}(\mathcal{B}_{t})$ as reward function (Section 2.2 ). (3) Behavior cloning. The trained selector is used to construct a keyframe augmented dataset $\mathcal{D}^{\dagger}$ , which is then used to train the BC policy $\pi_{\mathrm{BC}}(a_{t}\mid o_{t},\mathcal{B}_{t})$ (Section 2.3 ). Throughout the figure, blue tiles denote the current observation $o_{t}$ , orange tiles denote past frames in $h_{t}$ , hatched tiles are frames not selected into $\mathcal{B}_{t}$ .

Figure 2: Empirical evidence for Keyframe Mnemonics. (a) Adding oracle keyframes to the context buffer of the proxy policy leads to decrease in objective value across three problem domains (Corollary 1). (b) Active memory regret remains high for small masked-sampling support sizes

M

and converges to zero as the support size increases (Theorem 1). (c) Passive memory regret grows with

M

as the proxy policy fits the keyframe-based hypothesis. Note: Here,

M=1.0\times

corresponds to a calibrated value, chosen empirically (Appendix C.4 ), that is sufficient according to Theorem 1 and yields mnemonics for horizon-invariant BC policies, as shown in Figure 3 .

Figure 3: Horizon-invariance in synthetic and grid domains. Top row: Add, Scattered-Copy, T-Maze. Bottom row: MiniGrid-Memory, LTMB-Hallway. Each curve shows mean

\pm

std over 5 seeds. Band width is one std-dev. Vertical dashed line marks the maximum training horizon. Lower is better for MSE; higher is better for Success Rate. KM maintains near-perfect performance across orders of magnitude of horizon scaling. Most baselines degrade outside the training range.

Figure 4: Results across the 23 MIKASA-Robo tasks. Keyframe Mnemonics improves on the strongest baseline ( ELMUR ) by an average of

+13.9

absolute success rate, outperforming ELMUR on

17

of the

23

tasks. Results for tasks RememberShape3/5/9 , RememberColor3/5/9 and RememberShapeColor are aggregated, and the full table of results can be found in the Appendix D .

Panel of Figure 5 — Figure 5: Horizon-invariance in robot manipulation domain and results on Markovion tasks.

Figure 6: Horizon-invariance in real-world and qualitative keyframes across domains. From left-right: Real-world task SR over 20 trials per method. Performance of DP [ 5 ] and PTP [ 46 ] deteriorates while KM retains average $80\%$ SR at $20\times$ longer test horizons. Mnemonics in grid, robot manipulation and real-world domains. Frames with orange border shows discovered keyframes tracked by the memory buffer $\mathcal{B}_{t}$ , while the frame with blue border shows the current observation. Together, they provide sufficient context for BC policies to generate actions unambiguously. See the webpage https://keyframe-mnemonics.github.io for policy rollout videos.

Method / Stage	Time
KM (Ours)	10h 29m
Proxy	1h 15m
Selector	5h 46m
$\mathcal{D}^{\dagger}$ gen.	1h 23m
BC policy	2h 05m
PTP [ 46 ]	19h 34m

Figure 7: Training time.

Appendix figures & tables27 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 8: MiniGrid-Memory environment. The agent (red triangle) observes only a local ( $7\times 7$ ) view of the grid environment (dotted red box). It must navigate to the left room to observe the cue (green ball), then make a U-turn and traverse the corridor to the junction, where it must decide whether to move up or down. To succeed, the agent must move toward the same object observed in the left room.

Figure 9: LTMB-Hallway environment. The agent (red triangle) observes only a local ( $7\times 7$ ) view of the grid environment (dotted red box). It must remember the cue in the left room and then traverse the hallway while matching the remembered cue to the objects behind the locked doors. To succeed, it must open and enter the door containing the same object as the one in the left room. In the example above, the correct target is highlighted with a green circle.

	Train ( small ) ∗		Train ( large )		Test
Benchmark	Seq len	$\|\mathcal{D}\|$	Seq len	$\|\mathcal{D}\|$	Seq len	$\mathrm{pass}_{\mathrm{ms}}$
Add	$5$ – $10$	$1000$	$20$ – $30$	$5000$	$50$ – $100$	$300$
Scattered-Copy	$4$	$1000$	$40^{*}$	$5000$	$220$	$30$
T-Maze	$5$ – $10$	$100$	$55$ – $105$	$5000$	${\sim}1005$	$300$
MiniGrid-Memory	${\sim}21$	$2000$	—	—	${\sim}2012$	$300$
LTMB-Hallway	${\sim}46$	$50{,}000$	—	—	${\sim}425$	$10$

Appendix

Table 1: Training and test sequence lengths and number of demonstrations

|\mathcal{D}|

per dataset variant. MiniGrid-Memory and LTMB-Hallway use a single dataset (no large variant). In the “Seq len” columns, “

a

–

b

” denotes a range (lower–upper bound) and “

{\sim}n

” denotes the maximum sequence length. The last column

\mathrm{pass}_{\mathrm{ms}}

reports the number of offline rollout passes used to create the masked-sampling cache for proxy training (Appendix B.2 ).

Benchmark	$k$	$d_{\mathrm{feat}}$	Proxy / BC arch	$d_{\mathrm{hid}}$	$H$	$L$
Add	$3$	$32$	ReLU MLP	$256$	$1$	$3$
Scattered-Copy	$2/20$	$32$	ReLU MLP / Transformer	$256/64$	$1$	$3/4$
T-Maze	$2$	$16$	ReLU MLP	$256$	$1$	$3$
MiniGrid-Memory	$2/1$	$64$	Transformer	$128/64$	$1$	$4$
LTMB-Hallway	$2/1$	$64$	Transformer	$128/64$	$8/1$	$4$

Appendix

Table 2: KM module architectures per benchmark.

k

: memory budget;

d_{\mathrm{feat}}

: selector feature dimension;

d_{\mathrm{hid}}

: proxy / BC hidden width;

H

: action-chunk horizon;

L

: proxy / BC hidden layers. Cells with two values

a/b

are proxy / BC policy. Architecture of proxy and BC policy can be different per-domain basis.

	Selector PPO							Proxy
Benchmark	$n_{\mathrm{envs}}$	steps ppo	$n_{\mathrm{steps}}$	batch	$n_{\mathrm{epochs}}$	$\eta_{\pi}$	$\lambda$	$\eta_{f}$	batch
Add	$4$	$2\!\times\!10^{5}$	$2048$	$64$	$10$	$3\!\times\!10^{-4}$	$10^{-6}$	$10^{-3}$	$256$
Scattered-Copy	$4$	$2\!\times\!10^{5}$	$2048$	$64$	$10$	$3\!\times\!10^{-4}$	$10^{-3}$	$10^{-3}$	$256$
T-Maze	$16$	$2\!\times\!10^{5}$	$512$	$64$	$10$	$3\!\times\!10^{-4}$	$10^{-6}$	$10^{-3}$	$256$
MiniGrid-Memory	$32$	$3\!\times\!10^{5}$	$512$	$128$	$5$	$3\!\times\!10^{-4}$	$10^{-2}$	$10^{-3}$	$512$
LTMB-Hallway	$32$	$2\!\times\!10^{5}$	$1024$	$256$	$5$	$3\!\times\!10^{-4}$	$10^{-2}$	$10^{-3}$	$512$

Appendix

Table 3: KM selector optimization settings per benchmark. PPO uses Adam with the SB3 [ 60 ] defaults except where shown (

n_{\mathrm{envs}}

parallel selector environments, total environment steps ppo , rollout length

n_{\mathrm{steps}}

, batch size,

n_{\mathrm{epochs}}

per update, learning rate

\eta_{\pi}

, and reward normalisation constant

\lambda

from Equation 5 ). The value of

\lambda

depends on the type of loss function used for reward objective. The proxy is trained with AdamW (learning rate

\eta_{f}

, batch size).

Benchmark	LSTM	DLSTM	DT	DMamba	TrXL	ELMUR	KM (Ours)
Add	$100\mathrm{K}$	$740\mathrm{K}$	$182\mathrm{K}$	$394\mathrm{K}$	$182\mathrm{K}$	$4.4\mathrm{M}$	$\mathbf{268\mathrm{K}}$
Scattered-Copy	$102\mathrm{K}$	$742\mathrm{K}$	$184\mathrm{K}$	$396\mathrm{K}$	$184\mathrm{K}$	$4.4\mathrm{M}$	$\mathbf{366\mathrm{K}}$
T-Maze	$100\mathrm{K}$	$740\mathrm{K}$	$246\mathrm{K}$	$395\mathrm{K}$	$246\mathrm{K}$	$4.4\mathrm{M}$	$\mathbf{266\mathrm{K}}$
MiniGrid-Memory	$641\mathrm{K}$	$1.9\mathrm{M}$	$2.0\mathrm{M}$	$1.5\mathrm{M}$	$997\mathrm{K}$	$4.7\mathrm{M}$	$\mathbf{1.44\mathrm{M}}$
LTMB-Hallway	$641\mathrm{K}$	$1.9\mathrm{M}$	$2.0\mathrm{M}$	$1.5\mathrm{M}$	$997\mathrm{K}$	$4.7\mathrm{M}$	$\mathbf{1.45\mathrm{M}}$

Appendix

Table 4: Trainable parameter counts per benchmark. KM totals sum the proxy, selector and BC policy modules. The baseline architectural choices that determine these counts (layer widths, head counts, embedding dimensions) are inherited from the RATE [ 45 ] open-source implementation. Units: K =

10^{3}

, M =

10^{6}

Benchmark	Proxy $f_{\theta}$	Selector $\pi_{\mathrm{s}}$	BC policy $\pi_{\mathrm{BC}}$	Total
Add	$134\mathrm{K}$	$128$	$134\mathrm{K}$	$\mathbf{268\mathrm{K}}$
Scattered-Copy	$157\mathrm{K}$	$2.7\mathrm{K}$	$206\mathrm{K}$	$\mathbf{366\mathrm{K}}$
T-Maze	$133\mathrm{K}$	$48$	$133\mathrm{K}$	$\mathbf{266\mathrm{K}}$
MiniGrid-Memory	$957\mathrm{K}$	$242\mathrm{K}$	$245\mathrm{K}$	$\mathbf{1.44\mathrm{M}}$
LTMB-Hallway	$959\mathrm{K}$	$242\mathrm{K}$	$245\mathrm{K}$	$\mathbf{1.45\mathrm{M}}$

Appendix

Table 5: Decomposition of the Keyframe Mnemonics parameter budget into the proxy

f_{\theta}

, the selector policy

\pi_{\mathrm{s}}

, and the BC policy

\pi_{\mathrm{BC}}

Add
	LSTM(B)	LSTM(T)	DLSTM	DT	DMamba	TrXL	ELMUR	KM (Ours)
Train MSE $\downarrow$	$\mathbf{4{\scriptscriptstyle\,\cdot\,10^{-6}}\!\pm\!1{\scriptscriptstyle\,\cdot\,10^{-6}}}$	$0.006\!\pm\!0.003$	$0.167\!\pm\!0.106$	$\mathbf{8.1{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!6.8{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$\mathbf{2.0{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!1.3{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$\mathbf{2.2{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!1.2{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$0.114\!\pm\!0.014$	$\mathbf{3.4{\scriptscriptstyle\,\cdot\,10^{-5}}\!\pm\!3.6{\scriptscriptstyle\,\cdot\,10^{-5}}}$
Test MSE $\downarrow$	$0.008\!\pm\!0.004$	$0.040\!\pm\!0.024$	$0.368\!\pm\!0.186$	$0.148\!\pm\!0.083$	$0.160\!\pm\!0.086$	$0.133\!\pm\!0.064$	$0.157\!\pm\!0.047$	$\mathbf{3.2{\scriptscriptstyle\,\cdot\,10^{-5}}\!\pm\!3.3{\scriptscriptstyle\,\cdot\,10^{-5}}}$
Scattered-Copy
Train Copy SR $\uparrow$	$\mathbf{1.000\!\pm\!0.000}$	$0.100\!\pm\!0.035$	$0.000\!\pm\!0.000$	$0.538\!\pm\!0.303$	$0.000\!\pm\!0.000$	$0.018\!\pm\!0.016$	$0.000\!\pm\!0.000$	$\mathbf{1.000\!\pm\!0.000}$
Test Copy SR $\uparrow$	$0.160\!\pm\!0.021$	$0.106\!\pm\!0.030$	$0.000\!\pm\!0.000$	$0.000\!\pm\!0.000$	$0.000\!\pm\!0.000$	$0.000\!\pm\!0.000$	$0.000\!\pm\!0.000$	$\mathbf{1.000\!\pm\!0.000}$

Appendix

Table 6: Aggregated benchmark results across all five tasks. Cells show mean

\pm

std across

5

seeds for Add, Scattered-Copy, and T-Maze, and

3

seeds for MiniGrid-Memory and LTMB-Hallway. For Add, the metric is final-step MSE (

\downarrow

better); for all others, it is success rate (

\uparrow

better). “—” marks the LSTM (BPTT) variant on the grid tasks, which we did not train because it required significantly more compute than the other baselines; only TBPTT was trained. See Table 1 for train and test sequence lengths and the number of demonstrations

|\mathcal{D}|

used for training.

Figure 10: MIKASA-Robo Environments. A collage of six representative tasks from the MIKASA-Robo memory-intensive robot manipulation benchmark. The benchmark evaluates robot policies across comprehensive memory categories, including object memory, spatial memory, sequential memory, and memory capacity. More details can be found in the reference paper [ 7 ] .

Module	Architecture	$d$	$h$	$L$	$H$	Params
$\pi_{\mathrm{s}}$ (selector)	ResNet-18 + MLP $[500,\,500]$	256	—	—	—	15.80 M
$f_{\theta}$ (proxy)	ResNet-18 + Transformer	512	4	4	8	24.13 M
$\pi_{\mathrm{BC}}^{\text{SST}}$	Regression transformer	1024	16	4	1	62.13 M
$\pi_{\mathrm{BC}}^{\text{FM}}$	Flow-matching cross-attn transformer	1024	16	4	8	120.20 M

Appendix

Table 7: Module architectures, dimensions and parameter counts.

d

: hidden width,

h

: attention heads,

L

: transformer layers,

H

: action-chunk length. The selector and proxy are training-only; only one BC-policy variant is deployed per task.

Stage	Optimizer	LR	Weight decay	Schedule
Selector policy $\pi_{\mathrm{s}}$ (PPO)	Adam	$5\,{\times}\,10^{-5}$	$0$	constant
Proxy policy $f_{\theta}$	AdamW	$1\,{\times}\,10^{-4}$	$0$	constant
BC policy $\pi_{\mathrm{BC}}^{\text{SST}}$	AdamW	$5\,{\times}\,10^{-5}$	$0.05$	constant
BC policy $\pi_{\mathrm{BC}}^{\text{FM}}$	AdamW	$5\,{\times}\,10^{-5}$	$0.05$	constant

Appendix

Table 8: Optimizer settings for the three training stages.

Hyperparameter	Value	Hyperparameter	Value	Hyperparameter	Value
Total environment steps	$5\,{\times}\,10^{6}$	PPO epochs per update	$5$	Discount $\gamma$	$0.99$
Parallel envs ( $n_{\text{envs}}$ )	$32$	Entropy coefficient	$0.01$	GAE $\lambda_{\text{GAE}}$	$0.95$
Rollout length ( $n_{\text{steps}}$ )	$1024$	Clip range	$0.2$	Value-function coef. $c_{\text{vf}}$	$0.5$
Minibatch size	$256$	Reward floor $\lambda$	$10^{-5}$	Max grad norm	$0.5$

Appendix

Table 9: PPO hyperparameters used to train the selector policy

\pi_{\mathrm{s}}

Environment	RATE	DT	BC	CQL	DP	ELMUR	KM (Ours)
Shell Game Touch	0.92 $\pm$ 0.01	0.53 $\pm$ 0.07	0.28 $\pm$ 0.01	0.16 $\pm$ 0.04	0.18 $\pm$ 0.02	0.96 $\pm$ 0.02	0.94 $\pm$ 0.02
Shell Game Push	0.78 $\pm$ 0.06	0.62 $\pm$ 0.14	0.27 $\pm$ 0.01	0.25 $\pm$ 0.01	0.22 $\pm$ 0.03	0.87 $\pm$ 0.04	0.99 $\pm$ 0.01
Shell Game Pick	0.02 $\pm$ 0.01	0.00 $\pm$ 0.00	0.01 $\pm$ 0.01	0.00 $\pm$ 0.00	0.01 $\pm$ 0.00	0.66 $\pm$ 0.01	0.78 $\pm$ 0.08
Intercept Slow	0.23 $\pm$ 0.02	0.40 $\pm$ 0.02	0.37 $\pm$ 0.06	0.25 $\pm$ 0.01	0.33 $\pm$ 0.05	0.61 $\pm$ 0.04	0.48 $\pm$ 0.06
Intercept Medium	0.32 $\pm$ 0.02	0.56 $\pm$ 0.01	0.31 $\pm$ 0.14	0.03 $\pm$ 0.01	0.68 $\pm$ 0.02	0.64 $\pm$ 0.03	0.78 $\pm$ 0.01
Intercept Fast	0.30 $\pm$ 0.04	0.36 $\pm$ 0.04	0.03 $\pm$ 0.02	0.02 $\pm$ 0.02	0.21 $\pm$ 0.05	0.54 $\pm$ 0.02	0.59 $\pm$ 0.04

Appendix

Table 10: Success rate (mean

\pm

sem over 3 seeds, 100 episodes per seed) on the MIKASA-Robo benchmark. Bold marks every method whose mean is within one standard deviation of the best per row.

Add (final-step MSE $\downarrow$ )
	LSTM(B)	LSTM(T)	DLSTM	DT	DMamba	TrXL	ELMUR	KM (Ours)
$h=10$	$\mathbf{5.0{\scriptscriptstyle\,\cdot\,10^{-6}}\!\pm\!2.3{\scriptscriptstyle\,\cdot\,10^{-6}}}$	$3.4{\scriptscriptstyle\,\cdot\,10^{-3}}\!\pm\!3.5{\scriptscriptstyle\,\cdot\,10^{-3}}$	$0.258\!\pm\!0.216$	$\mathbf{6.5{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!6.2{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$\mathbf{2.5{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!1.8{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$\mathbf{3.9{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!2.4{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$0.143\!\pm\!0.114$	$\mathbf{3.1{\scriptscriptstyle\,\cdot\,10^{-5}}\!\pm\!3.4{\scriptscriptstyle\,\cdot\,10^{-5}}}$
$h=20$	$\mathbf{7.7{\scriptscriptstyle\,\cdot\,10^{-5}}\!\pm\!6.9{\scriptscriptstyle\,\cdot\,10^{-5}}}$	$6.8{\scriptscriptstyle\,\cdot\,10^{-3}}\!\pm\!5.6{\scriptscriptstyle\,\cdot\,10^{-3}}$	$0.330\!\pm\!0.366$	$\mathbf{9.2{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!7.7{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$\mathbf{2.1{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!1.9{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$\mathbf{1.8{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!1.4{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$0.131\!\pm\!0.090$	$\mathbf{2.8{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!5.7{\scriptscriptstyle\,\cdot\,10^{-4}}}$
$h=30$	$\mathbf{3.1{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!1.9{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$1.1{\scriptscriptstyle\,\cdot\,10^{-2}}\!\pm\!8.3{\scriptscriptstyle\,\cdot\,10^{-3}}$	$0.267\!\pm\!0.348$	$\mathbf{7.0{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!7.8{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$\mathbf{2.2{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!1.5{\scriptscriptstyle\,\cdot\,10^{-4}}}$	$\mathbf{1.4{\scriptscriptstyle\,\cdot\,10^{-4}}\!\pm\!6.1{\scriptscriptstyle\,\cdot\,10^{-5}}}$	$0.141\!\pm\!0.052$	$\mathbf{7.7{\scriptscriptstyle\,\cdot\,10^{-5}}\!\pm\!1.1{\scriptscriptstyle\,\cdot\,10^{-4}}}$
$h=10^{3}$	$0.246\!\pm\!0.097$	$0.127\!\pm\!0.027$	$0.531\!\pm\!0.455$	$0.408\!\pm\!0.439$	$0.387\!\pm\!0.247$	$0.296\!\pm\!0.216$	$0.797\!\pm\!0.938$	$\mathbf{3.8{\scriptscriptstyle\,\cdot\,10^{-5}}\!\pm\!4.8{\scriptscriptstyle\,\cdot\,10^{-5}}}$
$h=10^{4}$	$0.494\!\pm\!0.320$	$0.351\!\pm\!0.142$	$0.218\!\pm\!0.195$	$0.447\!\pm\!0.398$	$0.594\!\pm\!0.359$	$0.634\!\pm\!0.387$	$0.759\!\pm\!0.465$	$\mathbf{3.7{\scriptscriptstyle\,\cdot\,10^{-5}}\!\pm\!4.0{\scriptscriptstyle\,\cdot\,10^{-5}}}$

Appendix

Table 11: Horizon-invariance results across synthetic and grid benchmarks. Cells report mean

~\pm~

std of the relevant metric: final-step MSE for Add (

\downarrow

better), Copy SR for Scattered-Copy and Episode SR for T-Maze, MiniGrid-Memory and LTMB-Hallway (

\uparrow

better). Each benchmark lists the horizons (or environment size for MiniGrid-Memory) evaluated. “—” marks cells for baselines in extreme-horizon experiments as we run it only for Keyframe Mnemonics, or for the LSTM(BPTT) variant on grid tasks, which were trained only with LSTM(TBPTT).

Figure 11: Hover time is varied across evaluations resulting in catastrophic failure for the baselines, while KM maintains it’s performance even at $250$ timesteps of hover. Training data had 15 timesteps of hover.

$\Delta T$	PTP [ 46 ]	ELMUR [ 14 ]	KM (Ours)
$15$	$0.66\pm 0.03$	$0.61\pm 0.07$	$\mathbf{0.72\pm 0.07}$
$35$	$0.72\pm 0.04$	$0.00\pm 0.00$	$\mathbf{0.69\pm 0.11}$
$55$	$\mathbf{0.73\pm 0.05}$	$0.00\pm 0.00$	$0.69\pm 0.11$
$75$	$0.55\pm 0.07$	$0.00\pm 0.00$	$\mathbf{0.69\pm 0.08}$
$95$	$0.43\pm 0.06$	$0.00\pm 0.00$	$\mathbf{0.65\pm 0.08}$
$115$	$0.31\pm 0.05$	$0.00\pm 0.00$	$\mathbf{0.66\pm 0.09}$

Appendix

Table 12: RememberColor3 -

\Delta T

success rate under cue-to-action delays. Training delay is

\Delta T=15

; all larger values are out-of-distribution. Mean

~\pm~

sem over

3

seeds,

100

episodes per seed.

$\Delta T$	Learned selector	Random selector ( $\pi_{\mathrm{rnd}}$ )
15	72.3 $\pm$ 5.8	31.3 $\pm$ 3.1
35	69.0 $\pm$ 9.2	28.7 $\pm$ 2.9
55	69.0 $\pm$ 8.8	28.3 $\pm$ 2.1
75	68.7 $\pm$ 6.6	35.3 $\pm$ 1.2
95	65.0 $\pm$ 6.4	28.7 $\pm$ 4.0
115	66.3 $\pm$ 7.6	28.3 $\pm$ 4.1

Appendix

Table 13: Learned selector against the random selector

\pi_{\mathrm{rnd}}

on MIKASA-Robo RememberColor3-

\Delta T

. Reported mean

\pm

sem success rate (%)

\uparrow

Domain	Metric	$\gamma=0$	$\gamma=0.5$	$\gamma=0.9$	$\gamma=0.99$ (default)
T-Maze (active)	Episode SR $\uparrow$	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0
T-Maze (passive)	Episode SR $\uparrow$	48.8 $\pm$ 6.6	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0
Add	Final MSE $\downarrow$	1.84e-01 $\pm$ 1.77e-03	1.25e-02 $\pm$ 6.86e-03	1.28e-05 $\pm$ 5.67e-06	1.28e-05 $\pm$ 5.70e-06
Scattered-Copy	Copy SR $\uparrow$	0.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0

Appendix

Table 14: Effect of the PPO discount factor

\gamma

. Reported mean

\pm

sd.

Domain	Metric	$\kappa=0.02$	$\kappa=0.2$	$\kappa=2$	off
T-Maze	Episode SR $\uparrow$	0.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0
Add	Final MSE $\downarrow$	1.28e-05 $\pm$ 5.7e-06	1.28e-05 $\pm$ 5.7e-06	1.84e-01 $\pm$ 1.8e-03	1.84e-01 $\pm$ 1.8e-03
Scattered-Copy	Copy SR $\uparrow$	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	0.0 $\pm$ 0.0
MiniGrid-Memory	Success rate $\uparrow$	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	3.0 $\pm$ 0.0	3.7 $\pm$ 0.6
LTMB-Hallway	Success rate $\uparrow$	0.0 $\pm$ 0.0	0.0 $\pm$ 0.0	0.0 $\pm$ 0.0	99.7 $\pm$ 0.6
RememberColor3	Success rate $\uparrow$	25.0 $\pm$ 3.3	29.3 $\pm$ 6.2	34.7 $\pm$ 6.9	99.0 $\pm$ 0.0

Appendix

Table 15: Effect of the surprise coefficient across temperatures. Bold marks the best setting per domain. Reported mean

\pm

sd.

Domain	Metric	Proxy without $o_{t}$ (default)	Proxy with $o_{t}$
T-Maze	Episode SR $\uparrow$	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0
Add	Final MSE $\downarrow$	1.28e-05 $\pm$ 5.7e-06	2.20e-05 $\pm$ 1.3e-05
Scattered-Copy	Copy SR $\uparrow$	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0
MiniGrid-Memory	Success rate $\uparrow$	100.0 $\pm$ 0.0	33.3 $\pm$ 26.3
LTMB-Hallway	Success rate $\uparrow$	99.7 $\pm$ 0.6	93.0 $\pm$ 9.5
RememberColor3	Success rate $\uparrow$	99.0 $\pm$ 0.0	18.7 $\pm$ 5.9

Appendix

Table 16: Effect of including the current observation

o_{t}

in the proxy input. Reported mean

\pm

sd.

Domain	Metric	FIFO $\rightarrow$ FIFO	FIFO $\rightarrow$ LIFO (default)	LIFO $\rightarrow$ FIFO	LIFO $\rightarrow$ LIFO
T-Maze	Episode SR $\uparrow$	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0
Add	Final MSE $\downarrow$	1.28e-05 $\pm$ 5.7e-06	1.28e-05 $\pm$ 5.7e-06	1.28e-05 $\pm$ 5.7e-06	1.28e-05 $\pm$ 5.7e-06
Scattered-Copy	Copy SR $\uparrow$	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0
MiniGrid-Memory	Success rate $\uparrow$	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0
LTMB-Hallway	Success rate $\uparrow$	99.6 $\pm$ 0.5	99.6 $\pm$ 0.5	99.6 $\pm$ 0.5	99.6 $\pm$ 0.5
RememberColor3	Success rate $\uparrow$	93.7 $\pm$ 4.5	99.0 $\pm$ 0.0	98.7 $\pm$ 1.3	96.0 $\pm$ 2.2

Appendix

Table 17: Queue tie-break rule, train

\rightarrow

inference. Reported mean

\pm

sd.

$\Delta T$	FIFO $\rightarrow$ LIFO (default)	LIFO $\rightarrow$ FIFO
15	72.3 $\pm$ 5.8	75.0 $\pm$ 1.6
35	69.0 $\pm$ 9.2	66.0 $\pm$ 3.6
55	69.0 $\pm$ 8.8	66.0 $\pm$ 2.4
75	68.7 $\pm$ 6.6	59.3 $\pm$ 3.7
95	65.0 $\pm$ 6.4	59.3 $\pm$ 1.7
115	66.3 $\pm$ 7.6	58.3 $\pm$ 2.5

Appendix

Table 18: Queue tie-break rule on RememberColor3-

\Delta T

. Reported mean

\pm

sem success rate

\uparrow

Domain	Metric	$0.5\times$	$1\times$ (default)	$2\times$
T-Maze	Episode SR $\uparrow$	40.0 $\pm$ 23.2	100.0 $\pm$ 0.0	100.0 $\pm$ 0.0
Add	Final MSE $\downarrow$	9.55e-02 $\pm$ 2.38e-02	1.28e-05 $\pm$ 5.7e-06	2.33e-05 $\pm$ 2.03e-05
Scattered-Copy	Copy SR $\uparrow$	0.0 $\pm$ 0.0	100.0 $\pm$ 0.0	80.0 $\pm$ 44.7
MiniGrid-Memory	Success rate $\uparrow$	35.3 $\pm$ 26.4	100.0 $\pm$ 0.0	20.7 $\pm$ 27.2
LTMB-Hallway	Success rate $\uparrow$	0.0 $\pm$ 0.0	81.7 $\pm$ 15.0	99.3 $\pm$ 1.2
RememberColor3	Success rate $\uparrow$	94.7 $\pm$ 2.5	99.0 $\pm$ 0.0	100.0 $\pm$ 0.0

Appendix

Table 19: Effect of scaling the memory buffer size

k

. Reported mean

\pm

sd.

Figure 12: Trained selector recovers the keyframe oracle as a near-optimal function. Per-step selector priorities (bottom row of each pair) are compared against the oracle keyframe indicator (top row of each pair) on three synthetic memory benchmarks. Each benchmark is evaluated at three increasing episode horizons, where the longest horizon extends one order of magnitude beyond the training horizon. For each panel, we render $10$ evaluation episodes (rows) over time (columns), where each column corresponds to a single environment step and color encodes either the selector output ( $p_{t}\in[0,1]$ ) or the binary oracle indicator ( $1$ for keyframes and $0$ otherwise). The oracle definition is problem-specific, as discussed in Appendix C.1 . Across all cells, the selector priority agrees with the oracle. After squashing the raw selector outputs to $[0,1]$ , the mean selector priority is exactly $1.00$ on keyframes and $0.00$ on background frames. This empirical result shows that the selector policy $\pi_{\mathrm{s}}$ learns a near-optimal keyframe selection function for domains where our assumptions hold strongly.

Figure 13: BC policy activation map on mnemonics. Grad-CAM [ 65 , 66 ] activation maps of the behavior cloning policy $\pi_{\mathrm{BC}}$ on buffer images at decision-critical steps, where hot regions indicate inputs driving the predicted action. Frames with boundaries represent discovered keyframes, while the frame with a boundary shows the agent’s current observation. Across both grid and high-dimensional visual domains, the BC policy uses task-specific cue information from the discovered keyframes to generate actions. This highlights the utility of mnemonics for disambiguating policy actions under partial observability. See https://keyframe-mnemonics.github.io for activation map videos over policy rollouts.

Figure 14: Memory regret as a function of buffer capacity

k

. The

1\times

setting corresponds to the proxy training memory budget

k

in Table 2 . (a) Active regret decreases as buffer capacity becomes sufficient to retain the true keyframes, with domain-specific behavior when

k

is below the true keyframe count. (b) Passive regret remains low under insufficient capacity and decreases again for larger buffers, where random selection is more likely to retain keyframes.

Figure 15: Selector training PPO learning curves. The selector policy converges within $200$ K environment steps for synthetic domains and $5$ M steps for robot manipulation domains. For tasks where optimization is difficult, e.g., ChainOfColors , the learning curve is noisy and shows modest reward improvement, indicating that some of our assumptions may not hold strongly in the environment.

Explore similar work

Sep 24, 2026cs.RO

Training-free Behavior Cloning

Neural behavior cloning compresses demonstrations into large models, making individual actions difficult to trace and policy updates costly. Retrieval policies retain access to demonstrations but struggle with mismatch between recorded and live behavior. We introduce Behavior Predictive Control (BPC), which synthesizes policies without end-to-end policy training by combining an action-aware retrieval metric, a Hankel-based action-continuation prior, and a closed-form one-step residual correction. Inspired by behavioral systems theory, BPC predicts future actions by blending stored observation-action data that best reconstructs the recent runtime observation--action history. Across simulated benchmarks and real-robot deployments, BPC is competitive with learned policies such as

π_{0.5}

(surpassing it in some cases), while reducing policy fitting from hours to seconds on consumer GPUs and supporting closed-loop control upwards of 75 Hz on a Jetson Orin Nano. The retrieved demonstration windows and their coefficients also provide an intrinsic estimate of task progress. Retaining demonstrations within the deployed policy makes its predictions traceable to supporting trajectories and enables behavior revision through the demonstration bank.

Maximilian Adang, Timothy Chen, Lars Osterberg +2

Stanford University, Stanford, CA 94404, USA.

Robot Policy Learning Behavior Cloning

Ask PDF

Jun 8, 2026cs.RO

Difference-Aware Retrieval Policies for Imitation Learning

Parametric imitation learning via behavior cloning can suffer from poor generalization to out-of-distribution states due to compounding errors during deployment. We show that reusing the training data during inference via a semi-parametric retrieval-based imitation learning approach can alleviate this challenge. We present Difference-Aware Retrieval Policies for Imitation Learning (DARP), a semi-parametric retrieval-based imitation learning approach that addresses this limitation by reparameterizing the imitation learning problem in terms of local neighborhood structure rather than direct state-to-action mappings. Instead of learning a global policy, DARP trains a model to predict actions based on

k

-nearest neighbors from expert demonstrations, their corresponding actions, and the relative distance vectors between neighbor states and query states. DARP requires no additional assumptions beyond those made for standard behavior cloning -- it does not require additional data collection, online expert feedback, or task-specific knowledge. We demonstrate consistent performance improvements of 15-46% over standard behavior cloning across diverse domains, including continuous control and robotic manipulation, and across different representations, including high-dimensional visual features. Code and demos are available at https://weirdlabuw.github.io/darp-site/.

Quinn Pfeifer, Ethan Pronovost, Paarth Shah +3

1Paul G. Allen School of Computer Science & Engineering, University of Washington · 2Toyota Research Institute · 3Google DeepMind +1

Imitation Learning Interactive Imitation Learning Robot Imitation Learning Behavior Cloning Robot Policy Generalization Policy Learning

Ask PDF

Jun 25, 2026cs.RO

Scalable Behavior Cloning with Open Data, Training, and Evaluation

We introduce ABC, a fully open-source stack for manipulation with behavior cloning. At its core is ABC-130K: the largest open-source teleoperation dataset to date, featuring 3,500 hours of data spanning over 130K episodes across 195 diverse tasks. Furthermore, we open-source our accessible hardware setup, training infrastructure, and simulation pipeline. We also release 400 hours of sim-teleop data and provide a co-training recipe that produces correlated simulation and real-world evaluation, offering a reliable proxy for ablating model-design and training decisions before costly real-world evaluation. We explore various training recipes and compare common architectural choices for Diffusion Transformers (DiT) and Vision-Language-Action (VLA) models, grounding our findings in real-world evaluations. The resulting policies successfully execute dexterous tasks such as box folding and extracting credit cards from wallets. By providing a reproducible toolkit, we aim to place researchers on an equal footing, establishing the necessary foundation to learn the ABCs of Behavior Cloning together as a community.

Arthur Allshire, Himanshu Gaurav Singh, Ritvik Singh +15

UC Berkeley · MIT · XDOF +2

Robot Policy Learning Diffusion Transformer Robotic Data Collection Vision-Language-Action Models Behavior Cloning

Ask PDF