Structured State Space Models for In-Context Reinforcement Learning
Authors: Chris Lu
Yannick Schroecker
Albert Gu
Emilio Parisotto
Jakob Nicolaus Foerster
Satinder Singh
Feryal Behbahani
Venue: NeurIPS 2023 poster (2023)
Keywords: Reinforcement Learning
Meta-Learning
State Space Models
TL;DR: We show that S4 models enable fast and efficient long-horizon memory in RL

Abstract: Structured state space sequence (S4) models have recently achieved state-of-the-art performance on long-range sequence modeling tasks. These models also have fast inference speeds and parallelisable training, making them potentially useful in many reinforcement learning settings. We propose a  modification to a variant of S4 that enables us to initialise and reset the hidden state in parallel, allowing us to tackle reinforcement learning tasks. We show that our modified architecture runs asymptotically faster than Transformers in sequence length and performs better than RNN's on a simple memory-based task. We evaluate our modified architecture on a set of partially-observable environments and find that, in practice, our model outperforms RNN's while also running over five times faster. Then, by leveraging the model’s ability to handle long-range sequences, we achieve strong performance on a challenging meta-learning task in which the agent is given a randomly-sampled continuous control environment, combined with a randomly-sampled linear projection of the environment's observations and actions. Furthermore, we show the resulting model can adapt to out-of-distribution held-out tasks. Overall, the results presented in this paper show that structured state space models are fast and performant for in-context reinforcement learning tasks. We provide code at https://github.com/luchris429/s5rl.

Structured State Space Models for In-Context Reinforcement Learning
Chris Lu" FLAIR, University of Oxford
Yannick Schroecker DeepMind
Albert Gu DeepMind
Emilio Parisotto DeepMind
Jakoh Foerster FLAIR, University of Oxford
Satinder Singh
Feryal Behbahani DeepMind
Deep.ind
Abstract
Structured state space sequence (S4) models have recently achieved state-of-the-art performance long-range sequence modeling tasks. These models also have fast inference speeds and parallelisable training, making them potentially useful in many reinforcement learning settings. We propose modification of S4 that enables us to initialise and reset the hidden state in parallel, allowing to tackle reinforcement learning tasks. We show that our modified architecture runs asymptotically faster than Transformers in sequence length and performs better than RNN's on simple memory- -bsed task. We evaluate our modified architecture on set of partially our model outperforms RNN's while also running over five times faster. Then, by leveraging the model's ability to handle long-range sequences, we achieve strong performance on a challenging meta-learning task which the agent given aandomly-sappler continuous control environment, combined with a randomly-sanpled linear projection of the environment's observations and actions. Furthermore, we show the resulting model can adapt to out-of-distribution held-out tasks. Overall, the results presented in this paper show that structured state space models are fast and performant for in-context reinforcement learning tasks. We provide code at littt::///// //te:://htttut com/luchris429/s5rr
variant
environments and find that, in practice,
Introduction
Structured state space sequence (S4) models [2 and their variants such as S5 [8] have recently achieved impressive results in long-range sequence modelling tasks, far outperforming other popular sequence models such as the Transformer [42) and LSTM [16] on the Long-Range Arena bench- mark [41). Notably, S4 was the first architecture achieve non-trivial result the difficult Path-X task, which requires the ability to handle extremely long-range dependencies lengths 16k.
Furthermore, S4 models display number of desirable properties that are not directly tested by raw performance benchmarks. Unlike transformers, for which the per step runtime usually scales quadrat- ically with the sequence length, S4 models have highly-scalable inference runtime performance, asymptotically using constant memory and time per step with respect to the sequence length. While LSTMs and other RNNs also have this property, empirically, S4 models are far more performant while also being parallelisable across the sequence dimension during training.
While inference-tingg time is normally not included when evaluating sequence modelling benchmarks, has a large impact on the scalability and wallclock-time for reinforcement learning (RL) because
Work done during an internship at DeepMind. Contact: christopher.uu @xx.ter.ox.c..uukk
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Inference
Training
Parallel
Variable Lengths
bsuite Score
RNNs
O(1)
O(L)
No
Yes
No
Transformers
O(L2)
O(L2)
Yes
Yes
Yes
S5 with
0(1)
O(L)
Yes
No
N/A
S5 with 4
0(1)
O(I)
Yes
Yes
Yes
Table 1: The different properties of the different architectures. The asymptotic runtimes are in terms of the sequence length L assume a constant hidden size. The bsuite scores correspond to whether or not they achieve a perfect score in the median runs on the bsuite memory length environment.
the agent uses inference to collect data from the environment. Thus, transformers usually have poor runtime performance in reinforcement learning [33]. While transformers have become the default
architecture for many supervised acquence-modelling tasks [42]. RNNs are still widely-used for memory-based RL tasks [29].
The ability efficiently model contexts that are orders of magnitude larger may enable new possi- bilities in RL. This is particularly applicable in meta-reinforcemennt learning (Meta-RL), in which the agent is trained to adapt across multiple environment episodes. One approach Meta-RL. RL 8 43 uses sequence models to directly learn across these episodes, which can often result in trajectories that are thousands steps long. Most instances of RL. approaches, however, are limited to narrow task distributions and short adaptation horizons because of their limited effective memory length and slow training speeds.
Unfortunately, simply applying S4 models to reinforcement learning is challenging, This is because the most popular training paradigm on-policy RL with multiple actors involves collecting fixed- length environment trajectories, which often cross episode boundaries. RNNs handle episode boundaries by resetting the hidden state at those transitions when performing backpropagation through time. Unlike RNNs, S4 models cannot simply reset their hidden states within the sequence because they train using a fixed convolution kernel instead of using backpropagation through time.
recent modification to S4, called Simplified Structured State Space Sequence Models (S5), replaces this convolution with parallel scan operation (38. which we describe in Section 2 In this paper, we propose a modification to S5 that enables resetting its hidden state within trajectory during the training phase, which turn allows practitioners simply replace RNNs with S5 layers in existing frameworks. We then demonstrate S5's performance and runtime properties on the simple bsuite memory-length task [2] showing that S5 achieves a higher score than RNNs while also being nearly two times faster when using their provided baseline algorithm. We also re-implement and open source the recently- propposed Partially Observable Process Gym (POPGym) suite (27 pure JAX, resulting in end-to- evaluation speedups of over 30r. When evaluating our architecture on this suite, we show that outperforms GRU's while also running over six times faster, achieving state-of-the-art results on the "Repeat Hard" task, which all other architectures previously struggled to solve. We further show that the modified S5 architecture can tackle a long-context partially-obsers Meta-RL task with episode lengths of up 6400. Finally, we evaluate S5 challenging Meta-RL RL task in which the environment samples a random DMControl environment [40) and a random linear projection of the state and action spaces at the beginning of each episode. We show that the S5 agent achieves higher returns than LSTMs in this setting. Furthermore, we demonstrate that the resulting S5 agent performs well even random linear projections of the state and action spaces of out-of-distribution held-out tasks.
2 Background
2.1 Structured State Space Sequence Models
State Space Models (SSMs) have been widely used to model various phenomenon using first-order differential equations [4] At each timestep I, these models take an input signal u(t). This used to update a latent state z(f) which in turn computes the signal y(t). Some of the more widely-used

---

canonical SSMs are continuous-ting time linear SSMs, which are defined by the following equations:
=(t) Az(1) + Bu(t) v(() Ca(() + Du(()
where A, B. C, and are matrices of appropriate sizes. To model sequences with fixed step size A, one can discretise the SSM using various techniques, such as the zero-order hold method, to obtain a simple linear recurrence:
= Arn-1 + Bun
(2)
1/1 Can + Dun
where A, B, C, and D can be calculated as functions of A, C,D. and A.
S4 [12] proposed the of SSMs for modelling long sequences and various techniques to improve its stability, performance, and training speeds when combined with deep learning. For example, S4 models use special matrix initialisation to better preserve sequence history called HiPPO [[]]
One of the primary strengths of the S4 model that can be converted to both recurrent model, which allows for fast and memory efffireteffiieen ffference-time computation, and a convolutional model, which allows for efficient training that parallelisable across timesteps [3]
More recently, Smith et al. 38 proposed multiple simplifications to S4, called S5, One of its contribu- tions the use of parallel scans instead of convolution, which vastly simplifies S4's complexity and enables more flexible modifications. Parallel scans take advantage of the fact that the composition of associative operations can be computed in any order. Recall that for operation be associative, it must satisfy (x
Given associative binary operator and sequence of length N, parallel scan returns:
[ei,et
ex]
For example, when is addition, the parallel scan calculates the prefix-sum, which returns the rumning total of an input sequence. Parallel seans can be computed in (log(NN) time when given a sequence of length N, given N parallel processors.
S5's parallel scan is applied to initial elements co:N defined as:
ek (Ek,as (kDd) (A, Bux)
Where A, B, and are defined in Equation 2 S5's parallel operator is then defined as:
aj- (aj.a O 0j aj,a @ agp + 0g.b)
(5)
where @ is matrix maarix--mtrrx matrix multiplication and @ is matrix-vector multiplication. The parallel scan then generates the recurrence in the hidden state x, defined in Equation 2
(Ā, Bu1)
(A,r1)
(6)
E2 (A Ax + Buz) (A2,0))
(7)
C3 (A Azg + Buz) (A², xx))
(8)
Note that the model assumes a hidden state initialised to mo = by initialising the scan with eo (I,
2.2 Reinforcement Learning
Markov Decision Process (MDP) 139 is defined as a tuple (S, AS, which defines the environment. Here, S is the set of states, A the set of actions, R the reward function that maps from a given state and action to a real value R,, P defines the distribution of next-state transitions given state and action, and defines the discount factor. The agent's 's objective is to find policy π (a function which maps from given state distribution over actions) which maximises the expected discounted sum of returns.

---

Algorithm Pseudocode for the Multi-Enwironment Meta-Learning environment step.
Require:
Distribution of environments E ~ pE. fixed output observation dimension size O, and fixed action dimension size A. Agent action and Environment termination d
1: function StepEnvironment(a, d)
2: if the environment terminated (d) then
3:
Sample random environment E ~ PE
4: Initialise random observation projection matrix M, E RREx x0 where E E's observation size
5:
Initialise random action projection matrix
Ma E AA where E is E's action size
Reset E to receive initial observation 6:
7:
Apply the random observation projection matrix to the observation o' = Mo Append r 0 and d 0 to of to get o" Return d"
8:
9:
else
11:
Apply the projection matrix a' = M2 Step E using a' receive the next observation o, reward and done signal d.
13:
Apply the projection matrix o' = Moo
Append r and d to of get o" Return o", r. and d end if
15:
17: end function
E
E Ro
R(st, at)
In a Bartially-Observed Markov Decision Process (POMDP), the agent receives an observation Or Oss)) instead of directly observing the state. Because of this, the optimal policy n² does not depend just on the current observation o,, but (in generality) also on all previous observations Op.r and actions ag.t-
3 Method
We first modify to S5 to handle variable-length sequences, which makes the architecture more suitable for tackling POMDPs. We then introduce a challenging new Meta-RL setting that for broad generalisation capabilities.
3.1 Resettable S5
Implementations of on-policy policy gradient algorithms with parallel environments often collect fixed-length trajectory "rollouts" from the environment for training, despite the fact that the environ- ment episode lengths are often far longer and vary significantly. Thus, the collected rollouts (1) often begin within an episode and (2) may contain episode boundaries. Note that there are other, more complicated, approaches to rollout collection that can be used to collect full episode trajectories (2
To handle trajectory rollouts that begin in the middle an episode, sequence models must able to access the state of memory that was present prior to the rollout's collection. Usually, this is done by storing the RNN's hidden state at the beginning of each rollout to perform truncated backpropagation through time 44 This is challenging
do with transformers because they do not normally have an explicit recurrent hidden state, but instead simply retain the entire history during inference time. This is similarly challenging for S4 models since they assume that all hidden states initialised identically to perform a more efficient backwards pass.
To handle episode boundaries within the trajectory rollouts, memory-based models must be able to reset their hidden state, otherwise they would accessing memory and context from other episodes. RNNs can trivially reset their hidden state when performing backpropagation through time, and

---

transformers can mask out the invalid transitions. However, S4 models have no such mechanism to do this.
To resolve both of these issues, we modify S5's associative operator to include a reset flag that preserves the associative property, allowing S5 efficiently train over sequences of varying lengths and hidden state initialisations. We create a new associative operator @ that operates on elements eg defined:
fk (Et,as Ck,b; Ck,e) (A, Bug, dy.)
where de E { is the binary "done" signal for the given transition from the environment.
We define @ be:
(aja O aia ajja @ ai,b aic) if dee 0 a 0j (03,01 Oj.b aj.c)
if az,c
We prove that this operator is associative in Appendix A We now show that the operator computes the desired value. Assuming Ên,e = 1 corresponds to a "done" transition while all other timesteps before (10 = 0) and after (en+1:L,c 0) do not, we see:
6o 1 ... 1 c- = (A n-1 Ar + Bun
= (A n- ..)
① ... 1 en = (A, Bu,, 1)
= (A, Zn, 1)
60 Ex+1 = (A?, Ax, + Bux+1,
(A2 1))
Note that even if there are multiple "resets" within the rollout, the desired value will still be computed.
3.2
Multi-Enviromment Meta-Learning with Random Projections
Most prior work in Meta-RI. has only demonstrated the ability to adapt to small range of similar tasks 2] Furthermore, the action and observation spaces usually remain identical across different tasks, which severely limits the diversity of meta-learnnig environments. To achieve more general meta-learning, ideally the agent should learn from tasks of varying complexity and dynamics. Inspired by Kirsch et al. [21], we propose taking random linear projections of the observation space and action space to fixed size, allowing us use the same model for tasks of varying observation and action space sizes. Furthermore, randomised projections vastly increase the number of tasks in the meta-learning space. We can then evaluate the ability of our model to generalise to unseen held-out tasks.
We provide pseudocode for the environment implementation in Algorithm
4 Experiments
4.1
Memory Length Environment
First, we demonstrate our modified S5's improved training speeds in performance in the extremely simple memory length environment proposed in bsuite [32].
The environment is based on the well-known 't-maze" environment 31 in which the agent receives cue on the first timestep, which corresponds to the action the agent should take some number of steps in the future receive reward. We run experiments using bsuite's ' actor-critic baseline while swapping out the LSTM for Transformer self-attention blocks or S5 blocks [4 and using Gymnax for faster environment rollouts [22].
We show the results in Figure general, S5 displays a better asymptotic runtime than Transformers while far outperforming LSTMs in both performance and speed. Note that while S5 is theoretically
5

---

Peatoo Morners Usoge Vectorizen Ruyss
Madia Vraus Acmass Sends
201
12:40
1208 2008
2002
1180 1813
2005
(a)
(b)
(c)
Figure 1: Evaluating S5, LSTM, and Self-Attention across different Bsuite memory lengths in terms of (a) memory usage, (b) runtime, and (c) return. Error bars report the standard error of the mean across 5 seeds, Runs were performed on a single NVIDIA A100.
POPGym
POPGym Normelised MMER
1.0
with with
0.6
Meralted
150
2000
4000
6000
8000
10000
12000
Undates
Seconds
(a)
(b)
Figure 2: (a) Results across implemented environments in POPGym's 's suite. Scores normalised by the max-MMER per-environment. The shaded region represents the standard error of the mean across eight seeds. (b) The runtime for single seed averaged across the environments for each architecture. Note that our implementation is end-to-end compiled to run entirely on single NVIDIA A40.
(log(NN) in the backwards pass during training time (given enough processors), it is still bottle- necked by rollout collection from the environment, which takes O(N) time. Because of the poor runtime performance of transformers for long sequences, we did not collect results for them in the following experiments.
4.2 POPGym Environments
We evaluate our S5 architecture on environments from the Partially Observable Process Gym (POP- Gym) [27] suite, set of simple environments designed to benchmark memory in deep RL. To increase experiment throughput on limited compute budget, we carefully re-implemented environments from the POPGym suite entirely in JAX K by leveraging existing implementations of CartPole and Pendulum in Gymnax [22]. Py Torch does not support the associative scan operation, so we could directly use the S5 architecture in POPGym's 's RLLib benchmark.
Morad al. (27) evaluated several architectures and found the Gated Recurrent Unit (GRU) the most performant. Thus, we compare our results to the GRU architecture proposed in the original POPGym benchmark. Note that POPGym's 's reported results use RLLib's 24 implementation of PPO, which makes several non-standard code-level implementation decisions. For example, it uses a dynamic KL-divergence coefficient on top of the clipped surrogate objective of PPO 37 feature that does not appear in most PPO implementations (9] We instead use recurrent PPO implementation that more closely aligned with StableBaselines3 355 and CleanRL's [8 recurrent
be
PPO implementations. We include more discussion and the hyperparameters in Appendix

---

Heta State was Carticle Vetaa ring NAEE
Stardons CortPes Drckation Notomerres
Sttraees Carticia Nata Rundine
309
6.55
0.00
Upd0ns
643
(b)
(C)
Figure 3: (a) Performance and (b) runtime on randomly- projjected StatelessCardee Pareees across seeds. The shaded region represents the standard error. Shows performance the end of training across different trials. We evaluate on 32 trials even though we only train on 16. GRU's appear to have overfit the training horizon while S5 models continue perform well. The error bars represent the standard error across 4 seeds. Runs were performed on single NVIDIA A100.
We show the results in Figure 2 and Appendix C Notably, the S5 architecture performs well on the challenging "Repeat Previous Hard" task, far outperforming all architectures tested in Morad et al. 27 Furthermore, the S5 architecture also runs over six times faster than the GRU architecture.
4.3
Handomly-Projectee CartPole In-Context
We first demonstrate the long- horizon in-context learning capabilities of the modified S5 architecture by meta-learning across random projections POPGym's 's Stateless CartPole (CartPole without velocity observations) task. More specifically, we perform the observation and action projections described in Section 3.2 and report the results in Figure 3
Agents are given 16 trials per episode randonly-sampled linear projections of StatelessCartPolee observation and action spaces. We find that S5 with @ outperforms GRU's while running twice as quickly. Furthermore, learning and show that the modified S5 architecture, unlike the GRU, can continue to learn even past its training horizon, which was previously shown using Transformers in Adaptive Agent Team et al.
show that performance improves across trials, demonstrating in-context
4.4 Multi-Enviromment Meta- Reinffreeerre Learning
We run our S5 architecture, an LSTM baseline, and a memory-less fully-connected network in the environment described Section 3.2 For these experiments, we use Muesli [15] as our policy optimisation algorithm. We randomly project the environment observations to fixed size of 12 and randomly project from an ac- tion space of size 2 to the corresponding envi- ronment's action space. We selected all of the DMControl environments and tasks that had ob- servation and action spaces of size equal to or less than those values and split them into train and test set environments.
Mean Across Training Envs
800
FC
LSTM
600
400
50
125 Frames (10 Million)
75
100
150
175
200
We use the S5 architecture described in Smith al. (38), which consists of multiple stacked layers of S5 blocks. For both S5 and LSTM ar- chitectures, we found that setting the trajectory length equal to the maximum episode length 1 achieved the best performance in this setting. We provide more detailed description of our
Figure 4: The mean of the return across all of the training environments. The shaded regions repre- sent the range of returns reported across the three seeds. The environment observations and actions randomly projected described in Algorithm
7

---

ball in ctppccupppaapp catch
cartpale swingup
finger sager_spinn
1000
600
Returss
55
500
1STM
200
0
50
100
150
200
50
100
150
200
50
100
150
200
finger turn turn_hard
point mass easy
reacher_ pasy
750
750
Returns
750
500
500
500
250
250
100 150
100
200
Frames (10 Million)
Frames (10 Million)
Frames (10 Million)
Figure 5: Results across the different environments of the training distribution. The shaded regions represent the range of returns reported across the three seeds. The environment observations and actions are randomly projected as described in Algorithm] T
architecture and hyperparameters in the supple- mentary material.
In- In-Distribution Training Results We meta-train the model across the six DMControl environments in Figure Slau show the mean performance across them in Figure Note that while these environments are in the training distribution, they are still being evaluated on unseen random linear projections of the state and action spaces and are not given task labels. Thus, the agent must infer from the reward and obfuscated dynamics which environment it is in. In this setting, S5 outperforms LSTMs in both sample efficiency and ultimate performance.
Evaluation Results: Af- ter training the model the DMControl tasks in Figure we next evaluate the trained model random linear projections of five held-owt DMControl tasks, without any extra fine- fine-tuning. We present the mean score across these tasks in Figure O and the results for each task in Figure 7 While the agent displays impressive transfer performance to some tasks with unseen rewards and dynamics, it fails to successfully transfer to
Mean Across All Test Envs
55
500
LSTM
400
Return
300
200
100
completely unseen task in the test set, pendu- lum swingup, which has an unseen observation space, action space, and reward dynamics.
150
175
Frames (10 Million)
Figure 6: The mean of the return across all of the unseen held-out environments. The shaded regions represent the range of returns reported across the three seeds. The environment observations and actions are randomly projected as described in Al- gorithm
5 Related Work
S4 models have previously been shown to work in number previous settings, including audio generation [10] and video modeling [28]. Con- current work investigated S4 models in other reinforcement learning settings. David et al. [6] investigated S4 models in a largely offline RL setting, with some modifications for online finetuning of the model. Morad et al. [27] investigated sequence models for POMDPs and found that naively using S4 models did not perform well.
Most works in memory- based meta-RL have not focused on generalisation to diittr-iittiitttion tasks, but instead focused on maximising performance on the training task distribution [8] 43 2). Many past works that increased generalisation in meta-RL did so by restricting the model architecture class or architecture to symmetric models (20), loss functions [17] target values (30), or drift functions [25]
In contrast, this work achieves generalisation by vastly increasing the task distribution without limiting the expressivity of the underlying model, which eliminates the need for
restrictions.
8

---

cartpole balance
finger easy
pendulum swingup
100
750
Returs
55
200
FC
500
1.5M
250
50
100
150
200
0
50
100
150
200
50
100
150
200
point mass, hard
reacher
haad hard
6.00
Returss
400
200
200
Frames (10 Million)
Framnes (10 Million)
Figure 7: Results when evaluating on held-out DMControl tasks. The shaded regions represent the range of returns reported across the three seeds. The environment observations and actions are randomly projected described in Algorithm
Some works perform long-horizon meta meta-reinforcemmen learning through the use of evolution strate- gies
This is because RNNs and Transformers have historically struggled to model very long sequences due to computational constraints and vanishing or exploding gradients [26). However, evolution strategies are notoriously sample inefficient and computationally expensive.
25
Other works have investigated different sequence model architectures for memory- based reinforce- ment learning. Ni al. [29] showed that using well-tuned RNNs can be particularly effective compared to many more complicated methods in POMDPs. Parisotto et al. [34] investigated the use of transformer-basee models for memory-based reinforcement learning environments.
Sequence models have also been used for a number of other tasks in RL. For example, Transformers have been used for offline reinforcement learning [4], multi-task behavioural cloning (36), and algorithm distillation 23 Concurrent work used transformers also demonstrate out-of-distributiin generalisation in meta-reinforcement learning by leveraging a large task space 11.
Conclusion and Limitations
In this paper, we investigated the performance of the recently-proposed S5 model architecture in reinforcement learning. S5 models are highly promising for reinforcement learning because of their strong performance in sequence modelling tasks and their fast and efficient runtime properties, with clear adnatages over RNNs and Transformers. After identifying challenges in integrating S5 models into existing recurrent reinforcement learning implementations, we made simple modification the method that allowed us to reset the hidden state within a training sequence.
We then showed the desirable properties S5 models in the bsuite memory length task. We demon- strated that S5 is asymptotically faster than Transformers in the sequence length. Furthermore, we also showed that S5 models nearly twice quickly LSTMs with the same number of parameters while outperforming them. We further evaluated our S5 architecture on environments in the POPGym suite (27) where match or outperform RNNs while also running nearly five times faster. We achieve strong results in the "Repeat Previous Hard" task, which previous models struggled to solve.
meta-learning setting in which meta-learn across random linear projections of the observation and action spaces of randomly sampled DMControl tasks. We show that S5 outperforms LSTMs in this setting. We then evaluate the models on held-out DMControl tasks and demonstrate dfitt---iittiuutiin performance to unseen tasks through in-context adaptation.
Finally, proposed
There are several possible ways to further investigate S5 models for reinforcement learning in future work. For one, it may be possible to learn or distill (23 entire reinforcement learning algorithms within an S5 model, given its ability scale to extremely long contexts. Another direction would be to investigate models for continuous time RI. settings [7]. While A, the discrete time between timesteps, is fixed for the original S4 model, S5 can in theory use a different A for each timestep.

---

Limitations: There are several notable limitations of this architecture and analysis. Firstly, im- plementing the associative scan operator is currently not possible using PyTorch, limiting us to tasks where short rollouts are sufficient to achieve
using the JAX K framework. Furthermore, good performance, S5 offers limited speedups, as rolling out across time is no longer a bottleneck. Finally. it was not possible to perform fully comprehensive hyperparameter sweep in our results in Section 4.4 because the experiments used significant amounts of compute.
Acknowledgments and Disclosure of Funding
Work funded by DeepMind. We would like to thank Antonio Orvieto, Robert Lange, Junhyok Oh, Greg Farquhar, Ted Moskovitz, and the rest of the Discovery Team DeepMind for their helpful discussions throughout the course of the project.
References
[1] Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Wradley-Schmies Michael Chang, Natalie Clay, Adrian Collister, Vib- havari Dasagi, Lucy Gonzalez, Karol Gregor, Edward Hughes, Sheleem Kashem, Maria Loks- Thompson, Hannah Openshaw, Jack Parker-Holder, Shreya Pathak, Nicolas Perez-Nieves, Nemanja Rakicevic, Tim Rocktäschel, Yannick Schroecker, Jakub Sygnowski, Karl Tuyls, Sarah York, Alexander Zacherl, and Lei Zhang. Human-timescale adaptation in an open-ended task space. arXiv e-prints, 2023.
[2] Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. A survey of meta-reinforcement learning. arXiv preprint arXiv:2301.08028, 2023.
[3] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Mundeerman-Minee and Qiao Zhang. JAX: composable transformations of Python+NumPP programs, 2018. URL
http://github. htt.huubbooomeegeeeeeeeeeeeeeee http:/uw...ooooggoo////////o/eee google/joeaeeoo
[4] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling, Advances in neural information processing systems, 34:15084-15097,
2021.
[5] Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder- decoder for statistical machine translation. arXiv preprint 1rXiv:1406.1007, 2014.
[6] Shmuel Bar David, Itamar Zimerman, Eliya Nachmani, and Lior Wolf. Decision s4: Efficient sequence-based rl via state spaces layers. In The Eleventh International Conference on Learning Representations, 2023. URL tikttps://ppereviiee
[7] Kenji Doya. Reinforcement learning in continuous time and space. Neural computation, 12(1): 219-245, 2000.
[8] Yan Duan, John Schulman, Xi Chen, Peter L. Bartlett, Ilya Sutskever, and Pieter Abbeel. RI2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv: 1.02799, 2016.
[9] Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry. Implementation matters in deep policy gradients: A case study on ppo and trpo. In International Conference on Learning Representations, 2020.
[10] Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state-space models. arXiv preprint 29Xiv:2202.09972) 2022.
[11] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474-1487, 2020.
10

---

[12] Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.003390 2021.
[13] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Combining recurrent, convolutional, and continuous-time models with linear state space layers. Advances in neural information processing systems, 34:572-585, 2021.
[14] James Hamilton. State-space models. Handbook of econometrics, 4:3039-3080, 1994.
[15] Matteo Hessel, Ivo Danihelka, Fabio Viola, Arthur Guez, Simon Schmitt, Laurent Sifre, Theo- phane Weber, David Silver, and Hado Van Hasselt. Muesli: Combining improvements policy optimization. In International Conference Machine Learning, pages 4214-4226. PMLR, 2021.
[16] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()): 1735-1780, 1997.
[17] Rein Houthooft, Yuhua Chen, Phillip Isola, Bradly Stadie, Filip Wolski. OpenAI Jonathan Ho, and Picter Abbeel. Evolved policy gradients. Advances in Neural Information Processing Systems, 31, 2018.
[18] Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mehta, and João G.M. Araújo. Cleanrl: High-quality single-file implementations of deep reinforcement learning algorithms. Journal of Machine Learning Research, 23(274): 2022. URL http:t http::////mm...ogggp/ppp////////
org/ pttppppppa//a1aaaaaaaaaa 1342111111111111 html
[19] Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent experience replay in distributed reinforcement learning. In International conference learning representations, 2018.
[20] Louis Kirsch, Sebastian Flennerhag, Hado van Hasselt, Abram Friesen, Junhyuk Oh, and Yutian Chen. Introducing symmetries to black box meta reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 7202-7210, 2022.
[21] Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. Gereral-purpose in- context learning by meta-learning transformers. arXiv preprint arXiv:2212.04448, 2022.
[22] Robert Tjarko Lange. gymnax: A JAX-based reinforcement learning environment library, 2022.
URL htsttp///ggihhbbb.co///bbeeettie ggymaa
[23] Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steiger- wald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:22100 14215, 2022.
[24] Eric Liang, Richard Liaw, Robert Nishihara, Philipp Moritz, Roy Fox, Ken Goldberg, Joseph Gonzalez, Michael Jordan, and Ion Stoica. Rllib: Abstractions for distributed reinforcement learning. In International Conference on Machine Learning, pages 3053-3062. PMLR, 2018.
[25] Chris Lu, Jakub Grudzien Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. arXiv preprint arXiv:2210005639 2022.
[26] Luke Metz, C Daniel Freeman, Samuel S Schoenholz, and Tal Kachman. Gradients are not all you need. arXiv preprint arXiv:2111.0580. 2021.
[27] Steven Morad, Ryan Kortvelesy. Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POP. Gym: Benchmarking partially observable reinforcement learning. The Eleventh International Conference on Learning Representations, 2023. URL
wwwperrviiiieeeeefrrueeee net/forum?
id=chDruttUTsOX
[28] Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Ré. S4nd: Modeling images and videos as multidimensional signals with state spaces. In Advances in Neural Information Processing Systems, 2022.
11

---

[29] Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free rl can be strong baseline for many pomdps. International Conference on Machine Learning, pages 16691-16723. PMLR, 2022.
[30] Junhyuk Oh, Matteo Hessel, Wojciech M Czarnecki, Zhongwen Xu, Hado P van Hasselt, Satinder Singh, and David Silver. Discovering reinforcement learning algorithms. Advances in Neural Information Processing Systems, 33: 33:1060-1070, 2020.
[31] John O'Keefe and Jonathan Dostrovsky. The hippocampus spatial map: Preliminary evidence from unit activity in the freely-moving rat. Brain research, 1971.
[32] Ian Osband, Yotam Doron, Mattco Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvari, Satinder Singh, al. Behaviour suite for reinforcement learning. arXiv preprint arXiv: 1908.03568, 2019.
[33] Emilio Parisotto and Ruslan Salakhutdinov. Efficient transformers in reinforcement learning using actor-learner distillation. arXiv preprint arXiv:2104.016655 2021.
[34] Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayaku- mar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, et al. Stabilizing transformers for reinforcement learning. In International conference on machine learning. pages 7487-7498. PMLR, 2020.
[35] Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines3: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268):1-8, 2021.
URL http://jmlr httpjimwrr.rrr.reeeeee vrg/papers/v22/
20- 1261.364hhh html
[36] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenherg, et al. A generalist agent, arXiv preprint arXiv:2205.00617.. 2022.
[37] John Schulman, F. Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv,
2017.
[38] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling, arXiv preprini arXiv:2208.0499 2022.
[39] Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. 2018.
[40] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: 1801 8006009 2018
[41] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long range arena: A benchmark for efficient transformers, arXiv preprint arXiv:2011.04006 2020.
[42] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
[43] Jane X Wang. Zeb Kurth-Nelson, Dhruva Tirumala, Hubert Soyer, Joel Z Leibo, Remi Munos, Charles Blundell, Dharshan Kumaran, and Matt Botvinick. Learning to reinforcement learn. arXiv preprint
2016.
[44] Ronald J Williams and David Zipser. Gradient-based learning algorithms for recurrent. Back- propagation: Theory, architectures, and applications, 433:17, 1995.
12

---

A Proof of Associativity of Binary Operator
Recall that defined as:
(aja ai,a aja & ai,b aic) (aj,a; aj.b; aj.c)
if a if a
dg 0j
This is equivalent to the following:
((a az)a: (a: az)o, ai,e) Jj
0j.c if Oj.c = 1
0g 1 ag
where is S5's binary operator defined in Equation 3 Note that associativity was proved in Smith al. 38
Using this, can prove the associativity of @
Let v. y, and z refer to three elements. We will prove that for all possible values of a, y, and 2, 4
retains associativity.
Case 1: Ze
v)
D
(10)
=y + 2
(11)
(y
(12)
Case 2: 2c =0 and ye
(x θ y ( 2 Note that (y/ + 2). thus,
(13)
(14)
(y 2)
(15)
Case 3: 2c 0 and Ve 0
(x + v) 0 ((x y)a (
v)., Ke) @
(16)
(((r
2)0= xe)
(17)
((
)a Ne)
(18)
ar + ((y
2)b ye)
(19)
a + (y ① 2)
(20)
13

---

B Hyperparameters
Table 2: Hyperparameters for training A2C on Bsuite
Parameter
Value
Adam Learning Rate
3e-4
Entropy Coefficient
0.0
Encoder Layer Sizes
[256, 256]
Number of Recurrent Layers
1
Size of Recurrent Layer
256
Discount y
0.99
TD À
0.9
Number of Environments
I
Unroll Length
32
Number of Episodes
10000
Activation Function
ReLU
Table 3: Hyperparameters for training PPO on POPGym
Parameter
Value
Adam Learning Rate
5e-5
Number of Environments
64
Unroll Length
1024
Number of Timesteps
15e6
Number of Epochs
30
Number Minibatches
8
Discount y
0.99
GAE X
1.0
Clipping Coefficient
0.2
Entropy Coefficient
0.0
Value Function Weight
1.0
Maximum Gradient Norm
0.5
Learning Rate Annealing
None
Activation Function
LeakyReLU
Encoder Layer Sizes
[128, 256]
Recurrent Layer Hidden Size
256
Action Decoder Layer Sizes
[128, 128]
Value Decoder Layer Sizes
[128, 128]
S5 Layers
4
S5 Hidden Size
256
S5 Discretization
ZOH
S5 A min
0.001
S5 A
0.1
14

---

Table Hyperparameters for training Muesli on Multi-Environmen Meta-RL. These experiments were run using 64 TPUv3's.
Parameter
Value
Adam Learning Rate
3e-4
Value Function Weight
1.0
Muesli Auxiliary Loss Weight
3.0
Muesli Model Unroll Length
1.0
Encoder Layer Sizes
(512,522
Number of Environments
1024
Discount
0.995
Rollout Length
1000
Offline Data Fraction
0.0
Total Frames
2c9
LSTM Hidden Size
512
Projected Observation Size
12
Projected Action Size
2
S5 Layers
10
S5 Hidden Size
256
S5 Discretization
ZOH
S5 A min
0.001
S5 A max
0.1
15

---

C POPGym Discussion
Stateless
Noisy Stateless CartPole Hard
Stateless Pendulum Hard
Noisy Stateless Pendulum Hard
Repeat Previous Hard
CartPole Hard
S5 (ours)
1.0 0.0
0.28 + 0.0
0.79 0.01
0.55 = 0.01
0.91 0.01
GRU (ours)
0.27 = 0.0
0.75 0.0
0.61 = 0.01
-
MLP (ours)
0.26
0.22 0.0
0.41 0.02
-0.88 + 0.00
GRU
1.0000
0.390 = 0.007
0.828 = 0.001
1657=
1.002
MLP
0.265 = 0.002
0.229 = 0.002
0.477 0.030
0.351 0.012
-0.486 0.002
IndRNN
1.000 = 0.000
0.404 = 0.005
1.044
0.521
-0.384
LMU
0.987 = 0.007
0.352 + 0.019
0.806 + 0.006
0.063
= 0.014
0.191
S4D
0.127 = 0.026
0.207 = 0.007
0.303 = 0.014
0.289 0.011
-0.491 +.000
FART
0.996 = 0.000
0.366 = 0.002
0.698 + 0.077
0.553 = 0.007
-0.485 + 0.001
Table 5: Results in POPGym's suite. The reported number the max mean episodic reward (MMER) used in Morad et al. 27 To calculate this, we take the mean episodic reward for each epoch, and then take the maximum over epochs. For our results above, the mean and standard deviation across eight seeds are reported. The results below are selected architectures from Morad et al. [27]. which also includes the best- perstrrrfrmingggg one from each environment. They report the mean and standard deviation across three trials.
Morad et al. 27 used RILib''s 24 implementation of PPO, which differs significantly from standard implementations of PPO. It uses dynamic KL-divergence coefficient on top of the clipped surrogate objective of PPO 37 Furthermore, they use advanced orchestration to return full episode trajectories, rather than using the more commonly-studied "stored state" [19] strategy.
Instead, we follow the design decisions outlined in the StableBaselines3 [35 and CleanRL's [8] recurrent PPO implementations. While this results in different results shown in Table 5 for the same architecture, it recovers similar performance across the environments. Notably, our architecture far outperforms the best performing architecture in Morad et al. 27 in the "RepeadPrewiousHHard" environment.
We used the learning rate, number of environments, unroll length, timesteps, epochs, and minibatches, GAE A, and model architectures from Morad et al. [27]. However, we used the standard clipping coefficient 0.2 instead of 0.3 account for the lack of dynamic KL-divergence coefficient. Note that we also adjusted the S5 architecture to contain approximately the same number parameters as the GRU implementation instead of matching the size of the hidden state, which was done in Morad et al. (27).
We did not evaluate our architecture across the full POPGym suite. To enable more rapid experimen- tation, we implement our algorithms and environments end-to-end JAX R 25). While the original POPGym results took 2 hours per trial with a GRU with a Quadro RTX 8000 and 24 CPUs, we could run our experiments using only 3 minutes per trial on an NVIDIA A40. Because of this, we selected environments from Morad et al. [27] to implement in JAX. We chose the CartPole, Pendulum, and Repeat environments because they are modified versions of existing environments in Lange [22). We found that the "Easy" and "Medium" versions of these environments were not informative, as most models perform well on them and only report the "Hard" difficulty results.
We attach our code in the supplementary materials.
16


Reviews:
[{"review_id": "X36z0LD46m", "reviewer": "Reviewer_R54Y", "summary": "The authors propose a modification of the S5 sequence architecture with a resettable hidden state that leads to a drop-in replacement of RNNs and Transformers in partially observed / memory-intensive RL tasks. They show that S5 exceeds baseline performance while being more computationally efficient to train than LSTMs due to S5’s parallelization and faster at inference than Transformers’ N^2 attention.", "questions": "Please clarify if I am wrong in thinking that the DMControl tasks are zero-shot. The evaluation procedure is not as clearly discussed as it often is in meta-RL papers, and the appendix did not give more details.\n\nWere there any experiments to dramatically increase the context length towards the range seen with S4 in LRA? I’m curious if there are insights on limitations or changes that would be needed to make this work in an RL setting. According to Table 5 the DM Control experiments used 64 TPUv3s, so I assume compute was not the limiting factor here.", "limitations": "There is no Limitations section or a clear discussion on limitations. However, the paper is proposing an architectural change to an existing method that has known limitations, and some of those limitations are improved by the S5 model.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "State space models are a natural fit for in-context meta-RL and long-term memory because of their unique combination of train-time parallelism and test-time constant inference. They are also thought to be more compatible with the continuous inputs seen in RL than other long-sequence domains like NLP that use discrete tokens. In-context RL has a sequence length barrier that goes somewhat unnoticed because common benchmarks do not require adaptation over more than a few steps. Applying S5 and future versions here may have great long-term benefits. \n\nThe experiments do clearly demonstrate that resettable S5 can replace LSTMs and Transformers in partially observed tasks.\n\nThe paper also proposes a meta-learning version of the DM Control suite using randomly projected state and action spaces. While this is a somewhat arbitrary way to expand the meta-distribution of tasks from a small set of popular benchmarks, it is probably better than the common alternative of turning similar environments into goal-conditioned tasks and then hiding the goal (Ant Goal, Cheetah Fwd-Back, Humanoid Dir, ...). This benchmark will be useful for future work.\n\nThe method could be applied to most agents that support S5's computation, and should be easy to use and reproduce.", "weaknesses": "S5 would be best suited for extending the sequence lengths used by current methods, and could unlock a new level of difficulty in long-context RL. However, the experiments here stop short of exploring those limits. Instead, sequence lengths are capped at ~1k where LSTMs/Transformers are feasible but slower and more expensive. Wall-clock efficiency is a nice advantage but could also be improved by other factors like using a more efficient base algorithm. If the goal was to show that S5 is a valid substitute at existing sequence lengths while leaving expansion for future work, it would have been simpler to use more standard benchmarks in the meta-RL portion and compare to external baselines. \n\nIn general, the narrative of the paper gets a bit lost in the gray area between more traditional long-term memory and “in-context” meta-learning. The need for long sequences is motivated by multi-episodic in-context learners like RL^2 (lines 36-42), but the method (Sec. 3.1 and Alg 1) does not distinguish between episode resets (which would not reset the sequence model’s hidden state) and task resets. The first two experiments are focused on standard long-term memory, while the third appears to only evaluate zero-shot generalization to a meta-distribution of tasks. It would have been more interesting to evaluate RL^2-style multi-episodic learning, which would naturally extend the sequence length to highlight S5’s core advantages.\n\nThe need for a resettable hidden state within S5 is presented as a widespread barrier to its application in RL. But this is primarily caused by the data collection and optimization implementations common to on-policy policy gradient algorithms with parallel actors. I think it would be helpful to explain this issue in more detail than is provided (lines 111-144, 43-48). For example, there would be no need for the modified architecture when substituting S5 for RNNs/Transformers in an off-policy in-context learner like in [Ni et al.](https://arxiv.org/abs/2110.05038) The issue may also be avoided by batching the policy updates into padded sequences that do not cross episode boundaries.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to thank the reviewer for their well-thought review. The reviewer brings up many good points that we would like to address. \n\nFirstly, we are glad that the reviewer finds that our paper clearly demonstrates that resettable S5 can replace LSTM’s and Transformers and that our proposed benchmark can be used for future works. We would like to also highlight that our code is provided in the supplementary material, which should further help with reproducibility and ease of use.\n\n## On Longer Contexts\n\n> sequence lengths are capped at ~1k where LSTMs/Transformers are feasible but slower and more expensive\n\n> It would have been more interesting to evaluate RL^2-style multi-episodic learning, which would naturally extend the sequence length to highlight S5’s core advantages.\n\nThese are fair points! We originally wished to find benchmarks where longer sequence lengths would be helpful; however, we struggled to find scalable benchmarks (i.e. beyond POPGym), likely because current architectures would not be able to take advantage of the increased context. Because of this, we introduced our new meta-learning version of the DMControl Suite. We had expected that more in-context adaptation trials would be required because of the large space of tasks; however, our architecture was still able to perform near-optimally with just one trial. *We have since developed and evaluated the architecture on a version of StatelessCartPole where we simply perform random linear projections of the observation.* This involves sequences of length of up to $6400$ and we analyze its performance across multiple adaptation trials. We encourage the reviewer to read our general rebuttal and attached plots in Figure 8.\n\n>Were there any experiments to dramatically increase the context length towards the range seen with S4 in LRA? I’m curious if there are insights on limitations or changes that would be needed to make this work in an RL setting.\n\nWe agree with the reviewer that long-range tasks, especially meta-RL tasks, would be ideal for investigating dramatically increased context lengths. Firstly, we struggled to find appropriate benchmarks that require many thousands of steps of context. Ni et al. [1] find that for the vast majority of common POMDPs (including meta-RL tasks), short context lengths are often optimal -- and anything over just $100$ is considered “long”. It’s unclear if this is because long contexts are generally not useful, or if current benchmarks were designed with existing architectures in mind. Secondly, in such settings, credit assignment and long-horizon discounting become increasingly challenging. For example, the common discount factor of 0.99 (and even 0.999) vanishes after several thousands of timesteps. We have since added a task that uses sequence length of up to $6400$ in Figure 8 of our general rebuttal.\n\n## Others\n\n> But this is primarily caused by the data collection and optimization implementations common to on-policy policy gradient algorithms with parallel actors. I think it would be helpful to explain this issue in more detail than is provided\n\nThis is a good point. The explanation in the current paper can be unclear (and it is impressive that the reviewer understands the nuances of these implementation details given that we did not fully explain them). We have since re-written this section to further explain why on-policy algorithms with parallel actors often make use of resettable hidden states and thank the reviewer for the suggestion. Indeed, in the off-policy setting, implementations have more control over data organization.\n\n> The issue may also be avoided by batching the policy updates into padded sequences that do not cross episode boundaries.\n\nYes, this is also a potential approach -- however, padded sequences often involve large amounts of inefficient computation and would be significantly slower. If one is running on-policy algorithms they are usually concerned with wallclock-time rather than just sample efficiency. We’ve updated the paper to discuss this!\n\n> the DMControl tasks are zero-shot. The evaluation procedure is not as clearly discussed as it often is in meta-RL papers, and the appendix did not give more details.\n\nYes, the DMControl tasks are zero-shot, with no fine-tuning involved. The agent is given only a single trial to adapt in-context. We have updated the manuscript to clarify this point further. Thank you for the suggestion!\n\n[1] Ni, Tianwei, Benjamin Eysenbach, and Ruslan Salakhutdinov. \"Recurrent model-free rl can be a strong baseline for many pomdps.\" arXiv preprint arXiv:2110.05038 (2021).\n\n---\n*We hope that most of the reviewer’s concerns have been addressed and, if so, they would consider updating their score. We’d be happy to engage in further discussions.*\n"}, {"review_id": "LfTmAZhF5f", "reviewer": "Reviewer_Zjtv", "summary": "Structured state space sequence (S4) models deliver good performance on long-range sequence modeling tasks, fast inference speed and parallelize training, making them suitable for many RL settings. The authors propose a modification to the recently proposed S5 architecture and apply it to RL tasks. Their proposed model outperforms Transformers in terms of runtime and memory complexity (on a toy task), and RNNs in terms of task performance. They aim to show the efficacy of SSMs for in-context adaptation to new task variations.", "questions": "- Figure 2:\n  - How do models perform beyond a sequence length of 512? S5 and self-attention reach the same level of performance at all lengths. How does this change for longer sequence lengths (at least up to 2048)? \n    Do Transformers outperform S5 beyond some context length? \n    No standard deviations, only single seed?\n  - How does the memory consumption differ between the compared methods?\n\n- Figure 3, Table 2:\n  - What sequence length is used? Is a long sequence length even required for these tasks?\n  - The performance gains of S5 over GRU come from the “Repeat Previous Hard” task. Why is this the case? On all other tasks, they perform the same. Seems like the benchmark is too easy to solve and not a suitable test-bed.\n  - Why are the standard deviations 0.0 for all methods on cartpole? Seems unlikely across 8 seeds (especially in the noisy setting).\n\n- Section 4.3:\n  - The setup for this experiment needs clarification: \n    - Why does this experiment demonstrate in-context learning abilities? Does the model always only observe observations from the current episode? \n    - Are you training on held-out tasks for 2B steps (as shown on the x-axis) or just doing inference? From the learning curves, it looks like weights are updated. Why is this “evaluation” then? Please clarify.\n    - How would an agent trained from scratch compare against the pre-trained one? Learning the new tasks within 2B steps should not be an issue. \n    - Are you resetting S5 at episode boundaries during evaluation? Wouldn’t it be beneficial to maintain the context across episodes to encourage in-context learning abilities? It seems like resetting is a disadvantage here. \n    - How would the agents perform if trained and evaluated without random projections? \n\n  - Figure 7: \n    - In general, the considered tasks are very similar to the pre-training tasks. Four of the five tasks are minor variations of them. Is it correct to consider them OOD? \n    - Importantly, the model fails completely on pendulum_swingup. This suggests that, the task distribution is too narrow. Conducting this experiment on a larger number of tasks or robot morphologies may be more insightful.\n    - Why is “finget_turn_easy” in the OOD tasks, but the hard variant “finget_turn_hard” in the pre-training tasks. Shouldn't it be the other way round?\n  - Missing ablation on the proposed modification to S5. How does performance change when using the proposed modification vs. without?\n  - Please report parameter counts of the compared methods for all experiments (only reported for POPGym in the Appendix). Do all architectures have approximately the same amount of parameters?\n ", "limitations": "- Limitations of the proposed architecture have not been discussed sufficiently. As there is no comparison against the Transformer architecture, it is hard to assess the limitations of the proposed architecture.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 3, "strengths": "- The authors are the first (to the best of our knowledge) to leverage S4-like models for RL. This is the main novelty of the paper, and we consider this an important contribution.\n- The authors propose a modification/fix for the S5 architecture to enable handling sequences of varying lengths in and RL setting.\n- They show that S4-based models have advantages over commonly used Transformers (runtime, memory complexity) and RNNs (task performance) on a toy task (Figure 2).\n", "weaknesses": "**Major concerns**:\n\nWhile the method to add a mechanism for resetting the state of S5 layers seems sound the proposed evaluation and in-context RL setting do not make sense to me.\nThe authors themselves state, that using S5 layers for RL is problematic as such models would be \"accessing memory and context from other episodes\".\nWhy then evaluate the proposed solution in the in-context setting, where accessing context (e.g. from other episodes) is required?\n\nFurthermore, I do not understand what the authors mean with \"in-context RL\", and apart from the conclusion the authors never even mention it again even though its in the title. \nInstead, the authors perform meta-RL experiments by training models on multiple tasks simultaneously, however even here I don't understand the evaluation protocol.\nI would not call fine-tuning for another 2 billion timesteps an \"evaluation\". \nHow does this experiment show \"in-context adaptation\"?\n\nFinally, the authors argue that comparison to transformer based models is not possible due to poor runtime performance. \nHowever, at least for the POPGym environment this should have been possible, seeing as the authors state in the appendix that the runtime per trial with GRU is only 3 minutes.\nAnd for in-context experiments a transformer-based baseline should definitely be included.\n\n**Minor comments:**\n\n- Algorithm 1 seems redundant and can be moved to the Appendix, as it only shows a meta-RL loop with random projections.\n- Figure 1 is too large for its information content.\n- Line 166: “mmemory”\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to first thank the reviewer for their extremely thorough review. We are glad that the reviewer finds that investigating S4-like models for RL is an important contribution, especially since they have not been widely-adopted or thoroughly investigated in RL.\n\n### In-Context Learning and Meta-Learning\n\n> Why evaluate...where accessing context (e.g. from other episodes) is required?\n> Are you resetting S5 at episode boundaries during evaluation?\n\nRL^2-like meta-RL methods, an “episode” can consist of multiple “trials” (the terminology on this is inconsistent between papers in the field). In meta-RL it is important to access context between different *trials* but not *episodes*. We have updated the manuscript to clarify this.\n\nIn our meta-RL setting we only allow the agent access to **one trial** because it already achieves near-optimal performance in the underlying tasks through within-trial adaptation. ***We have attached additional results to show multi-trial in-context adaptation in the general rebuttal  Figure 8.***\n\n> I do not understand what the authors mean with \"in-context RL\"...How does this experiment show \"in-context adaptation\"?\n\nWhile in-context learning is now commonly used to refer to few-shot learning in transformers, we are referring to *in-context changes in behavior (adaptation)* when referring to “in-context RL”. We hope Figure 8c makes this clearer, where we show the agent learns across trials!\n\n> Are you training on held-out tasks for 2B steps? \n> How would an agent trained from scratch compare?\n\nThose plots are for evaluation throughout *meta-training*. The held-out tasks are thus evaluated at different points throughout *meta-training*. During the evaluation, we are just doing inference (including on held-out tasks). An agent trained from scratch could not solve an out-of-distribution task in one episode.\n\n> The considered tasks are very similar to the pre-training tasks\n\nIt would require a very large collection of pre-training tasks (or inductive bias on the meta-learner) to enable any transfer to far out-of-distribution tasks. We believe that showing transfer to close, but still OOD tasks, is still a positive result that shows robustness and generalization.\n\n> Why is “finget_turn_easy” in the OOD tasks?\n\nThe “easy” and “hard” versions of tasks are variants of the task with different sets of parameters or dynamics. Transferring either way demonstrates OOD performance. The choice of which versions were “held-out” was arbitrarily selected.\n\n> How does performance change when using the proposed modification vs. without?\n\nThe experiments for the original DMControl suite are slow and expensive to run. We have instead performed a similar experiment using random linear projections of POPGym’s StatelessCartPole, which we share in the general response Fig 8. The results show that performance improves significantly when using the proposed modification.\n\n> comparison to transformer based models\n\nTransformer-based models are far too slow for us to run in the meta-DMControl setting and are an *uncommon architecture choice for on-policy RL* because of their runtime. While transformers can show fast training on supervised learning tasks, they are slow in reinforcement learning due to their poor inference speed, especially when using 2B frames.\n\n> same amount of parameters?\n\nOne can reconstruct the sizes from the hyperparameters.\n\nS5 Encoder: 1975040 parameters\n\nLSTM Encoder: 2099200 parameters\n\n### BSuite Results\n\n> How do models perform beyond...(at least up to 2048)?...only single seed?\n\nThe original bsuite memory length evaluation only goes up to $105$. $512$ is already nearly 5x that number. Credit assignment gets increasingly infeasible as the length gets longer, especially since bsuite evaluation uses a fixed number ($10000$) of episodes for training, a fixed $\\gamma=0.99$, and a basic A2C algorithm. We are unlikely to get *any meaningful information* about an architecture’s capabilities past 512 steps.\n\nFurthermore, we expect self-attention to grow *cubically* in runtime since the number of steps grows *linearly*, and the cost-per-step grows *quadratically*. Sequence length $4096$ would be expected to take several days for self-attention. We are currently running these experiments.\n\nIn the meantime, we have reported results for sequences up to $2048$ (Fig 9). Our original results reported the median score across $5$ seeds. Our updated ones in Figure 9c report the mean and standard error across $5$ seeds. \n\n> How does the memory consumption differ?\n\nWe report this in the rebuttal! Figure 9a.\n\n### POPGym Results\n\n> What sequence length is used? Is a long sequence length even required? \n\nA sequence length of 1024 is used. We follow the training hyperparameters outlined in POPGym [1]. Even if long sequence lengths are not required, S5 *still trains six times faster and outperforms the baselines*. POPGym is not designed to test learning across long sequences; however, it is useful for evaluating architectures for partially-observed RL. We have experiments in Figure 8 that require long sequences. \n\n> The performance gains of S5 over GRU come from the “Repeat Previous Hard” task. Why is this the case?\n\nWe’re not sure why S5 performs better. POPGym’s best performing architecture, the LMU (MMER of $0.19$ vs. S5’s $0.91$), shares the underlying theory of continuous time representations. \n\n> Why are the standard deviations 0.0...on cartpole? \n\nThe choice of hyperparameters leads to very consistent policy updates since they train for many epochs on large batch sizes. For StatelessCartPole, all methods achieve perfect score. For NoisyStatelessCartPole, there are very small standard deviations at more digits.\n\n[1] Steven Morad, et al. “POP- Gym: Benchmarking partially observable reinforcement learning.”\n\n---\n*We hope that most of the reviewer’s concerns have been addressed and, if so, they would reconsider their assessment. We’d be happy to engage in further discussions.*"}, {"review_id": "sOleuF6AnE", "reviewer": "Reviewer_R3ow", "summary": "This paper investigates the effectiveness of structured state-space sequence (S4) models and in particular its variant S5 in reinforcement learning settings. To apply S5 to reinforcement learning, the authors propose a modified associative operator that handles episodic resets, allowing S5 to train over sequences spanning multiple episodes. Across a suite of partially observed RL environments from bsuite and POPGym, S5 outperforms RNN and Transformer baselines while being significantly more efficient at training and inference. To further evaluate the S5's ability to perform in-context RL and generalize out-of-distribution, the authors introduce a multi-environment meta-learning task based on DM Control where each meta-episode features different random projections of observation and action spaces. Not only does S5 outperform MLP and LSTM on in-distribution tasks, but it demonstrates better generalization to held-out tasks as well.", "questions": "- What is the reason behind the improved OOD generalization of the S5 model compared to baselines?\n- From Table 2, the proposed operator does not significantly outperform the vanilla S5 operator. Can you include more comparisons with the vanilla operator? E.g. by evaluating the vanilla S5 model on the meta-learning tasks.\n- According to [1], S4D is the worst-performing method out of all baselines. Is there a reason behind the huge improvement in asymptotic performance?\n\n[1] Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POP- Gym: Benchmarking partially observable reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.\n\n\n", "limitations": "The authors have not adequately addressed the limitations in the paper. I would imagine the limitations to be similar to those of state-space models. For example, they generally perform worse than RNN on partially observed RL settings, and they are hard to scale to higher dimensions.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- The paper proposes a novel associative operator which allows the S5 model to handle episodic resets in RL settings. In doing so, they demonstrate the effectiveness of the S5 model for partially-observed and meta-RL tasks, both in terms of asymptotic performance and training/inference speed.\n- The meta-RL environment with random observation and action projections can be used as an evaluation benchmark for future work.", "weaknesses": "- The proposed method is a rather marginal modification to the S5 model. And according to the results in Table 2, the new operator does not significantly outperform the vanilla S5 operator.\n- It would be fine to thoroughly evaluate an existing method in a new setting, making it an empirical analysis paper. However, the experiments in this paper are insufficient. The bsuite memory length is a toy environment. And there is a lack of baselines and variety of tasks in the POPGym experiments. It is elusive why the authors include MLP but not other baselines from [1].\n- A major benefit of state-space models is their ability to handle significantly longer sequences than RNNs and Transformers, but the paper does not demonstrate an RL setting where this can be useful.\n- The meta-learning setting, while novel, feels a bit contrived. It is unclear what the implications of this setting are.\n\n[1] Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POPGym: Benchmarking partially observable reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their insightful feedback. We are happy to hear that the reviewer agrees that we “demonstrate the effectiveness of the S5 model for partially-observed and meta-RL tasks, both in terms of asymptotic performance and training/inference speed” and that “the meta-RL environment with random observation and action projections can be used as an evaluation benchmark for future work.” We would like to address the reviewer’s concerns.\n\n> The proposed method is a rather marginal modification to the S5 model...in Table 2, the new operator does not significantly outperform the vanilla S5 operator\n> From Table 2, the proposed operator does not significantly outperform the vanilla S5 operator.\n\nIn Table 2, we show that the new operator significantly outperforms the vanilla S5 operator in $3/5$ of the environments, and matches in the remaining $2/5$. In the most challenging environment, the new operator raises the score from $0.76$ to $0.91$. We believe this is very significant, and visualize it more clearer in Figure 10 of our rebuttal.\n\nWe show this further in the shared rebuttal. The Meta-StatelessCartPole (Figure 8a and 8b) experiments further demonstrate its significance on a task that requires longer memory.\n\nFurthermore, state-space models in on-policy RL in general have not been rigorously studied and are not widely-used. The operator is just one of the contributions. We hope this paper and these results can encourage further adoption of this architecture to speedup and improve future research involving memory-based architectures for RL.\n\n> It is elusive why the authors include MLP but not other baselines from [1].\n\nWe report GRU results since it was generally the most consistent best-performing baseline evaluated in [1]. It is also one of the most widely-used. We include the MLP to demonstrate the performance vs. runtime tradeoff in Figure 3.\n\n> A major benefit of state-space models is their ability to handle significantly longer sequences than RNNs and Transformers, but the paper does not demonstrate an RL setting where this can be useful.\n\nIndeed, that is *a* benefit of state-space models. There are *many other benefits* of state-space models, many of which we show in the paper! In particular, we get *significant speedups* (over 6x faster on POPGym!) while also *significantly outperforming the baselines*. \n\n***We have also attached further additional results to show a scenario where longer sequence adaptation of up to $6400$ is useful. Please read the general rebuttal for plots (in Figure 8) and experimental setup***.\n\n> The meta-learning setting, while novel, feels a bit contrived. It is unclear what the implications of this setting are.\n\nMost existing RL settings do not require long sequence lengths (arguably *because* existing architectures cannot handle them), so we designed this meta-learning setting to construct one that does. \n\nIn the future, a more scaled up version of this setting with many more tasks and data augmentation could result in a *general* meta-learned reinforcement learning agent that can adapt to much farther OOD tasks.\n\n> What is the reason behind the improved OOD generalization of the S5 model compared to baselines?\n\nWe believe it is because the S5 model can incorporate information over longer sequences, which we demonstrated in the Bsuite experiments (Figure 2). The ability to incorporate longer contexts could have led to a more general adaptation mechanism than one that can only adapt to recent transitions.\n\n> According to [1], S4D is the worst-performing method out of all baselines. Is there a reason behind the huge improvement in asymptotic performance?\n\nWe hypothesize that it is because their implementation only uses one layer of S4D. While for most recurrent networks, stacking additional recurrent layers does not help in many tasks, including RL (see the architectures used for [2] and [3]), for state-space models it is extremely important for performance since it is a completely linear recurrence (and can only achieve non-linearity through depth).\n\n>limitations to be similar to those of state-space models. For example, they generally perform worse than RNN on partially observed RL settings, and they are hard to scale to higher dimensions.\n\nWe agree with the reviewer and have since added a limitation section, which includes implementation challenges (e.g. the lack of associative scans in PyTorch) and potentially diminishing speedup returns on tasks with short sequences or other computational bottlenecks (such as vision-based tasks). However, we are not sure why the reviewer believes that SSM’s perform worse than RNN’s on partially observed RL settings or that they are hard to scale to higher dimensions. Our experiments show that they perform significantly better on partially observed RL settings. Other concurrent work [5] has shown similar results in Offline RL and have scaled them to higher dimensions [4].\n\n---\n*We hope that most of the reviewer’s concerns have been addressed and, if so, they would reconsider their assessment. We’d be happy to engage in further discussions.*\n\n---\n\n[1] Steven Morad, Ryan Kortvelesy, Matteo Bettini, Stephan Liwicki, and Amanda Prorok. POP- Gym: Benchmarking partially observable reinforcement learning. In The Eleventh International Conference on Learning Representations, 2023.\n\n[2] OpenAI, C. Berner, et al. \"Dota 2 with large scale deep reinforcement learning.\" arXiv preprint arXiv:1912.06680 2 (2019).\n\n[3] Vinyals, Oriol, et al. \"Grandmaster level in StarCraft II using multi-agent reinforcement learning.\" Nature 575.7782 (2019): 350-354.\n\n[4] Deng, Fei, Junyeong Park, and Sungjin Ahn. \"Facing off World Model Backbones: RNNs, Transformers, and S4.\" arXiv preprint arXiv:2307.02064 (2023).\n\n[5] David, Shmuel Bar, et al. \"Decision S4: Efficient Sequence-Based RL via State Spaces Layers.\" The Eleventh International Conference on Learning Representations. 2022.\n"}, {"review_id": "38Anv4M4TW", "reviewer": "Reviewer_gUv9", "summary": "The authors propose a modification to S5 that enables \"resetting\" the recurrent state, allowing it to function as an RNN replacement in RL. They update the scan operator to utilize the `done` flag and use this to reset the recurrent state. They evaluate the resettable S5 on a portion of the POPGym suite and the bsuite memory task. They show that the S5 is able to solve long-term memory tasks previous methods were unable to solve.\n\nNext, they implement a metalearning approach from Kirsch et al. that enables learning across various action and observation space sizes by using random linear projections of the spaces to a fixed-size vector. They leverage S5's long-term memory to obtain good results across DMLab, on both in-distribution and out-of-distribution tasks.", "questions": "- The reuse of $a$ is confusing in the scan terms, since $a$ is also the subscript referring to the $A$ matrix. Perhaps a different term would be more clear?\n- Figure 3: This is across all POPGym envs or just a subset? From briefly reviewing the code I assume a subset.\n- I'm not an expert on meta rl, but it is very promising that their approach can do well on out-of-distribution tasks. \n- Experiments and implementation are the strength of this paper, I just wish they just put the S5 reset in the appendix or footnote", "limitations": "The authors do not have a limitations section. It is probably worth adding a short section on this.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and addresses a painful issue: that partially observable policies are inefficient to train. They show that S5 is the first model to solve the difficult RepeatHard task from POPGym. It is clear the author put a ton of work into implementation and put care into baseline implementations to ensure a fair comparison. The experimental setup is quite broad in the sense that it covers both POMDPs and metalearning.", "weaknesses": "#### Main weakness:\nPerhaps I am missing some information, and if so, please correct me. But the reset, the main contribution of the paper, appears to be a trivial change to S5. Line 94 already provides the initial state of vanilla S5, which is $e_0 = (I, 0)$. Their reset is just setting the recurrent state to the initial state when they receive a `done` flag from the environment. This is the standard for any recurrent model in RL, so I do not believe this counts as a novel contribution.\n\nWith a slight change of notation, the S5 scan operator is\n\n$f(a_t, b_t, a_{t-1}, b_{t-1}) = \\begin{bmatrix}\na_t \\odot a_{t-1} & \na_t \\otimes b_{t-1} + b_t \n\\end{bmatrix}$\n\nThe authors propose that at $t=0$, since we do not have $a_{t-1}, b_{t-1}$ we evaluate $f$ using $a_{t-1} = I$ and $b_{t-1} = 0$\n\n$f(a_t, b_t, I, 0) = \\begin{bmatrix}\na_t & \nb_t \n\\end{bmatrix}$\n\nSo to reset: just put the predefined vanilla S5 initial state $e_0$ and into S5.\n\n#### Other Weaknesses\n- The authors show very promising performance on Repeat Hard and a few stateless cartpole environments. That said, POPGym has something like 50 total tasks, so it is a bit disappointing that they picked the easiest control tasks, as I imagine S5 would do much better than GRUs on harder tasks.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We would like to first thank the reviewer for their detailed and technical review. We are glad that the reviewer finds our approach promising and the experiments fair and broad.\n\n### On the Reset\n\n>the reset, the main contribution of the paper, appears to be a trivial change to S5\n\nIt’s not immediately obvious why the parallel reset is non-trivial. Indeed, for any recurrent model in RL, resetting the hidden state during inference (and training) is standard. However, it’s not clear how to do this when one is *parallelizing across the time dimension* during *training/backpropagation* given a batch of data of shape $(Minibatch Size, Sequence Length, Observation Size)$. For example, most S4-like models usually implement their backwards pass using *convolutions* (instead of recurrence) to parallelize across time -- it’s not clear how or if we can adjust the convolution kernel to account for “dones” from the environment.\n\nInstead, S5 uses *associative scans*, which leverages a *binary associative operator* to parallelize across the time dimension. While it’s immediately obvious that simple operators, such as multiplication and addition, can be implemented associatively, previous implementations have not shown how to implement “reset” function associatively. Our contribution is introducing such an operator (Line 132) that computes the desired output (Line 135) and is provably associative (Appendix A).\n\nWith this in mind:\n\n>So to reset: just put the predefined vanilla S5 initial state $e_0$ into S5\n\nWe assume the reviewer is asking why we cannot just insert $e_0$ where there are “dones”. We believe this would not compute the desired recurrence. While indeed (using the reviewer’s notation) $f(a_t, b_t, I, 0)$ computes the desired value, note that $f(I, 0, a_{t-1}, b_{t-1})$ (i.e. the other side of the associative operation) returns $[a_{t-1} b_{t-1}]$. Thus, when scanning over the sequence $(e_{t-1}, e_0, e_t)$, we would get $f(a_t, b_t, a_{t-1}, b_{t-1}) \\neq [a_t, b_t]$. This means it does not reset the hidden state. At no point would inserting $e_0$ actually *reset* the hidden state of the scan -- it would merely repeat the previous element. \n\nFurthermore, inserting elements into the scan would usually involve dynamic shapes or masking, which is often inefficient or undesirable when performing batch learning.\n\nWe have updated our manuscript to clarify why we cannot just use $e_0$ to perform the reset\n\n---\n\n*Because this was the reviewer’s primary concern, we hope that our response has clarified this contribution and that the reviewer will let us know of any further concerns on this topic or, otherwise, consider updating their score.*\n\n---\n\n> The reuse of $a$ is confusing in the scan terms\n\nGood catch! We’ve updated the paper to reflect this change\n\n### On the tasks\n\n> POPGym has something like 50 total tasks, so it is a bit disappointing that they picked the easiest control tasks\n> This is across all POPGym envs or just a subset?\n\nIndeed, this is just a subset of the POPGym envs. While POPGym does have many tasks, many of them are simply easier versions of the ones we evaluated on. For example, we evaluate only on the “Hard” version of these tasks while there are “Medium” and “Easy” versions that are not particularly informative. Unfortunately, there are no harder control tasks in POPGym beyond CartPole and Pendulum; however, we evaluated on the environment that most architectures struggled the most on: RepeatPreviousHard. The best-performing architecture in POPGym reported a score of 0.191 while we obtained a score of 0.91.\n\nWe’ve since added further additional results on a long-horizon version of POPGym StatelessCartPole that involves sequences of length up to $6400$ to increase the difficulty and context. We’ve included a brief explanation and analysis in the general response. \n\n> It is probably worth adding a short section on [limitations].\n\nGood point -- thanks for the feedback! We have included this in our manuscript."}]

Author rebuttal:
We are grateful to the reviewers for their insightful feedback. We appreciate the consensus that the proposed S5 architecture offers **clear advantages over standard RNNs and Transformers in partially-observed RL, both in terms of performance and runtime.** This is the key takeaway of our work, and we hope it will accelerate future research in RL.

$\color{red} R1$ (R54Y): “experiments do clearly demonstrate that resettable S5 can replace LSTMs and Transformers in partially observed tasks”

$\color{green} R2$ (Zjtv): “They show that S4-based models have advantages over commonly used Transformers (runtime, memory complexity) and RNNs (task performance) on a toy task.”

$\color{blue} R3$ (R3ow): “they demonstrate the effectiveness of the S5 model for partially-observed and meta-RL tasks, both in terms of asymptotic performance and training/inference speed.”

$\color{magenta} R4$ (gUv9): “S5 is able to solve long-term memory tasks previous methods were unable to solve.”

We are glad that reviewers found the associative operator to be a “novel” ($\color{blue} R3$) modification that enables “drop-in replacement of RNN’s” ($\color{red} R1$) by “handling sequences of varying lengths” ($\color{green} R2$), although there are concerns with its significance ($\color{blue} R3$, $\color{magenta} R4$). We added Figure 10 to the general rebuttal to demonstrate its impact, along with results in Figure 8.

Some reviewers also found the meta-learning setup to be useful for future work ($\color{red} R1$, $\color{blue} R3$) and the results “very promising” ($\color{magenta} R4$), although some understandably have concerns over its clarity ($\color{green} R2$) and limited single-trial sequence length ($\color{red} R1$, $\color{green} R2$, $\color{blue} R3$) of $1000$, which we address below.

The results on POPGym demonstrated over *6x speedups* and better performance (being the “first model to solve the difficult RepeatHard task from POPGym” [$\color{magenta} R4$]) over the state-of-the-art GRU. Reviewers had concerns over the lack of Transformer-based baselines ($\color{green} R2$), lack of difficult tasks ($\color{green} R2$, $\color{blue} R3$, $\color{magenta} R4$), and short context lengths ($\color{red} R1$), which we address below.

## Meta StatelessCartPole

To address concerns about sequence length, in-context learning, and POPGym tasks, *we have run additional experiments combining POPGym’s StatelessCartPole task with randomly-projected observations.* We allow the agent to have $16$ trials per episode. We show in-context learning results in Figure 8(c) and demonstrate that it can even maintain performance to $32$ trials, evaluating on sequence lengths of up to *6400* -- which is a very long sequence length for reinforcement learning. We show training results and runtime in Figure 8(a) and 8(b) -- S5 still outperforms GRU’s while running significantly faster. 

We hope this addresses the reviewer’s concerns around the lack of multiple trials, long context lengths, and challenging tasks.

## Transformers

To further address $\color{green} R2$’s concerns about the lack of transformer-based baselines, we included longer context results in bsuite Figure 9 for lengths up to $2048$ and include memory usage statistics. Note that bsuite’s original evaluation only extends to a length of $105$, and is likely unsuitable for such long sequences (since it uses a fixed episode budget). We are currently running even longer sequences; however, they are expected to take days to complete and will not change the takeaway.

Furthermore, we ran a Transformer baseline on POPGym environments that have *fixed episode lengths*. We then modify the trajectory collection such that it *always collects exactly one episode per rollout* while maintaining an identical total batch size. Without this constraint, it is challenging to efficiently run Transformers in on-policy RL while maintaining its memory. We perform the same for other architectures across three seeds and report the mean and standard error below. Transformers are fast in these settings since the horizons are short and they can parallelize across time (unlike LSTM’s).

| Architecture | stateless pendulum hard MMER | noisy stateless pendulum hard MMER | repeat previous hard MMER |
|---|---|---|---|
|S5|**0.805±0.003**|0.575±0.002|**0.986±0.000**|
|Transformer|**0.796±0.005**|0.544±0.003|-0.462±0.003|
|GRU|**0.808±0.005**|**0.628±0.001**|-0.457±0.013|

| Architecture | stateless pendulum hard (s) | noisy stateless pendulum hard (s) | repeat previous hard (s) |
|---|---|---|---|
|S5|**297.56±0.45**|**321.18±0.15**|**308.50±0.31**|
|Transformer|2018.25±0.22|706.49±0.10|863.64±0.20|
|GRU|970.82±11.60|1925.91±32.59|1476.80±16.37|


## Misc

We’ve also made numerous other changes and additions to our manuscript from the reviewer’s suggestions, including: 

- Clarifying “the data collection and optimization implementations common to on-policy policy gradient algorithms with parallel actors” ($\color{red} R1$)

- The training and evaluation procedure of the DMControl multi-task setup ($\color{red} R1$, $\color{green} R2$)

- What we mean by “in-context learning” ($\color{green} R2$)

- Various fixes to spelling, notation, and resizing of figures ($\color{green} R2$, $\color{magenta} R4$)

- An explicit and thorough limitations section ($\color{red} R1$,$\color{green} R2$,$\color{blue} R3$,$\color{magenta} R4$)

- Clarifying why standard naive resetting does not work ($\color{magenta} R4$)

- Parameter counts for different architectures ($\color{green} R2$)

We hope these modifications adequately address the concerns raised by the reviewers, and we are confident they strengthen our manuscript's overall quality. To reiterate: We show that SSM’s provide *clear advantages in partially-observed RL* and this can accelerate future work in RL.

Meta-review: This paper develops a modified version of S5 architecture as a drop-in replacement of RNNs and Transformers in POMDP and memory-intensive RL tasks, which improves the long-range sequence modeling and inference cost. All the reviewers appreciate the significance of the problem and the contribution/novelty of the paper. And the questions raised by the reviewers have been well addressed during the rebuttal period.
Decision: Accept (poster)