ELMUR: External Layer Memory with Update/Rewrite for Long-Horizon RL

TL;DR

ELMUR tackles long-horizon, partially observable decision-making by extending transformers with layer-local external memory and a principled memory-update mechanism. Each layer includes a memory track that interacts with the token track via mem2tok and tok2mem cross-attention, and memory content is refreshed through an LRU policy with convex blending, enabling recurrence across segments and extending retention horizons by up to $100{,}000\times$ the native attention window. Theoretical results establish exponential forgetting bounds and memory boundedness, linking the blending factor $\lambda$, memory slots $M$, and segment length $L$ to retention horizons and stability. Empirically, ELMUR achieves perfect retention on the synthetic T-Maze with corridors up to one million steps, substantially improves robotic memory tasks, and attains top performance on a large POPGym suite, demonstrating robust generalization under partial observability. Together, these findings position structured, layer-local external memory as a simple yet scalable solution for reliable long-horizon decision making in sequential control and imitation-learning contexts, with practical implications for real-world robotic systems.

Abstract

Real-world robotic agents must act under partial observability and long horizons, where key cues may appear long before they affect decision making. However, most modern approaches rely solely on instantaneous information, without incorporating insights from the past. Standard recurrent or transformer models struggle with retaining and leveraging long-term dependencies: context windows truncate history, while naive memory extensions fail under scale and sparsity. We propose ELMUR (External Layer Memory with Update/Rewrite), a transformer architecture with structured external memory. Each layer maintains memory embeddings, interacts with them via bidirectional cross-attention, and updates them through an Least Recently Used (LRU) memory module using replacement or convex blending. ELMUR extends effective horizons up to 100,000 times beyond the attention window and achieves a 100% success rate on a synthetic T-Maze task with corridors up to one million steps. In POPGym, it outperforms baselines on more than half of the tasks. On MIKASA-Robo sparse-reward manipulation tasks with visual observations, it nearly doubles the performance of strong baselines. These results demonstrate that structured, layer-local external memory offers a simple and scalable approach to decision making under partial observability.

Figures (8)

• Figure 1: ELMUR overview. Each transformer layer is augmented with an external memory track that runs in parallel with the token track. Tokens attend to memory through a mem2tok block, while memory embeddings are updated from tokens through a tok2mem block. LRU block selectively rewrites memory via replacement or convex blending, ensuring bounded yet persistent storage. This design enables token-memory interaction and long-horizon recall beyond the attention window.
• Figure 2: LRU-based memory management in ELMUR. Each layer maintains $M$ memory slots, initialized with random vectors (green). As new segments arrive, tokens write updates into empty slots (purple) by full replacement. Once all slots are filled, the least recently used slot is refreshed via a convex update with parameter $\lambda$ that blends new content with the previous memory (grey). Anchors below each row indicate the timestep of the most recent update. This scheme ensures bounded capacity while preserving long-horizon information.
• Figure 3: Success rate on the T-Maze task as a function of inference corridor length. ELMUR achieves a 100% success rate up to corridor lengths of one million steps. In this figure, the context length is $L=10$ with $S=3$ segments; thus ELMUR carries information across horizons 100,000 times longer than its context window.
• Figure 4: Generalization of ELMUR across T-Maze lengths. Each cell shows success rate (mean $\pm$ standard error) for training vs. validation lengths. ELMUR transfers perfectly: models trained on shorter sequences retain 100% success up to $9600$ steps. Training lengths were split into three equal segments.
• Figure 5: ELMUR compared to DT on all 48 POPGym tasks. Each model was trained with three independent runs, validated over 100 episodes each. Bars show the mean performance with 95% confidence intervals computed over these three means. ELMUR achieves consistent improvements over DT, with the largest gains on memory-intensive puzzles.