CoChat
Sign Up
PreprintCited by 2Reinforcement Learning in RoboticsDomain Adaptation and Few-Shot LearningAI-based Problem Solving and Planning

Dreaming of Many Worlds: Learning Contextual World Models Aids Zero-Shot Generalization

Open access
Prasanna, Sai, Farid, Karim, Rajan, Raghu, Biedenkapp, André·2024·arXiv
DOI 10.48550/arxiv.2403.10967arXiv 2403.10967
Open this paper in CoChat — free accountOpen PDF

Abstract

Zero-shot generalization (ZSG) to unseen dynamics is a major challenge for creating generally capable embodied agents. To address the broader challenge, we start with the simpler setting of contextual reinforcement learning (cRL), assuming observability of the context values that parameterize the variation in the system's dynamics, such as the mass or dimensions of a robot, without making further simplifying assumptions about the observability of the Markovian state. Toward the goal of ZSG to unseen variation in context, we propose the contextual recurrent state-space model (cRSSM), which introduces changes to the world model of Dreamer (v3) (Hafner et al., 2023). This allows the world model to incorporate context for inferring latent Markovian states from the observations and modeling the latent dynamics. Our approach is evaluated on two tasks from the CARL benchmark suite, which is tailored to study contextual RL. Our experiments show that such systematic incorporation of the context improves the ZSG of the policies trained on the "dreams" of the world model. We further find qualitatively that our approach allows Dreamer to disentangle the latent state from context, allowing it to extrapolate its dreams to the many worlds of unseen contexts. The code for all our experiments is available at https://github.com/sai-prasanna/dreaming_of_many_worlds.

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2403.10967

Cites

Statistics

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

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?

  • DreamerV3 struggles to generalize across diverse latent contexts due to its limited ability to retain critical environmental information.

    “Unlike DreamerV3, which struggles to generalize across diverse latent contexts due to its limited ability to retain critical environmental information [Prasanna et al., 2024]”

    From Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization (2025) · cites this paper as Prasanna et al., 2024

    Section 5.1.2 (page unavailable) reports that explicit context conditioning often improves zero-shot generalization over the hidden-context Dreamer baseline, though results vary by task and observation modality. Section 4.2.1 (page unavailable) says that if context is not retained in the latent state, generalization could be difficult; it does not establish that DreamerV3 generally fails to retain critical environmental information or that this is the demonstrated cause of its performance. (LLM, conf=0.85)

    Experimental — accuracy not yet measured

  • Existing contextual RL methods often require explicit context variables, limiting their use when contexts are latent or hard to measure.

    “Contextual Markov Decision Processes (cMDP) model this challenge, but existing methods often require explicit context variables (e.g., friction, gravity), limiting their use when contexts are latent or hard to measure.”

    From Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization (2025) · cites this paper as Prasanna et al., 2024

    Section 4.2.1 describes supplying an inferred or true context variable as a predominant approach in model-free contextual RL, while Section 4.2.2 contrasts an approach that does not provide context. This supports the concern about relying on available context variables, but the paper does not explicitly say that existing methods are limited when context is hard to measure. Its 'Importance of Context' example also treats friction as a latent state, not a context variable. (LLM, conf=0.85)

    Experimental — accuracy not yet measured

  • DALI's context encoder is essential for efficient context inference and robust generalization.

    “We theoretically prove this encoder is essential for efficient context inference and robust generalization.”

    From Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization (2025) · cites this paper as Prasanna et al., 2024

    Section 4.2.2 (page unavailable) suggests that Dreamer’s observation encoder could infer context more efficiently, and Section 5.1.2 reports that explicit context conditioning aids zero-shot generalization. The paper does not establish that DALI’s context encoder is essential for either outcome, or theoretically prove that necessity. (LLM, conf=0.85)

    Experimental — accuracy not yet measured

  • Explicit context conditioning approaches scale poorly and break down in unstructured environments where contexts are latent or difficult to measure.

    “While effective in controlled settings, such approaches scale poorly and break down in unstructured environments where contexts are latent or difficult to measure.”

    From Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization (2025) · cites this paper as Prasanna et al., 2024

    Section 5.1.2 reports that explicit context conditioning improves performance in the evaluated settings, and ‘Zero-Shot Generalization in MBRL’ says the method assumes observable context. The paper does not show that these approaches scale poorly or break down in unstructured environments when context is latent or difficult to measure. The indexed passages provide no page numbers. (LLM, conf=0.78)

    Experimental — accuracy not yet measured

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Cites

No references could be resolved yet.

Research works printed but not linked: 47

  1. R. Agarwal, M. Schwarzer, P. Samuel Castro, A. C. Courville, and M. G. Bellemare. Deep reinforce- ment learning at the edge of the statistical precipice. In M. Ranzato, A. Beygelzimer, K. Nguyen, P. Liang, J. Vaughan, and Y. Dauphin (eds.), Proceedings of the 35th International Conference on Advances in Neural Information Processing Systems (NeurIPS'21).
  2. Ankesh Anand, Jacob C Walker, Yazhe Li, Eszter Vértes, Julian Schrittwieser, Sherjil Ozair, Theo- phane Weber, and Jessica B Hamrick. Procedural generalization by planning with self-supervised world models. In International Conference on Learning Representations, 2021.
  3. M. Andrychowicz, B. Baker, M. Chociej, R. Józefowicz, B. McGrew, J. Pachocki, A. Petron, M. Plap- pert, G. Powell, A. Ray, J. Schneider, S. Sidor, J. Tobin, P. Welinder, L. Weng, and W. Zaremba. Learning dexterous in-hand manipulation. International Journal of Robotics Research, 39(1), 2020.
  4. P. J. Ball, C. Lu, J. Parker-Holder, and S. Roberts. Augmented world models facilitate zero-shot dynamics generalization from a single offline environment. In M. Meila and T. Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning (ICML'21), volume 139 of Proceedings of Machine Learning Research, pp. 619-629. PMLR, 2021.
  5. A. G. Barto, R. S. Sutton, and C. W. Anderson. Neuronlike adaptive elements that can solve difficult learning control problems. IEEE Transactions on Systems, Man, and Cybernetics, SMC-13:834- 846, 1983.
  6. Charlie Beattie, Joel Z. Leibo, Denis Teplyashin, Tom Ward, Marcus Wainwright, Heinrich Küttler, Andrew Lefrancq, Simon Green, Víctor Valdés, Amir Sadik, Julian Schrittwieser, Keith Anderson, Sarah York, Max Cant, Adam Cain, Adrian Bolton, Stephen Gaffney, Helen King, Demis Hassabis, Shane Legg, and Stig Petersen. Deepmind lab. ArXiv, abs/1612.03801, 2016. URL https: //api.semanticscholar.org/CorpusID:3221395. · arXiv
  7. J. Beck, R. Vuorio, E. Z. Liu, Z. Xiong, L. Zintgraf, C. Finn, and S. Whiteson. A survey of meta- reinforcement learning. arXiv:2301.08028 [cs.LG], 2023. · arXiv
  8. Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling. The arcade learning environ- ment: An evaluation platform for general agents. Journal of Artificial Intelligence Research, 47: 253-279, 2013.
  9. C. Benjamins, T. Eimer, F. Schubert, A. Mohan, S. Döhler, A. Biedenkapp, B. Rosenhan, F. Hutter, and M. Lindauer. Contextualize me -the case for context in reinforcement learning. Transactions on Machine Learning Research, 2023.
  10. M. Beukman, D. Jarvis, R. Klein, S. James, and B. Rosman. Dynamics generalisation in rein- forcement learning via adaptive context-aware policies. In Proceedings of the 37th International Conference on Advances in Neural Information Processing Systems (NeurIPS'23).
  11. A. Biedenkapp, D. Speck, S. Sievers, F. Hutter, M. Lindauer, and J. Seipp. Learning domain- independent policies for open list selection. In M. Katz, H. Palacios, and V. Gómez (eds.), Work- shop on Bridging the Gap Between AI Planning and Reinforcement Learning (PRL@ICAPS'22), 2022.
  12. K. Chua, R. Calandra, R. McAllister, and S. Levine. Deep reinforcement learning in a handful of trials using probabilistic dynamics models. In S. Bengio, H. Wallach, H. Larochelle, K. Grau- man, N. Cesa-Bianchi, and R. Garnett (eds.), Proceedings of the 31st International Conference on Advances in Neural Information Processing Systems (NeurIPS'18), pp. 4759-4770.
  13. S. Du, A. Krishnamurthy, N. Jiang, A. Agarwal, M. Dudik, and J. Langford. Provably efficient RL with rich observations via latent state decoding. In K. Chaudhuri and R. Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning (ICML'19), volume 97, pp. 1665-1674. Proceedings of Machine Learning Research, 2019.
  14. Y. Duan, J. Schulman, X. Chen, P. Bartlett, I. Sutskever, and P. Abbeel. RL$ˆ2$: Fast reinforcement learning via slow reinforcement learning. arXiv:1611.02779 [cs.AI], 2016. · arXiv
  15. A. Escontrela, G. Yu, P. Xu, A. Iscen, and J. Tan. Zero-shot terrain generalization for visual locomotion policies. arXiv:2011.05513 [cs.RO], 2020. · arXiv
  16. J. Guo, M. Gong, and D. Tao. A relational intervention approach for unsupervised dynamics gen- eralization in model-based reinforcement learning. In Proceedings of the International Conference on Learning Representations (ICLR'22), 2022. Published online: iclr.cc.
  17. William H Guss, Cayden Codel, Katja Hofmann, Brandon Houghton, Noboru Kuno, Stephanie Milani, Sharada Mohanty, Diego Perez Liebana, Ruslan Salakhutdinov, Nicholay Topin, et al. The minerl 2019 competition on sample efficient reinforcement learning using human priors. arXiv preprint arXiv:1904.10079, 2019. · arXiv
  18. T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In J. Dy and A. Krause (eds.), Proceedings of the 35th International Conference on Machine Learning (ICML'18), volume 80. Proceedings of Machine Learning Research, 2018.
  19. D. Hafner, T. Lillicrap, I. Fischer, R. Villegas, D. Ha, H. Lee, and J. Davidson. Learning latent dynamics for planning from pixels. In K. Chaudhuri and R. Salakhutdinov (eds.), Proceedings of the 36th International Conference on Machine Learning (ICML'19), volume 97, pp. 2555-2565. Proceedings of Machine Learning Research, 2019.
  20. D. Hafner, T. P. Lillicrap, J. Ba, and M. Norouzi. Dream to control: Learning behaviors by latent imagination. In Proceedings of the International Conference on Learning Representations (ICLR'20), 2020. Published online: iclr.cc.
  21. D. Hafner, T. P. Lillicrap, M. Norouzi, and J. Ba. Mastering atari with discrete world models. In Proceedings of the International Conference on Learning Representations (ICLR'21), 2021. Published online: iclr.cc.
  22. D. Hafner, J. Pasukonis, J. Ba, and T. P. Lillicrap. Mastering diverse domains through world models. arXiv:2301.04104 [cs.AI], 2023. · arXiv
  23. A. Hallak, D. Di Castro, and S. Mannor. Contextual markov decision processes. arXiv:1502.02259 [stat.ML], 2015. · arXiv
  24. N. Hansen, H. Su, and X. Wang. TD-MPC2: Scalable, robust world models for continuous control. In International Conference on Learning Representations (ICLR'24), 2024. Published online: iclr.cc.
  25. D. Kingma and M. Welling. Auto-encoding variational bayes. In Proceedings of the International Conference on Learning Representations (ICLR'14). CBLS, 2014.
  26. R. Kirk, A. Zhang, E. Grefenstette, and T. Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research (JAIR), 76:201-264, 2023.
  27. R. Koppejan and S. Whiteson. Neuroevolutionary reinforcement learning for generalized helicopter control. In G. Raidl et al (ed.), Proceedings of the 11th Genetic and Evolutionary Computation Conference (GECCO'09). Morgan Kaufmann Publishers, 2009.
  28. K. Lee, Y. Seo, S. Lee, H. Lee, and J. Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In H. Daume III and A. Singh (eds.), Proceedings of the 37th International Conference on Machine Learning (ICML'20), volume 98, pp. 5757-5766. Proceedings of Machine Learning Research, 2020.
  29. L. C. Melo. Transformers are meta-reinforcement learners. In K. Chaudhuri, S. Jegelka, L. Song, C. Szepesvári, G. Niu, and S. Sabato (eds.), Proceedings of the 39th International Conference on Machine Learning (ICML'22), volume 162 of Proceedings of Machine Learning Research, pp. 15340-15359.
  30. PMLR, 2022.
  31. A. Modi, N. Jiang, S. Singh, and A. Tewari. Markov decision processes with continuous side infor- mation. In Algorithmic Learning Theory (ALT'18), volume 83, pp. 597-618, 2018.
  32. A. Nagabandi, I. Clavera, S. Liu, R. S. Fearing, P. Abbeel, S. Levine, and C. Finn. Learning to adapt in dynamic, real-world environments through meta-reinforcement learning. In Proceedings of the International Conference on Learning Representations (ICLR'19), 2019. Published online: iclr.cc.
  33. J. Parker-Holder, R. Rajan, X. Song, A. Biedenkapp, Y. Miao, T. Eimer, B. Zhang, V. Nguyen, R. Calandra, A. Faust, F. Hutter, and M. Lindauer. Automated reinforcement learning (AutoRL): A survey and open problems. Journal of Artificial Intelligence Research (JAIR), 74:517-568, 2022.
  34. X. B. Peng, M. Andrychowicz, W. Zaremba, and P. Abbeel. Sim-to-real transfer of robotic con- trol with dynamics randomization. In International Conference on Robotics and Automation, (ICRA'18), pp. 1-8. IEEE, 2018.
  35. C. Perez, F. P. Such, and T. Karaletsos. Generalized hidden parameter mdps: Transferable model- based rl in a handful of trials. In F. Rossi, V. Conitzer, and F. Sha (eds.), Proceedings of the AAAI Conference on Artificial Intelligence, pp. 5403-5411. Association for the Advancement of Artificial Intelligence, AAAI Press, 2020.
  36. K. Rakelly, A. Zhou, C. Finn, S. Levine, and D. Quillen. Efficient off-policy meta-reinforcement learning via probabilistic context variables. In K. Chaudhuri and R. Salakhutdinov (eds.), Pro- ceedings of the 36th International Conference on Machine Learning (ICML'19), volume 97, pp. 5331-5340. Proceedings of Machine Learning Research, 2019.
  37. B. Schölkopf, F. Locatello, S. Bauer, N. R. Ke, N. Kalchbrenner, A. Goyal, and Y. Bengio. Towards causal representation learning. arXiv:2102.11107 [cs.LG], 2021. · arXiv
  38. S. Sodhani, F. Meier, J. Pineau, and A. Zhang. Block contextual mdps for continual learning. In R. Firoozi, N. Mehr, E. Yel, R. Antonova, J. Bohg, M. Schwager, and M. J. Kochenderfer (eds.), Learning for Dynamics and Control Conference, (L4DC'22), volume 168 of Proceedings of Machine Learning Research, pp. 608-623. PMLR, 2022.
  39. Y. Tassa, Y. Doron, A. Muldal, T. Erez, Y. Li, D. L. Casas, D. Budden, A. Abdolmaleki, J. Merel, A. Lefrancq, T. Lillicrap, and M. Riedmiller. Deepmind control suite. arXiv preprint arXiv: 1801.00690, 2018. · arXiv
  40. J. Tobin, R. Fong, A. Ray, J. Schneider, W. Zaremba, and P. Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In International Conference on Intelligent Robots and Systems (IROS'17), pp. 23-30, 2017.
  41. J. Wang, Z. Kurth-Nelson, H. Soyer, J. Leibo, D. Tirumala, R. Munos, C. Blundell, D. Kumaran, and M. Botvinick. Learning to reinforcement learn. In G. Gunzelmann, A. Howes, T. Tenbrink, and E. Davelaar (eds.), Proceedings of the 39th Annual Meeting of the Cognitive Science Society. cognitivesciencesociety.org, 2017.
  42. L. Wen, S. Zhang, H. E. Tseng, and H. Peng. Dream to adapt: Meta reinforcement learning by latent context imagination and MDP imagination. arXiv:2311.06673 [cs.LG], 2023. · arXiv
  43. P. Wu, A. Escontrela, D. Hafner, P. Abbeel, and K. Goldberg. Daydreamer: World models for physical robot learning. In K. Liu, D. Kulic, and J. Ichnowski (eds.), Conference on Robot Learning (CoRL'22) 2022, volume 205 of Proceedings of Machine Learning Research, pp. 2226-2240. PMLR, 2022.
  44. K. J. Young, A. Ramesh, L. Kirsch, and J. Schmidhuber. The benefits of model-based generaliza- tion in reinforcement learning. In A. Krause, E. Brunskill, K. Cho, B. Engelhardt, S. Sabato, and J. Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning (ICML'23), volume 202 of Proceedings of Machine Learning Research, pp. 40254-40276.
  45. PMLR, 2023.
  46. A. Zhang, S. Sodhani, K. Khetarpal, and J. Pineau. Learning robust state abstractions for hidden- parameter block MDPs. In Proceedings of the International Conference on Learning Representa- tions (ICLR'21), 2021a. Published online: iclr.cc.
  47. B. Zhang, R. Rajan, L. Pineda, N. Lambert, A. Biedenkapp, K. Chua, F. Hutter, and R. Calandra. On the importance of Hyperparameter Optimization for model-based reinforcement learning. In A. Banerjee and K. Fukumizu (eds.), Proceedings of the 24th International Conference on Artifi- cial Intelligence and Statistics (AISTATS'21), pp. 4015-4023. Proceedings of Machine Learning Research, 2021b. 0.1 0.2 0.3 0.4 0.5 0.75 1.0 1.25 1.5 1.6 1.7 1.8 1.9 2.0 Actuator Strength 19.6 18.62 17.64 16.66 15.68 14.7 12.25 9.81 7.35 4.9 3.92 3.185 2.45 1.715 0.98 Gravity 22 25 28 252 708 738 762 713 294 22 26 30 343 834 844 810 744 416 23 27 34 500 898 887 832 758 22 26 38 605 911 896 840 763 875 23 27 31 57 175 788 936 937 908 881 848 804 754 711 937 27 178 861 890 745 716 630 574 28 353 870 839 692 661 610 551 779 44 584 779 709 609 603 569 552 653 168 473 563 585 557 560 537 529 dmc Walker -Pixel observation -default -Return 200 400 600 800 0.1 0.2 0.3 0.4 0.5 0.75 1.0 1.25 1.5 1.6 1.7 1.8 1.9 2.0 Actuator Strength 19.6 18.62 17.64 16.66 15.68 14.7 12.25 9.81 7.35 4.9 3.92 3.185 2.45 1.715 0.98 Gravity 21 24 34 492 802 825 841 853 560 21 25 44 632 856 866 874 878 697 22 25 66 763 890 889 896 891 23 26 85 812 902 901 901 898 887 23 25 47 183 473 863 920 927 921 915 913 910 903 894 928 26 530 864 914 876 861 822 795 30 681 879 893 825 807 771 748 866 62 774 846 801 738 719 686 668 726 348 621 658 642 641 615 605 604 dmc Walker -Pixel observation -hidden-context -Return 200 400 600 800 0.1 0.2 0.3 0.4 0.5 0.75 1.0 1.25 1.5 1.6 1.7 1.8 1.9 2.0 Actuator Strength 19.6 18.62 17.64 16.66 15.68 14.7 12.25 9.81 7.35 4.9 3.92 3.185 2.45 1.715 0.98 Gravity 23 25 37 597 832 845 861 861 680 22 24 44 745 868 870 880 870 800 23 26 61 837 885 885 889 882 24 27 88 868 888 888 893 887 908 24 26 46 253 620 909 923 921 905 909 909 906 900 902 930 27 651 905 918 891 883 871 842 31 794 910 906 868 862 833 806 892 79 850 883 851 793 773 759 732 793 347 672 719 702 677 675 659 646 dmc Walker -Pixel observation -concat-context -Return 200 400 600 800 0.1 0.2 0.3 0.4 0.5 0.75 1.0 1.25 1.5 1.6 1.7 1.8 1.9 2.0 Actuator Strength 19.6 18.62 17.64 16.66 15.68 14.7 12.25 9.81 7.35 4.9 3.92 3.185 2.45 1.715 0.98 Gravity 22 24 38 571 822 828 818 805 695 23 24 48 785 874 874 872 859 845 23 28 73 881 905 911 908 886 23 28 112 912 920 916 914 895 936 24 28 56 323 705 931 943 942 942 940 934 924 913 894 947 27 699 910 935 905 886 855 789 31 784 908 918 877 855 819 757 897 81 808 860 862 794 784 721 688 795 430 663 683 707 653 646 617 587 dmc Walker -Pixel observation -cRSSM -Return 200 400 600 800

Cited by

Sorted by citations

  • Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization

    Röder, F., Benad, J., Eppe, M., & Banerjee, P. (2025) · Advances in Neural Information Processing Systems 38 · doi:10.52202/085713-2438

  • Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization

    Röder, F., Benad, J., Eppe, M., & Banerjee, P. K. (2025) · arXiv

    “We also evaluate against the context-aware baselines cRSSM-S and cRSSM-D, which use the same world model and actor-critic architecture but directly incorporate the ground-truth context c [Prasanna et al., 2024].”

Authors
Prasanna, S., Farid, K., Rajan, R., & Biedenkapp, A.
Affiliations
University of Freiburg (DE)
ORCID
Venue
arXiv
Publisher
arXiv
Published
2024
ISSN
2331-8422 (ISSN-L)
DOI
10.48550/arxiv.2403.10967
arXiv
2403.10967
OpenAlex
W4393023599
Cited by
0
References
0
Type
posted content
Language
en
Open access
Licence
Locations
  • arXiv (Cornell University) (submittedVersion, cc-by-sa)
  • arXiv (Cornell University)
Funders
Albert-Ludwigs-Universität Freiburg
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 › AI-based Problem Solving and Planning
Keywords
zero-shot generalization, world models, recurrent state-space models, dreamers, embodied agents

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