# Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Deep RL algorithms and representations, one of 15 directions in Deep reinforcement learning methods.
- Builds on 46 of its 70 references that are in the map.
- First posted on arXiv 10 Jul 2025.

## Placed in

- Deep reinforcement learning methods / [Deep RL algorithms and representations](https://mlsanity.com/deeprl/deep-rl-algorithms-and-representations/)

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Playing Atari with Deep Reinforcement Learning](https://mlsanity.com/paper/1312.5602/) (2013)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Ray: A Distributed Framework for Emerging AI Applications](https://mlsanity.com/paper/1712.05889/) (2017)
- [Leveraging Procedural Generation to Benchmark Reinforcement Learning](https://mlsanity.com/paper/1912.01588/) (2019)
- [Sim-to-Real Transfer in Deep Reinforcement Learning for Robotics: a Survey](https://mlsanity.com/paper/2009.13303/) (2020)
- [Tune: A Research Platform for Distributed Model Selection and Training](https://mlsanity.com/paper/1807.05118/) (2018)
- [Towards Continual Reinforcement Learning: A Review and Perspectives](https://mlsanity.com/paper/2012.13490/) (2020)
- [CAD2RL: Real Single-Image Flight without a Single Real Image](https://mlsanity.com/paper/1611.04201/) (2016)
- [Contextual Markov Decision Processes](https://mlsanity.com/paper/1502.02259/) (2015)
- [Deep Reinforcement Learning for Robotics: A Survey of Real-World Successes](https://mlsanity.com/paper/2408.03539/) (2024)
- [Active Domain Randomization](https://mlsanity.com/paper/1904.04762/) (2019)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [DARLA: Improving Zero-Shot Transfer in Reinforcement Learning](https://mlsanity.com/paper/1707.08475/) (2017)
- [One Policy to Control Them All: Shared Modular Policies for Agent-Agnostic Control](https://mlsanity.com/paper/2007.04976/) (2020)
- [Replay-Guided Adversarial Environment Design](https://mlsanity.com/paper/2110.02439/) (2021)
- [Why Generalization in RL is Difficult: Epistemic POMDPs and Implicit Partial Observability](https://mlsanity.com/paper/2107.06277/) (2021)
- [A Tutorial on Meta-Reinforcement Learning](https://mlsanity.com/paper/2301.08028/) (2023)
- [Partially Observable Markov Decision Processes in Robotics: A Survey](https://mlsanity.com/paper/2209.10342/) (2022)
- [Understanding Domain Randomization for Sim-to-real Transfer](https://mlsanity.com/paper/2110.03239/) (2021)
- [My Body is a Cage: the Role of Morphology in Graph-Based Incompatible Control](https://mlsanity.com/paper/2010.01856/) (2020)
- [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)
- [Efficient Differentiable Simulation of Articulated Bodies](https://mlsanity.com/paper/2109.07719/) (2021)
- [Data-efficient Co-Adaptation of Morphology and Behaviour with Deep Reinforcement Learning](https://mlsanity.com/paper/1911.06832/) (2019)
- [On the Importance of Exploration for Generalization in Reinforcement Learning](https://mlsanity.com/paper/2306.05483/) (2023)
- [A Review of Differentiable Simulators](https://mlsanity.com/paper/2407.05560/) (2024)
- [Estimation Considerations in Contextual Bandits](https://mlsanity.com/paper/1711.07077/) (2017)

## Abstract (arXiv)

Deep reinforcement learning (DRL) has achieved remarkable success across multiple domains, including competitive games, natural language processing, and robotics. Despite these advancements, policies trained via DRL often struggle to generalize to evaluation environments with different parameters. This challenge is typically addressed by training with multiple contexts and/or by leveraging additional structure in the problem. However, obtaining sufficient training data across diverse contexts can be impractical in real-world applications. In this work, we consider contextual Markov decision processes (CMDPs) with transition and reward functions that exhibit regularity in context parameters. We introduce the context-enhanced Bellman equation (CEBE) to improve generalization when training on a single context. We prove both analytically and empirically that the CEBE yields a first-order approximation to the Q-function trained across multiple contexts. We then derive context sample enhancement (CSE) as an efficient data augmentation method for approximating the CEBE in deterministic control environments. We numerically validate the performance of CSE in simulation environments, showcasing its potential to improve generalization in DRL.