# Training on Irrelevant States Implies Data Augmentation: Generalization in Contextual MDPs

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 66 of its 85 references that are in the map.
- First posted on arXiv 4 Oct 2024.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Offline Reinforcement Learning: Tutorial, Review, and Perspectives on Open Problems](https://mlsanity.com/paper/2005.01643/) (2020)
- [Domain Randomization for Transferring Deep Neural Networks from Simulation to the Real World](https://mlsanity.com/paper/1703.06907/) (2017)
- [Hindsight Experience Replay](https://mlsanity.com/paper/1707.01495/) (2017)
- [Exploration by Random Network Distillation](https://mlsanity.com/paper/1810.12894/) (2018)
- [DeepMind Control Suite](https://mlsanity.com/paper/1801.00690/) (2018)
- [IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures](https://mlsanity.com/paper/1802.01561/) (2018)
- [Deep Learning and the Information Bottleneck Principle](https://mlsanity.com/paper/1503.02406/) (2015)
- [Sim-to-Real Transfer of Robotic Control with Dynamics Randomization](https://mlsanity.com/paper/1710.06537/) (2017)
- [Unifying Count-Based Exploration and Intrinsic Motivation](https://mlsanity.com/paper/1606.01868/) (2016)
- [Domain Generalization with MixStyle](https://mlsanity.com/paper/2104.02008/) (2021)
- [Reinforcement Learning with Augmented Data](https://mlsanity.com/paper/2004.14990/) (2020)
- [Deep Recurrent Q-Learning for Partially Observable MDPs](https://mlsanity.com/paper/1507.06527/) (2015)
- [A Survey of Data Augmentation Approaches for NLP](https://mlsanity.com/paper/2105.03075/) (2021)
- [Quantifying Generalization in Reinforcement Learning](https://mlsanity.com/paper/1812.02341/) (2018)
- [Minigrid & Miniworld: Modular & Customizable Reinforcement Learning Environments for Goal-Oriented Tasks](https://mlsanity.com/paper/2306.13831/) (2023)
- [Distributed Prioritized Experience Replay](https://mlsanity.com/paper/1803.00933/) (2018)
- [First return, then explore](https://mlsanity.com/paper/2004.12919/) (2020)
- [Count-Based Exploration with Neural Density Models](https://mlsanity.com/paper/1703.01310/) (2017)
- [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)
- [Randomized Prior Functions for Deep Reinforcement Learning](https://mlsanity.com/paper/1806.03335/) (2018)
- [A Survey of Zero-shot Generalisation in Deep Reinforcement Learning](https://mlsanity.com/paper/2111.09794/) (2021)
- [Exploration in Deep Reinforcement Learning: A Survey](https://mlsanity.com/paper/2205.00824/) (2022)
- [Goal-Conditioned Reinforcement Learning: Problems and Solutions](https://mlsanity.com/paper/2201.08299/) (2022)
- [Deep Exploration via Randomized Value Functions](https://mlsanity.com/paper/1703.07608/) (2017)
- [Generalization in Reinforcement Learning by Soft Data Augmentation](https://mlsanity.com/paper/2011.13389/) (2020)
- [RIDE: Rewarding Impact-Driven Exploration for Procedurally-Generated Environments](https://mlsanity.com/paper/2002.12292/) (2020)

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

- [How Ensembles of Distilled Policies Improve Generalisation in Reinforcement Learning](https://mlsanity.com/paper/2505.16581/) (2025)

## Abstract (arXiv)

In the zero-shot policy transfer (ZSPT) setting for contextual Markov decision processes (CMDP), agents train on a fixed, finite set of contexts and must generalize to new ones. Recent work has demonstrated that training on additional states, even if they are irrelevant for solving the current context, can improve generalization to unseen contexts. In this paper, we demonstrate that training on these states can indeed improve generalization, but can come at a cost of reducing the accuracy of the learned value function, which should hurt generalization. We hypothesize and demonstrate that increasing the agent's coverage by training on these additional states while also increasing the accuracy improves generalization even further. Inspired by this, we propose a simple approach Explore-Go that leverages existing pure exploration strategies in a new way: by introducing a pure exploration phase at the start of each training episode. Unlike previous approaches that apply exploration strategies for the purpose of improving generalization, our approach can be combined with both on- and off-policy algorithms. We demonstrate the effectiveness of Explore-Go when combined with several popular algorithms and show an increase in test-time performance across several generalization benchmarks, even partially observable ones. With this, we hope to provide practitioners with a simple modification that can significantly improve the generalization of their agents.