Procedural Generalization by Planning with Self-Supervised World Models
Open accessAbstract
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.
Integrity signals
Registry
Cites
Statistics
Claims the paper makes
Claims other papers make about this one
Each is checked against this paper’s own text: does it say what the citing paper says it does?
Generalising to new rewards (as is necessary for Meta-World) is not currently as amenable to MBRL approaches.
“although the performance there is not as impressive, showing that generalising to new rewards (as is necessary for Meta-World) is not currently as amenable to MBRL approaches.”
From A Survey of Zero-shot Generalisation in Deep Reinforcement Learning (2023) · cites this paper as Anand et al. (2021)
Very little work has applied MBRL to ZSG benchmarks, with Anand et al. (2021) being an initial example.
“Very little work has applied MBRL to ZSG benchmarks, with Anand et al. (2021) being an initial example.”
From A Survey of Zero-shot Generalisation in Deep Reinforcement Learning (2023) · cites this paper as Anand et al. (2021)
Anand et al. (2021) is the first example investigating how well standard model-based reinforcement learning approaches generalise.
“Anand et al. (2021) is the first example investigating how well standard MBRL approaches generalise”
From A Survey of Zero-shot Generalisation in Deep Reinforcement Learning (2023) · cites this paper as Anand et al., 2021
Cites
Learning Transferable Visual Models From Natural Language Supervision
Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al. (2021) · arXiv · doi:10.48550/arxiv.2103.00020 · Cited by 15,293 · see citing works
Proximal Policy Optimization Algorithms
Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017) · arXiv · doi:10.48550/arxiv.1707.06347 · Cited by 10,625 · see citing works
Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks
Finn, C., Abbeel, P., & Levine, S. (2017) · arXiv · doi:10.48550/arxiv.1703.03400 · Cited by 8,512 · see citing works
Representation Learning with Contrastive Predictive Coding
Oord, A. V. D., Li, Y., & Vinyals, O. (2018) · arXiv · doi:10.48550/arxiv.1807.03748 · Cited by 6,242 · see citing works
Self-citationLearning Deconvolution Network for Semantic Segmentation
Noh, H., Hong, S., & Han, B. (2015) · 2015 IEEE International Conference on Computer Vision (ICCV) · doi:10.1109/iccv.2015.178 · Cited by 4,422 · see citing works
Bandit Based Monte-Carlo Planning
Kocsis, L., & Szepesvari, C. (2006) · European Conference on Machine Learning · doi:10.1007/11871842_29 · Cited by 4,073 · see citing works
A general reinforcement learning algorithm that masters chess, shogi, and Go through self-play
Silver, D., Hubert, T., Schrittwieser, J., Antonoglou, I., Lai, M., Guez, A., Lanctot, M., Sifre, L., Kumaran, D., Graepel, T., et al. (2018) · Science · doi:10.1126/science.aar6404 · Cited by 3,562 · cited 1× · see citing works
“MuZero is an appealing candidate for investigating generalization because it already incorporates many ideas that are thought to be useful for generalization and transfer, including replay [61], planning [65], and model-based representation learning [22, 51].”
BackgroundSelf-citationDomain randomization for transferring deep neural networks from simulation to the real world
Tobin, J., Fong, R., Ray, A., Schneider, J., Zaremba, W., & Abbeel, P. (2017) · 2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS) · doi:10.1109/iros.2017.8202133 · Cited by 3,124 · cited 2× · see citing works
“In robotics, collecting and training over a diverse range of data has Preprint been found to be critical particularly in the sim2real literature where the goal is to transfer a policy learned in simulation to the real world [3, 67].”
BackgroundBuilding machines that learn and think like people
Lake, B. M., Ullman, T. D., Tenenbaum, J. B., & Gershman, S. J. (2017) · Behavioral and Brain Sciences · doi:10.1017/s0140525x16001837 · Cited by 2,550 · see citing works
Efficient Selectivity and Backup Operators in Monte-Carlo Tree Search
Coulom, R. (2006) · Computers and Games · doi:10.1007/978-3-540-75538-8_7 · Cited by 1,819 · see citing works
Learning dexterous in-hand manipulation
Andrychowicz, O. M., Baker, B., Chociej, M., Józefowicz, R., McGrew, B., Pachocki, J., Petron, A., Plappert, M., Powell, G., Ray, A., et al. (2020) · The International Journal of Robotics Research · doi:10.1177/0278364919887447 · Cited by 1,643 · see citing works
Recurrent World Models Facilitate Policy EvolutionPreprint
Ha, D. R., & Schmidhuber, J. (2018) · Neural Information Processing Systems · doi:10.48550/arxiv.1809.01999 · Cited by 1,592 · see citing works
Dayan, P., Hinton, G. E., Neal, R. M., & Zemel, R. S. (1995) · Neural Computation · doi:10.1162/neco.1995.7.5.889 · Cited by 1,207 · see citing works
A simple framework for contrastive learning of visual representations
Chen, T. (2024) · TIB Data Manager · doi:10.57702/iuuvgtaz · Cited by 1,140 · see citing works
Dyna, an Integrated Architecture for Learning, Planning, and Reacting
Sutton, R. S. (1991) · ACM SIGART Bulletin · doi:10.1145/122344.122377 · Cited by 777 · cited 2× · see citing works
“Third, the model can also be used to train a policy or value function by simulating training data [66] or constructing more informative losses [21, 30], again enabling faster learning (possibly entirely in simulation).”
BackgroundMethodCurious model-building control systems
Schmidhuber, J. (1991) · [Proceedings] 1991 IEEE International Joint Conference on Neural Networks · doi:10.1109/ijcnn.1991.170605 · Cited by 733 · see citing works
Model-based Reinforcement Learning: A Survey
Moerland, T. M., Broekens, J., Plaat, A., & Jonker, C. M. (2023) · Foundations and Trends® in Machine Learning · doi:10.1561/2200000086 · Cited by 668 · see citing works
Deep visual foresight for planning robot motion
Finn, C., & Levine, S. (2017) · 2017 IEEE International Conference on Robotics and Automation (ICRA) · doi:10.1109/icra.2017.7989324 · Cited by 599 · see citing works
Learning to reinforcement learn
Wang, J. X., Kurth-Nelson, Z., Tirumala, D., Soyer, H., Leibo, J. Z., Munos, R., Blundell, C., Kumaran, D., & Botvinick, M. (2016) · arXiv · doi:10.48550/arxiv.1611.05763 · Cited by 579 · see citing works
Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning
Yu, T., Quillen, D., He, Z., Julian, R., Narayan, A., Shively, H., Bellathur, A., Hausman, K., Finn, C., & Levine, S. (2019) · arXiv · doi:10.48550/arxiv.1910.10897 · Cited by 450 · see citing works
Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context VariablesPreprint
Rakelly, K., Zhou, A., Quillen, D., Finn, C., & Levine, S. (2019) · arXiv (Cornell University) · doi:10.48550/arxiv.1903.08254 · Cited by 371 · see citing works
Learning to Adapt in Dynamic, Real-World Environments through Meta-Reinforcement Learning
Nagabandi, A., Clavera, I., Liu, S., Fearing, R. S., Abbeel, P., Levine, S., & Finn, C. (2018) · arXiv · doi:10.48550/arxiv.1803.11347 · Cited by 370 · see citing works
Learning Continuous Control Policies by Stochastic Value GradientsPreprint
Heess, N., Wayne, G., Silver, D., Lillicrap, T., Tassa, Y., & Erez, T. (2015) · arXiv (Cornell University) · doi:10.48550/arxiv.1510.09142 · Cited by 359 · see citing works
Implicit Quantile Networks for Distributional Reinforcement Learning
Dabney, W., Ostrovski, G., Silver, D., & Munos, R. (2018) · arXiv · doi:10.48550/arxiv.1806.06923 · Cited by 331 · see citing works
Reinforcement Learning with Augmented Data
Laskin, M., Lee, K., Stooke, A., Pinto, L., Abbeel, P., & Srinivas, A. (2020) · arXiv · doi:10.48550/arxiv.2004.14990 · Cited by 299 · see citing works
Mastering Atari with Discrete World Models
Hafner, D., Lillicrap, T., Norouzi, M., & Ba, J. (2020) · arXiv · doi:10.48550/arxiv.2010.02193 · Cited by 247 · see citing works
Benchmarking Model-Based Reinforcement Learning
Wang, T., Bao, X., Clavera, I., Hoang, J., Wen, Y., Langlois, E., Zhang, S., Zhang, G., Abbeel, P., & Ba, J. (2019) · arXiv · doi:10.48550/arxiv.1907.02057 · Cited by 226 · see citing works
Generalization in Reinforcement Learning with Selective Noise Injection and Information BottleneckPreprint
Igl, M., Ciosek, K., Li, Y.-Z., Tschiatschek, S., Zhang, C., Devlin, S., & Hofmann, K. (2019) · Neural Information Processing Systems · doi:10.48550/arxiv.1910.12911 · Cited by 206 · see citing works
Self-citationRelational Deep Reinforcement LearningPreprint
Zambaldi, V., Raposo, D. C., Santoro, A., Bapst, V., Li, Y., Babuschkin, I., Tuyls, K., Reichert, D., Lillicrap, T., Lockhart, E., et al. (2018) · arXiv (Cornell University) · doi:10.48550/arxiv.1806.01830 · Cited by 153 · see citing works
Self-citationImagination-Augmented Agents for Deep Reinforcement LearningPreprint
Racanière, S., Weber, T., Reichert, D., Buesing, L., Guez, A., Rezende, D. J., Badia, A. P., Vinyals, O., Heess, N., Li, Y., et al. (2017) · arXiv (Cornell University) · Cited by 144 · see citing works
Self-citationGotta Learn Fast: A New Benchmark for Generalization in RLPreprint
Nichol, A., Pfau, V., Hesse, C., Klimov, O., & Schulman, J. (2018) · arXiv (Cornell University) · doi:10.48550/arxiv.1804.03720 · Cited by 136 · see citing works
Illuminating Generalization in Deep Reinforcement Learning through Procedural Level Generation
Justesen, N., Torrado, R. R., Bontrager, P., Khalifa, A., Togelius, J., & Risi, S. (2018) · OAR@UM (University of Malta) · doi:10.48550/arxiv.1806.10729 · Cited by 117 · see citing works
On Robustness and Transferability of Convolutional Neural Networks
Djolonga, J., Yung, J., Tschannen, M., Romijnders, R., Beyer, L., Kolesnikov, A., Puigcerver, J., Minderer, M., D'Amour, A., Moldovan, D., et al. (2021) · 2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR) · doi:10.1109/cvpr46437.2021.01619 · Cited by 98 · see citing works
Analogues of mental simulation and imagination in deep learning
Hamrick, J. B. (2019) · Current Opinion in Behavioral Sciences · doi:10.1016/j.cobeha.2018.12.011 · Cited by 83 · see citing works
Self-citationCobbe, K., Hilton, J., Klimov, O., & Schulman, J. (2020) · arXiv · doi:10.48550/arxiv.2009.04416 · Cited by 83 · see citing works
Bootstrap Latent-Predictive Representations for Multitask Reinforcement LearningPreprint
Guo, Z. D., Pires, B. Á., Piot, B., Grill, J.-B., Altché, F., Munos, R., & Azar, M. G. (2020) · arXiv (Cornell University) · doi:10.48550/arxiv.2004.14646 · Cited by 78 · see citing works
Improving Generalization in Meta Reinforcement Learning using Learned ObjectivesPreprint
Kirsch, L., van Steenkiste, S., & Schmidhuber, J. (2019) · arXiv (Cornell University) · doi:10.48550/arxiv.1910.04098 · Cited by 69 · see citing works
Contrastive Learning of Structured World Models
Kipf, T., van der Pol, E., & Welling, M. (2020) · Data Archiving and Networked Services (DANS) · Cited by 63 · see citing works
Shaping Belief States with Generative Environment Models for RLPreprint
Gregor, K., Rezende, D. J., Besse, F., Wu, Y. X., Merzić, H., & van den Oord, A. (2019) · arXiv (Cornell University) · doi:10.48550/arxiv.1906.09237 · Cited by 54 · see citing works
Structured agents for physical constructionPreprint
Bapst, V., Sánchez‐González, Á., Doersch, C., Stachenfeld, K., Kohli, P., Battaglia, P., & Hamrick, J. B. (2019) · arXiv (Cornell University) · doi:10.48550/arxiv.1904.03177 · Cited by 49 · see citing works
Self-citationAutomatic Data Augmentation for Generalization in Reinforcement LearningPreprint
Răileanu, R., Goldstein, M., Yarats, D., Kostrikov, I., & Fergus, R. (2021) · arXiv (Cornell University) · Cited by 44 · see citing works
Monte-Carlo Tree Search as Regularized Policy OptimizationPreprint
Grill, J.-B., Altché, F., Tang, Y., Hubert, T., Valko, M., Antonoglou, I., & Munos, R. (2020) · arXiv (Cornell University) · doi:10.48550/arxiv.2007.12509 · Cited by 42 · see citing works
Neural Predictive Belief RepresentationsPreprint
Guo, Z. D., Azar, M. G., Piot, B., Pires, B. Á., & Munos, R. (2018) · arXiv (Cornell University) · doi:10.48550/arxiv.1811.06407 · Cited by 40 · see citing works
Leveraging Procedural Generation to Benchmark Reinforcement Learning
Cobbe, K., Hesse, C., Hilton, J., & Schulman, J. (2020) · International Conference on Machine Learning · Cited by 40 · see citing works
Composable Planning with AttributesPreprint
Zhang, A., Lerer, A. K., Sukhbaatar, S., Fergus, R., & Szlam, A. D. (2018) · arXiv (Cornell University) · doi:10.48550/arxiv.1803.00512 · Cited by 38 · see citing works
Natural Environment Benchmarks for Reinforcement LearningPreprint
Zhang, A., Wu, Y., & Pineau, J. (2018) · arXiv (Cornell University) · doi:10.48550/arxiv.1811.06032 · Cited by 36 · see citing works
Measuring and characterizing generalization in deep reinforcement learning
Witty, S., Lee, J. K., Tosch, E., Atrey, A., Clary, K., Littman, M. L., & Jensen, D. D. (2021) · Applied AI Letters · doi:10.1002/ail2.45 · Cited by 35 · see citing works
On the role of planning in model-based deep reinforcement learningPreprint
Hamrick, J. B., Friesen, A. L., Behbahani, F., Guez, A., Viola, F., Witherspoon, S., Anthony, T., Buesing, L., Veličković, P., & Weber, T. (2020) · arXiv (Cornell University) · doi:10.48550/arxiv.2011.04021 · Cited by 34 · see citing works
Self-citationOnline and Offline Reinforcement Learning by Planning with a Learned ModelPreprint
Schrittwieser, J., Hubert, T., Mandhane, A., Barekatain, M., Antonoglou, I., & Silver, D. (2021) · arXiv (Cornell University) · doi:10.48550/arxiv.2104.06294 · Cited by 33 · see citing works
Self-citationJiang, M., Grefenstette, E., & Rocktäschel, T. (2020) · arXiv · doi:10.48550/arxiv.2010.03934 · Cited by 33 · see citing works
Learning and Planning in Complex Action SpacesPreprint
Hubert, T., Schrittwieser, J., Antonoglou, I., Barekatain, M., Schmitt, S., & Silver, D. (2021) · arXiv (Cornell University) · doi:10.48550/arxiv.2104.06303 · Cited by 33 · see citing works
Self-citationDeep Reinforcement and InfoMax LearningPreprint
Mazoure, B., Combes, R. T. D., Doan, T., Bachman, P., & Hjelm, R. D. (2020) · arXiv (Cornell University) · doi:10.48550/arxiv.2006.07217 · Cited by 32 · see citing works
CoBERL: Contrastive BERT for Reinforcement LearningPreprint
Banino, A., Badia, A. P., Walker, J., Scholtes, T., Mitrovic, J., & Blundell, C. (2021) · arXiv (Cornell University) · doi:10.48550/arxiv.2107.05431 · Cited by 20 · see citing works
Self-citationCausally Correct Partial Models for Reinforcement LearningPreprint
Rezende, D. J., Danihelka, I., Papamakarios, G., Ke, N. R., Jiang, R., Weber, T., Gregor, K., Merzić, H., Viola, F., Wang, J., et al. (2020) · arXiv (Cornell University) · doi:10.48550/arxiv.2002.02836 · Cited by 16 · see citing works
Self-citationMetaCURE: Meta Reinforcement Learning with Empowerment-Driven ExplorationPreprint
Zhang, J., Wang, J., Hu, H., Chen, T., Chen, Y., Fan, C., & Zhang, C. (2020) · arXiv (Cornell University) · doi:10.48550/arxiv.2006.08170 · Cited by 15 · see citing works
Improving Generalization in Meta Reinforcement Learning using Neural ObjectivesPreprint
Kirsch, L., van Steenkiste, S., & Schmidhuber, J. (2020) · arXiv (Cornell University) · Cited by 14 · see citing works
Cross-Trajectory Representation Learning for Zero-Shot Generalization in RLPreprint
Mazoure, B., Ahmed, A. M., MacAlpine, P., Hjelm, R. D., & Kolobov, A. (2021) · arXiv (Cornell University) · doi:10.48550/arxiv.2106.02193 · Cited by 9 · see citing works
Representation Learning for Out-Of-Distribution Generalization in Reinforcement LearningPreprint
Dittadi, A., Träuble, F., Wüthrich, M., Widmaier, F., Gehler, P., Winther, O., Locatello, F., Bachem, O., Schölkopf, B., & Bauer, S. (2021) · arXiv (Cornell University) · Cited by 4 · see citing works
Gradient Surgery for Multi-Task Learning
Yu, T., Kumar, S., Gupta, A., Levine, S., Hausman, K., & Finn, C. (2020) · arXiv · doi:10.48550/arxiv.2001.06782
Flagged first, then by citations
16 of the 75 references couldn’t be matched to a known paper.
Cited by
Sorted by citations
A Survey of Zero-shot Generalisation in Deep Reinforcement Learning
Kirk, R., Zhang, A., Grefenstette, E., & Rocktäschel, T. (2023) · Journal of Artificial Intelligence Research · doi:10.1613/jair.1.14174 · Cited by 193 · see citing works
“Anand et al. (2021) is the first example investigating how well standard MBRL approaches generalise, and we look forward to seeing more work in this area.”
BackgroundMethodInfluentialA Survey of Generalisation in Deep Reinforcement LearningPreprint
Kirk, R., Zhang, A., Grefenstette, E., & Rocktäschel, T. (2021) · arXiv (Cornell University) · Cited by 63 · see citing works
World and Human Action Models towards gameplay ideation
Kanervisto, A., Bignell, D., Wen, L., Grayson, M., Georgescu, R., Macua, S. V., Tan, S. Z., Rashid, T., Pearce, T., Cao, Y., et al. (2025) · Nature · doi:10.1038/s41586-025-08600-3 · Cited by 25 · see citing works
Learning few-shot imitation as cultural transmission
Bhoopchand, A., Brownfield, B., Collister, A., Dal Lago, A., Edwards, A., Everett, R., Fréchette, A., Oliveira, Y. G., Hughes, E., Mathewson, K. W., et al. (2023) · Nature Communications · doi:10.1038/s41467-023-42875-2 · Cited by 14 · see citing works
“To address this, we measure cultural transmission in held-out test tasks and with human expert demonstrators 40,41 , similar to the familiar train-test dataset split in supervised learning 42 .”
Mutschler, C., Kontes, G. D., & Rietsch, S. (2024) · Springer eBooks · doi:10.1007/978-3-031-64832-8_3
- Authors
- Anand, A., Walker, J., Li, Y., Vértes, E., Schrittwieser, J., Ozair, S., Weber, T., & Hamrick, J. B.
- Affiliations
- Université de Montréal (CA)
- Venue
- arXiv
- Publisher
- arXiv
- Published
- 2021
- ISSN
- 2331-8422 (ISSN-L)
- arXiv
- 2111.01587
- OpenAlex
- W3209153580
- Cited by
- 5
- References
- 75
- Citations by year
- 2025: 1 · 2024: 1 · 2023: 2 · 2021: 1
- Type
- posted content
- Language
- en
- Open access
- green
- arxiv.org
- Locations
- arXiv (Cornell University) (submittedVersion)
- arXiv (Cornell University) (cc-by)
- Topics
- Physical Sciences › Computer Science › Artificial Intelligence › Reinforcement Learning in Robotics
- Physical Sciences › Computer Science › Artificial Intelligence › Domain Adaptation and Few-Shot Learning
- Physical Sciences › Computer Science › Artificial Intelligence › Artificial Intelligence in Games
- Keywords
- model-based reinforcement learning, MuZero, planning, task generalization, transfer learning, Procgen, Meta-World, representation learning, data efficiency
About these checks
Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.
Statistics and arithmetic are recomputed, not judged.
Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.
What counts as a claim ›How we measure confidence ›How we check this ›
Are you an author, or do you see a problem with this page? Request a review
