Title: Structured State Space Models for In-Context Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: Reinforcement Learning

Meta-Learning

State Space Models
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

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 .

1 Introduction

Structured state space sequence (S4) models [12] and their variants such as S5 [38] 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 benchmark [41]. Notably, S4 was the first architecture to achieve a non-trivial result on the difficult Path-X task, which requires the ability to handle extremely long-range dependencies of lengths 16 k .

Furthermore, S4 models display a number of desirable properties that are not directly tested by raw performance benchmarks. Unlike transformers, for which the per step runtime usually scales quadratically 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-time is normally not included when evaluating on sequence modelling benchmarks, it has a large impact on the scalability and wallclock-time for reinforcement learning (RL) because

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 sequence-modelling tasks [42], RNNs are still widely-used for memory-based RL tasks [29].

The ability to efficiently model contexts that are orders of magnitude larger may enable new possibilities in RL. This is particularly applicable in meta-reinforcement learning (Meta-RL), in which the agent is trained to adapt across multiple environment episodes. One approach to Meta-RL, RL 2 [8, 43], uses sequence models to directly learn across these episodes, which can often result in trajectories that are thousands of steps long. Most instances of RL 2 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 in on-policy RL with multiple actors involves collecting fixedlength 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.

Arecent modification to S4, called Simplified Structured State Space Sequence Models (S5), replaces this convolution with a 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 a trajectory during the training phase, which in turn allows practitioners to 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 [32], 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-proposed Partially Observable Process Gym (POPGym) suite [27] in pure JAX, resulting in end-to-end evaluation speedups of over 30 x . When evaluating our architecture on this suite, we show that S5 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-observed Meta-RL task with episode lengths of up to 6400 . Finally, we evaluate S5 on a challenging Meta-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 on 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 [14]. At each timestep t , these models take an input signal u ( t ) . This is used to update a latent state x ( t ) which in turn computes the signal y ( t ) . Some of the more widely-used

canonical SSMs are continuous-time linear SSMs, which are defined by the following equations:

where A , B , C , and D are matrices of appropriate sizes. To model sequences with a fixed step size ∆ , one can discretise the SSM using various techniques, such as the zero-order hold method, to obtain a simple linear recurrence:

where ¯ A , ¯ B , ¯ C , and ¯ D can be calculated as functions of A , B , C , D , and ∆ .

S4 [12] proposed the use 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 a special matrix initialisation to better preserve sequence history called HiPPO [11].

One of the primary strengths of the S4 model is that it can be converted to both a recurrent model, which allows for fast and memory-efficient inference-time computation, and a convolutional model, which allows for efficient training that is parallelisable across timesteps [13].

More recently, Smith et al. [38] proposed multiple simplifications to S4, called S5. One of its contributions is 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 an operation · to be associative, it must satisfy ( x · y ) · z = x · ( y · z ) .

Given an associative binary operator · and a sequence of length N , parallel scan returns:

For example, when · is addition, the parallel scan calculates the prefix-sum, which returns the running total of an input sequence. Parallel scans can be computed in O (log( N )) time when given a sequence of length N , given N parallel processors.

S5's parallel scan is applied to initial elements e 0: N defined as:

Where ¯ A , ¯ B , and u k are defined in Equation 2. S5's parallel operator is then defined as:

where ⊙ is matrix-matrix multiplication and ⊗ is matrix-vector multiplication. The parallel scan then generates the recurrence in the hidden state x n defined in Equation 2.

Note that the model assumes a hidden state initialised to x 0 = 0 by initialising the scan with e 0 = ( I , 0) .

2.2 Reinforcement Learning

A Markov Decision Process (MDP) [39] is defined as a tuple ⟨S , A , R, P, γ ⟩ , 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 a state and action, and γ defines the discount factor. The agent's objective is to find a policy π (a function which maps from a given state to a distribution over actions) which maximises the expected discounted sum of returns.

Algorithm 1 Pseudocode for the Multi-Environment Meta-Learning environment step.

Require: Distribution of environments E ∼ ρ E , a fixed output observation dimension size O , and a fixed action dimension size A . Agent action a and Environment termination d

1: function StepEnvironment( a , d )

2: if the environment terminated ( d ) then

3: Sample random environment E ∼ ρ E

4: Initialise random observation projection matrix M o ∈ R E o × O where E o is E 's observation size

5: Initialise random action projection matrix M a ∈ R A × E a where E a is E 's action size

6: Reset E to receive an initial observation o

7: Apply the random observation projection matrix to the observation o ′ = M o o

8: Append r = 0 and d = 0 to o ′ to get o ′′

9: Return o ′′

10: else

11: Apply the projection matrix a ′ = M a a

12: Step E using a ′ to receive the next observation o , reward r , and done signal d .

13: Apply the projection matrix o ′ = M o o

14: Append r and d to o ′ to get o ′′

15: Return o ′′ , r , and d

16: end if

17: end function

In a Partially-Observed Markov Decision Process (POMDP), the agent receives an observation o t ∼ O ( s t ) instead of directly observing the state. Because of this, the optimal policy π ∗ does not depend just on the current observation o t , but (in generality) also on all previous observations o 0: t and actions a 0: 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 tests 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 environment 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 [24].

To handle trajectory rollouts that begin in the middle of an episode, sequence models must be 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 to 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 are 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 be 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 to efficiently train over sequences of varying lengths and hidden state initialisations. We create a new associative operator ⊕ that operates on elements e k defined:

where d k ∈ { 0 , 1 } is the binary 'done' signal for the given transition from the environment. We define ⊕ to be:

We prove that this operator is associative in Appendix A. We now show that the operator computes the desired value. Assuming e n,c = 1 corresponds to a 'done' transition while all other timesteps before it ( e 0: n -1 ,c = 0 ) and after it ( e n +1: L,c = 0 ) do not, we see:

Note that even if there are multiple 'resets' within the rollout, the desired value will still be computed.

3.2 Multi-Environment Meta-Learning with Random Projections

Most prior work in Meta-RL has only demonstrated the ability to adapt to a 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-learnning 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 a fixed size, allowing us to 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 1.

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 a cue on the first timestep, which corresponds to the action the agent should take some number of steps in the future to receive a reward. We run our experiments using bsuite's actor-critic baseline while swapping out the LSTM for Transformer self-attention blocks or S5 blocks [42] and using Gymnax for faster environment rollouts [22].

We show the results in Figure 1. In general, S5 displays a better asymptotic runtime than Transformers while far outperforming LSTMs in both performance and speed. Note that while S5 is theoretically

(a)

(b)

(c)

50

LSTM

Self-Attention

S5

Bsuite Memory Usage for 16 Vectorized Runs

50000

LSTM

S5

Bsuite Runtime

1.0

BSuite Mean Return Across 5 Seeds

40

Self-Attention

0.8

Memory Usage (GiB)

30

Runtime (Seconds)

40000

0.6

30000

Return

0.4

20

20000-

0.2

10

10000

0.0

0.2

0

0

LSTM

S5

Self-Attention

0

1000

2000

3000

4000

0

500

1000

1500

2000

4

8

16

32

64

128

256

512 1024 2048

Memory Length

Memory Length

Memory Length

(a)

(b)

POPGym Normalised MMER

POPGym Normalised MMER

1.0

1.0

S5 with ①

S5 with·

FC

0.8

0.8

GRU

MMER

MMER

S5 with

Normalised

0.6

S5 with·

Normalised

0.6

FC

GRU

0.4

0.4

0.2

0.2

0

50

100

150

200

0

2000

4000

6000

8000

10000

12000

Updates

Seconds

O (log( N )) in the backwards pass during training time (given enough processors), it is still bottlenecked 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 (POPGym) [27] suite, a set of simple environments designed to benchmark memory in deep RL. To increase experiment throughput on a limited compute budget, we carefully re-implemented environments from the POPGym suite entirely in JAX [3] by leveraging existing implementations of CartPole and Pendulum in Gymnax [22]. PyTorch does not support the associative scan operation, so we could directly use the S5 architecture in POPGym's RLLib benchmark.

Morad et al. [27] evaluated several architectures and found the Gated Recurrent Unit (GRU) [5] to be the most performant. Thus, we compare our results to the GRU architecture proposed in the original POPGym benchmark. Note that POPGym'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] - a feature that does not appear in most PPO implementations [9]. We instead use a recurrent PPO implementation that is more closely aligned with StableBaselines3 [35] and CleanRL's [18] recurrent PPO implementations. We include more discussion and the hyperparameters in Appendix B.

(a)

(b)

(c)

Meta Stateless CartPole Meta-Training MMER

Meta Stateless CartPole Evaluation Performance

S5 with

Meta Stateless CartPole Meta-Training Runtime

S5 with *

GRU

7000

0.45

6

FC

6000

0.40

5000

0.35

MMER

Return

0.30

寸

Seconds

4000

0.25

3000

E

2000

0.20

0.15

S5 with •

Meta-Training Horizon

S5with

1000

0.10

GRU

FC

0

100

200

300

Updates

400

500

600

700

0

GRU

S5 with

S5 with •

FC

0

5

10

Trial Number

15

20

25

30

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 Randomly-Projected CartPole In-Context

We first demonstrate the long-horizon in-context learning capabilities of the modified S5 architecture by meta-learning across random projections of POPGym'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 on randomly-sampled linear projections of StatelessCartPole's observation and action spaces. We find that S5 with ⊕ outperforms GRU's while running twice as quickly. Furthermore, we show that performance improves across trials, demonstrating in-context 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. [1].

4.4 Multi-Environment Meta-Reinforcement Learning

We run our S5 architecture, an LSTM baseline, and a memory-less fully-connected network in the environment described in Section 3.2. For these experiments, we use Muesli [15] as our policy optimisation algorithm. We randomly project the environment observations to a fixed size of 12 and randomly project from an action space of size 2 to the corresponding environment's action space. We selected all of the DMControl environments and tasks that had observation and action spaces of size equal to or less than those values and split them into train and test set environments.

We use the S5 architecture described in Smith et al. [38], which consists of multiple stacked layers of S5 blocks. For both S5 and LSTM architectures, we found that setting the trajectory length equal to the maximum episode length 1 k achieved the best performance in this setting. We provide a more detailed description of our

0

25

50

75

100

125

150

175

200

Frames (10 Million)

200

400

600

800

Return

Mean Across All Training Envs

S5

FC

LSTM

0

50

100

150

200

0

500

1000

Returns

ball_in_cup_catch

S5

FC

LSTM

0

50

100

150

200

200

400

600

cartpole_swingup

0

50

100

150

200

0

250

500

750

finger_spin

0

50

100

150

200

Frames (10 Million)

250

500

750

Returns

finger_turn_hard

0

50

100

150

200

Frames (10 Million)

0

250

500

750

point_mass_easy

0

50

100

150

200

Frames (10 Million)

250

500

750

reacher_easy

architecture and hyperparameters in the supplementary material.

In-Distribution Training Results Wemeta-train the model across the six DMControl environments in Figure 5 and show the mean performance across them in Figure 4. 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.

0

25

50

75

100

125

150

175

200

Frames (10 Million)

100

200

300

400

500

Return

Mean Across All Test Envs

S5

FC

LSTM

Out-of-Distribution Evaluation Results: After training the model on the DMControl tasks in Figure 5, we next evaluate the trained model on random linear projections of five held-out DMControl tasks, without any extra fine-tuning . We present the mean score across these tasks in Figure 6 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 a completely unseen task in the test set, pendulum swingup, which has an unseen observation space, action space, and reward dynamics.

5 Related Work

S4 models have previously been shown to work in a number of previous settings, including audio generation [10] and video modeling [28]. Concurrent 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 out-of-distribution 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 hand-crafted restrictions.

0

50

100

150

200

0

200

400

Returns

cartpole_balance

S5

FC

LSTM

0

50

100

150

200

0

250

500

750

finger_turn_easy

0

50

100

150

200

0

50

100

pendulum_swingup

0

50

100

150

200

Frames (10 Million)

0

200

400

600

Returns

point_mass_hard

0

50

100

150

200

Frames (10 Million)

0

200

400

600

reacher_hard

Some works perform long-horizon meta-reinforcement learning through the use of evolution strategies [17, 25]. 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.

Other works have investigated different sequence model architectures for memory-based reinforcement learning. Ni et 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-based 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 to also demonstrate out-of-distribution generalisation in meta-reinforcement learning by leveraging a large task space [1].

6 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 advantages over RNNs and Transformers. After identifying challenges in integrating S5 models into existing recurrent reinforcement learning implementations, we made a simple modification to 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 demonstrated that S5 is asymptotically faster than Transformers in the sequence length. Furthermore, we also showed that S5 models run nearly twice as quickly as LSTMs with the same number of parameters while outperforming them. We further evaluated our S5 architecture on environments in the POPGym suite [27], where we 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.

Finally, we proposed a new meta-learning setting in which we 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 out-of-distribution performance to unseen tasks through in-context adaptation.

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 to scale to extremely long contexts. Another direction would be to investigate S5 models for continuous-time RL settings [7]. While ∆ , the discrete time between timesteps, is fixed for the original S4 model, S5 can in theory use a different ∆ for each timestep.

Limitations: There are several notable limitations of this architecture and analysis. Firstly, implementing the associative scan operator is currently not possible using PyTorch, limiting us to using the JAX [3] framework. Furthermore, on tasks where short rollouts are sufficient to achieve good performance, S5 offers limited speedups, as rolling out across time is no longer a bottleneck. Finally, it was not possible to perform a 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 at DeepMind for their helpful discussions throughout the course of the project.

References

Adaptive Agent Team, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmieg, Michael Chang, Natalie Clay, Adrian Collister, Vibhavari Dasagi, Lucy Gonzalez, Karol Gregor, Edward Hughes, Sheleem Kashem, Maria LoksThompson, 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.

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.

James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http://github.com/google/jax .

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.

Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoderdecoder for statistical machine translation. arXiv preprint arXiv:1406.1078 , 2014.

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 https://openreview.net/forum?id=kqHkCVS7wbj .

Kenji Doya. Reinforcement learning in continuous time and space. Neural computation , 12(1): 219-245, 2000.

Yan Duan, John Schulman, Xi Chen, Peter L Bartlett, Ilya Sutskever, and Pieter Abbeel. Rl2: Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:1611.02779 , 2016.

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.

Karan Goel, Albert Gu, Chris Donahue, and Christopher Ré. It's raw! audio generation with state-space models. arXiv preprint arXiv:2202.09729 , 2022.

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.

Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00396 , 2021.

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.

James D Hamilton. State-space models. Handbook of econometrics , 4:3039-3080, 1994.

Matteo Hessel, Ivo Danihelka, Fabio Viola, Arthur Guez, Simon Schmitt, Laurent Sifre, Theophane Weber, David Silver, and Hado Van Hasselt. Muesli: Combining improvements in policy optimization. In International Conference on Machine Learning , pages 4214-4226. PMLR, 2021.

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

Rein Houthooft, Yuhua Chen, Phillip Isola, Bradly Stadie, Filip Wolski, OpenAI Jonathan Ho, and Pieter Abbeel. Evolved policy gradients. Advances in Neural Information Processing Systems , 31, 2018.

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):1-18, 2022. URL http://jmlr.org/papers/v23/21-1342.html .

Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent experience replay in distributed reinforcement learning. In International conference on learning representations , 2018.

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.

Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose incontext learning by meta-learning transformers. arXiv preprint arXiv:2212.04458 , 2022.

Robert Tjarko Lange. gymnax: A JAX-based reinforcement learning environment library, 2022. URL http://github.com/RobertTLange/gymnax .

Michael Laskin, Luyu Wang, Junhyuk Oh, Emilio Parisotto, Stephen Spencer, Richie Steigerwald, DJ Strouse, Steven Hansen, Angelos Filos, Ethan Brooks, et al. In-context reinforcement learning with algorithm distillation. arXiv preprint arXiv:2210.14215 , 2022.

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.

Chris Lu, Jakub Grudzien Kuba, Alistair Letcher, Luke Metz, Christian Schroeder de Witt, and Jakob Foerster. Discovered policy optimisation. arXiv preprint arXiv:2210.05639 , 2022.

Luke Metz, C Daniel Freeman, Samuel S Schoenholz, and Tal Kachman. Gradients are not all you need. arXiv preprint arXiv:2111.05803 , 2021.

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. URL https://openreview.net/forum? id=chDrutUTs0K .

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.

Tianwei Ni, Benjamin Eysenbach, and Ruslan Salakhutdinov. Recurrent model-free rl can be a strong baseline for many pomdps. In International Conference on Machine Learning , pages 16691-16723. PMLR, 2022.

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:1060-1070, 2020.

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

Ian Osband, Yotam Doron, Matteo Hessel, John Aslanides, Eren Sezener, Andre Saraiva, Katrina McKinney, Tor Lattimore, Csaba Szepesvari, Satinder Singh, et al. Behaviour suite for reinforcement learning. arXiv preprint arXiv:1908.03568 , 2019.

Emilio Parisotto and Ruslan Salakhutdinov. Efficient transformers in reinforcement learning using actor-learner distillation. arXiv preprint arXiv:2104.01655 , 2021.

Emilio Parisotto, Francis Song, Jack Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant Jayakumar, 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.

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.org/papers/v22/ 20-1364.html .

Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06175 , 2022.

John Schulman, F. Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv , abs/1707.06347, 2017.

Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. arXiv preprint arXiv:2208.04933 , 2022.

Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction . 2018.

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.00690 , 2018.

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.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

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 arXiv:1611.05763 , 2016.

Ronald J Williams and David Zipser. Gradient-based learning algorithms for recurrent. Backpropagation: Theory, architectures, and applications , 433:17, 1995.

A Proof of Associativity of Binary Operator

Recall that ⊕ is defined as:

This is equivalent to the following:

where · is S5's binary operator defined in Equation 5. Note that · 's associativity was proved in Smith et al. [38]. Using this, we can prove the associativity of ⊕ .

Let x , y , and z refer to three elements. We will prove that for all possible values of x , y , and z , ⊕ retains associativity.

Case 1: z c = 1

Case 2: z c = 0 and y c = 1

Case 3: z c = 0 and y c = 0

B Hyperparameters

C POPGym Discussion

Morad et al. [27] used RLLib's [24] implementation of PPO, which differs significantly from standard implementations of PPO. It uses a 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 [18] 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 S5 architecture far outperforms the best performing architecture in Morad et al. [27] in the 'RepeatPreviousHard' environment.

We used the learning rate, number of environments, unroll length, timesteps, epochs, and minibatches, GAE λ , and model architectures from Morad et al. [27]. However, we used the standard clipping coefficient ϵ of 0 . 2 instead of 0 . 3 to account for the lack of a dynamic KL-divergence coefficient. Note that we also adjusted the S5 architecture to contain approximately the same number of 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 experimentation, we implement our algorithms and environments end-to-end in JAX [3, 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\n#### 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.\n\nQUESTIONS:\n- 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\n\nLIMITATIONS:\nThe authors do not have a limitations section. It is probably worth adding a short section on this.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\n- 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\nLIMITATIONS:\nThe 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nStructured 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.\n\nSTRENGTHS:\n- 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\nWEAKNESSES:\n**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\nQUESTIONS:\n- 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\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nState 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.\n\nWEAKNESSES:\nS5 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.\n\nQUESTIONS:\nPlease 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.\n\nLIMITATIONS:\nThere 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}