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

Jul 16, 2026

Upload by Jiven

In the zero-shot policy transfer (ZSPT) setting for contextual Markov decision processes (CMDP), agents learn from a fixed, finite set of contexts and must generalize to unseen ones. While prior work suggests that training on additional, even irrelevant, states can improve generalization, this paper shows a value-function accuracy trade-off. Increasing context coverage while preserving low value error further boosts generalization. It introduces Explore-Go, which adds a pure exploration phase at each episode start and improves test-time performance across multiple benchmarks and algorithms.

LanguageEnglish

CategoryResearch & Report

FormatPDF

Pages35

Reading Time1 h 28 min

Introduction

Summary

Contribution(s)

Explore-Go Method

Experiments and Benchmarks

Related Settings and Trade-offs

arXiv :2410 .03565v4 [ cs .LG] 13 Jul 2026 Training on Irrelevant States Implies Data Augmentation: Generalization in Contextual MDPs Max Weltevrede, Caroline Horsch, Matthijs T. J. Spaan, Wendelin Böhmer Keywords: Generalization, Contextual MDP, Reinforcement Learning Summary In the zero-shot policy transfer (ZSPT) setting for contextual Markov decision processes (MDP), agents train on a fixed, finite set of contexts and must generalize to new ones. Recent work has argued and 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. Contribution(s) 1. We demonstrate that increasing the agent’s coverage through exploration of the training contexts can improve generalization but can also result in a larger error between the learned and optimal value. Additionally, we demonstrate that generalization can be further improved when expanding coverage while also taking care to maintain low value error. Context: We analyze the error of a DQN (Mnih et al., 2015) agent in Four Rooms (ChevalierBoisvert et al., 2023) between the estimated Q-value and the optimal value, averaged over all states in the training contexts. We indirectly demonstrate that generalization can be further improved by proposing a method that increases the coverage of training contexts, while also achieving low value error, and generalizes significantly better than any of the baselines we compare against. These results do not fully isolate or prove the proposed trade-off between coverage and value error, but demonstrate a key phenomenon that is mitigated with our proposed approach. 2. We propose a simple approach called Explore-Go that leverages existing pure exploration techniques in a novel way, leading to improved generalization performance. It artificially increases the number of contexts on which the agent trains by leveraging an exploration phase at the beginning of each episode. Context: The goal of Explore-Go is not to introduce a new way of exploring, but rather a different way of using existing exploration techniques. It can be combined with both on-policy and off-policy approaches. This in contrast with previous work that also uses exploration to improve generalization, which can only be used in combination with off-policy methods (Jiang et al., 2023) . Although increased training diversity has been associated with improved generalization before, we contextualize this in the ZSPT setting in a novel way. 3. We demonstrate that Explore-Go can improve test-time performance when combined with on-and off-policy methods on several generalization benchmarks, including both full and partially observable ones. Context: We show results for Explore-Go in combination with DQN, PPO (Schulmanet al., 2017), APPO (Petrenko et al., 2020), SAC (Haarnoja et al., 2018) and RAD (Laskin et al., 2020) and evaluate it on several environments from Minigrid, DMC (Tassa et al., 2018) and VizDoom (Wydmuch et al., 2019) . Although we believe our approach can be beneficial in many different CMDPs, our results are mostly demonstrated in environments with state-reaching objectives. Most of our analysis is focused on the fully observable setting, but we provide evidence for Explore-Go’s practical utility by demonstrating good results in the partially observable VizDoom as well. Training on Irrelevant States Implies Data Augmentation: Generalization in Contextual MDPs Max Weltevrede 1 , Caroline Horsch 1 , Matthijs T. J. Spaan 1 , Wendelin Böhmer 1 [m.r.weltevrede@tudelft.nl](m.r.weltev

FAQs

What is the main problem addressed in this paper?

The paper studies how agents in contextual MDPs generalize to new contexts when they train on a fixed, finite set under the zero-shot policy transfer (ZSPT) setting.

How can training on irrelevant states affect generalization?

Training on additional states can improve generalization, but it may reduce the accuracy of the learned value function, which would normally be expected to hurt generalization.

What is Explore-Go and how does it improve performance?

Explore-Go uses a pure exploration phase at the start of each training episode to effectively increase training coverage. Combined with on-policy and off-policy methods, it improves test-time performance on multiple generalization benchmarks, including partially observable settings.

Jul 16, 2026

Upload by Jiven