SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Chris Lu*
FLAIR, University of Oxford
Yannick Schroecker
Albert Gu
Emilio Parisotto
Jakob Foerster
Feryal Behbahani
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 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) because
*Work done during an internship at DeepMind. Contact: christopher.lu@exeter.ox.ac.uk
Inference	Training	Parallel	Variable Lengths	bsuite Score
RNNs	O(1)	O(L)	No	Yes	No
Transformers	O(L2)	O(L2)	Yes	Yes	Yes
S5 with •	O(1)	O(L)	Yes	No	N/A
S5 with ⊕	O(1)	O(L)	Yes	Yes	Yes
Table 1: The different properties of the different architectures. The asymptotic runtimes are in terms of the sequence length  $L$  assume a constant hidden size. The bsuite scores correspond to whether or not they achieve a perfect score in the median runs on the bsuite memory length environment.
the agent uses inference to collect data from the environment. Thus, transformers usually have poor runtime performance in reinforcement learning [33]. While transformers have become the default architecture for many supervised 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,  $\mathrm{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  $\mathrm{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.1 Structured State Space Sequence Models
State Space Models (SSMs) have been widely used to model various phenomena 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:
$$
\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} = \bar {\mathbf {A}} x _ {n - 1} + \bar {\mathbf {B}} u _ {n} \tag {2}
$$
$$
y _ {n} = \bar {\mathbf {C}} x _ {n} + \bar {\mathbf {D}} u _ {n}
$$
where  $\bar{\mathbf{A}},\bar{\mathbf{B}},\bar{\mathbf{C}}$  , and  $\bar{\mathbf{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:
$$
\left[ e _ {1}, e _ {1} \bullet e _ {2}, \dots , e _ {1} \bullet e _ {2} \bullet \dots \bullet e _ {N} \right] \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} = \left(e _ {k, a}, e _ {k, b}\right) := (\bar {\mathbf {A}}, \bar {\mathbf {B}} u _ {k}) \tag {4}
$$
Where  $\bar{\mathbf{A}}$ ,  $\bar{\mathbf{B}}$ , and  $u_{k}$  are defined in Equation 2. S5's parallel operator is then defined as:
$$
a _ {i} \cdot a _ {j} = \left(a _ {j, a} \odot a _ {i, a}, a _ {j, a} \otimes a _ {i, b} + a _ {j, b}\right) \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} = (\bar {\mathbf {A}}, \bar {\mathbf {B}} u _ {1}) \quad = (\bar {\mathbf {A}}, x _ {1}) \tag {6}
$$
$$
e _ {1} \bullet e _ {2} = \left(\bar {\mathbf {A}} ^ {2}, \bar {\mathbf {A}} x _ {1} + \bar {\mathbf {B}} u _ {2}\right) = \left(\bar {\mathbf {A}} ^ {2}, x _ {2}\right) \tag {7}
$$
$$
e _ {1} \bullet e _ {2} \bullet e _ {3} = \left(\bar {\mathbf {A}} ^ {2}, \bar {\mathbf {A}} x _ {2} + \bar {\mathbf {B}} u _ {3}\right) = \left(\bar {\mathbf {A}} ^ {3}, x _ {3}\right) \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)$ .
2.2 Reinforcement Learning
A Markov Decision Process (MDP) [39] is defined as a tuple  $\langle S, \mathcal{A}, R, P, \gamma \rangle$ , which defines the environment. Here,  $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.
Algorithm 1 Pseudocode for the Multi-Environment Meta-Learning environment step.
Require: Distribution of environments  $E\sim \rho_{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\sim \rho_{E}$    
4: Initialise random observation projection matrix  $M_o\in \mathbb{R}^{E^o\times O}$  where  $E^o$  is  $E$  's observation size   
5: Initialise random action projection matrix  $M_a\in \mathbb{R}^{A\times 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^{\prime} = M_{o}o$    
8: Append  $r = 0$  and  $d = 0$  to  $o^\prime$  to get  $o''$    
9: Return  $o''$    
10: else   
11: Apply the projection matrix  $a^\prime = M_ia$    
12: Step  $E$  using  $a^\prime$  to receive the next observation o, reward r, and done signal d.   
13: Apply the projection matrix  $o^{\prime} = M_{o}o$    
14: Append  $r$  and  $d$  to  $o^{\prime}$  to get  $o''$    
15: Return  $o''$  , r, and d   
16: end if   
17: end function
$$
\mathbb {E} [ \mathbf {R} ^ {\gamma} | \pi ] = \mathbb {E} _ {s _ {0} \sim d, a _ {0: \infty} \sim \pi , s _ {1: \infty} \sim P} \Bigl [ \sum_ {t = 0} ^ {\infty} \gamma^ {t} R (s _ {t}, a _ {t}) \Bigr ]
$$
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}$ .
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  $\oplus$  that operates on elements  $e_k$  defined:
$$
e _ {k} = \left(e _ {k, a}, e _ {k, b}, e _ {k, c}\right) := (\bar {\mathbf {A}}, \bar {\mathbf {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} := \left\{ \begin{array}{l l} (a _ {j, a} \odot a _ {i, a}, a _ {j, a} \otimes a _ {i, b} + a _ {j, b}, a _ {i, c}) & \text {i f} a _ {j, c} = 0 \\ (a _ {j, a}, a _ {j, b}, a _ {j, c}) & \text {i f} a _ {j, c} = 1 \end{array} \right.
$$
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:
$$
\begin{array}{l} e _ {0} \oplus \dots \oplus e _ {n - 1} = (\bar {\mathbf {A}} ^ {n - 1}, \bar {\mathbf {A}} x _ {n - 2} + \bar {\mathbf {B}} u _ {n - 1}, 0) \\ = (\bar {\mathbf {A}} ^ {n - 1}, x _ {n - 1}, 0) \\ e _ {0} \oplus \dots \oplus e _ {n} = (\bar {\mathbf {A}}, \bar {\mathbf {B}} u _ {n}, 1) \\ = (\bar {\mathbf {A}}, x _ {n}, 1) \\ \end{array}
$$
$$
\begin{array}{l} e _ {0} \oplus \dots \oplus e _ {n + 1} = (\bar {\mathbf {A}} ^ {2}, \bar {\mathbf {A}} x _ {n} + \bar {\mathbf {B}} u _ {n + 1}, 1) \\ = (\bar {\mathbf {A}} ^ {2}, x _ {n + 1}, 1) \\ \end{array}
$$
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-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.
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
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.
(a)
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.
$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.
(b)
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.
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  $\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].
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  $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 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
architecture and hyperparameters in the supplementary material.