{"about":{"site":"https://codewithpapers.app","non_affiliation":"Code with Papers and Syntology are not affiliated with, endorsed by, or sponsored by Papers with Code, Meta, or the pwc-archive mirror.","licence":"CC BY-SA 4.0","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","attribution":"https://codewithpapers.app/attribution","modified":"archive material modified by Syntology; see the attribution page"},"url":"/paper/structured-state-space-models-for-in-context-1","title":"Structured State Space Models for In-Context Reinforcement Learning","arxiv_id":"2303.03982","date":"2023-03-07","proceeding":"NeurIPS 2023 11","authors":["Chris Lu","Yannick Schroecker","Albert Gu","Emilio Parisotto","Jakob Foerster","Satinder Singh","Feryal Behbahani"],"abstract":"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.","url_abs":"https://arxiv.org/abs/2303.03982v3","url_pdf":"https://arxiv.org/pdf/2303.03982v3.pdf","source":{"archive":"pwc-archive (Hugging Face), CC BY-SA 4.0","snapshot":"2025-07-28","licence_url":"https://creativecommons.org/licenses/by-sa/4.0/legalcode","row_kind":"abstracts"},"code_links":[{"paper_slug":"structured-state-space-models-for-in-context-1","repo_url":"https://github.com/luchris429/popjaxrl","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":0,"framework":"jax","reach":null},{"paper_slug":"structured-state-space-models-for-in-context-1","repo_url":"https://github.com/luchris429/s5rl","is_official":1,"mentioned_in_paper":1,"mentioned_in_github":0,"framework":"jax","reach":{"status":"ok"}}],"tasks":[{"task_slug":"continuous-control","task_name":"Continuous Control"},{"task_slug":null,"task_name":"In-Context Reinforcement Learning"},{"task_slug":"meta-learning","task_name":"Meta-Learning"},{"task_slug":"reinforcement-learning","task_name":"Reinforcement Learning"},{"task_slug":"state-space-models","task_name":"State Space Models"},{"task_slug":"continuous-control","task_name":"continuous-control"},{"task_slug":"reinforcement-learning-2","task_name":"reinforcement-learning"}],"methods":[],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"syntology_url":null,"atlas_url":"https://app.syntology.ai/?focus=2303.03982","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2303.03982"}},"developers":"https://syntology.ai/developers","read_at":"2026-09-24T18:15:14+00:00","read_at_is":"when the build read Syntology's graph, not when any sample ran","claim":"Per-sample execution status on synthesized fixtures; not a correctness claim about the paper. Samples come from repositories linked to the paper, official or community; repo_kind says which.","repos":[{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/luchris429/popjaxrl","reach":null},{"provenance":"external:paperswithcode_snapshot_2025-07-28","url":"https://github.com/luchris429/s5rl","reach":{"status":"ok"}}],"summary":{"ran_draft_wrong":2,"ran_fixture":1},"by_repo_kind":{"official":{"samples":3,"ran":3,"repositories":1}},"repo_kind_vocabulary":{"official":"The archive marks this repository official for the paper","named_in_paper":"The archive records that the paper mentions this repository; it is not marked official","listed":"In the archive's code links for this paper, not marked official and not recorded as mentioned in the paper","found_in_text":"Syntology found this repository in the paper's own text; whether it is the authors' implementation is not asserted","community":"Not in the archive's code links for this paper; a community repository Syntology harvested"},"n_pointer_only_for_licence":3,"samples":[{"code_sha256_prefix":"ec8a321400aae8d2","entry":"init_VinvB","repo":"luchris429/popjaxrl","repo_kind":"official","path":"algorithms/s5.py","file_url":"https://github.com/luchris429/popjaxrl/blob/HEAD/algorithms/s5.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"MISDECLARED","metamorphic_tier":"invariant","behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"ec8a321400aae8d2"}},{"code_sha256_prefix":"a4bdc24302b8ab98","entry":"init_log_steps","repo":"luchris429/popjaxrl","repo_kind":"official","path":"algorithms/s5.py","file_url":"https://github.com/luchris429/popjaxrl/blob/HEAD/algorithms/s5.py","link_basis":"first_harvest_node","language":"python","status":"ran_fixture","verification_level":1,"contract_check":"RAISES","metamorphic_tier":"invariant","behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"a4bdc24302b8ab98"}},{"code_sha256_prefix":"8b1a3bc2f6b64dce","entry":"log_step_initializer","repo":"luchris429/popjaxrl","repo_kind":"official","path":"algorithms/s5.py","file_url":"https://github.com/luchris429/popjaxrl/blob/HEAD/algorithms/s5.py","link_basis":"first_harvest_node","language":"python","status":"ran_draft_wrong","verification_level":1,"contract_check":"OUTPUT_MISDECLARED","metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"NONE","inline_ok":false,"mcp_get_code":{"code_sha256":"8b1a3bc2f6b64dce"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}