# Structured State Space Models for In-Context Reinforcement Learning

arXiv: https://arxiv.org/abs/2303.03982 · HTML: https://mlsanity.com/paper/2303.03982/

## Key facts

- Cited by 111 papers in the map (data through 14 Oct 2026).
- Placed in Core SSM architectures for language, one of 10 directions in State-space sequence models (S4, Mamba).
- Builds on 30 of its 45 references that are in the map.
- First posted on arXiv 7 Mar 2023.

## Placed in

- State-space sequence models (S4, Mamba) / [Core SSM architectures for language](https://mlsanity.com/ssm/core-ssm-architectures-for-language/)

## Builds on (most cited first, 30 of 30)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translation](https://mlsanity.com/paper/1406.1078/) (2014)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Simplified State Space Layers for Sequence Modeling](https://mlsanity.com/paper/2208.04933/) (2022)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layers](https://mlsanity.com/paper/2110.13985/) (2021)
- [HiPPO: Recurrent Memory with Optimal Polynomial Projections](https://mlsanity.com/paper/2008.07669/) (2020)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [In-context Reinforcement Learning with Algorithm Distillation](https://mlsanity.com/paper/2210.14215/) (2022)
- [Implementation Matters in Deep Policy Gradients: A Case Study on PPO and TRPO](https://mlsanity.com/paper/2005.12729/) (2020)
- [It's Raw! Audio Generation with State-Space Models](https://mlsanity.com/paper/2202.09729/) (2022)
- [Human-Timescale Adaptation in an Open-Ended Task Space](https://mlsanity.com/paper/2301.07608/) (2023)
- [Behaviour Suite for Reinforcement Learning](https://mlsanity.com/paper/1908.03568/) (2019)
- [Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPs](https://mlsanity.com/paper/2110.05038/) (2021)
- [A Tutorial on Meta-Reinforcement Learning](https://mlsanity.com/paper/2301.08028/) (2023)
- [General-Purpose In-Context Learning by Meta-Learning Transformers](https://mlsanity.com/paper/2212.04458/) (2022)
- [Discovered Policy Optimisation](https://mlsanity.com/paper/2210.05639/) (2022)
- [Evolved Policy Gradients](https://mlsanity.com/paper/1802.04821/) (2018)
- [Gradients are Not All You Need](https://mlsanity.com/paper/2111.05803/) (2021)
- [Discovering Reinforcement Learning Algorithms](https://mlsanity.com/paper/2007.08794/) (2020)
- [POPGym: Benchmarking Partially Observable Reinforcement Learning](https://mlsanity.com/paper/2303.01859/) (2023)
- [Muesli: Combining Improvements in Policy Optimization](https://mlsanity.com/paper/2104.06159/) (2021)
- [Efficient Transformers in Reinforcement Learning using Actor-Learner Distillation](https://mlsanity.com/paper/2104.01655/) (2021)
- [Introducing Symmetries to Black Box Meta Reinforcement Learning](https://mlsanity.com/paper/2109.10781/) (2021)

## Built on by (most cited first, 30 of 111)

- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [Cobra: Extending Mamba to Multi-Modal Large Language Model for Efficient Inference](https://mlsanity.com/paper/2403.14520/) (2024)
- [The Hidden Attention of Mamba Models](https://mlsanity.com/paper/2403.01590/) (2024)
- [Theoretical Foundations of Deep Selective State-Space Models](https://mlsanity.com/paper/2402.19047/) (2024)
- [Dynamic Chunking for End-to-End Hierarchical Sequence Modeling](https://mlsanity.com/paper/2507.07955/) (2025)
- [JaxMARL: Multi-Agent RL Environments and Algorithms in JAX](https://mlsanity.com/paper/2311.10090/) (2023)
- [State Space Model for New-Generation Network Alternative to Transformers: A Survey](https://mlsanity.com/paper/2404.09516/) (2024)
- [XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAX](https://mlsanity.com/paper/2312.12044/) (2023)
- [Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Tasks](https://mlsanity.com/paper/2410.23208/) (2024)
- [AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agents](https://mlsanity.com/paper/2310.09971/) (2023)
- [Generalization to New Sequential Decision Making Tasks with In-Context Learning](https://mlsanity.com/paper/2312.03801/) (2023)
- [Mastering Memory Tasks with World Models](https://mlsanity.com/paper/2403.04253/) (2024)
- [Memory, Benchmark & Robots: A Benchmark for Solving Complex Tasks with Reinforcement Learning](https://mlsanity.com/paper/2502.10550/) (2025)
- [MambaLRP: Explaining Selective State Space Sequence Models](https://mlsanity.com/paper/2406.07592/) (2024)
- [Facing Off World Model Backbones: RNNs, Transformers, and S4](https://mlsanity.com/paper/2307.02064/) (2023)
- [Reward Is Enough: LLMs Are In-Context Reinforcement Learners](https://mlsanity.com/paper/2506.06303/) (2025)
- [Introspective Tips: Large Language Model for In-Context Decision Making](https://mlsanity.com/paper/2305.11598/) (2023)
- [Understanding the differences in Foundation Models: Attention, State Space Models, and Recurrent Neural Networks](https://mlsanity.com/paper/2405.15731/) (2024)
- [In-Context Reinforcement Learning for Variable Action Spaces](https://mlsanity.com/paper/2312.13327/) (2023)
- [Retrieval-Augmented Decision Transformer: External Memory for In-context RL](https://mlsanity.com/paper/2410.07071/) (2024)
- [Discovering Temporally-Aware Reinforcement Learning Algorithms](https://mlsanity.com/paper/2402.05828/) (2024)
- [NAVIX: Scaling MiniGrid Environments with JAX](https://mlsanity.com/paper/2407.19396/) (2024)
- [Slot State Space Models](https://mlsanity.com/paper/2406.12272/) (2024)
- [Memory Gym: Towards Endless Tasks to Benchmark Memory Capabilities of Agents](https://mlsanity.com/paper/2309.17207/) (2023)
- [AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers](https://mlsanity.com/paper/2411.11188/) (2024)
- [XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning](https://mlsanity.com/paper/2406.08973/) (2024)
- [When Meta-Learning Meets Online and Continual Learning: A Survey](https://mlsanity.com/paper/2311.05241/) (2023)
- [A Novel State Space Model with Local Enhancement and State Sharing for Image Fusion](https://mlsanity.com/paper/2404.09293/) (2024)
- [The Overcooked Generalisation Challenge: Evaluating Cooperation with Novel Partners in Unknown Environments Using Unsupervised Environment Design](https://mlsanity.com/paper/2406.17949/) (2024)
- [From Markov to Laplace: How Mamba In-Context Learns Markov Chains](https://mlsanity.com/paper/2502.10178/) (2025)

## Abstract (arXiv)

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.