# Augmented World Models Facilitate Zero-Shot Dynamics Generalization From a Single Offline Environment

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

## Key facts

- Cited by 42 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 53 of its 59 references that are in the map.
- First posted on arXiv 12 Apr 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 53)

- [Simple and Scalable Predictive Uncertainty Estimation using Deep Ensembles](https://mlsanity.com/paper/1612.01474/) (2016)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Conservative Q-Learning for Offline Reinforcement Learning](https://mlsanity.com/paper/2006.04779/) (2020)
- [OpenAI Gym](https://mlsanity.com/paper/1606.01540/) (2016)
- [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)
- [D4RL: Datasets for Deep Data-Driven Reinforcement Learning](https://mlsanity.com/paper/2004.07219/) (2020)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Off-Policy Deep Reinforcement Learning without Exploration](https://mlsanity.com/paper/1812.02900/) (2018)
- [Soft Actor-Critic Algorithms and Applications](https://mlsanity.com/paper/1812.05905/) (2018)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [CURL: Contrastive Unsupervised Representations for Reinforcement Learning](https://mlsanity.com/paper/2004.04136/) (2020)
- [Solving Rubik's Cube with a Robot Hand](https://mlsanity.com/paper/1910.07113/) (2019)
- [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)
- [Behavior Regularized Offline Reinforcement Learning](https://mlsanity.com/paper/1911.11361/) (2019)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [MOReL : Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2005.05951/) (2020)
- [Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels](https://mlsanity.com/paper/2004.13649/) (2020)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Generalizing to Unseen Domains via Adversarial Data Augmentation](https://mlsanity.com/paper/1805.12018/) (2018)
- [Efficient Off-Policy Meta-Reinforcement Learning via Probabilistic Context Variables](https://mlsanity.com/paper/1903.08254/) (2019)
- [Challenges of Real-World Reinforcement Learning](https://mlsanity.com/paper/1904.12901/) (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)
- [Model-Ensemble Trust-Region Policy Optimization](https://mlsanity.com/paper/1802.10592/) (2018)
- [VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning](https://mlsanity.com/paper/1910.08348/) (2019)

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

- [Evolving Curricula with Regret-Based Environment Design](https://mlsanity.com/paper/2203.01302/) (2022)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [RAMBO-RL: Robust Adversarial Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2204.12581/) (2022)
- [Synthetic Experience Replay](https://mlsanity.com/paper/2303.06614/) (2023)
- [Challenges and Opportunities in Offline Reinforcement Learning from Visual Observations](https://mlsanity.com/paper/2206.04779/) (2022)
- [Revisiting Design Choices in Offline Model-Based Reinforcement Learning](https://mlsanity.com/paper/2110.04135/) (2021)
- [DARA: Dynamics-Aware Reward Augmentation in Offline Reinforcement Learning](https://mlsanity.com/paper/2203.06662/) (2022)
- [Policy-Guided Diffusion](https://mlsanity.com/paper/2404.06356/) (2024)
- [One Risk to Rule Them All: A Risk-Sensitive Perspective on Model-Based Offline Reinforcement Learning](https://mlsanity.com/paper/2212.00124/) (2022)
- [Model-Based Offline Reinforcement Learning with Pessimism-Modulated Dynamics Belief](https://mlsanity.com/paper/2210.06692/) (2022)
- [GTA: Generative Trajectory Augmentation with Guidance for Offline Reinforcement Learning](https://mlsanity.com/paper/2405.16907/) (2024)
- [Dynamics Generalisation in Reinforcement Learning via Adaptive Context-Aware Policies](https://mlsanity.com/paper/2310.16686/) (2023)
- [On Pathologies in KL-Regularized Reinforcement Learning from Expert Demonstrations](https://mlsanity.com/paper/2212.13936/) (2022)
- [Improving Zero-shot Generalization in Offline Reinforcement Learning using Generalized Similarity Functions](https://mlsanity.com/paper/2111.14629/) (2021)
- [Dreaming of Many Worlds: Learning Contextual World Models Aids Zero-Shot Generalization](https://mlsanity.com/paper/2403.10967/) (2024)
- [Generalization of Reinforcement Learning with Policy-Aware Adversarial Data Augmentation](https://mlsanity.com/paper/2106.15587/) (2021)
- [A Relational Intervention Approach for Unsupervised Dynamics Generalization in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2206.04551/) (2022)
- [Simulating the Visual World with Artificial Intelligence: A Roadmap](https://mlsanity.com/paper/2511.08585/) (2025)
- [Representation Learning in Deep RL via Discrete Information Bottleneck](https://mlsanity.com/paper/2212.13835/) (2022)
- [State-Covering Trajectory Stitching for Diffusion Planners](https://mlsanity.com/paper/2506.00895/) (2025)
- [Learning General World Models in a Handful of Reward-Free Deployments](https://mlsanity.com/paper/2210.12719/) (2022)
- [Hindsight PRIORs for Reward Learning from Human Preferences](https://mlsanity.com/paper/2404.08828/) (2024)
- [Learn the Time to Learn: Replay Scheduling in Continual Learning](https://mlsanity.com/paper/2209.08660/) (2022)
- [Adapting World Models with Latent-State Dynamics Residuals](https://mlsanity.com/paper/2504.02252/) (2025)
- [The Edge-of-Reach Problem in Offline Model-Based Reinforcement Learning](https://mlsanity.com/paper/2402.12527/) (2024)
- [Pessimistic Value Iteration for Multi-Task Data Sharing in Offline Reinforcement Learning](https://mlsanity.com/paper/2404.19346/) (2024)
- [Is Single-View Mesh Reconstruction Ready for Robotics?](https://mlsanity.com/paper/2505.17966/) (2025)
- [Zero-Shot Reinforcement Learning Under Partial Observability](https://mlsanity.com/paper/2506.15446/) (2025)
- [Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization](https://mlsanity.com/paper/2508.20294/) (2025)
- [Imagined Autocurricula](https://mlsanity.com/paper/2509.13341/) (2025)

## Abstract (arXiv)

Reinforcement learning from large-scale offline datasets provides us with the ability to learn policies without potentially unsafe or impractical exploration. Significant progress has been made in the past few years in dealing with the challenge of correcting for differing behavior between the data collection and learned policies. However, little attention has been paid to potentially changing dynamics when transferring a policy to the online setting, where performance can be up to 90% reduced for existing methods. In this paper we address this problem with Augmented World Models (AugWM). We augment a learned dynamics model with simple transformations that seek to capture potential changes in physical properties of the robot, leading to more robust policies. We not only train our policy in this new setting, but also provide it with the sampled augmentation as a context, allowing it to adapt to changes in the environment. At test time we learn the context in a self-supervised fashion by approximating the augmentation which corresponds to the new environment. We rigorously evaluate our approach on over 100 different changed dynamics settings, and show that this simple approach can significantly improve the zero-shot generalization of a recent state-of-the-art baseline, often achieving successful policies where the baseline fails.