Title: Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts
Venue: NeurIPS 2025
Accepted: yes
Keywords: Reinforcement Learning

Zero-Shot Generalization

Contextual MDPs

Robust Machine Learning
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nCLAIMS: \n\nProposed CEBE based on the perturbation over CMDPs: The analyses surrounding how CEBE helps OOD generalization are technically sound, despite quite strong structural assumptions;\n\nDerived CSE as a data augmentation approach to approximate CEBE (for deterministic control): CSE is simple and straightforward, as presented in Sec. 4;\n\nThe authors claim that CEBE can help agents' generalization by training even on a single task. CSE is validated in experiments with 1-D contexts: 1-D contexts missed the opportunity to better convince the readers about the scalability and the potential of the proposed method, as discussed below;\n\nNOVELTY: \n\nthe CSE method is simple and straightforward, and possibly can be combined with other techniques to yield better generalization\n\nCSE seems to work well with even stochastic rewards, despite not mathematically proven\n\nSTRENGTHS AND WEAKNESSES:\nThe assumptions of explicit and differentiable context limit the applicability of the method.\n\nQUESTIONS:\nZhao et al. 2020 was inappropriately cited in Sec. 5 line 223. Domain randomization was indeed mentioned in the survey, but authors should instead properly acknowledge the origin of the method, instead of citing a survey that mentioned the method, let alone the fact the survey was already cited several times before this.\n\nLIMITATIONS:\nThe experimental contexts seem to be all 1D. Any additional experiment that goes beyond scalar context values could be a powerful argument for the scalability of the method.\n\nI think figure 1 can be better designed / scaled. Currently, it takes too much space in the main manuscript and it shows quite little.\n\nThe work also would benefit from a comparison with \"Consciousness-Inspired Spatio-Temporal Abstractions for Better Generalization in Reinforcement Learning\" (by Bengio's group), which heavily used discrete combinatorial contexts to demonstrate the generalization abilities of agents and the scalability of the methods.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-14"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a method to enhance zero-shot generalisation of RL methods by exploiting regularity in the set of environments the policies are trained in. The regularity is modelled through contextual markov decision process (CMDPs), and specifically regularity in the reward and transition functions wrt. the context parameters. The work introduces a contex-enhanced bellman equation (CEBE) which improves generalisation from training in a single context by updating the model on similar contexts through exploiting the regularity. The effectiveness of CEBE is demonstrated theoretically and empirically, and a data augmentation method (CSE) is derived from CEBE that approximates CEBE in control environments. CSE is compared empirically to baselines in these environments and performs well.\n\nSTRENGTHS AND WEAKNESSES:\n# Strengths\n\n* The paper is well-written, clear and easy to understand\n* The theoretical work is clear and an interesting contribution\n* the topic of the work is important and significant, the contribution novel, and the methodology for improving generalisation interesting and promising.\n* The empirical validation of the theory in figure 1 is beneficial, and the method being relatively close to the theoretical one is also beneficial in aiding understanding of the method's success.\n\n# Weaknesses\n\nThe first weakness of the contribution is in the somewhat minimal empirical justification of the method's benefits. Relatively few environments are experimented on, and only 1 meaningful baseline is compared too, which makes it difficult to assess whether the proposed method is an improvement over the current state-of-the-art in this space. A more thorough comparison against baselines in a wider set of environments would be beneficial.\n\nA second weakness is that the method relies on having direct access to the context and to the derivates of the reward and transition function wrt. the context. These assumptions are quite limiting as to the applicability of the method, and hence it would be beneficial to discuss whether they can be mitigated through further approximations, or whether they are fundamental to the method.\n\n# Summary\n\nOverall, I think the paper is just worthy of acceptance currently. Addressing the two weaknesses above would help me raise my score.\n\nQUESTIONS:\ncovered above.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper tackles zero-shot context generalization in contextual Markov decision processes (CMDPs) when only one (or very few) training context(s) are available. The main technical idea is to linearize the CMDP’s transition and reward functions around the training context, yielding the Context-Enhanced Bellman Equation (CEBE). Solving CEBE produces a Q-function that the authors prove is a first-order approximation of the true Q-function in nearby, unseen contexts. From CEBE they derive Context Sample Enhancement (CSE). This simple data-augmentation rule perturbs tuples in the replay buffer using analytic context gradients, letting standard off-policy RL algorithms train “as if” they had sampled additional contexts. Theoretical results establish that CEBE yields a first-order approximation to the Q function under smoothness assumptions, and show that an $\\epsilon$-optimal policy for CEBE is $(\\epsilon+\\delta)$-optimal for the original CMDP. Experiments range from tabular Cliff-Walking to classic-control, a toy linear system, and Mujoco benchmarks. Across tasks, CSE matches or outperforms vanilla training and closely tracks an oracle LDR baseline.\n\nSTRENGTHS AND WEAKNESSES:\n- Strengths: Theorems quantify approximation error and policy sub-optimality with proofs provided. Paper is well structured with clarity. This paper addresses the problem of generalization from scarce contexts relevant to sim-to-real robotics, meta-RL and safe deployment and shows that analytic structure can substitute for expensive domain randomization.\n- Weaknesses: The analysis requires strong assumptions that the agent must observe $\\delta_c R^C$ and $\\delta_c T^C$ and the transition is deterministic. Real-world simulators rarely expose these gradients. The impact of this paper may depend on availability of context gradients; without them, benefits may vanish. In terms of originality, contribution is an extension to context rather than a fundamentally new paradigm.\n\nQUESTIONS:\n- In many simulators like MuJoCo, the gradients are not exposed. Can CSE leverage automatic differentiation through the physics engine, or will finite-difference estimates suffice?\n- What is “Sample-enhanced transition and reward” in line 142? same as Context-enhanced transition and reward?\n- Just out of curiosity. I understand you assume some deterministic transitions for the theory. Can you validate CSE experimentally with the case with stochastic transitions? Also, Lipschitz assumption may not hold in real-world CMDPs. Do you have any examples with that?\n- Your theory gives an $\\mathcal{O}(\\|c-c_0\\|^2)$ error, but experiments did not quantify how far from $c_0$ that bound remains tight. Could you plot $\\| Q_{\\text{CE}}-Q^\\star \\|$ versus context distance or derive an explicit radius $⁡r_{\\max}$ beyond which linearization degrades?\n- All experiments couple CSE with SAC. Does the technique transfer to policy-gradient methods without a replay buffer (e.g., PPO/TRPO), on-policy actor–critic, or value-based DQN in discrete domains? Even one small-scale test would clarify the generality of CSE beyond off-policy continuous-control.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}