Title: Recurrent Action Transformer with Memory
Venue: NeurIPS 2025
Accepted: no
Keywords: Offline RL

Memory

Transformers

POMDP
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nWhile being widely used across many domains, transformer architectures have a core limitation: they require feeding in the whole context. This limitation makes it difficult to deal with very long-horizon tasks where it's computationally infeasible to feed in the whole sequence length. \nIn offline RL tasks where the offline trajectories are too long to be fed to the transformer model, the model will fail to retain long-horizon information beyond its context length. \n\nThe paper addresses this limitation of transformers by introducing a new architecture named Recurrent Action Transformers (RATE), which has a sort of recurrent memory that caches previous information. Thus, allows for retrieving information beyond the current sequence length. RATE also has a memory retention valve (MRV), which controls the flow of information between two consecutive hidden states to ensure that past information can be retained. \n\nThe paper provides an extensive set of experiments including some memory-specific tasks such as T-Maze and VizDoom and other general tasks including 4 atari games some mujoco environments. The experiments also have comparisons with a wide range of baselines: transformer-based, RNN-based, and some classic offline RL baselines. The results suggest that RATE provides a competitive performance compared to other baselines.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths:\n- The RATE algorithm is novel and interesting. \n- The algorithm presentation is very clear. \n- The theorem statement and derivations are clearly explained, and as far as I can check, the derivation is correct.\n- The experiments cover a wide range of environments and baselines. \n- The paper is well-written and easy to follow.\n### Weaknesses:\n- While the original (and most) transformer architectures have the quadratic complexity in the context length, several works have addressed this problem before and proposed architectures that have linear complexity (such as the linear transformers work). The introduction is missing references to that line of research, with the discussion only focused on the quadratic complexity. For example, see [1,3] for linear transformers, and [2] for using linear transformers in RL. \n- One of the motivations of the paper was that the quadratic complexity of transformers is computationally expensive. However, the proposed Memory Retention Valve (MRV) depends on the softmax operation, so the quadratic complexity still exists. While the quadratic complexity is now dependent on m (the memory dimension), i assume that for problems with very long horizons, one would need to have a large value of m so the MRV algorithm won’t be efficient in such cases. \n- The proposed algorithm doesn’t seem to improve much over other baselines (Figure 3 and 4).\n\n[1] Katharopoulos, Angelos, et al. \"Transformers are rnns: Fast autoregressive transformers with linear attention.\" International conference on machine learning. PMLR, 2020.\n\n[2] Pramanik, Subhojeet, et al. \"AGaLiTe: Approximate Gated Linear Transformers for Online Reinforcement Learning.\" Transactions on Machine Learning Research.\n\n[3] Irie et al. NeurIPS 2021. Going Beyond Linear Transformers with Recurrent Fast Weight Programmers.\n\nQUESTIONS:\n- One of the main motivations was overcoming the quadratic computational complexity of transformers. Did you have any measures of the computational complexity of RATE compared to other architectures (TrXL, for example)? Is RATE faster? Uses less FLOPs?\n\n- For problems where there is a long-term memory dependency, how big does one need to set m for the M_n matrix? And how does the computational complexity of RATE depend on m? Is it less expensive to set m to a large number rather than just feeding the whole context to a regular transformer? \n\n- In Figure 2, what do the x-axis and the y-axis refer to? Could you explain how to interpret the attention weights here? What does each little box in the grid refer to?\n\n- How accurate is the assumption about the alpha-alignment? Did you verify it empirically? \n\n- In algorithm 1, are the action/observation/return-to-go encoders MLP layers? \n- It seems that RATE still has some issues with long-term dependence in Figure 5. Is there a reason why this is the case? Does increasing the memory dimension help with this issue?\n- I think it is important to add to RMT to Figure 6 as well, since it’s the main competitor to RATE and has a very similar performance. I know that the RMT figure is in the appendix, but since it looks very similar to RATE and it is a strong baseline, I think it should be in the main paper too, or at least some acknowledgment that it performs similarly. \n- Why is RMT missing from Table 2?\n- In RQ1, can you explain the reasoning behind replacing the memory with a noisy memory? Why can’t we just evaluate without the MRV component?\n\nI am willing to increase my score once these questions are clarified.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n- This paper proposes a transformer based recurrent algorithm designed for offline RL called RATE, which applies memory embeddings, caching of hidden states and a Memory Retention Valve architecture. This helps tackle the problem that decision transformers scale quadratically with sequence length. The MRV is a component that helps to filter the memory that is passed on to the next segment, preventing error accumulation.\n- Authors conduct experiments on a number of RL environments (VIzDoom, T-Maze, Minigrid etc.) and RATE shows ability to memorize information during “memory intensive tasks”. Other experiments are conducted on standard Atari and MuJoCo tasks in which results are somewhat mixed.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n- Purely based on the architecture, which is a combination of a transformer and an RNN, it makes sense that RATE would be able to learn information from earlier in the trajectory.\n- The paper does exhaustive experiments against a large number of baselines.\n\n**Weaknesses**\n\n- The writing of the paper could be improved quite a bit. There is very little description/motivation/background in some places: A few examples: 1) It’s not clear to me why M_n is concatenated both at the beginning and at the end of S_n. 2) It’s not very clear what Figure 2 is trying to convey – it doesn’t add much to the text and leaves the reader wondering about what is supposed to be looked at. 3) Theorem 1 seems out of the blue and not well motivated nor referenced. 4) The key contribution seems to be the MRV module since the Recurrent Memory Transformer \\[1\\] seems to have the same basic idea – where the contribution starts and ends is not articulated very clearly. 5) It’s not clear why the baselines change from experiment to experiment.\n\n- The results show improved memorization for “memory intensive tasks” but less so for “reactive tasks”. Thus, the setting for which this algorithm outperforms baselines is quite narrow – specifically, it outperforms baselines in the POMDP + Offline RL setting with “memory intensive” tasks. Also, I am not sure I agree with the characterizations “memory intensive tasks” vs “reactive tasks” – they don’t appear to be well-defined. If state transitions aren’t Markov, I would expect “memory” to be involved in the optimal solution one way or another  – but the author appears to be referring to distinct clearly definable events (such as remembering which way to turn in a maze) by “memory intensive tasks”. As far as I know, in these enviornments, the memory items happen at one specific point in the trajectory, rather than a long trajectory of things to remember. This also seems to be a very limited setting.\n\n- Overall, while the environments are relatively toy, which is generally the case for POMDP papers, the experiments and baselines are quite exhaustive. My score weighs the fact that the experiments are exhaustive, while the algorithmic improvement occurs for a narrow setting and the results are mixed for standard deepRL tasks. Furthermore, the writing and the disclosure of this paper’s relationship to existing works could be substantially improved.\n\n[1]: https://arxiv.org/abs/2207.06881\n\nQUESTIONS:\n- Can you compare your paper’s contributions vs. the Recurrent Memory Transformer paper (Bulatov et. al.)?\n- Why are the reported baselines for each table in Table 1, Table 2 and Table3 different? Could you clarify?\n\nLIMITATIONS:\nThe setting in which RATE is necessary is very limited (offline RL and very long horizon partially observable environments).", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces **RATE**, a memory‑augmented Transformer architecture for offline RL. RATE combines learned memory embeddings, a recurrent state cache, and a newly proposed memory‑retention valve (MRV) block to preserve long‑term information efficiently and stably. The approach directly targets tasks with pronounced long‑term memory requirements, filling an important gap in current RL research. Empirical results demonstrate strong performance across a broad set of benchmarks, and the ablation study convincingly isolates the contribution of each key component.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n- The paper is generally well written and provides a solid mathematical justification for the memory‑preserving component\n- The experimental section offers a thorough comparison with strong baselines, including recent SOTA methods. Reported results—whether the new model surpasses or merely matches the baselines—appear consistent and believable.\n- The ablation study effectively demonstrates the contribution of each key module.\n\n# Weakness and Questions\n\n- I understand that Section 3.1 aims to justify the memory‑preserving mechanism, and I grasp its main message. However, several questions arise:\n\n    a) Lines 141–152 are somewhat unclear. Specifically, I do not follow the claim in line 149\n\n    >  there exists a row $V_j$ for which: $⟨V_j W_M , M_{n,i} ⟩ ≥ \\alpha $.\n\n    What exactly is $V_j$? Why should this condition hold empirically?\n\n    b) I am unsure about the broader motivation for Theorem 1 (On memory loss bounds). Beyond guaranteeing that the model retains a fraction of the memory, what additional insight does the proof offer—e.g., guidance on choosing the memory length $m$?\n\n    c) If $m$ is large, the fraction $1-\\sqrt{2(1-\\frac{\\alpha}{m})}$ becomes negative (since $m$ is the number of memory tokens and $\\alpha\\in(0,1]$). Because the Frobenius norm must be non‑negative, the right‑hand side of Eq. (2) would then be invalid.\n\n- Some claims in Section 3 need stronger support—either by citing prior work or by linking directly to ablation results. Otherwise, the reader cannot easily verify them.\nExample: line 127:\n    > Naively forwarding memory embeddings leads to error accumulation or overwriting of relevant information. \n\n    line 130\n    > Unlike static recurrence, it preserves sparse, long-range information\n\n- What causes the imbalance in the ViZDoom‑Two‑Colors task? Why does it harm DT and TrXL but not RATE? Please clarify the task setup; I could not find these details in the paper.\n\n# Suggestions to the writing.\n- Consider moving the proof in Section 3.1 to the appendix and retaining only the theorem in the main text; the full derivation feels cumbersome in the main body.\n- Expand Section 6 with a direct comparison between RATE and the baselines to highlight the proposed method’s unique contributions.\n- Place Section 6 before the experiments so readers first understand each baseline and why it might fail in subsequent tests.\n\nQUESTIONS:\nIf the questions in the Weaknesses section can be addressed well, I am willing to increase my score to 5.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-01"}