Research map / Reinforcement learning / Visual and representation RL
Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes the context-enhanced Bellman equation (CEBE) and context sample enhancement (CSE), a data augmentation that linearizes dynamics and rewards in context to generalize from a single training context, with theory showing first-order accuracy.
Problem
Zero-shot generalization of deep RL policies to unseen context parameters when only a single (or few) training contexts are available in smooth contextual MDPs.
Approach
Linearize transition and reward functions in the context parameter around a training context to define a context-enhanced Bellman equation, and implement it via CSE, a data augmentation that transforms samples from the base context into samples from nearby contexts using context gradients (deterministic transitions).
Key results
- In tabular Cliffwalking, CEBE Q-function approximation error scales as O(||c-c0||^2), with fitted log-log slope about 2 for both reward choices.
- On SimpleDirection, CSE performs similarly to local domain randomization (LDR) and much better than the baseline SAC when varying the first context parameter.
- On PendulumGoal, both CSE and LDR consistently outperform the baseline; CSE beats LDR when goal torque exceeds 0.6.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use CSE instead of baseline SAC | SAC · SimpleDirection | return | Helped |
| Use CSE instead of baseline SAC | SAC · PendulumGoal | return | Helped |
| Use CSE instead of baseline SAC | SAC · CheetahVelocity | return | Helped |
| Use CSE instead of baseline SAC | SAC · AntDirection | return | Helped |
| Use CSE instead of LDR | SAC · SimpleDirection | return | No clear effect |
| Use CSE instead of LDR | SAC · PendulumGoal | return | Mixed |
| Use CSE instead of LDR | SAC · CheetahVelocity | return | Mixed |
| Use CSE instead of LDR | SAC · AntDirection | return | Mixed |
| Use baseline SAC instead of CSE | SAC · CheetahVelocity, goal velocity greater than $2.6$ | return | Helped |
Surprising findings
- CSE outperforms LDR in some PendulumGoal contexts (goal torque above 0.6), despite LDR using exact data from perturbed contexts.
- On AntDirection, CSE achieves lower returns than LDR in the region [3.5, 5].
Limitations
- CSE is restricted to deterministic transitions; stochastic case only handled theoretically via projection/importance sampling ideas.
- Analysis assumes fully observable states and contexts and access to context derivatives of transitions and rewards.
- Experiments are simulation-only with a relatively small set of environments.
Experimental setup
Datasets: SimpleDirection Cliffwalking PendulumGoal CartGoal CheetahVelocity AntDirection
Benchmarks: Cliffwalking SimpleDirection PendulumGoal CheetahVelocity AntDirection
Metrics: approximation error (sup norm) mean return with 95% confidence interval
Baselines: SAC local domain randomization
Builds on: Soft Actor-Critic domain randomization sample enhancement
Closest papers by meaning
- Zero-shot Model-based Reinforcement Learning using Large Language Models (ICLR 2025 · similarity 0.90)
- XLand-100B: A Large-Scale Multi-Task Dataset for In-Context Reinforcement Learning (ICLR 2025 · similarity 0.89)
- Model-Based Transfer Learning for Contextual Reinforcement Learning (NeurIPS 2024 · similarity 0.89)
- Towards Large-Scale In-Context Reinforcement Learning by Meta-Training in Randomized Worlds (NeurIPS 2025 · similarity 0.88)
- In-Context Reinforcement Learning for Variable Action Spaces (ICML 2024 · similarity 0.88)
- Vintix: Action Model via In-Context Reinforcement Learning (ICML 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.