When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment

Reinforcement learning (RL) algorithms face two distinct challenges: learning effective representations of past and present observations, and determining how actions influence future returns. Both challenges involve modeling long-term dependencies. The Transformer architecture has been very successful to solve problems that involve long-term dependencies, including in the RL domain. However, the underlying reason for the strong performance of Transformer-based RL methods remains unclear: is it because they learn effective memory, or because they perform effective credit assignment? After introducing formal definitions of memory length and credit assignment length, we design simple configurable tasks to measure these distinct quantities. Our empirical results reveal that Transformers can enhance the memory capability of RL algorithms, scaling up to tasks that require memorizing observations $1500$ steps ago. However, Transformers do not improve long-term credit assignment. In summary, our results provide an explanation for the success of Transformers in RL, while also highlighting an important area for future research and benchmark design. Our code is open-sourced at https://github.com/twni2016/Memory-RL

Paper

Similar papers

Peer review

Reviewer 2AdR6/10 · confidence 5/52023-07-04

Summary

The paper proposes a set of metrics to measure and isolate the memory dependency of partially observable environments. The authors illustrate that some currently existing benchmarks do not sufficiently isolate the memory of an agent. Based on their analysis they propose two versions of a T-Maze environment, one of which appropriately isolates memory. Further, the authors investigate the effect of memory architecture on tasks that require long- and short-term memory. Transformer-based policies outperform recurrent policies on tasks that require long-term memory dependencies, while there seems to be no benefit on tasks that require short-term memory dependencies.

Strengths

**Significance:** I fully agree with the authors that prior works often conflate different effects and do not sufficiently isolate the memory component. Therefore metrics that help disentangle memory from other environmental effects are crucial. **Originality:** The authors propose novel metrics that evaluate and isolate different aspects of POMDPs in RL, such as credit assignment and long-term memory. The proposed metrics provide an important measure on the effect of memory and will be useful to facilitate future benchmark design and evaluation of new algorithms. **Quality:** Quality of theoretical contributions is high.

Weaknesses

**Scalability:** I credit the authors for mentioning the complexity for computing c^M as limitation, but it seems that the computation of other metrics is also limited. For example, l_{mem}(\pi^\asterisk) requires access to the optimal policy. For the minimalistic TMazes it is straightforward to obtain this policy, how is it obtained for more complex tasks such as Psychlab? Does it involve human experts? How is it computed for procedurally generated environments, where the optimal policy varies between levels? How would this scale to more complex environments, given that human demonstrations may not necessarily be optimal? **Relevance of proposed environments:** The authors propose two environments, namely Passive T-Maze and Active T-Maze to disentangle credit assignment from memory. The Minigrid benchmark suite already provides a T-Maze environment that should exhibit the same characteristics as the proposed Active T-Maze, namely MiniGrid-Memory [1]. Passive T-Maze could be useful, but it is very minimalistic and only evaluates whether a small bit of information can be stored and carried across long timespans in an agent's memory. There is no mention of the observation space, but providing, say, image-based observation for Passive T-Maze, would increase its complexity and enable application of vision-based methods. **Interpretation of results:** The authors mention that LSTM starts to falter at a memory length of 250 for Passive T-Maze. Figure 2, however, shows that LSTM can solve the same environment with a memory length of 750. This suggests that some other effects might influenced LSTM training, such as hyperparameter tuning. Maybe the authors can improve the results of LSTM by further tuning of the method? **Clarity:** The authors should clearly state that the paper focuses on long-term memory and does not consider other aspects, such as memory capacity, or robustness of memory to noise. **Missing relevant work:** There is prior work that derives theoretical bounds on approximation error of history-based methods [2]. Further, other history-based approaches such as a hierarchical Transformer memory [3] and pretrained models as a memory module in RL [4], are not cited. If applicable, adding these method in their experiments would strengthen the paper and might yield some more interesting findings. [1] Maxime Chevalier-Boisvert, Lucas Willems, and Suman Pal. Minimalistic Gridworld Environment for OpenAI Gym, 2018. Publication Title: GitHub repository. https://minigrid.farama.org/environments/minigrid/MemoryEnv/ [2] Gandharv Patil et al., On learning history based policies for controlling markov decision processes. 2022. [3] Andrew Lampinen et al., Towards mental time travel: a hierarchical memory for reinforcement learning agents. NeurIPS 2021 [4] Fabian Paischer et al., History compression via language models in reinforcement learning. ICML 2022

Questions

- The authors conclude that Transformers exhibit worse sample efficiency than LSTM on short-term memory tasks. Is there an intuition why that would be the case? On Passive Visual Match, which requires short-term memory according to the authors, a GPT reaches higher return compared to LSTM. Could it be that the drop in performance is due to different observation/action spaces (continuous vs discrete) instead of long-term vs short-term memory? - In line 304 the authors claim that the success on long-term memory tasks for GPT2 is consistent with the tendency to perform well in the supervised learning setup on large datasets. The authors draw an analogy here which lacks support. How do long-term memory tasks relate to the supervised learning setup on large datasets? Do the authors imply here, that there are only long-term dependencies present in large datasets used for supervised learning? Also, Morad et al. 2023 specifically advises against drawing comparisons between supervised learning and RL. - What is the observation space for the Active/Passive T-Mazes? - In line 270 it would be good to explain what the actual task is in Passive Visual Match, prior to discussion on the reward function. - Figure 4 left shows superiority of GPT2 over LSTM, for Passive Visual Match the difference is not as pronounced, a significance test would benefit the interpretation of these results. - Figure 5: it would be good to clarify what is meant by "optimal policies that lack (long-term) credit assignment" I assume it refers to a markovian policy, since the return on T-Maze is 0.5. - Do all figures show mean and standard deviation? This should be mentioned in the figure captions. - The last three sentences in Definition 1 could be moved to Definition 2, since they are about memory length and not context length. - Equation in line 103: I think there should be $t - l_{ctx}(\pi)+1:k$ instead of $t - l_{ctx}(\pi)+1:t$ in the subscript after the conditional independence, right? - The symbol k is reassigned repeatedly which is a bit confusing, the authors may consider using different symbols. - Will the authors make the code for reproducing their results and computing their metrics publicly available?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

Limitations

The authors have adressed limitations regarding computation of the metrics in more complex environments. However, it is not clear from the paper, how well the other metrics scale with complexity of the environment, or how the optimal policies are obtained, i.e. does that require human demonstration? In the case of procedurally generated environments (which are commonly used nowadays) would that require extensive "labeling" by humans?

Reviewer kd6F7/10 · confidence 4/52023-07-05

Summary

This work explores the impact of memory and credit assignment in decision transformer architectures, presenting several significant claims. Regarding memory, the authors establish an upper bound on the memory length required for an optimal policy. In terms of credit assignment, they provide a lower bound on the number of future steps needed. Memory length is defined as the number of steps that an action distribution depends on, representing a more recent and shorter dependency than the entire input length. This length, is defined in relation to the value function, transition, and reward. Credit assignment length of the policy is defined as the number of future steps required for a greedy action to yield a higher discounted sum of rewards compared to a non-greedy action.

Strengths

- This definitions of memory length and credit assignment length appear intuitive and sound. - Theorem 1 appears clear and mathematically sound. - A toy environment is presented to illustrate scenarios that are either heavily reliant on memory or credit assignment. This experiment is intuitive and sound. In the memory task, transformers perform optimally in long memory length tasks, while LSTMs struggle with long-term memory. However, in more complex tasks, the results are not as clear-cut. These results are significant to the DT community. - For credit assignment tasks, transformers do not outperform LSTMs and exhibit poor performance at longer credit assignment lengths. They also demonstrate lower sample efficiency compared to LSTMs. Again, this results is significant to the DT community. - Overall, the experiments are well-designed, clear, and straightforward, providing a substantial contribution to the field of decision transformer research. This work is highly novel and is the first to specifically investigate the benefits of transformers in either memory or credit assignment.

Weaknesses

- Can the authors provide any intuition on the failure of transformers and LSTM and transformers in the long term credit assignment tasks? I find it interesting that they both have the same performance trajectory. - It would be nice to see experimental evaluation of Transformers and LSTM in additional environments other than the toy tasks for better context.

Questions

Above

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

4 excellent

Limitations

Yes

Reviewer naYb8/10 · confidence 2/52023-07-06

Summary

This work aims to answer how the parametrization of policy in RL affects the performance of sparse-reward tasks that require memory and credit assignment. In particular, the authors are interested in when transformers perform well. The authors define in the POMDP settings their own criteria of how to quantify policy memory, and credit assignment. The authors go through a list of previous benchmarks to clarify which benchmarks require which, and design their own experiment which decouples the two effects. The authors conclude that while transformes do well in pure memory tasks, they do not show better credit assignment capabilities compared to LSTMs.

Strengths

- The question that was posed was quite unique and interesting. - The authors do a convincing empirical study of which settings transformers perform better in, and the conclusions are aligned with empirical intuition. - Conducted analysis of previous benchmarks is extensive, and the authors have a good experiment design.

Weaknesses

- The motivations for memory length is reasonable as a quantiative metric, but I wonder if there are other quantitative criteria for judging memory and credit assignment. For policies with memory (i.e. recurrent), the definition seems less relevant since the policy can remember parts of all the information starting from time 0. On the other hand, the credit assignment length definition was novel and I found it quite relevant.

Questions

- One of the key discussion that seem to be missing is whether or not the policy has the ability to succesfully abstract states that matter. For example, both in works along the lines of MDP compression with bisimulation, or Approximate Information States, they hypothesize a reduced-order model of the POMDP that the policy could possibly have that is sufficient for predicting rewards. - I wonder if aside from credit assignment and history, if transformers are predicting states that truly matter for the task as well (which is partly a question of credit assignment, but less of a question of "when" something mattered.) For example, suppose a pixel-domain example that controls an agent, but the background constantly has something going on (e.g. humans may play tennis in the park when ducks are flying in the background). In this case excelling at pure memory is actually disadvantageous because the ducks are completely irrelevant to the task. - In operations research there are classes of "submodular functions", and these have provably efficient greedy algorithms that are quantifiably suboptimal. I wonder how this is connected to credit assignment length.

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors have adequately addressed limitations.

Reviewer mXpQ7/10 · confidence 3/52023-07-07

Summary

In this paper the authors consider the effectiveness of transformers for solving two kinds of RL problems: credit assignment and memory. To achieve their goal the authors develop rigorous definitions for memory and credit assignment so that different RL tasks can be understood and compared in these terms. The authors then review many well known RL problems from the literature and classify them according to their definitions. Going one step further the authors then propose a new RL task where the role of memory and credit assignment can be disambiguated. Finally, the authors perform empirical experiments on many of the previously analyzed RL tasks to understand when transformers shine and when they fall short.

Strengths

The paper is accessible, rigorous, and timely. Great work.

Weaknesses

1. If I were to name anything it would be I wish the authors could have provided a little more on what they think potential solutions could be to the credit assignment problem when using transformers. The paper stands on its own without this, but the authors have clearly thought about the problem and their insights could be valuable.

Questions

1. I'm curious given the time and thought the authors have put into this problem what they think may be a solution. Can we simply slap on a specific submodule to transformers to augment their credit assignment capabilities? Or do you think we need to develop fundamentally new architectures?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

Limitations

The authors have adequately addressed limitations.

Reviewer 2AdR2023-08-11

I greatly appreciate the additional experiments on LSTMs given the limited time window for the rebuttal. Overall the response clarified most of my concerns which is why I decided to increase my rating to 6.

Reviewer naYb2023-08-14

I'd like to thank the authors for the response, my scores are the same.

Reviewer mXpQ2023-08-14

Thank you for your responses. I remain positive about the paper.

Program Chairsdecision2023-09-21

Decision

Accept (oral)

© 2026 NYSGPT2525 LLC