PaperlayerView exact revision on arXiv

Source-linked AI summary

Gated Memory Policy: In-Context Memorization and Adaptation

Yihuai Gao, Jeff Jinyun Liu, Shuang Li, Shuran Song

arXiv:2604.18933v2cs.ROcs.AI

TL;DR

Robotic manipulation ranges from Markovian tasks to tasks requiring in-trial memorization or cross-trial adaptation, but simply extending policy histories can cause overfitting and distribution shift. GMP learns when to recall history with a gate and what to recall with cross-attention, and reports a 30.1% improvement on non-Markovian tasks while maintaining competitive Markovian performance.

  • Problem

    Simply extending visuomotor policy histories can cause overfitting and distribution shift, while some manipulation tasks require recalling information within or across trials.

    Source: §1 · p. 2

  • Method

    GMP uses a calibrated binary gate to control history recall and cross-attention to retrieve relevant history, with diffusion noise applied to historical actions.

    Source: §1 · p. 2, §3.2 · p. 4

  • Results

    30.1% improvement on non-Markovian tasks, while maintaining competitive performance on Markovian benchmarks.

    Source: §1 · p. 2

  • Takeaways & Limitations

    GMP outperforms long-history baselines on MemMimic in simulation and real-world settings while avoiding performance degradation on Markovian tasks.

    Source: §6 · p. 9

    GMP’s attention window is finite, so it may lose access to critical information over extended sequences.

    Source: §5 · p. 8

Abstract

from arXiv · show

Robotic manipulation tasks exhibit varying memory requirements, ranging from Markovian tasks that require no memory to non-Markovian tasks that demand in-context memorization of historical information within a single trial or in-context adaptation based on the outcomes of multiple past 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/.

1 Introduction

Robotic tasks range from requiring no memory to needing information recalled within or across trials, but simply extending policy histories can make learning unreliable and inefficient. GMP selectively recalls history, and the paper reports improved non-Markovian performance while maintaining competitive Markovian performance.

1 Introduction

  • Robotic manipulation ranges from Markovian tasks needing minimal history to in-trial memorization and cross-trial adaptation from repeated interactions.Cross-trial adaptation can require inferring unobservable object or environment properties.Source: §1 · p. 1
  • Naively extending histories can degrade Markovian performance, increase overfitting and distribution shift, and raise computational costs.Long histories can also make policies less reactive when history distracts from current observations.Source: §1 · p. 2
  • GMP learns when to recall history with a calibrated gate and what to recall with cross-attention, while MemMimic evaluates in-trial and cross-trial memory.The paper also reports a simulation and real-world empirical study of memory-related design choices.Source: §1 · p. 2
  • 30.1% improvement on non-Markovian tasks, while maintaining competitive performance on Markovian benchmarks, is the reported overall result for GMP.The paper evaluates GMP on MemMimic and RoboMimic.Source: §1 · p. 2

2 Related Work

Prior robot policies store history as trajectories, traces, semantic representations, or latent memories, often with task-specific design choices. Long-context methods address memory needs but can trade off generalization, action-observation alignment, or reactivity.

2 Related Work

  • Prior memory policies use action parameterization or trajectory tracking, including keyframe heatmaps [9], object trajectories [10], and visual trace overlays [11].Other approaches store environment dynamics [12], LLM-generated plans [13], or observation and prompt embeddings.Source: §2 · p. 2
  • Structured memory often requires manual keypoint, keyframe, or language-planning design, whereas GMP learns recall from raw images and action trajectories.The paper presents this as a task-agnostic alternative intended to generalize across tasks.Source: §2 · p. 3
  • RNN policies have been explored [1] [15], and shorter or absent histories often perform better on Markovian tasks [16] [18] [19] [20], while non-Markovian tasks need long context [21] [22] [23].Source: §2 · p. 3
  • Long-context approaches use memory banks [24] [25] [26], keyframe selection [27], attention aggregation [28], history noise [29] [30], or auxiliary prediction [8] [31] [32], but can trade off alignment or reactivity.GMP attends to both action and image tokens with a cross-attention module whose complexity is linear in history length.Source: §2 · p. 3

3 Method

GMP builds on a diffusion-transformer policy, then conditions it on cached historical action and image tokens through cross-attention. A calibrated binary gate controls recall, while diffusion-scheduled noise augments history actions.

3.1 Background: Transformer-based Diffusion Policy

  • Extending history in a transformer policy risks overfitting and makes self-attention over H historical tokens scale quadratically as O(H^2).The diffusion-policy backbone predicts future actions from current image observations and robot proprioception, then denoises them at test time.Source: §3.1 · p. 3

3.2 Gated Memory Policy

  • A binary memory gate controls whether history cross-attention is applied, alongside noised history actions and cached history tokens for inference.The figure describes these as the gated attention module’s three design features.Source: Figure 2 · p. 3
  • GMP encodes historical action trajectories with an MLP and samples one image per trajectory, pooling visual patches into a token.A sliding window of trajectory tokens is cached so history can be retrieved without recomputing it.Source: §3.2 · p. 4
  • Diffusion-scheduled noise is injected into past actions as history conditioning to reduce overfitting and improve robustness and generalization.The augmented history action is one diffusion step cleaner than the future action being predicted.Source: §3.2 · p. 4
  • A learned gate modulates history’s contribution at each deployment step because useful history may be needed only at occasional critical moments.The binary gate is computed from current image and proprioception and controls the history cross-attention output before the feed-forward layer.Source: §3.2 · p. 4
  • Match Color tests visual memory after bin colors are shuffled, while Discrete Place Back tests recall of a cube’s original bin after a two-second hold.Source: Figure 4 · p. 5, Figure 5 · p. 5
  • Gate calibration labels validation timesteps memory-required when adding history significantly reduces action prediction error, then freezes gate weights before full-data retraining.End-to-end gate training is difficult because regularization weights that suppress unnecessary recall can also suppress useful history.Source: §3.2 · pp. 4–5

4 Experiments

Experiments test GMP on in-trial and cross-trial memory tasks, memory-free RoboMimic tasks, and MIKASA-Robo. The reported results cover benchmark comparisons, history selection, gating, and inference cost.

Benchmarks

  • MemMimic evaluates in-trial memorization and cross-trial adaptation in simulation and real-world tasks; RoboMimic and MIKASA-Robo provide additional evaluation settings.The baseline comparison includes Diffusion Policy, Past Token Prediction, and BC-RNN with varying history lengths.Source: §4 · p. 5

Markovian tasks

  • On RoboMimic, GMP maintains competitive performance on Markovian tasks, while Long-hist DP (120 steps) performs worse than Mid-hist DP (16 steps) and No-hist DP.RoboMimic evaluation includes Tool Hang, Square, and Transport; the paper attributes GMP's competitive performance to its gating mechanism.Source: Figure 10 · p. 7, §4 · p. 7

Benchmarks

  • 26.6% average improvement over MemoryVLA on MIKASA-Robo is reported, alongside stronger performance than long-history baselines on complex real-world and cross-trial MemMimic tasks.MIKASA-Robo evaluation covers five tasks; the MemMimic comparison distinguishes simpler tasks from complex real-world and cross-trial tasks.Source: Figure 11 · p. 7, §4 · p. 7

History attention

  • History cross-attention attends to the first-seen color frame during Match Color and to prior overshoot and undershoot outcomes during Iterative Pushing.These examples show attention to relevant history without explicit supervision, including across trials when adapting a push.Source: §4 · p. 7

Memory gate

  • Even on tasks needing history, the gate is off for 73% of Match Color and 58% of Iterative Pushing, and opens when needed.On Markovian tasks, the gate is usually off when action-prediction errors with and without memory are similar.Source: §4 · p. 7
  • Validation-set calibration selects moments when memory helps; straight-through gate training risks overuse, while strong regularization can suppress memory and impair non-Markovian tasks.Calibration compares action-prediction error with and without memory on validation samples.Source: §4 · pp. 7–8

Inference cost

  • GMP reduces inference cost by moving history into cross-attention, whose cost scales linearly with history length, and skipping history attention when the gate is off.The comparison reports average inference time and standard deviation over 100 rounds on an NVIDIA RTX 3080 GPU; self-attention DiT cost scales quadratically with context length.Source: §4 · p. 8

Noise injection

  • The Iterative Pushing ablation compares diffusion-noised history actions with clean, random-level, and diffusion-forcing alternatives; clean history can be brittle under evaluation distribution shifts.The paper describes the proposed diffusion-noising strategy as improving policy robustness.Source: §4 · p. 8

5 Limitations and Future Work

GMP's current cross-attention has a finite history window, limiting access to information from sufficiently distant points in extended sequences.

Limitations and future work

  • GMP's finite attention window may prevent access to critical information from extended sequences; selective caching and token replacement are proposed for future work.The proposed strategies would dynamically manage the history buffer to retain high-value tokens from the distant past within a fixed attention budget.Source: §5 · p. 8

6 Conclusion

GMP combines gated memory recall, cross-attention, and diffusion noise on history actions to support in-context memorization and adaptation. It outperforms long-history baselines on MemMimic while avoiding performance degradation on Markovian tasks.

Conclusion

  • GMP combines cross-attention over history tokens, a calibrated binary memory gate, and diffusion noise on history actions to selectively recall task-relevant information.The conclusion attributes in-context memorization and adaptation, robustness to long histories, and low computational cost to this framework.Source: §6 · pp. 8–9
  • GMP outperforms long-history baselines on the MemMimic non-Markovian task suite in simulation and real-world settings, without degrading performance on Markovian tasks.The conclusion says the method supports in-context memorization and adaptation while maintaining low computational cost.Source: §6 · pp. 8–9

8 Supplementary Materials

The supplement reports long-history recall performance, an overlapping-trajectory training method, and the task and evaluation details for MemMimic.

8.1 How Long History Can GMP Recall?

  • 99.0%±1.0% success is reported on Match Color after training on 500 episodes with a 6000-frame image and action buffer.Inference takes 0.16s for 8 denoising steps on an NVIDIA RTX 5090 GPU; the authors state that this history length suffices for short-term memory, while a higher-level VLM could handle long-term memory more efficiently via text representations.Source: §8.1 · p. 13

8.2 Overlapped Trajectory Training

  • Overlapped Trajectory Training batches overlapping subsequences so each image is encoded once per mini-batch and its feature reused across subsequences.Memory scales with unique images per batch rather than image occurrences across subsequences, while the vision encoder trains end-to-end on long-history dependencies.Source: §8.2 · p. 13

8.3 MemMimic Details

  • MemMimic's in-trial tasks test visual or spatial recall, including restoring a shuffled-color cube, returning a held cube, and placing a cup back after a human flips it.The tasks span simulation and real-world settings, including an in-the-wild task with unseen environments and diverse objects.Source: §8.3 · pp. 13–14
  • MemMimic's cross-trial tasks test adaptation across repeated attempts with fixed but unknown object properties, including pushing, flinging, and casting.Episode success is based on the last M of N trials; the tasks vary in trial counts and success criteria.Source: §8.3 · p. 14
  • Simulation tasks use MuJoCo and 100 evaluation episodes per checkpoint; results report the best checkpoint's mean and standard deviation across three training seeds.In-domain real-world tasks use a UR5 and WSG50 gripper, while the in-the-wild T3 task uses an ARX X5 with the same UMI setup.Source: §8.3 · p. 14

T1: Match Color (Sim)

The section evaluates memory-demanding place-back tasks in simulation and real continuous settings, contrasting history-based baselines and GMP. Long-history proprioception can fail under distribution shift, while sufficient history can solve the color-recall task.

T1: Match Color (Sim)

  • T1 varies cube position and four-bin color assignments across 96 configurations, and its pause requires more than two seconds of memory.The final color scheme is fixed after the robot lifts the cube to reduce the configuration count from 2304.Source: T1: Match Color (Sim) · p. 14
  • Long-hist DP completes T1 with 100% success, while No-hist DP gets stuck pausing and Mid-hist DP returns the cube to a random bin.Source: T1: Match Color (Sim) · p. 15

T2: Discrete Place Back (Sim)

  • T2 uses four configurations instead of T1’s color-varied bins, with data collection, training, and baseline failure modes otherwise identical to T1.Source: T2: Discrete Place Back (Sim) · p. 15

T3’: Continuous Place Back (Real)

  • Long-hist DP’s proprioceptive histories go out of distribution, yielding 0% success in continuous place back, while GMP achieves higher performance and faster execution than the baselines.The baselines also frequently trigger the UR5 emergency stop when pressing against the table.Source: T3’: Continuous Place Back (Real) · p. 15

T3: In-the-wild Flip and Place Back (Real)

The section reports real-world in-the-wild flipping and simulated iterative pushing, both involving behavior that depends on earlier actions or observations. GMP uses long context for pushing, where longer-history baselines struggle to maintain successful behavior across trials.

T3: In-the-wild Flip and Place Back (Real)

  • The in-the-wild dataset includes 1,274 cup place-back trajectories, 588 flipping trajectories, and 304 combined trajectories collected across 12 cups and more than 30 scenes.Source: T3: In-the-wild Flip and Place Back… · p. 15
  • Failures in challenging unseen environments include cups falling into a bench gap or being dropped early when surface height is misestimated; Long-hist DP more readily overfits after flipping.In other environments, the policy completes the task with high success.Source: T3: In-the-wild Flip and Place Back… · p. 16

T4: Iterative Pushing (Sim)

  • Long-hist DP and Long-hist PTP can reach the optimal pushing velocity but struggle to maintain it across trials, while shorter-history baselines lack enough memory to cover the last trial.GMP’s 240-timestep history includes two complete trials, and it moves closer to the optimal velocity by referring to the last trial.Source: T4: Iterative Pushing (Sim) · pp. 16–17

T5: Iterative Flinging (Sim)

These appendix sections detail iterative flinging and casting, evaluation protocols, and a data-related limitation in MIKASA-Robo. GMP learns iterative refinement in simulated flinging, while the casting baseline accounts describe failure to consistently choose or maintain the needed velocity.

T5: Iterative Flinging (Sim)

  • Despite imperfect cloth simulation and a relaxed success criterion, GMP learns to refine actions from past experience and outperforms other long-context baselines on iterative flinging.The criterion requires at least 10 of 13 cloth-edge grid points to touch the target area.Source: T5: Iterative Flinging (Sim) · p. 17
  • No-hist DP, Mid-hist DP, and Mid-hist PTP achieve 54% success on medium-mass cloths but fail on extremely light or heavy cloths.Source: T5: Iterative Flinging (Sim) · p. 17

T6: Iterative Casting (Real)

  • Across 80 training episodes, data collection randomly chooses between Teflon fiber (low friction; optimal casting velocity 1.4m/s) and electrical tape (high friction; optimal casting velocity 2.2m/s).Source: T6: Iterative Casting (Real) · p. 17
  • No-hist DP casts at medium velocity (1.8m/s), failing in 95% of episodes; Long-hist DP fails to include past proprioception under real-world noise (0% success) and has a 20% success rate.Source: T6: Iterative Casting (Real) · pp. 17–18

8.4 RoboMimic Evaluation Details

  • RoboMimic evaluation covers Square, Transport, and Tool Hang, reporting the best success rate across checkpoints evaluated on 100 episodes each.Source: §8.4 · p. 18

8.5 MIKASA-Robo Evaluation Details

  • MIKASA-Robo uses 250 training episodes per task, which the authors say is insufficient for remember-color tasks, especially RememberColor9; more episodes can improve performance.The PPO oracle actions are noisy, and the authors note that rule-based or teleoperated data can have lower action noise.Source: §8.5 · p. 18

8.6 Baseline Details

This section specifies the compared history configurations and illustrates memory-gate calibration. The calibration labels memory on when memory-conditioned prediction errors are lower by the chosen threshold, and off when errors are similar, including on RoboMimic tasks.

8.6 Baseline Details

  • No-hist DP uses only the latest image and proprioception, while Mid-hist DP extends both inputs to 16 steps and Long-hist DP to 120 steps with every eighth image sampled.Source: §8.6 · pp. 18–19
  • Mid-hist PTP predicts 16 past and 16 future actions; Long-hist PTP predicts 120 past and 16 future actions, while BC-RNN retains unlimited history without resetting its state.Source: §8.6 · pp. 18–19
  • GMP matches Long-hist DP’s history configuration for cross-attention, and extends history length further for iterative tasks T4–T6.Source: §8.6 · p. 19

8.7 Visualization of the Memory Gate Calibration

  • The gate label is on when the no-memory prediction error exceeds the memory-policy error by a threshold of 10.0; similar errors on RoboMimic indicate memory can remain off.The gate can activate before a critical moment because the error is computed over the predicted action trajectory.Source: §8.7 · pp. 19–20

8.8 Why Not Use Continuous Gating

Jointly trained continuous gates face a regularization trade-off: too little harms Markovian tasks, while too much hurts non-Markovian tasks. The gate is not necessary for every task, so the authors recommend evaluating the ungated policy first.

8.8 Why Not Use Continuous Gating

  • Without regularization, the continuous gate tends to stay on and degrades Markovian performance; high regularization hurts non-Markovian tasks, making one weight difficult to choose across tasks.The figure describes the same pattern: no regularization encourages the gate to stay on, while high regularization encourages it to be off.Source: §8.8 · p. 20, Figure 19 · p. 20
  • A calibrated continuous-gated policy works similarly to the binary-gated policy, but cannot skip unnecessary memory recall to speed inference.The continuous values limit the ability to skip recall when it is unnecessary.Source: §8.8 · p. 20

8.9 Training Cost and Calibration Overhead

  • On Iterative Pushing, training to the best success rate required 1.2h for π, 6.7h for πmem, 6.4h for Rollout, 1.3h for Gate, and 31.2h for GMP on H100 GPUs.The passage notes that calibration can use an early checkpoint and a subset of the validation set.Source: §8.9 · p. 20
  • The authors recommend first evaluating the ungated policy, then training a gate when there are clear signs of overfitting, as in RoboMimic Tool Hang.Ungated policies already perform decently on tasks that do not require high-precision actions, including MIKASA-Robo and Continuous Place Back.Source: §8.9 · p. 20
Loading 2604.18933v2…