Recurrent Action Transformer with Memory
Authors: Egor Cherepanov
Aleksei Staroverov
Alexey Kovalev
Aleksandr Panov
Venue: ICLR 2026 Poster (2026)
Keywords: RL
Offline RL
Memory
Transformers
POMDP
Primary area: reinforcement learning
TL;DR: The paper proposes Recurrent Action Transformer with Memory - a transformer model with recurrent memory and a procedure for training it for memory-intensive environments in an Offline RL setting.

Abstract: Transformers have become increasingly popular in offline reinforcement learning (RL) due to their ability to treat agent trajectories as sequences, reframing policy learning as a sequence modeling task. However, in partially observable environments (POMDPs), effective decision-making depends on retaining information about past events - something that standard transformers struggle with due to the quadratic complexity of self-attention, which limits their context length. One solution to this problem is to extend transformers with memory mechanisms. We propose the Recurrent Action Transformer with Memory (RATE), a novel transformer-based architecture for offline RL that incorporates a recurrent memory mechanism designed to regulate information retention. We evaluate RATE across a diverse set of environments: memory-intensive tasks (ViZDoom-Two-Colors, T-Maze, Memory Maze, Minigrid-Memory, and POPGym), as well as standard Atari and MuJoCo benchmarks. Our comprehensive experiments demonstrate that RATE significantly improves performance in memory-dependent settings while remaining competitive on standard tasks across a broad range of baselines. These findings underscore the pivotal role of integrated memory mechanisms in offline RL and establish RATE as a unified, high-capacity architecture for effective decision-making over extended horizons.

Published as a conference paper at ICLR 2026
RECURRENT ACTION TRANSFORMER WITH MEMORY
Egor Cherepanoy 1.2 Aleksei Staroverov 1., Alexey K. Kovaley! AAleksandr I. Panov! 1.2 !AXXX, MMIRAI
cherepanov@axxx. tech
ABSTRACT
Transformers have become increasingly popular in offline reinforcement learn- ing (RL) due to their ability to treat agent trajectories as sequences, reframing policy learning as sequence modeling task. However, in partially observable environments (POMDPs), effective decision-making depends on retaining in- formation about past events something that standard transformers struggle with due to the quadratic complexity of self-attention, which limits their con- text length. One solution to this problem is to extend transformers with memory mechanisms. We propose the Recurrent Action Transformer with Memory (RATE), a novel transformer-based architecture for offline RI. that incorporates recurrent memory mechanism designed to regulate information retention. We evaluate RATE across diverse set of environments: memory-intensive tasks WizDoom-Two-Collee T-Maze, Memory Maze, Minigrid-Memory, and POP- Gym), as well as standard Atari and MuJoCo benchmarks. Our comprehen- sive experiments demonstrate that RATE significantly improves performance in memory demoor-deeendent settings while remaining competitive on standard tasks across a broad range of baselines. These findings underscore the pivotal role of inte- grated memory mechanisms in offline RI. and establish RATE as unified, high- capacity architecture for effective decision makiii over extended horizons. Code: https://sites: google. wwweggl....mm///ee wwwwwwwwwwwwwwwwwww....wwwwwewwe model/.
INTRODUCTION
Originally developed for Natural Language Processing (NLP), transformers (Vaswani et al., 2017) have recently demonstrated strong performance across a wide range of Reinforcement Learning (RL) settings (Agarwal et al., 2023; Li et al., 2023). They have been successfully applied to online (Parisotto al., 2020; Esslinger al., 2022), offline (Chen al., 2021; Jiang et al., 2023; Wu et al., 2023; Zhuang et al., 2024; Wang et al., 2025), model-based (Chen et al., 2022; Robine al., 2023), and in-context RL (Polubarov et al., 2025; Grigsby et al., 2024; Schmied al., 2024). In particular, transformers show promise for tackling long- Long-horizon credit assignment and operating in memory- intensive environments (Ni et al., 2023; Grigsby et al., 2024; Esslinger et al., 2022; Parisotto et al., 2020), provided the full trajectory fits within the model context. Despite their success, transformers face fundamental limitations when applied to long sequences due to the quadratic complexity of self-attention (Keles al., 2023), which restricts their applicability in fong-horizon inference tasks. While various techniques have been proposed to extend the context window (Dai et al., 2019; Bulatov al., 2022), these approaches often suffer from training instability (Zhang al., 2022) or rely on task-specific sparse attention patterns that may not generalize well beyond NLP (Beltagy et al., 2020; Zaheer et al., 2020). Semory-auggennt
transformers offer a promising alternative by enabling access to past information without expanding the context length. Motivated by advances in memory mechanisms for NLP models (Dai et al., 2019; Bulatov et al., 2022), we investigate how such approaches can be adapted to RL. Unlike NLP, RL involves structured and modality-riee -rich inputs observations, actions, and rewards - that require domain-specific encoding, and frequently exhibit high sparsity in both reward signal and observations.
In RL, memory usually refers either to using past information within an episode (Lampinen et al., 2021; Ni et al., 2023), to transferring experience across environments (Kang et al., 2023; Team et al., 2023), aiding generalization, sample efficiency, and Meta-RL (Duan et al., 2016; Wang et al., 2016), and we focus on the former.

---

Published as a conference paper at ICLR 2026
RATE training
Ceuael Transloreer
MYV
Coural Tranafonrer
Gauael Thanaforme
WWV
Segrart
Segment
Segmant 9
RATE model
Mo+1
Causal Transformer
MRV
Segment S.
Daak) ... (m 1y
Ols 1:)xK; Pis 1)xx)
Ma+1
Returns to Enooder
Observations Encoder
Actions Encoder
Memory Retention
R
Multi-Heae Hell--HHdd Attention
Ma+1
Vodel
Moe1
-1.41
++1.92
Figure 1: Recurrent Action Transformer with Memory (RATE). The model processes trajectory divided into n segments Sn with memory embeddings Mn. where I denotes returns-to-go (future rewards), are observations, a are actions, and M,, are memory embeddings attached to each segment Sn to retain important historical information.
We introduce the Recurrent Action Transformer with Memory (RATE; Figure 1), memory- augmented transformer that incorporates three complementary mechanisms: learned memory em- beddings, recurrent caching of past hidden states, and a novel Memory Retention Valve (MRV) for selective information flow. We empirically show that memory mechanisms effectively preserve information from previous steps, allowing the model to use past information when making deci- sions in the present. MRV is designed to control the process of updating memory embeddings and prevent the loss of important information when processing long sequences, thus enabling processing of highly sparse tasks. To assess the effectiveness of our memory mechanisms, we conduct extensive experiments across a diverse set memory-intensive environments, including (Sorokin et al., 2022), Memory Maze (Pasukonis et al., 2022), Minigrid- Memory (Chevalier-F Bhhisvlier-Boisvertt et al., 2023), Passive T-Maze (Ni et al., 2023), and POPGym (Morad al., 2023a), well as standard RL benchmarks such Atari (Bellemare et al., 2013) and Mu- JoCo (Fu et al., 2021). We also study the impact of memory on the performance of the proposed model. RATE interpolates and extrapolates well outside the transformer context and is able to retain important information for a long time when operating in highly sparse environments.
Our main contributions are follows:
1. We propose Recurrent Action Transformer with Memory (RATE), transformer for offline RL that combines three complementary memory mechanisms: (i) memory embeddings, (ii) caching of hidden states, and (iii) Memory Retention Valve (MRV), which uses cross-attention to retain key information over long horizons (Section 3).
2. We conduct extensive evaluations on memory-intensive tasks including ViZDoom Two- Colors, Memory Maze, RATE consistently outperforms strong baselines (Section 4.1).
POPGym, and Passive T-Maze showing that
3. We further show that RATE matches or surpasses standard baselines on the Atari and MuloCo benchmarks, demonstrating strong generalization across task types and highlighting the model's versatility (Section 4...
2 BACKGROUND
Offline RL. In RL (Sutton & Barto, 2018), a task is formalized as Markov Decision Process (MDP):
AS P, R). where E S are states, E AA are actions, P(s'|s, a) is a transition function, and r R(ss a) is a reward function. States satisfy the Markov property: P(st+1|s:) = P(s (1,-1|11..... 8t). A trajectory 7 of length T is sequence (Sp, 0u. To...., . Si -1, 1,,, 11), where R(s,, 0z)
is the immediate reward at the timestep f. The returnto-go (Chen et al., 2021) R. ΣI=1 Ft is
2

---

Published as a conference paper at ICLR 2026
RATE
RATE
RATE
DT
Seginent
Segment 2
Segment
To S'ofat
B Pooaa of Wast
Figure 2: Attention maps of RATE and DT on the T-Maze (Ni et al., 2023) task with corridor length T = 8. DT is trained on full 8-step trajectories, while RATE processes the sequence in three segments of length recurrently, passing information between segments through memory embeddings.
the sum of future rewards from I. The goal is to learn a policy maximizing the expected return. While online RL iteratively collects trajectories through environment interaction, offline RL uses a fixed dataset of trajectories, making it suitable for scenarios where environment interaction is costly or risky. A popular offline RL, method, Decision Transformer (DT) (Chen et al., 2021), models returnconditioned trajectories with a GPT--style architecture, avoiding value estimation. However, its fixed context window limits performance in tasks with delayed rewards or long-term dependencies, motivating -uggeedd models.
POMDP. In real-world, agents often receive partial observations rather than full states, breaking the Markov property. For instance, robot using only camera input agent relying on past context. Such cases are modeled as Partially Observable MDPs (POMDPs): (S, A, O, P. R, Z), where o E O observations and Z = P(o s', a a) defines the observation function. Since single observations are insufficient, agents must use history to infer useful state representations.
3 RECURRENT ACTION TRANSFORMER WITH MEMORY
Transformers excel at sequence modeling, including offlin RL. (Chen et al., 2021; Janner al., 2021), but struggle with long-horizon tasks due to fixed context and quadratic attention cost, In memory tasks, agents must recall information seen thousands of steps ear- lier something models like DT cannot do once cues fall outside context. We propose the Recurrent Ac- tion Transformer with Memory (RATE), which in- troduces segment-level recurrence and dynamic mem- ory control. RATE processes trajectories segments, using lightweight memory and learnable Memory Retention Valve (MRV) to decide what to retain or discard. In T-Maze (Ni et al., 2023), the agent receives a one-bit cue ou at the first step indicating whether turn left or right at the end of maze. Solv- ing the task requires remembering this cue despite sparse rewards. DT fails once og leaves the context, making retrieval at inference impossible. Figure 2 shows this: DT attends to oo only when it fits the con- text, while RATE segments the input and propagates the memory embeddings, preserving the cue to the end and enabling explicit memory retention.
Algorithm RATE
Require: R E R ,0 E RdoXT 1: R R
; RT
ā + Encoder.(0)
ā Encoder,
To-T- + {(Rg, or, à
3: Mn + Mo ~ N0, 1)
for n in [0, T//K//KK - do
5: Sn + (n+1))
Sn 5 (Ma Sn, Mn)
an;Ma+ & Trans Trans fornerormerrS
8: Mn+1 t- MRV(M Mn+1)
ān - L(an, an), Mn+1
end for
Algorithm Memory Retention Valve
Require: Mn. Mn E R 1: Qn Wh.
2: K₁ + Ma+1 Wh
K
3: VA < Mn+1 W T
4: Mit + softmax
Q.KT
V
RATE combines memory embeddings (Bulatov et al., 2022), cached hidden states' (Dai et al., 2019), and a novel MRV to handle long and sparse sequences. The architecture is shown in Figure 1. Let trajectory
vd
5: Mn+1 + concatt (M
(M
b+ ......... M m+1)
6: Mat1 + MT. Output: Mn+1
In our setting, cached hidden states refers the mechanism intraduced in Transformer-XL. (Dai al., 2019), which the hidden activations computed for preceding segments are stored and reused as an extended key-value context when processing the next segment. Concretely, instead of recomputing all past representations from scratch, the model concatenates the fixed, non- trainable hidden states from earlier segments with the current segment's inputs, thereby enabling segment-level recurrence and information flow across boundaries without hackpropagating gradients through the cached states.
3

---

Published as a conference paper at ICLR 2026
Tu:T-1 of length T be represented by triplets (R,, oz, oz), where R, is the return-to-go or the obser- vation, and a, the action. Each modality is encoded using modality-speciie specific encoders (Algorithm 1):
R3 = Encoder ((c)) or = Encoder.(ot). āg = Encoder, (ar). The encoded sequence is split into N T//K non-overlapping segments S, of length K. Thus, the effective context is Ken = x K, well beyond standard attention limits. Each segment is prepended and appended with the same memory embeddings Mn. E Rmmdd where m is the number of memory tokens and d is the embedding dimension. This design follows from the use of causal self-attention in the decoder: the prefix copy of M,, provides read access, since every token in the segment Sn can attend backward to the incoming memory, while the suffix copy provides write access, since these memory tokens appear after the segment in the causal ordering and allow the final layers to attend forward into S,, to produce updated memory. Using only the prefix would make memory readable but not updatable, whereas using only the suffix would prevent the segment from accessing previously stored information. Both copies are therefore required for RATE's recurrent memory mechanism:
Sn concat (M,, Sn, Mn) E R (2(3
Each segment is then processed by the transformer,
ony M Transformer(S.))
and the resulting memory Mn+1 is refined via the MRV before being passed to the next segment.
Naively forwarding memory embeddings leads to error accumulation or overwriting of relevant information. To address this, introduce the Memory Retention Valve (MRV), a cross- crtss-attention module that filters new memory tokens through the lens of the previous ones (Algorithm 2):
MRV(M,, Mn+1) 1 = FFN CetultiHead(((0 (Query = MR, Key = Ma+1, Value = Ma+1))
This mechanism allows M,, to control what to retain or overwrite when updating M,+1. Unlike static recurrence", it preserves sparse, long-range information. RATE overcomes DT's limits by extending context with recurrence, preserving early cues via MRV, and retaining key events in sparse settings. As result, RATE solves tasks where DT fails, generalizes beyond training, and remains competitive on standard MDPs.
Attention pattern analysis. Figure compares attention maps of RATE and DT on T-Maze sequence. DT (right) attends only within a fixed window, focusing on recent tokens while losing early cues like Og- RATE (left) segments the input and uses memory tokens to propagate information across segments. These tokens retain access to on even in later segments, demonstrating RATE's ability to model long-ram rang-rnge dependencies beyond the context window through structured memory.
3.1 PRESERVATION PROPERTIES OF MRV
We formalize the intuition that the trass-attentionn-ueel MRV prevents catastrophic overwriting of memory by preserving alignment between consecutive memory states. All vectors are row-vectors. We use |[ for the Frobenius norm and || ||2 for the & norm.
Let M, E I mmd and Ma+1 Rd denote the incoming and updated memory embeddings segment n, where m is the number of memory tokens and d is the model dimension. We assume that each row of M,, nz-ommallie::: ||MM,ill 1. The MRV computes the next memory state as:
Q = M, K MM WK, V Ma+ Wv, A softmax
QKT vd
M,+1 AVWM- M-
a-alignment condition.
The memory embeddings are said to satisfy a-alignment if there exists a constant E (0, 1] such that for every row Mn.i, there exists row V,, for which: (VyWar, Mn.i) 2
"Static recurrence denotes the practice of forwarding eached hidden states from one segment to the next without any gating content-based filtering. For instance. Transformer- XL uses this mechanism by directly reusing past hidden states as extended key-value context, which increases the effective horizon but provides no control over what information is retained or overwritten
'Throughout the paper, alignment refers only to geometric notion: two vectors are "a-aligned" when the angle hetween them does not exceed specified threshold. This has no relation alignment preference tuning in large language models.
4

---

Published as a conference paper at ICLR 2026
Nit' plarattei1B.00
Wifheed plaar
1010 201
10000 2081
Iad Pl
Great Fles
FAT
10L
Text ffller
pllers
Sugert Denund
Tarpet Deward
Tarpet
Figure 3: Comparison of RATE with transformer baselines (DT, RMT, TrXL) on ViZDoom-Twe Colors trained on the first Timin 90 steps of the episode: with (a) and without (b) pillar the first 45 steps of the episode; caleulated at environment steps 0 89 (c) and 90 179 (d) with pillar in the first 45 steps; depending on the return-to-ge go (e, f. g). Episode timeout- 2100 steps.
a. This implies that the angle between Vg W M and Mn, is most arccos a. Empirically, this condition holds in trained models, as the transformer tends to preserve useful memory content and avoids orthogonal rotations between segments.
Theorem (On memory loss bounds). Let each memory row be fz-normalized, the a-alignment condition hold, and A softmax
QK" vd
be the MRV attention matrix. Then:
|M++1 -MllFF S /
/2 22(1
=) m
||M++ill 2 2 (
-IMM (4)
In words: at least a (1 (1-v2) 22-12(11- # part of the memory is guaranteed to be preserved after single MRV update (Equation 4 (right)), and the memory loss is upper bounded by Equation 4 (left).
Proof. Since each row of the attention matrix A is a probability distribution, we have E, Aij = for every i. By the pigeonhole principle, there exists an index j" such that Aij- 2 1
By assumption, for each Mn.s there exists V, such that (V, Mn,d) 2 a. In particular, this holds for j": (V3- WM., M, Mn,1) 2 a. Using the MRV definition M = E, V, W.,, M. we write:
(Ma+1,5 Mn,1)
As Mn.a) 2 Aij- (Vz- WM, Mn.i) m
(5)
Let 0g be the angle between Mn+1, and Mn Since both vectors are fz-normallzed, we have: (M
cos 0, Ma Ma Using the identity ||u v113 2(1 cos 0) for unit vectors:
||M++ Mn,il1 < (1- _) ||M Mn,illa V2(1 m)
(6)
Summing over all memory tokens and applying the previous bound: || Mn 17
Emh || -n,i.s-MM,illl Consequently, since || v due to row normalization, we conclude: ||MM+1 MnllF z
Mn,il12 1/2 S 2m (1 =) which simplifies to: ||M Malle S
V 2 (1 m)
(1 1 -
We now derive the lower bound Equation 4 (left) using the reverse triangle inequality. For any matrices Mr mmd we have: ||M,+ 2 |M || - M| M, F. Substituting
the upper bound from Equation 4 (right): [|M,+1 M. S 2 m) ||,,,r we obtain:
||M 2 (1 22 (1- )) ||Mallr, which completes the proof of Equation 4.
4 EXPERIMENTAL EVALUATION
We designed our experiments to achieve two main goals: (a) to showcase the strengths of the RATE model in memory-intensive environments (T-Maze, Memory, POPGym), and (b) to assess its effectiveness in standard MDPs, demonstrating its versatility across domains.
Memory Maze, Minigrid-
Baselines. To evaluate the performance of RATE, we compare it against diverse set of baselines spanning several categories: transformer-based models including Decision Transformer (DT) (Chen al., 2021), Recurrent Memory Transformer (RMT) (Bulatov al., 2022) and Transformer-A XL (TrXL) (Dai et al., 2019) specially adapted by us for offline RL, and Long-Short Short Decision Transformer (LSDT) (Wang et al., 2025); classic baselines such as Behavior Cloning with an MLP backbone
5

---

Published as a conference paper at ICLR 2026
Average Return
Total
(BC-MLP) and Conservative Q- Learning (Kumar et al., 2020) with an MLP backbone (CQL-MLP); current models including Behav- ior Cloning with an I.TMM back- bone (Hochreiter & Schmidhuber, 1997) (BC-LSTM), CQL with LSTM (CQL-LSTM), Decision LSTM (DL. STM) (Siebenborn al., 2022), and its GRU-based variant (Chung et al., 2014) (DGRU); and state space model baseline, Decision Mamba (DMamba) (Ota, 2024; Lv et al., 2024).
Red Rillars
Green Pillars
Imbalance
RATE Bog
0.0s
Bag
1.34
1,07
0.59
RMT
COL-LSTM
BC-LSTM
TKL
DLSTM
DGRU
LSDT
Dr
DMamba
COL-MLP
BCNLP
Random
= 150. The top plot shows average return across all episodes (yel- low), and separately for red (red) and green (green) pillars. The bottom plot shows the imbalance metric absolute differ- ence between red and green performance. Lower imbalance indicates more consistent behavior and is as important as average return.
Figure 4: ViZDoom- Teooom-Two-CCle Two-Colors results with m
Memory-intensive tasks. We evalu- ate RATE in tasks that require agents to retain information over time Figure 9; full details are in Appendix C. ViZDoom-Two-Colon
the agent must recall briefly visible pillar color to collect matching items; T-Maze: cue the start indicates the correct turn at the end, testing sparse long-term memory; Minigrid-Memory: like T-Maze. but the clue must be located first, combining memory and credit assignment (Ni et al., 2023); Memory Maze: the agent searches for objects matching a changing target color, requiring spatial memory; POPGym: suite of 48 partially observable tasks (Morad al., 2023a) designed to probe different aspects of memory.
Success Rate
1.0*
4.1 EXPERIMENTAL RESULTS
0.8
WWZDoom-Tw-Coloon Figure
shows training with 150 and in- ference up to 2100 steps, where the pillar disappears at step 90. RATE achieves the highest return and lowest imbalance between the red and green pillars, indicating strong and consis- tent memory use. Figure further tests transformer models trained with Timin = 90 on their ability to retain early
0.6
0.4
RATE (Ours)
DT
DLSTM
RMT
LSDT
BC-MLP
DGRU
DMamba
Persistent Agent
0.2
CQL-LSTM
BC-LSTM
Random Agent
TXXL
CQL-MLP
0.0
10a
102
10a
104
With the pillar present (a), RATE again yields the highest and most stable return. DT and TrXL un- derperform and show higher imbalance. Removing the pillar (b) degrades all models, confirming reliance on the initial cue. DT's unchanged performance across (a) and (b) highlights its failure to leverage long-term dependencies.
T-Maze Inference Corridor Length
Figure 5: T-Maze generalization task.
This limitation is clearer in Figure 3 (c. d), which separates performance within and beyond the 90-step context. DT's return drops by nearly 50% in red-pillar episodes once the cue leaves the window, while memory models (RATE, RMT, TrXL. remain stable, demonstrating their ability to retain and use information over long horizons.
Figure (e, f, g) shows model performance across target reward levels. RATE consistently outper- forms all baselines overall (e), and this advantage is even clearer when separating red (f) and green (g) pillar episodes. While other models show large disparities, RATE maintains stable performance across both conditions, demonstrating effective use of initial cues and validating the strength of its memory architecture.
T-Maze. Figure shows the model generalization in Passive T-Maze as inference length grows from 9 to 9600 steps. All models were trained on episodes up to 900 steps; extrapolation beyond this requires long-horizon generalization. RATE achieves 100% success across all in-distribution lengths and performs well even at 9600-step inference, corresponding to trajectories of 3 x 9600 = 28800 tokens due to the (R, o, a) triplets. This highlights RATE'S ability to retain and leverage sparse cues

---

Published as a conference paper at ICLR 2026
Trainjng Seeuuence Langth Truiniing Sequence Lanoth Trainjnng Serquence Latrgsh
RATE Success Rate
over extremely long horizons. Other transformers (e.g., DT. LSDT) match RATE training-length sequences but degrade sharply beyond. DT collapses to 50% even at moder- ate lengths due to its lack of mem- ory. Memory-angmented models like RMT generalize slightly further but deteriorate. TrXL performs simi- larly to DT, suggesting hidden-state caching alone is insufficient for long- range recall of sparse information. RNNs and SSMs (c.g., BC-LSTM,, DMamba) show flat curves and fail to learn from sparse long sequences.
1M X1 W
DT Success Rate
BC-LSTM Success Rate
T10 23 10
RATE both interpolates within train- ing and extrapolates well beyond, key strength for solving sparse POMDPs. Notably, poor performance of some memory baselines in Fig- ure 5 is due to difficulty modeling long sequences during training, not just generalization failure: even for Tral S. TTramn they may fail. However, when trained on shorter sequences, some models learn gen- eralizable behaviors. Figure 6 visualizes inference performance for RATE (top). DT (middle), and BC-LSTM (bottom) across training/validation lengths. The black dashed line separates in-distribution (Twal > Tuain). From Figure 6 (bottom), BC-LSTM generalizes well when trained on short sequences (S 150), but degrades as training lengths grow, reaching ~0.5 when trained on T 2 600, likely due vanishing gradients or limited capacity (Pascanu et al., 2013; Trinh et al., 2018). DT (Figure 6 (middle)) handles long training sequences via attention, but fails on longer validation sequences due to fixed context. In contrast, RATE (Figure 6 (top)) maintains high success across all validation lengths, enabled by its combination of attention and recurrent memory, which overcomes the limitations of both DT and RNNs.
Vulidation Sequence Langth
Figure 6: Heatmaps of success rates on T-Maze tasks. The black dashed line separates in-distribution inference (with Toal I Ttrain) from out-of-distribution inference (with Toal > Tirain). Results for other baselines can be found in Appendix, Figure 11.
(Taal < Tuain) from mon-of-distributte
Aonigrid-Memory Figure 7 presents average returns on Somigrid-Memoryy where all models were trained on grids of fixed size 41 x 41 and evaluated on wide range of unseen grid sizes from 11 x 11 to 501 x 501. RATE achieves consistently high performance across the entire spectrum, demonstrating both strong interpola- tion and extrapolation capabilities. While TrXL also performs well on average, its variance is no- tably higher, indicating sensitivity to grid scale.
Average Returns Across Grid Sizes 11x11-501x501
0.8
0.6
RATE
DT
RNT
TriL
BC.LSTM
CQL-LSTM
DORU
DMamba
LSDT
Random
Memory Maze. Table 1 presents results the Memory Maze task. RATE achieves higher average episode returns by effectively capturing implicit structure, such as maze layout. For reference, the dataset's 's average return is 4.69. All models were trained on 90-step trajectory subsequences, while full episodes span 1000 steps.
Figure 7: Minigrid-Memor generalization task.
POPGym. To further assess generalization and memory capabilities, we evaluated models on all 46 tasks from the POPGym benchmark suite, which covers a wide range of partially observable RL scenarios. The benchmark is split into 33 memory puzzle tasks and 15 reactive POMDP tasks.
Table 1: Average return + SEM in the Memory Maze (9 x 9) environment (ep. length: 1000 steps).
Method Random
BC-LSTM CQL-LSTM
DT
RMT
TrXL
RATE
Return
0.00++0.00
4.75=10.15
0.19
6.83+10.11
7.27
7.12
7.64:1041
7

---

Published as a conference paper at ICLR 2026
Table 2 reports average normalized scores across all tasks and subsets. RATE achieves the highest overall score (9.54), outperforming all baselines. On the challenging memory tasks, RATE main- tains positive average score (0.45), while all other models fall below zero indicating con- sistent failure exploit long-term dependencies. Notably, DT scores -3.49 and BC-MLP drops
Table Aggregated average returns on 48 POP. Gym tasks, split into memory and reactive subsets.
Tasks
Rand. BC-MLP DT BC-LSTM RATE
All (48)
-12.2
-6.8 5.8
9.0
9.5
Memory (33)-14.6
-11.9
-3.5
-0.2
0.5
Reactive (15) 2.3
5.1 9.3
9.1
9.1
11.91, showing the limitations of both context-limites transformers and non-recurrent policies.
On reactive tasks, all models perform better, but the gap between memory- based and non memory models narrows. RATE, DT, and BC-LSTM show almost the same results, suggesting that the greatest performance gains from RATE's memory mechanisms occur on memory puzzle tasks. For simpler reactive POMDPs, lightweight memory mechanisms appear sufficient. These results also underscore RATE'S s ability to generalize across both puzzle and reactive settings, confirming that its memory architecture does not hinder performance in simpler tasks while offering clear benefits in those with temporal dependencies. More details are provided in Appendix, Table 9.
Table Normalized scores on MuJoCo tasks from the D4RL benchmark (Fu al., 2021). Although RATE is designed for memory-intensive environments, it performs competitively and often surpasses methods tailored for standard MDP control. Top-1 and Top-2 results are highlighted.
Dataset
Euviromment
CQL
TT
DMamba (Ota, 2024)
DMamba (lxv et 2024)
MambaDM
RATE (ours)
ME
HalfCheetah
91.6
86.81:3
91.8208
95.0:02
91.9166
93.5
86.5111
87.4
Hopper
105.4
107.6:18
105.5
110.0:17
110.5:63
ME
Walker2d
108.8
108.1 ::
107.4:09
101.9 688
108.3
111.6:12
108.8141
108.7 and
HalfChectah
44.4
42
45.0-9.1
42,8+01
43.5a03
M
Hopper
58.0
67.6::14
63.4114
61.1256
83.5-14
98.5:81
85.7
77.4:14
M
Walker2a
72,5
74.02.14
64.9.22
79.0 28
78.2
80.3:
78.2164
8017 Ln.
MR
Hal[Cheetah
45.5
36 6600
40.8=06
41.9-25
39.6:18.1
40.8:204
39.1
39.0=00
MR
Hopper
95.0
82.7
87.3:23
91.5:A6
82.6-14.5
$9.1
86.1:25
83.7
Walker2d
77.2
666110
82.6249
79.
73.4:34
73.7
Averuge
77.6
74.7
74.3
78.9
75.8
83.2
79.0
78.5
Atari and MuJoCo. We evaluate RATE on standard RL benchmarks: Atari games and MuJoCo control tasks (Table 3, Table 4). For comparison, we include results from recent state-of-the-art methods: Decision Mamba (DMamba) (Ota, 2024; Lv al., 2024), Mamba as Decision Maker (MambaDM) (Cao et al., 2024), Conservative Q Q-Lerning (CQL) (Kumar et al., 2020), Trajectory Transformer (TT) (Janner al., 2021), and TAP (Jiang al., 2023), reported in their original papers. Results show that RATE matches or outperforms specialized offline RI algorithms across both benchmarks. Combined with its strong performance on memory-intensin tasks, this highlights RATE's versatility general-purpee offline RL model. See Appendix E for full training details and Table 10 for the evaluation protocol.
5 ABLATION STUDY
We conduct comprehensive ablation study to assess the contributions of individual components and architectural choices in RATE, structured around three key research questions.
1. How do different components of RATE influence performance memory tasks? (RQ1)
2. What is the upper-bound results RATE can achieve with access to perfect memory? (RQ2)
3. What role does the MRV play, and which configuration most effective? (RQ3)
Further ablations exploring key transformer parameters, memory tokens number, and sequence segmentation strategies are provided Appendix F and Appendix G.
Impact of RATE components. To assess the contribution of individual memory mechanisms in RATE, we performed efference-tiee ablations by replacing memory components with random noise. In T-Maze (K 30, N = 3 segments), corrupting memory embeddings M sharply reduced performance to 50% success (see Figure 8, right). The agent still reached the decision point but failed
RQI:
8

---

Published as a conference paper at ICLR 2026
Table 4: Raw scores on Atari games. RATE outperforms DT in out of 4 environments.
DMamba (Ota, 2024)
Environment
CQL
BC
DT
MambaDM
RATE (Ours)
Breakout 62.5
42.8
76.9
70.6
106.9
111.0
Qbert
14013.2
2862.0
2215.8
5786.0.12955.
10052.5:11165
12486.91
SeaQuest
782.2
992.1
1129.31
992.1157.77
1286.0.:2..0
1037.9
Pong
18.8
17.1+229
1.6:15.3
18.4=008
18.
to turn correctly - showing retained naviga- tion skills but lost the initial cue. Thus, mem-
at CE TrE00000
1.0
RATE-3 noised
webod tdoe taching
0.8
ory embeddings act as dedicated storage for task-relevant information, while transformer lay-
Total rewerdd
40
0.6
ers encode general behavior. In ViZDoom- Two-Colors (see Figure 8, left), adding noise separately to embeddings and cached hidden states showed performance was more sensitive to hidden-state corruption, highlighting their role continuous rewards and long depen- dencies. Overall, memory embeddings matter most for sparse, discrete decision points (c.g., T-Maze), while cached representations are cru- cial in dense, continuous-feedban tasks like ViZDoom.
30-
Succes
0.4
20
10
0.2
All
Red
Green
0.0
200 400 600 800
Test pillars
Test corridor length
Figure 8: Effect of memory corruption RATE at inference. (left) ViZDoom: performance drops when memory tokens cached states noised. (right) T-Maze: SR degrades when memory em- beddings are corrupted.
RQ2: Performance upper-bound estimate. To estimate the upper-bound performance achiev- able by RATE, we introduce OracleDT - a vari- ant of Decision Transformer augmented with perfect prior knowledge about the environment. Specifically, OracleDT receives an additional input vector E IR Ixd_model prepended and
Table 5: Performance comparison between DT, RATE, and OracleDT. OracleDT an oracle- informed variant used solely approximate the upper bound and is not feasible baseline.
T-Maze
Success Rate
OracleDT
DT
RATE
appended the context sequence, i.e., 5' concat (v, v). This vector encodes one bit information known in advance. In T-Maze, e represents the initial clue (v, = 0 if left, v; = right); in ViZDoom- Two-Colors, encodes the pillar color (u; for red, ;;- = for green). This setup mirrors context augmented with perfectly trained mem- ory embeddings, i.e., concat (M, S, M), where M encodes all relevant information. As a result, OracleDT provides an empirical upper bound on achievable performance when key information is available explicitly. In such settings, we expect the relation R| R[OracleDT] 2 R[RATE] > R RDT hold (see Table 5). Since this privileged information is not generally accessible during training, OracleDT is not a viable baseline but serves as useful reference. The gap between OracleDT and RATE quantifies the effectiveness of RATE'S memory mechanisms in autonomously discovering, storing, and utilizing task-relevant information.
T 90
1.00.1..00
1.00_0.000
1.00.1..00
T 480
1.00:10.00
0.50=0.00
0,90=0007
of environment-critce
900
1.00:10.00
0.50=0.00
0.90=0007
WiZDoom-Two-Clllne
Total Reward
56.5-+0088
24.8+:14
41.5++10
Red Pillars
55.3++16
7.2+0.4
38.2+5.1
Green Pillars 57.2.10.5
42.3113.
44.7.15.8
RQ 3. Memory Retention Valve scheme ablation. In the T-Maze environment, we observed that without MRV, RATE's performance deteriorates on long corridors (L >> K), eventually reaching SR = 50% (see Table 6). This degradation occurs because critical information to be remembered goes into memory embeddings when processing the first segment of the sequence, and then it must be retrieved when making decisions on the last segment. At the same time, due to the recurrent structure of the architecture, memory embeddings continue to be updated during the processing of intermediate segments when no new information needs to be memorized, causing important information from memory embeddings Retention Valve (MRV) and evaluated five variants: MRV-CA-1: Cross-attention mechanism where updated embeddings (Mx+1) query incoming ones (M,); MRV-CA-2: Reversed variant where incoming embeddings (M,) query updated ones (M,-1); MRV-G: Gating mechanism inspired by GRU-based (Chung et al., 2014) memory processing with
leak out. To address this information loss, we introduced the Memory
GTyXL (Parisotto al., 2020); MRV-GRU:

---

Published as a conference paper at ICLR 2026
hidden states; MRV-LSTM: LSTM-
Table 6: Ablation of MRV configurations in Maze 33==00 5=100) Baseline without MRV is marked f. Default: MRV-CA-2.
based (Hochreiter & Schmidhuber, 1997)
memory processing with cell states.
Among all tested configurations, MRV-CA-2 demonstrated best performance (see Table 6).
Model
150
360
600
900
w/o MRV!
1.00 =0.00
0.66
#0.06
0.65 =007
0.61 #0007
This cross-attention scheme uses incoming memory tokens (M,) as queries and updated tokens (Ma+1) ) as keys and values. This config-
MRV-CA-2
1.00 =0,0
0.95 #0006
0.90 =0007
0,90 =0007
MRV-G
0.86
0.77
0.66 =0.07
0.65 =0.0
MRY-GRU MRV-LSTM
0.99 =001
0.74 =0.07
0.56 =0.11
0.55 #0.12
0.85 =006
0.64 10.00
0.51 =0.11
0.47 10:11
uration, referred to simply as MRV throughout the paper, effectively controls information flow
MRV-CA-1
2:01
0.51 2000
20000
0.49 200
through memory. By allowing the model to selectively update its memory based on the relevance of new information, it prevents loss of important context over long sequences.
RELATED WORK
Transformers in RL: Transformers have been applied to online (Parisotto et al., 2020; Lampinen
al., 2021; Morad et al., 2023b; Le et al., 2024), offline (Chen et al., 2021; Janner al., 2021; Wang et al., 2025), and model- based RL (Chen et al., 2022). Prior work often assumes compact observations or known dynamics (Lee al., 2022; Jiang et al., 2023), whereas RATE targets long-horizon credit assignment and memory in partially observable environments, using DT (Chen et al., 2021) as baseline. The Long-Short Decision Transformer (LSDT) (Wang al., 2025) augments DT with dual context windows but still lacks explicit, learnable memory. Fast and Forgetful Memory (FFM) (Morad et al., 2023b) and Stable Hadamard Memory (SHM) (Le al., 2024) instead explore lightweight recurrent slots with greater stability. RNNs in RL: Recurrent models like LSTM (Hochreiter & Schmidhuber, 1997) and GRU (Chung al., 2014) have long supported memory in RL. DLSTM (Siebenborn et al., 2022) replaces transformers with LSTM for sequential control, but RNNs often struggle with long-term dependencies, especially under sparse rewards (Ni et al., 2023). SSMs in RI: SSMs such as S4 (Gu et al., 2021) and Mamba (Gu & Dao, 2023) offer efficient alternatives to attention, showing strong offline RL results (Bar-David et al., 2023; Ota, 2024; Lv et al., 2024; Cao et al., 2024), though their ability to handle aeemory-intensing generalization remains unclear. Memory Memory-Augmented Transformers: Extensions like Transformer- XL. (Dai et al., 2019), Compressive Transformer (Rae et al., 2019), and RMT (Bulatov et al., 2022) extend context via caching or compression. RATE combines token-level memory, hidden-state caching, and a novel MRV gate. Approximate Gated Linear Transformer (Pramanik al., 2023) replaces full attention with a gated, low-rank recurrent update that approximates outer- proerrtrouuctt memory via cosine features, enabling efficient long-range credit assignment at constant cost. Auggmeetdd Decision Transformer (RA-DT) (Schmied et al., 2024) augments DT with an external retrieval memory that stores past sub-trajectories, retrieves relevant ones by vector search, reweights them by utility, and integrates them through cross-attention to guide action prediction in sparse-reward RL.
LIMITATIONS
While RATE is tailored for long-horizon, memory -intenive tasks, its complexity may be unnecessary in fully observable or short-term settings where simpler recurrent models suffice. Nonetheless, RATE matches or exceeds their performance across all tasks. Future work may explore adaptive variants that scale memory based task complexity.
8 CONCLUSION
We propose the Recurrent Action Transformer with Memory (RATE), a transformer-based architecture for offline RL that combines attention with recurrence for long-horizon decision-making RATE integrates memory embeddings, hidden state caching, and a Memory Retention Valve (MRV) to selectively retain critical information across segments. RATE achieves -fate-of-the-art results on memory-intensive tasks such as T-Maze, Minigrid-Memon and 48 POPGym tasks, generalizing up to 9600-step sequences and outperforming both recurrent and transformer baselines. Theoretical analysis shows that MRV guarantees lower-bounded memory preservation across updates, and ablation studies confirm its importance for oong-horizon stability. Despite its memory focus, RATE also performs competitively on standard benchmarks like Atari and MuJoCo, demonstrating broad versatility. These results establish RATE as a unified, ceneeal-purpose offline RL model that excels across both short and long temporal contexts.
ViZDoom--Two-Coloss Memory Maze,
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
The study was supported by the Ministry of Economic Development of the Russian Federation (agreement No. 139-15-2025-013, dated June 20, 2025, IGK 000000C3 13255440009)))
REPRODUCIBILITY STATEMENT
We have taken several measures to ensure the reproducibility of our results. Model details: A full description of the RATE architecture, including pseudocode for both the model and the Memory Retention Valve (MRV), is provided in Section and Algorithms - 2. Theoretical results: Formal assumptions and complete proofs for our preservation theorem are given in Section Experimental setup: Details of environments, training procedures, and evaluation protocols are reported in Section 4, with additional specifications thyperparameters, dataset preprocessing, random seeds, and hardware setup) in Appendix E and Appendix C. Baselines: All baseline implementations are either drawn from widely used open-soure libraries or re-imme re-implemented with hyperparameters matched to their original publications, described in Section 4 and Appendix G. Code and data: An anonymous repository with the implementation of RATE, training scripts, and configuration files submitted as supplementary material. Together, these resources allow for full replication of our theoretical analyses and empirical results.
REFERENCES
Pranav Agarwal, Aamer Abdul Rahman, Pierre-Luc St-Charles, Simon JD Prince, and Samira Ebrahimi Kahou. Transformers in reinforcement learning: survey. arXiv preprint NXXv:2307.00979 2023.
Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi. An optimistic perspective offline reinforcement learning. In International Conference on Machine Learning, pp. 104-114. PMLR, 2020.
Shmuel Bar-David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision s4: Efficient sequence-basee based via state spaces layers. arXiv preprint arXiv:2306.05167, 2023.
Edward Beeching, Christian Wolf, Jilles Dibangoye, and Olivier Simonin. Deep reinforcement learning on budget: 3d control and reasoning without a supercomputer. CoRR, Abs/1904.00106 2019. URL http: //tp://aaiiii org/ alt///////// 1904..94.000 C1806.
Mare G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ-
ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05/1 205500000.0,,,,,,,, 2020.
Aydar Bulatov, Yury Kuratov, and Mikhail Burtsev. Recurrent memory transformer. Advances in Neural Information Processing Systems, 35: 35:11075 1111109991999, 2022.
Jiahang Cao, Qiang Zhang, Ziqing Wang, Jiaxu Wang, Hao Cheng, Yecheng Shao, Wen Zhao, Gang Han, Yijie Guo, and Renjing Xu. Mamba as decision maker: Exploring multi-scall sequence modeling in offline reinforcement learning. arXiv preprint arXiv:2406.000/3, 2024.
Chang Chen, Yi-Fu Wu, Jaesik Yoon, and Sungjin Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.0944 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 15044415599777 2021.
Maxime Chevalier- Chevalier-Boisvert Bolun Dai, Mark Towers, Rodrigo de Lazcano, Lucas Willems, Salem Lahlou, Suman Pal, Pablo Samuel Castro. and Jordan Terry Minigrid & miniworld: Modular & customizable reinforcement learning environments for goal-oriented tasks. CoRR, abs/2306.1383) 2023.
11

---

Published as a conference paper at ICLR 2026
Junyoung Chung. Caglar Gulcehre, KyungHyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555, 2014.
Zihang Dai, Zhilin Yang, Yiming Yang, Jaime Carbonell, Quoc Le, and Ruslan Salakhutdi- nov. Transformer-xl: Attentive language models beyond fixed-length context. arXiv preprint urXiv:/901.022 2019.
Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 1.027799 2016.
Kevin Esslinger, Robert Platt, and Christopher Amato. Deep transformer g-networks for partially observable reinforcement learning. arXiv preprint arXiv:2206.01078, 2022.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep
data-driven reinforcement learning, 2021.
Jake Grigsby, Linxi Fan, and Yuke Zhu. AMAGO: Scalable in-context reinforcement learning for
adaptive agents. In The Twelfth International Conference on Learning Representations, 2024. URL https: wpwpreveeeeewweeeeeeeeeeeeeeeewe net / forum? ddM6xWo..dwff.
Albert Gu and Tri Dao. Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.007522 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:21 1221111.00000 00396, 2021.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning 2019.
behaviors by latent imagination. arXiv preprint
Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()): 1785-1780, 1997.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 34:1 14:11273-11 1286, 2021.
Zhengyao Jiang, Tianjun Zhang, Michael Janner, Yueying Li, Tim Rocktäschel, Edward Grefenstette, and Yuandong Tian. Efficient planning in compact latent action space. In The Eleventh International Conference on Learning Representations, 2023, URL https:////// // /opprrveie..... net/ / forum? (ill==AA7NNERR 7N7VVuuu..
Jikun Kang, Romain Laroche, Xindi Yuan, Adam Trischler, Xue Liu, and Jie Fu. Think before you act: Decision transformers with internal working memory. arXiv preprint arXiv:2305. 16338, 2023.
Feyza Duman Keles, Pruthuvi Mahesakya Wijewardena, and Chinmay Hegde. On the computational complexity of self-attention. In International conference on algorithmic learning theory, pp. 597-619. PMLR, 2023.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative g-learning for offline reinforcement learning. Advances in Neural Information Processing Systems, 33:1 179-1199, 2020.
Andrew Lampinen, Stephanic Chan, Andrea Banino, and Felix Hill. Towards mental time travel: a hierarchical memory for reinforcement learning agents. Advances in Neural Information Processing Systems, 34:28188 112128229215 2021.
Hung Le, Kien Do, Dung Nguyen, Sunil Gupta, and Svetha Venkatesh. Stable hadamard memory: Re- vitalizing eeemory-augmente agents for reinforcement learning, arXiv preprint arXiv:2410.10132 2024.
Kuang- Huci Lee, Ofir Nachum, Mengjiao Sherry Yang, Lisa Lee, Daniel Freeman. Sergio Guadar- rama. Ian Fischer, Winnie Xu, Eric Jang. Henryk Michalewski, et al. Multi-gama decision transformers. Advances in Neural Information Processing Systems, 35:27921 2022.
12

---

Published as a conference paper at ICLR 2026
Wenzhe Li, Hao Luo. Zichuan Lin, Chongjie Zhang, Zongqing Lu, and Deheng Ye. A survey on transformers reinforcement learning. Transactions on Machine Learning Research, 2023. Survey
ISSN 2835-8856. URL Certification.
Qi Lv, Xiang Deng, Gongwei Chen, Michael Y Wang, and Liqiang Nie. Decision mamba: A multi-grained state space model with self-evolution regularization for offline RL. In The Thirty- eighth Annual Conference Neural Information Processing Systems, 2024. URL https: / /openrevieww net ittp:rrrreeeeeettooooiiiiiioaae
Volodymyr Mnih. Playing atari with deep reinforcement learning. arXiv preprint arXiv: 2013.
Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. Popgym: Benchmarking partially observable reinforcement learning. arXiv preprint arXiv:2303.0/859, 2023a.
Steven Morad, Ryan Kortvelesy, Stephan Liwicki, and Amanda Prorok. Reinforcement learning with fast and forgetful memory. Advances in Neural Information Processing Systems, 36:72008-72029, 2023b.
Tianwei Ni, Michel Ma, Benjamin Eysenbach, and Pierre-Luc Bacon. When do transformers shine in RL? decoupling memory from credit assignment. Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL https: htttp:://opnneeveeww www.e/www.eeeeeeeeeeeeeeeeeeeeee wwt. formm/oormmmaa APGXENKt 6h.
Toshihiro Ota. Decision mamba: Reinforcement learning via sequence modeling with selective state spaces. arXiv preprint arXiv:2403.1992. 2024.
Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International conference on machine learning, pp. 7487-7498. PMLR, 2020.
Razvan Pascanu, Tomas Mikolov, and Yoshua Bengio. On the difficulty of training recurrent neural networks. In International conference on machine learning, pp. 1310-1318. Pmir, 2013.
Jurgis Pasukonis, Timothy Lillicrap, and Danijar Hafner. Evaluating long-temn memory in 3d mazes. arXiv preprint 2tte:22200.3888. 2022.
Marco Pleines, Matthias Pallasch, Frank Zimmer, and Mike Preuss. Transformerxl as episodic
memory in proximal policy optimization. Github Repository, 2023. URL httph:tp////c /ttp://gttttt///// comm//arrcooeeter/epp.ddddo wew.deddeeeetttaaaooeeeeeeeeeene
Andrey Polubarov, Nikita Lyubaykin, Alexander Derevyagin, Ilya Zisman, Denis Tarasov, Alexander Nikulin, and Vladislav Kurenkov. Vintix: Action model via in-context reinforcement learning. arXiv preprini arXiv:2501. 2501.22500199000 19400, 2025.
Subhojeet Pramanik, Esraa Elelimy, Marlos C Machado, and Adam White. Agalite: Approximate gated linear transformers for online reinforcement learning. arXiv preprint 1rXiv:2310.15719, 2023.
Jack W Rae, Anna Potapenko, Siddhant M Jayakumar, and Timothy P Lillicrap. Compressive transformers for long-range sequence modelling. arXiv preprini 19Xiv:1911.05300 2019.
Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. In The Eleventh International Conference Learning Representations, 2023. URL
www.reiiieeeeeee net/ 2
Thomas Schmied, Fabian Paischer, Vihang Patil, Markus Hofmarcher, Razvan Pascanu, and Sepp Hochreiter. Retrieval- auggeeene decision transformer: External memory for in-context rl. arXiv preprint arXiv:2410.07071, 2024.
13

---

Published as a conference paper at ICLR 2026
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Max Siebenborn, Boris Belousov, Junning Huang, and Jan Peters. How crucial is transformer in decision transformer? arXiv preprint arXiv:2211.14655, 2022. URL https: //txiv..r/g// abs/2211. 14655.
Artyom Sorokin, Nazar Buzun, Leonid Pugachev, and Mikhail Burtsev. Explain my surprise: Learning efficient long-term memory by predicting uncertain outcomes. 07 2022. doi: 10.485500arXi..2207. 13649.
R.S. Sutton and A.G. Barto. Reinforcement Learning, second edition: An Introduction. Adaptive Computation and Machine Learning series. MIT Press, 2018. ISBN 9780262039246.
Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmies Michael Chang, Natalie Clay, Adrian Collister, al. Human-timescale adaptation in an open-ended task space. arXiv preprint arXiv:2301.07600, 2023,
Trieu Trinh, Andrew Dai, Thang Luong, and Quoe Le. Learning longer-term dependencies in mns with auxiliary losses. International Conference on Machine Learning, pp. 4965-4974 PMLR, 2018.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Jane Wang. Zeb Narth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint 161 1,0763,, 2016.
Jincheng Wang. Penny Karanasou, Pengyuan Wei, Elia Gatti, Diego Martinez Plasencia, and Dimitrios Kanoulas. Long-shon decision transformer: Bridging global and local dependencies for generalized decision-maie making. In The Thirteenth International Conference on Learning Representations, 2025.
URL https:tpp/// /wwpereeeiiieeeeeeeeee iew. net. forum? 4-
Yueh-Hua Wu, Xiaolong Wang, and Masashi Hamaya. Elastic decision transformer. Advances in neural information processing systems,
11532-18550, 2023.
Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontanon, Philip Pham, Anirudh Ravula, Qifan Wang, Yang, et al. Big bird: Transformers for longer sequences. Advances in neural information processing systems, 33:17283-17297, 2020.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christopher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre- pre-trained transformer language models. arXiv preprint arXiv:2205. 01068, 2022.
Zifeng Zhuang, Dengyun Peng, Jinxin Liu, Ziqi Zhang, and Donglin Wang. Reinformer: Max-return sequence modeling for offline rl. arXiv preprint S.iv:220050084000 2024.
14

---

Published as a conference paper at ICLR 2026
Table of Contents
A Discussion: Are RNNs Still Better for Memory?
15
B Decision Transformer
16
Environments
16
C.I Memory-intensive environments.
Standard benchmarks
D Action Associative Retrieval
19
E
Training
21
E.1
ViZDoom Tiooomm-To--ooss..
21
E.2
Passive T-Maze
21
E.3
Memory Maze
21
E.4
Minigrid-Memooy
................................
................................
22
POPGym Suite
22
Atari and MuloCo
F
Additional ablation studies
22
F.1
Additional ViZDoom TrZDoom-TTo-ooe Two- ablation
23
F2
Curriculum Learning
24
F3
Supplemental MRV ablation
25
F4
Ablation on number of segments and segment length
26
G Transformer Ablation Studies
26
Recommendations for Hyperparameter Settings
26
Technical details
29
A DISCUSSION: ARE RNNS STILL BETTER FOR MEMORY?
Our experiments provide a systematic comparison between recurrent and transformer-base archi- tectures in memory-intensive tasks. When trained short sequences, recurrent models such as BC-LSTM perform competitively. For example, in the T-Maze environment, BC-LSTM achieves perfect success rates when trained on sequences up 150 steps, effectively capturing short-term dependencies via its internal state dynamics.
However, this advantage quickly fades training sequences grow longer. Increasing the training horizon from 150 to 600 steps causes BC-LSTM's performance to collapse to 50% success rate across all inference lengths even those shorter than the training context indicating difficulty with gradient stability and information retention over long spans (Figure 6). In contrast, RATE maintains consistently high performance under the same conditions, demonstrating stronger scalability with sequence length. RATE generalizes robustly to inference horizons up to 9600 steps (28,800 tokens), reflecting the effectiveness of its hybrid memory design. The architecture combines token-based recurrence with gated memory updates via the Memory Retention Valve (MRV), enabling reliable propagation of sparse information across long temporal distances.
These findings extend more complex environments. In (Figure 4, Table 1), RATE significantly outperforms BC-LSTM. ViZDoom, RATE maintains balanced performance across red and green cues, whereas BC-LSTM exhibits instability and higher variance. In Memory Maze, RATE achieves substantially higher returns, benefiting from its capacity to encode and retrieve rempoallempporl patterns
and Memory Maze
long episodes.
In conclusion, while RNNs remain effective for short-range temporal dependencies, their performance degrades in long-horizon, searse-reward, and generalization-ccittiall settings. RATE bridges this gap
15

---

Published as a conference paper at ICLR 2026
VizDoom-Two-Colln Minigrid-Memory
T-Maze
Memory-Maze
POPGym-46
cluer
Figure 9; Memory-intensive environments used to evaluate RATE memory mechanisms.
by integrating attention with recurrence, offering scalable and robust memory solution. These results underscore the architectural promise of combining transformer attention with recurrent dynamics for long-term tasks in RL.
B DECISION TRANSFORMER
Algorithm 3 Decision Transformer
Decision Transformer (DT) (Chen et al., 2021) is an algorithm for offline RL that reduces the RL task to a sequence modeling task. DT, the scheme of which is presented in Algorithm 3, the trajectory is not divided into segments as in RATE. Instead, random fragments of length K are sampled from the trajectory, since origi- designed to work only
Require: R R ,0 RR..TT a RIxT
1: R Rxxd Encodern(R)
E Rxxd +- Encoder.(0)
Rxd Encoder,
2: To..r + {(Rg. o0, āo), ..... āo).....
( (RT, or, ar)}
3: n = random(0, -
nally this architecture with MDP. The predicted actions à are sampled autoregressively.
4: â, + Trans Trranfforerr( (f..n
Output: ân - L(an: 0,,)
C ENVIRONMENTS
C.I MEMORY-INTENSIVE ENVIRONMENTS
In this section, we provide an extended description of the environments used in this paper, as well as the methodology used to collect the trajectories. Table 7 summarizes the observations type, rewards type, and actions type for each of the environments considered in this paper.
C.1.1 VIZDoOM-TWO-COLLOO
used a modified ViZDoom- DirDoom-Tww-Coorss. environment from (Sorokin et al., 2022) to assess the model's 's memory abilities. The agent initially having 100 hit points (HP) is placed in inner walls filled with acid. At each step in the environment, the agent loses a fixed amount health (10/32 HP per step). In the center of the environment, there is pillar of either green red color, which disappears after 45 environment steps. Throughout the environment, objects of two colors (green and red) generated. When the agent interacts with object of the same color the pillar, it gains an increase in health of +25 and a reward of +1. When the agent interacts with an object of the opposite color, loses similar amount of health. The agent receives +0.02 for each step it survives. The episode ends when the agent has zero health. Thus, the agent needs remember the color of the pillar select items of the correct color, even if the pillar is out of sight or has disappeared. The agent does not receive information about its current health or rewards, as these observations essentially convey the same information as the color of the pillar but persist beyond step 45.
without
additional reward of
collected a dataset of 5000 trajectories of 90 steps in length using a trained A2C (Beeching et al., 2019) agent (an agent trained with a non-disappearing pillar). The average reward for these 90 steps 4.46. When collecting trajectories, to ensure that the agent saw the pillar before disappeared, the agent always appeared facing the pillar in the same place midway between the pillar and the nearest wall. In order to successfully complete this task, the agent needs to remember the color of the pillar. This environment tests the long-term memory mechanism, since the agent needs to
16

---

Published as a conference paper at ICLR 2026
Table Description of observations and reward functions for the considered environments.
Environment
Obs. Type
Rew. Type
Act. Space
Obs. Details
ViZDoom-Two Two-Colors
Image
Continuous
Discrete
First-person view Low-dimensional vector
T-Maze
Vector
Sparse & Discrete
Discrete
Memory Maze
Image
Sparse & Discrete
Discrete
First-person view 3x3 grid centered agent Vector or 2D grid
Minigrid- Memigrid-Memrry
Image
Sparse Discrete/Continuen
Discrete
POPGym
Vector/Image
Discrcte/Continuon
Action Assoc. Retrieval
Vector
Sparse & Discrete
Discrete
Symbolic vector input Full game screen
Atari
Image
Sparse & Discrete
Discrete
MuJoCo
Vector
Continuous
Continuous
Low-dimensional state voctor
retain information about the pillar for time much longer than the pillar has been the environment. Using only short-term memory and, for example, collecting the next item of the same color as the previous collected item, it will not be possible for the agent to survive for a long time, as this policy is extremely unstable. This is due to the fact that in the training dataset the agent occasionally makes a mistake and picks up an object of the opposite color. Thus, irrelevant information about the desired color may enter the transformer context and the agent will start collecting items of opposite color, which will quickly lead to failure.
C.1.2 T-MAZE
To investigate agent's ' long-term memory on very long environments (the inference trajectory length much longer than the effective context length Keff) we used modified version of the T-Maze environment (Ni et al., 2023). The agent's objective in this environment is to navigate from the beginning of the T-shaped maze to the junction and choose the correct direction, based on a signal given at the beginning of the trajectory using four possible actions signal, represented as the clue variable and equals to zero everywhere except the first observation, dictates whether the agent should turn up (clue = 1) down (clue 1). Additionally, constraint on the episode duration T = L + 2, where the maximum duration is determined by the length of the corridor L to the junction, adds complexity to the problem. To address this, binary flag, represented as the flag variable, which is equal to 1 one step before the junction and 0 otherwise, indicating the arrival of the agent at the junction, is included in the observation vector. Additionally, a noise channel is added to the observation vector, with random integer values from the set {-1, 0, +1}.
(left, up, right, down); }. This
The observation vector is thus defined as o [y, clue, flag, noise]. where y represents the vertical coordinate. The reward is given only at the end of the episode and depends on the correctness of the agent's 's turn at the junction, being for correct turn and 0 otherwise. This formulation deviates from the traditional Passive T-Maze environment (Ni et al., 2023) (different observations and reward functions) and presents a more intricate set of conditions for the agent to navigate and learn within the given time constraint.
The dataset consists of 2000 of trajectories for each segment of length 30 (i.e. 6000 trajectories for the Kess x 30 90) and consists only of successful episodes. An artificial oracle with priori information about the environment was used to generate the dataset.
C.1.3 MEMORY MAZE
In this first-person view 3D environment (Pasukonis et al., 2022), the agent appears in a randomly generated maze containing several objects of different colors at random locations. The agent's task is to find an object of the same color in the maze as the outline around its observation image. After the agent finds an object of the desired color and steps on it, the color of the outline changes and the agent must find another object. The agent receives +1 reward for stepping on the correct object. Otherwise, it receives no reward. The duration of an episode is a fixed number and is equal to 1000. Thus, the agent's 's task is to find as many objects of the desired color possible in limited time. The agent's effectiveness in this environment depends on its ability to memorize the structure of the maze and the location of objects in it in order to find the desired objects faster. Using the Dreamer model (Hafner et al., 2019) to collect dataset of 5000 trajectories only achieved an average award of 4.7 per episode, i.e., rather sparse dataset.
17

---

Published as a conference paper at ICLR 2026
C.1.4 MINIGRID-MEMORT
Minigrid-Memonn Schevalier-Boisverr al., 2023) is 2D grid environment designed test an agent's long-term memory and credit-assignment (Ni et al., 2023). The environment map T- shaped maze with small room with object inside it at the beginning of the corridor. The agent appears at a random coordinate in the corridor. The agent's task is to reach the room with the object and memorize it, then reach the junction at the end of the maze and make turn in the direction where the same object is located as in the room at the beginning of the maze. A reward r = 1 - 0.9 x + is given for success, and 0 for failure. The episode ends after any agent turns junction or after limited amount of time (95 steps) has elapsed. The agent's 's observations are limited to a 3 x 3 size frame. 10000 trajectories with grid size 41x41 were collected using PPO (Schulman al., 2017) with Transformer-Xi XL. (TrXL.) (Pleines et al., 2023) with a context length equal to the maximum episode duration.
C.1.5 POPGYM
POPGym (Morad et al., 2023a) is a benchmark suite consisting of 46 diverse partially observable environments designed to isolate different aspects of memory use and generalization in reinforcement learning. The tasks include both short-horizon reactive scenarios and long-horizon memory puzzles that require the agent to remember information across extended delays or infer hidden states from past observations. The environments vary in observation modality (image vs. vector), reward sparsity, and temporal dependencies. For our dataset, we followed the original POPGym evaluation protocol and used a PPO (Schulman al., 2017) agent with GRU (Chung al., 2014) backbone (PPO-GRU), which showed the best performance in the original benchmark. We collected trajectories using this policy for all 46 environments. The collected dataset reflects the diverse difficulty and memory requirements of the benchmark and serves as a challenging testbed for evaluating general -urposs memory architectures like RATE.
C.2 STANDARD BENCHMARKS
C.2.1 ATARI GAMES
For the Atari game environments (Bellemare et al., 2013), we used the same dataset as in DT, namely the DQN replay dataset with grayscale state images (Agarwal al., 2020). This dataset contains 500 thousand of the 50 million steps of an online DQN (Mnih, 2013) agent for each game. We use the following set of games: SeaQuest, Breakout, Pong and Qbert.
C.2.2 MUJoCo.
ap {0,1) 1}
Despite the fact that memory is not required in decision making in control environments like MuJoCo (Fu et al., 2021), conducted additional experiments in this envi- ronment to compare with DT. For the continuous control tasks, selected a standard MuJoCo locomotion envi- ronment and a set of trajectories from the D4RL bench- mark (Fu et al., 2021). Since chose DT and TAP as the main models for comparison on this data, we focused on the environments used both works (HalfCheetah, Hop- per, and Walker). We used three different dataset settings: 1) Medium - million timesteps generated by a "medium" policy that achieves about third of the score an expert policy; 2) Medium-Replay the replay buffer of an agent trained with the performance of a medium policy (about 200k-400k timesteps our environments); 3) Medium-Expert- - million timesteps generated by the medium policy concatenated with million timesteps generated by expert policy. The scores for the Mu- JoCo experiments are normalized such that 100 represents an expert policy, following the benchmark protocol outlined in (Fu et al., 2021). The performance metrics for Conservative Q-Learning (CQL) and Trajectory Autoencoding Planner (TAP) are reported from the TAP paper (Jiang et al., 2023), and for DT from the DT paper (Chen et al., 2021), 1), as they the same dataset and evaluation protocol.
So
S
a1.. =
aN E {0,1}
Figure 10: Action Associative Retrieval.
18

---

Published as a conference paper at ICLR 2026
RATE Success Rate
DT Success Rate
RATE
DT
CQL-LSTM Success Rate
COL-MLP Success Rate
2220
Volidatton Sequence Langth
Validaaton Sequence Langth
CQL-LSTM
CQL-MLP
DGRU Success Rate
DLSTM Success Rate 1
Malidation Sequence Length
Validation Sequeecee Length
DGRU
DLSTM
BC MCMMLP Success Rate
DMamba
BC-MLP
LSDT
BC-LSTM
XX 22
Traight
Validiatoo Sequeecee Lengtt
Validadie Sequence Langth
LSDT
BC-LSTM
RMT Success Rate
TrXL Success Rate
Tratgin
Vatsaation Sequence Length
Va Valiaciinn
RMT
TrXL
Figure 11: Results for all models in the T-Maze generalization task.
D ACTION ASSOCIATIVE RETRIEVAL.
As shown in Figure 6, DT has SR = 50% for inference at corridor lengths longer than the transformer context length. This is due to the fact that even a DT trained on balanced data has a slight bias in the predicted probability towards one of the two required actions, which leads to the fact that when
19

---

Published as a conference paper at ICLR 2026
1.0
100
1.0
100
0.8
0.8
Probability actiion
Number of train stepp, %
Probability acttionn
Number of train stepps %
0.6
0.6
0.2
0.2
DT, train
RATE, train R: 0.97
R: 1.00
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Probability action
Probability of action
1.0
100
1.0
100
0.8
0.8
Probability action
Number of train stepp, %
Number of train steps %
0.6
Probability
0.4
0.2
0.2
DT, val
RATE, val
R: -0.97
R: 0.80
0.0
0.0
0.0
0.2
0.4
0.6
0.8
1.0
0.0
0.2
0.4
0.6
0.8
1.0
Probability action
Probability of action
Figure 12: Experimental results with RATE and DT in the AAR environment. The graphs show the 10-runs average results of training on trajectories of length T 90 and validation trajectories of length = 180, for RATE with Keff = x 30 = 90 and for DT with K 90.
I. >>> the agent constantly produces only one action: up or down. In turn, the presence of memory in the agent allows combat this problem.
To check how the agent's performance changes during training, we design an Action Associative Retrieval (AAR) Figure 10 environment.
There are two states in this environment: So and Sj. The agent appears in state So and by performing the action dg E {0, 1} moves to state Sy. Next, the agent must take N steps move from state S1 to state Sj by performing action = 2 (no op.). At the end of the episode, the agent must perform the same action that moved it from state S to state in order move from state S₁ state Sg- Thus, the action E (0, 1, 2}. Agent observations the index of the current state, flag 10, is a flag equal to in case the next step requires returning to the initial state and equal to 0 otherwise, noise E {-1, 0, +1} } is the noise channel. The agent receives a +1 reward if returns to the initial state So by performing the action that took out from the Sy to the S1, and -1 in other cases. The training dataset consists of oracle- generated 6000 trajectories with positive reward.
[state, flag, noise]. where state {0. 1)
More formally, like) tasks under the condition that:
talk about the presence of memory in an agent when solving AAR (T-Maze-
No
N; Σ
Vt K
= = alao =a))
pi(a: a a' = =')>>
No
N
This condition means that the agent has memory, the sum the average conditional probabilities over all experiments will be greater than one, i.e., these probabilities are independent of each other.
20

---

Published as a conference paper at ICLR 2026
Provided that the sum of these probabilities is less than equal to one, the agent will choose at best the same target action in most experiments, even if another action is required.
where a", a' E A two mutually exclusive actions leading to reward; is the step at which the final action is required; Ni are the number of experiments in environments where target action a a and a, = a², respectively.
In the results Figure 12, the first 1% of training steps was removed because it corresponds to the beginning of the training and is unrepresentative. Blue dots correspond to the beginning of training, red dots to the end of training. As can be seen from Figure 12, during training, the probabilities pa aa a²) and pi =aa = a²) the training trajectories have strong positive correlation (R 1.00 and R = 0.97). where R - correlation coefficient, This indicates that within-context (effective context) DT and RATE models are able to predict both a and a actions equally well.
At the same time, during validation, for the RATE model this pattern preserved the red points corresponding to the probabilities of choosing actions a and a' are in the upper right part of the graph, positive correlation persists (RRATE of red dots is skewed toward choosing action a and action a with equal probabilities equal to 0.5. Thus, in sum, these probabilities are less or equal to one, as evidenced by strong negative correlation (Re -0.97). The results confirm the inability of DT to generalize trajectories whose lengths exceed the context length and the ability RATE to handle such tasks.
0.80). On the other hand, in the DT case, the cluster
E TRAINING
This section provides additional details on the training process of the baselines considered in the paper. We treated the inclusion of the ford-forward network (FFN) block RATE's transformer decoder hyperparameter, as RATE performed slightly better without FFN in some environments. In contrast, other transformer-based baselines were trained with the standard transformer decoder including FFN.
E.1 VIZDOOM-TWO-COLL
Since the pillar disappears at time t=45, all trajectories span from t=0 to t=90 to ensure that the cue remains available during training. In this setting, we compare DT with context length K=90 to RATE, RMT, and TrXL models using K K330 and N=3 segments. Thus, RATE processes sequences of the same total length Ka==Nxx Nx K =90 but accesses only K=30 tokens at time. Additionally, we experiments with N=3, K-50, and 1-150 to validate model robustness under longer and more complex configurations.
E.2 PASSIVE T-MAZE
We trained models on sequences of length Train E {9, 30, 90, 150, 300, 600, 900} and evaluated them Twa E {9, 30, 90, 150, 300, 600, 900, 1200, 2400, 4800, 9600}. For RATE, each sequence was split into N = 3 segments, yielding a context length of K = Twan/3. All training trajectories started from notation MODEL-N, where N = 3 indicates segmentation into three recurrent blocks (e.g., RATE-3 trained on full sequences length = with K 30). This convention is used throughout the ablation studies.
ensuring the cue was always included. In what follows, we adopt the
E.3 MEMORY MAZE
To train RATE, DT, RMT, and TrXL on Memory Maze, we used the same approach for ViZDoom- Two-Colors environment, but instead of using fixed trajectories starting at = O, we sampled consecutive 90-step subsequences from the original 1000-step trajectories. Each subsequence was sampled with a stride of 90 steps, resulting in approximately 11 training sequences per original trajectory. As in the of K = 90 and for RATE, RMT, and TrXL with a context length of K 30 and number of segments N 3, i.e., effective context length Kaff N K 3 30- 90.
TarooomTTw--olooo case, training for DT was performed with context length
21

---

Published as a conference paper at ICLR 2026
E.4 MINIGRID-MN MINIGRID-MEMORY
To train baselines in this environment, we used only mazes of fixed size 41 41, ensuring consistent corridor length during training. For evaluation, models were validated on mazes ranging from 11 x 11 501 501, where corridor lengths vary within each grid, enabling assessment of both interpolation and extrapolation capabilities. All training trajectories used an episode timeout of 96 steps, while validation trajectories across all maze sizes used longer timeout of 500 steps. As in T-Maze, each trajectory began at i = O, ensuring the cue was always observed. During training, RATE used context length of K = 30 with N = segments, while other baselines (except RMT and TrXL) used K 90.
E.5 POPGYM SUITE
POPGym (Morad al., 2023a) comprises 46 tasks of varying memory complexity, including both memory puzzles and reactive POMDPs. Since episode lengths vary widely across tasks - from as short as 12 steps as long as 1000 ensured consistent and fair memory evaluation for RATE by setting the context length K- = 7/3 and using N = 3 segments for every environment, where T denotes the maximum episode length of each task. This uniform configuration allowed RATE to process full trajectories with recurrent segmentation, ensuring its memory capacity was equally tested across tasks of different lengths and difficulties.
E.6 ATARI AND MuJoCo
When training RATE on Atari games and MuJoCo control tasks, sequences of length T 90 (Atari) and T = 60 (MuJoCo) were sampled randomly from the original trajectories in the dataset. These trajectories were then divided into 3 segments of length K 30 (Atari) and K 20 (MuJoCo). forming an effective context of length Kefs = N x K 90 (60 for MuJoCo).
For Atari, we used the identical experimental design described in the DT paper (Chen et al., 2021). It is worth noting that we presented raw scores for Atari, rather than gamer- normalzzed scores as described in the DT paper. Table shows the results for Atari environments. RATE outperforms DT significantly in environments like Breakout and Qbert. We attribute this to the observation that, although these environments do not explicitly demand memory, intricate dynamics from the past exert greater influence agent behavior than in environments such as SeaQuest. Actions executed in the past notably alter the present state of the environment in Breakout and Qbert, whereas in SeaQuest, such actions hold little significance. For instance, the emergence of enemies and divers in SeaQuest is entirely independent of the agent's ' prior actions.
For MuJoCo, our findings suggest that the conventional strategy of utilizing return is not suitable for our segment- scheme. The issue arises during the trajectory, where the agent's return persistently diminishes. However, the true value of the agent's state the onset and conclusion of the episode could remain unchanged, provided the agent's policy performs consistently well. To rectify this discrepancy, we propose novel evaluation strategy for MuJoCo tasks. In this approach, each segment commences with the maximum return, simulating the scenario where the agent initiates the trajectory anew. This method effectively mitigates the aforementioned issue, enhancing the accuracy of our evaluation process. Our MuloCo experiments in Table 3 show that this benefits performance significantly for some environments. Thus, using RATE allowed us obtain the best metrics for MuloCo in 3/9 cases compared to the other baselines. RATE also outperforms DT in 9/9 tasks.
F ADDITIONAL ABLATION STUDIES
To determine the optimal hyperparameters associated with memory mechanisms, additional ablation studies were performed in ViZDeom-Two-Colore and T-Maze environments, and the results are presented in Figure 14 and Figure 13 (right). From the ablation studies results, it was found that for environments like ViZDoom- Two-Colors with continuous reward signal and image observations, the best results can be obtained using number of cached memory tokens mem_ _en = (K x + 2 x tom_mem_tokenn x N, where K context length and N number of segments.
22

---

Published as a conference paper at ICLR 2026
Table 8: RATE hyperparameters for different experiments. Leaky ReLU used in Atari.Pong The listed hyperparameters for ViZDoom- De-DDoo-Two-CCC Two-Colors and T-Maze correspond to the experiments with Train 150, while for POPGym, they reflect the settings used in the POPGym-Concentratin task.
Hyperparumeter
NiZDoum2C
Memory
Minigrid-Menne
POPGym
MuJoCo
Memory-speckee specifif porameters
Number memory tokens
15
10
30
Number of eached tokens
100
360
0
180
100
360
60
Number heads
MRV activation
ReLU
ReLU
ReLU
ReLU
ReLU
ReLU
ReLUU
Transformer archlecture
layers
Number of altcation heads
8
8
2
Embedding dimension
64
64
32
Context length
30
50
18
Number segments
3
3
Skip doc FEN
False
True
True
False
True
True
True
Regularization
Hidden dropout
0.2
0.5
0.2
0.3
0.1
0.2
0.2
Attention dropout
0.05
0.2
0.1
0.05
0.05
0.05
Weight decay
0.001
0.1
0,001
0.001
0.001
0.1
0.1
Troining confguration
Max epachs
150
80
200
500
200
10
10
64
32
Loss function
CE
CE
CE
CE
CE
CE
MSE
Optimizer
AdamW
AdamW
Adam
AdamW
AdamW
AdamW
Adam
Learning rate
3e
3e-4
le-4
le.4
3e-4
3e-4
6e-5
Grad norm clip
5.0
1.0
1.0
5.0
5.0
1.0
1.0
Cosine decay
True
False
False
Linear warmup (8).
True
True
Thie
True
Thee
True
Thie
(0.9, 0.999)
0.95)
(0.9, 0.999)
(0.9, 0.999)
(0.9, 0.999)
(0.9, 0.95)
(0.9, 0.95)
the other hand, for environments with sparse events like T-Maze, it has been found that using caching of hidden states of previous tokens (mem_1e len > 0) prevents remembering important information.
Pillars
Red Piliars
50
RATE
2.38 =0
2.00
Lm 200
30x 900
0.50
Total Reward
40
0.35
0.00g
30
20
0.04
10
12
35
10
0.0
Noise
Figure 13: (left) Investigating the RATE memory tokens noise effect the ViZDoom- ViZDoom-Two--ollo (right) Results of RATE-3 (trained on corridor lengths & 90) ablation studies in the T-Maze environ- ment. n a__head____ ca number of MRV attention heads, num _emm_mmm____m okenss number memory tokens.
F.I ADDITIONAL VIZDOOM TERDOOM-TTOCOLOORR ABLATION
The effect of combining of memory tokens with noise is shown in Figure 13 (left). The noise was applied as convex combination: With unchanged caching of hidden states from previous steps at growth of the noise parameter at first there a decrease of performance at inference on green pillars (up to a = 0.5), and only then a decrease of performance at inference that memory embeddings is trained to record mostly information about red pillars, which helps to combat bias in the training data.
(1-a) memory Soteyyttoesssseeee + X noise.
red pillars. This phenomenon
be explained by the fact
23

---

Published as a conference paper at ICLR 2026
Table 9: Performance on POPGym tasks (mean over three runs, 100 seeds each).
Dataset Average
Enviroument
RATE
DT
Random
BC-MLP
BC-LSTM
Return
AutoencedeEasy-v00
-0.29 + 0.00
-0.17
0,00
-0.50
-0.47 0.00
00.00
-0.26
AutoencodeMedimm-W
-0.47 0.00
0.00
-0.49 0.00
-0.50
-0.49
-0.47
-0.48
AutoencodeHard
-0.46
0.00
-0.50
-0.50
-0.44
HanleshipEasy-ve
-0.81 + U.02
-0.93
0.03
-0.46
1.6600.00 0.01
-1.00-0.00
-0.49+0.01
-0.35
MaddishipMedium--x0
0.91 0.00
-0.91 + 0.03
-0.39
-1.00=0.00
-0.1
-0.43
BattleshipHar--OO
0.01
-0.1 0.00
-1.00
-0.67
-0.10
Enncentrationnassyyvo
-0.6
0.01 -0.05
-0.19
-0.92
0.00
-0.14
-0.22
inncentraioonMedium--o
0.00
0.00
-0.84 0.00
-0.88
0.00
0.00
-0.87
Ho
0.25
-0,25
-0.19 0.00
-0.92
0..19
-0.41
CounRccallEass--0
0.01
-0.46
0.01
0.01
0.00
-0.92 0.00
0.05
0.22
CounfeecallMciinn-s00
-0.47
-0.75
0.03
-0.88 0.00
-0.88 0.00
0..0
-0.48
CoumtRecallHadd-v0
-0.54 0.00
0.02 -0.81
-0.93
-0.92 0.00
0.00
0.55
Higherl Lowerfasy-v0
0.50 0.00
0.50
+
0.47 0.00
0.50
0.51
HigherLoowrMeedium
0.50=
0.00
-0.01 = 0.00
0.00
0.50
0.49
Higherl LowerHanddvv
0.52 0.00
0.51 0.00
0.01 + 0.01
0.50
0.51 + 0.01
0.49
LahyrinthEscuppeaag
0.95 0.00
0.80
-0.39 0.00
0.72
0.02
0.95
abbytinhhxxcue-Mdddum
-0.81
0.882
-0.44
0000
0.00
-0.44
LabyrinhExcupeHarda
-0.56
0.,7 0.01
0.01
-0.3
00.09
-0.49
LabyrinthExplae Eplore
0.95
0.88 0.06
0.01
0.01
0.00
0.96
AdlyrinthExploreMMMeo
0.79 0.00
0.77
-0.73 0.00
0,26
0.71
0.79
LabyrinthExpooreHarkka
0.88
0.00
-0.61
0.05
0.82
0.87
Mine
wineSSeeperEasy-uu
0.15 0.03
0.03
-0.26 = 0.03
00.1
0.20
0.28
MineSweeperMediung
-0.44 0.00
-0.40 = 0.01
-0.43 = 0.00
-0.49
-0.35
-0.27
MineSwesperHHate
-0.20 0.00
-0.37 + 0.00
-0.39 + 0.01
0..6
-0.00
0.37 0.01
0.27
0.02 + 0.00
0.05 0.00
0.02
0.62
0.22
0.27
= 0.00
0.00
0.17
0.43
0.32 - 0.01
0.35 + 0.01
0.01 + 0.00
0.21 + 0.01
0.14
0.59
NoisyPosittionOniiy PosiitiotOOllyCCaoollllaaassee
0.88 = 0.03
0.87 0.02
0.11 + 0.00
0.23 - 0.00
0.04
0.98
Noisy PoisyPositiooOnnlyyee PoniiiooOOnillCaroPoeeeeiit Mario
0.18 0.01
0.17
0.11 = 0.00
0.00
0.22
0.36
Noisy CanaiiiooOOiiyyaaPPeeeaar----
0.33 0.01
0.12 0.00
0.00
0.25
0.57
0.87 0.00
0.84
0.27 + 0.01
0.31 = 0.00
0.88
0.90
0.60 = 0.01
0.56 = 0.01
0.26 - 0.00
0.28 = 0.00
0.66
0.67
Noisy PositiinnOOiyyeemmalottt vO
0.68 0.00
0.63 + 0.01
0.27 + 0.01
0.30 0.00
0.72 0.00
0.73
AnliionOnlyCarPPlleeee
0.93 0.03
1.00 0.00
0.12
0.00
0.00
PositionOnly CanteennnnCaaaeeeeeeaeaa CarPPoleMeduum
0.05 0.01
0,03 = 0.00
0.04 = 0.00
0.05 0.00
0.06 0.06:00.00
1.00
PositionOnly CarnnnnlyCCeeeaaartt -i)
0.07 0.00
0.34 + 0.08
0.06 + 0.00
0.09 0.00
0.12+
1.00
PositionOnly PeaitinnOnlyPenuuluuEEaaa PendulumEEsy-v0
0.54 0.02
0.51 + 0.03
0.27 + 0.00
0
0.91 +00000000
0.92
PositionOnly PeiitionOnlyPenduluuMMee PendulumMedium--00
0.17 = 0.01
0.49 0.01
0.26 = 0.00
0.28
0.00
0,82
PositionOnly PcndulumHard-v0
0.49
0.00
0.26
0.30
0.00
0.88
DepeaFustEEasy
1.00 0.00
0.16
-0.49 = 0.01
-0.50
1.00 0.00
1,00
RepeutFirsMMedume
0.10 0.02
0.12
-0.50 0.00
-0.50 0.00
0..00 0.00
0.99
RepeatFnssHard-v0
0.99 0.99=0.00
-0.21
-0.50
-0.50
0.99 0.01
1,00
RepeatPreviousasyy Easy---
100
-0.49
-0.52 0.00
1.00
Repeal/revvoossMeddum-
-0.46 0.00
-0.47 + 0.00
-0.51 = 0.00
-0.48 0.00
-0.45 0.00
-0.48
ReppattrevviousHann--dd
00.88 0.01
-0.38 + 0.00
-0.50 + 0.01
-0.50 0.00
-0.38 0.00
-0.39
VelecityOOlyCCamtoleeaaee
1.00 0.00
1.00 + 0.00
0.11 + 0.00
0.99
0.00
1.00
1.00
VelocityOniyCCmtooeeMciii
1.00 = 0.00
0.96 0.02
0.01 0.00
0.63 0.00
0.00
0.99
VelnetiyOniyCantPlleeti
1.00 0.00
1.00 0.00
0.06 + 0.00
0.83 + 0.01
1.00 0.00
1.00
F.2 CURRICULUM LEARNING
Since in the T-Maze environment, the number of actions at the junction relates the number of actions when moving straight along the corridor as J. and tends to 0 as L increases, there is a significant imbalance in the agent's action distribution, which can cause problems when performing rare class (turning actions) prediction. Theoretically, this situation can be remedied through curriculum learning.
Curriculum learning (CL) is a technique in which a model is trained on examples of increasing difficulty. In this approach, the model is first trained on the set of trajectories Q qi of length K 1, then the trained model re-trained on the set of trajectories Q2 = qi U (2, where the set 42 is formed by trajectories of length K x 2, and so on (in order of increasing complexity of the trajectories). Thus, for the N segments considered during training, the set QN U % is used.
In the T-Maze environment, DT, RATE, RMT, and TrXL were trained with and without curriculum learning because this approach theoretically produces better results. However, it is important to note
24

---

Published as a conference paper at ICLR 2026
Table 10: Experimental setup and evaluation metries across different environments. Nruns denotes the number of model runs; N seedds denotes the number of inference episodes with different seeds; sem denotes standard error of the mean, and std denotes standard deviation.
Experiment Setup
Results
Environment
Nruns
seeds
Metric
Notation
Memory-intensive environments
WiZDoom-Two-Colorg
6
100
Return
meanissem
T-Maze
4
100
Success Rate meantsem
Memory Maze
3
100
Return
meantssm
Minigrid-Memory
3
100
Return mean-tsem
POPGym
3
100
Return
sean-tsem
Diagnostic environment
Action Associative Retrieval
10
Success Rate
meanttsem
Thee Rewernd
40
100
20
Total
Racs
number cached tokens (mem. Jen)
number of memory tokens (num mem tokens) with mam. Jen = (90 + 2 mem. tokens) N
40
30
Greens
Total
Reds
Greers
number of memory tokena (num. mem. tokens) with mem, Jen
number of MRV attention heads (n_ headheed.coo
60
10
No
Leaky RULU
RacUU
EUU
MRV activation
Figure 14: Results of RATE ablation studies in the ViZDoom- ViZDoom-Two-Colorr environment.
that the T-Maze task is successfully solved by the RATE model without using curriculum learning, and even vice versa its slightly degraded performance long corridors. However, with respect to TrXL, the use of CL yielded slightly better results. The work showed that using CL does not achieve significantly better performance on the T-Maze task. The results of using the CL on the T-Maze environment are presented in Figure 16 (left), and the results of applying noise to memory embeddings to assess its importance are presented in Figure 16 (right).
F.3 SUPPLEMENTAL MRV ABLATION
One of the options for implementing the memory tokenization gating mechanism approach similar to the one proposed in Gated Transforer-2 XL (GTrXL) (Parisotto et al., 2020) work. Thus, the MRV-G scheme inspired by the gating mechanism from GTrXL and implemented follows:
r = (M, W, + Mn+1 U,)
(8)
== a MntiUs bias)
25

---

Published as a conference paper at ICLR 2026
Table 11: RATE encoders for each part of (R, o, a) triplets. We use an Embedding layer for encoding discrete actions and a Linear layer for continuous ones. - channels / kernel sizes / padding. For POPGym tasks with grid-based observations (e.g.. MineSweeper and Battleship), we encoded the grid using a token dictionary followed by a linear encoder to produce a fixed-length vector. Actions were encoded using an embedding layer for discrete control tasks, while linear layer was used for continuous control environments (e.g., PositionOnly PanOOnyleennuumi
Encoder Configuration
Environment
Return
Observation
Conv. params*
Action
Image-basee environments
ViZDoom Twon--Two-Cooo Two-Colors
Linear
Conv2D 3
(32, 64, (8)
Embedding Emhedding Emhedding Embedding
Memory Maze
Linear
Conv2D 3
(32, 64, (8, 3)
Minigrid-Memon
Linear
Conv2D
(32, 64, 6 3)/0
Atari
Linear
Conv2D x 3
(32, 64, 64) (8, 4, 3/0
Vector-based environments
T-Maze
Linear
Linear
Embedding
MuJoCo
Linear
Linear
Linear
Action Associative Retrieval
Linear
Linear
Embedding
POPGym
Linear
Linear
Embedding / Linear
h = tan, + (Ma+1 r)),)
(10)
Mn+1 = (M,, z) + = x h)
(11)
The results of the RATE (trained on corridor lengths of < 150) inference on the T-Maze environment
with these MRV configurations are shown in Figure 17 and in Table 6. The results presented
in Figure 17 confirm the high stability of RATE when using roos-attenntion-bsee MRV (MRV-CA-2),
as well as the model's s ability to hold important information in memory embeddings when inference on long tasks.
F.4 ABLATION ON NUMBER OF SEGMENTS AND SEGMENT LENGTH
Partitioning the trajectorics into fixed-length segments allows the RATE model to train on long trajectories without increasing the context size, which makes the parameters N (the number of segments into which the training trajectories are divided) and K (the context length, i.e., the size of a single segment) critical because they determine the length of the effective context Keff K N. Figure 18 presents the results of ablation studies for parameters N and K at fixed Kefs 90.
G TRANSFORMER ABLATION STUDIES
Transformer core hyperparameters. This section presents the results of ablation studies on the main hyperparameters of the RATE transformer. The RATE configuration for the T-Maze environment specified in Table was chosen for the ablation studies. The ablation studies focus on understanding the impact of key hyperparameters by systematically varying one parameter while keeping others constant. The results are shown in Figure 20, Figure 21, and Figure 22.
Feed-Forward Network. For RATE, the inclusion of the decoder feed-forwar block is treated as a tunable hyperparameter. In most environments, we disable it, as doing so often leads to better performance Figure 19. However, for ViZDoom-Two-Colbn and Minigrid-Memory, retaining the feed-forward block yields slightly improved results, and thus enabled in those settings.
found that
H RECOMMENDATIONS FOR HYPERPARAMETER SETTINGS
Transformer based models require careful hyperparameter tuning, and the addition of memory mechanisms in RATE introduces few more components. However, configuring RATE remains
26

---

Published as a conference paper at ICLR 2026
MRV-CA-2 (preferred)
MRV-G
MHA
M Mn+1
tanh
M
W
M 72
M Mn+1
MRV-CA-1
W
U
W
U
MHA
Mn+1
X
Y
M n+
M n
M 12
Mn+1
MRV-LSTM
MRV-GRU
Impull
LSTM
M Mn+1
Input GRU
h
M n+1
Mn
Cn
M Mn+1
M Mn+1
h
Figure 15: Memory Retention Valve configurations used in the ablation study. MRV-CA-2: cross- iiteetiio-bbsed MRV which uses attention mechanism to control the updating of memory embed- dings and which is used in the work as the main mechanism. MRV-CA-1: uses the same mechanism MRV-CA-2 but the updated memory embeddings M fed Query, and the incoming memory embeddings Mn are fed to Key and Value. MRV-G: gated MRV which uses a gating mechanism similar the one used in Gated Transformer-Xi (Parisotto et al., 2020). MRV-GRU: uses GRU (Chung et al., 2014) block to process updated memory embeddings with hidden states. MRV-LSTM: uses a LSTM (Hochreiter & Schmidhuber, 1997) block to process updated memory embeddings with cached states.
nuived with CL
Succes Rate
Success Rate
0.4
0.4
TXL-3 - DT:3
-A-
RATE-3
w/o CL
0
100
200
300
400
500
600
700
800
900
100
200
300
400
500
600
700
800 900
Test corridor length
Test corridor ength
Figure 16: (left). Results with and without the of curriculum learning and (right) results of replacing RATE memory tokens with white noise at inference in T-Maze.
largely similar to tuning a standard transformer. Based on extensive empirical evaluation, we provide the following practical guidelines to simplify the setup process.
Step-by-step configuration:
Segment setup. Divide each trajectory into 3 segments. For trajectory of length T, set the context length K T//3.
27

---

Published as a conference paper at ICLR 2026
1.0
0.8
06
Succes Rate
0.6
0.4
Data for -5
MRV-CA-1
MRVALSTM
DT
MRV-G
0.2
0.0
100
200
300
400
500
600
700
800
900
Inference corridor length
Figure 17: Results of RATE inference with different MRV configurations on the T-Maze environment. Training was performed with the number of segments N and context length K 30, i.e. on trajectories of length <. 150. MRV-CA-2 is the final MRV configuration that is used throughout the work and is designated as MRV.
RATE/T-Maze: K ablation
RATE/T-Maze: Transfomerr ablation
1.0
1.0
Succes Rete
Succes Rate
0.4
T--t K-
+
k-1, K-30
1-110 K-d
0.4
Tez Keds
+
tws, Ke18
Panaaiant Agen
0.2
0.2
DT
Persistent Agent
wio FFN in Decoder
100
200
300
400
500
600
700
800
900
0.0
100
200
300
400
500
600
700 800 900
Test cornidr langth
Test corridor length
Ablation of segment size K and segment count N with fixed effective context in the decoder. Kett Kt x NJ 90.
Figure 18:
Figure 19: Ablation of ford-forward block usage
2. Memory configuration. Use the following default parameters for RATE's memory mecha- nisms:
num_nem_tokens=$
n head_cca=1
mrv_act ReLU
mem_len-
- -(3 x K + x _um__mem_______ nem_tckens) x for dense reward environments (e.g., ViZDoom-TTwoColort Minigrid-Memory)
0 for sparse reward environments (e.g., T-Maze)
3. Transformer core. Set the standard architecture parameters (number of layers, atten- tion heads, embedding dimension, etc.) based on the task complexity and computational constraints.
4. Memory tuning. After adjust, fine-tune memory-related parameters if needed (e.g., num sam__mem_token mem mem_len, dropout rates).
This configuration provides strong default setup and has consistently performed well across all evaluated tasks.
28

---

Published as a conference paper at ICLR 2026
sotist Mazor Marter Semrer locts 2000
BATERTMeee: Number of burnloreer attentin reads sblatoo
RATE Vede Terafomer fests ation
+40
peminart..p Agent
230
668
lingh
Figure 20: Results of ablation by the number of layers of the RATE model in T-Maze environ- ment.
Figure 21: Results of ablation by the number of attention heads of the RATE model in T-Maze environment.
Figure 22: Results of ablation by the features sizes of the RATE model T-Maze environment.
Table 12: Comparison of RATE and DT Model Parameters. RATE has 1.0-7.7% less parameters compared to DT due to the fact that RATE does not use feed-forward network in the transformer decoder by default.
Environment
RATE
DT
difr,
T-Maze
1,723,840
1,775,488
-2.91
ViZDoom- ViZDoom-Two-Colod Two-Colors
4,537,504
4,672.032
-2.88
Minigrid-Ment Memory
2,000,864
2,051,872
-2.49
Memory Maze
1,639,840
1,673,696
-2.02
POPGym
6,760,192
6,827,008
-0.98
I TECHNICAL DETAILS
Table 12 and Table 13 shows the technical parameters of the training models. Note that the difference between the number of DT and RATE parameters small. Training RATE with trajectory splitting into N segments allows ~ N smaller GPU memory size usage than for DT. The training was conducted using a single NVIDIA A100 80 Gb graphics card.
Table 13: Computational efficiency comparison between RATE and DT models across different memory-intent intensive environments. We report three key metrics: (I) training time per epoch (meantstd, in seconds), (2) inference latency per step (meanisem, in milliseconds), and (3) GPU memory footprint (in MiB). Lower values indicate better efficiency.
RATE
DT
Environment
Train (s)
Test (ms)
Size (MiB)
Train (s)
Test (ms)
Size (MiB)
T-Maze
16.17=2.75
7.20++.31
3,148
95.75+0.49
10.69+0.11 14
8,608
ViZDoom- Taoom------olorss
77.44+3.56
10.35=0.52
7,750
68.1811.56
10.45=0.41
14,046
Minigrid-Memory
33.74=2.65
9.94+2.24
4,102
16.77=1.37
10.43+2.86
4,298
Memory Maze
110.26+2.97
38.9820.62
6,638
82.6941.56
40.3620.46
10,386
POPGym
3.37+0.25
8.91+0.37
5,948
3.64+0.53
8. 8.98+0.32
10,696
29


Reviews:
[{"review_id": "JosbuqaS5R", "reviewer": "Reviewer_o12m", "summary": "The authors focus on solving POMDPs using transformers, in offline RL scenarios. They propose the Recurrent Action Transformer with Memory (RATE), which implements various modules described below. In particular, the authors focus on efficiently processing long contexts.\n\nThe proposed Memory Retention Valve (MRV) allows information to flow between sliding attention windows. Normal transformers are $O(n^2)$ space complexity, where $n$ is the length of the trajectory. However, by integrating recurrence through the memory retention valve, one can decouple $n$ from the trajectory length. \n\nThe proposed method breaks a trajectory into fixed-length, non-overlapping segments. Memory tokens $M$ span these fixed-length segments, enabling context beyond the sliding window. Each segment goes through self attention to produce embedding tokens $S$. I will admit it is not entirely clear to me what happens next. It appears that the current memory tokens $M$ for the segment $n$ are concatenated twice to either end of the current embedding tokens $S$ (Algorithm 1, line 6). But these are all initialized to 0 (Algorithm 1 line 3).\n\nThe authors go on to prove how quickly inputs vanish through the MRV. To do so, they make the assumptions that the rows of each memory token $M$ are $ell_2$-normalized. I will admit the proof was a bit difficult to follow and I did not check it thoroughly.\n\nFinally, the authors compare their method to more standard models such as the Decision Transformer, Transformer XL, Decision Mamba, GRU, and LSTM. They perform significant evaluation and analysis on T-Maze and related tasks, finding that the proposed method achieves state of the art results. They report similar results on Minigrid and POPGym baselines. \n\nFinally, they perform ablations on their model by replacing various memory model components with random noise. They perform an upper-bound estimate by turning POMDPs into MDPs and training the Decision Transformers on them. Finally, they provide ablation for various gating mechanisms for the MRV.", "questions": "> unlike static recurrence, it preserves sparse, long-range information. \n\nCan you explain what static recurrence means here?\n\n> prevents catastrophic overwriting of memory by preserving alignment between consecutive memory states.\n\nSimilarly, what is \"alignment\" referring to here?\n\nIn algorithm 1:\n\n> 6: $\\mathrm{concat}(M_n, S_n, M_n)$\n\nIs this a bug? Should it be $\\mathrm{concat}(M_n, S_n, M_{n+1})$?\n\n> $\\hat{a}_n \\rightarrow \\mathcal{L}(a_n, \\hat{a}_n)$\n\nI am not familiar with this notation. I know $\\leftarrow$ is assignment but what does $\\rightarrow$ mean here? \n\n> **Output:**\n\nHow are you returning from inside a for-loop at each iteration? Is this a generator?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 3, "presentation": 1, "contribution": 3, "strengths": "Transformers struggle with long contexts, and in RL this is especially painful due to poor sample efficiency. Introducing recurrence into a transformer is a good way to fix this issue.\n\nIn general, the authors provide a good amount of evidence for their method. They compare against other memory models and even training algorithms. Critically, they compare against TrXL which is similar to their method. They provide ablations for each part of their proposed method, and even test different gating methods for the MRV.", "weaknesses": "Although the experiments and models appear well-done, I had a lot of trouble reading this paper. I think the text is quite confusing and could use a lot of polish. Please see the comments below, this is not an exhaustive list but just some representative examples:\n- Figure 1 is too blurry and small to read, even after zooming to the maximum allowed\n- Experimental setup and results in the approach (section 3, lines 131-139) seem misplaced\n- The text surrounding the MRV is difficult to understand (lines 141-160, algorithm 1)\n    - Using the same terms to refer to different quantities is the main confusion: MRV takes $M_n, M_{n+1}$ as input, but also outputs $M_{n+1}$\n    - Could use a figure to help the readers understand \n- Section 3.1 is jarring and also a bit difficult to read, I think it should be moved to the appendix\n- Figure 3 left figures are too small and blurry to read at max zoom\n- Figure 6 is too small and blurry to read at max zoom\n\nI cannot recommend acceptance without a substantial rewrite. That said, I want to emphasize that I think there is some good work here. I just wish it was more accessible.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7BkBnY41Ew", "reviewer": "Reviewer_tfiE", "summary": "The paper proposes augmenting transformers with a memory mechanism to tackle the problem of training RL agents in partially observable environments, where the transformer-based policy suffers from trajectory length. Experiments across various benchmark sets showed that the proposed model improves performance in settings that require memory to solve.", "questions": "Minors: The algorithm 1 can be clearer in term of describing the outputs. The output of the algorithm 1 is inside the for loop.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The proposed mechanism, Memory Retention Valve (MRV), is sound and empirically proven to be effective.\n- Figure 2 showed a clear visualisation of RATE and DT when processing information in the task T-Maze.\n- Experiments compared the approach with sufficient baselines across different types, such as transformer-based models, recurrent models, and state space models.\n- Ablation study, especially the first part with noise mem. tokens and noise caching empirically showed the contributions of components in RATE. Different architectural choices are also analysed in the last ablation study section.", "weaknesses": "See in the questions section.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "1zMEAcST2D", "reviewer": "Reviewer_xQWj", "summary": "This paper introduces RATE (Recurrent Action Transformer with Memory), a transformer-based architecture for offline reinforcement learning (RL) that explicitly integrates memory mechanisms to improve performance in partially observable and long-horizon tasks. RATE processes trajectories in segments and passes information between them through memory embeddings, cached hidden states, and a novel Memory Retention Valve (MRV)—a cross-attention module controlling what information to retain or overwrite. The authors provide a theoretical bound guaranteeing partial memory preservation between updates and show strong results across memory-intensive environments (T-Maze, ViZDoom-Two-Colors, Memory Maze, Minigrid-Memory, POPGym) while remaining competitive on standard Atari and MuJoCo benchmarks. Extensive ablations isolate the contributions of MRV and memory embeddings and demonstrate RATE’s robustness in sparse, long-horizon POMDPs.", "questions": "1. How large is the compute/memory overhead of MRV and recurrent caching relative to standard DT or RMT during training and inference?\n\n2. Can MRV’s gating dynamics be visualized—e.g., which segments’ information it chooses to retain or forget?\n\n3. How sensitive is performance to the number of memory tokens m and segment length K?\n\n4. How does RATE behave under noisy or corrupted observation sequences beyond the artificial noise ablation?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- Strong motivation and clear contribution. The paper identifies a concrete weakness of Decision Transformer (DT)—limited context for long-term dependencies—and addresses it with a principled, integrated memory mechanism rather than ad hoc recurrence.\n\n- Theoretical component. Theorem 1 provides a nontrivial formal result bounding memory preservation under α-alignment, adding rigor beyond empirical evidence.\n\n- Comprehensive experimental coverage. RATE is evaluated on a wide spectrum of domains, from classic MDPs to highly memory-dependent tasks, showing consistent advantages and strong extrapolation (up to 9600 steps in T-Maze).", "weaknesses": "- Incremental novelty. RATE mainly extends prior memory-augmented transformers (RMT, TrXL, Compressive Transformer) by combining existing mechanisms; MRV, while effective, is a variant of cross-attention gating rather than a fundamentally new idea.\n\n- Theoretical analysis scope. The preservation theorem bounds norm differences but does not connect to RL performance or convergence guarantees, limiting theoretical depth.\n\n- Computational and scaling costs. The paper does not quantify the runtime or memory overhead of segment recurrence and MRV compared with DT or RMT, which is crucial for practical adoption.\n\n- Conceptual overlap with existing work. RATE’s motivation and mechanism strongly resemble RMT (2022) with an added gating step; the paper could better articulate its distinct advantages and design rationale.\n\n- Lack of reference: It is noted that this paper compare their method with \"Decision Mamba\". However, there is another existing work which also adopts mamba as its backbone [1].\n\n[1] Lv, Qi, et al. \"Decision mamba: A multi-grained state space model with self-evolution regularization for offline RL.\" Advances in Neural Information Processing Systems 37 (2024): 22827-22849.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Gh9zDPPEsz", "reviewer": "Reviewer_kH4W", "summary": "This paper introduces the Recurrent Action Transformer with Memory (RATE), a transformer-based architecture for offline reinforcement learning (RL) in partially observable and memory-intensive environments. RATE incorporates three memory mechanisms includinlearned memory embeddings, recurrent caching of hidden states, and the Memory Retention Valve (MRV), a cross-attention module designed for selective retention of key information across long sequences. The method is evaluated on a diverse suite of tasks, including ViZDoom-Two-Colors, T-Maze, Minigrid-Memory, Memory Maze, POPGym, as well as Atari and MuJoCo benchmarks. Experimental results and ablation studies are provided, alongside theoretical analysis on memory preservation by the MRV.", "questions": "-\tThe Decision Transformer (DT) underperformance on memory-dependent tasks such as T-Maze may stem not from memory limitations, but rather from its lack of trajectory stitching ability; that is, its inability to combine underperforming trajectories into coherent long-horizon trajectories. In cases where underperforming trajectories are more frequent than successful ones, DT will often follow the underperforming paths, further degrading performance.  The results would be more convincing if the authors compared against (applying on top of) more advanced DT architectures, such as Reinformer [1] or Elastic DT (EDT) [2], which explicitly address the trajectory stitching issue.\n-\tIn Table 13, RATE includes extra components such as MRV and memory embeddings, yet it often shows lower training time, inference latency, and GPU memory usage compared to DT. Could the authors clarify why adding these components does not increase computational cost, and how RATE achieves higher efficiency despite the additional modules?\n\n[1] Zhuang, Z., Peng, D., Liu, J., Zhang, Z., & Wang, D. (2024). Reinformer: Max-return sequence modeling for offline rl. arXiv preprint arXiv:2405.08740.\n\n[2] Wu, Y. H., Wang, X., & Hamaya, M. (2023). Elastic decision transformer. Advances in neural information processing systems, 36, 18532-18550.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 1, "strengths": "-\tThe paper presents a well-motivated and thoughtfully engineered hybrid memory mechanism.\n-\tThe empirical study is extensive, covering a comprehensive battery of memory-intensive RL tasks (ViZDoom-Two-Colors, T-Maze, Minigrid-Memory, Memory Maze, POPGym), and standard MDPs (Atari, MuJoCo).\n-\tThe paper provides a theoretical analysis characterizing", "weaknesses": "- Although the motivation is clear, it seems like the performance of RATE relied heavily on previously introduced memory embeddings and cached hidden states, as shown in the ablation study RQ1. This raises a question about the degree of novelty and contribution of the MRV component itself versus inherited advantages from prior architectures.  \n- If the sole contribution of the paper is the Memory Retention Valve (MRV), the manuscript should state this explicitly and carefully articulate the hypothesis for why MRV produces the reported gains. Concretely, the authors should explain the mechanisms by which MRV improves performance (e.g., how it interacts with learned memory embeddings and cached hidden states to prevent overwriting or improve retrieval)\n- Minor: The paper is quite dense, which makes some figures difficult to read due to the small font and cramped layout. While the inclusion of many experiments is commendable, improving figure clarity would help accessibility. In Figure 5, the y-axis label at 0.0 overlaps with the x-axis tick at 〖10〗^1. Slightly adjusting the axis limits or label offsets may resolve this issue.\n- Minor: The authors could include one or two sentences explaining what is meant by “cached hidden states”, as this is not a widely standardized term. Providing a brief description in the main text (rather than requiring readers to refer back to the original paper) would improve readability and self-containment.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)