mlsanityData to 14 Oct 2026

Structured State Space Models for In-Context Reinforcement Learning

arXiv 2303.03982 ↗First posted 7 Mar 2023cs.LGcited by 111

On the map: State-space sequence models (S4, Mamba) → Core SSM architectures for language

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.

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.

Builds on

30 of 45 references.

Closest ancestors · 2

Same direction, newest first

  1. 2022Simplified State Space Layers for Sequence Modeling9 Aug
  2. 2021Efficiently Modeling Long Sequences with Structured State Spaces31 Oct

Other references · 28

  1. 2017Attention Is All You Needcited by 60,361
  2. 2017Proximal Policy Optimization Algorithmscited by 15,215
  3. 2014Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translationcited by 4,057
  4. 2021Decision Transformer: Reinforcement Learning via Sequence Modelingcited by 1,546
  5. 2018DeepMind Control Suitecited by 956
  6. 2022A Generalist Agentcited by 808
  7. 2020Long Range Arena: A Benchmark for Efficient Transformerscited by 700
  8. 2021Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layerscited by 608
  9. 2020HiPPO: Recurrent Memory with Optimal Polynomial Projectionscited by 605
  10. 2016RL2: Fast Reinforcement Learning via Slow Reinforcement Learningcited by 503
  11. 2016Learning to reinforcement learncited by 377
  12. 2019Stabilizing Transformers for Reinforcement Learningcited by 270
16 more
  1. 2022In-context Reinforcement Learning with Algorithm Distillationcited by 190
  2. 2020Implementation Matters in Deep Policy Gradients: A Case Study on PPO and TRPOcited by 186
  3. 2022It's Raw! Audio Generation with State-Space Modelscited by 171
  4. 2023Human-Timescale Adaptation in an Open-Ended Task Spacecited by 138
  5. 2019Behaviour Suite for Reinforcement Learningcited by 132
  6. 2021Recurrent Model-Free RL Can Be a Strong Baseline for Many POMDPscited by 122
  7. 2023A Tutorial on Meta-Reinforcement Learningcited by 108
  8. 2022General-Purpose In-Context Learning by Meta-Learning Transformerscited by 105
  9. 2022Discovered Policy Optimisationcited by 104
  10. 2018Evolved Policy Gradientscited by 88
  11. 2021Gradients are Not All You Needcited by 86
  12. 2020Discovering Reinforcement Learning Algorithmscited by 77
  13. 2023POPGym: Benchmarking Partially Observable Reinforcement Learningcited by 62
  14. 2021Muesli: Combining Improvements in Policy Optimizationcited by 53
  15. 2021Efficient Transformers in Reinforcement Learning using Actor-Learner Distillationcited by 35
  16. 2021Introducing Symmetries to Black Box Meta Reinforcement Learningcited by 25

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2023: 15; 2024: 41; 2025: 38; 2026: 17.

111 papers.

Sort:

By area:

  1. 2023Supervised Pretraining Can Learn In-Context Reinforcement Learningcited by 139
  2. 2024Cobra: Extending Mamba to Multi-Modal Large Language Model for Efficient Inferencecited by 97
  3. 2024The Hidden Attention of Mamba Modelscited by 91
  4. 2024Theoretical Foundations of Deep Selective State-Space Modelscited by 74
  5. 2025Dynamic Chunking for End-to-End Hierarchical Sequence Modelingcited by 72
  6. 2023JaxMARL: Multi-Agent RL Environments and Algorithms in JAXcited by 63
  7. 2024State Space Model for New-Generation Network Alternative to Transformers: A Surveycited by 59
  8. 2023XLand-MiniGrid: Scalable Meta-Reinforcement Learning Environments in JAXcited by 55
  9. 2024Kinetix: Investigating the Training of General Agents through Open-Ended Physics-Based Control Taskscited by 45
  10. 2023AMAGO: Scalable In-Context Reinforcement Learning for Adaptive Agentscited by 41
  11. 2023Generalization to New Sequential Decision Making Tasks with In-Context Learningcited by 39
  12. 2024Mastering Memory Tasks with World Modelscited by 38
38 more
  1. 2025Memory, Benchmark & Robots: A Benchmark for Solving Complex Tasks with Reinforcement Learningcited by 36
  2. 2024MambaLRP: Explaining Selective State Space Sequence Modelscited by 32
  3. 2023Facing Off World Model Backbones: RNNs, Transformers, and S4cited by 31
  4. 2025Reward Is Enough: LLMs Are In-Context Reinforcement Learnerscited by 31
  5. 2023Introspective Tips: Large Language Model for In-Context Decision Makingcited by 28
  6. 2024Understanding the differences in Foundation Models: Attention, State Space Models, and Recurrent Neural Networkscited by 26
  7. 2023In-Context Reinforcement Learning for Variable Action Spacescited by 24
  8. 2024Retrieval-Augmented Decision Transformer: External Memory for In-context RLcited by 21
  9. 2024Discovering Temporally-Aware Reinforcement Learning Algorithmscited by 18
  10. 2024NAVIX: Scaling MiniGrid Environments with JAXcited by 18
  11. 2024Slot State Space Modelscited by 17
  12. 2023Memory Gym: Towards Endless Tasks to Benchmark Memory Capabilities of Agentscited by 16
  13. 2024AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformerscited by 15
  14. 2024XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learningcited by 14
  15. 2023When Meta-Learning Meets Online and Continual Learning: A Surveycited by 12
  16. 2024A Novel State Space Model with Local Enhancement and State Sharing for Image Fusioncited by 11
  17. 2024The Overcooked Generalisation Challenge: Evaluating Cooperation with Novel Partners in Unknown Environments Using Unsupervised Environment Designcited by 11
  18. 2025From Markov to Laplace: How Mamba In-Context Learns Markov Chainscited by 11
  19. 2023minimax: Efficient Baselines for Autocurricula in JAXcited by 10
  20. 2023On the Long Range Abilities of Transformerscited by 10
  21. 2024Do Transformer World Models Give Better Policy Gradients?cited by 10
  22. 2024Efficient Recurrent Off-Policy RL Requires a Context-Encoder-Specific Learning Ratecited by 10
  23. 2024Pretraining Decision Transformers with Reward Prediction for In-Context Multi-task Structured Bandit Learningcited by 10
  24. 2024Sable: a Performant, Efficient and Scalable Sequence Model for MARLcited by 10
  25. 2025Accelerating Model-Based Reinforcement Learning with State-Space World Modelscited by 10
  26. 2024Rethinking Transformers in Solving POMDPscited by 9
  27. 2024Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learningcited by 9
  28. 2025Learning to Dissipate Energy in Oscillatory State-Space Modelscited by 9
  29. 2023Hieros: Hierarchical Imagination on Structured State Space Sequence World Modelscited by 8
  30. 2024S4Fusion: Saliency-aware Selective State Space Model for Infrared Visible Image Fusioncited by 8
  31. 2024Artificial Generational Intelligence: Cultural Accumulation in Reinforcement Learningcited by 8
  32. 2024In-context learning and Occam's razorcited by 8
  33. 2025Mixture-of-Experts Meets In-Context Reinforcement Learningcited by 8
  34. 2025Revisiting associative recall in modern recurrent modelscited by 8
  35. 2024Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning?cited by 7
  36. 2024Mamba-FSCIL: Dynamic Adaptation with Selective State Space Model for Few-Shot Class-Incremental Learningcited by 7
  37. 2024MambaPlace:Text-to-Point-Cloud Cross-Modal Place Recognition with Attention Mamba Mechanismscited by 7
  38. 2025Learning Fine-grained Domain Generalization via Hyperbolic State Space Hallucinationcited by 7

50 most cited of 111 All 111 (JSON)

Similar, not cited

Close in wording, no citation either way

  1. 2023Decision S4: Efficient Sequence-Based RL via State Spaces Layers
  2. 2025Advancing Intelligent Sequence Modeling: Evolution, Trade-offs, and Applications of State-Space Architectures from S4 to Mamba
  3. 2024Overcoming Slow Decision Frequencies in Continuous Control: Model-Based Sequence Reinforcement Learning for Model-Free Control
  4. 2025Gym4ReaL: A Suite for Benchmarking Real-World Reinforcement Learning
  5. 2021Offline Reinforcement Learning as One Big Sequence Modeling Problem
  6. 2026Improving and Accelerating Offline RL in Large Discrete Action Spaces with Structured Policy Initialization
  7. 2024State Space Models as Foundation Models: A Control Theoretic Overview
  8. 2022On Uncertainty in Deep State Space Models for Model-Based Reinforcement Learning