# Procedural Generalization by Planning with Self-Supervised World Models

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

## Key facts

- Cited by 31 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 66 of its 78 references that are in the map.
- First posted on arXiv 2 Nov 2021.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)

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

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Representation Learning with Contrastive Predictive Coding](https://mlsanity.com/paper/1807.03748/) (2018)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Gradient Surgery for Multi-Task Learning](https://mlsanity.com/paper/2001.06782/) (2020)
- [Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning](https://mlsanity.com/paper/1910.10897/) (2019)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [Learning Dexterous In-Hand Manipulation](https://mlsanity.com/paper/1808.00177/) (2018)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Learning to reinforcement learn](https://mlsanity.com/paper/1611.05763/) (2016)
- [Implicit Quantile Networks for Distributional Reinforcement Learning](https://mlsanity.com/paper/1806.06923/) (2018)
- [Planning to Explore via Self-Supervised World Models](https://mlsanity.com/paper/2005.05960/) (2020)
- [Learning Deconvolution Network for Semantic Segmentation](https://mlsanity.com/paper/1505.04366/) (2015)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [Data-Efficient Reinforcement Learning with Self-Predictive Representations](https://mlsanity.com/paper/2007.05929/) (2020)
- [Deep Visual Foresight for Planning Robot Motion](https://mlsanity.com/paper/1610.00696/) (2016)
- [Learning to Adapt in Dynamic, Real-World Environments Through Meta-Reinforcement Learning](https://mlsanity.com/paper/1803.11347/) (2018)
- [Contrastive Learning of Structured World Models](https://mlsanity.com/paper/1911.12247/) (2019)
- [Representation Learning via Invariant Causal Mechanisms](https://mlsanity.com/paper/2010.07922/) (2020)
- [Prioritized Level Replay](https://mlsanity.com/paper/2010.03934/) (2020)

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

- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [CycleResearcher: Improving Automated Research via Automated Review](https://mlsanity.com/paper/2411.00816/) (2024)
- [What is Essential for Unseen Goal Generalization of Offline Goal-conditioned RL?](https://mlsanity.com/paper/2305.18882/) (2023)
- [Grounding Multimodal Large Language Models in Actions](https://mlsanity.com/paper/2406.07904/) (2024)
- [Efficient World Models with Context-Aware Tokenization](https://mlsanity.com/paper/2406.19320/) (2024)
- [On the Importance of Exploration for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2306.05483/) (2023)
- [The Generalization Gap in Offline Reinforcement Learning](https://mlsanity.com/paper/2312.05742/) (2023)
- [HarmonyDream: Task Harmonization Inside World Models](https://mlsanity.com/paper/2310.00344/) (2023)
- [On the Effectiveness of Fine-tuning Versus Meta-reinforcement Learning](https://mlsanity.com/paper/2206.03271/) (2022)
- [The Benefits of Model-Based Generalization in Reinforcement Learning](https://mlsanity.com/paper/2211.02222/) (2022)
- [On the link between conscious function and general intelligence in humans and machines](https://mlsanity.com/paper/2204.05133/) (2022)
- [Dreaming of Many Worlds: Learning Contextual World Models Aids Zero-Shot Generalization](https://mlsanity.com/paper/2403.10967/) (2024)
- [AMAGO-2: Breaking the Multi-Task Barrier in Meta-Reinforcement Learning with Transformers](https://mlsanity.com/paper/2411.11188/) (2024)
- [Distilling Reinforcement Learning Algorithms for In-Context Model-Based Planning](https://mlsanity.com/paper/2502.19009/) (2025)
- [Value Memory Graph: A Graph-Structured World Model for Offline Reinforcement Learning](https://mlsanity.com/paper/2206.04384/) (2022)
- [Learning General World Models in a Handful of Reward-Free Deployments](https://mlsanity.com/paper/2210.12719/) (2022)
- [Vintix: Action Model via In-Context Reinforcement Learning](https://mlsanity.com/paper/2501.19400/) (2025)
- [Learning Robust Real-Time Cultural Transmission without Human Data](https://mlsanity.com/paper/2203.00715/) (2022)
- [Rethinking Value Function Learning for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2210.09960/) (2022)
- [Composing Task Knowledge with Modular Successor Feature Approximators](https://mlsanity.com/paper/2301.12305/) (2023)
- [Discovering Hierarchical Achievements in Reinforcement Learning via Contrastive Learning](https://mlsanity.com/paper/2307.03486/) (2023)
- [CompilerDream: Learning a Compiler World Model for General Code Optimization](https://mlsanity.com/paper/2404.16077/) (2024)
- [A Look at Value-Based Decision-Time vs. Background Planning Methods Across Different Settings](https://mlsanity.com/paper/2206.08442/) (2022)
- [Investigating the role of model-based learning in exploration and transfer](https://mlsanity.com/paper/2302.04009/) (2023)
- [Context and Diversity Matter: The Emergence of In-Context Learning in World Models](https://mlsanity.com/paper/2509.22353/) (2025)
- [Contrastive Abstraction for Reinforcement Learning](https://mlsanity.com/paper/2410.00704/) (2024)
- [A Game-Theoretic Perspective of Generalization in Reinforcement Learning](https://mlsanity.com/paper/2208.03650/) (2022)
- [Analogy making as amortised model construction](https://mlsanity.com/paper/2507.16511/) (2025)
- [Vintix II: Decision Pre-Trained Transformer is a Scalable In-Context Reinforcement Learner](https://mlsanity.com/paper/2604.05112/) (2026)
- [Direct Advantage Estimation for Scalable and Sample-efficient Deep Reinforcement Learning](https://mlsanity.com/paper/2606.20411/) (2026)

## Abstract (arXiv)

One of the key promises of model-based reinforcement learning is the ability to generalize using an internal model of the world to make predictions in novel environments and tasks. However, the generalization ability of model-based agents is not well understood because existing work has focused on model-free agents when benchmarking generalization. Here, we explicitly measure the generalization ability of model-based agents in comparison to their model-free counterparts. We focus our analysis on MuZero (Schrittwieser et al., 2020), a powerful model-based agent, and evaluate its performance on both procedural and task generalization. We identify three factors of procedural generalization -- planning, self-supervised representation learning, and procedural data diversity -- and show that by combining these techniques, we achieve state-of-the art generalization performance and data efficiency on Procgen (Cobbe et al., 2019). However, we find that these factors do not always provide the same benefits for the task generalization benchmarks in Meta-World (Yu et al., 2019), indicating that transfer remains a challenge and may require different approaches than procedural generalization. Overall, we suggest that building generalizable agents requires moving beyond the single-task, model-free paradigm and towards self-supervised model-based agents that are trained in rich, procedural, multi-task environments.