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

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

## Key facts

- Cited by 16 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 37 of its 43 references that are in the map.
- First posted on arXiv 16 Mar 2024.

## 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 37)

- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [Learning Dexterous In-Hand Manipulation](https://mlsanity.com/paper/1808.00177/) (2018)
- [Sim-to-Real Transfer of Robotic Control with Dynamics Randomization](https://mlsanity.com/paper/1710.06537/) (2017)
- [TD-MPC2: Scalable, Robust World Models for Continuous Control](https://mlsanity.com/paper/2310.16828/) (2023)
- [RL2: Fast Reinforcement Learning via Slow Reinforcement Learning](https://mlsanity.com/paper/1611.02779/) (2016)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [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)
- [Learning to Adapt in Dynamic, Real-World Environments Through Meta-Reinforcement Learning](https://mlsanity.com/paper/1803.11347/) (2018)
- [Contextual Markov Decision Processes](https://mlsanity.com/paper/1502.02259/) (2015)
- [Towards Causal Representation Learning](https://mlsanity.com/paper/2102.11107/) (2021)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Provably efficient RL with Rich Observations via Latent State Decoding](https://mlsanity.com/paper/1901.09018/) (2019)
- [Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2005.06800/) (2020)
- [Automated Reinforcement Learning (AutoRL): A Survey and Open Problems](https://mlsanity.com/paper/2201.03916/) (2022)
- [On the Importance of Hyperparameter Optimization for Model-based Reinforcement Learning](https://mlsanity.com/paper/2102.13651/) (2021)
- [Markov Decision Processes with Continuous Side Information](https://mlsanity.com/paper/1711.05726/) (2017)
- [Contextualize Me -- The Case for Context in Reinforcement Learning](https://mlsanity.com/paper/2202.04500/) (2022)
- [Augmented World Models Facilitate Zero-Shot Dynamics Generalization From a Single Offline Environment](https://mlsanity.com/paper/2104.05632/) (2021)
- [The MineRL 2020 Competition on Sample Efficient Reinforcement Learning using Human Priors](https://mlsanity.com/paper/2101.11071/) (2021)
- [Transformers are Meta-Reinforcement Learners](https://mlsanity.com/paper/2206.06614/) (2022)
- [Procedural Generalization by Planning with Self-Supervised World Models](https://mlsanity.com/paper/2111.01587/) (2021)

## Built on by (most cited first, 16 of 16)

- [Is Single-View Mesh Reconstruction Ready for Robotics?](https://mlsanity.com/paper/2505.17966/) (2025)
- [A Mechanistic View on Video Generation as World Models: State and Dynamics](https://mlsanity.com/paper/2601.17067/) (2026)
- [Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization](https://mlsanity.com/paper/2508.20294/) (2025)
- [Out-of-Distribution Generalization with a SPARC: Racing 100 Unseen Vehicles with a Single Policy](https://mlsanity.com/paper/2511.09737/) (2025)
- [Craftium: Bridging Flexibility and Efficiency for Rich 3D Single- and Multi-Agent Environments](https://mlsanity.com/paper/2407.03969/) (2024)
- [GRAM: Generalization in Deep RL with a Robust Adaptation Module](https://mlsanity.com/paper/2412.04323/) (2024)
- [MetaMind: General and Cognitive World Models in Multi-Agent Systems by Meta-Theory of Mind](https://mlsanity.com/paper/2603.00808/) (2026)
- [Adaptive Shielding for Safe Reinforcement Learning under Hidden-Parameter Dynamics Shifts](https://mlsanity.com/paper/2506.11033/) (2025)
- [Efficient Generation of Diverse Cooperative Agents with World Models](https://mlsanity.com/paper/2506.07450/) (2025)
- [Contextual Intelligence The Next Leap for Reinforcement Learning](https://mlsanity.com/paper/2604.02348/) (2026)
- [Latent Q-Barrier Shielding for Safe In-Context Reinforcement Learning](https://mlsanity.com/paper/2605.25267/) (2026)
- [Zero-shot Generalization in Inventory Management: Train, then Estimate and Decide](https://mlsanity.com/paper/2411.00515/) (2024)
- [Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts](https://mlsanity.com/paper/2507.07348/) (2025)
- [Context-Aware Model-Based Reinforcement Learning for Autonomous Racing](https://mlsanity.com/paper/2510.11501/) (2025)
- [Vehicle Dynamics Embedded World Models for Autonomous Driving](https://mlsanity.com/paper/2512.02417/) (2025)
- [Dynamic Context Scheduling: Learning Beyond the Static Universe](https://mlsanity.com/paper/2608.20799/) (2026)

## Abstract (arXiv)

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.