Source-linked AI summary
Structured State Space Models for In-Context Reinforcement Learning
Chris Lu, Yannick Schroecker, Albert Gu, Emilio Parisotto, Jakob Foerster, Satinder Singh, Feryal Behbahani
TL;DR
The paper asks whether structured state-space models can provide efficient memory and long-context processing for reinforcement learning. It modifies S5 to initialize and reset hidden states in parallel, then evaluates it on memory, partially observable, and projected Meta-RL tasks, where it reports strong performance and faster execution than recurrent baselines.
Problem
Reinforcement learning needs scalable sequence memory, but transformers are costly at inference time and prior Meta-RL sequence methods often have limited effective memory and adaptation horizons.
Method
The paper modifies S5 with resettable hidden states so it can train over fixed-length rollouts containing variable episode boundaries, and evaluates it in projected Meta-RL.
Results
S5 matches or outperforms recurrent models across the evaluated tasks while running nearly five times faster on POPGym and adapting to held-out tasks.
Takeaways & Limitations
Structured state-space models are fast and performant for in-context reinforcement-learning tasks, including long-context adaptation and out-of-distribution held-out tasks.
Takeaways & Limitations
The architecture requires JAX because its associative scan is not currently possible in PyTorch, and speedups are limited when short rollouts are sufficient.
Abstract
from arXiv · showhide
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/popjaxrl.
1 Introduction
S4 and S5 offer strong long-range sequence modeling with scalable inference and parallel training, motivating their use in reinforcement learning. The paper modifies S5 to reset hidden states within trajectories and evaluates its performance, speed, and long-context adaptation.
- Motivation: S4 models use scalable inference with constant memory and time per step in sequence length while remaining parallelisable across the sequence dimension during training.RNNs share the asymptotic inference property, but S4 models are reported as more performant and trainable in parallel.
- Motivation: Long-context Meta-RL can require trajectories thousands of steps long, but prior RL2 approaches are often limited by effective memory length and slow training.
- Motivation: S5 is modified to reset hidden states within training trajectories, addressing episode boundaries that fixed-kernel S4 models cannot handle directly.The modification targets existing recurrent reinforcement-learning frameworks.
- Contributions: The paper reports higher memory-task scores than RNNs, over six-times-faster POPGym execution, and strong performance on the difficult Repeat Previous Hard task.
2 Background
The paper develops S5 from state-space models and replaces convolution with associative parallel scans for flexible, efficient sequence processing. It situates these models in partially observable reinforcement learning, where policies may require histories of observations and actions.
- Structured State Space Models: A discretized state-space model updates a latent state from each input signal and produces an output using matrices determined by the continuous dynamics and step size.
- Structured State Space Models: S4 can operate recurrently for fast, memory-efficient inference or convolutionally for parallel training across timesteps.
- Structured State Space Models: S5 replaces S4’s convolution with parallel scans, simplifying modifications while exploiting associative composition for efficient sequence processing.Parallel scans compute prefix operations by reorganizing associative compositions.
- Reinforcement Learning: In a POMDP, the agent receives observations rather than states, so an optimal policy generally depends on prior observations and actions.
3 Method
The method modifies S5 for variable-length rollouts and episode-boundary resets while introducing a Meta-RL setting with randomly projected task spaces. These changes target the memory-state requirements of partially observable and broadly varying reinforcement-learning tasks.
- Resettable S5: The method modifies S5 to handle variable-length sequences, making it more suitable for partially observable environments.
- Resettable S5: Fixed-length on-policy rollouts can begin mid-episode or cross episode boundaries, requiring models to preserve prior memory while resetting at termination.RNNs store rollout-initial hidden states and reset at boundaries, whereas unmodified S4 lacks both mechanisms.
- Resettable S5: The modified S5 incorporates a reset flag into an associative operator, enabling parallel training over variable-length sequences and different hidden-state initializations.The binary done signal marks environment transitions requiring resets, and the operator is proven associative.
- Randomly Projected Meta-RL: The proposed Meta-RL setting uses random linear projections of observation and action spaces to broaden task diversity beyond fixed spaces and similar task distributions.The setting is intended to test broader generalisation across task complexity and dynamics.
4 Experiments
Across memory, partially observable, and meta-reinforcement-learning experiments, the modified S5 architecture combines strong performance with favorable runtime and generalization, including on randomly projected tasks.
- 4.1 Memory Length Environment: In the bsuite memory-length task, S5 has better asymptotic runtime than Transformers and far outperforms LSTMs in both performance and speed.The comparison covers memory usage, runtime, and return across different memory lengths.
- 4.1 Memory Length Environment: Theoretical training speed does not remove the O(N) rollout-collection bottleneck, and Transformers are omitted from later experiments because of poor long-sequence runtime.S5’s backward pass is theoretically O(log(N)) with enough processors.
- Experimental setup: The experiments use JAX-reimplemented POPGym environments and compare S5 with GRU, while other meta-learning experiments use randomly projected observations and actions.DMControl experiments use stacked S5 blocks and set trajectory length to the maximum episode length of 1k.
- 4.2 POPGym Environments: S5 performs well on POPGym’s challenging Repeat Previous Hard task, far outperforming tested baselines, while running over six times faster than GRU.Results are reported across implemented POPGym environments, with runtime averaged across environments.
- 4.3 Randomly-Projected CartPole In-Context: S5 outperforms GRU on randomly projected StatelessCartPole while running twice as quickly and continues learning beyond the 16-trial training horizon.Evaluation extends to 32 trials, where GRU appears to overfit the training horizon while S5 continues to perform well.
- 4.4 Multi-Environment Meta-Reinforcement Learning: S5 outperforms LSTMs in both sample efficiency and ultimate performance when meta-training across six DMControl environments with unseen random projections and no task labels.The agent must infer the environment from rewards and obfuscated dynamics.
- 4.4 Multi-Environment Meta-Reinforcement Learning: Without extra fine-tuning, the trained model transfers impressively to some held-out tasks but fails on completely unseen Pendulum Swingup with unseen spaces and reward dynamics.The evaluation covers five held-out DMControl tasks with random linear projections.
5 Related Work
Prior work applies sequence models to memory-based RL and meta-RL, but long-horizon adaptation and out-of-distribution generalization remain important challenges. This paper is situated among approaches using S4-family models, RNNs, Transformers, and broader task distributions.
- S4 models have also been studied for audio, video, and reinforcement learning, although naive S4 use performed poorly in one POMDP investigation.
- Most memory-based meta-RL work maximizes performance on training distributions, whereas other generalization approaches restrict architectures, losses, targets, or drift functions.
- Held-out DMControl evaluation uses random projections of environment observations and actions, with shaded regions showing return ranges across three seeds.
- Long-horizon meta-RL remains difficult because RNNs and Transformers struggle with very long sequences, while evolution strategies are sample inefficient and computationally expensive.
- Sequence models have been used for memory-based RL, offline RL, multitask behavioral cloning, algorithm distillation, and out-of-distribution meta-RL.
6 Conclusion and Limitations
The paper adapts S5 for reinforcement learning by enabling hidden-state resets within training sequences, then evaluates its speed, memory performance, POPGym performance, and meta-RL generalization. The authors report strong performance and faster execution, while noting framework, task-length, and compute limitations.
- The proposed S5 modification enables hidden-state resets within training sequences, allowing S5 layers to replace RNNs in existing recurrent reinforcement-learning frameworks.
- S5 matches or outperforms RNNs on POPGym while running nearly five times faster, and outperforms LSTMs in the proposed projected-task meta-learning setting.
- S5 is asymptotically faster than Transformers in sequence length and runs nearly twice as quickly as equally parameterized LSTMs while outperforming them on the bsuite memory-length task.
- The meta-learning setting samples DMControl tasks and random linear projections of their observation and action spaces, and S5 adapts to held-out tasks through in-context adaptation.
- Limitations include dependence on JAX because PyTorch cannot implement associative scan, limited speedups for short rollouts, and no comprehensive hyperparameter sweep due to compute costs.
A Proof of Associativity of Binary Operator
The proof establishes associativity of the paper’s reset-aware binary operator by rewriting its cases in terms of S5’s associative operator. It then applies associativity step by step to three arbitrary elements.
- The operator’s case definition either combines components through S5’s binary operator or selects the later element when its reset indicator is one.
- An equivalent formulation expresses the operator using S5’s binary operator and a conditional choice between the combined state and the later element.
- The proof reduces associativity of the reset-aware operator to associativity of S5’s binary operator.
- For arbitrary x, y, and z, the expansion transforms (x ⊕ y) ⊕ z into x ⊕ (y ⊕ z) by applying the binary operator’s associativity.
B Hyperparameters
The appendix lists training hyperparameters for A2C on Bsuite, PPO on POPGym, and Muesli on the multi-environment Meta-RL setting. The Meta-RL experiments used 64 TPUv3 devices.
- Table 2 lists the hyperparameters used to train A2C on Bsuite.
- Table 3 lists the hyperparameters used to train PPO on POPGym.
- Table 4 lists Muesli hyperparameters for multi-environment Meta-RL experiments run using 64 TPUv3 devices.
C POPGym Discussion
The discussion compares the authors’ POPGym implementation choices with prior work and reports strong S5 performance under those choices. Evaluation coverage was narrowed to selected hard environments rather than the full POPGym suite.
- Results: S5 far outperforms Morad et al.’s best architecture in the “RepeatPreviousHard” environment while recovering similar performance across environments.The comparison uses the authors’ StableBaselines3 and CleanRL recurrent PPO design decisions.
- Implementation: The implementation follows StableBaselines3 and CleanRL recurrent PPO designs rather than Morad et al.’s differing PPO setup.Morad et al. used a dynamic KL-divergence coefficient and full episode trajectories, whereas the authors follow recurrent PPO conventions described here.
- Limitations: Evaluation excluded the full POPGym suite and reported only selected CartPole, Pendulum, and Repeat environments at hard difficulty.The authors selected these environments for faster JAX-based experimentation because easier versions were not informative.