When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment
Authors: Tianwei Ni
Michel Ma
Benjamin Eysenbach
Pierre-Luc Bacon
Venue: NeurIPS 2023 oral (2023)
Keywords: Memory-based RL
Transformers
Credit Assignment
Online RL
Model-free RL
TL;DR: Transformers can help learn long-term memory but not long-term credit assignment in online 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.

enhance the When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment
Tianwei Ni
Michel Ma
Mila, Université de Montréal
Mila, Université de Montréal michel. mawmila. quebec
tianwei. nimmila. @mila.quebec
Benjamin Eysenbach
Pierre-Lue Bacon
Princeton University eysenbach@prrincetonn edu
Mila, Université de Montréal pierre-luc bacon quebec
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 depen-- dencies. 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-basd 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 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 important area for future research and benchmark design. Our code is open-sourced"
RL methods
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-basse 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 et al. (2023); Agarwal et al. (2023) for the recent surveys).
methods can achieve excellent results in offline RL (Chen et al., 2021; Janner
However, the underlying reasons why Transformers achieve excellent results in the RL setting remain a mystery.
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 distant past event at the current time (Blankenship, 1938; Dempster, 1981).
'https://githan thb.. httppp///gtthbb.ccc///m/////////
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

while credit assignment is the ability to determine when the actions that deserve current credit oc- curred (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-tent memory and credit assignment. Second, there RL (Osband et al., 2020). The absence of quantifiable measures creates ambiguity in the concept of "long-term" often found in prior works.
lack of rigorous definition of memory and credit assignment in
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 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 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 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- RL algorithms with LSTMs or Transformers on 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 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) Mo (O.A, P, R. R,,, an agent receives an observation o, E O at step t E (1.....T), takes an action a A based on the observed history h1s == (O1:t. 01:--1) H, and receives a reward Te ~ (e az) and the next observation 0++1 ~ P(- hig, at). The initial observation h1S == 01 follows P(or). The total horizon is T E N+ [+oo) and the discount factor 7 E [0. 1] (less than 1 for infinite horizon). In a Markov decision process (MDP) Ms = (S, P. R. y, T), the observation of and history h1 are replaced by the state s S. In the most generic case, the agent composed of a policy (a, 1x))) and value function Q" (hi,, at). The optimal value function Q* 4x) satisfies Q" (ha:,, ax) = E[re hi:e, (elty/EEr-PP(ha. a deterministic optimal policy x² (Iax) = argmax,, o, Q"(h1,,a,, 4z). Let IT be the space of all optimal policies for POMDP M.
[maxerx: Q" ("h1a++1, a(+1)] and induces
"The classic notion of POMDPs (Kaelbling al., 1998) includes state space, but we omit it because we assume that is unknown.
For any k E {l...., t}. let ki denote the sequence (o.-k..... o.). and similarly for a ket. In particular, let hise of to align with the MDP setting, and let hre = 0 for t' > and hea = his for t' < 1.
2

---

Below, we provide definitions for context-basee 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 letx(1) observations and actions as inputs, expressed as [a, (a Here, (TT) N represents the policy context length. For simplicity, we refer to context-based policies as "policies". Let I1 be the space of all policies with context length of k. Markovian policies form II, while Hr. the largest policy space, contains an optimal policy. Recurrent policies, recursively taking one input time, belong to TT,.; while Transformer-based policies have limited context length. Context-based value function Qn. of a POMDP M is the expected return conditioned on the past n observations and actions, denoted as Q% (h definition, Q* (hist; 0z) = (h12, at), Vhia,
Ge) = Er.M |E= 7 1h-n++1:: at]. By
The sum of n-step rewards. Given POMDP M and a policy x. the expectation of the discounted sum of n--tep rewards denoted as G (Pise at) F..M [ p his (t) By definition, GT (brat, de)
E [rx has, as] (immediate reward), and G -- (h1:, 0z) Q" (h13, ax) (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 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 me meemara. For a POMDP M, meward the smallest N such that the expected reward conditioned on recent n observations is the same the one conditioned on onfull history, i.c., Eir, hig, oz] = E[r, ht-n+1st ar], Vt, hi: at.
m
Definition 1.B (Transition memory length mmrnsit) For POMDP M, m is the smallest n E N such that next observations are conditionally independent of the rest history given the past n. observations and actions, i.e., Ot+1 1L h1.1, at hh-n+1:1, Ot-
Next, policy memory length defined based on the intuition that while 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 (mmm(T) The policy memory length (mem(() (T) of a policy т the minimum horizon n E (0....,,etx((T)) (T)) such that its actions are conditionally independent of the rest history given the pasl n observations and actions, i.e., 0z 1I. (n-lax(()+111 ht-n+ 1st-
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 L.D (Value memory length of policy (T)). A policy has its value memory length Ivauue (T) as the smallest such that the context-based value function is equal to the value function, ie., (f at) = Q" (hiat, 0t), Vt, his, Ot-
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.
M 1mem (T²) < L.M.. (TT") < max(m Teward: mi Mi = m.M
VIT" E arem (1)
xEll
3

---

Thm. suggests that the constant of mM can serve proxy to analyze the policy and value m lengths required in task, which are often challenging to directly compute, as we will demonstrate in Sec. 3.*
2.2 Credit Assignment Lengths
Classically, temporal credit assignment is defined with backward view (Minsky, 1961; Sutton, 1984) how previous actions contributed to current reward. The backward view is specifically studied through gradient back proppppaaattion 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 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(b1a: R)). Given a policy π in task, the credit assignment length of a history hia quantifies the minimal temporal distance n², such that a greedy action at is strictly better than any non-greedy action at in terms of their n-step rewards G Formally, let A == argmaxa, Q" (hi:t 0z),, define
chb15 == min
{n 3az E S.L. C. (hia,,))) 02) (7 d'), Val, & A7}
(2)
15n5T -441 141
Definition 2.B (Credit assignment length of a policy c(T) and of task M) Credit assignment length of a policy и represents the worst case of credit assignment lengths of all visited histories. For- mally, let dg be the history occupancy distribution of n, define c(π) == h e(hi1:; x)-
Further, we define the credit assignment length of a task M, denoted as e,M as the minimal credit assignment lengths over all optimal policies', i.e. e. := mink EMT c(TT).
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)) = T. This is because the credit assignment length of the initial state 81 is always c($1;T)= = T, as the performance gain of any greedy action at 81 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 (s:) + + +a(; 75+(+)+) - (s:).. This formula back-propagates the information of current reward Tz and next value V 15(+1) to current value (sz). 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 on memory lengths (Thm. 1) and credit assignment length cM (Def. 2.B), we can perform a quantitative analysis of various environments. These environments include both abstract problems and concrete benchmarks. this section, we aim address two questions: (1) Do prior environments actually require long-term memory credit assignment? (2) Can we disentangle memory from credit assignment?
*Thm. 1 also applies to MDPs, where it provides quick proof of the well-known fact that optimal policies and values of MDPs can be Markovian: mi MM are both at most 1, the policy and value memory lengths are also at most
'The n is guaranteed to exist because greedy action must be the best in T - + 1-step rewards (i.e. Q-value).
"If Ai = A. iee. all actions are greedy, then the set of n contains all integers between and T' - + 1.
II there multiple optimal policies, c(x²) may be not unique. See Appendix A for example.
4

---

collect Table 1: Estimated memory and credit assignment lengths required in prior tasks, defined in Sec. 2. The top block shows tasks designed memory, while the bottom one shows tasks designed for credit assignment. All tasks have horizon of T and discount factor of E (0, 1). The T column shows the largest value used in prior work. The terms "long" and "short" are relative to horizon T. We manually estimate these lengths from task definition and mark tasks purely evaluating long-term memory or credit assignment black.
Task M
T
(mem(xt))
mM
e.
Reacher-pomdp (Yang and Nguyen, 2021)
50
long
long
short 1
Memory Cards (Esslinger ct al., 2022)
50
long
long
TMaze Long (Noise) (Beck et al., 2019)
100
T
T
1
Memory Length (Osband et al., 2020)
100
T
T
1
Mortar Mayhem (Pleines et al., 2023)
135
long
long
S
Autoencode (Morad et al., 2023)
311
T
T
1
Numpad (Parisotto et al., 2020)
500
long
long
short
PsychI.ab (Fortunato et al., 2019)
600
T
T
short
Passive Visual Match (Hung et al., 2018)
600
T
I
18
Memory
Repeat First (Morad et al., 2023)
831
T T
1
Ballet (Lampinen et al., 2021)
1024
2 464
short
Passive Visual Match (Sec. 5.1; Our experiment)
1030
T
T
I 18
17 MiniGrid-Memoor (Chevalier-Boisver et al., 2018)
1445
< 51
T
I
51 1
Passive T-Maze (Eg. 1; Ours)
1500
T
T
152 Memory Maze (Pasukonis et al., 2022)
4000
long
long
< 225 T T
HeavenHell (Esslinger al., 2022)
20
T
T-Maze (Bakker, 2001)
70
T
T
Goal Navigation (Fortunato et al., 2019)
120
T
T
T-Maze (Lambrechts al., 202
200
T
T
T
Spot the Difference (Fortunato et al., 2019)
240
T
T
PyBulllt-P benchmark (Ni et al., 2022)
1000
2
short
PyBullet-V benchmark (Ni et al., 2022)
1000
short
short
short
Umbrella Length (Osband et al., 2020)
100
T
Credit Assignment
Push-r-bump (Yang and Nguyen, 2021)
50
long
long
long
Key-to-Door (Raposo et al., 2021)
90
short
T
T
Delayed Catch (Raposo al., 2021)
280
T
T
Active T-Maze (Eg. 2; Ours)
500
T
T
T
Key-to-Door (Sec. 5.2; Our experiment)
530
short
T
T
Active Visual Match (Hung et al., 2018)
600
T
T
T
Episodic MuJoCo (Ren et al., 2022)
1000
T
T
3.1 How Long Do Prior Tasks Require Memory and Credit Assignment?
To study memory 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, representative problems and benchmarks with a similar analysis in Sec. 3.2, to figure out the quantities M. related to memory lengths and eM related to credit assignment length.
Regarding abstract problems, we summarize the analysis in Appendix 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 applics to episodic MuJoCo (Ren ct al., 2022) and Catch with delayed rewards (Raposo et al., 2021) l). We show that it has policy memory length <mem(x") at most I while requiring M and 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-Mede Medina et al., 2019) or execution (Derman et al., 2021) or observation (Katsikopoulos and Engelbrecht, 2003) with n steps all require m.M m of n, while only delayed rewards and execution may be challenging in credit assignment.
As for concrete benchmarks, Table summarizes our analysis, with detailed descriptions of these environments in Appendix First, some tasks designed for long-term memory (Beck al., 2019; Osband et al., 2020; Hung al., 2018; Yang and Nguyen, 2021; Esslinger al., 2022) credit
5

---

assignment (Osband et al., 2020) indeed solety evaluate the respective capability, Although some memory tasks such as Numpad (Parisotto et al., 2020) and Memory Maze (Pasukonis al., 2022) do require long-term memory with short-term credit assignment, the exact memory and credit assignment contrast, our Passive T-Maze introduced in Sec. 3.2, provides clear understanding of the memory and credit assignment lengths, allowing us to successfully scale the memory length up to 1500.
lengths are hard to quantify due to the complexity.
Conversely, other works (Bakker, 2001; Lambrechts et al., 2021; Esslinger et al., 2022; Hung et al., 2018; Raposo et al., 2021; Ren 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 as- signment independent concepts by providing two examples, all with a finite horizon of I and no discount. The examples have distinct pairs of (max(r.M M.M ) M) (T, 1), (T, T), re- spectively.
G1
0
S
J
G2
Corridor Length
We adapt T-Maze (Fig. 1), a POMDP envi- ronment rooted in neuroscience (e'Keefe and Dostrovsky, 1971) and AI (Bakker, 2001; Os- band et al., 2020), to allow different credit signment lengths. Passive and active concepts borrowed from Hung al. (2018).
Figure I: T-Maze task. In the passive version, the agent starts at the oracle state ("S" and "O" the same state), while the active version requires that the agent navigate left from the initial state to collect information from the oracle state.
Example 1 (Passive T-Maze). T-Maze has a 4-direction action space {c, 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 G1 and G2. O has horizontal position of 0, and J has L. The agent observes null N except at states J, O G1, Gz. The observation space (N, O,,, G2} G2,G2] is two-dimensional. At state O, the agent can observe the goal position G E [G1, G2). uniformly sampled the beginning of an episode. The initial state S. The transition is deterministic, and the agent remains in place it hits the wall. Thus, the horizontal position of the agent 2t E NN is determined by its previous action sequence 01:t-1-
In a Passive T-Maze M Mpassive O = S, allowing the agent to observe G initially, and L = -1 2()=+12(111 for t S T T- and Rr ((((bbrrr))))) 1 = G). The unique optimal policy xt" moves right for T 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
Rewards are given by Re
Example 2 (Active T-Maze). In an Active T-Maze Mactiee O is one step left of and L - 2. The unique optimal policy moves left to reach O, then right to J. The rewards R₁ 0, Re az) = ((re:22t1)-11 for 2 < t < T - 1, and Rr ar) = 1 G). The optimal Markovian and worst policies have the same expected returns as in the passive setting.
T-2
M Mpsssive 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 I' 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, Mactive also requires 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 RI., 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 RI, with representative works shown in Table 1. The work most closely related to ours is bsuite (Osband et al., 2020), which

---

Phase 1
Phase 2
Phase
Phase
Phase
Phase
Figure 2: Pixel-based tasks, Passive Visual Match (left) and Key-to-Door (right), evaluated in experiments. Each task has three phases, where adjust the length of Phase 2 and keep the lengths of the others constant. The agent (in beige) only observes nearby grids (within white borders) and cannot pass through the walls (in black). In Passive Visual Match, the agent observes a randomized target color (blue in this case) in Phase 1, collects apples (in green) in Phase 2, and reaches the grid with the matching color in Phase 3. In Key-to-Door, the agent reaches the key (in brown) in Phase 1, collects apples in Phase 2, and reaches the door (in eyan) in Phase 3.
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 '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 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 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.
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-Dooe (Raposo 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-baa based RL on standard POMDP benchmarks that only require short-term dependencies in Sec. 5.3, with focus on their sample efficiency.
We focus 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, the RL algorithm to better control the exploration strategy. In pixel-based tasks, we use CNNs as observation embedders and SAC-Discret (Christodoulou, 2019) 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.
DDQN (van Hasselt et al., 2016) with epsilon-greedy exploration as
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. (left), Transformer-basse agents consistently solve the task requiring memory lengths from 50 up to 1500. To the best of our knowledge, this achievement
7

---

T-Mase
Paesive Vanual Masch
GPT
0,50
LSTM
Optinal agenl w/o mennory
0
200
400
600
800
1000
1200
1000
600
(Eazy) <- Memory length -> (Laard)
(Easy) e Menoory length (laad))
Figure 3: Transformer-based RL outperforms LSTM-based RL in tasks (purely) requiring
long-term memory. Left: results in Passive T-Mazes with varying memory lengths from 50 to 1500; Right: results in Passive Visual Match with varying memory lengths from 60 to 1000. We also show the performance of the optimal Markovian policies. Fach data point in the figure represents the final performance of an agent with the memory length indicated on the x-axis. All the figures in this paper show the mean and its 95% confidence interval over 10 seeds.
Pleetrr I-blion
Zimbrr E-31ap imath
Dtaebrr T-btipx
mermery
1.0
1.0
055
m0m
60
AV
-0.
-05
-B5
-10
-10
-10
XII
XIIP
X01
Figure 4: Transformer-based agent can solve long-term memory low-dimensional tasks with good sample efficiency. In Passive T-Maze with long memory length (from left to right: 500, 750, 1000, 1500), Transformer-basee agent is significantly better than LSTM-base agent.
may well set a new record for the longest memory length of a task that an RL agent can solve. In contrast, LSTM--asee 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-baser 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 al., 2018) has a more complex reward function: zero rewards for observing the color in Phase 1, immediate rewards for apple picking in Phase 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 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. (right), Transformer- outperform LSTM-bbssd o in success rates when the tasks require memory length over 1 120, yet their success rates are close to that of the optimal Markovian policy (around 1/3). Interestingly, Transformer-base 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.
agents slightly
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 memory- based 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-hh based and LSST--based agents fail to solve the task when the credit assignment (and memory) length extends to just 250. The failure of Transformer-basee credit assignment. This is because they possess sufficient memory capabilities (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., 202 1), where both Transformer-base and LSTM-
agents on Active T-Mazes suggests the bottleneck of
solve this task
'It is worth noting that at the memory length of 1000, LSTM-based agents encountered NaN gradient issues in around 30% of the runs, which were not shown in the plots.
8

---

Passiee Vieaal blask Lanuary loack
Fander Vieea Vaic my cerwary leph In8
Zuaii Vinall Skadd massaay kaghi Ine
200
LSTu
200
0.4
07
130
0.1
100
11.0
022
0.4
0.5
11.5
1.0
01.
0.2
0.5
588
100
Figure 5: Transformer-based agent is more sample-efficient in long-term memory high- dimensional tasks. In Passive Visual Match with long memory length (750 on the left; 1000 on the right), Transformer-based agents show a slight advantage over LSTM-based agents in the memory capability (as measured by success rate), while being much higher in total rewards (as measured by return).
Ariive T-Mae
Koyor-Door
1.0
1.0
0.5
returs
3.5
0.0
GPP
LSTM
0.5
Openaal agent w/o creans assignment
0.0
100
200
300
400
(Rag)) e- Credit ase pp meeedt kengh
(Harl)
(Erey) 4- Chedit resignment longth (Hari)
Figure 6: Transformer-based RL improves temporal credit assignment compared to LSTM- based RL, but its advantage diminishes in long-term scenarios. Left: results in Active T-Mazes with varying credit assignment lengths from 20 to 500; Right: results in Key-to-Door with varying credit assignment lengths from 60 to 500. We also show the performance of the optimal policies that lack (long-term) credit assignment - taking greedy actions to maximize immediate rewards. Each data point in the figure represents the final performance of an agent in task with the credit assignment length indicated on the x-axis, averaged over 10 seeds.
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 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 RI., 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 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.
Mour mork moga a. Inggl
Ariuen mage In
0.23
0.3
-3.0
130 730
130
216
316
210
(Eag/ Crolk wwggaard D4glh (lae)
(2a000 Cledl awwggnase Deggla (1lx)))
Figure 7: Increasing the number of layers (and heads) in Transformers aids in temporal, but not long-term, credit assignment. The left two figures show the final results from Active T-Maze and Key-to-Door, varying the number of layers and heads in Transformers from 1, 2 to 4. The right two figures show the associated learning curves in Active T-Maze with different credit assignment lengths.
5.3 Transformers in RL Raise Sample Complexity in Certain Short-Terne 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 "-p" tasks) or velocities (referred to as "-V" tasks) of the agents are revealed, with dense rewards. As 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

---

X10
X19
Chectal-P
Eoppan-7
X303
0.0
0.0
0.5
Lo
1.5
0.0
0.3
1.0
1.5
0.0
0.3
1.0
00
355
1.0
XXL
con,ttope
XI0
enrutrpe
Auc-V
X380
CheetahVV
Hoppee-V
X30
Walka-V
0.75
0.70
0.00
0.5
10
x
XX0
Figure 8: Transformer-based RL is sample-inefficien compared to LSTM-based RL in most PyBullet occlusion tasks. The top row shows tasks with occluded velocities, while the bottom row shows tasks with occluded positions. Results are averaged over 10 seeds.
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-bics 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
We provide 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 train in RL agents. The worst compute complexity of eM can be ((S)) T for a finite MDP, while the complexity of m-M can be O([]O) for a finite POMDP. In addition, cM 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 does not take into account the memory capacity (i.e., the number of bits) and robustness observation noise (Beck et al., 2019).
On the empirical side, we evaluate Transformers on 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 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-bast 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 short- term dependencies. While Transformers are powerful tools in RL, especially for long-term memory tasks, they are not a universal solution to RI. challenges. Our results underscore the ongoing need for careful task design and the continued advancement of core RL algorithms.
RL. While Transformer-basee agents excel in tasks requiring
10

---

Acknowledgements
This research was enabled by the computational resources provided by the Calcul Québec (calcul quebec. (calculquebec. ca) and the Digital Research Alliance of Canada (alliancecan. ca). We thank Pierluca D' D'Oro and Zhixuan Lin for their technical help. We thank Jurgis Pasukonis and Zhixuan Lin for correcting some factual errors. We also thank anonymous reviewers, Pierluca D'Oro, Siddhant Sahu, Giancarlo Kerg, Guillaume Lajoie, and Ryan D' O'Oazzio for constructive discussion.
References
Agarwal, 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
J. A. Arjona-Medina, M. Gillhofer, M. Widrich, Unterthiner, J. Brandstetter. and S. Hochreiter. RUDDER: return decomposition for delayed rewards. Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, 2019. 5, 17
B. Bakker. Reinforcement learning with long short-term memory. In Advances Neural Information Processing Systems 14 [Neural Information Processing Systems: Natural and Synthetic, NIPS 2001, December 3-8, 2001, Vancouver, British Columbia, Camada), 2001. 5, 6, 18, 19
J. 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. 10,
M. G. Bellemare, Y. Naddaf, J. Veness, and M. Bowling. The arcade learning environment: An evaluation platform for general agents. J. Artif. Intell. Res., 2013.
A. B. Blankenship. Memory span: a review of the literature. Psychological Bulletin, 35(1):1, 1938.
C. Chen, Y. J. Yoon, and S. Ahn. Transdrcamer: Reinforcement learning with transformer world models. arXiv preprint arXiv:2202.00 09481, 2022.
L. Chen, K. La, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforeement learning via sequence modeling. In Advances in Neural Information Processing Systems 34: Annual Conference Neural Information Processing Systems 2021, NeurlPS 2021, December
6-14, 2021, virtual, 2021. 1, 7.
M. Chevalier- Boisvert, L. Willems, and S. Pal. Minimalistic gridworld environment for openai gym. https: // //github. www.//oai/mcc///c/caoc.oo.oooeee som-mmniirid, 2018. 5, 18
Christodoulou. Soft actor-critic for discrete action settings. arXiv preprint 191.:21910007207, 2019. 7,
F. N. Dempster. Memory span: Sources of individual and developmental differences. Psychological Bulletin, 89 (1)63, 1981.
E. 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
K. Esslinger, R. Platt, and C. Amato, Deep transformer q-networks for partially observable reinforcement learning. arXiv preprint arXiv:2206.01073 2022. 5, 6, 7, 19
J. Ferret, Marinier, M. Geist, and Pietquin. Self-attentional credit assignment for transfer reinforcement leamning In Proceedings of the Twenty-Ninn Ninth International Joint Conference on Artificial Intelligence, IJCAI 2020, 2020. 7
M. Fortunato, M. Tan, R. Faulkner, Hansen, A. P. Badia, G. Buttimore, C. Deck, J. 7. Leibo, and C. Blundell. Generalization reinforcement learners with working and episodic memory. In Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NewrIPS 2019, December 8-14, 2079, 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. Stockhoim, Sweden, July 10-15, 2018, 2018. 23
S. Gershman and N. D. Daw. Reinforcement learning and episodic memory in humans and animals: an integrative framework. Annual review of psychology, 68: 101-118, 2017.
J. Grigsby, L. Fan, and Y. Zhu. Amago: Scalable in-context reinforcement learning for adaptive agents. arXiv preprint 2.Xix:3310009777 2023. 21
D. 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
11

---

S. Hochreiter and J. Schmidhuber. Long short-term memory. Neural computation, 98): 17800 1997. 22 J. Humplik, A. Galashov, L. Hasenclever, P. A. Ortega, Y. W. Tch, and N. Heess. Meta reinforcement learning as task inference. arXiv preprint arXiv:I905. 19045064444,, 2019. 19
C. Hung, T. P. Lillicrap, J. Abramson, Wu, M. Mirza, F. Carnevale, A. Ahuja, and G. Wayne. Optimizing agent hehavior over long time scales by transporting value. arXiv preprint arXiv: 1810.06721, 2018. 2, 5, 6, 8, 18, 20
B. 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, Quiambao, P. Pastor, L. Luu, K. Lec, Y. Kuang, S. Jesmonth, N. J. Joshi, K. Jeffrey, R. Ruano, Hsu, K. Gopalakrishnan, B. David, A. Zeng, and C. K. Fu. Do as can, not as I say: Grounding language robotic affordances. In Conference on Robot Learning, CoRL 2022, 14-18 December 2022, Auckland, New Zealand,
2022. 1
M. 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, NeurlPS 2021, December 6-14, 2021, virtual, 2021. 1
L. P. Kaclbling, M. I.. I Littman, and A. R. Cassandra. Planning and acting in partially observable stochastic domains. Artif. Intell.. 1998. 2
J. Kaplan, S. McCandlish, T. Henighan, T. B. Brown, B. Chess, R. Child, S. Gray, A. Radford, Wu, and D. Amodei. Scaling laws for neural language models. arXiv preprint arXiv:2001 0,8361, 2020. 9
K. Katsikopoulos and S. E. Engelbrecht, Markov decision processes with delays and asynchronous cost
collection. IEEE Trans. Autom. Control., 2003. 5, 18
N. R. Ke, A. Goyal, O. Bilaniuk. Binas, M. C. Mozer, C. Pal, and 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, NewrlPS 2018, December 3-8, 2018, Montréal, Canada, 2018. 4
G. 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.00000 2021. 5, 6.
19
A. 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 2212222895510 2021. 1,5, 19
D. Lee. S. Zhang, A. Fischer, and Y. Bengio. Difference target propagation. In Muchine Learning and Knowledge Discovery in Databases European Conference, ECML PKDD 2015, Porto, Portugal, September 7-11, 2015. Proceedings, Part I, 2015. 4
K. Lee, O. Nachum, M. Yang. L. Lec, D. Freeman, S. Guadarrama, Fischer, W. Xu, E. Jang, H. Michalewski, and Mordatch, Multi-gam decision transformers. NeurlPS, 2022. 1
W. Li, H. Luo, Z. Lin, C. Zhang, Z. Lu, and D. Ye. A survey on transformers in reinforcement learning. arXiv preprint Ar0i X0102201 2023.
T. 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
Y. 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.1344: 2019. 7,
L. C. Melo. Transformers are reifforennr learners. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, 2022.
L. Meng, R. Gorbet, and D. Kulic. Memory-based deep reinforcement learning for pomdps. In JEEE/RSJ International Conference on Intelligent Robots and Systems, IROS 2021, Prague, Czech Republic, September 27- Oct. 1, 2021, 2021.9 9
T. 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. Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, 2021. 2, 10, 20
V. Micheli, E. Alonso, and F. Fleuret. Transformers sample efficient world models. arXiv preprint arXiv:2209.000 00588, 2022.
M. Minsky. Steps toward artificial intelligence. Proceedings the IRE, 49(1):8 30, 1961. 4
12

---

S. D. Morad, R. Kortvelesy, M. Bettini, S. Liwicki, and A. Prorok. Popgym: Benchmarking partially observable reinforcement learning. arXiv preprint (rXiv:2303.0185/ 2023. 5, 7, 10, 20
T. 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, 7. 10, 21
J. O'Keefe and J. Dostrovsky. The hippocampus as a spatial map: Preliminary evidence from unit activity in the freely-moving rat. Bruin research, 1971. 6
I. 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
L. 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
E. Parisotto, H. F. Song, J. W. Rac, R. Pascanu, Ç. Gülşchre, S. M. Jayakumar, M. Jaderberg, R. L. Kaufman, A. Clark, S. Noury, M. M. Botvinick, N. Heess, and R. Hadsell, Stabilizing transformers for reinforcement leamning In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, 2020. 5, 6, 7, 19
R. 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
J. Pasukonis, T. P. Lillicrap, and D. Hafner. Evaluating long-term memory in 3d mazes. arXiv preprint arXiv:2210.13388 2022. 6,6, 19
M. Pleines, M. Pallasch, F. Zimmer, and M. Preuss. Memory gym: Partially observable challenges to memory- based agents. In The Eleventh International Conference on Learning Representations, 2023. 5, 7, 20
A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 1(8)9, 2019. 1, 22
D. 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.1245 2021. 5, 6, 7, 20
Z. 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
J. Robine, M. Höftmann, T. Uelwer, and S. Harmeling. Transformer-bssed world models are happy with 100k interactions. arXiv preprint arXiv:2333.077 1090,.0700900000 2023. 1
R. S. Sutton. Temporal credit assignment in reinforcement learning. University of Massachusetts Amherst, 1984. 2, 2,5
R. S. Sutton. Learning to predict by the methods of temporal differences. Mach. Learn., 1988.
R. 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
E. Todorov, T. Erez, and Y. Tassa. Mujaco: A physics engine for model-based control. In 2012 IEEE/RS) International Conference on Intelligent Robots and Systems, IROS 2012, Vilamoura, Algarve, Portugal, October 7-12, 2012, 2012. 2
H. 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
A. Vaswani, N. Shazeer, N. Parmar, J. Uszkorcit, L. Jones, A. N. Gomez, L. Kaiser, and Polosukhin. Attention all you need. 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,
M. L. Waskom. seaborn: statistical data visualization. Journal of Open Source Software, 6(60))3021, 2021. doi: 10.21 035/) URL https: / //doi. orgg/1.. 21105/ /os..03021........ 24
T. 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 Statc-of-the-an natural language processing. arXiv preprint 2019. 22
Z. Yang and H. Nguyen. Recurrent off policy baselines for memory- based continuous control. arXiv preprint arXiv:2100. 12628, 2021. 5, 19, 20, 21,
13

---

D. Yarats, I. Kostrikov, and R. Fergus. Image augmentation is all you need: Regularizing deep reinforvement learning from pixels. In 9th International Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021a. 21
D. Yarats, A. Zhang. 1. Kostrikov, B. Amos, J. Pineau, and R. Fergus. Improving sample efficiency in model- free reinforcement learning from images. In Thirty-Fifth AAAI Conference Artificial Intelligence. AAAI 2021, Thirty-Thiss Third Conference on Innovative Applications of Artificial Intelligence, IAAI 2021, Eleventh Sympoxium on Educational Advances in Artificial Intelligence FAAI 2021, Virtual Event, Fehruary 2-9, 2021, 2021b. 21,22
Zheng, Zhang, and Grover. Online decision transformer. International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, 2022. 1,7
E. A. Zilli and M. E. Hasselmo, Modeling the role of working memory and episodic memory in behavioral tasks. Hippocampus, 18(2):133-209 2008 2
14

---

Appendix
Table of Contents
A Proofs
15
B Details on Environments Related to Memory and Credit Assignment
17
Abstract Problems
B.2
Concrete Benchmarks for Memory
18
B.3
Concrete Benchmarks for Credit Assignment
20
Experiment Details
21
Memory-Basse RL Implementation
21
C.2 Additional Results
23
C,3 Plotting Details
24
C,4 Training Compute
24
D Broader Impacts
24
A Proofs
Lemma (Upper bound of value memory length). For any policy n.
I. (α) S max(m eransit mmm (T) 1)
(3)
Proof of Lemma 1.
We show this by expanding the sum of rewards in Q-value function:
Q" (hi.2, a1) E ET.
'r' Trs hiee ae
(4)
i=
T-1
=E(r: hi:e ael +
i+1-1
(5)
reward memory
P(o+1 Mis, a.) x(ai+1 h14+1) Blri+1 h1441, di41]
(6)
transition memory
policy memory
rewand memory
T-1
(A) rt b / mM Hl:t 0t
+
411
(7)
/ P(osti mi- hi-ma manat +1as ai (ar+1 1 11 (1)+ Pi41 h Monned +2:41, 05+1
do 1st
d0
(8)
R
rr Amin(t- milwued +1, +1 teamoit +1,t-1 (x)+2,1 sowaed +2)) Og
(9)
= Er.M
Σ 1 max--maxn(n (m
M
Amem (1:)=1)::: aq
(10)
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. I.D), the shortest history match the value, we have the upper bound: Wt
(aiue( S max (m-M
rewards my M
iransit: Immm (T.) - 1).
15

---

Lemma 2
(Lower bound of value memory length). For an optimal policy x² argmin (ine((T);; we have lnnem( (TT) l (=T))
Proof. Proof by contradiction. Suppose L' valuu ( lmomm (T²).. We can construct deterministic
policy x' that has context length of 1MM
value (1+*
x/ ( value (x (+))+1t) argmax Q
value (x--) (hy (t az)
(11)
argmax Q" (h1:0: ae), vt, his
(12)
This is greedy policy derived from " value function. The policy memory length of x' is strictly shorter than (TT), ), by definition and assumption:
mme (T') (e)x((T)
4mmm(("
(13)
Similar to policy improvement theorem (Sutton and Barto, 2018, Chap 4.2) in Markov policies, can show that is at least as good as π², thus contradicting the assumption that x is an optimal policy that has the shortest policy memory length. To see this, for any and h1:0
v (lea)) I max Q (hi:e,a))
(14)
(4) max Q (Mi T' (hy-MM (om)))) (11)+12))
(15)
(B) EM ra 1Q= (h1,41,,041) hra ~ n', att m²)
(16)
(c)
s Em re + 72" (114+1,0441) hig aest1 T
(17)
(18)
T
< EM
Σ
his a
(19)
E. .M
tr hist
V= (hi:)
(20)
where (A) follows the definition of n², (B) uses the Bellman equation and (C) uses the property of greedy policy x'.
Proof of Theorem 1. By Lemma 2, value (TT) 2 (x*). On the other hand, by Lemma 1,
value (T*) max(m Teward miransit Imem (x²) 1). Thus,
M
M muem (T*) < max(m Miraniity Ammmmm")) - 1)
(21)
This implies
mmem (Tt - 1 max max
zeward: (iransit)
(22)
I (o < max(m M
miransit)
(23)
c(T+") may be not unique among all optimal policies. For example, consider an MDP described in Fig. 9, and two optimal deterministic Markov policy, ni, n only different at the state P2: лī (P2) while Tī (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 be better than suboptimal action & at the next 2 steps for πミ. while at the next 3 steps for -.- Thus, c(Tī) = 3 and c(T)) = 2.
16

---

P3
x.y
P4
x.y
-1
+1
x.y
P1
P2
P7
x.y
+5
P5 +1
x.y
P6 -
Figure 9: An MDP that shows the credit assignment length over optimal policies is not constant. Each node state with its name on the top and its reward the bottom. The initial state Pl. The action space is {z, y). The finite horizon T 2 4 and 1.
B Details on Environments Related to Memory and Credit Assignment
B.1 Abstract Problems
Example 3 (n-order MDPs). It a POMDP where the recent n observations form the current state, Le. the transition is P(o++1 Ot-n+1:1, a4) and the reward is (o Qt).
Example (Episodic reward). is a POMDP with finite horizon of The reward function is episodic in that R² 0 for Vt T.
Example 5 (Decomposed episodic reward (Ren et al., 2022)), It has an episodic reward (Eg. 4), decomposed into the sum of Markovian rewards: RT (h1.T (T) E r' (o,, at). where r' O A + R a Markovian reward function. The transition P(o ot, ae) is Markovian. The smallest policy memory length for optimal policy lmem (TT" can be shown (see below).
Proof of decomposed episodic reward problem requiring lmmm (xx) = at most. Suppose the pol- icy is non-stationary, composed of (xt) For any time step s T in the POMDP M,
T
Q" @"(haa,aj ar) = -E.M
Σ
h1:1.
(24)
A=
EErT hi
r' (ox, ax)
h1a a
(25)
T Σ k=t+1
r'ko + Bammmooss--((laa)...
roo,, as) his at
(26)
By taking argmax,
argmax Q" (hia a.) argmaxr'(o,, P
Σ L
ror, ax)
h at
(27)
Now we show that there exists an optimal policy that is Markovian by backward induction. When t T, (27) reduces to
argmax Q" (IILT, aT) = argmax r' (OT. uT)
(28)
or
thus Markovian policy (or) argmax, r' r'To aT) is optimal. Now assume #++1:T Marko- vian and optimal, consider the step in (27), all future observations and actions (0t+1:7: ((+1.T)
rely on (o,, a;) by the property of Markovian transition and induction, thus optimal action aj can only rely on
Example 6 (Delayed rewards of n steps (Arjona-Medinn) et al., 2019)). The transition P(o oz, az)
Markovian, and the reward (o at-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 (Delayed execution of steps (Derman et al., 2021)). The transition P Ot, at-n); and the reward is Rg Re(o,, a.-n).
17

---

Table 2: The memory and credit assignment lengths required in the related abstract problems, using our notion. Assume all tasks have a horizon of T.
Abstract Problem MM
mem (x*)
m M reward
M Mernsiit
CM
MDP n-order MDP (Eg. 3)
0.
0.1
0,
1,
T
0, n
0, n.
0,
1.T
POMDP
[0, T
[0,
0.
[1,
Episodic reward (Eg. 4)
T
T
0,
T
Decomposed episodic reward (Eg. 5)
T
0,
T
Delayed rewards of steps (Eg. 6)
[0.
N
[0,
1]
n, T
Delayed exccution of n steps (Eg. 7)
n. I
Delayed observation of steps (Eg. 8)
[0,T]
Example 8 (Delayed observation of n steps (Katsikopoulos and Engelbrecht, 2003)). The transition P(o Oz, at-n), and the reward Re(or (Oc, Ot-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 Passive T-Maze, except that the agent is not penalized if it moves left instead of right in the corridor. The modification results in credit assignment length closer T', The authors test RL algorithms with the largest horizon of 70.
Sndevalier-Boisverr et al., 2018). This task has the same structure as our Active T-Maze. In the worst-case scenario, the optimal policy 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 x 17 51 for a 17 x 17 maze (MiniGride -iiddMeooryyyS Saltamoryyy/Raanonm..a 7Random-vO). The reward memory length can equal the horizon of 1445.
MiniGrid-Memory
traverse the horizontal corridor back and
Passive Visual Match (Hung et al., 2018). The environment consists a partially observable 7 x 11 grid-world, where the agent can only observe 5 x 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 recall temporally distant events. Since rewards relatively immediate with respect to the agent's 's actions, the credit assignment length of optimal policy is upper bounded by the length of the shortest path needed 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 Passive T-Maze. In TMaze Long, the agent is forced to take the forward action until reaching the junction (i.e., the action space 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 and the credit assignment length 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 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 correctly memorize the contents of the images for future actions. Images are passively observed by the agent, thus climinating the need for long-term credit assignment, and instead focusing long-term memory.
Spot the Difference (Fortunato 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 go from one room 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 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 needed to incorporate these actions into optimal policy.
18

---

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 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 x N number pad, and the goal of the agent is to activate up to N 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 N2 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 sol by supervised learning. Moreover, in this task, the observation i.i.d. sampled each time step, making a transition memory length of only 1. The authors set the horizon at 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, 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 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 's memory length is least 464 steps, and the reward memory length is capped by the horizon of 1024. The agent's 's actions are constrained to the post-dance phase, resulting in 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 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 choose the paired card that matches the observed card. It will get an immediate reward of O if it makes the correct choice, otherwise, 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 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 me = 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(N2) steps. Nevertheless, the optimal policy memory length can be much longer than O steps and close to the horizon T'. This is because the optimal policy may need to memorize the initial subsequence of observations infer the maze plan when taking current actions.
19

---

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 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, 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).
agent first observes sequence of cards,
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 memory length upper bounded by the area of the board, i.e. 12a = 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 considered as one of the hardest tasks in this benchmark. MM requires the agent to first memorize 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 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 al., 2019; Ren et al., 2022). It is a decomposed episodic reward problem (Fg. 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-Dook (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 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 to receive terminal reward. The credit assignment length is thus long. The horizon is 50.
20

---

C Experiment Details
C.1 Memory-Based RL Implementation
Our implementation builds upon prior work (Ni et al., 2022), providing a strong baseline variety of POMDP tasks, including Key-to-Door and PyBullet tasks used in this work. We denote this implemen- tation' as HOMDP-baselly All parameters of our agents in experiments trained end-to-end with model-free fooee--freeAa RL algorithms from scratch. The following paragraphs provide detailed descriptions.
Critic Heads
9f (Ac..))
Critic Head
0,
Olsery. Embeader
4g X
Actian Ensbedder
Model
(c(he))
4,
Action Embrrdurr
Sequence Model
MLP B
o2(h,..))
Reword
T1
Embeddee
Reword Embeddee
MLP
x(a, (h)
T2
Actor Head
Figure 10: Agent architectures for memory-based DDQN (left) and TD3 (right) in our implementa- tion. While rewards are not used as inputs in our experiments, they are included in this figure for completeness. Compared to SUDD--baseling (Ni et al., 2022), we have removed shortcuts from or to the MLPs and shared the same sequence model between the actor and critic.
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 POMPP-basel by sharing sequence model between the actor and critic. Previously, JOMDP-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-basee RL and related to sharing CNN encoders in visual RL (Yarats al., 2021a,b). Below formalize the issue and introduce solution.
In the separate encoder setting, we consider an actor and critic with distinet history encoders før and føa, mapping a history h into latent states and &Q. respectively. The history encoder parameters are denoted as o and dqq The actor is parameterized by TT (h)) and the critic is pa- rameterized by (foo (h).a) with v and w denoting their respective MLP parameters. Considering the DDPG algorithm (Lillicrap al., 2016) for training 10 and tuple of data (h, a, drrr", the loss of critic parameters is
1 Lo( = (u (h), - a,
(29)
where (ht a, o..). = Qa0 (N'), (f (W))) and 9 denotes the -rapdint version of 0 and ht' = o').. The loss of actor parameters is
pg
Iig (0x, v) = - Qwo (h), ( (h)))-
(30)
The total loss of actor and critic parameters is thus L (#r, 0q w. Lo(ØQ. + L VQn. The gradients of Lopp are
Væg Sep (QulJea (b), QTT a, V QulSao (oeg
(31)
V V (Qw (l), Qh a, V. (h),
(32)
de Loop VaQu(fee (h), Tx(fs. (h) V T for fo.
(33)
V, Vulsop - V.Qu (h), ( (1))) V. (Sa (h))
(34)
pattps://githhe htt.ggitub..oo.etteeeeeeeeeeeeee wtw/:/ni2016/6//1111.. httm/ttm12066//onttt hwtp:////ppppaa///e http///////pppppaaaaeaeieeeeee baselines 1h analysis can be generalized to stochastic actors by the reparametrization trick.
It denotes current history, current action, next observation, and current reward.
21

---

Now we consider an actor and a critic with shared encoder denoted as fø. mapping a history h into a latent state 2. The history encoder has parameters @. The actor parameterized by ग. (f and the critic is parameterized by (w( a). The total loss of actor and critic
Lata (1, = Lo(0, (fol)), (fo())
(35)
has such gradients:
V s =
- 10.1://00000000"aa..000000004000
V.
-. (36)
gradient
L = (00 -
x (
-
(37)
extra gradicnt term
Comparing the two gradient formulas across the two settings, we find the extra gradient terms using the same value of ø for ox and ΦQ in Lep:
Vd. Lep 2, w. v) + Vo LepLsapl0,00000 d, V VoLaha(O,w, m. Voll(
(38)
V 1e w.
(39)
To remove these extra gradient terms, we propose freezing the critic parameters @ and w in the actor loss for the shared encoder setting, resulting in a modified loss function Loha-ours (Ø, w, 1):
Luha-ours (6, w. v) Lo w) Q= (f) (f))
(40)
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). is worth noting that these works take an additional step of freezing the encoder parameters present in actions, resulting in an expression like Qu. (h), (f (h))). (A))) 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 implement our approach Tsba-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 input of sequence encoders is based on a history of / observations and 1 actions (with the first action zero- paro-padded). The history is embedded follows according to the POMDP-baselines state-based tasks (T-Maze 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.
to context length of 64. The training
For
For LSTMs (Hochreiter and Schmidhuber, 1997), we train LSTMs with 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 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 stack of N layers with H/ hhadedd 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 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 reported in Fig. 3 and Fig. 6. See Table 4 for details.
is used for regularization
For PyBullet tasks, we follow PaMMP-basella wwsebaseeeeeeeeee to use single-layer LSTMs and Transformers.
22

---

Table 4: Transformer hyperparameters used for all experiments.
Table 3: LSTM hyperparameters used for experiments.
Hyperparameter
Value
Hyperparameter
Value
State embedder
Obs. embedding size
64
State
Obs. embedding size
32
Act. embedding size
64
embedder
Act. embedding size
16
No. channels
(8, 16)
No. channels
(8, 16)
Pixel embedder
Kernel size
Pixel
Kernel size
Stride
embedder
Stride
Obs. embedding size
100
Obs. embedding size
100
Act. embedding size
Act. embedding size
0
Dropout
0.1
Hidden size
128 1, 1,2,4 4
LSTM
GPT
No. heads (H)
1,
No. layers
No. layers (N)
1,
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 outper- forms SAC-Discrete in our preliminary experiments. We use epsilon-greedy exploration strategy in DDQN with a linear schedule, where the ending epsilon is & with 7 being the episode length. This en- sures that the probability of always taking deterministic actions throughout an episode asymptotically approaches constant:
T []
lim (1 (1- T-++
lim T-++0
1 T
1
= R 0.368
(41)
where is the base of the natural logarithm. This approach critical solving T-Maze tasks which strictly require a deterministic policy. Table 5 summarizes the details of RL agents.
Table 5: RL agent hyperparameters used in all experiments.
Hyperparameter
Value
Network hidden size
(256, 256)
Discount factor (r)
0.99
Target update rate
0.005
Replay buffer size
106
Learning rate
0.0003
Batch size
64
epsilon greedy schedule schedule seceodulesstpps
linear(1.0, 4, techeeule_steppee
DDQN
0.1* num num_emimmsseeeeeee _am___ppsddssss
SAC-Discrete
entropy temperature
0.1
Table 6: Training hyperparameters in our experiments.
Tasks
context_length
num_episodes
Passive T-Maze Active T-Maze
1500
8000
500
4000
Passive Visual Match
1000
10000
Key-to-Door
500
4000
C.2
Additional Results
Fig. 11 shows the learning curves of training Transformers with varying numbers of layers and heads.
Similar the scaling experiments Transformers demonstrated in Sec. 5.2, conduct 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).
23

---

Evpy ny) modt calga ae imghh of
Key codlt cudgem longk of
1.00
0.73
0.76
0.30
0.35
6:35
11:00
1000
stope
stope
stape
Figure 11: Learning curves of scaling the number of layers and heads in Transformers in Key-to-Door tasks, associated with Fig. 7.
Active T Maze
Key to Door
0.8
0.5
0.0
0.0
0.4
-1.0
8
100
200
300
400
500
100
200
300
500
(Rary] Crvdit
length (Rard)
(Eagy)
Credit oongmmment Jnntth (llrd])
Figure 12: Scaling the number of layers in LSTMs does not affect performance much.
C.3
Plotting Details
For all the learning curve plots in this paper, we use seaborn. 1 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 See. 5.2, because the PyBullet experiments (Sec. 5.3) are relatively cheap to compute. Each experiment run was carried out on single A100 GPU and a single CPU core. For Transformer-basee based RI,, the GPU memory usage is approximately proportional to the square of the context lengths, with maximum usage of 4GB for Passive T-Maze with 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 sem__ppssdesss * context_lengthet * update_frequene frequency ffeffuuurrucc..o set as 0.25, is the ratio of parameter for both Transformers and LSTMs. The multi-layer Transformers, the training time is roughly update w.r.t. environment step. Additionally, for 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 and days to train Transformer- based and LSTM- RL, respectively. In Passive Visual Match with memory length of 1000, it took around days for both Transformer based and L.ST--baedd 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. 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.
24


Reviews:
[{"review_id": "NHhvqRqWXw", "reviewer": "Reviewer_mXpQ", "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.", "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.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 4, "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> What potential solutions could be to the credit assignment problem when using Transformers?\n\nWe think the credit assignment problem might be easier to solve with Transformers by using them to explicitly redistribute reward signals [Liu et al., 2019, Ferret et al., 2020]. Specifically, Transformers can efficiently redistribute distant reward signals to the time step when the corresponding action occurs. Their strength in handling long-term memory makes them suitable for learning such temporal dependencies with appropriate objectives beyond model-free RL.\n\n> 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?\n\nWe conduct new experiments on *multi-layer* Transformers in credit assignment tasks. Although they aid in medium-term credit assignment, they do not help with long-term credit assignment. Please see the **general response** for details. This result suggests that simply adding a specific submodule to Transformers is unlikely to improve long-term credit assignment, indicating the potential need for new architectures.\n\nPlease let us know if there are any further questions."}, {"review_id": "YZWWsLGupw", "reviewer": "Reviewer_naYb", "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.\n", "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.\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.", "limitations": "The authors have adequately addressed limitations.", "rating": 8, "confidence": 2, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "- 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", "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.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n> 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.\n\nRecurrent policies can *take* all the information starting from time 0 as inputs, but not necessarily *remember* it starting from that point.  As mentioned in L88-91, the context lengths of recurrent policies can be infinite, but their memory lengths can be very short, due to training issues such as gradient vanishing or explosion.  In this sense, our definition of memory lengths remains relevant to recurrent policies, with context lengths being infinite.\n\n> Lacking discussion on state and history abstraction, for example, bisimulation and approximate information states. \n\nThank you for pointing it out. This work focuses on the architectural aspect (e.g., LSTMs and Transformers) of RL rather than the objective aspect (e.g., different abstraction methods). We view these two aspects as orthogonal and believe they can complement each other to enhance RL performance. We will include some related work on abstraction and clarify our focus in the camera-ready version.\n\n>  Are Transformers predicting states that truly matter for the task? \n\nWe evaluate model-free RL agents equipped with memory architectures. Unlike model-based RL, which explicitly predicts next states, model-free RL purely aims to maximize returns. Therefore, we do not think Transformer-based RL agents used in our experiments face the issue of predicting irrelevant parts of states, regardless of their memory capability.\n\n> In operations research there are classes of \"submodular functions\", and these have provably efficient greedy algorithms that are quantifiably suboptimal. How is this connected to credit assignment length? \n\nAfter examining the definition of submodular functions, we don't see a direct connection between it and credit assignment length.  We are open to further insights or clarification on this connection, as it may lead to an intriguing avenue for future research."}, {"review_id": "ZTbT29HQWZ", "reviewer": "Reviewer_kd6F", "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.\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", "questions": "Above", "limitations": "Yes", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "- 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.", "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. \n- It would be nice to see experimental evaluation of Transformers and LSTM in additional environments other than the toy tasks for better context.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> This work explores the impact of memory and credit assignment in decision transformer architectures.\n\nWe would like to clarify that our experiments evaluate the *online model-free RL* setting, not the offline RL setting that is related to Decision Transformers (DT). Nevertheless, extending our evaluation to include DT is an interesting direction for future work!\n\n> Can the authors provide any intuition on the failure of Transformers and LSTM in the long-term credit assignment tasks?\n\nWe find that in Active T-Maze task that has a credit assignment length of 250, the Transformer-based agent does not reach the Oracle (return lower than 1), but still reaches the Junction (return higher than 0). This indicates that Transformer-based agent may fail to explore enough to reach the Oracle, although the exploration should be relatively easy (just taking the initial action to move left). In LSTM-based agent, the return can *sometimes* be negative, indicating that it does not consistently reach the Junction and may face additional issues beyond credit assignment. Similar issues seem to occur in the Key-to-Door tasks where both agents do not reach the key in the initial phase. \n\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\nWe agree with this limitation, and we plan to evaluate agents on more complicated environments in future work."}, {"review_id": "TgIxu2ss2S", "reviewer": "Reviewer_2AdR", "summary": "The 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.", "questions": "- 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", "limitations": "The 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?\n", "rating": 6, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "**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.", "weaknesses": "**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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\n> Does computing memory lengths involve human demonstrations?\n\nWe here clarify that computing these metrics *does not* require any human demonstrations. For some tasks like T-Maze, we obtain the precise numbers with careful analysis of the task structures. For other complex tasks like Memory Maze and proc-gen environments, as written in L576-579, it's hard for human experts to determine *exact* lengths by problem definition, although some bounds can still be derived, as shown in Table 1.\nWe will clarify how we compute the metrics in the camera-ready version. \n\n> Will you release code?\n\nYes, we have released our code in the supplementary material. We do not use a program to compute our metrics. \n\n> How does computing memory metrics scale with the task?\n\nThe complexity depends on the history space, which grows exponentially in a finite POMDP. \nThe worst-case complexity for reward and value memory lengths is $O(T (|\\mathcal O| + |\\mathcal A|)^T |\\mathcal A|)$.\nThe transition and policy memory lengths' complexity is further multiplied by the cardinality of observation and action spaces.\n\n> The MiniGrid-Memory task already shares the same properties as the proposed Active T-Maze. \n\nThank you for pointing this out! We will add MiniGrid-Memory task into Table 1. While this task has the same properties as Active T-Maze, it is easier in terms of memory and credit assignment lengths.\nMiniGrid-Memory (sized $13\\times 13$) has its memory and credit assignment lengths of optimal policies $\\le 3\\times 13 = 39$, much smaller than the lengths of 250 in Active T-Maze.\nIn addition, MiniGrid-Memory's $147$-dim observations and sparse rewards add unnecessary complexity to evaluation. \nActive T-Maze simplifies the problem with 2-dim observations and dense penalties.  \nLastly, combined with Passive T-Maze, Active T-Maze can be used to evaluate the bottleneck of memory or credit assignment. These advantages motivate us to propose Active T-Maze. \n\n\n> What is the observation space of T-Mazes? \n\nThese tasks have 2-dim discrete observations indicating the position of an agent: Oracle, Start, Junction, Goal candidates, Corridors. We have revised the paper to include these details.\n\n> Passive T-Maze is useful but minimalistic. It only evaluates if a small bit can be stored and carried across long timespans in memory. \n\nYou are correct, and that was indeed our intention. We chose this minimalistic task to focus on \"memory length\" rather than \"memory capacity\" in terms of bits. With no prior evidence that RL agents can maintain a memory spanning $1500$ steps, this task precisely targets this capability. To avoid confusion, we have updated our abstract to state \"memorizing observations 1500 steps ago\" instead of \"memorizing 1500 observations\".\n\n> How about trying an image-based version of Passive T-Maze?\n\nIn fact, we have conducted experiments on an image-based version of Passive T-Maze, the Passive Visual Match.\n\n> The paper should clarify its focus on long-term memory, not capacity or robustness.\n\nYes, we will do it. \n\n> Missing related work [2,3,4]. If applicable, running these methods?\n\nThanks for pointing them out. We find that the hierarchical chunk attention memory [3] is related and will include it into the camera-ready version. It's worth noting that [3] uses observation reconstruction for sparse-reward tasks, unlike our model-free RL. Due to time limit, we cannot run experiments on this method. \n[2] and [4] are remotely related -- [2] focuses on recurrent-based RL  *convergence* and [4] uses pre-trained and frozen language models for RL.\n\n\n> The authors say GPT is less sample-efficient than LSTM in short-term memory tasks, but GPT outperforms LSTM in Passive Visual Match. Could the drop be due to observation/action spaces (continuous vs discrete)?\n\nYes, it is possible. As stated in L278 and L292, we find that GPTs are sample-efficient on Passive Visual Match, but not PyBullet, though both require short-term memory.\n\n>  Any intuition why GPT is less sample-efficient than LSTM on short-term memory tasks?  \n\nOur intuition is that GPTs have fewer inductive bias than LSTMs, thus generally require more data to learn. \n\n> How do long-term memory tasks relate to the SL setup on large datasets? \n\nWe did not relate them. \n\n> Do the authors imply that there are only long-term dependencies in large datasets for SL?  \n\nNo. \n\n> Morad et al. 2023 advises against comparing SL and RL.\n\nWe think **purely long-term memory tasks** in RL and SL are related. For example, Copy task [Arjovsky et al., ICML 2016]  is like an SL version of Passive T-Maze. LSTMs failed to solve the long-term Copy task due to gradient training issue, mirroring our findings in Passive T-Maze. \n\n> In L270 it'd be good to explain what the actual task is.\n\nWe will add these descriptions in the camera-ready version. \n\n> Fig 5: what is optimal policy that lacks (long-term) credit assignment?\n\nSuch a policy maximizes its actions depending on *immediate* rewards only, according to definition of the credit assignment length. In Active T-Maze task, such a policy happens to be Markovian, but it can have memory in general. \nThis is a good spot, and we will clarify it in the camera-ready. \n\n\n> Do all figures show mean and std dev? This should be mentioned in captions.\n\nAll figures are generated by seaborn.lineplot that shows the mean and its 95\\% confidence interval, but not std dev. We will add these descriptions to the captions. \n\n> The last 3 sentences in Def. 1 can be moved to Def. 2, as they are on memory length and not context length.\n\nWe will separate the last three sentences from the paragraph to create a new paragraph for Def. 2. \n\n> Question on equation in L103?\n\nActually, these two are equivalent. For random variables $X,Y,Z$, $X\\perp Y \\mid Z$ is equivalent to $X \\perp Y,Z \\mid Z$, where $X = a_t$, $Y = h_{t-l_{ctx}(\\pi)+1:t-k}$, and $Z = h_{t-k+1:t}$ in this case. \n\n> The symbol k is reassigned, maybe using other symbols?\n\nWe will rename these symbols. "}]

Author rebuttal:
# General Response

First, we would like to thank all four reviewers for their positive and constructive feedback on our work! 
During the rebuttal period, we conducted additional experiments to address some of the questions posed by the reviewers.

**mXpQ  (see Figure 1 in the PDF)**

> Can we simply slap on a specific submodule to Transformers to augment their credit assignment capabilities?

We showed that (single-layer) Transformers cannot help long-term credit assignment. To further confirm this conclusion, we ran *multi-layer* Transformers on these tasks.
In both Active T-Maze and Key-to-Door, we find that both 2-layer and 4-layer Transformers greatly outperform 1-layer Transformer in tasks with medium-term credit assignment (length $\le 200$), while still fail to perform long-term credit assignment (length $\ge 250$). 
This finding suggests that it is *unlikely* that simply increasing the size of Transformers will enable them to excel in long-term credit assignment.

**2AdR (see Figure 2 in the PDF)** 

> LSTM starts to falter at a memory length of 250 for Passive T-Maze. But Fig 2 shows that LSTM can solve it with a memory length of 750. This suggests that hyperparameter tuning might influence LSTM training. Will the results of LSTM be improved by further tuning?

First, we want to clarify that although LSTM reaches returns higher than $0.5$ in Passive T-Maze with a memory length of $750$, it does not *solve* the task, requiring the return to be $1.0$. 

Your point on tuning may be valid for medium-term memory lengths. But for long-term memory tasks, our new results show that tuning is unlikely to help LSTM.
During the rebuttal period, we conduct new experiments on LSTM-based agent on the tasks with memory lengths of 1250 and 1500. LSTM-based agent reaches return below 0.5. 
The results confirm our conclusion that LSTM-based agent cannot solve long-term memory tasks.  

**2AdR**

> Can you do a significance test to show superiority of GPT2 over LSTM in Figure 4?

Following the reviewer's suggestion, we did a Welch’s t-test on Passive Visual Match success rates (Figure 4 in the main paper). The p-values for the memory length of 500 and 750 are 0.698 and 0.038, respectively. 
These results indicate that there is significant evidence to reject the null hypothesis at the 750 memory length, but not at 500. Therefore, we clarify that the advantages observed with GPT at a memory length of 500 predominantly pertain to short-term memory, rather than long-term memory.


Meta-review: This paper aims to understand the success of transformer-based methods for reinforcement learning. In particular, they aim to understand the role of transformers in addressing two challenges (1) credit assignment (understanding how actions influence future rewards), and (2) memory (representing prior observations); both challenges entail modeling long-term dependencies. To separate these challenges, the authors design a suite of experimental benchmarks, and find transformers are effective at enhancing memory, but do not improve credit assignment.

Reviewers found the paper's topic to be timely, and found the results to be convincing and likely to be of significant interest to the reinforcement learning community. In addition, they found the paper to be accessible and well-written, and unanimously advocated for its acceptance.
Decision: Accept (oral)