Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Memory Retention Is Not Enough to Master Memory Tasks in Reinforcement Learning

Oleg Shchendrigin, Egor Cherepanov, Alexey K. Kovalev, Aleksandr I. Panov

TL;DR

This paper targets memory rewriting in reinforcement learning under partial observability, introducing two diagnostic benchmarks—Endless T-Maze and Color-Cubes—that require continual memory updating. It systematically compares recurrent, transformer-based, and structured memory architectures, revealing a consistent hierarchy where explicit adaptive forgetting (as in LSTM-based policies) yields superior rewriting and generalization, while attention-based and cache-like memories struggle when rewriting is essential. The work provides detailed training/evaluation protocols, reports nuanced results across retention, rewriting, and generalization tasks, and argues for designing memory systems as active, adaptable belief states rather than static histories. Together, these contributions push toward memory mechanisms that balance stable retention with dynamic updating, with implications for robust, memory-aware RL in non-stationary environments.

Abstract

Effective decision-making in the real world depends on memory that is both stable and adaptive: environments change over time, and agents must retain relevant information over long horizons while also updating or overwriting outdated content when circumstances shift. Existing Reinforcement Learning (RL) benchmarks and memory-augmented agents focus primarily on retention, leaving the equally critical ability of memory rewriting largely unexplored. To address this gap, we introduce a benchmark that explicitly tests continual memory updating under partial observability, i.e. the natural setting where an agent must rely on memory rather than current observations, and use it to compare recurrent, transformer-based, and structured memory architectures. Our experiments reveal that classic recurrent models, despite their simplicity, demonstrate greater flexibility and robustness in memory rewriting tasks than modern structured memories, which succeed only under narrow conditions, and transformer-based agents, which often fail beyond trivial retention cases. These findings expose a fundamental limitation of current approaches and emphasize the necessity of memory mechanisms that balance stable retention with adaptive updating. Our work highlights this overlooked challenge, introduces benchmarks to evaluate it, and offers insights for designing future RL agents with explicit and trainable forgetting mechanisms. Code: https://quartz-admirer.github.io/Memory-Rewriting/

Memory Retention Is Not Enough to Master Memory Tasks in Reinforcement Learning

TL;DR

This paper targets memory rewriting in reinforcement learning under partial observability, introducing two diagnostic benchmarks—Endless T-Maze and Color-Cubes—that require continual memory updating. It systematically compares recurrent, transformer-based, and structured memory architectures, revealing a consistent hierarchy where explicit adaptive forgetting (as in LSTM-based policies) yields superior rewriting and generalization, while attention-based and cache-like memories struggle when rewriting is essential. The work provides detailed training/evaluation protocols, reports nuanced results across retention, rewriting, and generalization tasks, and argues for designing memory systems as active, adaptable belief states rather than static histories. Together, these contributions push toward memory mechanisms that balance stable retention with dynamic updating, with implications for robust, memory-aware RL in non-stationary environments.

Abstract

Effective decision-making in the real world depends on memory that is both stable and adaptive: environments change over time, and agents must retain relevant information over long horizons while also updating or overwriting outdated content when circumstances shift. Existing Reinforcement Learning (RL) benchmarks and memory-augmented agents focus primarily on retention, leaving the equally critical ability of memory rewriting largely unexplored. To address this gap, we introduce a benchmark that explicitly tests continual memory updating under partial observability, i.e. the natural setting where an agent must rely on memory rather than current observations, and use it to compare recurrent, transformer-based, and structured memory architectures. Our experiments reveal that classic recurrent models, despite their simplicity, demonstrate greater flexibility and robustness in memory rewriting tasks than modern structured memories, which succeed only under narrow conditions, and transformer-based agents, which often fail beyond trivial retention cases. These findings expose a fundamental limitation of current approaches and emphasize the necessity of memory mechanisms that balance stable retention with adaptive updating. Our work highlights this overlooked challenge, introduces benchmarks to evaluate it, and offers insights for designing future RL agents with explicit and trainable forgetting mechanisms. Code: https://quartz-admirer.github.io/Memory-Rewriting/
Paper Structure (32 sections, 1 equation, 6 figures, 7 tables)

This paper contains 32 sections, 1 equation, 6 figures, 7 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Markov Decision Processes.
  4. Partially Observable Markov Decision Processes.
  5. Dynamic Memory and Rewriting.
  6. Memory-intensive environments and dynamic correlations.
  7. Related Work
  8. Memory-augmented RL agents.
  9. Existing benchmarks for memory testing
  10. Benchmarking Memory Rewrite
  11. Endless T-Maze.
  12. Color-Cubes.
  13. Experiments
  14. Baselines
  15. Training and Evaluation Protocol
  16. ...and 17 more sections

Figures (6)

  • Figure 1: Illustration of memory rewriting vs. retention. The agent first observes a cue, stores it in memory, and acts according to that cue (jump for $\uparrow$). Later, a new cue appears indicating a change in required behavior ($\downarrow$). In the rewriting case (top), the agent overwrites its previous memory with the new cue and successfully adapts its action. In the retention case (bottom), the agent fails to update its memory, continues to act based on the old cue, and consequently performs the wrong action, leading to failure.
  • Figure 2: T-Maze vs. Endless T-Maze. In the classic T-Maze (left), the agent receives a cue at the start and must remember it until the junction. Our proposed Endless T-Maze (right) extends this setup by chaining junctions: at each corridor, a fresh cue overrides all previous ones, and continual rewriting rather than simple retention is necessary to success.
  • Figure 3: Visualization of key principles of the Color-Cubes environment, such as Initialization, Stochastic Teleportation, Successful Interaction.
  • Figure 4: Intermediate progress of SHM and GTrXL agents in the Endless T-Maze. The figure show the number of successfully passed corridors over time for different training configurations of corridors number $n$ and length $l$. Despite low overall success rates, both agents can partially progress through several corridors before failing, indicating limited but nonzero capacity for short-term memory rewriting. Dashed red lines denote the target number of corridors for each validation configuration to reach success.
  • Figure 5: Baseline comparison on Endless T-Maze under interpolation and extrapolation conditions, that is, where corridor lengths and fixed sampling are the same during training and validation, and in each validation task the parameter of the number of corridors varies in accordance with the experiment protocol. The result for Success Rate is mean$\pm$sem.
  • ...and 1 more figures