Paper deep dive

Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning

Prabin Kumar Rath, Omkar Patil, Nakul Gopalan

Year: 2026Venue: arXiv preprintArea: cs.AIType: PreprintEmbeddings: 196

Intelligence

Status: succeeded | Model: Gemma-4-26B-A4B | Prompt: intel-v1 | Confidence: 92%

Last extracted: 10/9/2026, 7:58:11 AM

Summary

The paper introduces Keyframe Mnemonics (KM), a self-supervised method for Behavior Cloning (BC) in non-Markovian environments. KM discovers information-critical observations (keyframes/mnemonics) by training a proxy policy with masked-sampling to exploit simplicity bias, then using the proxy's loss as a reward to train a selector policy via PPO. The final BC policy conditions on these keyframes, achieving horizon-invariant performance with 100% success rates in synthetic domains and significant improvements on robot manipulation benchmarks.

Entities (7)

Keyframe Mnemonics · method · 98%Behavior Cloning · task · 95%Proxy Policy · model-component · 92%Selector Policy · model-component · 92%MIKASA-Robo · benchmark · 90%PPO · algorithm · 90%Simplicity Bias · concept · 88%

Relation Signals (6)

Keyframe Mnemonics → solves → Behavior Cloning

confidence 95% · we propose Keyframe Mnemonics... We then train a BC policy

Keyframe Mnemonics → uses → Proxy Policy

confidence 92% · Proxy policy training... on which the proxy is trained

Keyframe Mnemonics → uses → Selector Policy

confidence 92% · The selector is trained using PPO

Keyframe Mnemonics → evaluatedon → MIKASA-Robo

confidence 90% · evaluate on memory-intensive robot manipulation benchmark MIKASA-Robo

Selector Policy → trainedwith → PPO

confidence 90% · The selector is trained using PPO

Keyframe Mnemonics → leverages → Simplicity Bias

confidence 88% · leverages the intrinsic simplicity bias of neural networks

Cypher Suggestions (0)

No Cypher suggestions yet.

Abstract

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 this https URL.

PDF

Open source PDF →Open local PDF →

Trouble viewing inline? Open PDF directly →

Full Text

193,068 characters extracted from source content.

Expand or collapse full text

Self-Supervised Keyframe Discovery for Horizon-Invariant Behavior Cloning Prabin Kumar Rath Affiliation: Arizona State University Email: prath4@asu.edu Omkar Patil Affiliation: Arizona State University Email: opatil3@asu.edu Nakul Gopalan Affiliation: Arizona State University Email: ng@asu.edu 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 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%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× longer horizons on a real robot. Code and videos are available at https://keyframe-mnemonics.github.io. 1 Introduction The Markov assumption for decision processes is that the agent’s current state contains sufficient information to determine the optimal action with respect to its goal or objective [1]. In practice, however, real-world environments are non-Markovian and partially observable, where decision-making often requires retaining (“remembering”) information over time to provide context for future actions [2, 3]. Recent advancements in policy learning have increasingly relied on behavior cloning (BC) [4, 5, 6], in which policies are highly dependent on short observation contexts with a working memory spanning only a few seconds into the past. This makes such BC policies brittle on tasks that require remembering information cues over long horizons [7]. Recurrent models have been the architecture of choice for behavior cloning in such memory intensive partially observable environments and are typically trained using Back-Propagation Through Time (BPTT) [8, 9]. However, BPTT is difficult to scale, as it requires sequential gradient accumulation over episode timesteps [10]. Truncated-BPTT tries to mitigate this issue but trades off information from distant past observations [11]. Moreover, recurrent models suffer from hidden-state collapse due to repeated representation updates over time [12, 10]. Attention-based architectures such as transformers [13] address many representation learning challenges, but their finite context length limits retention of the full history in partially observable domains. Furthermore, increasing context through frame stacking incurs quadratic compute complexity [13], degrades at extreme horizons [14], and is prone to overfitting spurious correlations with visual observations [15]. Many hybrid architectures have been proposed that improve information retention through dedicated memory mechanisms [14, 16, 17, 18, 19, 20, 21] or by learning recurrent summaries [22, 23, 24]. However, these methods remain susceptible to state collapse beyond the horizons encountered during training, often leading to severe performance degradation on tasks that require generalization to longer horizons. Figure 1: Keyframe Mnemonics three-stage training for horizon-invariant behavior cloning. (1) Proxy policy training. A random binary mask Ψk _k is applied to the observation history hth_t to assemble a memory buffer ℬtB_t of size at most k, on which the proxy fθ:ℬ→f_θ:B\!→\!A is trained to predict the expert action ata_t under a negative log-likelihood loss ℒL. Combinatorial support of training data (Lemma 1) biases fθf_θ toward keyframe-based hypotheses (Theorem 1, Proposition 1), so the frozen objective Jθ∗(ℬt)=ℒ⁡(fθ∗(ℬt),at)J_θ (B_t)=L(f_θ^*(B_t),a_t) scores keyframe-rich buffers favourably (Corollary 1; Section 2.1). (2) Selector policy training. The selector πs _s is trained using PPO to populate ℬtB_t with keyframes using Jθ∗(ℬt)J_θ (B_t) as reward function (Section 2.2). (3) Behavior cloning. The trained selector is used to construct a keyframe augmented dataset †D , which is then used to train the BC policy πBC(at∣ot,ℬt) _BC(a_t o_t,B_t) (Section 2.3). Throughout the figure, blue tiles denote the current observation oto_t, orange tiles denote past frames in hth_t, hatched tiles are frames not selected into ℬtB_t. To address these challenges in behavior cloning, we propose Keyframe Mnemonics (KM), a policy learning method that decomposes the original non-Markovian decision-making problem into two subproblems: (i) discovering observations (keyframes) that contain information relevant to the action distribution, and (i) training a BC policy on demonstrations augmented with the identified keyframes. Our keyframe discovery process is self-supervised and leverages the intrinsic simplicity bias of neural networks [25, 26, 27] to train a reward function that assigns high rewards to information-rich keyframes in the policy’s context buffer. The reward learning process is driven by the mutual information advantage of keyframes with the expert action distribution, which is amplified through a novel training method that uses randomly sampled past observations as training samples for the reward function. We then train a selector policy using reinforcement learning to optimize the learned reward function and identify a set of decision-relevant keyframes, which we call Mnemonics (“something that assists memory”). Finally, we train BC policies conditioned on these mnemonics to match the expert action distribution. The trained policy is deployed together with the selector, which identifies keyframes at runtime and provides context to disambiguate actions under partial observability. Under certain assumptions, this architectural decomposition enables horizon-invariance i.e., memory retention beyond the training horizon, allowing BC policies to generalize on longer time scales where existing hybrid memory methods fail. We extensively evaluate Keyframe Mnemonics on synthetic memory benchmarks (Add [28], Scattered Copy [22], T-Maze [29]), grid worlds (MiniGrid Memory [30], LTMB-Hallway [31]), memory-intensive robot manipulation benchmark MIKASA-Robo [7], and additionally on a real robot, demonstrating stable performance without degradation far beyond the training horizon. Our contributions are: • We propose a novel self-supervised algorithm to discover decision-relevant keyframes in non-Markovian expert demonstrations without explicit keyframe annotations (Figure 1, Algorithm 1). • We provide a theoretical intuition of our method under task-structure assumptions, derive information retention guarantees, and support our hypotheses with empirical evidence. • We demonstrate horizon-invariant performance of mnemonic-conditioned BC policies, with MLP policies achieving 100%100\% SR in synthetic and grid memory domains, Transformer policies improving absolute SR by 13.9%13.9\% on the MIKASA-Robo benchmark, and a flow-matching policy retaining 80%80\% SR at 20×20× longer test horizons on a real robot. 2 Methodology We assume a partially observable Markov decision process (POMDP) [2] ℳ=(,,,,Ω,ρ0)M=(S,A,O,T, , _0), where ρ0∈Δ⁡() _0∈ (S) denotes the initial state distribution. At timestep t∈1,…,Tt∈\1,…,T\ with horizon T∈ℕ+T ^+, the unobserved state is denoted by st∈s_t , the agent receives an observation ot∈o_t , and takes an action at∈a_t based on a belief over the state, bt=P⁡(st∣ht)b_t=P(s_t h_t), where ht=o1,…,oth_t=\o_1,…,o_t\ denotes the history of observations up to time t. Here, ⁡(st+1∣st,at)T(s_t+1 s_t,a_t) and Ω⁡(ot+1∣st+1,at) (o_t+1 s_t+1,a_t) denote the transition dynamics and observation function respectively. We omit the reward function ℛR and discount factor γ from the notation, as they are not available in standard behavior cloning. Additionally, we assume access to an expert demonstration dataset =τii=1ND=\ _i\_i=1^N, where each trajectory is given by τi=(ot,at)t=1Ti _i=\(o_t,a_t)\_t=1^T_i. We make the following structural assumptions over this dataset, where I⁡(⋅,⋅)I(·;·) and H⁡(⋅)H(·) denote mutual information and entropy, respectively. Assumption 1 (Approximate belief sufficiency). The trajectory observations admit a partition into two disjoint subsets: a sparse subset keyX_key of size C≪TC T (keyframes) and the remaining T−CT-C observations bgX_bg (background frames). The keyframes key,t⊂htX_key,t⊂ h_t up to time t contain observations with invariant task information and causal cues that remain consistent across the demonstration dataset, thereby forming an approximate sufficient statistic for the belief state btb_t. In contrast, the background subset bg,t⊂htX_bg,t⊂ h_t contributes minimal information toward the belief state. This assumption imposes a sparse keyframe structure on observations in the demonstration dataset, which we empirically observe in several memory intensive partially observable domains evaluated in Section 3. Assumption 2 (Action disambiguation by keyframes). Any ambiguity in the action distribution is primarily captured by the keyframes keyX_key, with the following relationship in the mutual information I⁡(at,key,t)>I⁡(at,bg,t)I(a_t;X_key,t)>I(a_t;X_bg,t). That is, keyframes strictly provide more information for disambiguating the action than the background frames. Section 3 provides empirical evidence consistent with this assumption in the evaluated memory-intensive POMDPs. We formally define Mnemonics as the set of keyframes up to time t, denoted by key,t⊂htX_key,t⊂ h_t, which provide sufficient information for a behavior cloning policy to generate actions unambiguously. Our goal is to learn a selector policy πs _s that identifies keyframes from a continuous stream of observations and adds them to a memory buffer ℬtB_t, maximizing its overlap with key,tX_key,t. These mnemonics are then used to condition a behavior cloning policy πBC _BC that models the expert action distribution. 2.1 Leveraging Simplicity-bias for Keyframe Discovery We formulate keyframe discovery as a self-supervised buffer optimization problem. Specifically, we define an objective J⁡(⋅,f)J(·,f) over a proxy policy f mapping observation frames in the buffer, which collectively represent the belief state, to expert actions, f:ℬ→f:B , and optimize the buffer ℬB subject to a memory budget k much smaller than the observation history (k≪Tk T). For J⁡(⋅,f)J(·,f) to serve as an effective objective, it must provide a stronger learning signal (i.e., lower objective value) when ℬB contains a higher proportion of keyframes than background frames. In this section, we introduce a training strategy for the proxy policy f that increases the complexity of the input–output mapping, thereby exploiting the simplicity bias of deep neural networks to encourage learning a hypothesis that favors keyframes over background frames due to their greater informational content. The resulting biased model can then be used to learn a selector πs _s for keyframe discovery. Proxy Policy Training: To train the proxy policy fθf_θ, a neural network parameterized by θ, we construct input buffers using a random binary mask Ψk∈0,1|ht| _k∈\0,1\^|h_t| that independently selects past observations from the history hth_t, with cardinality bounded by the memory budget k. We call this masked-sampling past observations that constructs the training-time memory buffer ℬt=oi∈ht∣Ψk(i)=1B_t=\o_i∈ h_t _k(i)=1\, which serves as input to the proxy policy fθ:ℬ→f_θ:B with |ℬt|≤k|B_t|≤ k, where ℬB denotes the space of all possible memory buffers. Algorithm 1 shows the masked-sampling method for training the proxy policy fθf_θ. Treating (ℬt,at)(B_t,a_t) as i.i.d. samples, we now analyze the hypothesis space ℋH explored by fθf_θ during training through the lens of Kolmogorov complexity [32], denoted K⁡(f)K(f) for f∈ℋf 11 1 For notational simplicity, we use the same variable f to denote both the proxy policy and the hypothesis learned by the proxy policy, and use Kolmogorov complexity as a measure of hypothesis description length [33]. . Let ℋ∗⊂ℋH^* denote the set of interpolating hypotheses achieving zero empirical risk on the training dataset induced by masked-sampling. Assume ℋ∗H^* partitions into an invariant class ℋinH_in, corresponding to hypotheses conditioned only on keyframes, and a non-invariant class ℋ¬inH_ , which depend on background frames. Let fkey∈ℋinf_key _in denote keyframe-based and fbg∈ℋ¬inf_bg _ denote background-based hypotheses defined by mappings fkey:key→f_key:X_key and fbg:bg→f_bg:X_bg , where the inputs are constrained to at most k frames, consistent with fθf_θ. Algorithm 1 Masked-sampling training of the proxy policy fθf_θ 1: dataset D, training memory budget k, proxy policy fθf_θ 2: for each training step do 3: Sample τ∼τ and timestep t∈1,…,|τ|t∈\1,…,|τ|\ ⊳ uniform sampling of timestep 4: n←min⁡(Binomial⁡(t, 0.5),k)n← \! (Binomial(t,\,0.5),\;k ) ⊳ buffer occupancy under budget k 5: Draw ℐ⊂1,…,tI⊂\1,…,t\ with |ℐ|=n|I|=n ⊳ uniform sampling without replacement 6: Ψk(i)←[i∈ℐ] _k(i) 1[\,i \,] for i=1,…,ti=1,…,t ⊳ induced binary mask on hth_t 7: ℬt←oi∈ht∣Ψk(i)=1B_t←\\,o_i∈ h_t _k(i)=1\,\ ordered by i ⊳ |ℬt|≤k|B_t|≤ k 8: a^t←fθ(ℬt) a_t← f_θ(B_t) ⊳ predict expert action 9: θ←θ−η∇θℒ(a^t,at)θ←θ-η\, _θ\,L( a_t,\,a_t) ⊳ update proxy using NLL loss 10: end for Overview: First, we show that masked-sampling of past observations during training of fθf_θ induces combinatorially large support of training pairs (Lemma 1), which artificially increases the complexity of background-based hypotheses fbgf_bg and effectively restricts the feasible hypothesis space (Theorem 1). Assuming that deep neural networks exhibit an implicit bias toward simpler hypotheses, we show that keyframe-based hypotheses fkeyf_key are learned with high probability (Proposition 1). Finally, we demonstrate that the prediction error of the trained proxy policy provides an optimizable signal to identify keyframe observations (Corollary 1). Lemma 1 (Combinatorial cardinality of the sample space). The total cardinality M of the space of distinct timestep-indexed training pairs (ℬt,at)(B_t,a_t) is given by: M=∑i=1N(Ti+∑j=1k(Ti+1j+1)).M= _i=1^N (T_i+ _j=1^k T_i+1j+1 ). (1) Masked-sampling induces a combinatorially large finite support of training pairs, whose cardinality M admits a closed-form expression via the Hockey-stick identity [34]. (Proof in Appendix A). □ Assumption 3 (Realizability of conditional programs by interpolating hypotheses). For Z∈key,bgZ∈\key,bg\, let fZ∗=argminf∈ℋZ:f⁡(Z,ℓ)=Aℓ∀ℓK(f)f_Z^*= _f _Z:\,f(X_Z, )=A_ \ ∀ K(f) denote the minimum-complexity interpolating hypothesis using input source Z. We assume that the shortest conditional program mapping the training input sequence Z,1:MX_Z,1:M to the action sequence A1:MA_1:M is realizable by the hypothesis class ℋH up to an additive constant: |K(fZ∗)−K(A1:M∣Z,1:M)|≤cZ, |K(f_Z^*)-K(A_1:M _Z,1:M) |≤ c_Z, where cZc_Z is an overhead constant independent of the sample size M. This assumption combines the standard interpretation of conditional Kolmogorov complexity as the shortest program mapping inputs to outputs [35, 32] with the finite-sample expressivity of overparameterized neural networks, which can interpolate or memorize arbitrary training labels [36, 37]. Theorem 1 (Expected gap in hypothesis complexity). Under the assumption that training pairs are sampled i.i.d. from a sufficiently large support of size M (Lemma 1), the expected Kolmogorov complexity of the minimum-complexity background-based hypothesis fbg∗f^*_bg is strictly greater than that of the keyframe-based hypothesis fkey∗f^*_key. ⁡[K⁡(fbg∗)]>⁡[K⁡(fkey∗)].E[K(f_bg^*)]>E[K(f_key^*)]. (2) Using results from [35, 32] (Assumption 3), we show that, under the stated structural assumptions, the complexity gap between background-based and keyframe-based hypotheses grows linearly with the masking-induced support size M according to M⁡(I⁡(A,key)−I⁡(A,bg))M (I(A;X_key)-I(A;X_bg) ). Thus, any information advantage of keyframes is amplified by the large support size M, rendering background-based hypotheses intractable for the proxy policy fθf_θ. (Proof in Appendix A). □ Proposition 1 (Conditional simplicity-bias preference). Under the simplicity-bias assumption and the complexity gap established in Theorem 1, Stochastic Gradient Descent (SGD) is biased toward the invariant hypothesis class ℋinH_in. Given a masking-induced sample distribution with support size M (Lemma 1), the probability of selecting an invariant hypothesis is approximated by PSGD(ℋin∣)≈11+2−[M⁡(I⁡(A,key)−I⁡(A,bg))−⁡(1)].P_SGD(H_in )≈ 11+2^- [M (I(A;X_key)-I(A;X_bg) )-O(1) ]. (3) Neural networks and SGD are biased toward Kolmogorov-simple functions [26, 38, 27]. Hence, the positive hypothesis complexity gap from Theorem 1 drives PSGD(ℋin∣)P_SGD(H_in ) close to 11, biasing fθf_θ toward learning an invariant keyframe-based hypothesis. (Proof in Appendix A). □ Corollary 1 (Keyframe recovery from the frozen loss objective). Let us define a subset-selection objective function Jθ∗(ℬ)=⁡[ℒ⁡(fθ∗(ℬ),A)]J_θ^*(B)=E\! [L(f_θ^*(B),A) ] (4) over a frozen predictor fθ∗f_θ^* obtained under Proposition 1 after training with negative log-likelihood loss ℒL, where fθ∗f_θ^* represents a near-optimal invariant hypothesis in ℋinH_in. Consider candidate buffer subsets k=ℬt⊆ht:|ℬt|≤kC_k=\B_t h_t:|B_t|≤ k\, and let ℬopt∈arg⁡minℬ∈kJθ∗(ℬ)B_opt∈ _B _kJ_θ^*(B). If keyframe subsets ℬ⊂keyB _key have a positive action-information margin over all non-keyframe candidate subsets, and this margin dominates the approximation error of the frozen predictor fθ∗f_θ^*, then the loss objective Jθ∗(ℬ)J_θ^*(B) recovers the keyframes up to information-equivalent subsets. In particular, if keyX_key is unique under the memory budget k, then ℬopt=keyB_opt=X_key. Optimizing Jθ∗(ℬ)J_θ^*(B) is equivalent to minimizing the information gap I⁡(A,key)−I⁡(A,ℬ)I(A;X_key)-I(A;B) over kC_k. Under Assumption 2, this drives ℬB toward keyX_key. (Proof in Appendix A). □ In summary, we assume a structural partition of dataset observations based on a mutual information gap with the action distribution and propose a training procedure to learn an objective function that can be optimized to discover informative keyframes. This mechanism is analogous to data augmentation along the temporal dimension, where our masking strategy increases the difficulty of learning spurious correlations and biases the model toward invariant temporal mappings. Theorem 1 provides theoretical intuition for this phenomenon; see Section 3.1, where we provide empirical evidence across synthetic benchmark domains. 2.2 Selector Optimization for Keyframe Identification Given the learned objective function, we train a selector policy πs _s to identify keyframes. The selector serves two purposes: (i) extract keyframes from the demonstration dataset for constructing mnemonics that can be used to train a downstream policy πBC _BC, and (i) identify keyframes from a continuous stream of observations during inference to populate the memory buffer for πBC _BC. However, note that the objective Jθ∗(ℬ)J_θ^*(B) is defined over buffer inputs rather than individual observations. Furthermore, the number of possible buffer states is combinatorially large (Lemma 1), making exhaustive evaluation non-trivial. To address this, we formulate keyframe identification as a deterministic Markov Decision Process (MDP) over the demonstration dataset, modeling the buffer ℬB as a priority queue and using a selector policy πs _s to assign priority scores that update the buffer. This reformulates the problem from optimizing over buffer states to learning a selector policy that prioritizes keyframe observations. The selection MDP is defined by the tuple ⟨,,,R,γ⟩ ,P,K,R,γ over offline demonstration trajectories τi∈ _i , where τi=(ot,at)t=1Ti _i=\(o_t,a_t)\_t=1^T_i. At each timestep t∈1,…,Tit∈\1,…,T_i\, the selector policy πs(pt∣ot) _s(p_t o_t) receives an observation ot∈o_t and outputs a priority action pt∈p_t , which is used to update the buffer ℬB, modeled as a priority queue with memory budget k. The MDP tuple is specified as follows. State (O), is the current observation oto_t from a trajectory τi∈ _i . An Action (P) is a priority score pt∈[0,1]p_t∈[0,1]. Transition (K): The transition is deterministic, with the next state given by the subsequent observation ot+1o_t+1 in τi _i. The buffer is updated as ℬt=ℬt−1(ot,pt)B_t=B_t-1(o_t,p_t) according to an eviction rule that determines whether oto_t is inserted into or rejected from the buffer. We use a Discount Factor (γ) =0.99=0.99. The Reward (R) function depends on the updated buffer ℬtB_t and a normalization constant λ∈(0,1)λ∈(0,1) (typically a value close to 00). R⁡(ℬt)=Ti⋅eJθ∗j=t/κ∑j=1TieJθ∗j/κ⏟surprise⋅log⁡(Jθ∗(ℬt)+λ)log⁡(λ)⏟reward.R(B_t)= T_i· e^J^j=t_θ^*/κ _j=1^T_ie^J^j_θ^*/κ_surprise· (J_θ^*(B_t)+λ ) (λ )_reward. (5) We use a normalized logarithmic reward which increases sharply toward 11 as the objective value approaches 00, providing a strong optimization signal even for small changes in the memory buffer. This formulation encourages πs _s to assign higher priorities to observations that minimize the learned objective. The reward is further scaled by a surprise [22] coefficient weighted by episode length, which captures action uncertainty. Intuitively, this assigns higher reward weight to timesteps with larger objective values, which corresponds to higher action entropy under the proxy policy. Such timesteps are typically decision-critical, where actions depend strongly on the memory buffer. Hence, emphasizing them focuses optimization on resolving keyframes where they matter most. The surprise coefficients are precomputed for each trajectory using a temperature (κ)-scaled softmax over the expected objective values across timesteps. The expected objective values are computed over a few epochs of masked-sampled input buffers by evaluating the objective Jθ∗(ℬ)J_θ^*(B) at each timestep in the demonstration dataset D. We optimize πs _s using Proximal Policy Optimization (PPO) [39], which enables stable training while promoting exploration through entropy regularization. The memory buffer employs different priority strategies during training and inference. During training, a FIFO queue penalizes the selector policy for assigning low priority to observations that have long-term influence on actions, whereas during inference, a LIFO queue preserves high-priority observations, ensuring their retention over long horizons. See Appendix B.1 for algorithms on priority queue eviction rules used during training and inference. 2.3 Training BC Policy with Keyframe Mnemonics A trained selector policy is used to generate mnemonics, represented by the buffer state ℬtB_t with |ℬt|≤k|B_t|≤ k, which approximates the keyframes up to time t, i.e., key,t⊂htX_key,t⊂ h_t. We now present a conditional guarantee for context retention in memory buffer assuming an optimal selector, which enables mnemonics-conditioned BC policies to generalize to an infinite horizon (Theorem 2). Theorem 2 (Infinite-horizon keyframe retention). Assume a POMDP in which the expert behavior strictly satisfies Assumption 1 and 2. Let πs∗ _s^* be an optimal selector minimizing the frozen objective Jθ∗(ℬt)J_θ^*(B_t) (Corollary 1), and suppose a unique keyframe set keyX_key exists that has the highest mutual information with the action and fits within the memory budget k. Then, as T→∞T→∞, πs∗ _s^* assigns higher priority to keyframes than to background frames, ensuring that the priority queue retains the keyframe set for the entire horizon T without eviction, i.e., ℬT=keyB_T=X_key. An optimal selector minimizes the learned objective Jθ∗(ℬt)J_θ^*(B_t) at every time step t, which corresponds to recovering the keyframe set key,tX_key,t at each t. (Proof in Appendix A). □ Ideally, the memory buffer should retain the full context needed to disambiguate actions at each timestep t. However, in complex domains such as robot manipulation, our assumptions may not hold strongly, and actions may have some dependency on the background observation set. To address this, we additionally include the current observation oto_t and construct an augmented dataset †D with trajectories τi†=(ℬt,ot,at)t=1Tiτ _i=\(B_t,o_t,a_t)\_t=1^T_i, which is used to train the BC policy πBC(at∣ot,ℬt) _BC(a_t o_t,B_t). See Appendix C and D for details on policy training. 3 Experiments and Result Analysis In this section, we present experiments validating the effectiveness of the Keyframe Mnemonics (KM) method. We first provide empirical evidence supporting our theoretical analysis of keyframe discovery from Section 2.1. We then evaluate mnemonic-conditioned BC policies along two central questions: (1) How effectively does the selector identify keyframes across diverse memory benchmarks, ranging from synthetic domains to robot manipulation tasks? (2) Can mnemonic-conditioned BC policies generalize beyond the training horizon and exhibit horizon-invariant behavior cloning? 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×M=1.0× 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 ± 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. 3.1 Synthetic Domain Benchmarks and Convergence Analysis Domains. We evaluate on three synthetic memory sequence-to-sequence tasks and two grid-based environments. These domains satisfy Assumptions 1 and 2, making them suitable for empirical analysis of our method under the stated assumptions. Add [28] requires the policy to emit the sum of two marked floating point numbers from a sequence. Scattered Copy [22] interleaves a small set of n relevant tokens with noise tokens, and the policy must recover the relevant tokens in same order as observed. In Active T-Maze [29] the agent must remember an oracle observation from the start of a long corridor to make a binary decision at the final junction. MiniGrid-Memory [30] is a similar memory task with a more diverse grid-based observation space. LTMB-Hallway [31] is a visual memory benchmark where the agent must match a key observation seen at the episode start to one of several targets uniformly placed along a long hallway. The policy must compare each observed door with the remembered visual cue and select the matching one. See Appendix C.1 for details on ground truth keyframes in each domain. Proxy policy objective. Figure 2(a) shows the learned objective value Jθ∗(ℬ)J_θ^*(B) as a function of the number of keyframes in the input buffer. Across three synthetic tasks, we observe a decrease in the objective value as more keyframes are added to the buffer, supporting Corollary 1. Memory regret metrics. Let ℒtπL^π_t denote the per-timestep proxy policy loss when the buffer is populated by a selection policy π. We compare the trained selector πs _s against an oracle selector π⋆π that assigns priority 11 to keyframes and 00 otherwise, and a passive baseline πrnd _rnd that samples priorities from ⁡[0,1]U[0,1]. For memory retention analysis, we define active memory regret as ractive=⁡[1T∑t=1T(ℒtπs−ℒtπ⋆)]r_active=E\! [ 1T _t=1^T(L _s_t-L^π _t) ] and passive memory regret as rpassive=⁡[1T∑t=1T(ℒtπrnd−ℒtπ⋆)]r_passive=E\! [ 1T _t=1^T(L _rnd_t-L^π _t) ]. rpassiver_passive characterizes the penalty incurred by not actively selecting what to retain, while ractiver_active measures how close the trained selector is to the oracle. The desired regime for long horizon information retention is rpassive≫ractive≈0r_passive r_active≈ 0. We vary the masked-sampling support size M across three synthetic domains and report the memory regret incurred using the learned objective in Figure 2(b,c). With sufficiently large calibrated support, the active regret drops to zero, while passive regret gradually increases with M, providing empirical evidence for Theorem 1. Figure 4: Results across the 23 MIKASA-Robo tasks. Keyframe Mnemonics improves on the strongest baseline (ELMUR) by an average of +13.9+13.9 absolute success rate, outperforming ELMUR on 1717 of the 2323 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. Baselines and horizon-invariance experiment. We compare KM (mnemonic-conditioned MLP/Transformer policies) against seven competitive sequence models, all trained on BC datasets with favorable parameter budgets. The set comprises LSTM (BPTT/TBPTT) [28], an LSTM trained with backpropagation through time, Decision-LSTM (DLSTM) [40], Decision Transformer (DT) [41], Decision-Mamba (DMamba) a DT-like state-space model [42], Transformer-XL (TrXL) with relative position encoding and a recurrent memory cache [43], ELMUR, a long-horizon mixture-of-experts memory transformer [14], and Random, a stochastic baseline that we include as a noise floor. For each task domain, we evaluate across a range of horizons that includes both in-distribution and out-of-distribution settings, spanning orders of magnitude beyond training distribution. Results in Figure 3 shows that KM maintains near-perfect performance across all horizons on every task, whereas all baseline policies collapse beyond the training range. To probe the limits of generalization we also evaluated KM at two further horizons 10510^5 and 10610^6. KM retained equivalent performance to its in-distribution behavior on all three synthetic tasks; see Appendix E, Table 11. 3.2 Robot Manipulation Benchmarks and Real-world Experiment Memory benchmark. We evaluate the performance of our method on robot-manipulation memory tasks. MIKASA-Robo [7] is a recent benchmark of 32 tabletop manipulation tasks specifically designed to evaluate the object, spatial, sequential, and capacity memory capabilities of policies. In each task the agent observes critical information at moments when the cues are visible, and must then act on that information after the cues are no longer in the observation. For example, in RememberColor3, the agent is briefly shown a colored cube as a cue; after the cue is removed from the observation, the robot must reach to and touch the matching cube among three differently-colored distractors, requiring the policy to retain the cue across the gap between observation and action. MIKASA-Robo thus evaluates the policy’s ability to develop a relevant context from past observations to inform its action at the current state. To further test the memory-retention capabilities of policies, that is, how long the policy can hold the right context for decision-making, we evaluate the policies on RememberColor3-ΔT T with cue-to-choice time delays ΔT T larger than any seen during training. Baselines. For memory benchmarking, we compare our method on all 32 tasks of the MIKASA-Robo against a broad set of baselines including sequence models (Decision Transformer [41]), offline RL (Conservative Q-Learning [44]), Behavior Cloning, Diffusion Policies [5] and the recent state of the art method ELMUR [14]. Here, RATE [45] and ELMUR [14] serve as strong hybrid memory architecture baselines for comparison on tasks that demand memory capabilities from policies. For evaluation of memory retention capabilities, we compare our method with ELMUR and PTP [46]. PTP sub-samples frames from past observations for policy conditioning during training, in comparison to our learned approach of selecting relevant keyframes. Experimental setup. We train Keyframe Mnemonics on the default dataset released by MIKASA-Robo. All results are reported as mean and standard deviation over 33 different seeds with 100100 evaluations per seed. For task designed to evaluate memory retention in manipulation tasks RememberColor3-ΔT T, we train all baselines with the same dataset and a ΔT=15 T=15. During evaluation, we evaluate each baseline from ΔT=15 T=15 to ΔT=150 T=150 at spread of 2020 timesteps, and then at ΔT=200,250 T=200,250. We refer the reader to the Appendix D for more details on the training configuration and the hyperparameters. Results on MIKASA-Robo. Figure 4 summarizes the absolute success-rate improvement of Keyframe Mnemonics over the strongest published baseline, ELMUR, on the 2323 MIKASA-Robo tasks. Keyframe Mnemonics outperforms ELMUR on 21 of the 23 memory-intensive visual robot manipulation tasks and improves average absolute SR by 13.9%; see full table of results in Appendix D. These results show that our self-supervised approach to keyframe identification yields policies that effectively leverage past observational cues from the environment to inform their actions. (a) Memory retention on RememberColor3-ΔT T under cue-to-choice delays beyond the training horizon (marked at 15). ELMUR collapses immediately, PTP degrades gradually, while Keyframe Mnemonics sustains its training-time success rate up to 250 steps. (b) Keyframe Mnemonics matches the diffusion-policy (DP) [5] baseline on three Robomimic tasks (Lift, Can, Square) where memory provides no benefit, confirming it does not regress in the no-memory regime. 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%80\% SR at 20×20× longer test horizons. Mnemonics in grid, robot manipulation and real-world domains. Frames with orange border shows discovered keyframes tracked by the memory buffer ℬtB_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. Horizon-invariance experiment. Figure 5(a) evaluates how long each policy can hold the relevant cue once it is no longer visible. All three methods are trained at ΔT=15 T=15, but their out-of-distribution behavior diverges sharply: ELMUR collapses to near-zero success by ΔT=35 T=35, PTP retains roughly half its training performance through ΔT=75 T=75 before degrading, while Keyframe Mnemonics sustains its training-time success rate of ∼70% \!70\% across the full evaluated range up to ΔT=250 T=250. This highlights the advantages of our learned approach to keyframe identification over sub-sampling of frames from the past, as performance of PTP deteriorates once the cue recedes from its observational context, while Keyframe Mnemonics keeps holding it in the policy’s context. No regression on Robomimic tasks. Figure 5(b) shows that Keyframe Mnemonics matches the performance of diffusion-policy baseline trained without keyframes on three standard Robomimic manipulation tasks (Lift, Can, Square), confirming that the keyframe-conditioning does not deteriorate performance on Markovian benchmarks. Figure 7: Training time. Method / Stage Time KM (Ours) 10h 29m Proxy 1h 15m Selector 5h 46m †D gen. 1h 23m BC policy 2h 05m PTP [46] 19h 34m ELMUR [14] 1h 23m Real-world robot experiment. We evaluate KM on the RememberColor2 task, where the robot observes a red or green cube for 3–5 s. The cue is then removed, and after a delay, two cubes appear 30 cm apart. The robot must reach for the previously observed color. We collect 30 VR-teleoperated demonstrations (15 per color) with ∼5 5 s delays and evaluate at delays up to 20×20× longer. Results shown in Figure 6. Compute and training time. Table 7 reports training time for RememberColor3-ΔT T on a single workstation with an NVIDIA RTX 3090 GPU and 12th-Gen Intel Core i9 CPU, with KM decomposed by training stages. KM trains faster than PTP, which conditions on subsampled frame buffers, while ELMUR is the fastest to train. We also trained ELMUR for ∼ 10 hours to match KM’s training budget, but it remained at 0%0\% SR for ΔT>15 T>15 and failed to generalize beyond the training horizon even with further training. 4 Related Work Explicit memory modeling. Behavior cloning policies with dedicated memory have been developed using embedding caching [14], learned memory compression models [22, 23, 17, 16, 20, 21], temporal state aggregation [18, 46], and hybrid combinations of these approaches [19, 24]. However, these methods generally construct compressed state representations that can blur information over long horizons. In contrast, our method preserves information-rich keyframes that remain unaffected by horizon length. Other approaches use 2D point tracking [47] or 3D workspace mapping [48], but are limited by the performance of the underlying vision foundation model. BC with keyframe annotations. Prior work on keyframe-conditioned behavior cloning has relied on manual keyframe annotations [31, 49, 50, 51] or task-specific prompts for VLM-based selection [52]. However, such approaches require additional supervision derived from oracle reasoning for each task. In contrast, our keyframe identification method is self-supervised, grounded in statistical learning principles, and requires no task-level oracle annotations. Therefore, we do not compare against methods that require additional data beyond expert demonstrations for policy training. Learning to manage policy context. Closer to our work, RL has been used to learn keyframe selection policies for optimizing context buffers [53, 54]. However, these methods have only been studied in simple grid domains and MuJoCo control environments, and their scalability to complex robot manipulation tasks remains unclear. To the best of our knowledge, Keyframe Mnemonics is the first self-supervised keyframe identification method for non-Markovian behavior cloning. Please see Appendix H for an extended related work section. 5 Conclusion We presented a novel method for enabling horizon-invariant behavior cloning in partially observable memory intensive environments. Our approach uses a training strategy that samples random past observations to train a reward model, which is then optimized with reinforcement learning to learn a selector policy that discovers decision-relevant keyframes in a fully self-supervised manner. We provided theoretical analysis and intuition for our algorithm, along with empirical evidence supporting the proposed theory. Furthermore, we evaluated standard behavior cloning policy architectures conditioned on learned mnemonics, achieving state-of-the-art performance across diverse memory-intensive benchmarks, ranging from synthetic tasks to complex robot manipulation environments. We hope our work provides a step toward more reliable long-horizon behavior cloning by enabling policies to retain compact, decision-relevant context over extended horizons. Limitations. The memory budget k is a hyperparameter that is difficult to predetermine, which can lead to multiple information-equivalent buffer states and limit the number of correct keyframes retained in the buffer. In such cases, allowing the inference-time priority queue to retain more observations can provide richer context for the policy. Each demonstration used in this work must contain at least one keyframe. This limits the current algorithm in multi-task or “play demonstration” settings, where keyframes may vary significantly across tasks or may not be present in every trajectory. We aim to extend the method to support these settings in future work. References [1] Martin L Puterman. Markov decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014. [2] Leslie Pack Kaelbling, Michael L Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99–134, 1998. [3] Mikko Lauri, David Hsu, and Joni Pajarinen. Partially observable markov decision processes in robotics: A survey. IEEE Transactions on Robotics, 39(1):21–40, 2022. [4] Tony Z Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn. Learning fine-grained bimanual manipulation with low-cost hardware. arXiv preprint arXiv:2304.13705, 2023. [5] Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023. [6] Kento Kawaharazuka, Jihoon Oh, Jun Yamada, Ingmar Posner, and Yuke Zhu. Vision-language-action models for robotics: A review towards real-world applications. IEEE Access, 13:162467–162504, 2025. doi: 10.1109/ACCESS.2025.3609980. [7] Egor Cherepanov, Nikita Kachaev, Alexey K Kovalev, and Aleksandr I Panov. Memory, benchmark & robots: A benchmark for solving complex tasks with reinforcement learning. arXiv preprint arXiv:2502.10550, 2025a. [8] Paul J Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 78(10):1550–1560, 2002. [9] Maxime Chevalier-Boisvert, Dzmitry Bahdanau, Salem Lahlou, Lucas Willems, Chitwan Saharia, Thien Huu Nguyen, and Yoshua Bengio. Babyai: A platform to study the sample efficiency of grounded language learning. arXiv preprint arXiv:1810.08272, 2018. [10] Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pages 1310–1318. Pmlr, 2013. [11] Corentin Tallec and Yann Ollivier. Unbiasing truncated backpropagation through time. arXiv preprint arXiv:1705.08209, 2017. [12] Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, Jürgen Schmidhuber, et al. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies, 2001. [13] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017. [14] Egor Cherepanov, Alexey K Kovalev, and Aleksandr I Panov. Elmur: External layer memory with update/rewrite for long-horizon rl. arXiv preprint arXiv:2510.07151, 2025b. [15] Pim De Haan, Dinesh Jayaraman, and Sergey Levine. Causal confusion in imitation learning. Advances in neural information processing systems, 32, 2019. [16] Haoquan Fang, Markus Grotz, Wilbert Pumacay, Yi Ru Wang, Dieter Fox, Ranjay Krishna, and Jiafei Duan. Sam2act: Integrating visual foundation model with a memory architecture for robotic manipulation. arXiv preprint arXiv:2501.18564, 2025. [17] Myungkyu Koo, Daewon Choi, Taeyoung Kim, Kyungmin Lee, Changyeon Kim, Younggyo Seo, and Jinwoo Shin. Hamlet: Switch your vision-language-action model into a history-aware policy. arXiv preprint arXiv:2510.00695, 2025. [18] Huiwon Jang, Sihyun Yu, Heeseung Kwon, Hojin Jeon, Younggyo Seo, and Jinwoo Shin. Contextvla: Vision-language-action model with amortized multi-frame context. arXiv preprint arXiv:2510.04246, 2025. [19] Hao Shi, Bin Xie, Yingfei Liu, Lin Sun, Fengrong Liu, Tiancai Wang, Erjin Zhou, Haoqiang Fan, Xiangyu Zhang, and Gao Huang. Memoryvla: Perceptual-cognitive memory in vision-language-action models for robotic manipulation. arXiv preprint arXiv:2508.19236, 2025. [20] Nhat Chung, Taisei Hanyu, Toan Nguyen, Huy Le, Frederick Bumgarner, Duy Minh Ho Nguyen, Khoa Vo, Kashu Yamazaki, Chase Rainwater, Tung Kieu, et al. Rethinking progression of memory state in robotic manipulation: An object-centric perspective. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 40, pages 3407–3415, 2026. [21] Liufan Tan, Jiale Li, and Gangshan Jing. Memoact: Atkinson-shiffrin-inspired memory-augmented visuomotor policy for robotic manipulation. arXiv preprint arXiv:2603.18494, 2026. [22] Artyom Sorokin, Nazar Buzun, Leonid Pugachev, and Mikhail Burtsev. Explain my surprise: Learning efficient long-term memory by predicting uncertain outcomes. Advances in Neural Information Processing Systems, 35:36875–36888, 2022. [23] Yulin Zhou, Yuankai Lin, Fanzhe Peng, Jiahui Chen, Kaiji Huang, Hua Yang, and Zhouping Yin. Mtil: Encoding full history with mamba for temporal imitation learning. IEEE Robotics and Automation Letters, 2025. [24] Yuheng Lei, Zhixuan Liang, Hongyuan Zhang, and Ping Luo. Vpwem: Non-markovian visuomotor policy with working and episodic memory. arXiv preprint arXiv:2603.04910, 2026. [25] Devansh Arpit, Stanisław Jastrzębski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S Kanwal, Tegan Maharaj, Asja Fischer, Aaron Courville, Yoshua Bengio, et al. A closer look at memorization in deep networks. In International conference on machine learning, pages 233–242. PMLR, 2017. [26] Guillermo Valle-Perez, Chico Q Camargo, and Ard A Louis. Deep learning generalizes because the parameter-function map is biased towards simple functions. arXiv preprint arXiv:1805.08522, 2018. [27] Chris Mingard, Henry Rees, Guillermo Valle-Pérez, and Ard A Louis. Deep neural networks have an inbuilt occam’s razor. Nature Communications, 16(1):220, 2025. [28] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9(8):1735–1780, 1997. [29] Bram Bakker. Reinforcement learning with long short-term memory. Advances in neural information processing systems, 14, 2001. [30] Maxime Chevalier-Boisvert, Bolun Dai, Mark Towers, Rodrigo Perez-Vicente, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro, and J K Terry. Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. Advances in Neural Information Processing Systems, 36:73383–73394, 2023. [31] William Yue, Bo Liu, and Peter Stone. Learning memory mechanisms for decision making through demonstrations. arXiv preprint arXiv:2411.07954, 2024. [32] Ming Li, Paul Vitányi, et al. An introduction to Kolmogorov complexity and its applications, volume 3. Springer, 2008. [33] Peter D Grünwald. The minimum description length principle. MIT press, 2007. [34] Blaise Pascal. Traité du triangle arithmetique: avec quelques autres petits traitez sur la mesme matiere. Chez Guillaume Desprez, 1978. [35] Peter D Grünwald, Paul MB Vitányi, et al. Algorithmic information theory. Handbook of the Philosophy of Information, pages 281–320, 2008. [36] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understanding deep learning requires rethinking generalization. arXiv preprint arXiv:1611.03530, 2016. [37] Chulhee Yun, Suvrit Sra, and Ali Jadbabaie. Small relu networks are powerful memorizers: a tight analysis of memorization capacity. Advances in neural information processing systems, 32, 2019. [38] Chris Mingard, Guillermo Valle-Pérez, Joar Skalse, and Ard A Louis. Is sgd a bayesian sampler? well, almost. Journal of Machine Learning Research, 22(79):1–64, 2021. [39] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. [40] Max Siebenborn, Boris Belousov, Junning Huang, and Jan Peters. How crucial is transformer in decision transformer? arXiv preprint arXiv:2211.14655, 2022. [41] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084–15097, 2021. [42] Qi Lv, Xiang Deng, Gongwei Chen, Michael Y Wang, and Liqiang Nie. Decision mamba: A multi-grained state space model with self-evolution regularization for offline rl. Advances in neural information processing systems, 37:22827–22849, 2024. [43] Zihang Dai, Zhilin Yang, Yiming Yang, Jaime G Carbonell, Quoc Le, and Ruslan Salakhutdinov. Transformer-xl: Attentive language models beyond a fixed-length context. In Proceedings of the 57th annual meeting of the association for computational linguistics, pages 2978–2988, 2019. [44] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in neural information processing systems, 33:1179–1191, 2020. [45] Egor Cherepanov, Alexey Staroverov, Dmitry Yudin, Alexey K Kovalev, and Aleksandr I Panov. Recurrent action transformer with memory. arXiv preprint arXiv:2306.09459, 2023. [46] Marcel Torne, Andy Tang, Yuejiang Liu, and Chelsea Finn. Learning long-context diffusion policies via past-token prediction. arXiv preprint arXiv:2505.09561, 2025. [47] Jingjing Chen, Hongjie Fang, Chenxi Wang, Shiquan Wang, and Cewu Lu. History-aware visuomotor policy learning via point tracking. arXiv preprint arXiv:2509.17141, 2025. [48] Remo Steiner, Alexander Millane, David Tingdahl, Clemens Volk, Vikram Ramasamy, Xinjie Yao, Peter Du, Soha Pouya, and Shiwei Sheng. mindmap: Spatial memory in deep feature maps for 3d action policies. arXiv preprint arXiv:2509.20297, 2025. [49] Longxin Kou, Fei Ni, Yan Zheng, Jinyi Liu, Yifu Yuan, Zibin Dong, and Jianye Hao. Kisa: A unified keyframe identifier and skill annotator for long-horizon robotics demonstrations. In Forty-first International Conference on Machine Learning, 2024. [50] BinXu Wu, TengFei Zhang, Chen Yang, JiaHao Wen, HaoCheng Li, JingTian Ma, Zhen Chen, and JingYuan Wang. Sage: State-aware guided end-to-end policy for multi-stage sequential tasks via hidden markov decision process. arXiv preprint arXiv:2509.19853, 2025. [51] Yipeng Chen, Wentao Tan, Lei Zhu, Fengling Li, Jingjing Li, Guoli Yang, and Heng Tao Shen. Non-markovian long-horizon robot manipulation via keyframe chaining. arXiv preprint arXiv:2603.01465, 2026. [52] Max Sobol Mark, Jacky Liang, Maria Attarian, Chuyuan Fu, Debidatta Dwibedi, Dhruv Shah, and Aviral Kumar. Bpp: Long-context robot imitation learning by focusing on key history frames. arXiv preprint arXiv:2602.15010, 2026. [53] Burak Han Demirbilek, Alper Demir, and Faruk Polat. External visual memory with autoencoder-based intrinsic motivation for reinforcement learning under partial observability. In 2025 IEEE 37th International Conference on Tools with Artificial Intelligence (ICTAI), pages 796–800. IEEE, 2025. [54] Geraud Nangue Tasse, Matthew Riemer, Benjamin Rosman, and Tim Klinger. Beyond sliding windows: Learning to manage memory in non-markovian environments. arXiv preprint arXiv:2512.19154, 2025. [55] Vladimir Vapnik. The nature of statistical learning theory. Springer science & business media, 2013. [56] Thomas M Cover. Elements of information theory. John Wiley & Sons, 1999. [57] Tilmann Gneiting and Adrian E Raftery. Strictly proper scoring rules, prediction, and estimation. Journal of the American statistical Association, 102(477):359–378, 2007. [58] Christopher M Bishop and Nasser M Nasrabadi. Pattern recognition and machine learning, volume 4. Springer, 2006. [59] Ian Goodfellow, Yoshua Bengio, Aaron Courville, and Yoshua Bengio. Deep learning, volume 1. MIT press Cambridge, 2016. [60] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1–8, 2021. URL http://jmlr.org/papers/v22/20-1364.html. [61] Giacomo De Palma, Bobak Kiani, and Seth Lloyd. Random deep neural networks are biased towards simple functions. Advances in Neural Information Processing Systems, 32, 2019. [62] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. [63] Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022. [64] Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Luc Bacon. When do transformers shine in rl? decoupling memory from credit assignment. Advances in Neural Information Processing Systems, 36:50429–50452, 2023. [65] Ramprasaath R Selvaraju, Michael Cogswell, Abhishek Das, Ramakrishna Vedantam, Devi Parikh, and Dhruv Batra. Grad-cam: visual explanations from deep networks via gradient-based localization. International journal of computer vision, 128(2):336–359, 2020. [66] Aditya Chattopadhay, Anirban Sarkar, Prantik Howlader, and Vineeth N Balasubramanian. Grad-cam++: Generalized gradient-based visual explanations for deep convolutional networks. In 2018 IEEE winter conference on applications of computer vision (WACV), pages 839–847. IEEE, 2018. [67] Preetum Nakkiran, Gal Kaplun, Dimitris Kalimeris, Tristan Yang, Benjamin L Edelman, Fred Zhang, and Boaz Barak. Sgd on neural networks learns functions of increasing complexity. arXiv preprint arXiv:1905.11604, 2019. [68] Harshay Shah, Kaustav Tamuly, Aditi Raghunathan, Prateek Jain, and Praneeth Netrapalli. The pitfalls of simplicity bias in neural networks. Advances in Neural Information Processing Systems, 33:9573–9585, 2020. [69] Hojoon Lee, Dongyoon Hwang, Donghu Kim, Hyunseung Kim, Jun Jet Tai, Kaushik Subramanian, Peter R Wurman, Jaegul Choo, Peter Stone, and Takuma Seno. Simba: Simplicity bias for scaling up parameters in deep reinforcement learning. arXiv preprint arXiv:2410.09754, 2024. [70] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 conference of the North American chapter of the association for computational linguistics: human language technologies, volume 1 (long and short papers), pages 4171–4186, 2019. [71] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 16000–16009, 2022. [72] Ajay Sridhar, Jennifer Pan, Satvik Sharma, and Chelsea Finn. Memer: Scaling up memory for robot control via experience retrieval. arXiv preprint arXiv:2510.20328, 2025. [73] Marcel Torne, Karl Pertsch, Homer Walke, Kyle Vedder, Suraj Nair, Brian Ichter, Allen Z Ren, Haohuan Wang, Jiaming Tang, Kyle Stachowicz, et al. Mem: Multi-scale embodied memory for vision language action models. arXiv preprint arXiv:2603.03596, 2026. Appendix AProofs⋅· A BAlgorithms⋅· B B.1Buffer priority queue⋅· B.1 B.2Proxy policy training⋅· B.2 B.3Selector RL environment⋅· B.3 B.4End-to-end training and inference⋅· B.4 CSynthetic and Grid Domain Details⋅· C C.1Synthetic environments⋅· C.1 C.2Grid environments⋅· C.2 C.3Training dataset⋅· C.3 C.4Training Keyframe Mnemonics policies⋅· C.4 C.5Baselines and benchmarking details⋅· C.5 C.6Parameter-budget comparison⋅· C.6 C.7Benchmark result analysis⋅· C.7 C.8Hardware and wall-clock⋅· C.8 DRobot Manipulation Domain Details⋅· D D.1MIKASA-Robo environments⋅· D.1 D.2Pipeline overview⋅· D.2 D.3Architectures and parameter counts⋅· D.3 D.4Optimizer and schedule⋅· D.4 D.5Stage-1: Proxy policy fθf_θ⋅· D.5 D.6Stage-2: Selector policy πs _s⋅· D.6 D.7Stage-3: BC policy πBC _BC⋅· D.7 D.8Camera-image data augmentation⋅· D.8 D.9MIKASA-Robo all results⋅· D.9 D.10Hardware and wall-clock⋅· D.10 EExperimental Details⋅· E E.1Horizon-invariance experiments on synthetic and grid domains⋅· E.1 E.2Horizon-invariance experiment on RememberColor3-ΔT T⋅· E.2 E.3Reproducibility⋅· E.3 FAblation Study⋅· F F.1Learned selector against a random selector⋅· F.1 F.2Selector discount factor⋅· F.2 F.3Surprise coefficient⋅· F.3 F.4Current observation in proxy training⋅· F.4 F.5Queue tie-break rule⋅· F.5 F.6Memory buffer size⋅· F.6 GAdditional Analysis⋅· G G.1Interpretability of discovered keyframes⋅· G.1 G.2Selector priority behavior⋅· G.2 G.3BC Policy activation map⋅· G.3 G.4Sensitivity to memory budget k⋅· G.4 G.5PPO reward plots⋅· G.5 G.6Limitation discussion⋅· G.6 HExtended Related Works⋅· H Appendix A Proofs Lemma 1 (Combinatorial cardinality of the sample space). The total cardinality M of the space of distinct timestep-indexed training pairs (ℬt,at)(B_t,a_t) is given by: M=∑i=1N(Ti+∑j=1k(Ti+1j+1)).M= _i=1^N (T_i+ _j=1^k T_i+1j+1 ). Proof. Assume a stochastic training procedure that samples a trajectory τi _i, a timestep t∈[1,Ti]t∈[1,T_i], and a mask Ψk _k selecting a subset of j≤kj≤ k observations from the history hth_t without replacement. The masked-sampling process generates an observation subset representing the buffer ℬB by selecting j observation frames, where j∼min⁡(Binomial⁡(t,0.5),k)j (Binomial(t,0.5),k). Because this spans j∈0,…,min⁡(t,k)j∈\0,…, (t,k)\, the sample space includes all subsets of size up to k, including the empty set for j=0j=0. For any timestep t, the number of unique combinations of size j≥1j≥ 1 is (tj) tj evaluating to 00 if t<jt<j. Summing these over all timesteps t yields the total combinations for size j: ∑t=1Ti(tj)=(Ti+1j+1), _t=1^T_i tj= T_i+1j+1, which simplifies via the combinatorial Hockey-stick identity [34]. For the case where j=0j=0, there is exactly (t0)=1 t0=1 unique combination at each timestep. Summing this across all TiT_i timesteps adds exactly TiT_i combinations: ∑t=1Ti(t0)=Ti. _t=1^T_i t0=T_i. Summing the combinations for all possible mask sizes j∈0,…,kj∈\0,…,k\ gives the sample space cardinality for a single trajectory: Mi=Ti+∑j=1k(Ti+1j+1).M_i=T_i+ _j=1^k T_i+1j+1. Finally, summing across all N independent trajectories in D yields the total cardinality M of the training sample space: M=∑i=1N(Ti+∑j=1k(Ti+1j+1)).M= _i=1^N (T_i+ _j=1^k T_i+1j+1 ). □ Theorem 1 (Expected gap in hypothesis complexity). Under the assumption that training pairs are sampled i.i.d. from a sufficiently large support of size M (Lemma 1), the expected Kolmogorov complexity of the minimum-complexity background-based hypothesis fbg∗f^*_bg is strictly greater than that of the keyframe-based hypothesis fkey∗f^*_key. ⁡[K⁡(fbg∗)]>⁡[K⁡(fkey∗)].E[K(f_bg^*)]>E[K(f_key^*)]. Proof. We first analyze complexity using established equivalence results from algorithmic information theory [35]. Let P⁡(X,Y)P(X,Y) denote a computable joint probability distribution generating data pairs (X,Y)(X,Y) with marginal entropy H⁡(Y)H(Y). According to Theorem 10 in Grünwald et al. [35], the expected Kolmogorov complexity K⁡(⋅)K(·) of a random variable equals its Shannon entropy up to an additive term bounded by the algorithmic complexity of the data-generating distribution itself: ⁡[K⁡(Y)]=H⁡(Y)+O⁡(K⁡(P)).E[K(Y)]=H(Y)+O(K(P)). (6) Similarly, Theorem 11 in Grünwald et al. [35] establishes that expected algorithmic mutual information I⁡(Y:X)I(Y:X) equals the probabilistic Shannon mutual information I⁡(Y,X)I(Y;X) up to an additive term bounded by the same algorithmic complexity: ⁡[I⁡(Y:X)]=I⁡(Y,X)+O⁡(K⁡(P)).E[I(Y:X)]=I(Y;X)+O(K(P)). (7) The algorithmic conditional complexity is defined via algorithmic mutual information (Equation 20 in Grünwald et al. [35]) as K⁡(y∣x∗)=K⁡(y)−I⁡(y:x)K(y x^*)=K(y)-I(y:x). Taking the expectation over the data distribution and substituting Equations 6 and 7, we derive expected conditional Kolmogorov complexity in terms of conditional Shannon entropy 22 2 For notational simplicity, we suppress the distinction between conditioning on X and on its shortest program X∗X^*, absorbing the resulting additive terms into the O⁡(⋅)O(·) overhead.: ⁡[K⁡(Y)]−⁡[I⁡(Y:X)] [K(Y)]-E[I(Y:X)] =H⁡(Y)−I⁡(Y,X)+O⁡(K⁡(P)), =H(Y)-I(Y;X)+O(K(P)), (8) ⁡[K⁡(Y∣X)] [K(Y X)] =H⁡(Y∣X)+O⁡(K⁡(P)). =H(Y X)+O(K(P)). We now apply this to our setting, where P=P⁡()P=P(D) denotes the distribution induced by the POMDP and expert policy that generate dataset D, and we take K⁡()=K⁡(P⁡())K(D)=K(P(D)) to represent the constant algorithmic overhead of this data-generating process. During training, (,A)(X,A) pairs are sampled from a finite support of cardinality M induced by the masked-sampling procedure (Lemma 1) 33 3 Note that we use X to denote the input domain of interpolating hypotheses f∈ℋf , while ℬB denotes a buffer subset used as input to the neural network fθf_θ; in both cases, the input size is bounded by k.. Extending Equation 8 to the entire sequence of M samples, the expected conditional Kolmogorov complexity of the sequence of actions given the sequence of sampled frames is bounded by the joint Shannon entropy plus the one-time algorithmic overhead of the data-generating distribution: [K(A1:M∣1:M)]=H(A1:M∣1:M)+O(K()).E[K(A_1:M _1:M)]=H(A_1:M _1:M)+O(K(D)). (9) Because the M sampled training pairs are treated as i.i.d., the joint conditional Shannon entropy scales linearly, H(A1:M∣1:M)=M⋅H(A∣)H(A_1:M _1:M)=M· H(A ). Substituting this yields: [K(A1:M∣1:M)]=M⋅H(A∣)+O(K()).E[K(A_1:M _1:M)]=M· H(A )+O(K(D)). (10) The conditional Kolmogorov complexity K⁡(y∣x)K(y x) is strictly defined as the length of the shortest program on a universal conditional prefix machine that outputs y when x is provided as input [35]. Therefore, for any interpolating hypothesis fZf_Z satisfying fZ(Z,ℓ)=Aℓf_Z(X_Z, )=A_ for all ℓ , where Z∈key,bgZ∈\key,bg\, there exists a fixed evaluator program which, given a description of fZf_Z and the input sequence Z,1:MX_Z,1:M, perfectly reconstructs the action sequence A1:MA_1:M. Hence, the conditional sequence complexity is upper-bounded by the hypothesis complexity: K(A1:M∣Z,1:M)≤K(fZ∗)+O(1).K(A_1:M _Z,1:M)≤ K(f_Z^*)+O(1). Assumption 3 states that, for each input source Z∈key,bgZ∈\key,bg\, the minimum-complexity interpolating hypothesis fZ∗f_Z^* realizes the corresponding shortest conditional program up to an additive constant: |K(fZ∗)−K(A1:M∣Z,1:M)|≤cZ. |K(f_Z^*)-K(A_1:M _Z,1:M) |≤ c_Z. Equivalently, unpacking the absolute value gives the two-sided bound K(A1:M∣Z,1:M)−cZ≤K(fZ∗)≤K(A1:M∣Z,1:M)+cZ.K(A_1:M _Z,1:M)-c_Z≤ K(f_Z^*)≤ K(A_1:M _Z,1:M)+c_Z. To analyze the complexity gap K⁡(fbg∗)−K⁡(fkey∗)K(f_bg^*)-K(f_key^*), we derive a conservative lower bound by using a lower bound for the background hypothesis and an upper bound for the keyframe hypothesis. This yields K⁡(fbg∗) K(f_bg^*) ≥K(A1:M∣bg,1:M)−cbg, ≥ K(A_1:M _bg,1:M)-c_bg, K⁡(fkey∗) K(f_key^*) ≤K(A1:M∣key,1:M)+ckey. ≤ K(A_1:M _key,1:M)+c_key. Subtracting the keyframe upper bound from the background lower bound, and taking the expectation over the data-generating distribution, yields the expected hypothesis complexity gap: [K(fbg∗)]−[K(fkey∗)]≥[K(A1:M∣bg,1:M)]−[K(A1:M∣key,1:M)]−(ckey+cbg).E[K(f_bg^*)]-E[K(f_key^*)] [K(A_1:M _bg,1:M)]-E[K(A_1:M _key,1:M)]-(c_key+c_bg). Substituting the sequence complexities from Equation 8 and 10, we obtain: ⁡[K⁡(fbg∗)]−⁡[K⁡(fkey∗)]≥M⁡(I⁡(A,key)−I⁡(A,bg))−O⁡(K⁡())−(ckey+cbg).E[K(f_bg^*)]-E[K(f_key^*)]≥ M (I(A;X_key)-I(A;X_bg) )-O(K(D))-(c_key+c_bg). Under Assumption 2, the mutual information gap I⁡(A,key)−I⁡(A,bg)I(A;X_key)-I(A;X_bg) is strictly positive. Because the sample support M is combinatorially massive (Lemma 1), this linear scaling term strictly dominates the combined constant overheads O⁡(K⁡())+(ckey+cbg)O(K(D))+(c_key+c_bg). Thus, for sufficiently large M, we obtain: ⁡[K⁡(fbg∗)]>⁡[K⁡(fkey∗)].E[K(f_bg^*)]>E[K(f_key^*)]. □ Assumption 4 (Simplicity-biased SGD approximation). Conditioned on reaching an interpolating solution in ℋ∗H^*, SGD is approximated as selecting hypotheses with probability biased toward lower Kolmogorov complexity: PSGD(f∣)∝2−K⁡(f),f∈ℋ∗.P_SGD(f ) 2^-K(f), f ^*. This approximation follows the simplicity-bias view of neural networks, where low-complexity functions receive larger prior mass [26, 38, 27]. Proposition 1 (Conditional simplicity-bias preference). Under the simplicity-bias assumption and the complexity gap established in Theorem 1, Stochastic Gradient Descent (SGD) is biased toward the invariant hypothesis class ℋinH_in. Given a masking-induced sample distribution with support size M (Lemma 1), the probability of selecting an invariant hypothesis is approximated by PSGD(ℋin∣)≈11+2−[M⁡(I⁡(A,key)−I⁡(A,bg))−⁡(1)].P_SGD(H_in )≈ 11+2^- [M (I(A;X_key)-I(A;X_bg) )-O(1) ]. Proof. Under the Empirical Risk Minimization (ERM) framework [55], an overparameterized neural network can realize hypotheses in the interpolating set ℋ∗H^* that achieve zero empirical risk on D. To analyze which interpolating hypothesis is selected by SGD, we use the simplicity-bias approximation in Assumption 4, motivated by evidence that neural network function priors and SGD solutions are biased toward low-complexity functions [26, 38, 27]. The probability of selecting an invariant hypothesis is obtained by marginalizing over the invariant class: PSGD(ℋin∣)≈∑f∈ℋin2−K⁡(f)∑f∈ℋ∗2−K⁡(f).P_SGD(H_in )≈ _f _in2^-K(f) _f ^*2^-K(f). Since ℋ∗H^* is partitioned into the invariant class and the non-invariant class, ℋ∗=ℋin∪ℋ¬in,H^*=H_in _ , we can rewrite this as PSGD(ℋin∣)≈∑f∈ℋin2−K⁡(f)∑f∈ℋin2−K⁡(f)+∑f∈ℋ¬in2−K⁡(f).P_SGD(H_in )≈ _f _in2^-K(f) _f _in2^-K(f)+ _f _ 2^-K(f). Because the probability mass decays exponentially with complexity, we use the standard minimum-description approximation that each class-level mass is dominated, up to constant factors independent of M, by its minimum-complexity interpolating hypothesis. Let fkey∗=arg⁡minf∈ℋin⁡K⁡(f),fbg∗=arg⁡minf∈ℋ¬in⁡K⁡(f).f_key^*= _f _inK(f), f_bg^*= _f _ K(f). Then ∑f∈ℋin2−K⁡(f)≈2−K⁡(fkey∗),∑f∈ℋ¬in2−K⁡(f)≈2−K⁡(fbg∗). _f _in2^-K(f)≈ 2^-K(f_key^*), _f _ 2^-K(f)≈ 2^-K(f_bg^*). Substituting this approximation gives PSGD(ℋin∣)≈2−K⁡(fkey∗)2−K⁡(fkey∗)+2−K⁡(fbg∗).P_SGD(H_in )≈ 2^-K(f_key^*)2^-K(f_key^*)+2^-K(f_bg^*). Dividing numerator and denominator by 2−K⁡(fkey∗)2^-K(f_key^*) gives PSGD(ℋin∣)≈11+2−(K⁡(fbg∗)−K⁡(fkey∗)).P_SGD(H_in )≈ 11+2^- (K(f_bg^*)-K(f_key^*) ). By Theorem 1, the expected complexity gap between the non-invariant/background-based hypothesis and the invariant/keyframe-based hypothesis satisfies ⁡[K⁡(fbg∗)]−⁡[K⁡(fkey∗)]≥M⁡(I⁡(A,key)−I⁡(A,bg))−⁡(1).E[K(f_bg^*)]-E[K(f_key^*)]≥ M (I(A;X_key)-I(A;X_bg) )-O(1). Here, the constant overhead O⁡(K⁡())+ckey+cbgO(K(D))+c_key+c_bg is absorbed into the ⁡(1)O(1) term. Therefore, for a typical dataset sampled from the masking-induced distribution, we use the corresponding expected gap approximation K⁡(fbg∗)−K⁡(fkey∗)≈M⁡(I⁡(A,key)−I⁡(A,bg))−⁡(1).K(f_bg^*)-K(f_key^*)≈ M (I(A;X_key)-I(A;X_bg) )-O(1). Substituting this into the class-probability expression yields PSGD(ℋin∣)≈11+2−[M⁡(I⁡(A,key)−I⁡(A,bg))−⁡(1)].P_SGD(H_in )≈ 11+2^- [M (I(A;X_key)-I(A;X_bg) )-O(1) ]. Thus, when the keyframe information advantage is positive and the assumptions hold, the simplicity-bias approximation predicts a rapid increase in the probability of selecting an invariant hypothesis. This follows because the exponent scales linearly with the combinatorially large support size M. 44 4 If the assumptions fail, for example, when a lower-complexity non-invariant shortcut exists, this preference need not hold. □ Corollary 1 (Keyframe recovery from the frozen loss objective). Let us define a subset-selection objective function Jθ∗(ℬ)=⁡[ℒ⁡(fθ∗(ℬ),A)]J_θ^*(B)=E\! [L(f_θ^*(B),A) ] over a frozen predictor fθ∗f_θ^* obtained under Proposition 1 after training with negative log-likelihood loss ℒL, where fθ∗f_θ^* represents a near-optimal invariant hypothesis in ℋinH_in. Consider candidate buffer subsets k=ℬt⊆ht:|ℬt|≤kC_k=\B_t h_t:|B_t|≤ k\, and let ℬopt∈arg⁡minℬ∈kJθ∗(ℬ)B_opt∈ _B _kJ_θ^*(B). If keyframe subsets ℬ⊂keyB _key have a positive action-information margin over all non-keyframe candidate subsets, and this margin dominates the approximation error of the frozen predictor fθ∗f_θ^*, then the loss objective Jθ∗(ℬ)J_θ^*(B) recovers the keyframes up to information-equivalent subsets. In particular, if keyX_key is unique under the memory budget k, then ℬopt=keyB_opt=X_key. Proof. For negative log-likelihood loss, the Bayes-optimal expected loss conditioned on an input buffer subset ℬB is the conditional entropy, since the logarithmic score is a strictly proper scoring rule and the expected negative log-likelihood decomposes into conditional entropy plus a nonnegative KL divergence [56, 57]: inff⁡[ℒ⁡(fθ∗(ℬ),A)]=H⁡(A∣ℬ). _fE[L(f_θ^*(B),A)]=H(A ). Using the Shannon identity H⁡(A∣ℬ)=H⁡(A)−I⁡(A,ℬ)H(A )=H(A)-I(A;B) minimizing the Bayes-optimal negative log-likelihood is equivalent to maximizing the mutual information between the selected subset and the action [56]. Define the action-information margin of the keyframe set over competing candidate subsets as μk=infℬ∈k:ℬ≁key[I(A;key)−I(A;ℬ)], _k= _B _k:\,B _key [I(A;X_key)-I(A;B) ], where ℬ∼keyB _key denotes an information-equivalent subset. The condition μk>0 _k>0 extends Assumption 2 from the background subset to the full candidate search space kC_k, requiring that keyframes have strictly larger action information than any non-equivalent competing subset. Since Proposition 1 biases training toward the invariant class, suppose the frozen model fθ∗f_θ^* is a near-optimal keyframe-conditioned predictor. Define its approximation error on the keyframe input as ϵapp=Jθ∗(key)−H⁡(A∣key). _app=J_θ^*(X_key)-H(A _key). Since negative log-likelihood risk is lower-bounded by conditional entropy, ϵapp≥0 _app≥ 0. Thus, Jθ∗(key)=H⁡(A∣key)+ϵapp.J_θ^*(X_key)=H(A _key)+ _app. For any competing subset ℬ∈kB _k with ℬ≁keyB _key, every predictor using only ℬB has negative log-likelihood risk at least its Bayes risk by nonnegativity of the KL divergence [56, 57]: Jθ∗(ℬ)≥H⁡(A∣ℬ).J_θ^*(B)≥ H(A ). By the definition of μk _k, I⁡(A,key)−I⁡(A,ℬ)≥μk,I(A;X_key)-I(A;B)≥ _k, which is equivalent to H⁡(A∣ℬ)≥H⁡(A∣key)+μk.H(A )≥ H(A _key)+ _k. Therefore, Jθ∗(ℬ)≥H⁡(A∣key)+μk.J_θ^*(B)≥ H(A _key)+ _k. If the information margin dominates the approximation error, μk>ϵapp, _k> _app, then Jθ∗(key)=H⁡(A∣key)+ϵapp<H⁡(A∣key)+μk≤Jθ∗(ℬ)J_θ^*(X_key)=H(A _key)+ _app<H(A _key)+ _k≤ J_θ^*(B) for every non-equivalent competing subset ℬB. Hence, the minimizer of the frozen loss objective belongs to the keyframe equivalence class. If keyX_key is unique under the memory budget k, then ℬopt=key.B_opt=X_key. This keyframe recovery guarantee naturally extends to continuous action spaces optimized via Mean Squared Error (MSE). Under the assumption that the conditional distribution of the actions is Gaussian with a fixed variance, minimizing the MSE is equivalent to minimizing the negative log-likelihood up to affine constants [58, 59]. Consequently, the same argument applies with conditional differential entropy and mutual information replacing their discrete counterparts, so the frozen MSE objective recovers keyX_key under the same positive-margin and approximation-error conditions. □ Theorem 2 (Infinite-horizon keyframe retention). Assume a POMDP in which the expert behavior strictly satisfies Assumption 1 and 2. Let πs∗ _s^* be an optimal selector minimizing the frozen objective Jθ∗(ℬt)J_θ^*(B_t) (Corollary 1), and suppose a unique keyframe set keyX_key exists that has the highest mutual information with the action and fits within the memory budget k. Then, as T→∞T→∞, πs∗ _s^* assigns higher priority to keyframes than to background frames, ensuring that the priority queue retains the keyframe set for the entire horizon T without eviction, i.e., ℬT=keyB_T=X_key. Proof. We analyze the behavior of an optimal selector policy πs∗ _s^* over the temporal horizon T. The reward function R⁡(ℬt)R(B_t) is designed to increase as the frozen loss objective Jθ∗(ℬt)J_θ^*(B_t) decreases. Thus, maximizing expected return requires πs∗ _s^* to assign priority scores that configure the buffer ℬtB_t into the subset minimizing Jθ∗(ℬ)J_θ^*(B). By Assumption 2 and Corollary 1, the keyframe set keyX_key has a strictly positive action-information margin μk>0 _k>0 over all non-keyframe subsets. Consider a timestep t at which the keyframe set has already been inserted into the buffer. Let a new background observation ot∈bgo_t _bg be evaluated. If πs∗ _s^* assigns a priority ptp_t high enough to evict a keyframe, then the resulting buffer ℬtB_t has strictly lower mutual information with the action, decreasing by at least μk _k. By Corollary 1, this induces a strictly higher expected loss: Jθ∗(ℬt)≥H⁡(A∣key)+μk>Jθ∗(key),J_θ^*(B_t)≥ H(A _key)+ _k>J_θ^*(X_key), which penalizes the reward. Therefore, any such eviction is suboptimal, and πs∗ _s^* must assign lower priority to background observations than to keyframes. Once the keyframe set keyX_key is fully inserted into the buffer, this argument applies recursively. Hence, under the optimal selector policy, background observations never receive sufficient priority to evict keyframes. Thus, the priority queue preserves keyX_key indefinitely, yielding ℬT=keyB_T=X_key. □ Appendix B Algorithms Algorithm 2 FIFO priority queue — training-time eviction rule. 1: Parameters: training memory budget k, rejection threshold ηr _r 2: State: buffer ℬt←∅B_t← , insertion counter c←0c← 0 3: procedure Push(observation o, priority p) 4: if p<ηrp< _r then 5: return ⊳ reject priority below threshold 6: end if 7: c←c+1c← c+1 ⊳ monotone insertion timestamp 8: if |ℬt|<k|B_t|<k then 9: ℬt+1←ℬt∪(o,p,c)B_t+1 _t∪\(o,p,c)\ ⊳ free slot, insert directly 10: else 11: (o⋆,p⋆,c⋆)←arg⁡min(o′,p′,c′)∈ℬt⁡(p′,c′)(o ,p ,c )← _(o ,p ,c ) _t(p ,\,c ) ⊳ lowest priority; oldest on a tie 12: if p≥p⋆p≥ p then 13: ℬt+1←(ℬt∖(o⋆,p⋆,c⋆))∪(o,p,c)B_t+1← (B_t \(o ,p ,c )\ )∪\(o,p,c)\ 14: end if 15: end if 16: end procedure Algorithm 3 LIFO priority queue — inference-time eviction rule. 1: Parameters: inference memory budget k, rejection threshold ηr _r, no-repeat threshold ηn _n 2: State: buffer ℬt←∅B_t← , insertion counter c←0c← 0, last priority plast←0p_last← 0 3: procedure Push(observation o, priority p) 4: if p<ηrp< _r or |p−plast|<ηn|p-p_last|< _n then 5: return ⊳ reject low or quasi-repeated priority 6: end if 7: c←c+1c← c+1 8: if |ℬt|<k|B_t|<k then 9: ℬt+1←ℬt∪(o,p,c)B_t+1 _t∪\(o,p,c)\ 10: else 11: (o⋆,p⋆,c⋆)←arg⁡min(o′,p′,c′)∈ℬt⁡(p′,−c′)(o ,p ,c )← _(o ,p ,c ) _t(p ,\,-c ) ⊳ lowest priority; newest on a tie 12: if p≥p⋆p≥ p then 13: ℬt+1←(ℬt∖(o⋆,p⋆,c⋆))∪(o,p,c)B_t+1← (B_t \(o ,p ,c )\ )∪\(o,p,c)\ 14: end if 15: end if 16: plast←p_last← p ⊳ update even if no eviction occurred 17: end procedure B.1 Buffer priority queue The memory buffer ℬtB_t introduced in Section 2.2 is realised by a bounded priority queue of capacity k. The buffer design supports two key behaviors: (i) overflow forgetting, where consistently assigning high priorities causes queue overflow and removes keyframes from the buffer, and (i) background rejection, where assigning priorities below a rejection threshold actively ignores background observations. Both training and inference queue use the same insertion semantics - the selector emits a priority p∈[0,1]p∈[0,1] per observation, and the queue accepts or rejects based on p; the two phases differ only in how ties on the lowest priority are broken at eviction time. Both preserve the temporal ordering of frames in the buffer. Algorithms 2 and 3 state the two rules. We use ηr=0.5 _r=0.5 as the rejection threshold throughout, and ηn=0.05 _n=0.05 as the no-repeat threshold for the inference queue. Why FIFO at training time. The training rule (Algorithm 2) breaks ties in favour of evicting the oldest entry. Under this rule, a frame that has been assigned a uniformly high priority will still age out as new high-priority frames arrive, so the selector cannot rely on saturating priorities to retain information. This sharpens the gradient signal driving πs _s towards genuinely discriminative scores rather than constant-1 outputs. Why LIFO at inference time. The inference rule (Algorithm 3) breaks ties by evicting the newest entry. Thus, initial frames assigned high priority by the selector are preserved, enabling robust keyframe retention and maintaining the keyframe set over arbitrarily long horizons. The no-repeat threshold ηn _n further prevents near-tied priority observations from entering the queue, ensuring that it retains diversity. This especially helps with retaining keyframes when there exists multiple information-equivalent keyframe subsets in the demonstration dataset. Note that under a near-optimal selector and a sufficient memory budget k (Theorem 2), such that the buffer can store all true keyframes, both queues exhibit identical behavior. We observed this empirically in synthetic domains, where the discovered mnemonics remained consistent regardless of the queue strategy used during inference. B.2 Proxy policy training Section 2.1 introduces the proxy policy fθ:ℬ→f_θ:B trained by masked-sampling of past observations (Algorithm 1). In practice, we implement two equivalent realizations of this training procedure, selected based on the scale of the masked-sampling support M (Lemma 1). Online masked-sampling (robot manipulation). For MIKASA-Robo tasks we draw a fresh mask Ψk _k at every gradient step, as stated in Algorithm 1: a trajectory τ∼τ and timestep t∈1,…,Tτt∈\1,…,T_τ\ are sampled, the masked-sampling distribution n∼min⁡(Binomial⁡(t,0.5),k)n (Binomial(t,0.5),k) selects a memory buffer ℬt=oi∈ht∣Ψk(i)=1B_t=\o_i∈ h_t _k(i)=1\ of cardinality |ℬt|≤k _t ≤ k, and the proxy regresses the expert action ata_t under negative log-likelihood loss ℒL. Offline masked-sampling rollouts (synthetic and grid domains). For the synthetic and grid domains (Add, Scattered-Copy, T-Maze, MiniGrid-Memory, LTMB-Hallway), the support size M is small enough to precompute masked samples. We therefore run multiple offline passes over D and cache the random buffer-action pairs on disk, rather than resampling masks at every gradient step. In each pass, for each trajectory τi∈ _i , we iterate over t=1,…,Tit=1,…,T_i, and execute the masked-sampling procedure (Algorithm 1) for each timestep. The resulting pairs (ℬt,at)(B_t,a_t) are written to the cache, and the proxy fθf_θ is trained for multiple epochs on the cached pairs. The number of passes trades off cache size against Monte Carlo coverage of the masked-sampling support and is chosen empirically so that the total number of cached pairs exceeds the sufficient support size M required by Theorem 1. Although online masked-sampling provides the maximum possible cardinality support M, this offline counterpart helped us with more controlled experiments on M to empirically verify Theorem 1 as shown in Figure 2. See Appendix C.3 for more details. Lemma 2 (Equivalence to uniform selector priority queue). Let maskP_mask denote the distribution over buffer states ℬB induced by masked-sampling (Algorithm 1), and let queueP_queue denote the distribution induced by using a random selector πrand _rand that assigns uniform priorities p∈[0,1]p∈[0,1] to observations retained in a priority queue. Both procedures induce the same distribution over buffer states. Proof. Under maskP_mask, the buffer is drawn by picking a size n∼min⁡(Binomial⁡(t,0.5),k)n (Binomial(t,0.5),k) uniformly at random and then selecting n indices uniformly without replacement from 1,…,t\1,…,t\ (Lemma 1). Under queueP_queue, each past observation oi∈hto_i∈ h_t receives an i.i.d. priority pi∼Unif⁡(0,1)p_i (0,1), and the queue retains the top-k priorities exceeding the rejection threshold ηr _r. Because the priorities are exchangeable and independent of i, every k-subset of hth_t is selected with equal probability conditional on its cardinality, and the cardinality distribution is min(Binomial(t,Pr[p>ηr]),k) (Binomial(t, [p> _r]),k). With ηr=0.5 _r=0.5 we have Pr[p>ηr]=0.5 [p> _r]=0.5, recovering the masked-sampling cardinality distribution of Algorithm 1. Hence queue=maskP_queue=P_mask at the level of marginal over ℬtB_t. □ The above result establishes an equivalence between the masked-sampling procedure (Algorithm 1) and selector-driven buffer construction under a random selector policy. Intuitively, this shows that the selector optimization (Section 2.2) operates over the same combinatorial space of buffer states explored during masked-sampling training. Consequently, the learned objective Jθ∗J_θ^* remains aligned with the optimization problem solved by the selector policy. Because Jθ∗J_θ^* wraps a function that learns a keyframe-conditioned hypothesis favored by Theorem 1, optimizing the selector is naturally biased toward recovering keyframe subsets. B.3 Selector RL environment We provide the algorithmic specification of the deterministic MDP ⟨,,,R,γ⟩ ,P,K,R,γ described in Section 2.2, on which the selector policy πs _s is trained with PPO. Algorithm 4 states the per-step transition: the policy observes ot∈o_t , emits a priority pt∈=[0,1]p_t =[0,1], and the environment updates the priority queue ℬtB_t via the transition rule K (Algorithm 2). The post-update buffer is then scored by the frozen proxy fθ∗f_θ to compute the surprise-weighted log-reward R⁡(ℬt)R(B_t) from Equation 5. The episode terminates after TτT_τ steps and the next trajectory is loaded from D. Algorithm 4 Selector RL environment step Require: dataset D, frozen proxy fθ∗f_θ , memory budget k, normalisation constant λ, surprise coefficients βτ,t\ _τ,t\, queue eviction rule Push (Algorithm 2) State: current trajectory τ∈τ , step counter t, buffer ℬt−1B_t-1 1: procedure Step(ot,pto_t,p_t) ⊳ pt∼πs(⋅∣ot)p_t _s(· o_t) 2: at←τ.Action(t)a_t←τ. Action(t) ⊳ expert action from D 3: ℬt←Push(ℬt−1,ot,pt)B_t← Push(B_t-1,o_t,p_t) ⊳ priority-queue update 4: ℓt←ℒ⁡(fθ∗(ℬt),at) _t (f_θ (B_t),a_t ) ⊳ frozen proxy objective Jθ∗(ℬt)J_θ^*(B_t) 5: rt←βτ,t⋅log⁡(ℓt+λ)log⁡(λ)r_t← _τ,t· ( _t+λ) (λ) ⊳ Equation 5 6: t←t+1t← t+1 7: if t=Tτt=T_τ then ⊳ episode terminated 8: τ∼τ ⊳ load next trajectory 9: ℬt−1←∅B_t-1← , t←1t← 1 10: done←Truedone← True 11: else 12: done←Falsedone← False 13: end if 14: ot+1←τ.Obs(t)o_t+1←τ. Obs(t) 15: return (ot+1,rt,done)(o_t+1,r_t,done) 16: end procedure The surprise coefficient βτ,t _τ,t in Equation 5 is precomputed once before PPO training as βτ,t=Tτ⋅softmaxt(Jθ∗τ,t/κ) _τ,t=T_τ·softmax_t(J_θ ^τ,t/κ), where Jθ∗τ,t=Ψk[ℒ⁡(fθ∗(ℬt),at)]J_θ ^τ,t=E_ _k[L(f_θ (B_t),a_t)] is the expected masked-sampling loss evaluated at (τ,t)(τ,t). Surprise scaling is applied on a per-domain basis: we enable it for synthetic and grid domains, where the loss landscape over t is sharply peaked at decision-critical timesteps, and disable it (βτ,t≡1 _τ,t\!≡\!1) for robot manipulation domains, where the unscaled reward in Equation 5 empirically provides a more stable PPO training signal. At training time, the environment is vectorized across nenvsn_envs parallel trajectories drawn from D, and the proxy is evaluated on batches of buffer states for scalable RL training. B.4 End-to-end training and inference Algorithm 5 describes the end-to-end training procedure introduced in Sections 2.1–2.3. The proxy fθf_θ is trained using the masked-sampling algorithm (Algorithm 1), then frozen and used to score the reward R⁡(ℬt)R(B_t) for selector RL (Algorithm 4). The trained selector is then used to produce an augmented dataset †D by rolling πs _s on D with the LIFO eviction rule55 5 For synthetic domains where Theorem 2 holds strongly, we empirically observe that either FIFO or LIFO eviction can be used at inference time without affecting the discovered mnemonics; the LIFO rule is preferred for grid and robot manipulation domains where the assumptions hold less tightly., which is then used to train the BC policy πBC(at∣ot,ℬt) _BC(a_t o_t,B_t). See Appendix C and D for more details. Algorithm 5 Keyframe Mnemonics training Require: dataset D, training memory budget k, proxy fθf_θ, selector πs _s, BC policy πBC _BC, training epochs Ef,Es,EBCE_f,E_s,E_BC Returns: trained (fθ,πs,πBC)(f_θ, _s, _BC) 1: // Stage 1 -- proxy policy training (Section 2.1) 2: for e=1,…,Efe=1,…,E_f do 3: update θ on a batch sampled by Algorithm 1 on D 4: end for 5: freeze θ∗←θ ←θ; define Jθ∗(ℬ)=⁡[ℒ⁡(fθ∗(ℬ),a)]J_θ (B)=E[L(f_θ (B),a)] 6: // Stage 2 -- selector policy training (Section 2.2) 7: precompute surprise coefficients βτ,t\ _τ,t\ from Jθ∗J_θ 8: for e=1,…,Ese=1,…,E_s do 9: roll out the selector MDP (Algorithm 4) with FIFO eviction 10: update πs _s by PPO on the surprise-weighted reward R⁡(ℬt)R(B_t) 11: end for 12: // Stage 3 -- BC policy training (Section 2.3) 13: †←∅D ← 14: for τ∈τ do 15: roll πs _s on τ with LIFO eviction (Algorithm 3) 16: append (ℬt,ot,at)t=1Tτ\(B_t,o_t,a_t)\_t=1^T_τ to †D 17: end for 18: for e=1,…,EBCe=1,…,E_BC do 19: update πBC _BC on a batch from †D under the BC loss 20: end for 21: return (fθ∗,πs,πBC)(f_θ^*, _s, _BC) Algorithm 6 Keyframe Mnemonics inference Require: environment ℰE, trained selector πs _s, trained BC policy πBC _BC, training memory budget k, max rollout steps TmaxT_ , queue eviction rule Push (Algorithm 3) 1: o1←ℰ.Reset()o_1 . Reset() 2: ℬ0←∅B_0← 3: for t=1,…,Tmaxt=1,…,T_ do 4: pt←πs(ot)p_t← _s(o_t) ⊳ assign priority 5: ℬt←Push(ℬt−1,ot,pt)B_t← Push(B_t-1,o_t,p_t) ⊳ LIFO buffer update 6: at←πBC(ℬt,ot)a_t← _BC(B_t,o_t) 7: ot+1←ℰ.Step(at)o_t+1 . Step(a_t) 8: end for At inference time the trained (πs,πBC)( _s, _BC) pair is rolled out in the target environment as specified in Algorithm 6. At each step the selector emits a priority pt=πs(ot)p_t= _s(o_t) that updates the buffer ℬtB_t under the LIFO eviction rule of Algorithm 3 (matching the rule used to construct †D ), and the BC policy emits an action at=πBC(ℬt,ot)a_t= _BC(B_t,o_t) that is applied to the environment for at most TmaxT_ steps. Appendix C Synthetic and Grid Domain Details C.1 Synthetic environments In this section we show expert trajectories for the three synthetic memory benchmarks. Cell colors group conceptually related observations: / / mark the keyframes that an oracle selector must retain, marks background/noise observations that an optimal selector should reject, and marks the decision-critical query timestep at which the policy emits its action. Add [28]. A sequence of T noisy real-valued observations. Two positions t1,t2t_1,t_2 such that t1<t2t_1<t_2 are marked with binary flags (the “first (1,0)(1,0)” and “second (0,1)(0,1)” summands); at the final timestep t=T−1t=T-1 both flags are simultaneously set (1,1)(1,1), signalling the policy to emit the sum of the two marked values. The per-step observation is the triple ot=(vt, 1[t=t1], 1[t=t2])o_t=(v_t,\;1[t=t_1],\;1[t=t_2]) with vt∼⁡[0,1]v_t\! \!U[0,1], and the target action is the scalar at=vt1+vt2a_t=v_t_1+v_t_2 at t=T−1t=T-1 and at=0a_t=0 elsewhere. Example with T=6T=6, t1=1t_1=1, t2=3t_2=3: t 00 11 22 33 44 55 oto_t (0.42,0,0)(0.42,0,0) (0.83,1,0)(0.83,1,0) (0.17,0,0)(0.17,0,0) (0.65,0,1)(0.65,0,1) (0.91,0,0)(0.91,0,0) (0.0,1,1)(0.0,1,1) ata_t 00 00 00 00 00 1.481.48 The three keyframes key=o1,o3,o5X_key=\o_1,o_3,o_5\ comprise the two flag-marked summands ( , ) and the emit step ( ); the latter is itself a keyframe because the policy must use the emit token to know when to emit. The supervised target at the emit step is a5=vt1+vt2=0.83+0.65=1.48a_5=v_t_1+v_t_2=0.83+0.65=1.48; all other timesteps carry unmarked (vt,0,0)(v_t,0,0) observations and are background frames. Evaluation is done on sequences with larger T obtained by increasing the number of unmarked (vt,0,0)(v_t,0,0) background observations in the sequence. Scattered-Copy [22]. A subsequence of nrelevantn_relevant relevant tokens is scattered among nnoisen_noise noise tokens to form a source segment of length Lsrc=nrelevant+nnoiseL_src=n_relevant+n_noise; the source is followed by a recall segment of nrelevantn_relevant delimiter tokens E scattered among nblankn_blank blank tokens □ (length Lrec=nrelevant+nblankL_rec=n_relevant+n_blank). At each delimiter the policy must emit the next relevant token in the order it originally appeared, and a no-op token 0 at every other position. The set of tokens for relevant and noise are disjoint and remains fixed for the task during both training and evaluation. Example with nrelevant=3n_relevant=3, nnoise=3n_noise=3, nblank=5n_blank=5: t 00 11 22 33 44 55 66 77 88 99 1010 1111 1212 1313 oto_t N1 A N2 B N3 C □ □ E □ E □ □ E ata_t 0 0 0 0 0 0 0 0 A 0 B 0 0 C The six keyframes key=o1,o3,o5,o8,o10,o13X_key=\o_1,o_3,o_5,o_8,o_10,o_13\ comprise the three relevant source tokens ( / / ) and the three delimiter tokens ( ); the delimiters are themselves keyframes because the policy must use an E token to know it is at an emit position. The noise tokens Ni in the source and blank tokens in the recall are background frames. Evaluation is done by increasing nnoisen_noise and nblankn_blank in the sequence. Active T-Maze [29]. The agent is placed at the start position x=1x=1 in a corridor of length L with the oracle at x=0x=0 and the T-junction at x=L+1x=L+1. The goal g∈+1,−1 g∈\+1,-1\ is observable only at the oracle position (cue observed at x=0→(0,g)x=0→(0, g)); elsewhere along the corridor the observation is the ambiguous (0,0)(0,0). At the T-junction the observation is (1,0)(1,0) and the agent must choose to turn up or down based on the remembered cue. The action is the 2D displacement at∈(1,0),(−1,0),(0,1),(0,−1)a_t∈\(1,0),(-1,0),(0,1),(0,-1)\ (right/left/up/down). 1+1+ g=±1±1 s ⋅· ⋅· ⋅· J 1−1- where s is the start, J the T-junction, 1+,1−\1+,1-\ the two goal arms, and the three dots between s and J are a schematic ellipsis denoting a variable-length corridor of L ambiguous (0,0)(0,0) cells; the agent must traverse s→g→s→⋯→J→1±s\!→\! g\!→\!s\!→\!·s\!→\! J\!→\!1± correctly to succeed. Example with L=4L=4 and g=+1 g=+1 is t 00 11 22 33 44 55 66 oto_t (0,0)(0,0) (0,+1)(0,+1) (0,0)(0,0) (0,0)(0,0) (0,0)(0,0) (0,0)(0,0) (1,0)(1,0) ata_t (−1,0)(-1,0) (1,0)(1,0) (1,0)(1,0) (1,0)(1,0) (1,0)(1,0) (1,0)(1,0) (0,+1)(0,+1) The two keyframes key=o1,o6X_key=\o_1,o_6\ are the cue exposure at t=1t=1 ( ) and the T-junction observation at t=6t=6 ( ); the junction is itself a keyframe because the policy must use the (1,0)(1,0) token to know when to turn into a goal arm. At the junction the agent must emit (0,+1)(0,+1) if g=+1 g=+1 and (0,−1)(0,-1) if g=−1 g=-1. The ambiguous (0,0)(0,0) corridor observations at all other timesteps ( ) are background frames. Evaluation is done by stretching the gap of (0,0)(0,0) corridor observations between the cue g and the junction J. C.2 Grid environments Both grid benchmarks share the MiniGrid [30] interface. Figure 8: MiniGrid-Memory environment. The agent (red triangle) observes only a local (7×77× 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×77× 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. MiniGrid-Memory [30]. A green Key or Ball is placed in a start room attached to a long corridor; the corridor ends in a T-junction with one Key arm and one Ball arm (order randomised) (Figure 8). The agent spawns at a random cell along the corridor and must reach the arm matching the start-room object; stepping into the wrong arm terminates with task failure. The keyframe is the start-room object observation (cue); all corridor observations are background frames. Evaluation is done stretching the corridor length. LTMB-Hallway [31]. A colored target object (Key, Ball, or Box in one of six colors) is placed in the start room (Figure 9). A horizontal corridor leads past several branching vertical hallways, each containing random objects behind colored doors. Exactly one branch object matches the target in both shape and color, and the agent must open and enter the corresponding hallway door to succeed. Entering any other door results in task failure. The keyframe is the start-room target observation (cue), which must be remembered and matched with hallway targets, while the other corridor observations are background frames. Evaluation is performed by increasing the number of distractor hallways. C.3 Training dataset All training datasets contain expert demonstrations only, with no failure trajectories. This matches the requirements of Keyframe Mnemonics, which needs an expert dataset for training. For each baseline we trained two variants per benchmark, differing only in the training-distribution scale: a small variant matching the original demonstration distribution of D used for training Keyframe Mnemonics, and a large variant in which the trajectory horizon and the number of variations are scaled up. Training sequence lengths and number of demonstrations (|||D|, N in Lemma 1) used in both two variants are summarized in Table 1. Policies trained on both variants are evaluated on the same out-of-distribution test distribution, and the variant that achieves the higher test success rate (or lower test MSE for Add) is considered for benchmarking. Some baselines benefit from the larger distribution and from more diverse training data (e.g. DT and TrXL on Scattered-Copy), while others struggle on the large variant and generalize better when trained on the small dataset (e.g. LSTM-BPTT on T-Maze). See Appendix C.7 for more details. As discussed in Appendix B.2, we run multiple offline passes over the dataset D during masked-sampling training of the proxy policy. The number of passes is denoted by passmspass_ms in Table 1. We determine this value by starting from a small constant and increasing it by orders of 1010 until the requirement on M in Theorem 1 is satisfied. The x-axis multipliers in Figure 2 scale passmspass_ms; for example, M=0.1×M=0.1× for the Add domain means that passms=0.1×300=30pass_ms=0.1× 300=30 passes were used to train the proxy policy. Table 1: Training and test sequence lengths and number of demonstrations |||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 “∼n n” denotes the maximum sequence length. The last column passmspass_ms reports the number of offline rollout passes used to create the masked-sampling cache for proxy training (Appendix B.2). Train (small)∗ Train (large) Test Benchmark Seq len |||D| Seq len |||D| Seq len passmspass_ms Add 55–1010 10001000 2020–3030 50005000 5050–100100 300300 Scattered-Copy 44 10001000 40∗40^* 50005000 220220 3030 T-Maze 55–1010 100100 5555–105105 50005000 ∼1005 1005 300300 MiniGrid-Memory ∼21 21 20002000 — — ∼2012 2012 300300 LTMB-Hallway ∼46 46 50,00050,000 — — ∼425 425 1010 ∗ Dataset used for training Keyframe Mnemonics. C.4 Training Keyframe Mnemonics policies Keyframe Mnemonics (KM) uses three networks per benchmark (Algorithm 5): a proxy policy fθf_θ trained by masked-sampling, a selector πs _s policy trained with PPO, and a BC policy πBC _BC. Per-benchmark architecture and PPO hyperparameters are summarised in Tables 2 and 3. Table 2: KM module architectures per benchmark. k: memory budget; dfeatd_feat: selector feature dimension; dhidd_hid: proxy / BC hidden width; H: action-chunk horizon; L: proxy / BC hidden layers. Cells with two values a/ba/b are proxy / BC policy. Architecture of proxy and BC policy can be different per-domain basis. Benchmark k dfeatd_feat Proxy / BC arch dhidd_hid H L Add 33 3232 ReLU MLP 256256 11 33 Scattered-Copy 2/202/20 3232 ReLU MLP / Transformer 256/64256/64 11 3/43/4 T-Maze 22 1616 ReLU MLP 256256 11 33 MiniGrid-Memory 2/12/1 6464 Transformer 128/64128/64 11 44 LTMB-Hallway 2/12/1 6464 Transformer 128/64128/64 8/18/1 44 Table 3: KM selector optimization settings per benchmark. PPO uses Adam with the SB3 [60] defaults except where shown (nenvsn_envs parallel selector environments, total environment stepsppo, rollout length nstepsn_steps, batch size, nepochsn_epochs per update, learning rate ηπ _π, and reward normalisation constant λ from Equation 5). The value of λ depends on the type of loss function used for reward objective. The proxy is trained with AdamW (learning rate ηf _f, batch size). Selector PPO Proxy Benchmark nenvsn_envs stepsppo nstepsn_steps batch nepochsn_epochs ηπ _π λ ηf _f batch Add 44 ×1052\!×\!10^5 20482048 6464 1010 ×10−43\!×\!10^-4 10−610^-6 10−310^-3 256256 Scattered-Copy 44 ×1052\!×\!10^5 20482048 6464 1010 ×10−43\!×\!10^-4 10−310^-3 10−310^-3 256256 T-Maze 1616 ×1052\!×\!10^5 512512 6464 1010 ×10−43\!×\!10^-4 10−610^-6 10−310^-3 256256 MiniGrid-Memory 3232 ×1053\!×\!10^5 512512 128128 55 ×10−43\!×\!10^-4 10−210^-2 10−310^-3 512512 LTMB-Hallway 3232 ×1052\!×\!10^5 10241024 256256 55 ×10−43\!×\!10^-4 10−210^-2 10−310^-3 512512 In Table 2, note that our formulation allows flexibility in the memory budget k used for proxy and BC policy training, since the inference buffer size can be adjusted after selector training to store a different number of keyframes than what was used during proxy training. The memory budget for the proxy and BC policies is selected empirically for each benchmark, guided by the expected number of task-relevant keyframes. We try a small set of candidate values around this estimate and choose the smallest k that yields stable proxy policy convergence and produces qualitatively valid mnemonics after selector training. See Appendix G for detailed analysis. Also note that in Table 2, the proxy policy uses different network architectures across benchmarks. We find that Theorem 1 holds generally, irrespective of network architecture; for instance, the synthetic domains use MLPs whereas grid-world models use CNNs and Transformers. Our choice of ReLU MLPs in synthetic domains is motivated by prior empirical evidence showing that such models exhibit a stronger simplicity bias [61]. C.5 Baselines and benchmarking details The five benchmarks (Add, Scattered-Copy, T-Maze, MiniGrid-Memory, LTMB-Hallway) share a common trainer pipeline; only the per-baseline hyperparameters and per-task input/output dimensions differ. All baselines are trained with AdamW, a linear warm-up cosine learning-rate schedule (10% of total steps warmup, decaying to 1%1\% of the peak rate), and gradient clipping. Our baseline design and implementation choices are derived from the open-source code of RATE [45]. Additionally, we tuned the training recipe to obtain the best possible baseline results on our benchmarking domains. 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 = 10310^3, M = 10610^6. Benchmark LSTM DLSTM DT DMamba TrXL ELMUR KM (Ours) Add 100K100K 740K740K 182K182K 394K394K 182K182K 4.4M4.4M K268 K Scattered-Copy 102K102K 742K742K 184K184K 396K396K 184K184K 4.4M4.4M K366 K T-Maze 100K100K 740K740K 246K246K 395K395K 246K246K 4.4M4.4M K266 K MiniGrid-Memory 641K641K 1.9M1.9M 2.0M2.0M 1.5M1.5M 997K997K 4.7M4.7M 1.44M1.44 M LTMB-Hallway 641K641K 1.9M1.9M 2.0M2.0M 1.5M1.5M 997K997K 4.7M4.7M 1.45M1.45 M Table 5: Decomposition of the Keyframe Mnemonics parameter budget into the proxy fθf_θ, the selector policy πs _s, and the BC policy πBC _BC. Benchmark Proxy fθf_θ Selector πs _s BC policy πBC _BC Total Add 134K134K 128128 134K134K K268 K Scattered-Copy 157K157K 2.7K2.7K 206K206K K366 K T-Maze 133K133K 4848 133K133K K266 K MiniGrid-Memory 957K957K 242K242K 245K245K 1.44M1.44 M LTMB-Hallway 959K959K 242K242K 245K245K 1.45M1.45 M C.6 Parameter-budget comparison Table 4 compares the total trainable parameter count of every method per benchmark. For Keyframe Mnemonics we report the sum of the proxy, selector, and BC policy parameters. Despite using three networks, the KM total remains well below most of the baselines demonstrating that the architectural decomposition does not come at a parameter cost. Table 5 breaks the KM total down by module. The selector policy is compact (under 1%1\% of the total on the synthetic tasks and roughly 15%15\% on the grid tasks), since it only encodes a single observation at a time; most of the budget sits in the proxy and, the BC policy. C.7 Benchmark result analysis We report per-benchmark in-distribution and out-of-distribution performance for every method in Table 6. For each baseline we report the better of policy performance trained on small / large dataset variants (Appendix C.3); Keyframe Mnemonics is trained only on the small variant (except Scattered-Copy that uses large66 6 For Scattered-Copy, the proxy policy (Stage 1) and selector policy (Stage 2) are trained on the small variant, while the BC policy (Stage 3) is trained on the large variant.). The selected variant per (benchmark, baseline) pair is as follows. On Add, the small variant is used for LSTM(BPTT), DLSTM, and the large variant for LSTM(TBPTT), DT, DMamba, TrXL, and ELMUR. On Scattered-Copy, the small variant is used for LSTM(BPTT) and LSTM(TBPTT), and the large variant for all remaining baselines. On T-Maze, the large variant is used for LSTM(TBPTT), DT, and TrXL, and the small variant for all remaining baselines. The grid benchmarks (MiniGrid-Memory and LTMB-Hallway) use a single dataset. Table 6: Aggregated benchmark results across all five tasks. Cells show mean ± std across 55 seeds for Add, Scattered-Copy, and T-Maze, and 33 seeds for MiniGrid-Memory and LTMB-Hallway. For Add, the metric is final-step MSE (↓ better); for all others, it is success rate (↑ 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 |||D| used for training. LSTM(B) LSTM(T) DLSTM DT DMamba TrXL ELMUR KM (Ours) Add Train MSE ↓ ⋅ 10−±⋅ 10−4 \,·\,10^-6\!±\!1 \,·\,10^-6 ±0.0030.006\!±\!0.003 ±0.1060.167\!±\!0.106 8.1⋅ 10−±6.8⋅ 10−8.1 \,·\,10^-4\!±\!6.8 \,·\,10^-4 2.0⋅ 10−±1.3⋅ 10−2.0 \,·\,10^-4\!±\!1.3 \,·\,10^-4 2.2⋅ 10−±1.2⋅ 10−2.2 \,·\,10^-4\!±\!1.2 \,·\,10^-4 ±0.0140.114\!±\!0.014 3.4⋅ 10−±3.6⋅ 10−3.4 \,·\,10^-5\!±\!3.6 \,·\,10^-5 Test MSE ↓ ±0.0040.008\!±\!0.004 ±0.0240.040\!±\!0.024 ±0.1860.368\!±\!0.186 ±0.0830.148\!±\!0.083 ±0.0860.160\!±\!0.086 ±0.0640.133\!±\!0.064 ±0.0470.157\!±\!0.047 3.2⋅ 10−±3.3⋅ 10−3.2 \,·\,10^-5\!±\!3.3 \,·\,10^-5 Scattered-Copy Train Copy SR ↑ ±0.0001.000\!±\!0.000 ±0.0350.100\!±\!0.035 ±0.0000.000\!±\!0.000 ±0.3030.538\!±\!0.303 ±0.0000.000\!±\!0.000 ±0.0160.018\!±\!0.016 ±0.0000.000\!±\!0.000 ±0.0001.000\!±\!0.000 Test Copy SR ↑ ±0.0210.160\!±\!0.021 ±0.0300.106\!±\!0.030 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0001.000\!±\!0.000 T-Maze Train SR ↑ ±0.0001.000\!±\!0.000 ±0.0320.330\!±\!0.032 ±0.0550.530\!±\!0.055 ±0.0001.000\!±\!0.000 ±0.0320.550\!±\!0.032 ±0.0001.000\!±\!0.000 ±0.1560.682\!±\!0.156 ±0.0001.000\!±\!0.000 Test SR ↑ ±0.2360.828\!±\!0.236 ±0.0780.314\!±\!0.078 ±0.0580.478\!±\!0.058 ±0.0430.518\!±\!0.043 ±0.0290.450\!±\!0.029 ±0.0470.500\!±\!0.047 ±0.0350.018\!±\!0.035 ±0.0001.000\!±\!0.000 MiniGrid-Memory Train SR ↑ — ±0.2140.823\!±\!0.214 ±0.0670.747\!±\!0.067 ±0.0500.553\!±\!0.050 ±0.0060.203\!±\!0.006 ±0.0350.483\!±\!0.035 ±0.0250.653\!±\!0.025 ±0.0120.993\!±\!0.012 Test SR ↑ — ±0.0720.497\!±\!0.072 ±0.0690.530\!±\!0.069 ±0.0250.547\!±\!0.025 ±0.0000.030\!±\!0.000 ±0.0310.487\!±\!0.031 ±0.1000.300\!±\!0.100 ±0.0060.997\!±\!0.006 LTMB-Hallway Train SR ↑ — ±0.0640.037\!±\!0.064 ±0.0310.077\!±\!0.031 ±0.1120.217\!±\!0.112 ±0.0000.000\!±\!0.000 ±0.0150.123\!±\!0.015 ±0.0000.070\!±\!0.000 ±0.0001.000\!±\!0.000 Test SR ↑ — ±0.0060.003\!±\!0.006 ±0.0060.003\!±\!0.006 ±0.0060.013\!±\!0.006 ±0.0000.000\!±\!0.000 ±0.0000.010\!±\!0.000 ±0.0000.010\!±\!0.000 ±0.0060.993\!±\!0.006 Random baseline omitted; its test values are ±0.0610.499\!±\!0.061 test MSE on Add, ±0.0000.000\!±\!0.000 test copy SR on Scattered-Copy, ±0.0000.000\!±\!0.000 test SR on T-Maze, ±0.0100.030\!±\!0.010 test SR on MiniGrid-Memory, and ±0.0000.000\!±\!0.000 test SR on LTMB-Hallway. Memory blur issue. In T-Maze, sequence models such as LSTM (BPTT) and TrXL fit the training distribution well and achieve a train SR of 1.01.0, but collapse to near-chance performance on longer test horizons. We observe a similar trend in the Add domain. Recurrent memory baselines struggle in these domains because their hidden states or external memories are repeatedly updated over long horizons. Keyframe Mnemonics avoids these failure modes by retaining only keyframes which are relevant for the task which remain unaffected to change in horizon length. A more severe case is Scattered-Copy, which requires remembering multiple numbers and emitting them in the order they were observed. This task is particularly difficult for baselines that repeatedly update hidden states or external memories, as such updates blur information over time and fail to generalize to longer OOD sequences. In contrast, the KM selector learns to identify relevant keyframes during Stage 2 of training and stores only these keyframes in the buffer, making it easier for the BC policy to learn the copy operation over buffer elements while ignoring noise tokens in the sequence. This highlights a key advantage of our method; since the buffer stores individual tokens, it avoids information blur. Non-memory baseline. Among the baselines, Decision Transformer (DT) and DMamba do not include an explicit memory mechanism. We include them for completeness in our benchmark evaluation. The observed results are expected, as these models are limited by the context window length used during training; once the sequence length exceeds this window, relevant information falls out of context and the models can no longer predict the correct action. Consequently, they perform well on the training distribution but degrade to near-chance performance on the test distribution. C.8 Hardware and wall-clock All training runs were executed on a single NVIDIA RTX 3090 (2424 GB) GPU. The three-stage Keyframe Mnemonics pipeline (proxy training, selector PPO, BC policy training) completes across all five benchmarks in ∼3 3 h per seed, with the grid benchmarks accounting for the majority of the wall-clock time. The cumulative training time for all baselines on a single synthetic benchmark is under ∼30 30 minutes per seed. The cumulative baseline training time is ∼4 4 h per seed on MiniGrid-Memory and ∼6 6 h per seed on LTMB-Hallway. Appendix D Robot Manipulation Domain Details This section provides details on the hyperparameters and architectural choices used to produce the MIKASA-Robo results in Section 3. For the Robomimic domain experiments (Figure 5(b)), we use the same configuration as MIKASA-Robo. Values not stated here take the standard defaults of the underlying libraries. D.1 MIKASA-Robo environments A set of six representative tasks from the benchmark is shown in Figure 10. 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]. D.2 Pipeline overview Each MIKASA-Robo task is trained in three stages (Algorithm 5), following Sections 2.1–2.3 of the main paper. (i) The proxy policy fθf_θ is trained under masked-sampling of past observations, after which its frozen prediction loss serves as the keyframe-discovery objective Jθ∗(ℬ)J_θ^*(B). (i) The selector policy πs _s is trained with PPO against the reward derived from Jθ∗J_θ^*, learning to assign high priority to keyframes. (i) The BC policy πBC(at∣ot,ℬt) _BC(a_t o_t,B_t) is trained on the augmented dataset †D , in which every observation oto_t is paired with the priority queue ℬtB_t produced by πs _s. All three stages use only the offline expert demonstrations D; the simulator is never queried during proxy or selector training. The memory budget is fixed at k=3/2k=3/2 (proxy / BC policy) throughout, so that πBC _BC sees the current observation alongside two mnemonics drawn from the queue. D.3 Architectures and parameter counts All modules share a ResNet-18 visual encoder, initialised from ImageNet weights and then trained end-to-end. Camera frames are 128×128128×128 RGB; the proprioceptive state concatenates the TCP pose, joint positions, and joint velocities into a 2525-dimensional vector. Both overhead and wrist cameras are used. Table 7 summarises the architecture and parameter count of each module. 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. Module Architecture d h L H Params πs _s (selector) ResNet-18 + MLP [500, 500][500,\,500] 256 — — — 15.80 M fθf_θ (proxy) ResNet-18 + Transformer 512 4 4 8 24.13 M πBCSST _BC^SST Regression transformer 1024 16 4 1 62.13 M πBCFM _BC^FM Flow-matching cross-attn transformer 1024 16 4 8 120.20 M The single-step regression head (πBCSST _BC^SST) is the default BC-policy. We substitute a flow-matching transformer (πBCFM _BC^FM) only on the four Intercept* dynamics tasks (InterceptMedium and InterceptGrab Slow,Medium,Fast), where the higher-frequency control profile benefits from action chunking of length H=8H=8. The flow-matching backbone is a cross-attention transformer in which action tokens self-attend and then cross-attend to visual and proprioceptive context tokens. D.4 Optimizer and schedule The proxy and BC policies are trained with AdamW [62]; the selector is trained with PPO [39] using Adam (Table 8). Actions are predicted directly in the simulator’s [−1, 1][-1,\,1] range without output normalization. Table 8: Optimizer settings for the three training stages. Stage Optimizer LR Weight decay Schedule Selector policy πs _s (PPO) Adam 5× 10−55\,×\,10^-5 00 constant Proxy policy fθf_θ AdamW 1× 10−41\,×\,10^-4 00 constant BC policy πBCSST _BC^SST AdamW 5× 10−55\,×\,10^-5 0.050.05 constant BC policy πBCFM _BC^FM AdamW 5× 10−55\,×\,10^-5 0.050.05 constant D.5 Stage-1: Proxy policy fθf_θ The proxy policy is trained for 6060 epochs over the demonstration dataset D with batch size 128128. We use online masked-sampling (Appendix B.2) and at every gradient step we sample a random binary mask Ψk _k over the past observations hth_t with cardinality bounded by k, forming the masked memory buffer ℬt=oi∈ht∣Ψk(i)=1B_t=\o_i∈ h_t _k(i)=1\ (Section 2.1). Conditioned on ℬtB_t, the proxy is trained to reconstruct the expert action chunk at horizon H=8H=8 under L1 loss. D.6 Stage-2: Selector policy πs _s The selector policy is a PPO actor-critic with the hyperparameters of Table 9. The reward at every step is the normalised log objective defined in Section 2.2, evaluated on the post-update memory buffer ℬtB_t. We do not use surprise coefficient for robot manipulation domains. We clip the proxy loss inside the log-reward at λ=10−5λ=10^-5 to prevent unbounded rewards near the proxy’s convergence. The selector observes the same RGB and proprioceptive state input as the proxy policy, and emits a continuous priority score pt∈[0,1]p_t∈[0,1] per observation; the priority queue uses this score to update ℬtB_t from ℬt−1B_t-1 (Algorithm 2). We train the selector for 5× 1065\,×\,10^6 environment steps across nenvs=32n_envs=32 parallel selector environments. Table 9: PPO hyperparameters used to train the selector policy πs _s. Hyperparameter Value Hyperparameter Value Hyperparameter Value Total environment steps 5× 1065\,×\,10^6 PPO epochs per update 55 Discount γ 0.990.99 Parallel envs (nenvsn_envs) 3232 Entropy coefficient 0.010.01 GAE λGAE _GAE 0.950.95 Rollout length (nstepsn_steps) 10241024 Clip range 0.20.2 Value-function coef. cvfc_vf 0.50.5 Minibatch size 256256 Reward floor λ 10−510^-5 Max grad norm 0.50.5 Table 10: Success rate (mean± 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. Environment RATE DT BC CQL DP ELMUR KM (Ours) Shell Game Touch 0.92 ± 0.01 0.53 ± 0.07 0.28 ± 0.01 0.16 ± 0.04 0.18 ± 0.02 0.96 ± 0.02 0.94 ± 0.02 Shell Game Push 0.78 ± 0.06 0.62 ± 0.14 0.27 ± 0.01 0.25 ± 0.01 0.22 ± 0.03 0.87 ± 0.04 0.99 ± 0.01 Shell Game Pick 0.02 ± 0.01 0.00 ± 0.00 0.01 ± 0.01 0.00 ± 0.00 0.01 ± 0.00 0.66 ± 0.01 0.78 ± 0.08 Intercept Slow 0.23 ± 0.02 0.40 ± 0.02 0.37 ± 0.06 0.25 ± 0.01 0.33 ± 0.05 0.61 ± 0.04 0.48 ± 0.06 Intercept Medium 0.32 ± 0.02 0.56 ± 0.01 0.31 ± 0.14 0.03 ± 0.01 0.68 ± 0.02 0.64 ± 0.03 0.78 ± 0.01 Intercept Fast 0.30 ± 0.04 0.36 ± 0.04 0.03 ± 0.02 0.02 ± 0.02 0.21 ± 0.05 0.54 ± 0.02 0.59 ± 0.04 Intercept Grab Slow 0.09 ± 0.03 0.00 ± 0.00 0.28 ± 0.18 0.03 ± 0.00 0.03 ± 0.01 0.08 ± 0.03 0.30 ± 0.04 Intercept Grab Medium 0.09 ± 0.03 0.00 ± 0.00 0.11 ± 0.02 0.08 ± 0.04 0.03 ± 0.01 0.10 ± 0.04 0.12 ± 0.08 Intercept Grab Fast 0.14 ± 0.03 0.11 ± 0.03 0.09 ± 0.02 0.08 ± 0.03 0.18 ± 0.02 0.20 ± 0.01 0.37 ± 0.02 Rotate Lenient Pos 0.11 ± 0.04 0.01 ± 0.01 0.15 ± 0.03 0.16 ± 0.02 0.11 ± 0.02 0.24 ± 0.02 0.52 ± 0.05 Rotate Lenient Pos Neg 0.29 ± 0.03 0.05 ± 0.02 0.22 ± 0.01 0.12 ± 0.02 0.14 ± 0.05 0.26 ± 0.00 0.71 ± 0.06 Rotate Strict Pos 0.03 ± 0.02 0.05 ± 0.04 0.01 ± 0.00 0.03 ± 0.01 0.06 ± 0.02 0.16 ± 0.00 0.52 ± 0.02 Rotate Strict Pos Neg 0.08 ± 0.01 0.05 ± 0.03 0.04 ± 0.02 0.04 ± 0.02 0.15 ± 0.01 0.12 ± 0.02 0.42 ± 0.07 Take It Back 0.42 ± 0.24 0.08 ± 0.04 0.33 ± 0.10 0.04 ± 0.01 0.05 ± 0.02 0.78 ± 0.03 0.90 ± 0.04 Remember Color3 0.65 ± 0.04 0.01 ± 0.01 0.27 ± 0.03 0.29 ± 0.01 0.32 ± 0.01 0.89 ± 0.07 0.99 ± 0.01 Remember Color5 0.13 ± 0.03 0.07 ± 0.05 0.12 ± 0.01 0.15 ± 0.02 0.10 ± 0.02 0.19 ± 0.03 0.46 ± 0.07 Remember Color9 0.09 ± 0.02 0.01 ± 0.01 0.12 ± 0.02 0.15 ± 0.01 0.17 ± 0.01 0.23 ± 0.02 0.33 ± 0.03 Remember Shape3 0.21 ± 0.04 0.05 ± 0.04 0.31 ± 0.04 0.20 ± 0.10 0.32 ± 0.05 0.76 ± 0.01 0.63 ± 0.15 Remember Shape5 0.17 ± 0.04 0.04 ± 0.04 0.18 ± 0.01 0.15 ± 0.00 0.21 ± 0.04 0.27 ± 0.02 0.49 ± 0.08 Remember Shape9 0.05 ± 0.00 0.05 ± 0.02 0.10 ± 0.02 0.14 ± 0.01 0.11 ± 0.02 0.17 ± 0.03 0.17 ± 0.02 Remember Shape Color3x2 0.14 ± 0.02 0.04 ± 0.02 0.13 ± 0.02 0.11 ± 0.05 0.14 ± 0.02 0.19 ± 0.02 0.52 ± 0.07 Remember Shape Color3x3 0.08 ± 0.03 0.06 ± 0.06 0.09 ± 0.02 0.09 ± 0.02 0.16 ± 0.01 0.19 ± 0.02 0.27 ± 0.03 Remember Shape Color5x3 0.07 ± 0.02 0.01 ± 0.01 0.09 ± 0.01 0.09 ± 0.02 0.11 ± 0.03 0.13 ± 0.01 0.14 ± 0.06 Bunch Of Colors3 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 Bunch Of Colors5 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 Bunch Of Colors7 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 Seq Of Colors3 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 Seq Of Colors5 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 Seq Of Colors7 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 Chain Of Colors3 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 Chain Of Colors5 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 Chain Of Colors7 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 0.00 ± 0.00 D.7 Stage-3: BC policy πBC _BC Optimization budget. The regression head (πBCSST _BC^SST) is trained for 15,00015,000 gradient steps; the flow-matching variant (πBCFM _BC^FM) is trained for 100,000100,000 gradient steps. Batch size is 128128 throughout, and the final checkpoint is used for evaluation. Flow-matching scheduler. We adopt the Euler discrete flow-matching scheduler [63] with 200200 training timesteps sampled uniformly at random. At inference we run 55 Euler steps; we found this sufficient on every Intercept* task. Closed-loop control. The BC policy emits an action chunk of length H at every environment step, but only the first action of the chunk is executed before re-querying. D.8 Camera-image data augmentation We apply augmentation to all camera images in the buffer during BC policy training. Each of the following operations is applied independently with probability p=0.5p=0.5: brightness jitter [0.8, 1.2][0.8,\,1.2], contrast jitter [0.8, 1.2][0.8,\,1.2], saturation jitter [0.5, 1.5][0.5,\,1.5], hue jitter [−0.05, 0.05][-0.05,\,0.05], sharpness adjustment (factor 1.51.5), and a random affine transform with up to 10∘10 rotation, ±10%± 10\% translation, and ±10%± 10\% scaling. The proprioceptive state vector is not augmented. We find this augmentation crucial for training the BC policy to avoid causal confusion [15] induced by the larger context size. D.9 MIKASA-Robo all results Averaged over the 23 non-zero tasks, Keyframe Mnemonics outperforms Decision Transformer by +40.4%+40.4\%, Conservative Q-Learning by +42.5%+42.5\%, Behavior Cloning by +37.1%+37.1\%, RATE by +30.6%+30.6\%, the diffusion-policy baseline by +36.6%+36.6\%, and ELMUR by +13.9%+13.9\% absolute success rate. See Table 10 for complete results. D.10 Hardware and wall-clock Training runs were conducted on a mix of NVIDIA A4040 (4848 GB) and L4040S (4646 GB) GPU instances. Per-stage wall-clock, averaged across the 2323 non-zero MIKASA-Robo tasks and rounded, is approximately 1.51.5 h for the proxy fθf_θ (6060 epochs), ∼10 \!10 h for the selector πs _s (5× 1065\,×\,10^6 environment steps), and ∼1.5 \!1.5 h for the BC policy πBCSST _BC^SST (15,00015,000 gradient steps). End-to-end training therefore takes roughly half a day per SST task (∼13 \!13 h). The flow-matching variant πBCFM _BC^FM trains for 100,000100,000 gradient steps and takes proportionally longer at the BC stage (∼10 \!10 h), bringing the four Intercept* flow-matching tasks to roughly a full day each. Appendix E Experimental Details Table 11: Horizon-invariance results across synthetic and grid benchmarks. Cells report mean±~±~std of the relevant metric: final-step MSE for Add (↓ better), Copy SR for Scattered-Copy and Episode SR for T-Maze, MiniGrid-Memory and LTMB-Hallway (↑ 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). LSTM(B) LSTM(T) DLSTM DT DMamba TrXL ELMUR KM (Ours) Add (final-step MSE ↓ ) h=10h=10 5.0⋅ 10−±2.3⋅ 10−5.0 \,·\,10^-6\!±\!2.3 \,·\,10^-6 3.4⋅ 10−3±3.5⋅ 10−33.4 \,·\,10^-3\!±\!3.5 \,·\,10^-3 ±0.2160.258\!±\!0.216 6.5⋅ 10−±6.2⋅ 10−6.5 \,·\,10^-4\!±\!6.2 \,·\,10^-4 2.5⋅ 10−±1.8⋅ 10−2.5 \,·\,10^-4\!±\!1.8 \,·\,10^-4 3.9⋅ 10−±2.4⋅ 10−3.9 \,·\,10^-4\!±\!2.4 \,·\,10^-4 ±0.1140.143\!±\!0.114 3.1⋅ 10−±3.4⋅ 10−3.1 \,·\,10^-5\!±\!3.4 \,·\,10^-5 h=20h=20 7.7⋅ 10−±6.9⋅ 10−7.7 \,·\,10^-5\!±\!6.9 \,·\,10^-5 6.8⋅ 10−3±5.6⋅ 10−36.8 \,·\,10^-3\!±\!5.6 \,·\,10^-3 ±0.3660.330\!±\!0.366 9.2⋅ 10−±7.7⋅ 10−9.2 \,·\,10^-4\!±\!7.7 \,·\,10^-4 2.1⋅ 10−±1.9⋅ 10−2.1 \,·\,10^-4\!±\!1.9 \,·\,10^-4 1.8⋅ 10−±1.4⋅ 10−1.8 \,·\,10^-4\!±\!1.4 \,·\,10^-4 ±0.0900.131\!±\!0.090 2.8⋅ 10−±5.7⋅ 10−2.8 \,·\,10^-4\!±\!5.7 \,·\,10^-4 h=30h=30 3.1⋅ 10−±1.9⋅ 10−3.1 \,·\,10^-4\!±\!1.9 \,·\,10^-4 1.1⋅ 10−2±8.3⋅ 10−31.1 \,·\,10^-2\!±\!8.3 \,·\,10^-3 ±0.3480.267\!±\!0.348 7.0⋅ 10−±7.8⋅ 10−7.0 \,·\,10^-4\!±\!7.8 \,·\,10^-4 2.2⋅ 10−±1.5⋅ 10−2.2 \,·\,10^-4\!±\!1.5 \,·\,10^-4 1.4⋅ 10−±6.1⋅ 10−1.4 \,·\,10^-4\!±\!6.1 \,·\,10^-5 ±0.0520.141\!±\!0.052 7.7⋅ 10−±1.1⋅ 10−7.7 \,·\,10^-5\!±\!1.1 \,·\,10^-4 h=103h=10^3 ±0.0970.246\!±\!0.097 ±0.0270.127\!±\!0.027 ±0.4550.531\!±\!0.455 ±0.4390.408\!±\!0.439 ±0.2470.387\!±\!0.247 ±0.2160.296\!±\!0.216 ±0.9380.797\!±\!0.938 3.8⋅ 10−±4.8⋅ 10−3.8 \,·\,10^-5\!±\!4.8 \,·\,10^-5 h=104h=10^4 ±0.3200.494\!±\!0.320 ±0.1420.351\!±\!0.142 ±0.1950.218\!±\!0.195 ±0.3980.447\!±\!0.398 ±0.3590.594\!±\!0.359 ±0.3870.634\!±\!0.387 ±0.4650.759\!±\!0.465 3.7⋅ 10−±4.0⋅ 10−3.7 \,·\,10^-5\!±\!4.0 \,·\,10^-5 h=105h=10^5 — — — — — — — 3.8⋅ 10−±3.8⋅ 10−3.8 \,·\,10^-5\!±\!3.8 \,·\,10^-5 h=106h=10^6 — — — — — — — 1.6⋅ 10−±3.5⋅ 10−1.6 \,·\,10^-5\!±\!3.5 \,·\,10^-5 Scattered-Copy (Copy SR ↑ ) h=10h=10 ±0.1100.720\!±\!0.110 ±0.1100.080\!±\!0.110 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0001.000\!±\!0.000 h=20h=20 ±0.1670.360\!±\!0.167 ±0.1100.080\!±\!0.110 ±0.0000.000\!±\!0.000 ±0.3900.480\!±\!0.390 ±0.0000.000\!±\!0.000 ±0.0890.040\!±\!0.089 ±0.0000.000\!±\!0.000 ±0.0001.000\!±\!0.000 h=30h=30 ±0.1670.360\!±\!0.167 ±0.1100.080\!±\!0.110 ±0.0000.000\!±\!0.000 ±0.0890.040\!±\!0.089 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0001.000\!±\!0.000 h=103h=10^3 ±0.1790.120\!±\!0.179 ±0.1670.160\!±\!0.167 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0001.000\!±\!0.000 h=104h=10^4 ±0.0000.000\!±\!0.000 ±0.0890.040\!±\!0.089 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0001.000\!±\!0.000 h=105h=10^5 — — — — — — — ±0.0001.000\!±\!0.000 h=106h=10^6 — — — — — — — ±0.0001.000\!±\!0.000 T-Maze (Episode SR ↑ ) h=10h=10 ±0.0001.000\!±\!0.000 ±0.2190.360\!±\!0.219 ±0.1790.680\!±\!0.179 ±0.0001.000\!±\!0.000 ±0.2610.560\!±\!0.261 ±0.0001.000\!±\!0.000 ±0.2000.400\!±\!0.200 ±0.0001.000\!±\!0.000 h=20h=20 ±0.0001.000\!±\!0.000 ±0.1790.320\!±\!0.179 ±0.3290.560\!±\!0.329 ±0.0001.000\!±\!0.000 ±0.3350.520\!±\!0.335 ±0.0001.000\!±\!0.000 ±0.2280.280\!±\!0.228 ±0.0001.000\!±\!0.000 h=30h=30 ±0.0001.000\!±\!0.000 ±0.3350.280\!±\!0.335 ±0.1670.640\!±\!0.167 ±0.0001.000\!±\!0.000 ±0.2280.520\!±\!0.228 ±0.0001.000\!±\!0.000 ±0.2280.520\!±\!0.228 ±0.0001.000\!±\!0.000 h=103h=10^3 ±0.2830.800\!±\!0.283 ±0.2280.320\!±\!0.228 ±0.2280.320\!±\!0.228 ±0.2280.320\!±\!0.228 ±0.2610.360\!±\!0.261 ±0.3030.520\!±\!0.303 ±0.0000.000\!±\!0.000 ±0.0001.000\!±\!0.000 h=104h=10^4 ±0.2610.840\!±\!0.261 ±0.1100.320\!±\!0.110 ±0.1100.320\!±\!0.110 ±0.1100.320\!±\!0.110 ±0.1670.360\!±\!0.167 ±0.2190.440\!±\!0.219 ±0.0000.000\!±\!0.000 ±0.0001.000\!±\!0.000 h=105h=10^5 — — — — — — — ±0.0001.000\!±\!0.000 h=106h=10^6 — — — — — — — ±0.0001.000\!±\!0.000 MiniGrid-Memory (Episode SR ↑ ) h=17h=17 — ±0.4040.767\!±\!0.404 ±0.0000.800\!±\!0.000 ±0.1530.333\!±\!0.153 ±0.0000.500\!±\!0.000 ±0.0000.600\!±\!0.000 ±0.0580.467\!±\!0.058 ±0.0001.000\!±\!0.000 h=51h=51 — ±0.3210.667\!±\!0.321 ±0.2310.533\!±\!0.231 ±0.0580.533\!±\!0.058 ±0.0000.300\!±\!0.000 ±0.0000.600\!±\!0.000 ±0.1150.367\!±\!0.115 ±0.0001.000\!±\!0.000 h=101h=101 — ±0.2650.600\!±\!0.265 ±0.1730.500\!±\!0.173 ±0.1530.333\!±\!0.153 ±0.0000.100\!±\!0.000 ±0.1000.500\!±\!0.100 ±0.1530.333\!±\!0.153 ±0.0001.000\!±\!0.000 h=501h=501 — ±0.1730.500\!±\!0.173 ±0.1730.500\!±\!0.173 ±0.2310.433\!±\!0.231 ±0.0000.000\!±\!0.000 ±0.1000.300\!±\!0.100 ±0.1530.333\!±\!0.153 ±0.0001.000\!±\!0.000 h=1001h=1001 — ±0.2080.467\!±\!0.208 ±0.1730.500\!±\!0.173 ±0.1530.433\!±\!0.153 ±0.0000.000\!±\!0.000 ±0.1150.367\!±\!0.115 ±0.1730.300\!±\!0.173 ±0.0001.000\!±\!0.000 LTMB-Hallway (Episode SR ↑ ) h=8h=8 — ±0.0580.033\!±\!0.058 ±0.0580.033\!±\!0.058 ±0.0580.133\!±\!0.058 ±0.0000.000\!±\!0.000 ±0.0000.100\!±\!0.000 ±0.0000.100\!±\!0.000 ±0.0001.000\!±\!0.000 h=50h=50 — ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0001.000\!±\!0.000 h=100h=100 — ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0580.933\!±\!0.058 h=500h=500 — ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0580.967\!±\!0.058 h=1000h=1000 — ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0000.000\!±\!0.000 ±0.0580.933\!±\!0.058 Random baseline omitted; its test SR remains low across every (benchmark, horizon) cell, consistent with performance shown in Table 6. E.1 Horizon-invariance experiments on synthetic and grid domains Table 11 reports the consolidated horizon-invariance results underlying Figure 3. We evaluated the trained checkpoints from Appendix C.7, by instantiating the test environment with a target horizon h taken from a pre-defined set per benchmark: Add, Scattered-Copy, T-Maze episode length h∈10,20,30,1000,10000h∈\10,20,30,1000,10000\; Keyframe Mnemonics is additionally evaluated at h∈105,106h∈\10^5,10^6\ to test horizon-invariance behavior. MiniGrid-Memory environment max size h∈17,51,101,501,1001h∈\17,51,101,501,1001\. LTMB-Hallway corridor length h∈8,50,100,500,1000h∈\8,50,100,500,1000\. Per-cell metrics are mean±~±~std across 55 random seeds for the synthetic tasks and 33 random seeds for the grid tasks. For each horizon, every method is evaluated on 55 episodes for the synthetic tasks (Add, Scattered-Copy, T-Maze) and 1010 episodes for the grid tasks (MiniGrid-Memory, LTMB-Hallway) per seed. Within the training horizon range, several baselines achieve strong performance; for example, LSTM(BPTT), DT, and TrXL reach 1.01.0 SR on T-Maze up to h=30h=30. However, baseline methods collapse once the test horizon exceeds the training horizon by even one order of magnitude. In contrast, Keyframe Mnemonics consistently exhibits horizon-invariant performance. In particular, on synthetic tasks, it generalizes up to five orders of magnitude beyond the training horizon, confirming the horizon-invariance guarantee of Theorem 2 in regimes where the underlying assumptions hold strongly. E.2 Horizon-invariance experiment on RememberColor3-ΔT T Figure 11: Hover time is varied across evaluations resulting in catastrophic failure for the baselines, while KM maintains it’s performance even at 250250 timesteps of hover. Training data had 15 timesteps of hover. The horizon-invariance study (Figure 5(a)) trains every method at a fixed cue-to-action delay ΔT=15 T=15 and evaluates the trained policies, without retraining, at delays ΔT∈15,35,55,75,95,115,135,150,200,250 T∈\15,35,55,75,95,115,135,150,200,250\. The training set is built from the released RememberColor3-v0 demonstrations, whose default cue-to-action delay is ΔTorig=5 T_orig=5. Hover-augmented data. We re-replay each successful expert demonstration in a live ManiSkill simulator with the cue-to-action delay extended to ΔT=15 T=15 and a 1010-step extra hover phase spliced into the original action tape immediately after the cue disappears (Figure 11). During the hover phase the gripper is held and small net-zero noise is added to the arm for the entire delay time. Training on the hover dataset. We run the standard three-stage KM pipeline of Appendix D.2 on the hover trajectories. Every hyperparameter matches the corresponding entries for RememberColor3 in Tables 7 (k=3k=3, proxy chunk H=8H=8, SST policy with H=1H=1, 15,00015,000 BC gradient steps). We train PTP using their official implementation on our dataset with prediction horizon of 1616, observation length of 88 and a sub-sampling rate of 88 to roughly cover the entire episode of length 7070. This indicates why we see steeper deterioration in the performance of PTP on and after 7575 steps of time delay. We train ELMUR on the same trajectories using its published RememberColor3 recipe, with the context length extended to cover the full 7070-step hover episode. All other hyperparameters (architecture, optimizer, schedule, batch size, number of epochs) are left at the authors’ defaults. Horizon-invariance evaluation. The trained policy is evaluated at every target delay ΔT T by re-instantiating the simulator with the requested ΔT T; no retraining or fine-tuning is performed at any ΔT≠15 T≠ 15. Following the MIKASA-Robo convention, an episode is counted as a success only if the binary task-completion reward is reached before the time budget elapses. We report mean ± standard error across three independent environment seeds, 100100 episodes per seed. Results are plotted in Figure 5(a) and details can be found in Table 12. Table 12: RememberColor3-ΔT T success rate under cue-to-action delays. Training delay is ΔT=15 T=15; all larger values are out-of-distribution. Mean±~±~sem over 33 seeds, 100100 episodes per seed. ΔT T PTP [46] ELMUR [14] KM (Ours) 1515 0.66±0.030.66± 0.03 0.61±0.070.61± 0.07 0.72±0.070.72± 0.07 3535 0.72±0.040.72± 0.04 0.00±0.000.00± 0.00 0.69±0.110.69± 0.11 5555 0.73±0.050.73± 0.05 0.00±0.000.00± 0.00 0.69±0.110.69± 0.11 7575 0.55±0.070.55± 0.07 0.00±0.000.00± 0.00 0.69±0.080.69± 0.08 9595 0.43±0.060.43± 0.06 0.00±0.000.00± 0.00 0.65±0.080.65± 0.08 115115 0.31±0.050.31± 0.05 0.00±0.000.00± 0.00 0.66±0.090.66± 0.09 135135 0.23±0.030.23± 0.03 0.00±0.000.00± 0.00 0.66±0.090.66± 0.09 150150 0.17±0.030.17± 0.03 0.00±0.000.00± 0.00 0.65±0.090.65± 0.09 200200 0.07±0.020.07± 0.02 0.00±0.000.00± 0.00 0.64±0.080.64± 0.08 250250 0.07±0.030.07± 0.03 0.00±0.000.00± 0.00 0.65±0.090.65± 0.09 Results in Table 12 show that KM preserves horizon-invariant performance in memory-intensive robot manipulation with high-dimensional visual data. ELMUR matches the other methods at the training delay (ΔT=15 T=15), but its external memory is overwritten as the delay increases, causing success to collapse to zero for every ΔT≥35 T≥ 35. PTP retains its training-time performance up to roughly ΔT=55 T=55 and then degrades monotonically, dropping below 10%10\% by ΔT=200 T=200. In contrast, Keyframe Mnemonics remains nearly flat across the longer horizons, retaining ≈0.65≈ 0.65 success rate at ΔT=250 T=250, i.e., ∼⁣× 17\!× the training delay, without any architectural change or fine-tuning. E.3 Reproducibility Released code, per-task configurations, and trained checkpoints, together with launch scripts, are available at https://keyframe-mnemonics.github.io. Appendix F Ablation Study We report ablation results over 5 seeds for the synthetic domains and 3 seeds for the grid and robot domains unless noted, with 100 evaluation episodes per seed. All domains are evaluated at out-of-distribution horizons. F.1 Learned selector against a random selector We replace the learned selector with a random selector πrnd _rnd that assigns priorities p∼⁡(0,1)p (0,1) and train the BC policy using Stage 3 (Algorithm 5). The resulting policy performs worse at every delay interval, demonstrating that the learned selector is critical for the memory task. Table 13: Learned selector against the random selector πrnd _rnd on MIKASA-Robo RememberColor3-ΔT T. Reported mean± success rate (%) ↑ . ΔT T Learned selector Random selector (πrnd _rnd) 15 72.3 ± 5.8 31.3 ± 3.1 35 69.0 ± 9.2 28.7 ± 2.9 55 69.0 ± 8.8 28.3 ± 2.1 75 68.7 ± 6.6 35.3 ± 1.2 95 65.0 ± 6.4 28.7 ± 4.0 115 66.3 ± 7.6 28.3 ± 4.1 135 65.7 ± 7.4 33.0 ± 1.4 150 65.3 ± 7.6 29.0 ± 0.8 200 64.3 ± 6.9 31.7 ± 1.7 250 65.3 ± 7.6 26.7 ± 1.2 F.2 Selector discount factor We vary the discount factor γ of the PPO selector to optimize over a frozen proxy policy. At γ=0γ=0 the selector greedily optimizes the per-step proxy loss. For T-Maze we report the active (default) and passive variations [64]. Greedy optimization fails at γ=0γ=0 on every domain except active T-Maze, and performance recovers as γ increases. In active T-Maze, the action changes right after the keyframe is observed, so the immediate change in proxy loss is sufficient. In the other domains, the gap between the keyframe and the decision-critical action spans multiple steps, and a higher γ is necessary. Keyframe identification from the proxy loss is therefore a delayed credit-assignment problem, which explains our choice of using RL for selector training (Section 2.2). Table 14: Effect of the PPO discount factor γ. Reported mean± . Domain Metric γ=0γ=0 γ=0.5γ=0.5 γ=0.9γ=0.9 γ=0.99γ=0.99 (default) T-Maze (active) Episode SR ↑ 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 T-Maze (passive) Episode SR ↑ 48.8 ± 6.6 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 Add Final MSE ↓ 1.84e-01 ± 1.77e-03 1.25e-02 ± 6.86e-03 1.28e-05 ± 5.67e-06 1.28e-05 ± 5.70e-06 Scattered-Copy Copy SR ↑ 0.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 F.3 Surprise coefficient We change the temperature κ of the surprise coefficient in Equation 5. Off in Table 15 corresponds to an unscaled reward, which is the κ→∞κ→∞ limit where surprise coefficient equals 1. For every task we first train with the unscaled reward and add surprise scaling only if it improves performance. κ=0.2κ=0.2 has the widest usable range and is our default for the tasks that need scaling (Add, Scattered-Copy and MiniGrid-Memory). Table 15: Effect of the surprise coefficient across temperatures. Bold marks the best setting per domain. Reported mean± . Domain Metric κ=0.02κ=0.02 κ=0.2κ=0.2 κ=2κ=2 off T-Maze Episode SR ↑ 0.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 Add Final MSE ↓ 1.28e-05 ± 5.7e-06 1.28e-05 ± 5.7e-06 1.84e-01 ± 1.8e-03 1.84e-01 ± 1.8e-03 Scattered-Copy Copy SR ↑ 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 0.0 ± 0.0 MiniGrid-Memory Success rate ↑ 100.0 ± 0.0 100.0 ± 0.0 3.0 ± 0.0 3.7 ± 0.6 LTMB-Hallway Success rate ↑ 0.0 ± 0.0 0.0 ± 0.0 0.0 ± 0.0 99.7 ± 0.6 RememberColor3 Success rate ↑ 25.0 ± 3.3 29.3 ± 6.2 34.7 ± 6.9 99.0 ± 0.0 F.4 Current observation in proxy training Including oto_t during proxy training collapses MiniGrid-Memory and RememberColor3 (Table 16), as the proxy overfits to the current observation. This is consistent with Proposition 1, which requires that no low-complexity non-invariant shortcut exists in the proxy training data. The current observation is helpful for learning the final policy, but not for learning a reward function for keyframe identification. Table 16: Effect of including the current observation oto_t in the proxy input. Reported mean± . Domain Metric Proxy without oto_t (default) Proxy with oto_t T-Maze Episode SR ↑ 100.0 ± 0.0 100.0 ± 0.0 Add Final MSE ↓ 1.28e-05 ± 5.7e-06 2.20e-05 ± 1.3e-05 Scattered-Copy Copy SR ↑ 100.0 ± 0.0 100.0 ± 0.0 MiniGrid-Memory Success rate ↑ 100.0 ± 0.0 33.3 ± 26.3 LTMB-Hallway Success rate ↑ 99.7 ± 0.6 93.0 ± 9.5 RememberColor3 Success rate ↑ 99.0 ± 0.0 18.7 ± 5.9 F.5 Queue tie-break rule The memory buffer uses a FIFO queue (Algorithm 2) during selector training and a LIFO queue (Algorithm 3) at inference. The two differ in which slot is evicted when several observations are tied at the lowest priority. We vary this rule independently at training and inference. On the synthetic and grid domains, all four combinations are identical (Table 17). The rule matters on RememberColor3-ΔT T, where the default FIFO → LIFO beats its reverse at 9 of 10 retention intervals (Table 18). Table 17: Queue tie-break rule, train → inference. Reported mean± . Domain Metric FIFO → FIFO FIFO → LIFO (default) LIFO → FIFO LIFO → LIFO T-Maze Episode SR ↑ 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 Add Final MSE ↓ 1.28e-05 ± 5.7e-06 1.28e-05 ± 5.7e-06 1.28e-05 ± 5.7e-06 1.28e-05 ± 5.7e-06 Scattered-Copy Copy SR ↑ 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 MiniGrid-Memory Success rate ↑ 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 100.0 ± 0.0 LTMB-Hallway Success rate ↑ 99.6 ± 0.5 99.6 ± 0.5 99.6 ± 0.5 99.6 ± 0.5 RememberColor3 Success rate ↑ 93.7 ± 4.5 99.0 ± 0.0 98.7 ± 1.3 96.0 ± 2.2 Table 18: Queue tie-break rule on RememberColor3-ΔT T. Reported mean± success rate ↑ . ΔT T FIFO → LIFO (default) LIFO → FIFO 15 72.3 ± 5.8 75.0 ± 1.6 35 69.0 ± 9.2 66.0 ± 3.6 55 69.0 ± 8.8 66.0 ± 2.4 75 68.7 ± 6.6 59.3 ± 3.7 95 65.0 ± 6.4 59.3 ± 1.7 115 66.3 ± 7.6 58.3 ± 2.5 135 65.7 ± 7.4 60.3 ± 1.7 150 65.3 ± 7.6 59.3 ± 2.5 200 64.3 ± 6.9 58.0 ± 1.6 250 65.3 ± 7.6 57.7 ± 3.8 F.6 Memory buffer size We scale the memory buffer size k to 0.5×0.5× and 2×2× the default value and retrain all three stages. Halving k degrades performance across all domains, while doubling it helps on some domains and hurts on others. A larger k can make the KM policy harder to train, and our default values serve as a minimum-capacity floor, which were chosen per domain independently. Table 19: Effect of scaling the memory buffer size k. Reported mean± . Domain Metric 0.5×0.5× 1×1× (default) 2×2× T-Maze Episode SR ↑ 40.0 ± 23.2 100.0 ± 0.0 100.0 ± 0.0 Add Final MSE ↓ 9.55e-02 ± 2.38e-02 1.28e-05 ± 5.7e-06 2.33e-05 ± 2.03e-05 Scattered-Copy Copy SR ↑ 0.0 ± 0.0 100.0 ± 0.0 80.0 ± 44.7 MiniGrid-Memory Success rate ↑ 35.3 ± 26.4 100.0 ± 0.0 20.7 ± 27.2 LTMB-Hallway Success rate ↑ 0.0 ± 0.0 81.7 ± 15.0 99.3 ± 1.2 RememberColor3 Success rate ↑ 94.7 ± 2.5 99.0 ± 0.0 100.0 ± 0.0 Appendix G Additional Analysis G.1 Interpretability of discovered keyframes For a well-trained selector, the mnemonics buffer ℬtB_t should accumulate keyframes that are information-critical for the task and provide necessary context for the downstream BC policy to generate actions unambiguously. Figure 6 shows the buffer contents for our visual benchmark domains. We observe that the buffer keyframes qualitatively capture the belief state for each task by retaining observations with task-relevant cues. The discovery process is fully self-supervised, as the selector is trained only from expert demonstration state-action pairs without any additional temporal keyframe annotation. These qualitative results highlight the interpretability of Keyframe Mnemonics: the buffer contains prior observations that can be inspected to understand what the BC policy relies on when taking a particular action in a partially observable environment. G.2 Selector priority behavior In Figure 12 we evaluate trained selector policy πs _s across synthetic domains to qualitatively assess how well the learned priorities align with the oracle keyframes defined in Appendix C.1. We find that the selector learns a near-optimal selection function that agrees with the oracle, highlighting that the proxy objective provides a strong signal for keyframe discovery, consistent with Corollary 1. During Stage 2 RL training (Algorithm 4), the selector explores priority assignments and learns to commit high priorities to decision-relevant keyframes. This qualitatively explains the horizon-invariant results in Table 11, where BC policies achieve perfect performance over horizons orders of magnitude beyond training. Since the selector is Markovian and depends only on the current observation, its behavior remains invariant to horizon length, allowing it to admit only informative keyframes into the buffer at inference time. 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 1010 evaluation episodes (rows) over time (columns), where each column corresponds to a single environment step and color encodes either the selector output (pt∈[0,1]p_t∈[0,1]) or the binary oracle indicator (11 for keyframes and 00 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][0,1], the mean selector priority is exactly 1.001.00 on keyframes and 0.000.00 on background frames. This empirical result shows that the selector policy πs _s learns a near-optimal keyframe selection function for domains where our assumptions hold strongly. G.3 BC Policy activation map To understand whether the discovered keyframes are useful to the behavior cloning policy, we use Grad-CAM [65, 66], an explainable AI technique commonly used to visualize which parts of the input a policy focuses on when generating actions. In Figure 13, we visualize activation maps at decision-critical timesteps in grid and robot manipulation domains, where the BC policy must rely on memory-buffer information to generate the correct action. We observe that the BC policy attends to information cues from both the buffer and the current observation to infer the POMDP belief state and act correctly. In the grid domain, the heatmap is concentrated around the matching object cue. In RememberColor3, activations concentrate around the blue colored cube in both the keyframe and the current observation. In TakeItBack, activations appear around the initial cube location in the first keyframe, while the second keyframe captures temporal task-progress information just before the robot begins pushing the cube, helping disambiguate the state so that the policy pulls the cube back to its initial location. These results show that the keyframes discovered by our method are decision-relevant and provide necessary context for BC policies to disambiguate actions under partial observability, even in high-dimensional robot visual domains. Figure 13: BC policy activation map on mnemonics. Grad-CAM [65, 66] activation maps of the behavior cloning policy πBC _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. G.4 Regret on low memory budget k Figure 14 analyzes how memory regret (Section 3.1) changes with buffer capacity k for synthetic domains. In T-Maze, active regret remains high when k=1k=1 (⌈0.5×2⌉ 0.5× 2 ), which is below the true number of keyframes (22), because the buffer cannot retain all required keyframes. In Add, active regret becomes negative when k=2k=2 (⌈0.5×3⌉ 0.5× 3 ), also below the true keyframe count (33). This occurs because the selector prioritizes the two marked numbers over the emit step, whereas the oracle assigns equal priority to all three keyframes. Under the low memory budget, the oracle evicts the first marked step, leading to higher proxy loss. It shows that the selector can learn a non-uniform priority policy that favors keyframes with higher mutual information with the action. For Scattered-Copy, active regret remains close to zero since the oracle and selector perform similarly across memory budgets. Passive regret remains low when the buffer size is smaller than the true keyframe count because the proxy cannot converge to an invariant hypothesis under constrained memory. At the 2×2× buffer size, passive regret is lower than at 1×1×, since a larger buffer increases the probability that a random selector retains keyframes. Figure 14: Memory regret as a function of buffer capacity k. The 1×1× 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. G.5 PPO reward plots Figure 15 shows the learning curves for training the selector policy πs _s on the frozen proxy objective (Section 2.2) using PPO [39]. For synthetic domains where our assumptions hold strongly, we observe clear learning dynamics, with the selector converging within 200200K environment (Algorithm 4) steps . For robot manipulation domains, the selector converges within 55M steps. The plot for ChainOfColors illustrates a task that is more difficult to optimize, exhibiting a noisy learning curve. We observed that such noise typically appears when the proxy policy does not converge properly, either because it is undertrained or because our assumptions do not hold well in the environment. Figure 15: Selector training PPO learning curves. The selector policy converges within 200200K environment steps for synthetic domains and 55M 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. G.6 Limitation discussion Our Assumptions 1 is restrictive, as it requires keyframes to be individually identifiable from a single observation because the selector policy πs(pt∣ot) _s(p_t o_t) is Markovian. This formulation allows selector training to scale without recurrent modules, thereby avoiding backpropagation through time. Hence, our method is best suited for sparse, frame-localizable memory tasks. However, problems that do not satisfy the assumption may require multiple keyframes to represent the same underlying information, leading to multiple information-equivalent buffer states that achieve similar objective values under the proxy policy. In such cases, it becomes challenging to retain only one mnemonic set among several equivalent alternatives. Our Assumption 2 requires action ambiguity to be strongly correlated with a sparse set of keyframes. In other words, the action distribution must exhibit sufficient multimodality that can be disambiguated by the remembered keyframes. In tasks where actions are mostly local and exhibit low variance, actions may carry little mutual information with the buffer observations. In such a regime, the mutual information gap in Theorem 1 becomes negligible, and the proxy policy may fail to converge because the invariant hypothesis is difficult to learn from a sparse memory buffer. We observed this failure mode on the VizDoom-Two-Colors [22] task, where the agent observes a color cue and must later choose the corresponding target, but the discrete navigation actions remain largely local. We found that the proxy policy struggles to converge to the true invariant hypothesis when expert actions at decision-critical timesteps depend simultaneously on multiple unique keyframes. This occurs because the probability that masked-sampling selects all relevant keyframes during proxy training decreases as the number of action-relevant keyframes increases, making it difficult to learn tasks requiring high memory capacity. A similar issue arises in very long demonstrations with few keyframes, where keyframe sparsity further reduces the probability of sampling the relevant frames. A straightforward way to address this issue is to increase the training memory budget, which makes sampling keyframes into the buffer more likely and also increases the support size M, thereby enforcing Theorem 1 more effectively. Another promising direction is to train proxy policies on shorter low-memory-capacity tasks, learn specialized selectors for each, and compose them to support high-memory-capacity BC policies, which we leave for future work. The PPO training time dominates the computational cost in robot-domain experiments. Our current implementation uses SB3 [60], which introduces nontrivial CPU-GPU overhead and becomes the main bottleneck. The selector training environment (Algorithm 4) primarily requires proxy policy inference, which can in principle be executed entirely on the GPU. We therefore view this as an engineering limitation. In future work, we plan to use a GPU-native PPO implementation to accelerate selector training. Appendix H Extended Related Works Simplicity-bias in neural networks. Simplicity-bias is a well-studied property of neural networks, reflecting their tendency to prefer simpler patterns over more complex ones [25, 26, 67, 68, 27]. Recent work has leveraged simplicity bias to improve RL performance [69], but this area remains largely underexplored, particularly in how such biases can be effectively utilized. We present one of the first methods that explicitly leverage this bias for behavior cloning in POMDP domains. Masking-induced invariant feature learning. Prior works have used masking in self-supervised learning over language tokens [70] and spatial image patches [71] to learn invariant features from input data. Our keyframe discovery process instead applies masking over the temporal history of individual observations. We hypothesize that this induces a similar regularization effect: random input sampling increases the difficulty of learning complex, non-invariant features, thereby biasing the model toward invariant hypotheses that better capture true correlations over spurious ones. Language as a proxy for state tracking. Recent methods have used segment-level language-annotated demonstrations to finetune vision-language models (VLMs) that generate autoregressive subgoal commands for policy steering [72, 73]. While language provides high-level semantic guidance, it lacks the spatio-temporal detail present in visual observations. In contrast, our approach does not strictly require language annotations for training. Instead, we view Keyframe Mnemonics as complementary to such methods, with informative keyframes addressing partial observability and language supporting abstract reasoning over past events. NeurIPS Paper Checklist 1. Claims Question: Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope? Answer: [Yes] Justification: We have clearly stated the contributions in the introduction and have ensured that the Methodology, and the Experiments and Results Analysis sections back up the three contributions clearly. Guidelines: • The answer [N/A] means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A [No] or [N/A] answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 2. Limitations Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes] Justification: The paper has a limitations section where we have stated clearly the limitations of the current work. Guidelines: • The answer [N/A] means that the paper has no limitation while the answer [No] means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate “Limitations” section in their paper. • The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren’t acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 3. Theory assumptions and proofs Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes] Justification: In the methodology section and the appendix we have attached mathematical proofs of our claims which are correct in our opinion. Guidelines: • The answer [N/A] means that the paper does not include theoretical results. • All the theorems, formulas, and proofs in the paper should be numbered and cross-referenced. • All assumptions should be clearly stated or referenced in the statement of any theorems. • The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition. • Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material. • Theorems and Lemmas that the proof relies upon should be properly referenced. 4. Experimental result reproducibility Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)? Answer: [Yes] Justification: We are submitting self contained code to run our framework on two of the domains demonstrated in our results. We have aimed to provide parameters for our architectures where appropriate. We have not collected any data as part of this paper. Guidelines: • The answer [N/A] means that the paper does not include experiments. • If the paper includes experiments, a [No] answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not. • If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable. • Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed. • While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example (a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm. (b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully. (c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset). (d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 5. Open access to data and code Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We have released two self contained domains with our architecture that are runnable. This should allow the reviewers to verify our claims. We have also attached our webpage with more information and access to github. Guidelines: • The answer [N/A] means that paper does not include experiments requiring code. • Please see the NeurIPS code and data submission guidelines (https://neurips.c/public/guides/CodeSubmissionPolicy) for more details. • While we encourage the release of code and data, we understand that this might not be possible, so [No] is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark). • The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https://neurips.c/public/guides/CodeSubmissionPolicy) for more details. • The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc. • The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why. • At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable). • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 6. Experimental setting/details Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer) necessary to understand the results? Answer: [Yes] Justification: The results and Appendix B aim to provide complete experimental details for reproducibility. Guidelines: • The answer [N/A] means that the paper does not include experiments. • The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 7. Experiment statistical significance Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments? Answer: [Yes] Justification: The error bars are visible in each of the figures and claims are made with respect to these error bars. Guidelines: • The answer [N/A] means that the paper does not include experiments. • The authors should answer [Yes] if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g., negative error rates). • If error bars are reported in tables or plots, the authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 8. Experiments compute resources Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments? Answer: [Yes] Justification: Appendix C and D specified compute times for our methods and computational resources used. Guidelines: • The answer [N/A] means that the paper does not include experiments. • The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn’t make it into the paper). 9. Code of ethics Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.c/public/EthicsGuidelines? Answer: [Yes] Justification: This work has not collected any human data. We have listed computational costs and released runnable code to demonstrate reproducibility. Our work is in introductory stages so does not have comparable societal impact to existing generative AI technologies as of now. The goal is to improve memory of neural networks to eventually perform long horizon tasks on robots which would be more directly useful to our society. Guidelines: • The answer [N/A] means that the authors have not reviewed the NeurIPS Code of Ethics. • If the authors answer [No] , they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [N/A] Justification: The work on improving memory for neural networks is not practical yet for direct societal interaction. This work is on artificial benchmark and simulated robot tasks which is not meaningful to society at large in its current state. Guidelines: • The answer [N/A] means that there is no societal impact of the work performed. • If the authors answer [N/A] or [No] , they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate Deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 11. Safeguards Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pre-trained language models, image generators, or scraped datasets)? Answer: [N/A] Justification: The work is too preliminary to have a high risk of misuse in a societal scale. Guidelines: • The answer [N/A] means that the paper poses no such risks. • Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected? Answer: [Yes] Justification: We have cited the papers whose code and/or data we have used to run our experiments. All other code is our own. Guidelines: • The answer [N/A] means that the paper does not use existing assets. • The authors should cite the original paper that produced the code package or dataset. • The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., C-BY 4.0) should be included for each asset. • For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. • If this information is not available online, the authors are encouraged to reach out to the asset’s creators. 13. New assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [Yes] Justification: We are releasing code as part of the supplementary material and on our webpage. No human subjects data was collected or used as part of this project. Guidelines: • The answer [N/A] means that the paper does not release new assets. • Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 14. Crowdsourcing and research with human subjects Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)? Answer: [N/A] Justification: We did not collect any human subjects data. Guidelines: • The answer [N/A] means that the paper does not involve crowdsourcing nor research with human subjects. • Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper. • According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector. 15. Institutional review board (IRB) approvals or equivalent for research with human subjects Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained? Answer: [N/A] Justification: We did not collect any human subjects data. Guidelines: • The answer [N/A] means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 16. Declaration of LLM usage Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigor, or originality of the research, declaration is not required. Answer: [N/A] Justification: LLMs were not used to develop the core idea or experimental design of this research. Guidelines: • The answer [N/A] means that the core method development in this research does not involve LLMs as any important, original, or non-standard components. • Please refer to our LLM policy in the NeurIPS handbook for what should or should not be described.

Instant research discovery