Title: When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment
Venue: NeurIPS 2023
Accepted: yes
Keywords: Memory-based RL

Transformers

Credit Assignment

Online RL

Model-free RL
Abstract: 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.

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

Abstract

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

1 Introduction

In recent years, Transformers (Vaswani et al., 2017; Radford et al., 2019) have achieved remarkable success in domains ranging from language modeling to computer vision. Within the RL community, there has been excitement around the idea that large models with attention architectures, such as Transformers, might enable rapid progress on challenging RL tasks. Indeed, prior works have shown that Transformer-based methods can achieve excellent results in offline RL (Chen et al., 2021; Janner et al., 2021; Lee et al., 2022), online RL (Parisotto et al., 2020; Lampinen et al., 2021; Zheng et al., 2022; Melo, 2022; Micheli et al., 2022; Robine et al., 2023), and real-world tasks (Ouyang et al., 2022; Ichter et al., 2022) (see Li et al. (2023); Agarwal et al. (2023) for the recent surveys).

However, the underlying reasons why Transformers achieve excellent results in the RL setting remain a mystery. Is it because they learn better representations of sequences, akin to their success in computer vision and NLP tasks? Alternatively, might they be internally implementing a learned algorithm, one that performs better credit assignment than previously known RL algorithms? At the core of these questions is the fact that RL, especially in partially observable tasks, requires two distinct forms of temporal reasoning: (working) memory and (temporal) credit assignment . Memory refers to the ability to recall a distant past event at the current time (Blankenship, 1938; Dempster, 1981),

while credit assignment is the ability to determine when the actions that deserve current credit occurred (Sutton, 1984). These two concepts are also loosely intertwined - learning what bits of history to remember depends on future rewards, and learning to assign current rewards to previous actions necessitates some form of (episodic) memory (Zilli and Hasselmo, 2008; Gershman and Daw, 2017).

Inspired by the distinct nature of memory and credit assignment in temporal dependencies, we aim to understand which of these two concepts Transformers address in the context of RL. Despite empirical success in prior works, directly answering the question is challenging due to two issues. First, many of the benchmarks used in prior works, such as Atari (Bellemare et al., 2013) and MuJoCo (Todorov et al., 2012), require minimal memory and short-term credit assignment, because they closely resemble MDPs and their rewards are mostly immediate. Other benchmarks, such as Key-to-Door (Hung et al., 2018; Mesnard et al., 2021), entangle medium-term memory and credit assignment. Second, there is a lack of rigorous definition of memory and credit assignment in RL (Osband et al., 2020). The absence of quantifiable measures creates ambiguity in the concept of 'long-term' often found in prior works.

In this paper, we address both issues concerning memory and credit assignment to better understand Transformers in RL. First, we provide a mathematical definition of memory lengths and credit assignment lengths in RL, grounded in common understanding. We distinguish the memory lengths in reward, transition, policy, and value, which relate to the minimal lengths of recent histories needed to preserve the corresponding quantities. Crucially, our main theoretical result reveals that the memory length of an optimal policy can be upper bounded by the reward and transition memory lengths. For credit assignment, we adopt the forward view and define its length as the minimal number of future steps required for a greedy action to start outperforming any non-greedy counterpart.

The definitions provided are not only quantitative in theory but also actionable for practitioners, allowing them to analyze the memory and credit assignment requirements of many existing tasks. Equipped with these tools, we find that many tasks designed for evaluating memory also evaluate credit assignment, and vice versa, as demonstrated in Table 1. As a solution, we introduce toy examples called Passive and Active T-Mazes that decouple memory and credit assignment. These toy examples are configurable, enabling us to perform a scalable unit-test of a sole capability, in line with best practice (Osband et al., 2020).

Lastly, we evaluate memory-based RL algorithms with LSTMs or Transformers on our configurable toy examples and other challenging tasks to address our research question. We find that Transformers can indeed boost long-term memory in RL, scaling up to memory lengths of 1500 . However, Transformers do not improve long-term credit assignment, and struggle with data complexity in even short-term dependency tasks. While these results suggest that practitioners might benefit from using Transformer-based architectures in RL, they also highlight the importance of continued research on core RL algorithms. Transformers have yet to replace RL algorithm designers.

2 Measuring Temporal Dependencies in RL

MDPs and POMDPs. In a partially observable Markov decision process (POMDP) M O = ( O , A , P, R, γ, T ) 2 , an agent receives an observation o t ∈ O at step t ∈ { 1 , . . . , T } , takes an action a t ∈ A based on the observed history h 1: t := ( o 1: t , a 1: t -1 ) ∈ H t 3 , and receives a reward r t ∼ R t ( h 1: t , a t ) and the next observation o t +1 ∼ P ( · | h 1: t , a t ) . The initial observation h 1:1 := o 1 follows P ( o 1 ) . The total horizon is T ∈ N + ∪{ + ∞} and the discount factor is γ ∈ [0 , 1] (less than 1 for infinite horizon). In a Markov decision process (MDP) M S = ( S , A , P, R, γ, T ) , the observation o t and history h 1: t are replaced by the state s t ∈ S . In the most generic case, the agent is composed of a policy π ( a t | h 1: t ) and a value function Q π ( h 1: t , a t ) . The optimal value function Q ∗ ( h 1: t , a t ) satisfies Q ∗ ( h 1: t , a t ) = E [ r t | h 1: t , a t ]+ γ E o t +1 ∼ P ( | h 1: t ,a t ) [ max a t +1 Q ∗ ( h 1: t +1 , a t +1 ) ] and induces a deterministic optimal policy π ∗ ( h 1: t ) = argmax a t Q ∗ ( h 1: t , a t ) . Let Π ∗ M be the space of all optimal policies for a POMDP M .

Below, we provide definitions for context-based policy and value function, as well as the sum of n -step rewards given a policy. These concepts will be used in defining memory and credit assignment lengths.

Context-based policy and value function. A context-based policy π takes the recent l ctx ( π ) observations and actions as inputs, expressed as π ( a t | h t -l ctx ( π )+1: t ) . Here, l ctx ( π ) ∈ N represents the policy context length . For simplicity, we refer to context-based policies as 'policies'. Let Π k be the space of all policies with a context length of k . Markovian policies form Π 1 , while Π T , the largest policy space, contains an optimal policy. Recurrent policies, recursively taking one input at a time, belong to Π ∞ ; while Transformer-based policies have a limited context length. Context-based value function Q π n of a POMDP M is the expected return conditioned on the past n observations and actions, denoted as Q π n ( h t -n +1: t , a t ) = E π, M [∑ i = t γ i -t r i | h t -n +1: t , a t ] . By definition, Q π ( h 1: t , a t ) = Q π T ( h 1: t , a t ) , ∀ h 1: t , a t .

The sum of n -step rewards. Given a POMDP M and a policy π , the expectation of the discounted sum of n -step rewards is denoted as G π n ( h 1: t , a t ) = E π, M [ ∑ t + n -1 i = t γ i -t r i | h 1: t , a t ] . By definition, G π 1 ( h 1: t , a t ) = E [ r t | h 1: t , a t ] (immediate reward), and G π T -t +1 ( h 1: t , a t ) = Q π ( h 1: t , a t ) (value).

2.1 Memory Lengths

In this subsection, we introduce the memory lengths in the common components of RL, including the reward, transition, policy, and value functions. Then, we will show the theoretical result of the relation between these lengths.

The reward (or transition) memory length of a POMDP quantifies the temporal dependency distance between current expected reward (or next observation distribution ) and the most distant observation.

Definition 1.A ( Reward memory length m M reward ) . For a POMDP M , m M reward is the smallest n ∈ N such that the expected reward conditioned on recent n observations is the same as the one conditioned on full history, i.e ., E [ r t | h 1: t , a t ] = E [ r t | h t -n +1: t , a t ] , ∀ t, h 1: t , a t .

Definition 1.B ( Transition memory length m M transit ) . For a POMDP M , m M transit is the smallest n ∈ N such that next observations are conditionally independent of the rest history given the past n observations and actions, i.e ., o t +1 ⊥ ⊥ h 1: t , a t | h t -n +1: t , a t .

Next, policy memory length is defined based on the intuition that while a policy may take a long history of observations as its context, its action distribution only depends on recent observations. Therefore, policy memory length represents the temporal dependency distance between current action and the most distant observation.

Definition 1.C ( Policy memory length l mem ( π ) ) . The policy memory length l mem ( π ) of a policy π is the minimum horizon n ∈ { 0 , . . . , l ctx ( π ) } such that its actions are conditionally independent of the rest history given the past n observations and actions, i.e ., a t ⊥ ⊥ h t -l ctx ( π )+1: t | h t -n +1: t .

Lastly, we define the value memory length of a policy, which measures the temporal dependency distance between the current value and the most distant observation.

Definition 1.D ( Value memory length of a policy l M value ( π ) ) . A policy π has its value memory length l M value ( π ) as the smallest n such that the context-based value function is equal to the value function, i.e ., Q π n ( h t -n +1: t , a t ) = Q π ( h 1: t , a t ) , ∀ t, h 1: t , a t .

Now we prove that the policy and value memory lengths of optimal policies can be upper bounded by the reward and transition memory lengths (Proof in Appendix A).

Theorem 1 ( Upper bounds of memory lengths for optimal policies ) . For optimal policies with the shortest policy memory length, their policy memory lengths are upper bounded by their value memory lengths, which are further bounded by the maximum of reward and transition memory lengths.

Thm. 1 suggests that the constant of m M can serve as a proxy to analyze the policy and value memory lengths required in a task, which are often challenging to directly compute, as we will demonstrate in Sec. 3. 4

2.2 Credit Assignment Lengths

Classically, temporal credit assignment is defined with a backward view (Minsky, 1961; Sutton, 1984) how previous actions contributed to current reward. The backward view is specifically studied through gradient back-propagation in supervised learning (Lee et al., 2015; Ke et al., 2018). In this paper, we consider the forward view in reinforcement learning - how current action will influence future rewards. Specifically, we hypothesize that the credit assignment length n of a policy is how long the greedy action at the current time-step starts to make a difference to the sum of future n -step rewards.

Before introducing the formal definition, let us recall the common wisdom on dense-reward versus sparse-reward tasks. Dense reward is generally easier to solve than sparse reward if all other aspects of the tasks are the same. One major reason is that an agent can get immediate reward feedback in dense-reward tasks, while the feedback is delayed or even absent in sparse-reward tasks. In this view, solving sparse reward requires long-term credit assignment. Our definition will reflect this intuition.

Definition 2.A ( Credit assignment length of a history given a policy c ( h 1: t ; π ) ) . Given a policy π in a task, the credit assignment length of a history h 1: t quantifies the minimal temporal distance n 5 , such that a greedy action a ∗ t is strictly better than any non-greedy action a ′ t in terms of their n -step rewards G π n . Formally, let A ∗ t := argmax a t Q π ( h 1: t , a t ) 6 , define

Definition 2.B ( Credit assignment length of a policy c ( π ) and of a task c M ) . Credit assignment length of a policy π represents the worst case of credit assignment lengths of all visited histories. Formally, let d π be the history occupancy distribution of π , define c ( π ) := max h 1: t ∈ supp( d π ) c ( h 1: t ; π ) .

Further, we define the credit assignment length of a task M , denoted as c M , as the minimal credit assignment lengths over all optimal policies 7 , i.e ., c M := min π ∗ ∈ Π ∗ M c ( π ∗ ) .

As an illustrative example, consider the extreme case of a sparse-reward MDP where the reward is only provided at the terminal time step T . In this environment, the credit assignment length of any policy π is c ( π ) = T . This is because the credit assignment length of the initial state s 1 is always c ( s 1 ; π ) = T , as the performance gain of any greedy action at s 1 is only noticeable after T steps.

Our definition of credit assignment lengths has ties with the temporal difference (TD) learning algorithm (Sutton, 1988). With the tabular TD(0) algorithm, the value update is given by V ( s t ) ← V ( s t ) + α ( r t + γV ( s t +1 ) -V ( s t )) . This formula back-propagates the information of current reward r t and next value V ( s t +1 ) to current value V ( s t ) . When we apply the TD(0) update backwards in time for n steps along a trajectory, we essentially back-propagate the information of n -step rewards to the current value. Thus, the credit assignment length of a state measures the number of TD(0) updates we need to take a greedy action for that state.

3 Environment Analysis of Memory and Credit Assignment Lengths

By employing the upper bound m M on memory lengths (Thm. 1) and credit assignment length c M (Def. 2.B), we can perform a quantitative analysis of various environments. These environments include both abstract problems and concrete benchmarks. In this section, we aim to address two questions: (1) Do prior environments actually require long-term memory or credit assignment? (2) Can we disentangle memory from credit assignment?

3.1 How Long Do Prior Tasks Require Memory and Credit Assignment?

To study memory or credit assignment in RL, prior works propose abstract problems and evaluate their algorithms in benchmarks. A prerequisite for evaluating agent capabilities in memory or credit assignment is understanding the requirements for them in the tested tasks. Therefore, we collect representative problems and benchmarks with a similar analysis in Sec. 3.2, to figure out the quantities m M related to memory lengths and c M related to credit assignment length.

Regarding abstract problems, we summarize the analysis in Appendix B. For example, decomposable episodic reward assumes the reward is only given at the terminal step and can be decomposed into the sum of Markovian rewards. This applies to episodic MuJoCo (Ren et al., 2022) and Catch with delayed rewards (Raposo et al., 2021). We show that it has policy memory length l mem ( π ∗ ) at most 1 while requiring m M reward and c M of T , indicating that it is really challenging in both memory and credit assignment. Delayed environments are known as difficult to assign credit (Sutton, 1984). We show that delayed reward (Arjona-Medina et al., 2019) or execution (Derman et al., 2021) or observation (Katsikopoulos and Engelbrecht, 2003) with n steps all require m M reward of n , while only delayed rewards and execution may be challenging in credit assignment.

As for concrete benchmarks, Table 1 summarizes our analysis, with detailed descriptions of these environments in Appendix B. First, some tasks designed for long-term memory (Beck et al., 2019; Osband et al., 2020; Hung et al., 2018; Yang and Nguyen, 2021; Esslinger et al., 2022) or credit

assignment (Osband et al., 2020) indeed solely evaluate the respective capability. Although some memory tasks such as Numpad (Parisotto et al., 2020) and Memory Maze (Pasukonis et al., 2022) do require long-term memory with short-term credit assignment, the exact memory and credit assignment lengths are hard to quantify due to the complexity. In contrast, our Passive T-Maze introduced in Sec. 3.2, provides a clear understanding of the memory and credit assignment lengths, allowing us to successfully scale the memory length up to 1500 .

Conversely, other works (Bakker, 2001; Lambrechts et al., 2021; Esslinger et al., 2022; Hung et al., 2018; Raposo et al., 2021; Ren et al., 2022) actually entangle memory and credit assignment together, despite their intention to evaluate one of the two. Lastly, some POMDPs (Ni et al., 2022) only require short-term memory and credit assignment, which cannot be used for any capability evaluation.

3.2 Disentangling Memory from Credit Assignment

We can demonstrate that memory and credit assignment are independent concepts by providing two examples, all with a finite horizon of T and no discount. The examples have distinct pairs of (max( r M , m M transit ) , c M ) : ( T, 1) , ( T, T ) , respectively.

We adapt T-Maze (Fig. 1), a POMDP environment rooted in neuroscience (O'Keefe and Dostrovsky, 1971) and AI (Bakker, 2001; Osband et al., 2020), to allow different credit assignment lengths. Passive and active concepts are borrowed from Hung et al. (2018).

Example 1 ( Passive T-Maze ) . T-Maze has a 4-direction action space {L , R , U , D} . It features a long corridor of length L , starting from (oracle) state O and ending with (junction) state J , connecting two potential goal states G 1 and G 2 . O has a horizontal position of 0 , and J has L . The agent observes null N except at states J, O, G 1 , G 2 . The observation space { N,J,O,G 1 , G 2 } is two-dimensional. At state O , the agent can observe the goal position G ∈ { G 1 , G 2 } , uniformly sampled at the beginning of an episode. The initial state is S . The transition is deterministic, and the agent remains in place if it hits the wall. Thus, the horizontal position of the agent x t ∈ N is determined by its previous action sequence a 1: t -1 .

In a Passive T-Maze M passive , O = S , allowing the agent to observe G initially, and L = T -1 . Rewards are given by R t ( h 1: t , a t ) = 1 ( x t +1 ≥ t ) -1 T -1 for t ≤ T -1 , and R T ( h 1: T , a T ) = 1 ( o T +1 = G ) . The unique optimal policy π ∗ moves right for T -1 steps, then towards G , yielding an expected return of 1 . 0 . The optimal Markovian policy can only guess the goal position, ending with an expected return of 0 . 5 . The worst policy has an expected return of -1 . 0 .

Example 2 ( Active T-Maze ) . In an Active T-Maze M active , O is one step left of S and L = T -2 . The unique optimal policy moves left to reach O , then right to J . The rewards R 1 = 0 , R t ( h 1: t , a t ) = 1 ( x t +1 ≥ t -1) -1 T -2 for 2 ≤ t ≤ T -1 , and R T ( h 1: T , a T ) = 1 ( o T +1 = G ) . The optimal Markovian and worst policies have the same expected returns as in the passive setting.

M passive is specially designed for testing memory only. To make the final decision, the optimal policy must recall G , observed in the first step. All the memory lengths are T as the terminal reward depends on initial observation and observation at J depends on the previous action sequence. The credit assignment length is 1 as immediate penalties occur for suboptimal actions. In contrast, M active also requires a credit assignment length of T , as the initial optimal action L only affects n -step rewards when n = T . Passive and Active T-Mazes are suitable for evaluating the ability of pure memory, and both memory and credit assignment in RL, as they require little exploration due to the dense rewards before reaching the goal.

4 Related Work

Memory and credit assignment have been extensively studied in RL, with representative works shown in Table 1. The work most closely related to ours is bsuite (Osband et al., 2020), which

BC

中

口

一

Phase 1

Phase2

Phase3

Phase1

Phase2

Phase3

also disentangles memory and credit assignment in their benchmark. bsuite is also configurable with regard to the memory and credit assignment lengths to evaluate an agent's capability. Our work extends Osband et al. (2020) by providing formal definitions of memory and credit assignment lengths in RL for task design and evaluation. Our definitions enable us to analyze existing POMDP benchmarks.

Several prior works have assessed memory (Parisotto et al., 2020; Esslinger et al., 2022; Chen et al., 2022; Pleines et al., 2023; Morad et al., 2023) in the context of Transformers in RL. However, these prior methods use tasks that require (relatively short) context lengths between 50 and 831 . In addition, there are two separate lines of work using Transformers for better credit assignment. One line of work (Liu et al., 2019; Ferret et al., 2020) has developed algorithms that train Transformers to predict rewards for return redistribution. The other line of work (Chen et al., 2021; Zheng et al., 2022) proposes return-conditioned agents trained on offline RL datasets. Both lines are beyond the scope of canonical online RL algorithms that purely rely on TD learning, which we focus on.

5 Evaluating Memory-Based RL Algorithms

In this section, we evaluate the memory and credit assignment capabilities of memory-based RL algorithms. Our first experiment aims to evaluate the memory abilities by disentangling memory from credit assignment, using our Passive T-Maze and the Passive Visual Match (Hung et al., 2018) tasks in Sec. 5.1. Next, we test the credit assignment abilities of memory-based RL in our Active T-Maze and the Key-to-Door (Raposo et al., 2021) tasks in Sec. 5.2. Passive Visual Match and Key-to-Door tasks (demonstrated in Fig. 2) can be viewed as pixel-based versions of Passive and Active T-Mazes, respectively. Lastly, we evaluate memory-based RL on standard POMDP benchmarks that only require short-term dependencies in Sec. 5.3, with a focus on their sample efficiency.

We focus on model-free RL agents that take observation and action sequences as input, which serves as a simple baseline that can achieve excellent results in many tasks (Ni et al., 2022). The agent architecture is based on the codebase from Ni et al. (2022), consisting of observation and action embedders, a sequence model, and actor-critic heads. We compare agents with LSTM and Transformer (GPT-2 (Radford et al., 2019)) architectures. For a fair comparison, we use the same hidden state size ( 128 ) and tune the number of layers (varying from 1 , 2 , 4 ) for both LSTM and Transformer architectures. Since our tasks all require value memory lengths to be the full horizon, we set the context lengths to the full horizon T .

For the T-Maze tasks, we use DDQN (van Hasselt et al., 2016) with epsilon-greedy exploration as the RL algorithm to better control the exploration strategy. In pixel-based tasks, we use CNNs as observation embedders and SAC-Discrete (Christodoulou, 2019) as RL algorithm following Ni et al. (2022). We train all agents to convergence or at most 10M steps on these tasks. We provide training details and learning curves in Appendix C.

5.1 Transformers Shine in Pure Long-Term Memory Tasks

First, we provide evidence that Transformers can indeed enhance long-term memory in tasks that purely test memory, Passive T-Mazes. In Fig. 3 (left), Transformer-based agents consistently solve the task requiring memory lengths from 50 up to 1500 . To the best of our knowledge, this achievement

0

200

400

600

800

1000

1200

1400

(Easy)

Memory length

(Hard)

-

1

0

1

return

Passive T-Maze

GPT

LSTM

Optimal agent w/o memory

200

400

600

800

1000

(Easy)

Memory length

(Hard)

0

.

25

0

.

50

0

.

75

success

Passive Visual Match

→

→

0

1

2

3

4

env steps

10

6

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

return

Passive T-Maze w/ memory length of 500

0

2

4

6

env steps

10

6

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

return

Passive T-Maze w/ memory length of 750

0

2

4

6

8

env steps

10

6

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

return

Passive T-Maze w/ memory length of 1000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

env steps

10

7

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

return

Passive T-Maze w/ memory length of 1500

GPT

LSTM

×

×

×

×

may well set a new record for the longest memory length of a task that an RL agent can solve. In contrast, LSTM-based agents start to falter at a memory length of 250 , underperforming the optimal Markovian policies with a return of 0 . 5 . Both architectures use a single layer, yielding the best performance. We also provide detailed learning curves for tasks requiring long memory lengths in Fig. 4. Transformer-based agent also shows much greater stability and sample efficiency across seeds when compared to their LSTM counterparts.

On the other hand, the pixel-based Passive Visual Match (Hung et al., 2018) has a more complex reward function: zero rewards for observing the color in Phase 1, immediate rewards for apple picking in Phase 2 requiring short-term dependency, and a final bonus for reaching the door of the matching color, requiring long-term memory. The success rate solely depends on the final bonus. The ratio of the final bonus w.r.t. the total rewards is inversely proportional to the episode length, as the number of apples increases over time in Phase 2. Thus, the low ratio makes learning long-term memory challenging for this task. As shown in Fig. 3 (right), Transformer-based agents slightly outperform LSTM-based ones 8 in success rates when the tasks require a memory length over 120 , yet their success rates are close to that of the optimal Markovian policy (around 1 / 3 ). Interestingly, Transformer-based agents achieve much higher returns with more apples collected, shown in Fig. 5. This indicates Transformers can help short-term memory (credit assignment) in some tasks.

5.2 Transformers Enhance Temporal Credit Assignment in RL, but not Long-Term

Next, we present evidence that Transformers do not enhance long-term credit assignment for memorybased RL. As illustrated in Fig. 6 (left), in Active T-Mazes, by merely shifting the Oracle one step to the left from the starting position as compared to Passive T-Mazes, both Transformer-based and LSTM-based agents fail to solve the task when the credit assignment (and memory) length extends to just 250 . The failure of Transformer-based agents on Active T-Mazes suggests the bottleneck of credit assignment. This is because they possess sufficient memory capabilities to solve this task (as evidenced by their success in Passive T-Mazes), and they can reach the junction state without exploration issues (as evidenced by their return of more than 0 . 0 ). A similar trend is observed in pixel-based Key-to-Door tasks (Raposo et al., 2021), where both Transformer-based and LSTM-

0

2

4

6

env steps

10

6

100

150

200

250

return

Passive Visual Match w/ memory length of 750

0

2

4

6

env steps

10

6

0

.

1

0

.

2

0

.

3

0

.

4

success

Passive Visual Match w/ memory length of 750

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

env steps

10

7

100

200

300

return

Passive Visual Match w/ memory length of 1000

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

env steps

10

7

0

.

2

0

.

4

0

.

6

success

Passive Visual Match w/ memory length of 1000

GPT

LSTM

×

×

×

×

0

100

200

300

400

500

(Easy)

Credit assignment length

(Hard)

-

0

.

5

0

.

0

0

.

5

1

.

0

return

Active T-Maze

GPT

LSTM

Optimal agent w/o credit assignment

100

200

300

400

500

(Easy)

Credit assignment length

(Hard)

0

.

0

0

.

5

1

.

0

success

Key-to-Door

→

→

based agents struggle to solve the task with a credit assignment length of 250 , as indicated by the low success rates in Fig. 6 (right). Nevertheless, we observe that Transformers can improve credit assignment, especially when the length is medium around 100 in both tasks, compared to LSTMs.

On the other hand, we investigate the performance of multi-layer Transformers on credit assignment tasks to explore the potential scaling law in RL, as discovered in supervised learning (Kaplan et al., 2020). Fig. 7 reveals that multi-layer Transformers greatly enhance performance in tasks with credit assignment lengths up to 100 in Active T-Maze and 120 in Key-to-Door. Yet they fail to improve long-term credit assignment. Interestingly, 4 -layer Transformers have similar or worse performance compared to 2 -layer Transformers, suggesting that the scaling law might not be seamlessly applied to credit assignment in RL.

0

100

200

300

400

500

(Easy)

Credit assignment length

(Hard)

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

return

Active T-Maze

1

2

4

100

200

300

400

500

(Easy)

Credit assignment length

(Hard)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

success

Key-to-Door

0

1

2

3

4

env steps

10

6

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

return

Active T-Maze w/ credit assignment length of 100

1

2

4

0

2

4

6

env steps

10

6

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

return

Active T-Maze w/ credit assignment length of 250

1

2

4

→

→

×

×

5.3 Transformers in RL Raise Sample Complexity in Certain Short-Term Memory Tasks

Lastly, we revisit the standard POMDP benchmark used in prior works (Ni et al., 2022; Han et al., 2020; Meng et al., 2021). This PyBullet benchmark consists of tasks where only the positions (referred to as '-P' tasks) or velocities (referred to as '-V' tasks) of the agents are revealed, with dense rewards. As a result, they necessitate short-term credit assignment. Further, the missing velocities can be readily inferred from two successive positions, and the missing positions can be

0

.

0

0

.

5

1

.

0

1

.

5

env steps

×

10

6

0

1

2

3

return

×

10

3

Ant-P

0

.

0

0

.

5

1

.

0

1

.

5

env steps

×

10

6

-

1

0

1

2

3

return

×

10

3

Cheetah-P

0

.

0

0

.

5

1

.

0

1

.

5

env steps

×

10

6

0

1

2

return

×

10

3

Hopper-P

0

.

0

0

.

5

1

.

0

1

.

5

env steps

×

10

6

0

.

0

0

.

5

1

.

0

return

×

10

3

Walker-P

0

.

0

0

.

5

1

.

0

1

.

5

env steps

10

6

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

return

×

10

3

0

.

0

0

.

5

1

.

0

1

.

5

env steps

10

6

-

1

0

1

2

return

×

10

3

Cheetah-V

0

.

0

0

.

5

1

.

0

1

.

5

env steps

10

6

0

.

0

0

.

5

1

.

0

1

.

5

return

×

10

3

Hopper-V

0

.

0

0

.

5

1

.

0

1

.

5

env steps

10

6

0

1

2

3

4

return

×

10

2

Walker-V

GPT

LSTM

×

×

×

×

approximately deduced by recent velocities, thereby requiring short-term memory. Adhering to the implementation of recurrent TD3 (Ni et al., 2022) and substituting LSTMs with Transformers, we investigate their sample efficiency over 1.5M steps. We find that in most tasks, Transformer-based agents show worse sample efficiency than LSTM-based agents, as displayed in Fig. 8. This aligns with the findings in POPGym paper (Morad et al., 2023) where GRU-based PPO is more sample-efficient than Transformer-based PPO across POPGym benchmark that requires relatively short memory lengths. These observations are also consistent with the known tendency of Transformers to thrive in situations with large datasets in supervised learning.

6 Limitations

Weprovide quantitative definitions of memory and credit assignment lengths, but developing programs capable of calculating them for any given task remains an open challenge. This can often be desirable since sequence models can be expensive to train in RL agents. The worst compute complexity of c M can be O ( |S||A| T 2 ) for a finite MDP, while the complexity of m M can be O (( |O| + |A| ) T |O||A| T ) for a finite POMDP. In addition, c M does not take the variance of intermediate rewards into account. These reward noises have been identified to complicate the learning process for credit assignment (Mesnard et al., 2021). On the other hand, m M does not take into account the memory capacity ( i.e ., the number of bits) and robustness to observation noise (Beck et al., 2019).

On the empirical side, we evaluate Transformers on a specific architecture (GPT-2) with a relatively small size in the context of online model-free RL. As a future research direction, it would be interesting to explore how various sequence architectures and history abstraction methods may impact and potentially improve long-term memory and/or credit assignment.

7 Conclusion

In this study, we evaluate the memory and credit assignment capabilities of memory-based RL agents, with a focus on Transformer-based RL. While Transformer-based agents excel in tasks requiring long-term memory, they do not improve long-term credit assignment and generally have poor sample efficiency. Furthermore, we highlighted that many existing RL tasks, even those designed to evaluate (long-term) memory or credit assignment, often intermingle both capabilities or require only shortterm dependencies. While Transformers are powerful tools in RL, especially for long-term memory tasks, they are not a universal solution to all RL challenges. Our results underscore the ongoing need for careful task design and the continued advancement of core RL algorithms.

References

Agarwal, A. A. Rahman, P.-L. St-Charles, S. J. Prince, and S. E. Kahou. Transformers in reinforcement learning: A survey. arXiv preprint arXiv:2307.05979 , 2023. 1

A. Arjona-Medina, M. Gillhofer, M. Widrich, T. Unterthiner, J. Brandstetter, and S. Hochreiter. RUDDER: return decomposition for delayed rewards. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada , 2019. 5, 17

Bakker. Reinforcement learning with long short-term memory. In Advances in Neural Information Processing Systems 14 [Neural Information Processing Systems: Natural and Synthetic, NIPS 2001, December 3-8, 2001, Vancouver, British Columbia, Canada] , 2001. 5, 6, 18, 19

Beck, K. Ciosek, S. Devlin, S. Tschiatschek, C. Zhang, and K. Hofmann. Amrl: Aggregated memory for reinforcement learning. In International Conference on Learning Representations , 2019. 5, 10, 18

G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. J. Artif. Intell. Res. , 2013. 2

B. Blankenship. Memory span: a review of the literature. Psychological Bulletin , 35(1):1, 1938. 1

Chen, Y. Wu, J. Yoon, and S. Ahn. Transdreamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.09481 , 2022. 7

Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , 2021. 1, 7, 20

Chevalier-Boisvert, L. Willems, and S. Pal. Minimalistic gridworld environment for openai gym. https: //github.com/maximecb/gym-minigrid , 2018. 5, 18

Christodoulou. Soft actor-critic for discrete action settings. arXiv preprint arXiv:1910.07207 , 2019. 7, 23

N. Dempster. Memory span: Sources of individual and developmental differences. Psychological Bulletin , 89 (1):63, 1981. 1

Derman, G. Dalal, and S. Mannor. Acting in delayed environments with non-stationary markov policies. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 , 2021. 5, 17

Esslinger, R. Platt, and C. Amato. Deep transformer q-networks for partially observable reinforcement learning. arXiv preprint arXiv:2206.01078 , 2022. 5, 6, 7, 19

Ferret, R. Marinier, M. Geist, and O. Pietquin. Self-attentional credit assignment for transfer in reinforcement learning. In Proceedings of the Twenty-Ninth International Joint Conference on Artificial Intelligence, IJCAI 2020 , 2020. 7

Fortunato, M. Tan, R. Faulkner, S. Hansen, A. P. Badia, G. Buttimore, C. Deck, J. Z. Leibo, and C. Blundell. Generalization of reinforcement learners with working and episodic memory. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada , 2019. 5, 18, 19

Fujimoto, H. van Hoof, and D. Meger. Addressing function approximation error in actor-critic methods. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018 , 2018. 23

J. Gershman and N. D. Daw. Reinforcement learning and episodic memory in humans and animals: an integrative framework. Annual review of psychology , 68:101-128, 2017. 2

Grigsby, L. Fan, and Y. Zhu. Amago: Scalable in-context reinforcement learning for adaptive agents. arXiv preprint arXiv:2310.09971 , 2023. 21

Han, K. Doya, and J. Tani. Variational recurrent models for solving partially observable control tasks. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020 , 2020. 9

Hochreiter and J. Schmidhuber. Long short-term memory. Neural computation , 9(8):1735-1780, 1997. 22

Humplik, A. Galashov, L. Hasenclever, P. A. Ortega, Y. W. Teh, and N. Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv:1905.06424 , 2019. 19

Hung, T. P. Lillicrap, J. Abramson, Y. Wu, M. Mirza, F. Carnevale, A. Ahuja, and G. Wayne. Optimizing agent behavior over long time scales by transporting value. arXiv preprint arXiv:1810.06721 , 2018. 2, 5, 6, 7, 8, 18, 20

Ichter, A. Brohan, Y. Chebotar, C. Finn, K. Hausman, A. Herzog, D. Ho, J. Ibarz, A. Irpan, E. Jang, R. Julian, D. Kalashnikov, S. Levine, Y. Lu, C. Parada, K. Rao, P. Sermanet, A. Toshev, V. Vanhoucke, F. Xia, T. Xiao, P. Xu, M. Yan, N. Brown, M. Ahn, O. Cortes, N. Sievers, C. Tan, S. Xu, D. Reyes, J. Rettinghouse, J. Quiambao, P. Pastor, L. Luu, K. Lee, Y. Kuang, S. Jesmonth, N. J. Joshi, K. Jeffrey, R. J. Ruano, J. Hsu, K. Gopalakrishnan, B. David, A. Zeng, and C. K. Fu. Do as I can, not as I say: Grounding language in robotic affordances. In Conference on Robot Learning, CoRL 2022, 14-18 December 2022, Auckland, New Zealand , 2022. 1

Janner, Q. Li, and S. Levine. Offline reinforcement learning as one big sequence modeling problem. In Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurIPS 2021, December 6-14, 2021, virtual , 2021. 1

P. Kaelbling, M. L. Littman, and A. R. Cassandra. Planning and acting in partially observable stochastic domains. Artif. Intell. , 1998. 2

Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, J. Wu, and D. Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 , 2020. 9

V. Katsikopoulos and S. E. Engelbrecht. Markov decision processes with delays and asynchronous cost collection. IEEE Trans. Autom. Control. , 2003. 5, 18

R. Ke, A. Goyal, O. Bilaniuk, J. Binas, M. C. Mozer, C. Pal, and Y. Bengio. Sparse attentive backtracking: Temporal credit assignment through reminding. In Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada , 2018. 4

Lambrechts, F. De Geeter, N. Vecoven, D. Ernst, and G. Drion. Warming-up recurrent neural networks to maximize reachable multi-stability greatly improves learning. arXiv preprint arXiv:2106.01001 , 2021. 5, 6, 19

Lampinen, S. Chan, A. Banino, and F. Hill. Towards mental time travel: a hierarchical memory for reinforcement learning agents. Advances in Neural Information Processing Systems , 34:28182-28195, 2021. 1, 5, 19

Lee, S. Zhang, A. Fischer, and Y. Bengio. Difference target propagation. In Machine Learning and Knowledge Discovery in Databases - European Conference, ECML PKDD 2015, Porto, Portugal, September 7-11, 2015, Proceedings, Part I , 2015. 4

Lee, O. Nachum, M. Yang, L. Lee, D. Freeman, S. Guadarrama, I. Fischer, W. Xu, E. Jang, H. Michalewski, and I. Mordatch. Multi-game decision transformers. In NeurIPS , 2022. 1

Li, H. Luo, Z. Lin, C. Zhang, Z. Lu, and D. Ye. A survey on transformers in reinforcement learning. arXiv preprint arXiv:2301.03044 , 2023. 1

P. Lillicrap, J. J. Hunt, A. Pritzel, N. Heess, T. Erez, Y . Tassa, D. Silver, and D. Wierstra. Continuous control with deep reinforcement learning. In 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings , 2016. 21

Liu, Y. Luo, Y. Zhong, X. Chen, Q. Liu, and J. Peng. Sequence modeling of temporal credit assignment for episodic reinforcement learning. arXiv preprint arXiv:1905.13420 , 2019. 7, 20

C. Melo. Transformers are meta-reinforcement learners. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA , 2022. 1

Meng, R. Gorbet, and D. Kulic. Memory-based deep reinforcement learning for pomdps. In IEEE/RSJ International Conference on Intelligent Robots and Systems, IROS 2021, Prague, Czech Republic, September 27 - Oct. 1, 2021 , 2021. 9

Mesnard, T. Weber, F. Viola, S. Thakoor, A. Saade, A. Harutyunyan, W. Dabney, T. S. Stepleton, N. Heess, A. Guez, E. Moulines, M. Hutter, L. Buesing, and R. Munos. Counterfactual credit assignment in model-free reinforcement learning. In Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event , 2021. 2, 10, 20

Micheli, E. Alonso, and F. Fleuret. Transformers are sample efficient world models. arXiv preprint arXiv:2209.00588 , 2022. 1

Minsky. Steps toward artificial intelligence. Proceedings of the IRE , 49(1):8-30, 1961. 4

D. Morad, R. Kortvelesy, M. Bettini, S. Liwicki, and A. Prorok. Popgym: Benchmarking partially observable reinforcement learning. arXiv preprint arXiv:2303.01859 , 2023. 5, 7, 10, 20

Ni, B. Eysenbach, and R. Salakhutdinov. Recurrent model-free RL can be a strong baseline for many pomdps. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA , 2022. 5, 6, 7, 9, 10, 21

O'Keefe and J. Dostrovsky. The hippocampus as a spatial map: Preliminary evidence from unit activity in the freely-moving rat. Brain research , 1971. 6

Osband, Y. Doron, M. Hessel, J. Aslanides, E. Sezener, A. Saraiva, K. McKinney, T. Lattimore, C. Szepesvári, S. Singh, B. V. Roy, R. S. Sutton, D. Silver, and H. van Hasselt. Behaviour suite for reinforcement learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020 , 2020. 2, 5, 6, 7, 19, 20

Ouyang, J. Wu, X. Jiang, D. Almeida, C. L. Wainwright, P. Mishkin, C. Zhang, S. Agarwal, K. Slama, A. Ray, J. Schulman, J. Hilton, F. Kelton, L. Miller, M. Simens, A. Askell, P. Welinder, P. F. Christiano, J. Leike, and R. Lowe. Training language models to follow instructions with human feedback. In NeurIPS , 2022. 1

Parisotto, H. F. Song, J. W. Rae, R. Pascanu, Ç. Gülçehre, S. M. Jayakumar, M. Jaderberg, R. L. Kaufman, A. Clark, S. Noury, M. M. Botvinick, N. Heess, and R. Hadsell. Stabilizing transformers for reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event , 2020. 1, 5, 6, 7, 19

Pascanu, Ç. Gülçehre, K. Cho, and Y. Bengio. How to construct deep recurrent neural networks. In 2nd International Conference on Learning Representations, ICLR 2014, Banff, AB, Canada, April 14-16, 2014, Conference Track Proceedings , 2014. 23

Pasukonis, T. P. Lillicrap, and D. Hafner. Evaluating long-term memory in 3d mazes. arXiv preprint arXiv:2210.13383 , 2022. 5, 6, 19

Pleines, M. Pallasch, F. Zimmer, and M. Preuss. Memory gym: Partially observable challenges to memorybased agents. In The Eleventh International Conference on Learning Representations , 2023. 5, 7, 20

Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019. 1, 7, 22

Raposo, S. Ritter, A. Santoro, G. Wayne, T. Weber, M. M. Botvinick, H. van Hasselt, and H. F. Song. Synthetic returns for long-term credit assignment. arXiv preprint arXiv:2102.12425 , 2021. 5, 6, 7, 8, 20

Ren, R. Guo, Y. Zhou, and J. Peng. Learning long-term reward redistribution via randomized return decomposition. In The Tenth International Conference on Learning Representations, ICLR 2022, Virtual Event, April 25-29, 2022 , 2022. 5, 6, 17, 20

Robine, M. Höftmann, T. Uelwer, and S. Harmeling. Transformer-based world models are happy with 100k interactions. arXiv preprint arXiv:2303.07109 , 2023. 1

S. Sutton. Temporal credit assignment in reinforcement learning . University of Massachusetts Amherst, 1984. 2, 4, 5

S. Sutton. Learning to predict by the methods of temporal differences. Mach. Learn. , 1988. 4

S. Sutton and A. G. Barto. Reinforcement learning: An introduction . MIT press, 2018. 16

Thrun. Monte carlo pomdps. In Advances in Neural Information Processing Systems 12, [NIPS Conference, Denver, Colorado, USA, November 29 - December 4, 1999] , 1999. 19

Todorov, T. Erez, and Y. Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, IROS 2012, Vilamoura, Algarve, Portugal, October 7-12, 2012 , 2012. 2

van Hasselt, A. Guez, and D. Silver. Deep reinforcement learning with double q-learning. In Proceedings of the Thirtieth AAAI Conference on Artificial Intelligence, February 12-17, 2016, Phoenix, Arizona, USA , 2016. 7, 23

Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA , 2017. 1, 22

L. Waskom. seaborn: statistical data visualization. Journal of Open Source Software , 6(60):3021, 2021. doi: 10.21105/joss.03021. URL https://doi.org/10.21105/joss.03021 . 24

Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, and J. Brew. Huggingface's transformers: State-of-the-art natural language processing. arXiv preprint arXiv:1910.03771 , 2019. 22

Yang and H. Nguyen. Recurrent off-policy baselines for memory-based continuous control. arXiv preprint arXiv:2110.12628 , 2021. 5, 19, 20, 21, 22

Yarats, I. Kostrikov, and R. Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021 , 2021a. 21, 22

Yarats, A. Zhang, I. Kostrikov, B. Amos, J. Pineau, and R. Fergus. Improving sample efficiency in modelfree reinforcement learning from images. In Thirty-Fifth AAAI Conference on Artificial Intelligence, AAAI 2021, Thirty-Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, The Eleventh Symposium on Educational Advances in Artificial Intelligence, EAAI 2021, Virtual Event, February 2-9, 2021 , 2021b. 21, 22

Zheng, A. Zhang, and A. Grover. Online decision transformer. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA , 2022. 1, 7

A. Zilli and M. E. Hasselmo. Modeling the role of working memory and episodic memory in behavioral tasks. Hippocampus , 18(2):193-209, 2008. 2

Appendix

Table of Contents

A Proofs

Lemma 1 ( Upper bound of value memory length ) . For any policy π ,

Proof of Lemma 1. We show this by expanding the sum of rewards in Q -value function:

where line ( A ) follows from the definition of memory lengths in policy (Def. 1.C), reward (Def. 1.A), and transition (Def. 1.B). By the definition of value memory length (Def. 1.D), the shortest history to match the value, we have the upper bound: l M value ( π ) ≤ max( m M reward , m M transit , l mem ( π ) -1) .

Lemma 2 ( Lower bound of value memory length ) . For an optimal policy π ∗ ∈ argmin π ∈ Π ∗ M { l mem ( π ) } , we have l mem ( π ∗ ) ≤ l M value ( π ∗ ) .

Proof. Proof by contradiction. Suppose l M value ( π ∗ ) < l mem ( π ∗ ) . We can construct a deterministic policy π ′ that has context length of l M value ( π ∗ ) :

This is a greedy policy derived from π ∗ 's value function. The policy memory length of π ′ is strictly shorter than l mem ( π ∗ ) , by definition and assumption:

Similar to policy improvement theorem (Sutton and Barto, 2018, Chap 4.2) in Markov policies, we can show that π ′ is at least as good as π ∗ , thus contradicting the assumption that π ∗ is an optimal policy that has the shortest policy memory length. To see this, for any t and h 1: t ,

where ( A ) follows the definition of π ′ , ( B ) uses the Bellman equation and ( C ) uses the property of greedy policy π ′ .

Proof of Theorem 1. By Lemma 2, l M value ( π ∗ ) ≥ l mem ( π ∗ ) . On the other hand, by Lemma 1, l M value ( π ∗ ) ≤ max( m M reward , m M transit , l mem ( π ∗ ) -1) . Thus,

c ( π ∗ ) may be not unique among all optimal policies. For example, consider an MDP described in Fig. 9, and two optimal deterministic Markov policy, π ∗ 1 , π ∗ 2 , only different at the state P2: π ∗ 1 ( P2 ) = x while π ∗ 2 ( P2 ) = y . It is easy to compute the credit assignment length in this MDP, because only state P1 contains suboptimal action. At state P1, the optimal action y starts to be better than suboptimal action x at the next 2 steps for π ∗ 2 , while at the next 3 steps for π ∗ 1 . Thus, c ( π ∗ 1 ) = 3 and c ( π ∗ 2 ) = 2 .

This implies

P1

0

P2

0

P3

-

1

P5

+1

P4

+1

P7

+5

P6

-

1

x

y

x

y

x,y

x,y

x,y

x,y

x,y

B Details on Environments Related to Memory and Credit Assignment

B.1 Abstract Problems

Example 3 ( n -order MDPs) . It is a POMDP where the recent n observations form the current state, i.e . the transition is P ( o t +1 | o t -n +1: t , a t ) and the reward is R t ( o t -n +1: t , a t ) .

Example 5 (Decomposed episodic reward (Ren et al., 2022)) . It has an episodic reward (Eg. 4), decomposed into the sum of Markovian rewards: R T ( h 1: T , a T ) = ∑ T t =1 r ′ ( o t , a t ) , where r ′ : O × A → R is a Markovian reward function. The transition P ( o t +1 | o t , a t ) is Markovian. The smallest policy memory length for an optimal policy l mem ( π ∗ ) can be shown as 1 (see below).

Example 4 (Episodic reward) . It is a POMDP with a finite horizon of T . The reward function is episodic in that R t = 0 for ∀ t < T .

Proof of decomposed episodic reward problem requiring l mem ( π ∗ ) = 1 at most. Suppose the policy is non-stationary, composed of { π t } T t =1 . For any time step t ≤ T in the POMDP M ,

By taking argmax,

Now we show that there exists an optimal policy that is Markovian by backward induction. When t = T , (27) reduces to

thus Markovian policy π T ( o T ) = argmax a T r ′ ( o T , a T ) is optimal. Now assume π t +1: T are Markovian and optimal, consider the step t in (27), all future observations and actions ( o t +1: T , a t +1: T ) rely on ( o t , a t ) by the property of Markovian transition and induction, thus optimal action a ∗ t can only rely on o t .

Example 6 (Delayed rewards of n steps (Arjona-Medina et al., 2019)) . The transition P ( o t +1 | o t , a t ) is Markovian, and the reward R t ( o t -n +1 , a t -n +1 ) depends on the observation and action n steps before. Decomposed episodic rewards (Eg. 5) can be viewed as an extreme case of delayed rewards.

Example 7 (Delayed execution of n steps (Derman et al., 2021)) . The transition is P ( o t +1 | o t , a t -n ) , and the reward is R t ( o t , a t -n ) .

Example 8 (Delayed observation of n steps (Katsikopoulos and Engelbrecht, 2003)) . The transition is P ( o t +1 | o t , a t -n ) , and the reward is R t ( o t , a t -n : t ) .

We summarize our analysis in Table 2.

B.2 Concrete Benchmarks for Memory

T-Maze (Bakker, 2001). The task is almost the same as our Passive T-Maze, except that the agent is not penalized if it moves left instead of right in the corridor. The modification results in a credit assignment length closer to T . The authors test RL algorithms with the largest horizon of 70 .

MiniGrid-Memory (Chevalier-Boisvert et al., 2018). This task has the same structure as our Active T-Maze. In the worst-case scenario, the optimal policy can traverse the horizontal corridor back and forth, and then follow the vertical corridor to solve the task. Thus, the memory and credit assignment lengths of optimal policy is upper bounded by three times the maze side length, which is 3 × 17 = 51 for a 17 × 17 maze ( MiniGrid-MemoryS17Random-v0 ). The reward memory length can equal the horizon of 1445 .

Passive Visual Match (Hung et al., 2018) . The environment consists of a partially observable 7 × 11 grid-world, where the agent can only observe a 5 × 5 grid surrounding itself. The task is separated into three phases. In the first, the agent passively observes a randomly generated color. In the second, the agent must pick up apples with immediate reward. In the final phase, three random colored squares are displayed, and the agent must pick up the one corresponding to the color observed in the first phase, which tests the agent's ability to recall temporally distant events. Since rewards are relatively immediate with respect to the agent's actions, the credit assignment length of optimal policy is upper bounded by the length of the shortest path needed to traverse the grid-world, i.e . 7 + 11 = 18 .

TMaze Long and TMaze Long Noise (Beck et al., 2019). Both tasks are structurally similar to our Passive T-Maze. In TMaze Long, the agent is forced to take the forward action until reaching the junction ( i.e ., the action space is a singleton except at the junction). At the junction, the agent is required to choose the goal candidate matching the color it observed at the starting position. Therefore, the memory length is the horizon T and the credit assignment length is 1 since only the terminal action can be credited. TMaze Long Noise introduces uniform noise appended to the observation space. Both tasks have a horizon and corridor length of 100 .

PsychLab (Fortunato et al., 2019) . This environment simulates a lab environment in first person, and can be sub-categorized into four distinct tasks. In all tasks, one or many images are shown, and the overall objective of all tasks is to correctly memorize the contents of the images for future actions. Images are passively observed by the agent, thus eliminating the need for long-term credit assignment, and instead focusing on long-term memory.

Spot the Difference (Fortunato et al., 2019). In this task, two nearly identical rooms are randomly generated, and the agent must correctly identify the differences between the rooms. The agent must navigate through a corridor of configurable length to go from one room to another. While this task was designed to test an agent's memory, effective temporal credit assignment is also necessary for an optimal policy. Similar in flavor to Active Visual Match (Hung et al., 2018) , an agent must first explore the initial room to be able to identify the differences in the second room. Rewards for this exploratory phase are only given near the end of the task once the agent has identified all differences. Long-term credit assignment is needed to incorporate these actions into an optimal policy.

Goal Navigation (Fortunato et al., 2019). A random maze and goal state is generated at the start of each episode. The agent is rewarded every time it reaches the goal state, and is randomly re-initialized in the maze each time it does so. Crucial information about the maze needs to be memorized throughout the episode for optimal performance, while exploratory actions may also be taken early for better quicker navigation later. Consequently, this environment fails to cleanly disentangle temporal credit assignment from memory.

Numpad (Parisotto et al., 2020; Humplik et al., 2019). This is a visual RL problem with continuous control. In this task, there is an N × N number pad, and the goal of the agent is to activate up to N 2 pads in a specific order. The order is randomized at the beginning of each episode. The agent does not know the order, and has to explore to figure it out. The agent will receive an immediate reward of 1 if it hits the next pad correctly; otherwise, the agent must restart. As a result, the credit assignment length is short, although the exact value is hard to determine. The optimal policy needs to memorize the order of N 2 pads during the whole episode, thus the memory length is upper bounded by the full horizon of 500 .

Memory Length (Osband et al., 2020, A.6) . The task is similar to our Passive T-Maze in terms of reward memory and credit assignment lengths. However, we did not adopt this task because the agent cannot change the observation through its action, making the task solvable by supervised learning. Moreover, in this task, the observation is i.i.d. sampled at each time step, making a transition memory length of only 1 . The authors set the horizon at a maximum of 100 .

Reacher-pomdp (Yang and Nguyen, 2021). In this task, the goal position in Reacher is only revealed at the first step. The optimal policy has to memorize the goal until it is reached. The policy memory length is thus long, but not necessarily the full horizon. The task provides dense rewards, so the credit assignment length is short, although the exact value is unclear. The horizon is set to 50 .

T-Maze (Lambrechts et al., 2021). This task has the same structure as T-Maze (Bakker, 2001). The RL algorithms are tested with the largest horizon of 200 .

Ballet (Lampinen et al., 2021). Ballet is a 2D gridworld task with the agent situated at the center. In its most challenging variant, there are 8 dancers around the agent, each executing a 16-step dance sequentially. A 48-step interval separates two consecutive dances. Consequently, the agent observes the environment for a minimum of 16 ∗ 8 + 48 ∗ (8 -1) = 464 steps. Following the final dance, the agent receives a reward upon reaching the dancer performing a specific routine. Hence, the optimal policy's memory length is at least 464 steps, and the reward memory length is capped by the horizon of 1024 . The agent's actions are constrained to the post-dance phase, resulting in a short credit assignment length determined by the time taken to reach the correct dancer.

HeavenHell (Esslinger et al., 2022; Thrun, 1999). This task has the same structure as our Active T-Maze. The agent has to first move south to reach the oracle and then move north to reach the heaven. The reward is given only at the terminal step. The credit assignment length is the full horizon, which is around 20 . The Gym-Gridverse and Car Flag tasks in their work have the same structures as HeavenHell but feature more complex observation spaces.

Memory Cards (Esslinger et al., 2022). The task can be viewed as a discrete version of Numpad. There are N pairs of cards, with each pair sharing the same value. The agent can observe a random card's value at each step, after which the card value is hidden. The agent then takes action to choose the paired card that matches the observed card. It will get an immediate reward of 0 if it makes the correct choice, otherwise, a reward of -1 . The credit assignment length is thus 1 . The minimal memory length is around O ( N ) to observe all cards' values, but hard to determine exactly due to randomness. The finite horizon is set to 50 .

Memory Maze (Pasukonis et al., 2022). In this environment, the agent is placed in a randomly generated N × N maze with N = 15 at most. There are K colored balls in the maze, and the agent must pick up the K colored balls in a specified order (randomized in every episode). As for the reward memory, the reward is given only when the agent picks up the correct ball and solely depends on current observation and action, thus m M reward = 1 . Although their longest horizon T can be 4000 , the credit assignment length is upper bounded by the size of the maze. In the worst-case scenario, an optimal policy can perform a naive search on the entire maze and uncover the maze plan to obtain a reward, which costs O ( N 2 ) steps. Nevertheless, the optimal policy memory length can be much longer than O ( N 2 ) steps and close to the horizon T . This is because the optimal policy may need to memorize the initial subsequence of observations to infer the maze plan when taking current actions.

POPGym (Morad et al., 2023). This benchmark comprises 15 tasks from which we examine several representative ones:

The diagnostic POMDP task, Repeat First, provides a reward based on whether the action repeats the first observation. Thus, the credit assignment length is 1 . Although this task indeed has a reward memory length of T , the optimal policy memory length can be as low as 2 simply by recalling the previous optimal action. The horizon T of this task can be 16 decks ( 831 steps).

In another diagnostic POMDP task, Autoencode, an agent first observes a sequence of cards, and then must reproduce the exact sequence in a reverse order to earn rewards. Here, the credit assignment length is 1 , and the optimal policy memory length is at most twice the number of cards, which is 6 decks ( 311 steps).

The control POMDP tasks, Stateless Cartpole and Pendulum, occlude velocities, making them similar to the '-P' tasks we previously evaluated. These tasks require short-term dependencies.

The game POMDP task, Battleship, requires the agent to deduce the locations of all ships on a board without observing the board. This requires a memory length upper bounded by the area of the board, i.e . 12 2 = 144 , and short-term credit assignment. Similarly, the Concentration game requires the agent to maximize the number of card matches in an episode. Although the reward memory length of horizon T , the optimal policy memory length can be restricted to the number of cards, i.e . 2 ∗ 52 = 104 .

Overall, this benchmark indeed emphasizes the evaluation of memory capabilities with short-term credit assignment, where the memory length is limited by the maximal horizon of around 831 .

Memory Gym (Pleines et al., 2023). This benchmark has a pixel-based observation space and a discrete action space for controlling movement within a gridworld. The Mortar Mayhem (MM) task is considered as one of the hardest tasks in this benchmark. MM requires the agent to first memorize a sequence of five commands, and then execute the exact commands in order. A reward is given when the agent successfully executes one command. Thus, the credit assignment length for optimal policy is at most the area of the gridworld, which is 5 ∗ 5 = 25 . The memory length is at most the full horizon, which according to their Table 3 is at most 135 .

B.3 Concrete Benchmarks for Credit Assignment

Active Visual Match (Hung et al., 2018) . The agent has to open the door to observe the color in the first phase, picks apples up in the second phase, and goes to the door with the corresponding color in the final phase. Thus, the policy needs to memorize the color.

Episodic MuJoCo (Liu et al., 2019; Ren et al., 2022). It is a decomposed episodic reward problem (Eg. 5). The terminal reward is the sum of the original Markov dense rewards. They tried the maximal horizon as 1000 .

Umbrella Length (Osband et al., 2020, A.5) . It is an MDP and the initial action taken by the agent will be carried through the episode. The terminal reward depends on the terminal state, which is determined by the initial action. Thus, the credit assignment length is the full horizon, at most 100 . There are also some random intermediate rewards that are independent of actions.

Key-to-Door (Raposo et al., 2021; Mesnard et al., 2021; Chen et al., 2021). The agent has to reach the key in the first phase, picks apples up in the second phase, and goes to the single door. Thus, the policy has little memory, while the terminal reward depends on whether the agent picks the key up in the first phase. In this sense, it can be also viewed as a decomposed episodic reward problem with some noisy immediate rewards. Raposo et al. (2021) tried the maximal horizon as 90 . The horizon in Mesnard et al. (2021); Chen et al. (2021) is unknown but should be similar.

Catch with delayed rewards (Raposo et al., 2021). It is a decomposed episodic reward problem (Eg. 5). The terminal reward is given by how many times the agent catches the ball in an episode. They tried the maximal horizon as 280 .

Push-r-bump (Yang and Nguyen, 2021). The agent has to first explore to find out the correct bump, and then move towards it to receive a terminal reward. The credit assignment length is thus long. The horizon is 50 .

C Experiment Details

C.1 Memory-Based RL Implementation

Our implementation builds upon prior work (Ni et al., 2022), providing a strong baseline on a variety of POMDPtasks, including Key-to-Door and PyBullet tasks used in this work. We denote this implementation 9 as POMDP-baselines . All parameters of our agents in our experiments are trained end-to-end with model-free RL algorithms from scratch. The following paragraphs provide detailed descriptions.

Critic Heads

Ot

Embedder

Observ.

Critic Head

MLPA

qt(ht,at)

Ot

Observ.

Prev.

Embedder

at-T

Embedder

Action

Sequence

Model

MLP

qt(ht;)

Prev.

Sequence

MLPB

q²(ht,at)

at-T

Embedder

Action

Model

Reward

Embedder

Reward

MLP

π(atIht)

Embedder

Actor Head

Sharing the sequence model in actor-critic with frozen critic parameters in the actor loss. Our agent architecture, depicted in Fig. 10, simplifies the one in POMDP-baselines by sharing the sequence model between the actor and critic. Previously, POMDP-baselines found sharing the RNNs between actor and critic causes instability in gradients, thus they adopt the separate sequence encoders for actor and critic. We examine their code and find that it is mainly due to the gradient bias in the actor loss. This issue is also discussed in prior (Yang and Nguyen, 2021) and concurrent works (Grigsby et al., 2023) on memory-based RL and related to sharing CNN encoders in visual RL (Yarats et al., 2021a,b). Below we formalize the issue and introduce our solution.

In the separate encoder setting, we consider an actor and a critic with distinct history encoders f ϕ π and f ϕ Q , mapping a history h into latent states z π and z Q , respectively. The history encoder parameters are denoted as ϕ π and ϕ Q . The actor is parameterized by π ν ( f ϕ π ( h )) and the critic is parameterized by Q ω ( f ϕ Q ( h ) , a ) with ν and ω denoting their respective MLP parameters. Considering the DDPG algorithm (Lillicrap et al., 2016) for training 10 and a tuple of data ( h, a, o ′ , r ) 11 , the loss of critic parameters is

where Q tar ( h, a, o ′ , r ) := r + γQ ω ( f ϕ Q ( h ′ ) , π ν ( f ϕ π ( h ′ ))) and θ denotes the stopped-gradient version of θ and h ′ = ( h, a, o ′ ) . The loss of actor parameters is

The total loss of actor and critic parameters is thus L sep ( ϕ π , ϕ Q , ω, ν ) := L Q ( ϕ Q , ω ) + L π ( ϕ π , ν ) . The gradients of L sep are

Now we consider an actor and a critic with a shared encoder denoted as f ϕ , mapping a history h into a latent state z . The history encoder has parameters ϕ . The actor is parameterized by π ν ( f ϕ ( h )) and the critic is parameterized by Q ω ( f ϕ ( h ) , a ) . The total loss of actor and critic

has such gradients:

Comparing the two gradient formulas across the two settings, we find the extra gradient terms using the same value of ϕ for ϕ π and ϕ Q in L sep :

To remove these extra gradient terms, we propose freezing the critic parameters ϕ and ω in the actor loss for the shared encoder setting, resulting in a modified loss function L sha-ours ( ϕ, ω, ν ) :

This offers a theoretical explanation for the practice of detaching critic parameters in the actor loss, adopted in prior works (Yarats et al., 2021a,b; Yang and Nguyen, 2021). It is worth noting that these works take an additional step of freezing the encoder parameters present in actions, resulting in an expression like Q ω ( f ϕ ( h ) , π ν ( f ϕ ( h ))) . In our preliminary experiments conducted on PyBullet tasks, we observed negligible differences in the empirical performance between their method and ours. Consequently, to maintain consistency with our theoretical insights, we have chosen to implement our approach L sha-ours in all of our experimental setups.

Implementation of sequence encoders in our experiments. We train both LSTM and Transformer sequence encoders with a full context length equal to the episode length on all tasks, except for PyBullet tasks where we follow POMDP-baselines to use a context length of 64 . The training input of sequence encoders is based on a history of t observations and t actions (with the first action zero-padded). The history is embedded as follows according to the POMDP-baselines . For state-based tasks (T-Mazes and PyBullet), the observations and actions are individually embedded through a linear and ReLU layer. For pixel-based tasks (Passive Visual Match and Key to Door), the observations (images) are embedded through a small convolutional neural network. For all tasks, the sequence of embedded observations and actions are then concatenated along their final dimension to form the input sequence to the sequence model.

For LSTMs (Hochreiter and Schmidhuber, 1997), we train LSTMs with a hidden size of 128 , varying the number of layers from 1 , 2 , 4 . We find single-layer LSTM performs best in T-Mazes and Passive Visual Match, and two-layer best in Key-to-Door, which are reported in Fig. 3 and Fig. 6. See Table 3 for details.

For Transformers , we utilize the GPT-2 model (Radford et al., 2019) implemented by Hugging Face Transformer library (Wolf et al., 2019). Our Transformer is a stack of N layers with H -headed self-attention modules. It includes causal masks to condition only on past context. A sinusoidal positional encoding (Vaswani et al., 2017) is added to the embedded sequence. The same dropout rate is used for regularization on the embedding, residual, and self-attention layers. We tune ( N,H ) from (1 , 1) , (2 , 2) , (4 , 4) . We find ( N,H ) = (1 , 1) performs best in Passive T-Maze and Passive Visual Match, and (2 , 2) best in Active T-Maze and Key-to-Door, which are reported in Fig. 3 and Fig. 6. See Table 4 for details.

For PyBullet tasks, we follow POMDP-baselines to use single-layer LSTMs and Transformers.

Implementation of RL algorithms in our experiments. We use SAC-Discrete (Christodoulou, 2019) for pixel-based tasks, and TD3 (Fujimoto et al., 2018) for PyBullet tasks, following POMDP-baselines . We use DDQN (van Hasselt et al., 2016) for T-Maze tasks, which outperforms SAC-Discrete in our preliminary experiments. We use epsilon-greedy exploration strategy in DDQNwith a linear schedule, where the ending epsilon is 1 T with T being the episode length. This ensures that the probability of always taking deterministic actions throughout an episode asymptotically approaches a constant:

where e is the base of the natural logarithm. This approach is critical to solving T-Maze tasks which strictly require a deterministic policy. Table 5 summarizes the details of RL agents.

C.2 Additional Results

Fig. 11 shows the learning curves of training Transformers with varying numbers of layers and heads.

Similar to the scaling experiments on Transformers we demonstrated in Sec. 5.2, we conduct an ablation study on scaling the number of layers in LSTMs from 1, 2 to 4 in credit assignment tasks. Fig. 12 shows the results that multi-layer (stacked) LSTMs do not help performance, aligned with empirical findings in stacked LSTMs (Pascanu et al., 2014).

0

1

2

3

env steps

×

10

6

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

success

Key-to-Door w/ credit assignment length of 60

1

2

4

0

2

4

env steps

×

10

6

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

success

Key-to-Door w/ credit assignment length of 120

1

2

4

0

2

4

6

env steps

×

10

6

0

.

0

0

.

1

0

.

2

success

Key-to-Door w/ credit assignment length of 250

1

2

4

0

2

4

6

env steps

×

10

6

0

.

00

0

.

05

0

.

10

0

.

15

success

Key-to-Door w/ credit assignment length of 500

1

2

4

0

100

200

300

400

500

(Easy)

Credit assignment length

(Hard)

-

1

.

0

-

0

.

5

0

.

0

0

.

5

1

.

0

return

Active T-Maze

1

2

4

100

200

300

400

500

(Easy)

Credit assignment length

(Hard)

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

success

Key-to-Door

→

→

C.3 Plotting Details

For all the learning curve plots in this paper, we use seaborn.lineplot (Waskom, 2021) to show the mean and its 95% confidence interval across 10 seeds. For the aggregation plots, we also use seaborn.lineplot to report the final performance, which is averaged over the evaluation results during the last 5% of interactions.

C.4 Training Compute

Here we report the training compute of the experiments in Sec. 5.1 and Sec. 5.2, because the PyBullet experiments (Sec. 5.3) are relatively cheap to compute. Each experiment run was carried out on a single A100 GPU and a single CPU core. For Transformer-based RL, the GPU memory usage is approximately proportional to the square of the context lengths, with a maximum usage of 4GB for Passive T-Maze with a context length of 1500 . For LSTM-based RL, GPU memory usage is roughly linearly proportional.

In our tasks, the context length equals the episode length (and also memory length), thus the total training time is proportional to num_episodes * context_length**2 * update_frequency for both Transformers and LSTMs. The update_frequency , set as 0 . 25 , is the ratio of parameter update w.r.t. environment step. Additionally, for multi-layer Transformers, the training time is roughly linear to the number of layers. Table 6 summarizes the training hyperparameters.

In Passive T-Maze with a memory length of 1500 , it took around 6 and 4 days to train Transformerbased and LSTM-based RL, respectively. In Passive Visual Match with a memory length of 1000 , it took around 5 days for both Transformer-based and LSTM-based RL.

D Broader Impacts

Our work enhances understanding of long-term memory capability in Transformer-based RL. This improved memory capability, while offering advancements in RL, may pose privacy concerns if misused, potentially enabling systems to retain and misuse sensitive information. As this technology develops, strict data privacy measures are essential. However, negative impacts directly tied to our foundational research are speculative, as we propose no specific application of this technology.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a set of metrics to measure and isolate the memory dependency of partially observable environments.\nThe authors illustrate that some currently existing benchmarks do not sufficiently isolate the memory of an agent.\nBased on their analysis they propose two versions of a T-Maze environment, one of which appropriately isolates memory.\nFurther, the authors investigate the effect of memory architecture on tasks that require long- and short-term memory. \nTransformer-based policies outperform recurrent policies on tasks that require long-term memory dependencies, \nwhile there seems to be no benefit on tasks that require short-term memory dependencies.\n\nSTRENGTHS:\n**Significance:**\n\nI fully agree with the authors that prior works often conflate different effects and do not sufficiently isolate the memory component.\nTherefore metrics that help disentangle memory from other environmental effects are crucial.\n\n**Originality:**\n\nThe authors propose novel metrics that evaluate and isolate different aspects of POMDPs in RL, such as credit assignment and long-term memory.\nThe 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.\n\n**Quality:**\n\nQuality of theoretical contributions is high.\n\nWEAKNESSES:\n**Scalability:**\n\nI 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.\nFor example, l_{mem}(\\pi^\\asterisk) requires access to the optimal policy.\nFor 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?\nHow is it computed for procedurally generated environments, where the optimal policy varies between levels?\nHow would this scale to more complex environments, given that human demonstrations may not necessarily be optimal?\n\n**Relevance of proposed environments:**\n\nThe authors propose two environments, namely Passive T-Maze and Active T-Maze to disentangle credit assignment from memory.\nThe 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].\nPassive 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.\nThere 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.\n\n**Interpretation of results:**\n\nThe authors mention that LSTM starts to falter at a memory length of 250 for Passive T-Maze.\nFigure 2, however, shows that LSTM can solve the same environment with a memory length of 750.\nThis suggests that some other effects might influenced LSTM training, such as hyperparameter tuning.\nMaybe the authors can improve the results of LSTM by further tuning of the method?\n\n**Clarity:**\n\nThe 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.\n\n**Missing relevant work:**\n\nThere is prior work that derives theoretical bounds on approximation error of history-based methods [2].\nFurther, other history-based approaches such as a hierarchical Transformer memory [3] and pretrained models as a memory module in RL [4], are not cited.\nIf applicable, adding these method in their experiments would strengthen the paper and might yield some more interesting findings.\n\n[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/\n\n[2] Gandharv Patil et al., On learning history based policies for controlling markov decision processes. 2022.\n\n[3] Andrew Lampinen et al., Towards mental time travel: a hierarchical memory for reinforcement learning agents. NeurIPS 2021\n\n[4] Fabian Paischer et al., History compression via language models in reinforcement learning. ICML 2022\n\nQUESTIONS:\n- The authors conclude that Transformers exhibit worse sample efficiency than LSTM on short-term memory tasks.\nIs there an intuition why that would be the case?\nOn Passive Visual Match, which requires short-term memory according to the authors, a GPT reaches higher return compared to LSTM.\nCould 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?\n\n- 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.\nThe authors draw an analogy here which lacks support.\nHow do long-term memory tasks relate to the supervised learning setup on large datasets?\nDo the authors imply here, that there are only long-term dependencies present in large datasets used for supervised learning?\nAlso, Morad et al. 2023 specifically advises against drawing comparisons between supervised learning and RL.\n\n- What is the observation space for the Active/Passive T-Mazes?\n\n- 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.\n\n- 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.\n\n- Figure 5: it would be good to clarify what is meant by \"optimal policies that lack (long-term) credit assignment\"\nI assume it refers to a markovian policy, since the return on T-Maze is 0.5.\n\n- Do all figures show mean and standard deviation?\nThis should be mentioned in the figure captions.\n\n- The last three sentences in Definition 1 could be moved to Definition 2, since they are about memory length and not context length.\n\n- 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?\n\n- The symbol k is reassigned repeatedly which is a bit confusing, the authors may consider using different symbols.\n\n- Will the authors make the code for reproducing their results and computing their metrics publicly available?\n\nLIMITATIONS:\nThe authors have adressed limitations regarding computation of the metrics in more complex environments. \nHowever, 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?\nIn the case of procedurally generated environments (which are commonly used nowadays) would that require extensive \"labeling\" by humans?", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nMemory 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. \n\nCredit 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.\n\nSTRENGTHS:\n- This definitions of memory length and credit assignment length appear intuitive and sound.\n- Theorem 1 appears clear and mathematically sound.\n- 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.\n- 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.\n- Overall, the experiments are well-designed, clear, and straightforward, providing a substantial contribution to the field of decision transformer research.\n\nThis work is highly novel and is the first to specifically investigate the benefits of transformers in either memory or credit assignment.\n\nWEAKNESSES:\n- 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. \n- It would be nice to see experimental evaluation of Transformers and LSTM in additional environments other than the toy tasks for better context.\n\nQUESTIONS:\nAbove\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- The question that was posed was quite unique and interesting.\n- The authors do a convincing empirical study of which settings transformers perform better in, and the conclusions are aligned with empirical intuition.\n- Conducted analysis of previous benchmarks is extensive, and the authors have a good experiment design.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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.\n- 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. \n- 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.\n\nLIMITATIONS:\nThe authors have adequately addressed limitations.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

comments: 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?

LIMITATIONS:
The authors have adequately addressed limitations.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-07