Gated Memory Policy

TL;DR

Gated Memory Policy (GMP) is a novel visuomotor approach that learns when and what to recall from memory, significantly improving performance on non-Markovian robotic tasks by using a gating mechanism and cross-attention.

Contribution

GMP introduces a learned memory gate and a lightweight cross-attention module to enhance memory recall in robotic policies, addressing distribution shift and overfitting issues.

Findings

GMP improves success rates by 30.1% on the MemMimic benchmark.

GMP maintains competitive performance on Markovian tasks.

Injecting diffusion noise enhances robustness to noisy histories.

Abstract

Robotic manipulation tasks exhibit varying memory requirements, ranging from Markovian tasks that require no memory to non-Markovian tasks that depend on historical information spanning single or multiple interaction trials. Surprisingly, simply extending observation histories of a visuomotor policy often leads to a significant performance drop due to distribution shift and overfitting. To address these issues, we propose Gated Memory Policy (GMP), a visuomotor policy that learns both when to recall memory and what to recall. To learn when to recall memory, GMP employs a learned memory gate mechanism that selectively activates history context only when necessary, improving robustness and reactivity. To learn what to recall efficiently, GMP introduces a lightweight cross-attention module that constructs effective latent memory representations. To further enhance robustness, GMP injects diffusion noise into historical actions, mitigating sensitivity to noisy or inaccurate histories during both training and inference. On our proposed non-Markovian benchmark MemMimic, GMP achieves a 30.1% average success rate improvement over long-history baselines, while maintaining competitive performance on Markovian tasks in RoboMimic. All code, data and in-the-wild deployment instructions are available on our project website https://gated-memory-policy.github.io/.