Gated Memory Policy

1. Problem Statement

Robotic manipulation tasks exhibit a wide spectrum of memory requirements:

• Markovian Tasks: Require no history (e.g., simple pick-and-place).
• In-Trial Memory: Require recalling information within a single episode (e.g., holding an object while waiting for a signal).
• Cross-Trial Memory: Require aggregating information across multiple attempts to infer unobservable physical properties (e.g., friction, mass) and adapt actions (e.g., iterative casting).

The Core Challenge:
Naively extending the observation history of visuomotor policies (e.g., feeding long sequences of images and actions into a Transformer) leads to significant performance degradation, particularly on Markovian tasks. This occurs due to:

1. Distribution Shift: Noisy or unfamiliar frames in a long history shift the input distribution away from the training data.
2. Overfitting: Increased input dimensionality makes it difficult for models to generalize.
3. Computational Cost: Self-attention mechanisms scale quadratically ($O(H^2)$) with history length $H$, making long-context inference expensive.
4. Reactivity: Constantly attending to irrelevant history distracts the policy from current observations, slowing inference.

Existing solutions often rely on manual memory structures (keyframes, text plans) or fail to balance the trade-off between robustness on non-Markovian tasks and efficiency on Markovian ones.

2. Methodology: Gated Memory Policy (GMP)

GMP is a visuomotor policy framework based on Diffusion Transformers (DiT) that learns to selectively recall memory "on demand." It addresses the problem through three key innovations:

A. Learned Memory Gate (When to Recall)

Instead of always attending to history, GMP employs a binary memory gate ($\mu_t \in \{0, 1\}$) that decides whether to activate the history cross-attention module at each timestep.

• Mechanism: An MLP takes the current image observation ($I_t$) and proprioception ($P_t$) and outputs a gate value via a sigmoid function.
• Calibration Strategy: Since training the gate end-to-end is unstable (the policy tends to overuse history to minimize training loss), the authors propose a self-supervised calibration process:
  1. Train two separate policies on a training set: one with memory always off ($\pi$) and one with memory always on ($\pi_{mem}$).
  2. Evaluate both on a validation set to compute action prediction errors ($\delta_t$ and $\delta^{mem}_t$).
  3. Label a timestep as "memory-required" ($\mu_t=1$) if the error of the no-memory policy is significantly larger than the memory policy (i.e., $\delta_t > \theta \cdot \delta^{mem}_t$).
  4. Freeze the gate weights and retrain the final policy on the full dataset using these labels.
• Benefit: This ensures the policy ignores history for Markovian tasks (maintaining reactivity) but activates it only when critical historical context is needed.

B. Cross-Attention with Cached Tokens (What to Recall)

To avoid the quadratic computational cost of standard self-attention over long histories:

• Architecture: GMP uses a cross-attention module where the current observation queries a cache of historical tokens.
• Efficiency: Historical action trajectories and image features are cached (sliding window). The computational cost scales linearly ($O(H)$) with history length rather than quadratically.
• Representation: Images are encoded via a ViT and pooled into single tokens; actions are encoded via MLPs. This reduces the token count significantly.

C. Diffusion Noise Augmentation (Robustness)

To prevent the policy from overfitting to "clean" historical data and to handle noisy real-world histories:

• Strategy: During training, diffusion noise is injected into the historical action inputs.
• Schedule: At diffusion step $k$, the model predicts future actions conditioned on history actions at step $k-1$ (one step "cleaner").
• Effect: This forces the policy to learn robust representations that do not rely on perfect historical accuracy, mitigating sensitivity to accumulated errors during inference.

3. Key Contributions

1. Gated Memory Policy (GMP): A practical framework that dynamically switches between memory-dependent and memory-independent modes, achieving strong performance on both Markovian and non-Markovian tasks without sacrificing reactivity.
2. MemMimic Benchmark: A new suite of simulation and real-world benchmarks specifically designed to evaluate cross-trial memory (learning from trial-and-error) and in-trial memory, filling a gap left by existing benchmarks like RoboMimic.
3. Comprehensive Empirical Study: Demonstrates that GMP outperforms long-history baselines by 30.1% on non-Markovian tasks while maintaining competitive performance on Markovian tasks, with significantly lower inference latency.

4. Experimental Results

The authors evaluated GMP on MemMimic (non-Markovian) and RoboMimic (Markovian).

• Performance on Non-Markovian Tasks (MemMimic):

  • GMP achieved a 30.1% average success rate improvement over long-history baselines (e.g., Long-hist DP, Long-hist PTP).
  • Tasks:
    • Iterative Pushing/Casting: GMP successfully inferred unknown friction coefficients by attending to outcomes from previous trials, whereas baselines failed to adapt.
    • Match Color (with Delay): GMP maintained 99% success even with random delays of up to 600 seconds, proving its ability to handle long temporal gaps.
  • Attention Visualization: Attention weights correctly focused on critical past frames (e.g., the moment a bin color was first observed) while ignoring irrelevant intermediate steps.

• Performance on Markovian Tasks (RoboMimic):

  • Unlike naive long-history policies which suffer performance drops on simple tasks (e.g., Square, Transport), GMP maintained competitive success rates.
  • The memory gate remained off (µ=0) for the majority of timesteps (e.g., 73% in Match Color, 100% in Square), effectively acting as a standard diffusion policy when memory was unnecessary.

• Efficiency:

  • Inference Time: GMP's inference time remains constant and minimal when the gate is off. When the gate is on, the cost scales linearly with history length, whereas self-attention baselines scale quadratically.
  • Robustness: The noise injection strategy significantly improved success rates in tasks with high uncertainty (e.g., Iterative Pushing) compared to baselines trained on clean history.

5. Significance

This work addresses a fundamental bottleneck in robotic learning: the trade-off between memory capacity and generalization/efficiency.

• Paradigm Shift: It moves away from "longer is better" (blindly increasing context windows) toward "smart memory" (selective, on-demand retrieval).
• Real-World Applicability: By demonstrating success on real-world tasks involving unknown physical properties (friction, mass) and noisy environments, GMP provides a viable path for robots to learn from trial-and-error in unstructured settings.
• Scalability: The linear complexity of the cross-attention mechanism and the gating strategy make it feasible to deploy long-horizon policies on edge hardware without prohibitive computational costs.

In summary, Gated Memory Policy offers a robust, efficient, and adaptive solution for robotic manipulation, enabling agents to learn complex, history-dependent behaviors without compromising performance on simpler, memory-free tasks.