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

1 INTRODUCTION

• Short observation windows leave recurrent and transformer policies brittle to long-horizon dependencies and sparse signals, motivating hybrid memory-augmented transformers.The introduction frames the central question as equipping imitation-learning policies with efficient long-term memory for partially observable tasks.Source: §1 · p. 1
• ELMUR augments every transformer layer with persistent memory embeddings, bidirectional token–memory cross-attention, and LRU updates by replacement or convex blending.The design extends retention of task-relevant information up to 100,000× beyond the native attention window.Source: §1 · p. 2
• ELMUR ranks first on 21 of 23 MIKASA-Robo tasks and improves aggregate suite success rate by roughly 70% relative to the strongest prior method.It also achieves 100% success on T-Maze corridors up to one million steps and the top score on 24 of 48 POPGym tasks.Source: §1 · p. 2
• The paper reports theoretical bounds on forgetting, retention horizons, and stability of memory embeddings.This theoretical analysis is presented alongside empirical evaluation across synthetic, robotic, and puzzle/control tasks.Source: §1 · p. 2

2 BACKGROUND

• In a partially observable task, the agent cannot access the true state and must rely on observation-action history or an approximate learned memory state.Under full observability, the optimal policy can depend only on the current state.Source: §2 · pp. 2–3

3 METHOD

• Each ELMUR layer has coupled token and memory tracks: tokens process observations into actions, while memory persists across segments and the tracks read from and write to one another.The memory is layer-local rather than a cache of token activations.Source: §3 · pp. 3–4
• Segment-level recurrence splits trajectories into shorter segments and carries layer-local memory forward for reading and updating across segments.This replaces cached activations with external memory carried between segments.Source: §3 · p. 4
• Within a segment, tokens model local dependencies with causal self-attention and retrieve long-term context by querying memory through mem2tok cross-attention.The retrieved information is incorporated into token representations before feed-forward refinement.Source: §3 · p. 4
• Tokens update persistent memory through tok2mem cross-attention, with relative bias reversed to favor temporally aligned memory embeddings.The memory track enables writing new information rather than relying only on reading from memory.Source: §3 · p. 4
• Relative bias uses token positions and memory anchors, with opposite offset directions for reading and writing.The read path favors temporally close memories while keeping distant ones accessible; the write path guides updates toward embeddings aligned with writing tokens.Source: §3 · p. 5
• LRU management fills empty slots by full replacement, then refreshes the least recently used slot through convex blending controlled by λ.This policy uses bounded memory capacity while retaining context; λ controls the balance between overwriting and retention.Source: Figure 2 · p. 3, §3 · p. 5

4 THEORETICAL ANALYSIS

• The memory module uses a convex update with blending factor λ; empty slots receive full replacement, while other slots are updated and unselected slots remain unchanged.This is the update rule analyzed for retention and forgetting.Source: §4 · p. 6
• After k overwrites, an embedding retains (1−λ)^k of its initial content, while a write from τ updates earlier contributes λ(1−λ)^(τ−1).The analysis presents this as exponential forgetting under the convex update rule.Source: §4 · p. 6
• The effective retention horizon scales linearly with memory slots M and segment length L; smaller λ extends retention, while larger λ accelerates overwriting.The horizon is defined by when a stored contribution decays below a negligible threshold.Source: §4 · p. 6
• LRU leaves unselected embeddings unchanged until replacement, then their contributions decay exponentially at rate λ, giving a conservative retention-horizon lower bound.The passage notes that practical horizons are often longer.Source: §4 · p. 6
• With norm-bounded writes and initial memory, every memory embedding remains bounded in norm by C throughout training and inference.The stated reason is that each update is a convex combination of the previous embedding and a bounded new value.Source: §4 · p. 6

5 EXPERIMENTS

• The evaluation spans synthetic tasks (Ni et al., 2023), 48 POPGym puzzle/control tasks (Morad et al., 2023a), and robotic manipulation (Cherepanov et al., 2026a).The tasks are described as designed to test memory under partial observability.Source: §5 · p. 7
• The study frames its evaluation around memory under partial observability across synthetic, puzzle/control, and robotic manipulation tasks.The named task families include synthetic tasks, POPGym, and robotic manipulation.Source: §5 · p. 7
• RQ1 asks whether ELMUR retains information across horizons far beyond its attention window.Source: §5 · p. 7

5.1 BENCHMARKS AND BASELINES

• Baselines span transformer, state-space, and offline or generative approaches; online RL and real-robot experiments are omitted to avoid incomparable training budgets and experimental confounds.The comparison includes DT, RATE, DMamba, BC, CQL, and Diffusion Policy.Source: §5.1 · p. 7
• The evaluation covers T-Maze, robotic MIKASA-Robo, and 48 partially observable POPGym tasks, using them to assess retention, generalization, manipulation, cross-domain robustness, and ablations.The reported benchmark suite uses T-Maze cue recall, visual tabletop tasks, and diverse puzzles and control environments.Source: §5.1 · p. 7, §5.2 · p. 8

5.2 RESULTS

• 100% success on T-Maze corridors up to one million steps demonstrates retention nearly 100,000× beyond the attention window.The reported setup uses context length L=10 and S=3 segments.Source: §5.2 · p. 8, Figure 3 · p. 8
• 100% success across all reported T-Maze training and validation length pairs shows generalization from 3–300-step training contexts to validation lengths of 9–9600 steps.The training lengths were split into three equal segments.Source: §5.2 · p. 8, Figure 4 · p. 8
• On MIKASA-Robo, ELMUR achieves higher success rates than the baselines on the reported manipulation tasks.Table 1 reports mean ± standard error over three runs and 100 evaluation seeds.Source: §5.2 · p. 8, Table 1 · p. 8
• On POPGym, ELMUR scores 10.4 overall, ranks first on 24 of 48 tasks, and leads most on memory puzzles, scoring 1.2 versus 0.45 for RATE.It remains competitive on reactive tasks and runs faster per step than RATE and DT on T-Maze despite having more parameters.Source: §5.2 · p. 9, Figure 5 · p. 9, Table 2 · p. 9
• Memory capacity and LRU updates dominate the RememberColor3-v0 ablations: performance is near-perfect when M ≥ N, while under-capacity and shared memory degrade results.Intermediate blending factors are unstable, larger initialization mitigates collapse, and finer recurrence stresses capacity unless M scales accordingly.Source: §5.2 · p. 9, Figure 6 · p. 10
• All tested models reach 500 ± 0 return on CartPole-v1, indicating that adding memory does not harm performance on this fully observable control task.The passage also points to additional MDP-task results in the appendix.Source: §5.2 · p. 9

6 RELATED WORK

• Prior approaches include implicit recurrence, learned external memory, context caches, episodic buffers, and transformer adaptations such as RATE, Memformer, and Block-Recurrent Transformers.The passage distinguishes ELMUR by its per-layer memory, dedicated mem2tok/tok2mem cross-attention, and LRU updates.Source: §6 · p. 10

7 CONCLUSION

• ELMUR gives every layer external memory with bidirectional cross-attention and LRU updates; the conclusion reports formal guarantees on half-life and boundedness under convex blending.It reports experiments on T-Maze, 48 POPGym tasks, and MIKASA-Robo, and presents the framework as extensible for long-horizon sequential control.Source: §7 · p. 10

REPRODUCIBILITY STATEMENT

• The paper reports architecture details and proofs, experimental protocols and specifications, baseline implementations, and supplementary code and configuration files for replication.The stated materials cover both theoretical analysis and empirical findings.Source: REPRODUCIBILITY STATEMENT · pp. 10–11

A.1 PROOF OF EXPONENTIAL FORGETTING IN LRU UPDATES

• After k overwrites, the initial slot content has coefficient (1 − λ)^k, while a write from τ updates ago has coefficient λ(1 − λ)^(τ−1).The proposition derives these coefficients for a slot updated across successive segments.Source: §A.1 · p. 18
• The coefficients in the update expansion form a convex combination, so they represent fractions of contribution.Source: §A.1 · p. 18
• Smaller λ yields longer retention horizons, while larger λ overwrites past content more aggressively.Source: §A.1 · p. 18

A.2 BOUNDEDNESS OF MEMORY EMBEDDINGS

• The proof fixes a memory slot and considers its update rule.Source: §A.2 · p. 19
• The claim is proved by induction on i.Source: §A.2 · p. 19
• The induction step uses the triangle inequality together with the inductive hypothesis.Source: §A.2 · p. 19

A.3 MEMORY-INTENSIVE ENVIRONMENTS

• A POMDP is memory-intensive when every relevant decision depends on information separated by at least one intervening step, making reactive policies insufficient.The paper adopts an event–recall horizon framework, while noting alternative memory-demand definitions are broadly consistent.Source: §A.3 · p. 19

A.4 MEMORY-INTENSIVE ENVIRONMENTS

• T-Maze requires recalling a randomly revealed goal cue after a corridor, with sparse reward only for reaching the correct goal.Source: §A.4 · p. 19
• MIKASA-Robo evaluates visual tabletop manipulation with delayed hidden-cube color recall or returning a cube after its goal changes.The simulated arm has 7 degrees of freedom, and rewards are binary and given only on task success.Source: §A.4 · p. 20
• POPGym contains 48 partially observable environments spanning diagnostic memory puzzles and observation-limited control tasks.The suite tests symbolic memory skills and generalization in continuous control.Source: §A.4 · p. 20

A.5 TRAJECTORY CONDITIONING

• 76.4 average normalized score on D4RL MuJoCo tasks under (R, o, a) conditioning was only about 2% higher than with o-only conditioning.ELMUR matches or exceeds several established offline RL baselines on these MDP tasks.Source: §A.5 · p. 20

A.6 EXTENDED RELATED WORK

• Related work spans perception, sequence and planning transformers, VLA models, implicit and external memory, and RL and manipulation systems with distinct persistence limits.The cited discussion contrasts finite-context approaches with explicit memory stores and task-specific memory extensions.Source: §A.6 · pp. 21–22

A.7 3D VISUAL NAVIGATION

• On ViZDoom-Two-Colors, ELMUR matches the previous best baseline within statistical uncertainty on a visual navigation task requiring recall of an initial color cue.The cue is shown for 45 steps before the agent navigates among colored objects.Source: §A.7 · p. 23, Table 6 · p. 23

A.8 TRAINING PROCESS

• Training computes loss on each recursively processed segment and detaches memory between segments; the loss is cross-entropy for discrete-action tasks and MSE for continuous-action tasks.The appendix reports 1,000 expert demonstrations per MIKASA-Robo environment, 6,000 successful oracle-level T-Maze trajectories, and 3,000 PPO-GRU-expert trajectories per POPGym environment.Source: §A.8 · p. 23, Table 7 · p. 24

• The probes test whether memory embeddings encode the target color, the minimal sufficient statistic for correct recall-stage decisions.Source: §A.9 · p. 24
• The study trains independent post-hoc probes on memory vectors by layer and timestep to predict the target cube color, without allowing probes to influence the policy.The dataset uses memory embeddings and ground-truth colors from 1,000 successful evaluation episodes.Source: §A.9 · p. 24
• Across layers, the update schedule writes to a slot and then leaves it inactive, preserving its contents for the rest of the episode.The analysis interprets this as a sparse, well-timed LRU update rather than continuous rewriting.Source: Figure 10 · p. 26, §A.9 · p. 26
• After cue disappearance at t=5, first-layer accuracy reaches 100%; at t=10, the first three layers reach 100% and the last is near perfect, with essentially perfect first-three-layer accuracy persisting at t=20.Source: Figure 9 · p. 25, §A.9 · p. 26
• After the cue is written, one slot’s variance rises while unused slots remain nearly constant, and the active slot’s representation remains stable through the delay.The similarity trajectories likewise remain nearly constant after the write event.Source: §A.9 · pp. 26–27, Figure 11 · p. 27

A.10 PCA MEMORY ANALYSIS

• Across target colors, slots share a content space, while deeper layers form more structured, compact color encodings; after the cue, memory vectors converge to stable color-specific regions.The PCA views color embeddings by slot, layer, and timestep.Source: §A.10 · p. 28, Figure 13 · p. 29, Figure 14 · p. 29

A.11 CROSS-ATTENTION MAPS

• Tok2mem attention concentrates in a short cue-time burst on one slot, while mem2tok repeatedly reads that slot during the blank interval and decision phase.The attention maps display 16 heads across three slots and timesteps.Source: §A.11 · p. 28, Figure 15 · p. 30

A.12 LIMITATIONS

• Evaluation covers synthetic, POPGym, and simulated robotic tasks under imitation learning; online RL and real-robot deployment remain future directions.The fixed-blending LRU rule is another stated avenue for future adaptive variants.Source: §A.12 · p. 29
• ELMUR ranks first on 21 of 23 MIKASA-Robo tasks with nonzero success and exceeds RATE by roughly 70% in aggregate success across all 32 tasks.The comparison uses sparse, success-only rewards and RGB observations.Source: Table 8 · p. 31