mlsanityData to 14 Oct 2026

Mastering Memory Tasks with World Models

arXiv 2403.04253 ↗First posted 7 Mar 2024cs.LGcited by 38

On the map: World models → Latent dynamics models and model-based RL

also in 1 more
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.

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.

Builds on

45 of 71 references.

Closest ancestors · 6

Same direction, newest first

  1. 2023Transformer-based World Models Are Happy With 100k Interactions13 Mar
  2. 2023Mastering Diverse Domains through World Models10 Jan
  3. 2021Mastering Atari Games with Limited Data30 Oct
  4. 2020Mastering Atari with Discrete World Models5 Oct
  5. 2019Dream to Control: Learning Behaviors by Latent Imagination3 Dec
  6. 2018World Models27 Mar

Other references · 39

  1. 2017Attention Is All You Needcited by 60,361
  2. 2020Language Models are Few-Shot Learnerscited by 32,419
  3. 2017Proximal Policy Optimization Algorithmscited by 15,215
  4. 2014Learning Phrase Representations using RNN Encoder-Decoder for Statistical Machine Translationcited by 4,057
  5. 2022FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awarenesscited by 3,673
  6. 2016Pointer Sentinel Mixture Modelscited by 3,096
  7. 2021Efficiently Modeling Long Sequences with Structured State Spacescited by 2,465
  8. 2019Transformer-XL: Attentive Language Models Beyond a Fixed-Length Contextcited by 1,850
  9. 2020Big Bird: Transformers for Longer Sequencescited by 1,684
  10. 2020Linformer: Self-Attention with Linear Complexitycited by 1,474
  11. 2019Mastering Atari, Go, Chess and Shogi by Planning with a Learned Modelcited by 1,353
  12. 2018Learning Latent Dynamics for Planning from Pixelscited by 1,264
27 more
  1. 2018DeepMind Control Suitecited by 956
  2. 2018IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architecturescited by 883
  3. 2021Deep Reinforcement Learning at the Edge of the Statistical Precipicecited by 807
  4. 2018Speech Commands: A Dataset for Limited-Vocabulary Speech Recognitioncited by 752
  5. 2016Language Modeling with Gated Convolutional Networkscited by 702
  6. 2020Long Range Arena: A Benchmark for Efficient Transformerscited by 700
  7. 2021Combining Recurrent, Convolutional, and Continuous-time Models with Linear State-Space Layerscited by 608
  8. 2020HiPPO: Recurrent Memory with Optimal Polynomial Projectionscited by 605
  9. 2019Model-Based Reinforcement Learning for Ataricited by 502
  10. 2022On the Parameterization and Initialization of Diagonal State Space Modelscited by 424
  11. 2019Learning Deep Transformer Models for Machine Translationcited by 309
  12. 2019Stabilizing Transformers for Reinforcement Learningcited by 270
  13. 2022Long Range Language Modeling via Gated State Spacescited by 265
  14. 2018RUDDER: Return Decomposition for Delayed Rewardscited by 173
  15. 2019Behaviour Suite for Reinforcement Learningcited by 132
  16. 2023Structured State Space Models for In-Context Reinforcement Learningcited by 111
  17. 2020Objective Mismatch in Model-based Reinforcement Learningcited by 96
  18. 2023When Do Transformers Shine in RL? Decoupling Memory from Credit Assignmentcited by 63
  19. 2023POPGym: Benchmarking Partially Observable Reinforcement Learningcited by 62
  20. 2018Optimizing Agent Behavior over Long Time Scales by Transporting Valuecited by 58
  21. 2022Evaluating Long-Term Memory in 3D Mazescited by 43
  22. 2022Efficient Long Sequence Modeling via State Space Augmented Transformercited by 37
  23. 2022Deep Transformer Q-Networks for Partially Observable Reinforcement Learningcited by 36
  24. 2021Control-Oriented Model-Based Reinforcement Learning with Implicit Differentiationcited by 35
  25. 2021Synthetic Returns for Long-Term Credit Assignmentcited by 27
  26. 2023Decision S4: Efficient Sequence-Based RL via State Spaces Layerscited by 26
  27. 2022Simplifying and Understanding State Space Models with Diagonal Linear RNNscited by 23

Built on by

Cited by year (2026 to 14 Oct)

Citing papers per year: 2024: 17; 2025: 13; 2026: 8.

38 papers.

Sort:

By area:

  1. 2025V-JEPA 2: Self-Supervised Video Models Enable Understanding, Prediction and Planningcited by 809
  2. 2024Understanding World or Predicting Future? A Comprehensive Survey of World Modelscited by 166
  3. 2024State Space Model for New-Generation Network Alternative to Transformers: A Surveycited by 59
  4. 2025Memory, Benchmark & Robots: A Benchmark for Solving Complex Tasks with Reinforcement Learningcited by 36
  5. 2024GenRL: Multimodal-foundation world models for generalization in embodied agentscited by 31
  6. 2024Streaming Deep Reinforcement Learning Finally Workscited by 25
  7. 2024Making Large Language Models into World Models with Precondition and Effect Knowledgecited by 22
  8. 2024Slot State Space Modelscited by 17
  9. 2025SEEA-R1: Tree-Structured Reinforcement Fine-Tuning for Self-Evolving Embodied Agentscited by 15
  10. 2024Open-World Reinforcement Learning over Long Short-Term Imaginationcited by 13
  11. 2024UniZero: Generalized and Efficient Planning with Scalable Latent World Modelscited by 12
  12. 2025Episodic memory in AI agents poses risks that should be studied and mitigatedcited by 11
26 more
  1. 2024Rethinking Transformers in Solving POMDPscited by 9
  2. 2024KalMamba: Towards Efficient Probabilistic State Space Models for RL under Uncertaintycited by 9
  3. 2024Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learningcited by 9
  4. 2025medDreamer: Model-Based Reinforcement Learning with Latent Imagination on Complex EHRs for Clinical Decision Supportcited by 9
  5. 2024Is Mamba Compatible with Trajectory Optimization in Offline Reinforcement Learning?cited by 7
  6. 2026Paying Less Generalization Tax: A Cross-Domain Generalization Study of RL Training for LLM Agentscited by 6
  7. 2024Uncertainty Representations in State-Space Layers for Deep Reinforcement Learning under Partial Observabilitycited by 5
  8. 2024Unraveling the Complexity of Memory in RL Agents: an Approach for Classification and Evaluationcited by 5
  9. 2026A Definition and Roadmap for World Modelscited by 5
  10. 2024Drama: Mamba-Enabled Model-Based Reinforcement Learning Is Sample and Parameter Efficientcited by 4
  11. 2024FACTS: A Factored State-Space Framework For World Modellingcited by 4
  12. 2025EDELINE: Enhancing Memory in Diffusion-based World Models via Linear-Time Sequence Modelingcited by 3
  13. 2025The Yokai Learning Environment: Tracking Beliefs Over Space and Timecited by 3
  14. 2025Context and Diversity Matter: The Emergence of In-Context Learning in World Modelscited by 3
  15. 2025Beyond Sliding Windows: Learning to Manage Memory in Non-Markovian Environmentscited by 3
  16. 2026MetaMind: General and Cognitive World Models in Multi-Agent Systems by Meta-Theory of Mindcited by 3
  17. 2026Latent State Design for World Models under Sufficiency Constraintscited by 3
  18. 2026Composition of Memory Experts for Diffusion World Modelscited by 3
  19. 2025On Memory: A comparison of memory mechanisms in world modelscited by 2
  20. 2025Predictive Learning in Energy-based Models with Attractor Structurescited by 1
  21. 2024Adaptive World Models: Learning Behaviors by Latent Imagination Under Non-Stationaritycited by 0
  22. 2025Bi-directional Recurrence Improves Transformer in Partially Observable Markov Decision Processescited by 0
  23. 2025Vector Quantization in the Brain: Grid-like Codes in World Modelscited by 0
  24. 2026Differentiable Clone-Structured Causal Graphs for End-to-End Cognitive Map Learning from Image Sequencescited by 0
  25. 2026Reinformed Dreamer: An Asymmetric World Model Efficiently Trained through Latent Guidancecited by 0
  26. 2026Reinforcement Learning with Complex (valued) Memoriescited by 0

Similar, not cited

Close in wording, no citation either way

  1. 2026Memory Retention Is Not Enough to Master Memory Tasks in Reinforcement Learning
  2. 2023Facing Off World Model Backbones: RNNs, Transformers, and S4
  3. 2026The World Model Remembers, the Actor Forgets: Dream Rehearsal for Continual Model-Based RL
  4. 2024Augmenting Replay in World Models for Continual Reinforcement Learning
  5. 2026Memory-R2: Fair Credit Assignment for Long-Horizon Memory-Augmented LLM Agents
  6. 2026Coding Agent Memory Post-training: Unlocking the Memory Potential of Pre-trained File Operations for Long-Horizon Tasks via Reinforcement Learning
  7. 2026Agentic Memory: Learning Unified Long-Term and Short-Term Memory Management for Large Language Model Agents
  8. 2025Memory as Action: Autonomous Context Curation for Long-Horizon Agentic Tasks