{"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/mastering-memory-tasks-with-world-models","title":"Mastering Memory Tasks with World Models","arxiv_id":"2403.04253","date":"2024-03-07","proceeding":null,"authors":["Mohammad Reza Samsami","Artem Zholus","Janarthanan Rajendran","Sarath Chandar"],"abstract":"Current model-based reinforcement learning (MBRL) agents struggle with long-term dependencies. This limits their ability to effectively solve tasks involving extended time gaps between actions and outcomes, or tasks demanding the recalling of distant observations to inform current actions. To improve temporal coherence, we integrate a new family of state space models (SSMs) in world models of MBRL agents to present a new method, Recall to Imagine (R2I). This integration aims to enhance both long-term memory and long-horizon credit assignment. Through a diverse set of illustrative tasks, we systematically demonstrate that R2I not only establishes a new state-of-the-art for challenging memory and credit assignment RL tasks, such as BSuite and POPGym, but also showcases superhuman performance in the complex memory domain of Memory Maze. At the same time, it upholds comparable performance in classic RL tasks, such as Atari and DMC, suggesting the generality of our method. We also show that R2I is faster than the state-of-the-art MBRL method, DreamerV3, resulting in faster wall-time convergence.","url_abs":"https://arxiv.org/abs/2403.04253v1","url_pdf":"https://arxiv.org/pdf/2403.04253v1.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":"mastering-memory-tasks-with-world-models","repo_url":"https://github.com/chandar-lab/Recall2Imagine","is_official":1,"mentioned_in_paper":0,"mentioned_in_github":1,"framework":"jax","reach":{"status":"ok","spdx":"MIT"}}],"tasks":[{"task_slug":"model-based-reinforcement-learning","task_name":"Model-based Reinforcement Learning"},{"task_slug":"state-space-models","task_name":"State Space Models"}],"methods":[{"method_slug":"set","method_name":"SET"}],"datasets_introduced":[],"methods_introduced":[],"results":[],"syntology":{"atlas_url":"https://app.syntology.ai/?focus=2403.04253","mcp":{"get_harvested_code_for_paper":{"arxiv_id":"2403.04253"}},"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/chandar-lab/Recall2Imagine","reach":{"status":"ok","spdx":"MIT"}}],"summary":{"ran":4,"unverified":8},"by_repo_kind":{"official":{"samples":12,"ran":4,"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":0,"samples":[{"code_sha256_prefix":"ed94f4331173c9bc","entry":"cast_to_compute","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/jaxutils.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/jaxutils.py","link_basis":"first_harvest_node","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"ed94f4331173c9bc"}},{"code_sha256_prefix":"bbd25fa32893729e","entry":"pure","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/ninjax.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/ninjax.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"bbd25fa32893729e"}},{"code_sha256_prefix":"e5b1910ea05891b1","entry":"subsample","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/jaxutils.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/jaxutils.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"e5b1910ea05891b1"}},{"code_sha256_prefix":"67c6d2355562342f","entry":"tensorstats","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/jaxutils.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/jaxutils.py","link_basis":"harvester_set","language":"python","status":"ran","verification_level":1,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"67c6d2355562342f"}},{"code_sha256_prefix":"4b830637b16db4c8","entry":"convert","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/embodied/core/basics.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/embodied/core/basics.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"4b830637b16db4c8"}},{"code_sha256_prefix":"3360b32511bc42d4","entry":"format_","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/embodied/core/basics.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/embodied/core/basics.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"3360b32511bc42d4"}},{"code_sha256_prefix":"03dd934385b4d240","entry":"get_act","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/nets.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/nets.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"03dd934385b4d240"}},{"code_sha256_prefix":"1ad869363216c281","entry":"grad","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/ninjax.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/ninjax.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"1ad869363216c281"}},{"code_sha256_prefix":"ce9e0c5374238d06","entry":"init_mimo_ssm","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/ssm_nets.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/ssm_nets.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"ce9e0c5374238d06"}},{"code_sha256_prefix":"ad8b4b5ddd9df92c","entry":"init_siso_ssm","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/ssm_nets.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/ssm_nets.py","link_basis":"first_harvest_node","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"ad8b4b5ddd9df92c"}},{"code_sha256_prefix":"45b97b51433849a7","entry":"rng","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/ninjax.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/ninjax.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"45b97b51433849a7"}},{"code_sha256_prefix":"5ba274a4e18b46f3","entry":"treemap","repo":"chandar-lab/Recall2Imagine","repo_kind":"official","path":"recall2imagine/embodied/core/basics.py","file_url":"https://github.com/chandar-lab/Recall2Imagine/blob/HEAD/recall2imagine/embodied/core/basics.py","link_basis":"harvester_set","language":"python","status":"unverified","verification_level":0,"contract_check":null,"metamorphic_tier":null,"behaviour_fingerprint":false,"licence":"MIT","inline_ok":true,"mcp_get_code":{"code_sha256":"5ba274a4e18b46f3"}}]},"arxiv_metadata":null,"syntology_extracted_results":null}