Omkar Patil

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

Generative robot policies trained on demonstrations using behavior cloning often learn actions that are sub-optimal or misaligned with respect to the downstream task. Policy improvement approaches aim to bridge this gap and improve the cumulative reward with minimal interventions on the pre-trained policy. However, an impediment to their practical deployment is the requirement of cumbersome hyperparameter tuning specific to model or task families. In this work, we present an episodic, derivative-free latent policy improvement approach that works with little to no changes across model scales from MLPs to VLAs. We demonstrate that the performance of a pretrained, frozen diffusion or flow matching policy can be improved with respect to a downstream reward by swapping the sampling of initial noise from the prior distribution (typically isotropic Gaussian) with a well-chosen, constant initial noise input – a golden ticket. We show the prevalence of golden tickets by improving the policy performance of 46 out of 51 tasks across manipulation benchmarks, with absolute improvements in success rate by up to 79% for simulated tasks, and 28% within 60 search episodes for real-world tasks. Further, we find that the versatility of our approach opens up new avenues such as simultaneous policy improvement for multiple downstream rewards.