# Context-aware Dynamics Model for Generalization in Model-Based Reinforcement Learning

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

## Key facts

- Cited by 98 papers in the map (data through 14 Oct 2026).
- First posted on arXiv 14 May 2020.

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

- [Dropout Q-Functions for Doubly Efficient Reinforcement Learning](https://mlsanity.com/paper/2110.02034/) (2021)
- [Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability](https://mlsanity.com/paper/2107.06277/) (2021)
- [A Survey on Model-based Reinforcement Learning](https://mlsanity.com/paper/2206.09328/) (2022)
- [Generalizing to New Physical Systems via Context-Informed Dynamics Model](https://mlsanity.com/paper/2202.01889/) (2022)
- [Offline Reinforcement Learning with Reverse Model-based Imagination](https://mlsanity.com/paper/2110.00188/) (2021)
- [Contextualize Me -- The Case for Context in Reinforcement Learning](https://mlsanity.com/paper/2202.04500/) (2022)
- [Factored Adaptation for Non-Stationary Reinforcement Learning](https://mlsanity.com/paper/2203.16582/) (2022)
- [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)
- [A Survey on Physics Informed Reinforcement Learning: Review and Open Problems](https://mlsanity.com/paper/2309.01909/) (2023)
- [State Regularized Policy Optimization on Data with Dynamics Shift](https://mlsanity.com/paper/2306.03552/) (2023)
- [Double Check Your State Before Trusting It: Confidence-Aware Bidirectional Offline Model-Based Imagination](https://mlsanity.com/paper/2206.07989/) (2022)
- [Dynamics Generalisation in Reinforcement Learning via Adaptive Context-Aware Policies](https://mlsanity.com/paper/2310.16686/) (2023)
- [GOPlan: Goal-conditioned Offline Reinforcement Learning by Planning with Learned Models](https://mlsanity.com/paper/2310.20025/) (2023)
- [Context Shift Reduction for Offline Meta-Reinforcement Learning](https://mlsanity.com/paper/2311.03695/) (2023)
- [Model-based Meta Reinforcement Learning using Graph Structured Surrogate Models](https://mlsanity.com/paper/2102.08291/) (2021)
- [Live in the Moment: Learning Dynamics Model Adapted to Evolving Policy](https://mlsanity.com/paper/2207.12141/) (2022)
- [Meta-DT: Offline Meta-RL as Conditional Sequence Modeling with World Model Disentanglement](https://mlsanity.com/paper/2410.11448/) (2024)
- [Dreaming of Many Worlds: Learning Contextual World Models Aids Zero-Shot Generalization](https://mlsanity.com/paper/2403.10967/) (2024)
- [A Relational Intervention Approach for Unsupervised Dynamics Generalization in Model-Based Reinforcement Learning](https://mlsanity.com/paper/2206.04551/) (2022)
- [MAMBA: an Effective World Model Approach for Meta-Reinforcement Learning](https://mlsanity.com/paper/2403.09859/) (2024)
- [Robust Fast Adaptation from Adversarially Explicit Task Distribution Generation](https://mlsanity.com/paper/2407.19523/) (2024)
- [Improving Generalization in Meta-RL with Imaginary Tasks from Latent Dynamics Mixture](https://mlsanity.com/paper/2105.13524/) (2021)
- [Learning General World Models in a Handful of Reward-Free Deployments](https://mlsanity.com/paper/2210.12719/) (2022)
- [The Evolution of Reinforcement Learning in Quantitative Finance: A Survey](https://mlsanity.com/paper/2408.10932/) (2024)
- [PerSim: Data-Efficient Offline Reinforcement Learning with Heterogeneous Agents via Personalized Simulators](https://mlsanity.com/paper/2102.06961/) (2021)
- [Mind the Model, Not the Agent: The Primacy Bias in Model-based RL](https://mlsanity.com/paper/2310.15017/) (2023)
- [PEAC: Unsupervised Pre-training for Cross-Embodiment Reinforcement Learning](https://mlsanity.com/paper/2405.14073/) (2024)
- [Offline Imitation Learning with Model-based Reverse Augmentation](https://mlsanity.com/paper/2406.12550/) (2024)
- [Meta-Learning Parameterized Skills](https://mlsanity.com/paper/2206.03597/) (2022)

## Abstract (arXiv)

Model-based reinforcement learning (RL) enjoys several benefits, such as data-efficiency and planning, by learning a model of the environment's dynamics. However, learning a global model that can generalize across different dynamics is a challenging task. To tackle this problem, we decompose the task of learning a global dynamics model into two stages: (a) learning a context latent vector that captures the local dynamics, then (b) predicting the next state conditioned on it. In order to encode dynamics-specific information into the context latent vector, we introduce a novel loss function that encourages the context latent vector to be useful for predicting both forward and backward dynamics. The proposed method achieves superior generalization ability across various simulated robotics and control tasks, compared to existing RL schemes.