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

Fuente: arXiv
Enregistré dans:
Détails bibliographiques
Auteurs principaux: Chapman, James, Karhadkar, Kedar, Montufar, Guido
Format: Preprint
Publié: 2025
Sujets:
Accès en ligne:
Tags: Ajouter un tag
Pas de tags, Soyez le premier à ajouter un tag!
_version_ 1866908733265674240
author Chapman, James
Karhadkar, Kedar
Montufar, Guido
author_facet Chapman, James
Karhadkar, Kedar
Montufar, Guido
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2507_07348
institution arXiv
publishDate 2025
record_format arxiv
spellingShingle Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts
Chapman, James
Karhadkar, Kedar
Montufar, Guido
Machine Learning
I.2.6; I.2.8
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.
title Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts
topic Machine Learning
I.2.6; I.2.8
url https://arxiv.org/abs/2507.07348