Research map / Reinforcement learning / Visual and representation RL
Mastering Memory Tasks with World Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces Recall to Imagine (R2I), which replaces DreamerV3's GRU-based RSSM with a parallel-scan structured state space model (S3M) to improve long-term memory and credit assignment in model-based RL, reporting state-of-the-art on BSuite, POPGym, and Memory Maze with ~9x faster training.
Problem
Model-based RL agents struggle with long-term dependencies needed for memory and credit assignment tasks, and RNN- and Transformer-based world models are limited by vanishing gradients or quadratic cost on long sequences.
Approach
Integrate a diagonal MIMO state space model (S4-style SSM) computed with a parallel scan, with resettable hidden states, into DreamerV3's world model, and make the representation model non-recurrent so posterior states can be computed in parallel; the policy is conditioned on the SSM hidden state and/or output state.
Key results
- R2I achieves superhuman performance in Memory Maze 9x9, 11x11 and 13x13 mazes, establishing new state of the art over DreamerV3 and IMPALA.
- R2I outperforms DreamerV3 and model-free POPGym baselines on Autoencode, RepeatPrevious and Concentration tasks, and tracks memory up to roughly 100 steps on BSuite.
- R2I computes up to 9x faster than DreamerV3 in wall-clock time while matching DreamerV3 on Atari and DMC.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace DreamerV3's recurrent sequence model with S3M | DreamerV3 · BSuite Memory Length and Discounting Chain | success rate | Helped |
| Replace DreamerV3's recurrent sequence model with S3M | DreamerV3 · POPGym Autoencode Easy and Medium; RepeatPrevious Medium and Hard | performance | Helped |
| Replace DreamerV3's recurrent sequence model with S3M | DreamerV3 · POPGym Concentration Easy and Medium | performance | Mixed |
| Replace DreamerV3's recurrent sequence model with S3M | DreamerV3 · Memory Maze 9x9 | performance | No clear effect |
| Replace DreamerV3's recurrent sequence model with S3M | DreamerV3 · Memory Maze 11x11, 13x13, and 15x15 | performance | Helped |
| Replace DreamerV3's recurrent sequence model with S3M | DreamerV3 · Atari 100K and DMC | average performance | No clear effect |
| Remove the recurrent connection from the sequence model to the representation model (ablation) | R2I · Memory Maze | performance | Helped |
| Remove the recurrent connection from the sequence model to the representation model (ablation) | R2I · Memory Maze with probing head | performance | No clear effect |
| Remove the recurrent connection from the sequence model to the representation model (ablation) | R2I · standard RL benchmarks | performance | No clear effect |
| Condition the policy on hidden state x_t instead of output state h_t (ablation) | R2I · memory-intensive environments | performance | Helped |
| Condition the policy on output state h_t instead of hidden state x_t (ablation) | R2I · memory-intensive environments | policy learning stability | Hurt |
| Condition the policy on both output state h_t and hidden state x_t (ablation) | R2I · memory-intensive environments | policy learning stability | Hurt |
Surprising findings
- Conditioning the policy on the full state (including both h_t and x_t) can hurt performance due to non-stationary feature distributions, so the hidden-state policy is preferred in memory environments.
- Removing the recurrent connection from the sequence model to the representation model does not hurt, and even helps, memory performance.
Limitations
- Sequence length in world model training batches and imagination horizon are not extremely long.
- Does not yet combine SSMs with attention mechanisms (hybrid architectures left for future work).
- Hyperparameter tuning for POPGym was done only on network sizes for both R2I and DreamerV3.
Experimental setup
Datasets: BSuite POPGym Memory Maze Atari 100K DMC
Benchmarks: BSuite POPGym Memory Maze Atari 100K DMC
Metrics: success rate return wall-clock time
Baselines: DreamerV3 IMPALA PPO+GRU PPO+LSTM PPO+MLP PPO+PosMLP
Builds on: DreamerV3 S4 S5 Recurrent State-Space Model
Closest papers by meaning
- Stable Hadamard Memory: Revitalizing Memory-Augmented Agents for Reinforcement Learning (ICLR 2025 · similarity 0.88)
- Mem-T: Densifying Rewards for Long-Horizon Memory Agents (ICML 2026 · similarity 0.87)
- AMA-Bench: Evaluating Long-Horizon Memory for Agentic Applications (ICML 2026 · similarity 0.87)
- MemRL: Self-Evolving Agents via Runtime Reinforcement Learning on Episodic Memory (NeurIPS 2026 · similarity 0.87)
- Recurrent Reinforcement Learning with Memoroids (NeurIPS 2024 · similarity 0.87)
- Meta-Cognitive Memory Policy Optimization for Long-Horizon LLM Agents (NeurIPS 2026 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.