Recurrent Action Transformer with Memory

Thank you for your feedback. We have responded to your questions and comments and highlighted the changes in the text of the paper in blue.

W1. Sequence modeling. Thank you for your thoughtful concern. We would like to clarify several points:

Our work builds upon the sequential modeling framework established in the DT paper, which provides comprehensive experiments demonstrating the benefits of transformer architectures over standard offline RL approaches. The DT paper includes experiments with the Key-to-Door environment, which is memory-intensive as it requires the agent to memorize and utilize specific information after a long sequence of steps.

Our primary contribution focuses on enhancing existing transformer architectures with effective memory mechanisms rather than comparing different learning paradigms (Offline RL vs. BC). The advantage of using sequential modeling with transformers over LSTM architectures is empirically supported by our experiments. Specifically, we demonstrate that LSTMs struggle to propagate gradients beyond approximately 100 steps in a simple T-Maze task, even when trained on optimal trajectories. This limitation is consistent with findings from the NTM paper [1]. Consequently, when action decision at step 100 depend on information from step 0, LSTM-based approaches fail, while our RATE algorithm successfully maintains and utilizes this long-term information and works in both simple and memory-intensive environments.

This fundamental architectural advantage in handling long-term dependencies justifies our focus on improving transformer-based sequential modeling rather than exploring alternative learning frameworks.

W2. ViZDoom-Two-Colors bias. Memory mechanisms can combat such biases in the dataset (see Fig 3), so we felt it was not necessary to obtain a dataset balanced not only by pillar colors but also by object colors. In addition, when conducting experiments on T-Maze, the dataset for which is perfectly balanced, we obtained that for out-of-context inference DT, unlike baselines with memory, always acts as a persistent agent (the agent turns only up or only down).

W3. FFN in Decoder. For clarity, we performed additional ablation studies on the use of FFNs in the decoders of the considered transformer baselines (see Fig 20). Disabling FFN allows to obtain better results only for RATE, for other baselines the results deteriorate. Therefore, for fair comparison we disable FFN only for RATE and use FFN for DT, RMT and TrXL.

W4. Memory Maze data collection. In our dataset, the average total reward per trajectory is <r> ~ 7 with episode length T = 1000. Thus, the agent on average receives a reward +1 for 143 steps. Since we train on three consecutive segments of length 30, we cover a total of 90 environment steps. Thus, to ensure that the reward signal falls within the trajectory segment used for training, we introduce the condition presented in Appendix D. Moreover, we can sample segments from an arbitrary segment for this task.

W5. Atari and MuJoCo sampling horizons. If by sampling horizons you mean the length of an episode, we have the same length as in DT. In this case, we use only a part of the trajectory during training. Thus, in the original DT paper, a subsequence of length $K=50$ for Pong and $K=30$ for Breakout, Qbert, Seaquest is randomly sampled. In our paper, we sample a subsequence of length $K=30\times 3 = 90$ for RATE. If this does not answer your comment, please clarify what you meant.

Q1. Data collection policies. We use the optimal policy only for T-Maze. In other memory-intensive tasks, the policies used for trajectories collection are far from optimal.

Since we work in a sequential modeling paradigm, data quality is critical to us, so we use optimal or near optimal policies to collect the data (see Tab 4).

However, as mentioned earlier about ViZDoom, the memory mechanisms allow us to deal with bias in the data, as can be seen in Fig. 3 (c, d), where at t $\in$ [0, 90) (i.e. pillar inside the context) all baselines perform almost the same, but at t $\in$ [90, 189) only the DT result slips, since the pillar is no longer available in the context and the agent starts reproducing bias in the data.

Q2. BPTT. During training, we use BPTT to propagate gradients through memory tokens across previous segments, which gradients are not stopped between segments. The number of segments for backpropagation is equal to $N$, and ablation results for this hyperparameter are shown in Fig. 15.

Q3. TrXL performance in T-Maze. On T-Maze, TrXL performs worse than other memory-enhanced baselines because its memory mechanism (caching previous hidden states) mixes representations, preventing the explicit extraction of 1 bit of information. Ablation studies of RATE show that for sparse tasks like T-Maze, this memory mechanism degrades performance.

[1] Graves, A. (2014). Neural Turing Machines. arXiv preprint arXiv:1410.5401.

We are pleased that the reviewer noted the good performance of RATE and the robustness of our experimental results. We have answered the reviewer's questions below. Updates are marked in blue in the text of the article.

W1. Transformers in Offline RL. The works proposing the use of transformers in Offline RLs are mainly focused on MDPs. Our study demonstrates that while transformers like DT perform well on MDPs (e.g., Atari, MuJoCo), their naive application in memory-intensive environments results in severe performance degradation (Figure 5; Tables 2, 5), even failing to address memory tasks entirely (Figures 3, 4; Table 1; Figure 9). Introducing memory mechanisms to transformers resolves this issue, though it introduces challenges, such as excessive memory updating, which we address with MRV.

Memory Retention Valve. MRV prevents important information from being “washed out” during long recurrent segment processing. As explained in Section 5 (RQ3), the need for MRV arose from experiments on ViZDoom-Two-Colors and T-Maze. These tasks required critical information to be stored in memory embeddings after the first segment and remain unchanged. However, recurrent updates in subsequent segments risked overwriting this information. To address this, we introduced a gating mechanism to control memory updates, with a cross-attention-based design proving most effective.

W2. Hyperparameters. Good performance of RATE can be achieved by tuning just two additional parameters, despite the fact that, like any transformer-based model, it has many hyperparameters: the number of memory tokens (num_mem_tokens) and the number of MRV attention heads (n_head_ca). The third key parameter, the number of cached hidden states (mem_len) can be predefined based on the type of task.

In order to make it clearer how changing the main hyperparameters of RATE affects performance, we conducted new ablation studies on the additional RATE parameters: number of segments (N), context length (K), and the use of FFN in the decoder (skip_dec_ffn), and classic transformer parameters: number of attention heads (n_heads), number of transformer layers (n_layer), transformer features size (d_model) (see updates in Appendix F and Appendix H).

Ablation studies on num_mem_tokens, mem_len, n_head_ca, N, and K are presented in Appendix F. For environments with continuous rewards (e.g., ViZDoom-Two-Colors), optimal results are achieved when mem_len = (3xK + 2x num_mem_tokens) x N, where K is context length and N is the number of segments ( K_{eff} = N \times K ). For sparse environments (e.g., T-Maze), mem_len = 0 is most effective.

Other transformer parameters, such as the n_heads, n_layers, d_model, skip_dec_ffn, can be tuned independently of memory hyperparameters. A minimal starting configuration is num_mem_tokens = 5, n_head_ca = 2, and mem_len = 0 or (3xK + 2x num_mem_tokens) x N, depending on trajectory structure (see Table 8, Figures 10, 11). Additional results are provided in Appendix H. To clarify hyperparameter tuning, we’ve added guidelines in the Appendix I.

W3. Baselines tuning. To ensure a fair comparison, we used the same base transformer hyperparameters for DT, RMT, TrXL, and RATE. For RMT, TrXL, and RATE, we tuned parameters related to memory mechanisms. In Atari and MuJoCo experiments, we used the same transformer parameters as in the original DT paper.

In turn, for DLSTM, DGRU and DMamba tuning we also tuned by basic parameters and reported the results corresponding to the best configuration (see Appendix G). At the request of Reviewer hxTt, we have posted the code and results of tuning these baselines in the recurrent_baselines/ directory of our code repository, an updated link to which is available in the abstract of updated paper.

Q1. Why could the MRV improve the performance of RATE? Recurrent segment processing in RATE aims to update memory tokens to carry important information between segments. However, when memory tokens are updated for a long time, important information can be lost from them (e.g., on T-Maze we get information about a clue on the first segment, but it can be used for decision making only on segment 30, see Fig 4, Tab 6). To control the updating of information, we need a kind of valve that will keep important information inside the memory tokens, but will not interfere with the writing of new information. We have tried different ways to implement such a valve (Tab 6, Fig 14), and came to the conclusion that the best results can be obtained if we let the model itself determine how to implement the gating of memory tokens. The cross-attention architecture is the best for this purpose, as evidenced by our experiments.

Q2. Guide how to tune RATE. We are grateful to you for a really good idea. For ease of use of our model, we have added a section to the appendix (see Appendix I) with our recommendations on how best to tune RATE for different tasks.

Thank you so much for your prompt response! We really appreciate it.

Regarding your first concern about hyperparameters - any transformer architecture, including DT, contains many hyperparameters. Our RATE model contains only 5 more hyperparameters than DT: num_mem_tokens, mem_len, n_head_ca, mrv_act, sections. In our initial ablation studies, supplemented by your recommendation, we show that of these 5 additional hyperparameters only sections, i.e. number of segments, should be configured separately, while the others already in the initial configuration allow to obtain significantly better results than just DT (see Appendix F, Appendix H).

We emphasized the simplicity of RATE tuning in an additional section in Appendix I. Thus, RATE can be tuned almost in the same way as DT, and even with coarse tuning of memory-intensive parameters, it can solve memory-intensive problems.

In your review in Strengths 2, you noted that “RATE is validated on various memory-intensive environments. The results demonstrate RATE achieves good performance on different tasks and environments.”

In section 4.1. “memory-intensive environments” we define this term: “memory-intensive environments - environments where the agent requires memory to operate successfully”. In other words, these are POMDPs where the agent's decision making in the present depends on information from the past that can be retrieved using memory.

All our work aims to show that the proposed architecture successfully solves memory-intensive environments and common tasks. Thus, Fig. 3, Fig. 4, Fig. 5, Tab. 1, Tab. 2, show the benefits of using RATE to solve memory-intensive tasks, while Tab. 3 and Tab. 4 show that RATE is not inferior to SOTA solutions in the Offline RL domain when solving classical problems.

We really appreciate your quick comments and engagement in the discussion, and we really hope we can get our RATE proposal to you.

Thank you for recognizing our algorithm’s effectiveness, especially for memory-intensive tasks requiring efficient transformers. We value your feedback and have addressed your comments. In the updated version of the text, the changes are highlighted in blue.

Q1. Cached hidden states: Cached hidden states refer to saved hidden states from previous tokens (as in TrXL [1]). This mechanism combines current hidden states with cached ones, enabling access to past information. The mem_len parameter sets how many past token states are available for processing each segment. We’ve clarified this in the paper.

Q2. Positional Encoding: In our architecture, we employ relative sinusoidal positional embeddings, as in TrXL [1]. This design enables the model to incorporate previously computed hidden states when attending to memory. The approach is invariant to sequence shifts and supports generalization to sequences of varying lengths, which is important in RL.

However, we do not use positional encoding within the MRV block for the following reasons:

1. Memory tokens serve as an abstract repository of information, where their content, rather than their position, is of primary importance.

2. The number of memory tokens is typically small and fixed, making positional encoding less relevant.

3. The primary function of the MRV mechanism is to update the memory state through interactions among memory tokens, rather than rely on positional relationships.

Q3. ERNIE-Docs baseline: Thanks for the suggestion! We're eager to see how ERNIE-Docs' memory mechanisms perform in memory-intensive settings. Since there's no existing implementation for Offline RL, we'll work on it ourselves and aim to share results before the rebuttal period ends.

Q4. T-Maze and recurrent baselines: For the T-Maze environment, we everywhere specify the episode duration T through a tight bound on the corridor length L:T=L+2. Because of this constraint, the agent is guaranteed to receive a reward of 0 if it performs a single action outside the optimal policy. We added the performance of a random agent to Table 1 in the attached paper text.

In the T-Maze experiment shown in Table 1, we test the ability of the models to train on trajectories of a given length. This table for T-Maze shows results for models trained on data with trajectories of length T=3-90 (i.e., with corridor lengths L=1-88) and failed only on trajectories of length T=90 (L=88). It is important to note that all considered baselines process a triplet (R,o,a), i.e., 3 tokens, at each step of the environment, and thus trajectories of length T=90 have a length of 3x90=270 tokens.

The results show recurrent baselines fail to retain key information over 270-token episodes, while transformers excel at linking sparse information via attention.

To make the results of this experiment clearer, we performed additional experiments with baseline learning on trajectories of length T=9 and T=30 (with context K=9 and K=30, respectively):

The new results confirm the ability of recurrent baselines to capture sparse dependencies on small time horizons (T=30 corresponds to 90 tokens) and their inability to learn on longer sequences.

In the case where only observations instead of triplets are fed at each step, the situation changes slightly:

Here, unlike in the previous table, T=90 corresponds to 90 tokens rather than 270 tokens. In turn, the recurrent baselines also show an inability to solve the task at T=270 (as in the previous table, when T=90 corresponded to 270 tokens), indicating directly the limited time window available for recurrent baselines to memorize during training.

Q5. Hyperparameters: We fully agree with you, and therefore we performed additional ablation studies on the number of transformer layers, the number of attention heads, the model dimensionality, the use of FFN in the decoder, and the segment lengths under a fixed effective context. We have added the results to the updated text of our paper in Section F. Additional Ablation Studies and Section H. Transformer Ablation Studies.

The initial results of the ablation studies in Section F referred only to the new RATE parameters, such as number of memory tokens (num_mem_tokens), number of cached hidden states of previous tokens (mem_len), number of MRV's attention heads (n_head_ca), and MRV activation function (mrv_act). We answered about cached hidden states in Q1.

Q6. Attention maps: Thanks for the tip! We have added attention maps for RATE and DT to the Appendix, Section K.

[1] Dai, Z. (2019). Transformer-xl: Attentive language models beyond a fixed-length context. arXiv preprint arXiv:1901.02860.

Thank you for your valuable feedback. We have addressed your comments in the revised text, highlighting the changes in blue, and look forward to further discussion.

Reproducibility (W1-2, Q2)

We agree that reproducibility is crucial in RL, and you are right to emphasize its importance. While we cited the open-source implementations of DLSTM [1] and DMamba [2], we’ve updated our repository with a recurrent_baselines/ section, including code for all recurrent baselines and instructions for running them in T-Maze and ViZDoom-Two-Colors. The anonymized repository link can be found in the abstract of the updated version of the paper.

Environments (W3)

In our work, we consider not only T-Maze (which is a common memory test in RL and has been used in numerous papers, e.g. [3-6], to name a few) and ViZDoom-Two-Colors, but also other challenging environments such as Minigrid-Memory and Memory Maze, on which we perform extensive experiments (Fig 5, Tab 2).

Even the T-Maze environment is more complex than it seems, as the agent receives unique information only at the first and the last steps of an episode. The random agent gets SR= 0 starting from the corridor length L = 5. Since the episode length is defined as T=L+2 and any deviation from the optimal policy leads to failure (we have added the random agent results to Tab 1 and Fig 4). SR=0.5 corresponds to a persistent agent that successfully reaches the junction, but then makes a turn up or down randomly.

Speaking of episode length, we represent each step of the environment as a triplet of tokens (R, o, a), so that when processing a context of length K, we actually process 3K tokens: In T-Maze, we trained agents on trajectories of length L = 3 - 90 and validated on trajectories of length L = 3 - 900. Thus, at L = 900, the agent demonstrates information memorization at 900 steps and successfully processes 2700 tokens, well in excess of 100 steps. In ViZDoom-Two-Colors, the RATE agent trained on the first 90 steps of the environment then operates for ~1700 steps (episode timeout = 2100), processing ~1700 × 3 = 5100 tokens — again, far exceeding 100 steps.

LSTM, GRU, Mamba, and Transformer (W1, Q1-2)

In Tab 1, we show results in the T-Maze environment for baselines trained on trajectories of 90 steps and validated on corridors of length 90. As noted earlier, 90 steps correspond to 270 tokens.

To prove that all baselines were properly trained, we included the code for LSTM, GRU, and Mamba and conducted a hyperparameter sweep (see App G - Recurrent Baselines), as per your request. To confirm that these models can solve the T-Maze task in principle but fail on long sequences, we added results for T=K={9, 30, 90}. Results for the best configurations are shown below:

This demonstrates that recurrent baselines can solve the T-Maze task when trained on data with moderate corridor lengths but fail to retain the clue information for longer lengths, unlike a Transformer. This is because the transformer's attention mechanism can effectively capture dependencies in highly sparse data, which recurrent models cannot. DT achieves SR = 0.5 for T>K for any K, while recurrent networks can achieve SR > 0.5 in this setting. RATE combines the strengths of transformers (direct access to information in context) and recurrent networks (hidden states for information retrieval).

Novelty (W4, Q3)

Sequential updates in our model’s memory tokens can cause information loss, as shown in the T-Maze task (Tab 6). To address this, we introduce a memory valve (MRV) to regulate memory token updates. New valve mechanism (MRV) is exactly one of the contributions of our work. If you are referring to GTrXL [7] as "recurrent transformers with gating mechanisms", we tested their approach (MRV-G in our paper), which underperformed compared to our method (Tab 6). If not, we would appreciate references to other works on recurrent transformers with gating mechanisms for Offline RL tasks that you mentioned to support our future research.

[1] Siebenborn, M. et al. How crucial is transformer in decision transformer?NeurIPS 2022

[2] Ota, T. Decision mamba: Reinforcement learning via sequence modeling with selective state spaces.arXiv:2403.19925.

[3] Ni, T. et al. When do transformers shine in rl? decoupling memory from credit assignment.ICLR 2024

[4] Grigsby, J. et al. Amago: Scalable in-context reinforcement learning for adaptive agents.ICLR 2024

[5] Esslinger, K. et al.. Deep transformer q-networks for partially observable reinforcement learning. arXiv:2206.01078.

[6] Pramanik, S. et al. AGaLiTe: Approximate Gated Linear Transformers for Online Reinforcement Learning.TMLR 2024

[7] Parisotto, E. et al. Stabilizing transformers for reinforcement learning.ICML 2020