Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts
James Chapman, Kedar Karhadkar, Guido Montufar
TL;DR
The paper addresses zero-shot generalization in reinforcement learning under contextual MDPs with smooth context dependencies. It introduces the context-enhanced Bellman equation (CEBE), a first-order approximation of the Bellman equation for CMDPs linearized about a base context $c_0$, and proves that $Q^{CE}$ closely tracks $Q^{BE}$ in a neighborhood of $c_0$. To make this practical, the authors derive context sample enhancement (CSE), a data-augmentation scheme that generates context-perturbed samples from a single training context in deterministic settings. Empirical validation spans tabular and continuous-control domains (including PendulumGoal and MuJoCo tasks), showing that CEBE and CSE improve zero-shot generalization over vanilla training and domain randomization, with the approach being simple to implement and compatible with model-free and potentially model-based RL. Limitations include a focus on deterministic transitions and fully observable contexts, suggesting future work on high-dimensional contexts, partial observability, and offline RL extensions.
Abstract
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.