Question:
You are a highly experienced, conscientious, and fair academic reviewer, please help me review this paper. The review should be organized into nine sections:
1. Summary: A summary of the paper in 100-150 words.
2. Strengths/Weaknesses/Questions: The Strengths/Weaknesses/Questions of paper, which should be listed in bullet points, with each point supported by specific examples from the article where possible.
3. Soundness/Contribution/Presentation: Rate the paper's Soundness/Contribution/Presentation, and match this score to the corresponding description from the list below and provide the result. The possible scores and their descriptions are:
    1 poor
    2 fair
    3 good
    4 excellent
4. Rating: Give this paper an appropriate rating, match this rating to the corresponding description from the list below and provide the result. The possible Ratings and their descriptions are:
    1 strong reject
    2 reject, significant issues present
    3 reject, not good enough
    4 possibly reject, but has redeeming facets
    5 marginally below the acceptance threshold
    6 marginally above the acceptance threshold
    7 accept, but needs minor improvements
    8 accept, good paper
    9 strong accept, excellent work
    10 strong accept, should be highlighted at the conference
5. Paper Decision: It must include the Decision itself(Accept or Reject) and the reasons for this decision, based on the criteria of originality, methodological soundness, significance of results, and clarity and logic of presentation.

Here is the template for a review format, you must follow this format to output your review result:
**Summary:**
Summary content

**Strengths:**
- Strength 1
- Strength 2
- ...

**Weaknesses:**
- Weakness 1
- Weakness 2
- ...

**Questions:**
- Question 1
- Question 2
- ...

**Soundness:**
Soundness result

**Presentation:**
Presentation result

**Contribution:**
Contribution result

**Rating:**
Rating result

**Paper Decision:**
- Decision: Accept/Reject
- Reasons: reasons content

Please ensure your feedback is objective and constructive. The paper is as follows:

# Structured State Space Models for In-Context Reinforcement Learning

Chris Lu

FLAIR, University of Oxford

Work done during an internship at DeepMind. Contact: christopher.lu@exeter.ox.ac.uk

Yannick Schroecker

DeepMind

Albert Gu

DeepMind

Emilio Parisotto

DeepMind

Jakob Foerster

FLAIR, University of Oxford

Satinder Singh

DeepMind

Feryal Behbahani

DeepMind

###### 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](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 \(16k\).

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) becausethe 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 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.

A recent 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 \(30x\). 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

### 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

\begin{table}
\begin{tabular}{l|l l l l l}  & Inference & Training & Parallel & Variable Lengths & bsuite Score \\ \hline RNNs & **O(1)** & **O(\(L\))** & No & **Yes** & No \\ Transformers & O(\(L^{2}\)) & O(\(L^{2}\)) & **Yes** & **Yes** & **Yes** \\ S5 with \(\bullet\) & **O(1)** & **O(\(L\))** & **Yes** & No & N/A \\ S5 with \(\oplus\) & **O(1)** & **O(\(L\))** & **Yes** & **Yes** & **Yes** \\ \end{tabular}
\end{table}
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.

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

\[\dot{x}(t) =\mathbf{A}x(t)+\mathbf{B}u(t) \tag{1}\] \[y(t) =\mathbf{C}x(t)+\mathbf{D}u(t)\]

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

\[x_{n} =\mathbf{\bar{A}}x_{n-1}+\mathbf{\bar{B}}u_{n} \tag{2}\] \[y_{n} =\mathbf{\bar{C}}x_{n}+\mathbf{\bar{D}}u_{n}\]

where \(\mathbf{\bar{A}},\mathbf{\bar{B}},\mathbf{\bar{C}}\), and \(\mathbf{\bar{D}}\) can be calculated as functions of \(\mathbf{A},\mathbf{B},\mathbf{C},\mathbf{D}\), and \(\Delta\).

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 \(\bullet\) to be associative, it must satisfy \((x\bullet y)\bullet z=x\bullet(y\bullet z)\).

Given an associative binary operator \(\bullet\) and a sequence of length \(N\), parallel scan returns:

\[[e_{1},e_{1}\bullet e_{2},\cdots,e_{1}\bullet e_{2}\bullet\cdots\bullet e_{N}] \tag{3}\]

For example, when \(\bullet\) 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:

\[e_{k}=(e_{k,a},e_{k,b}):=(\mathbf{\bar{A}},\mathbf{\bar{B}}u_{k}) \tag{4}\]

Where \(\mathbf{\bar{A}}\), \(\mathbf{\bar{B}}\), and \(u_{k}\) are defined in Equation 2. S5's parallel operator is then defined as:

\[a_{i}\bullet a_{j}=(a_{j,a}\odot a_{i,a},a_{j,a}\otimes a_{i,b}+a_{j,b}) \tag{5}\]

where \(\odot\) is matrix-matrix multiplication and \(\otimes\) is matrix-vector multiplication. The parallel scan then generates the recurrence in the hidden state \(x_{n}\) defined in Equation 2.

\[e_{1} =(\mathbf{\bar{A}},\mathbf{\bar{B}}u_{1}) =(\mathbf{\bar{A}},x_{1}) \tag{6}\] \[e_{1}\bullet e_{2} =(\mathbf{\bar{A}}^{2},\mathbf{\bar{A}}x_{1}+\mathbf{\bar{B}}u_{ 2}) =(\mathbf{\bar{A}}^{2},x_{2})\] (7) \[e_{1}\bullet e_{2}\bullet e_{3} =(\mathbf{\bar{A}}^{2},\mathbf{\bar{A}}x_{2}+\mathbf{\bar{B}}u_{ 3}) =(\mathbf{\bar{A}}^{3},x_{3}) \tag{8}\]

Note that the model assumes a hidden state initialised to \(x_{0}=0\) by initialising the scan with \(e_{0}=(\mathbf{I},0)\).

### Reinforcement Learning

A Markov Decision Process (MDP) [39] is defined as a tuple \(\langle\mathcal{S},\mathcal{A},R,P,\gamma\rangle\), which defines the environment. Here, \(\mathcal{S}\) is the set of states, \(\mathcal{A}\) the set of actions, \(R\) the reward function that maps from a given state and action to a real value \(\mathbb{R}\), \(P\) defines the distribution of next-state transitions given a state and action, and \(\gamma\) defines the discount factor. The agent's objective is to find a policy \(\pi\) (a function which maps from a given state to a distribution over actions) which maximises the expected discounted sum of returns.

\[\mathbb{E}[\mathbb{R}^{\gamma}|\pi]=\mathbb{E}_{s_{0}\sim d,a_{0: \infty}\sim\pi,s_{1:\infty}\sim P}\Big{[}\sum_{t=0}^{\infty}\gamma^{t}R(s_{t},a_{ t})\Big{]}\]

In a Partially-Observed Markov Decision Process (POMDP), the agent receives an observation \(o_{t}\sim O(s_{t})\) instead of directly observing the state. Because of this, the optimal policy \(\pi^{*}\) 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.

### 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 \(\oplus\) that operates on elements \(e_{k}\) defined:

\[e_{k}=(e_{k,a},e_{k,b},e_{k,c}):=(\mathbf{\bar{A}},\mathbf{\bar{B}}u_{k},d_{k}) \tag{9}\]

where \(d_{k}\in\{0,1\}\) is the binary "done" signal for the given transition from the environment.

We define \(\oplus\) to be:

\[a_{i}\oplus a_{j}:=\begin{cases}(a_{j,a}\odot a_{i,a},a_{j,a}\otimes a_{i,b}+a_ {j,b},a_{i,c})&\text{if }a_{j,c}=0\\ (a_{j,a},a_{j,b},a_{j,c})&\text{if }a_{j,c}=1\end{cases}\]

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:

\[e_{0}\oplus\cdots\oplus e_{n-1} =(\mathbf{\bar{A}}^{n-1},\mathbf{\bar{A}}x_{n-2}+\mathbf{\bar{B}} u_{n-1},0)\] \[=(\mathbf{\bar{A}}^{n-1},x_{n-1},0)\] \[e_{0}\oplus\cdots\oplus e_{n} =(\mathbf{\bar{A}},\mathbf{\bar{B}}u_{n},1)\] \[=(\mathbf{\bar{A}},x_{n},1)\] \[e_{0}\oplus\cdots\oplus e_{n+1} =(\mathbf{\bar{A}}^{2},\mathbf{\bar{A}}x_{n}+\mathbf{\bar{B}}u_{n+ 1},1)\] \[=(\mathbf{\bar{A}}^{2},x_{n+1},1)\]

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

### 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-learning 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.

```
1:Initialization: \(\mathbf{\bar{A}}\), \(\mathbf{\bar{B}}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{A}}^{n}\), \(\mathbf{\bar{B}}^{n}\), \(\mathbf{\bar{\(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.

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

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.

Figure 2: (a) Results across implemented environments in POPGym’s suite. Scores are normalised by the max-MMER per-environment. The shaded region represents the standard error of the mean across eight seeds. (b) The runtime for a single seed averaged across the environments for each architecture. Note that our implementation is end-to-end compiled to run entirely on a single NVIDIA A40.

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.

### 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 \(\oplus\) 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].

### 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 \(1k\) achieved the best performance in this setting. We provide a more detailed description of our

Figure 4: The mean of the return across all of the training 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 Algorithm 1.

Figure 3: (a) Performance and (b) runtime on randomly-projected StatelessCartPole across 4 seeds. The shaded region represents the standard error. (c) Shows performance at 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 to the training horizon while S5 models continue to perform well. The error bars represent the standard error across 4 seeds. Runs were performed on a single NVIDIA A100.

architecture and hyperparameters in the supplementary material.

**In-Distribution Training Results** We meta-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.

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

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 1

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

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 \(\Delta\), the discrete time between timesteps, is fixed for the original S4 model, S5 can in theory use a different \(\Delta\) for each timestep.

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 as described in Algorithm 1

**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

* [1] 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 Loks-Thompson, Hannah Openshaw, Jack Parker-Holder, Shreya Pathak, Nicolas Perez-Nieves, Nemanja Rakicevic, Tim Rocktaschel, 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 Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL [http://github.com/google/jax](http://github.com/google/jax).
* [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 Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. _arXiv preprint arXiv:1406.1078_, 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 [https://openreview.net/forum?id=kqHkCVS7wbj](https://openreview.net/forum?id=kqHkCVS7wbj).
* [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. Rl2: Fast reinforcement learning via slow reinforcement learning. _arXiv preprint arXiv:1611.02779_, 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 Re. It's raw! audio generation with state-space models. _arXiv preprint arXiv:2202.09729_, 2022.
* [11] Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Re. Hippo: Recurrent memory with optimal polynomial projections. _Advances in neural information processing systems_, 33:1474-1487, 2020.

* Gu et al. [2021] Albert Gu, Karan Goel, and Christopher Re. Efficiently modeling long sequences with structured state spaces. _arXiv preprint arXiv:2111.00396_, 2021.
* Gu et al. [2021] Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Re. Combining recurrent, convolutional, and continuous-time models with linear state space layers. _Advances in neural information processing systems_, 34:572-585, 2021.
* Hamilton [1994] James D Hamilton. State-space models. _Handbook of econometrics_, 4:3039-3080, 1994.
* Hessel et al. [2021] 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.
* Hochreiter and Schmidhuber [1997] Sepp Hochreiter and Jurgen Schmidhuber. Long short-term memory. _Neural computation_, 9(8):1735-1780, 1997.
* Houthooft et al. [2018] 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.
* Huang et al. [2022] Shengyi Huang, Rousslan Fernand Julien Dossa, Chang Ye, Jeff Braga, Dipam Chakraborty, Kinal Mehta, and Joao G.M. Araujo. 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](http://jmlr.org/papers/v23/21-1342.html).
* Kapturowski et al. [2018] 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.
* Kirsch et al. [2022] 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.
* Kirsch et al. [2022] Louis Kirsch, James Harrison, Jascha Sohl-Dickstein, and Luke Metz. General-purpose in-context learning by meta-learning transformers. _arXiv preprint arXiv:2212.04458_, 2022.
* Lange [2022] Robert Tjarko Lange. gymax: A JAX-based reinforcement learning environment library, 2022. URL [http://github.com/RobertTLange/gymax](http://github.com/RobertTLange/gymax).
* Laskin et al. [2022] 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.
* Liang et al. [2018] 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.
* Lu et al. [2022] 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.
* Metz et al. [2021] Luke Metz, C Daniel Freeman, Samuel S Schoenholz, and Tal Kachman. Gradients are not all you need. _arXiv preprint arXiv:2111.05803_, 2021.
* Morad et al. [2023] 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. URL [https://openreview.net/forum?id=chDrutUTsoK](https://openreview.net/forum?id=chDrutUTsoK).
* Nguyen et al. [2022] Eric Nguyen, Karan Goel, Albert Gu, Gordon Downs, Preey Shah, Tri Dao, Stephen Baccus, and Christopher Re. S4nd: Modeling images and videos as multidimensional signals with state spaces. In _Advances in Neural Information Processing Systems_, 2022.

* Ni et al. [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.
* Oh et al. [2020] 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.
* O'Keefe and Dostrovsky [1971] 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.
* Osband et al. [2019] 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.
* Parisotto and Salakhutdinov [2021] Emilio Parisotto and Ruslan Salakhutdinov. Efficient transformers in reinforcement learning using actor-learner distillation. _arXiv preprint arXiv:2104.01655_, 2021.
* Parisotto et al. [2020] 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.
* Raffin et al. [2021] 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](http://jmlr.org/papers/v22/20-1364.html).
* Reed et al. [2022] 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.
* Schulman et al. [2017] John Schulman, F. Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. _ArXiv_, abs/1707.06347, 2017.
* Smith et al. [2022] Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. Simplified state space layers for sequence modeling. _arXiv preprint arXiv:2208.04933_, 2022.
* Sutton and Barto [2018] Richard S Sutton and Andrew G Barto. _Reinforcement learning: An introduction_. 2018.
* Tassa et al. [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.
* Tay et al. [2020] 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.
* Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. _Advances in neural information processing systems_, 30, 2017.
* Wang et al. [2016] 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.
* Williams and Zipser [1995] Ronald J Williams and David Zipser. Gradient-based learning algorithms for recurrent. _Backpropagation: Theory, architectures, and applications_, 433:17, 1995.

Proof of Associativity of Binary Operator

Recall that \(\oplus\) is defined as:

\[a_{i}\oplus a_{j}:=\begin{cases}(a_{j,a}\odot a_{i,a},a_{j,a}\otimes a_{i,b}+a_{j, b},a_{i,c})&\text{if }a_{j,c}=0\\ (a_{j,a},a_{j,b},a_{j,c})&\text{if }a_{j,c}=1\end{cases}\]

This is equivalent to the following:

\[a_{i}\oplus a_{j}:=\begin{cases}((a_{i}\bullet a_{j})_{a},(a_{i}\bullet a_{j}) _{b},a_{i,c})&\text{if }a_{j,c}=0\\ a_{j}&\text{if }a_{j,c}=1\end{cases}\]

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

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

**Case 1: \(z_{c}=1\)**

\[(x\oplus y)\oplus z =z \tag{10}\] \[=y\oplus z\] (11) \[=x\oplus(y\oplus z) \tag{12}\]

**Case 2: \(z_{c}=0\)** and \(y_{c}=1\)

\[(x\oplus y)\oplus z =y\oplus z \tag{13}\] \[\text{Note that }(y\oplus z)_{c} =1\text{ thus,}\] (14) \[=x\oplus(y\oplus z) \tag{15}\]

**Case 3: \(z_{c}=0\)** and \(y_{c}=0\)

\[(x\oplus y)\oplus z =((x\bullet y)_{a},(x\bullet y)_{b},x_{c})\oplus z \tag{16}\] \[=(((x\bullet y)\bullet z)_{a},((x\bullet y)\bullet z)_{b},x_{c})\] (17) \[=((x\bullet(y\bullet z))_{a},(x\bullet(y\bullet z))_{b},x_{c})\] (18) \[=x\oplus((y\bullet z)_{a},(y\bullet z)_{b},y_{c})\] (19) \[=x\oplus(y\oplus z) \tag{20}\]

## Appendix B Hyperparameters

\begin{table}
\begin{tabular}{l|l} Parameter & Value \\ \hline 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, 512] \\ Number of Environments & 1024 \\ Discount \(\gamma\) & 0.995 \\ Rollout Length & 1000 \\ Offline Data Fraction & 0.0 \\ Total Frames & 2e9 \\ LSTM Hidden Size & 512 \\ Projected Observation Size & 12 \\ Projected Action Size & 2 \\ S5 Layers & 10 \\ S5 Hidden Size & 256 \\ S5 Discretization & ZOH \\ S5 \(\Delta\) min & 0.001 \\ S5 \(\Delta\) max & 0.1 \\ \end{tabular}
\end{table}
Table 4: Hyperparameters for training Muesli on Multi-Environment Meta-RL. These experiments were run using 64 TPUv3’s.

## Appendix 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 \(\lambda\), and model architectures from Morad et al. [27]. However, we used the standard clipping coefficient \(\epsilon\) 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.

\begin{table}
\begin{tabular}{r|c c c c|c c}  & Stateless & Noisy Stateless & Stateless & Noisy Stateless & Repeat Previous \\  & CartPole Hard & CartPole Hard & Pendulum Hard & Pendulum Hard & Hard \\ \hline S5 (ours) & \(\mathbf{1.0\pm 0.0}\) & \(\mathbf{0.28\pm 0.0}\) & \(\mathbf{0.79\pm 0.01}\) & \(0.55\pm 0.01\) & \(\mathbf{0.91\pm 0.01}\) \\ GRU (ours) & \(\mathbf{1.0\pm 0.0}\) & \(0.27\pm 0.0\) & \(0.75\pm 0.0\) & \(\mathbf{0.61\pm 0.01}\) & \(-0.46\pm 0.01\) \\ MLP (ours) & \(0.26\pm 0.0\) & \(0.22\pm 0.0\) & \(0.41\pm 0.02\) & \(0.34\pm 0.01\) & \(-0.48\pm 0.00\) \\ \hline GRU & \(\mathbf{1.000\pm 0.000}\) & \(0.390\pm 0.007\) & \(\mathbf{0.828\pm 0.001}\) & \(\mathbf{0.657\pm 0.002}\) & \(-0.428\pm 0.002\) \\ MLP & \(0.265\pm 0.002\) & \(0.229\pm 0.002\) & \(0.477\pm 0.030\) & \(0.351\pm 0.012\) & \(-0.486\pm 0.002\) \\ IndRNN & \(\mathbf{1.000\pm 0.000}\) & \(\mathbf{0.404\pm 0.005}\) & \(0.804\pm 0.023\) & \(0.521\pm 0.109\) & \(-0.384\pm 0.013\) \\ LMU & \(0.987\pm 0.007\) & \(0.352\pm 0.019\) & \(0.806\pm 0.006\) & \(0.563\pm 0.014\) & \(\mathbf{0.191\pm 0.041}\) \\ S4D & \(0.127\pm 0.026\) & \(0.207\pm 0.007\) & \(0.303\pm 0.014\) & \(0.289\pm 0.011\) & \(-0.491\pm 0.001\) \\ FART & \(\mathbf{0.996\pm 0.000}\) & \(0.366\pm 0.002\) & \(0.698\pm 0.077\) & \(0.553\pm 0.007\) & \(-0.485\pm 0.001\) \\ \end{tabular}
\end{table}
Table 5: Results in POPGym’s suite. The reported number is 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 all 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-performing one from each environment. They report the mean and standard deviation across three trials.

Response:
## Review

### Summary

The paper presents a modified version of the S5 architecture, which allows for resetting the recurrent state, enabling it to serve as a drop-in replacement for RNNs and Transformers in reinforcement learning (RL) tasks. The authors demonstrate that their approach outperforms existing models in terms of performance and computational efficiency across various partially observed RL environments. Specifically, the resettable S5 model showcases superior long-range sequence modeling capabilities and improved inference speed. Additionally, the paper introduces a meta-learning task with random projections of state and action spaces, further validating the efficacy of S5 in in-context adaptation scenarios.

### Strengths

- The paper is the first to leverage S4-like models in RL, showcasing an important contribution.
- A novel associative operator is proposed to handle episodic resets in the S5 architecture.
- S5 outperforms RNNs and Transformers in terms of task performance and computational efficiency.
- The introduction of a meta-RL environment with random projections serves as a useful benchmark for future research.
- The experimental setup covers both POMDPs and meta-learning, indicating a broad applicability of the method.

### Weaknesses

- The evaluation does not sufficiently explore the potential of S5 for longer sequence lengths, which is a key benefit of state-space models.
- The experimental design lacks variety and depth, particularly in the choice of tasks used for evaluation.
- The main contribution regarding the reset mechanism appears to be a trivial modification rather than a novel advancement.
- There is insufficient discussion of the limitations of the proposed architecture and its comparison to existing models.
- The narrative around in-context RL is confusing, and the evaluation protocol lacks clarity.

### Questions

- What is the rationale behind the improved out-of-distribution (OOD) generalization of the S5 model compared to baselines?
- Can the authors clarify the evaluation procedure for the meta-learning tasks, particularly regarding training and held-out tasks?
- How do models perform beyond a sequence length of 512, especially in relation to S5's capabilities?
- Why is the performance gain of S5 primarily observed on the 'Repeat Previous Hard' task, and what implications does this have for the benchmark's validity?
- Is there a reason for the absence of transformer-based baselines in the evaluation?

### Soundness

**Score:** 2

**Description:** 2 = fair; while the methodology is generally sound, the contributions and evaluations are somewhat limited and lack depth.

### Presentation

**Score:** 2

**Description:** 2 = fair; the writing has clarity issues, particularly regarding the explanation of key concepts and evaluation procedures.

### Contribution

**Score:** 3

**Description:** 3 = good; the paper makes a notable contribution to the field by introducing a new application of the S5 architecture in RL, despite some concerns about the novelty of the modification.

### Rating

**Score:** 6

**Description:** 6 = accept with minor revisions; while the paper shows promise and addresses significant challenges in RL, it requires some clarifications and improvements in experimental design.

### Paper Decision

**Decision:** Accept (poster)

**Reasons:** The paper presents an innovative approach to leveraging the S5 architecture for reinforcement learning, demonstrating solid technical foundations and promising results. However, it requires clarity in its presentation and a more thorough exploration of its limitations and evaluations. The strengths outweigh the weaknesses, leading to a recommendation for acceptance with minor revisions.