SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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.
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.
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.
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.
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
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].
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.
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.
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.
[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 Rocktäschel, Yannick Schroecker, Jakub Sygnowski, Karl Tuyls, Sarah York, Alexander Zacherl, and Lei Zhang. Human-timescale adaptation in an open-ended task space. arXiv e-prints, 2023.
[2] Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. A survey of meta-reinforcement learning. arXiv preprint arXiv:2301.08028, 2023.
[3] James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL 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.
[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. R12: 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 Ré. 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 Ré. Hippo: Recurrent memory with optimal polynomial projections. Advances in neural information processing systems, 33: 1474-1487, 2020.
A Proof of Associativity of Binary Operator
Recall that  $\oplus$  is defined as:
This is equivalent to the following:
$$
a _ {i} \oplus a _ {j} := \left\{ \begin{array}{l l} ((a _ {i} \bullet a _ {j}) _ {a}, (a _ {i} \bullet a _ {j}) _ {b}, a _ {i, c}) & \text {i f} a _ {j, c} = 0 \\ a _ {j} & \text {i f} a _ {j, c} = 1 \end{array} \right.
$$
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$
$$
\begin{array}{l} (x \oplus y) \oplus z = z (10) \\ = y \oplus z (11) \\ = x \oplus (y \oplus z) (12) \\ \end{array}
$$
Case 2:  $z_{c} = 0$  and  $y_{c} = 1$
$$
(x \oplus y) \oplus z = y \oplus z \tag {13}
$$
Note that  $(y\oplus z)_{c} = 1$  thus, (14)
$$
= x \oplus (y \oplus z) \tag {15}
$$
Case 3:  $z_{c} = 0$  and  $y_{c} = 0$
$$
\begin{array}{l} (x \oplus y) \oplus z = ((x \bullet y) _ {a}, (x \bullet y) _ {b}, x _ {c}) \oplus z (16) \\ = \left(\left(\left(x \bullet y\right) \bullet z\right) _ {a}, \left(\left(x \bullet y\right) \bullet z\right) _ {b}, x _ {c}\right) (17) \\ = \left(\left(x \bullet (y \bullet z)\right) _ {a}, \left(x \bullet (y \bullet z)\right) _ {b}, x _ {c}\right) (18) \\ = x \oplus ((y \bullet z) _ {a}, (y \bullet z) _ {b}, y _ {c}) (19) \\ = x \oplus (y \oplus z) (20) \\ \end{array}
$$
Table 2: Hyperparameters for training A2C on Bsuite
Parameter	Value
Adam Learning Rate	3e-4
Entropy Coefficient	0.0
Encoder Layer Sizes	[256, 256]
Number of Recurrent Layers	1
Size of Recurrent Layer	256
Discount γ	0.99
TD λ	0.9
Number of Environments	1
Unroll Length	32
Number of Episodes	10000
Activation Function	ReLU
Table 3: Hyperparameters for training PPO on POPGym
Parameter	Value
Adam Learning Rate	5e-5
Number of Environments	64
Unroll Length	1024
Number of Timesteps	15e6
Number of Epochs	30
Number of minibatches	8
Discount γ	0.99
GAE λ	1.0
Clipping Coefficient ε	0.2
Entropy Coefficient	0.0
Value Function Weight	1.0
Maximum Gradient Norm	0.5
Learning Rate Annealing	None
Activation Function	LeakyReLU
Encoder Layer Sizes	[128, 256]
Recurrent Layer Hidden Size	256
Action Decoder Layer Sizes	[128, 128]
Value Decoder Layer Sizes	[128, 128]
S5 Layers	4
S5 Hidden Size	256
S5 Discretization	ZOH
S5 Δ min	0.001
S5 Δ max	0.1
Table 4: Hyperparameters for training Muesli on Multi-Environment Meta-RL. These experiments were run using 64 TPUv3's.
Parameter	Value
Adam Learning Rate	3e-4
Value Function Weight	1.0
Muesli Auxiliary Loss Weight	3.0
Muesli Model Unroll Length	1.0
Encoder Layer Sizes	[512, 512]
Number of Environments	1024
Discount γ	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 Δ min	0.001
S5 Δ max	0.1
C POPGym Discussion
Stateless CartPole Hard	Noisy Stateless CartPole Hard	Stateless Pendulum Hard	Noisy Stateless Pendulum Hard	Repeat Previous Hard
S5 (ours)	1.0 ± 0.0	0.28 ± 0.0	0.79 ± 0.01	0.55 ± 0.01	0.91 ± 0.01
GRU (ours)	1.0 ± 0.0	0.27 ± 0.0	0.75 ± 0.0	0.61 ± 0.01	-0.46 ± 0.01
MLP (ours)	0.26 ± 0.0	0.22 ± 0.0	0.41 ± 0.02	0.34 ± 0.01	-0.48 ± 0.00
GRU	1.000 ± 0.000	0.390 ± 0.007	0.828 ± 0.001	0.657 ± 0.002	-0.428 ± 0.002
MLP	0.265 ± 0.002	0.229 ± 0.002	0.477 ± 0.030	0.351 ± 0.012	-0.486 ± 0.002
IndRNN	1.000 ± 0.000	0.404 ± 0.005	0.804 ± 0.023	0.521 ± 0.109	-0.384 ± 0.013
LMU	0.987 ± 0.007	0.352 ± 0.019	0.806 ± 0.006	0.563 ± 0.014	0.191 ± 0.041
S4D	0.127 ± 0.026	0.207 ± 0.007	0.303 ± 0.014	0.289 ± 0.011	-0.491 ± 0.001
FART	0.996 ± 0.000	0.366 ± 0.002	0.698 ± 0.077	0.553 ± 0.007	-0.485 ± 0.001
Table 5: Results in POPGym's suite. The reported number 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.
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.