# Mastering Memory Tasks with World Models

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

## Key facts

- Cited by 38 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 45 of its 71 references that are in the map.
- First posted on arXiv 7 Mar 2024.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [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)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Pointer Sentinel Mixture Models](https://mlsanity.com/paper/1609.07843/) (2016)
- [Efficiently Modeling Long Sequences with Structured State Spaces](https://mlsanity.com/paper/2111.00396/) (2021)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [Big Bird: Transformers for Longer Sequences](https://mlsanity.com/paper/2007.14062/) (2020)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Speech Commands: A Dataset for Limited-Vocabulary Speech Recognition](https://mlsanity.com/paper/1804.03209/) (2018)
- [Language Modeling with Gated Convolutional Networks](https://mlsanity.com/paper/1612.08083/) (2016)
- [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)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [On the Parameterization and Initialization of Diagonal State Space Models](https://mlsanity.com/paper/2206.11893/) (2022)
- [Learning Deep Transformer Models for Machine Translation](https://mlsanity.com/paper/1906.01787/) (2019)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [Long Range Language Modeling via Gated State Spaces](https://mlsanity.com/paper/2206.13947/) (2022)
- [Mastering Atari Games with Limited Data](https://mlsanity.com/paper/2111.00210/) (2021)

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

- [V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planning](https://mlsanity.com/paper/2506.09985/) (2025)
- [Understanding World or Predicting Future? A Comprehensive Survey of World Models](https://mlsanity.com/paper/2411.14499/) (2024)
- [State Space Model for New-Generation Network Alternative to Transformers: A Survey](https://mlsanity.com/paper/2404.09516/) (2024)
- [Memory, Benchmark & Robots: A Benchmark for Solving Complex Tasks with Reinforcement Learning](https://mlsanity.com/paper/2502.10550/) (2025)
- [GenRL: Multimodal-foundation world models for generalization in embodied agents](https://mlsanity.com/paper/2406.18043/) (2024)
- [Streaming Deep Reinforcement Learning Finally Works](https://mlsanity.com/paper/2410.14606/) (2024)
- [Making Large Language Models into World Models with Precondition and Effect Knowledge](https://mlsanity.com/paper/2409.12278/) (2024)
- [Slot State Space Models](https://mlsanity.com/paper/2406.12272/) (2024)
- [SEEA-R1: Tree-Structured Reinforcement Fine-Tuning for Self-Evolving Embodied Agents](https://mlsanity.com/paper/2506.21669/) (2025)
- [Open-World Reinforcement Learning over Long Short-Term Imagination](https://mlsanity.com/paper/2410.03618/) (2024)
- [UniZero: Generalized and Efficient Planning with Scalable Latent World Models](https://mlsanity.com/paper/2406.10667/) (2024)
- [Episodic memory in AI agents poses risks that should be studied and mitigated](https://mlsanity.com/paper/2501.11739/) (2025)
- [Rethinking Transformers in Solving POMDPs](https://mlsanity.com/paper/2405.17358/) (2024)
- [KalMamba: Towards Efficient Probabilistic State Space Models for RL under Uncertainty](https://mlsanity.com/paper/2406.15131/) (2024)
- [Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learning](https://mlsanity.com/paper/2410.10132/) (2024)
- [medDreamer: Model-Based Reinforcement Learning with Latent Imagination on Complex EHRs for Clinical Decision Support](https://mlsanity.com/paper/2505.19785/) (2025)
- [Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning?](https://mlsanity.com/paper/2405.12094/) (2024)
- [Paying Less Generalization Tax: A Cross-Domain Generalization Study of RL Training for LLM Agents](https://mlsanity.com/paper/2601.18217/) (2026)
- [Uncertainty Representations in State-Space Layers for Deep Reinforcement Learning under Partial Observability](https://mlsanity.com/paper/2409.16824/) (2024)
- [Unraveling the Complexity of Memory in RL Agents: an Approach for Classification and Evaluation](https://mlsanity.com/paper/2412.06531/) (2024)
- [A Definition and Roadmap for World Models](https://mlsanity.com/paper/2607.06401/) (2026)
- [Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficient](https://mlsanity.com/paper/2410.08893/) (2024)
- [FACTS: A Factored State-Space Framework For World Modelling](https://mlsanity.com/paper/2410.20922/) (2024)
- [EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modeling](https://mlsanity.com/paper/2502.00466/) (2025)
- [The Yokai Learning Environment: Tracking Beliefs Over Space and Time](https://mlsanity.com/paper/2508.12480/) (2025)
- [Context and Diversity Matter: The Emergence of In-Context Learning in World Models](https://mlsanity.com/paper/2509.22353/) (2025)
- [Beyond Sliding Windows: Learning to Manage Memory in Non-Markovian Environments](https://mlsanity.com/paper/2512.19154/) (2025)
- [MetaMind: General and Cognitive World Models in Multi-Agent Systems by Meta-Theory of Mind](https://mlsanity.com/paper/2603.00808/) (2026)
- [Latent State Design for World Models under Sufficiency Constraints](https://mlsanity.com/paper/2605.01694/) (2026)
- [Composition of Memory Experts for Diffusion World Models](https://mlsanity.com/paper/2605.18813/) (2026)

## Abstract (arXiv)

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.