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

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

## Key facts

- Cited by 5 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 30 of its 39 references that are in the map.
- First posted on arXiv 27 Aug 2025.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Deep reinforcement learning methods / [RL for robot locomotion and manipulation](https://mlsanity.com/deeprl/rl-for-robot-locomotion-and-manipulation/)

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [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)
- [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)
- [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)
- [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)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [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)
- [Multi-Task Reinforcement Learning with Context-based Representations](https://mlsanity.com/paper/2102.06177/) (2021)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2005.06800/) (2020)
- [DensePhysNet: Learning Dense Physical Object Representations via Multi-step Dynamic Interactions](https://mlsanity.com/paper/1906.03853/) (2019)
- [Markov Decision Processes with Continuous Side Information](https://mlsanity.com/paper/1711.05726/) (2017)
- [Hardware Conditioned Policies for Multi-Robot Transfer Learning](https://mlsanity.com/paper/1811.09864/) (2018)
- [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)
- [Transformers are Meta-Reinforcement Learners](https://mlsanity.com/paper/2206.06614/) (2022)
- [Towards General-Purpose Model-Free Reinforcement Learning](https://mlsanity.com/paper/2501.16142/) (2025)
- [Trajectory-wise Multiple Choice Learning for Dynamics Generalization in Reinforcement Learning](https://mlsanity.com/paper/2010.13303/) (2020)
- [Context is Everything: Implicit Identification for Dynamics Adaptation](https://mlsanity.com/paper/2203.05549/) (2022)
- [Block Contextual MDPs for Continual Learning](https://mlsanity.com/paper/2110.06972/) (2021)
- [Dynamics Generalisation in Reinforcement Learning via Adaptive Context-Aware Policies](https://mlsanity.com/paper/2310.16686/) (2023)
- [HyperDynamics: Meta-Learning Object and Agent Dynamics with Hypernetworks](https://mlsanity.com/paper/2103.09439/) (2021)
- [Dreaming of Many Worlds: Learning Contextual World Models Aids Zero-Shot Generalization](https://mlsanity.com/paper/2403.10967/) (2024)

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

- [Overcoming Statistical Bias in Action-Controllable World Models](https://mlsanity.com/paper/2608.04653/) (2026)
- [HaM-World: Soft-Hamiltonian World Models with Selective Memory for Planning](https://mlsanity.com/paper/2605.05951/) (2026)
- [Event-Conditioned Diagnostics of Kinematic, Contact, and Object-Permanence Structure in Passive Object-State World Models](https://mlsanity.com/paper/2606.28455/) (2026)
- [When Does Test-Time Physical Diagnosis Pay? A Frozen Policy Buys Evidence It Never Reads](https://mlsanity.com/paper/2609.22299/) (2026)
- [Shaping Persistent Representations from Independent Interactions](https://mlsanity.com/paper/2609.34604/) (2026)

## Abstract (arXiv)

Real-world reinforcement learning demands adaptation to unseen environmental conditions without costly retraining. 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. We introduce Dynamics-Aligned Latent Imagination (DALI), a framework integrated within the Dreamer architecture that infers latent context representations from agent-environment interactions. By training a self-supervised encoder to predict forward dynamics, DALI generates actionable representations conditioning the world model and policy, bridging perception and control. We theoretically prove this encoder is essential for efficient context inference and robust generalization. DALI's latent space enables counterfactual consistency: Perturbing a gravity-encoding dimension alters imagined rollouts in physically plausible ways. On challenging cMDP benchmarks, DALI achieves significant gains over context-unaware baselines, often surpassing context-aware baselines in extrapolation tasks, enabling zero-shot generalization to unseen contextual variations.