Reproduce “Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts” on OpenResearch

About this paper

James Chapman, Kedar Karhadkar, Guido Montufar

A new perturbative framework, featuring the Context-Enhanced Bellman Equation (CEBE) and Context Sample Enhancement (CSE), allows deep reinforcement learning agents to achieve zero-shot generalization to unseen contexts using limited training data. This approach theoretically provides first-order accurate Q-function approximations and empirically enhances policy performance in diverse control tasks, frequently outperforming or matching local domain randomization.

Key insights

• The Context-Enhanced Bellman Equation (CEBE) provides a first-order accurate approximation ($O(||c-c_0||^2)$ error) of the true Q-function in a neighborhood of the training context, established through theoretical stability guarantees.
• Q-functions in CMDPs exhibit stability under small perturbations to transition dynamics and rewards, which can be formally quantified using total variation norms and Wasserstein distances.
• A policy optimized to be approximately optimal with respect to the CEBE will also be approximately optimal for the true Bellman equation, bridging the gap between theoretical value function approximation and practical policy performance.

Summary by alphaXiv.