Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts
Authors: James Chapman
Kedar Karhadkar
Guido Montufar
Venue: NeurIPS 2025 poster ()
Keywords: Reinforcement Learning
Zero-Shot Generalization
Contextual MDPs
Robust Machine Learning
Primary area: reinforcement_learning
TL;DR: We propose a data augmentation method for zero-shot generalization in reinforcement learning and provide a theoretical and empirical analysis of the method.

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.

Zero-Shot Context Generalization in Reinforcement Learning from Few Training Contexts
"James Chapman
*Kedar Karhadkar
Department of Mathematics
Department of Mathematics University of Califoria, Los Angeles
University of California, Los Angeles
Los Angeles, CA 90095 2 @math. ucla. edu
Los Angeles, CA 90095 kedar@math .ula..
Guido Montúfar
Departments of Mathematics and of Statistics & Data Science University of California, Los Angeles
Los Angeles, CA 90095 montufar@math. www..ccccccoooooo edu
Abstract
Deep reinforcement learning (DRL) has achieved remarkable success across mul- tiple 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 addi- tional structure in the problem. However, obtaining sufficient training data across diverse contexts can be impractical in real-world applications. In this work, we con- sider contextual Markov decision processes (CMDPs) with transition and reward functions that exhibit regularity in context parameters. We introduce the context- enhanced Bellman equation (CEBE) 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) an efficient data augmentation method for approximating the CEBE deterministic control environments. We nu- merically validate the performance of CSE in simulation environments, showcasing potential to improve generalization in DRL
Introduction
Deep reinforcement learning has been deployed successfully in many problems requiring decision making and multi-step optimization. Despite these successes, reinforcement learning still struggles to perform well at test time real-world scenarios. This is a well documented phenomenon with multiple causes (e.g., sim-to-real distribution shift, limited contextual data) (Zhao al. 2020 Ghosh et al. 2021 Kirk et al. 2023). This places aut-of-distribution generalization as a fundamental challenge of reinforcement learning.
One approach to address generalization in reinforcement learning is to use continual learning in which some training occurs at deployment in the testing environment (Khetarpal et al. 2022 Wang
al. 2023). However, this is not always feasible or desirable, as, for instance, safety considerations
'Equal Contribution.
Code: http:///sttccccccoeee//////eeeee
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

may prevent early agent deployment, and online adaptation may be prohibitively expensive (Gu al. 2024 Groshev al. 2023). Instead, we need ensure that the agent has good zero-shot generalization from only the training data. Another approach is to create many training environments to prevent over-optimization on specific context parameters (e.g., friction coefficients) as in Meta-RI. RI. and domain randomization (Beck et al. 2025 Chen et al.. 2022). This is commonly used in robotics where sim-to-real pipelines are necessary, and it has been found to produce policies that are more robust to parameter uncertainty (e.g., less sensitive to distribution shift) (Tang et al. 2024 Zhao ctt al. 2020p. However, the environment construction may be prohibitively expensive. For example, designing and building complex robots requires substantial engineering labor and expensive hardware.
Without sufficient structure and prior information, zero-shot generalization from only a few contexts is an impossible problem. However, in many domains we have knowledge of the underlying dynamics of the environment. For example, when solving control problems we may have equations for the dynamics and only uncertainty in the parameters of the equations (Kay) 2013). In such cases it may be possible to include knowledge about the general form of the dynamics into the architecture and/or training algorithm without ever sampling from these unseen contexts. Some works build architectures which try to capture such biases in the underlying dynamics (Wang et al. 2018 8 Kurin et al. 2021 Huang et al.) 2020). Although this approach tends to improve training performance, the generalization benefits have remained limited. The authors of these works note that the variance on the evaluation set is still relatively large. Attaining consistent improvements in
contexts typically requires developing complex architectures which are specifically tailored to the problem at hand (Hong et al. 2022 Chen et al.
2023 Xiong et al. 2023 Li et al. 2024). This is valuable approach, but is challenging to scale to larger sets of environments as it requires domain specific information.
1.1 Our Contributions
In this paper we develop an approach to incorporate structure across contexts into the training algorithm in order to improve generalization performance on previously unseen contexts. We introduce the context-enhanced Bellman equation (CEBE), which is the Bellman equation of a contextual Markov decision process (CMDP) with transition and reward functions linearized about a particular training context. This allows us to estimate data from the CMDP in a neighborhood of the training context. We also derive context sample enhancement (CSE) an efficient method to generate data from nearby contexts in the linearized CMDP that we can then use to optimize the CEBE in deterministic control environments. This allows us to enhance the data sampled from single training context to effectively train on nearby contexts. Our contributions are as follows:
1. We provide a theoretical analysis of context perturbations in contextual Markov decision pro- cesses (CMDPs), demonstrating the viability of our proposed CEBE method for improving out-of-distribution generalization.
2. We derive CSE as an efficient data augmentation method for deep reinforcement learning (DRL), enabling more robust policy learning when learning from samples generated in the training context of the CMDP.
3. We perform experiments with CEBE and CSE on a variety of RL environments that exhibit regularity in the transition and reward functions with respect to the context parameters, which demonstrate the utility of our methods to improve diittiibiititiiooonnn generalization in context-limited training scenarios.
This paper is organized as follows. Section provides an overview of relevant background. Section introduces CEBE and our theoretical analysis of CEBE. Section Sattiion[iiitoddo introduces CSE and provides an algorithm for data augmentation with CSE. Section shows our experimental evaluation of CEBE and CSE Section 6 offfers discussion and concluding remarks.
1.2 Related Works
Here we briefly discuss some of the most closely related works and defer more extensive overview of related works to Appendix A A related strategy for introducing knowledge of the dynamics into RL training has been considered in previous work focusing on the action space. Qiao et al. (2021) considered gradients with respect to actions and introduced sample enhancement and policy enhance- ment to augment samples from the replay buffer. We instead focus the context variable towards

---

addressing the problem of context generalization, for which we develop theory and experiments. In the stochastic setting, our proposed CEBE can be estimated using importance sampling. This would then be related to Tirinzoni et al. (2018), who considered importance sampling for improving out-of-distribution zero-shot generalization in RI.. That work considers exact transitions and rewards and must therefore consider information from nearby contexts.
Modi al. (2018) assume continuity in the context space and propose a covering algorithm based on zeroth order approximation of the policy in context space to derive PAC bounds. This requires access to many contexts, but we note that our first-order approximation provides better coverage of the context space. Malik et al. (202T) consider CMDPs which have transitions and rewards that are uniformly close in a neighborhood to provide lower bounds on the number of context queries required for generalization. Our setting differs because we only sample one context, we consider first-order information, and we obtain approximation results under more relaxed assumptions on the state and action spaces. We also provide a practical algorithm for achieving generalization comparable to domain randomization. Levy and Mansour (2022) derive polynomial sample complexity bounds when one has access an ERM oracle and is able to sample from nearby contexts. This differs from our work as we seek to approximate the dynamics and reward functions of nearby contexts without sampling.
2 Contextual Markov Decision Processes
Sequential decision making is often modeled as a Markov decision process (MDP), and deep reinforcement learning is employed solve complex MDPs. However, the MDP framework does not directly account for changes in the underlying context, and policies trained via DRL tend to overfit to context parameters. Contextual MDPs (CMDPs) extend MDPs by including context parameters explicitly in the MDP dynamics and rewards, which allows one to more directly account for their impact. The context parameters can include a variety of aspects, particularly parameters that influence the transition dynamics and parameters that specify task by modifying the reward function.
We recall the relevant definitions and notation, consistent with Beukman et al. (2023), that we use to formulate our results below. A CMDP is defined as
M = (C,
where C the context space, S is the state space, A is the action space, and 7 E [0, 1] is the discount factor. Here M' is function that takes any context E C to corresponding MDP
(c)) = (S, RR,
where T and R. are the transition and reward functions of the context-c MDP. For specific E C, T is a map S x A - A(S), where (SS) is the space of probability distributions on S. In general the state and action spaces could depend on but we restrict to common state and action space in this paper. We consider policies of the form π
- (AA). Given an initial state distribution, we define the expected discounted return in context e when using policy at |c, s) as
( c) = E Eep. TT T
Σ 7 R
(3)
$20
For set of contexts C C, we say that policy n is (U, (U)otimll if
((T ((p.). -,
for all E U and all policies p: S + (AA). Given test context distribution Diest with support S C C, the objective is to find a (S, <oottiml policy. In this paper, we consider the -uttof-distribution setting where the train and test context distributions, Dtrain and Diest have different supports.
Structural Assumptions We will consider the setting where the agent must generalize from single context and particular the policy must generalize to contexts that are unseen during training. For this to be possible we must add structure to the CMDP. We place the following assumptions:
1.
The state space S and action space A are metric spaces. The context space C is an open convex subset of a Banach space.
3

---

2. The reward function R S x A S - R is deterministic.
3. The partial derivatives 0. and ².T exist and are Lipschitz.
4. The context E C is fully observable and fixed over an episode.
5. When making transition (s.a) s', the agent observes (s, and 0 (s( (7 a).
Settings of this form commonly appear in physical environments where the underlying dynamics are determined by system of differential equations with physical parameters (e.g.. friction coefficient, elastic modulus, mass) that may be modeled in terms of context
3 Perturbative Theory of CMDPs
3.1 Context-Enhanced Bellman Equation (CEBE)
QBE(S, a, c) (xxTT=(s,a) [R + ( ,',
(5)
is a fundamental equation in RL which is used to train the Q-function. However, the transitions and rewards may not known all contexts. We fix a base training context ep and suppose that we have access to approximate transition and reward functions TCE and Ree satisfying:
= RC
R
TE
ICE R
for all in neighborhood of co. We refer to such approximations the context-enhaneee transitions and context-enhancee rewards, respectively. We define the context-cnhanced Bellman equation (CEBE) as the Bellman equation of the approximate CMDP:
QcE(s,a,C (grTTTr(s.a)) (Rêr (e ,)),
Note that Co this coincides with the original Bellman equation. While theoretical results further below apply in more general settings, in this paper we will let RoE be the Taylor approximation of R about co and primarily consider the following cases with corresponding definitions of TeE.
Deterministic Transitions For the setting of deterministic transitions, there exists function (s a) such that (s, a) = 8y((s.a) a Dirac delta distribution for each 8, a, c. In this case, we define the context-enhanced transition and reward functions
(e a)
RC = R + R (c co) + 0. R a
(8)
This is particularly useful in the online setting Here we additionally assume that the reward function is differentiable in the next state to compute its linear approximation (i.e., 0, R exists and is Lipschitz).
simply estimate the CEBE from samples.
Transitions with Differentiable Measure If + (s, a) is a differentiable map, then we can define the context-enhanced transition and reward functions by
(s a) P (s. a) + T
(s, a) (c cu))
(9)
RCg = R + 0 (c- (o)
(10)
where P(k) at of the signed measure μ.. This projection is required since the linear approximation may produce
T denotes projection onto the probability simplex and 1²t denotes the positive part
measure that negative in places. If Ac = c co is sufficiently small, the projection well-defined and sends the linear approximation probability measure.
3.2 CEBE Approximation of the Bellman Equation
In this subsection, prove that the context-enhanced transitions and rewards we introduced above lead to Q-function QCE which close to QBE. For our approximation results, we introduce the following notation spaces of functions and measures. If X and Y are metric spaces and

---

f: X + Y is Lipschitz, we denote its Lipschitz constant by L.. If X is a metric space, we let Lip(X) denote the Banach space of Lipschitz functions X equipped with the metric
- f(v)
[|]l = max
[(x)], sup sup =EX
d(x,y) v)
2.9EX
aFy
For a metric space X, let (eas(X) denote the space of finite signed Borel measures on X equipped with the total variation norm. If X is metric space and p E oo), let (X) denote the space of probability measures on X, equipped with the Wasserstein distance
(f u(a,,))
1/P
Wp(,, v)
inf yEl((a))
Here [(,, v/) denotes the space of probability distributions on X x X which have marginal distribu- tions y and v. For p oo, we define (1,, limp-+oo W v).
Now we present our main result, which establishes that Q-functions are stable under small perturba- tions the transition dynamics and rewards.
Theorem ((T, stasbiliy of the Q-function). Let R(), R??) Lip(S A) be reward functions with (R - ||
xo I ōR. Let T(), ( 8 A Wp(S) be transition functions with W W (s, a), ( (8, a)) S år,
sup (s,0)ESxA
and let π : C + Wp(A) be a Lipschitz policy. Let y E (0,1) be a discount factor with < maxi ((2) - r(2) [1- Let Q(1) and Q(2) denote solutions of the Bellman equation for (T(1), ,) and (T(2), 1((2), 7) respectively. Then
1 S - ôn OR++ 7(1 18T (2 [Lip((XA))
1- max (1, Lo (1 + Iox))
Proof sketch. Using that ((2) and (22) are solutions to their respective Bellman equations, we can decompose
= (R( R(?)) (7A()
+ -2)))(11))))))
where (1) and A12) are transition operators for the dynamics the two MDPs under policy и. The first term is small if R(')) and R222 are sufficiently close. The second term can be bounded in terms of the norms of (1) and Q - Q(2), The third term can bounded by showing that A and A
are close, and showing that the learned function (?2) is sufficiently smooth. Combining these bounds and using& triangle inequality, obtain the desired bound [|Q(1) Q(2) lloo-
The full proof for Theorem is contained in Appendix B.2 Next we apply the theorem in the setting where (1) and ((2) represent different contexts a CMDP. We first consider the case where the transition dynamics are deterministic and we context enhancement in and (8). For this case, we assume that and A are Banach spaces, so we can differentiate with respect to states and actions.
Theorem (Deterministic CEBE is first-order accurate). Consider Lipschitz policy S C (A)) and deterministic T. Suppose that the discount factor " satisfies
1
DT + |D2TT (1 Lx)
Let the context-enhanced transitions and rewards be defined by 7 and respectively. Then
||QER - 1i S ||c (lc-col
(ID2RIIo 7(1 + LI 2 (I + =IDRRl(=))
1-7
1 1- max DT oo L.))
This establishes that for fixed policy म, the Q-function of the CEBE approximates the Q-function of the original CMDP with O[Ic- 2) error, provided that the transition and reward functions are sufficiently smooth. prove Theorem in Appendix B.3 Next, we establish an analogue of Theorem 2 for the case where the transition map is stochastie.
B.3

---

Theorem 3 (Stochastic CEBE is first-order accurate). Consider a Lipschitz policy 8 x C - Wy(A) and stochastic T. Let co. C with ||c coll < || | and
1
diami()) + ||c - col La. ( Lz)
Let the context eateext-nhanced transitions and rewards be defined by ( and (10), respectively. Then
||c col2 (10ERII 3y dydim L. (x
||QER Qf I
1 1-7
1 1-y max diam(S) Lx))
Here LT denotes the Lipschitz constant Tf as a map S AxC + Meas(S).
We prove Theorem 3 in Appendix B.4 The above two theorems show that training with CEBE does not incur too much approximation error versus training with BE a neighborhood of the training context. In particular, we can use CEBE to approximately solve the Bellman equation. The following theorem complements these results by showing that a policy optimized using CEBE will also be close to optimal on the original Bellman equation.
Theorem
Let TCE be an L-Lipschitz policy which is (C, c)-otimall with respect to the CEBE, and suppose that there exists an I--lipschitz policy THE which is (CC,))-ptimad with respect to the BE. Suppose that
[ m;) (,,
QBE ( 1l <
for all с E C and all L-Lipschitz policies n, and suppose that the reward function R. is bounded. Then TCE is (C, 28 -optimal with respect to the Bellman equation.
The proof of the above theorem uses the observation that two functions are uniformly close, then a good optimizer of one is also good optimizer of the other. We provide a proof in Appendix B.5
4 Context Sample Enhancement
In this section, we show how one can estimate the CEBE from samples in the deterministic-traniinn case. Suppose we are given dataset 'D consisting of samples of the form (s, a, r, 8'). We introduce the following context sample enhancement CSE procedure, which takes sample and a context perturbation Ac
co inputs and returns context-enhanced sample:
CSE((s,a,,, 8'), Ac) (r R (s, a, 8') Ac Og. R ('),,s')..TT Ac, O.T Ac). (11)
If (s,a,r, 8) is sampled from the CMDP with context co, then CSE((s, (f, 8') from the approximate CMDP in context In particular, CSE allows us to sample from the approximate CMDP at perturbed contexts by performing data augmentation on samples from the original CMDP at base context. The samples generated via this procedure might have lower quality than those generated by exact domain randomization, due to approximation errors, but on the upside, when the derivatives are available, CSE offers a very easy to implement way to integrate structure information into the training. We provide regularization perspective on CSE in Appendix B.6 The full algorithm for CSE in deterministic environments is shown in Algorithm I
co) is sample
5 Experiments
In this section, we perform numerical experiments to test CEBE and CSE. For continuous control problems, we use simple feed-forward network and train with Soft Actor Critic (SAC) (Haarnoja
ct al. 2018 The inputs to the neural network are states with the context appended on (i.c., (s, c)). We compare against bascline training (i.e., vanilla SAC) and local domain randomization (LDR).
LDR is a popular method robotics and can provide significant generalization benefits (Jakobi et al.
1995 Levy and Wolf 2015 Sadeghi and Levine 2016 Tobin et al. 2017). When using LDR, one first samples a context that is a perturbation of the original training context and then produces a trajectory from the perturbed context in order to broaden the training context distribution. Since LDR has access to exact training data in a neighborhood of the original training context, we treat LDR as a idealized benchmark method for comparison.
In Section 5.1 we empirically demonstrate that the Q-function obtained by dynamic programming with the CEBE is first-order accurate, confirming our theory in Section In Section 5.2

---

Algorithm Off-policy RL algorithm with context sample enhancement
1: Given: CMDP M. training contexts Dirain, data collection iterations N, train iterations M, perturbation radius and off-policy RL algorithm ALGO.
2: Initialize policy æ. value functions Q, and replay buffer B.
number of trajectories from random poliey in CMDP (' with Dirain 4: for N iterations do
3: Collect
5: Sample Dirain
6: Collect a trajectory {(1;, C, az rt, 81-1)120 from Mc(c) using and store in buffer B
7: for Some number of training iterations do
8: Sample a batch {(s}, c², aas T2, Fi from buffer B
9: Generate perturbations Act E B(c', E) and compute (F. 3²) = CSE(x', Ac')
Update samples x² (s, Ac", a, (3) and train with ALGO on the updated batch
11: end for
end for
empirically demonstrate improved performance with compared to baseline training relatively simple continuous control environment. In Section 5.3 we test on goal-based classic control environments. In Section 5.4, we test on MuJoCo environments with context dependent tasks introduced the work Lee and Chung (202I). In all continuous control environments, we train 10 policies and compute the average return over 64 trajectories. We then report the mean and 95% confidence interval the mean returns for each policy. The full list of hyperparameters and experiment configurations are included in Appendix D
5.1 Tabular CEBE
We begin by examining a tabular setting where we show that the CEBE a first-order approximation of the Bellman equation. Using dynamic programming allows us to exactly solve the Bellman equation and the CEBE and thus avoid noise due to sampling and training in deep RL. We consider the tabular Cliffwalking from gymnasium (Towers et al. 2024), in which agent must navigate grid-world to goal state without first "falling" off the cliff into terminal state. At each step, the agent has some probability of slipping into an adjacent state. To introduce nonlinearity into the reward function, we consider the following two choices:
100
-10 I+
R
R goal
Raur
: = (1+
(12)
In each case, we let It be equal to Rinr the agent falls off the cliff, Reoal it achieves the goal state, and otherwise. In Figure we plot the approximation error ||Q - Qiell on a log-log scale for 100 choices of and fit line to the first 10 points show that Q has approximation error of O(Icc- cp )1²) We see the best-fit line has a slope R 2 for both choices of the reward function. This numerically demonstrates Theorem 3 and shows that O is first-order approximation of QéE-
True Erroc
Teae
10-1
2.057
30"1
10
10"4
30'
13.2
10
12.7
30"
(a) Reward in eq.
left
(b) Reward in eq. right
Figure 1: Approximation enror of CEBE on Cliffwalker environment with different rewards. This experiment uses 5 rows, 6 columns, e 0.1, and 7 = 0.9.

---

5.2 Simple Control Environments
We begin our study of continuous control problems with environments with linear transitions and rewards. We let SimpleDirection denote the environment with transition and reward functions
(8, ca [8 ((s,,,8) a,,', 8') = s''-c.
where = RR? -- and 80 ~ Uniform ([-1, This environment was chosen because it has a reward function similar to ones we will test later in the MuloCo environments, but has much more simple dynamics. SimpleDirection, the agent is incentivized to move in the direction by picking action a sign(c). Figure S shows the evaluation results for policies trained on SimpleDirection with Baseline, CSE, and LDR as we vary the context parameters. observe that CSE performs similarly to LDR and much better than the baseline. We provide additional analyses of this environment in Appendix
80
SAC CSE
SAC LDR
Mean Naturn
cc-
40
20
-1.00
-0.75 -0.50 -0225
0.00
0.25
0.50
0.75
1.00
Context
Figure 2: Comparison of training methods as we vary the first context parameter in SimpleDirection.
5.3 Classic Control
We consider a goal-based extension of the Pendulum environment (Towers et al. 2024). which denote PendulumGoal. In addition to the physical parameters, we add goal state to the context to increase the task complexity. We let c (g. m, L, 7). where g the gravitational accelcration, m the mass, / the pendulum length, and T E [-1, 1] the desired torque at the goal state, and set
Agoal A + 0.1 + 0.001 u2,
Rd = 2 sin
(13)
2
where u is the action and fooal = sin (-27 2r/mgl). Figure 3 we plot the evaluation results for each method baseline. In some contexts, CSE performs much better than LDR, e.g. when the goal torque is greater than 0.6. We present the results of varying the other context parameters in Appendix C We also present additional experiments with a goal-based CartPole environment (CartGoal) in Appendix
vary 9 and the goal torque Both CSE and LDR consistently outperform the
5.4 MuJoCo Environments
For this section, we use the goal-based MuJoCo environments Cheetah CheetahVelocity and AntDirection introduced in the work of and Chung (202)) These environments are based on HalfCheetah and Ant environments from Todorov et al. (2022). The Cheetah VelcctaV environment uses modified reward function that rewards the cheetah for running at a velocity specified by the context. The AntDirection environment uses modified reward function that rewards the ant for running in
direction specified by the context.
In Figure 4a we show the results of the Cheetah Velocit experiment. When sweeping over the goal velocity, CSE and LDR outperform the baseline for most of the contexts. CSE performs similar to LDR. When the goal velocity is between 0 and 2, CSE and LDR perform nearly as well as on the training context, while baseline starts to degrade. In Cheetah Velocity, although baseline performs best when the goal velocity is greater than 2.6, CSE still outperforms LDR. In AntDirection, shown in Figure 4b CSE performs similarly to LDR in most contexts, though it achieves lower returns in the region [3.5, 5].

---

-200
-400
-250
Mean Return
Mean Return
300
-350
-1000
SAC
-100
SAC CSE
CSE
-1200
SAC
LDR
SAC LDR
85
1.0
20
25
3.0
35
4.5
5.0
10
15 30-0.22
0.00 0.25
0.50
0.75 1.00
Context
Contoxt
Gravitational Acceleration
(b) Goal Torque
Figure 3: Comparison of training methods on PendulumGoal evaluating on different gravitational acceleration and goal torque parameters. All context parameters (g, m,,, 7) are perturbed during training.
1200
SAC
SAC
300
1000
SWC LDR
Mean Retur
Mean Return
800
Irain
1000
-1500
400
SAC
-200
SAC + CSE
-200
Train
-200
Context
Contaxt
(a) Cheetah CheetahVelocityy
(b) AntDirection
Figure Comparison of training methods on the goal-based MuJoCo environments.
Conclusion
We proposed a perturbative framework for CMDPs as an approach to improve zero-shot generalization reinforcement learning. In particular, we introduced an approximation of the Bellman equation called the context-enhanced Bellman equation CEBE and showed theoretically that it approximates the value functions of the true CMDP neighborhood of contexts. Moreover, we showed that optimizing policy based on our proposed CEBE produces a policy that is nearly optimal in the original CMDP. Using this framework we then introduced context sample enhancement CSE procedure to generate samples that provably approximate samples from unobserved contexts. Finally, we performed experiments in diverse simulation environments. The results suggest that CSE can serve as powerful method for improving generalization in DRL in smooth CMDPs.
We highlight that CSE is easy implement and be easily incorporated with other generalization methods. While we focused on model-free algorithms, CSE is compatible with model-based methods and may aid in training better world models. We think that CSE could potentially also be used to effectively sample from larger volume of the context space in nigh-dimensional context spaces where domain randomization may suffer from a curse of dimensionality. However, further work is needed to understand the sample complexity of CSE in comparison to domain randomization in this setting. Another potential avenue for future study is in applying CSE offline RL where obtaining new samples is impractical. Aside from improving generalization, one could explore different ways to leverage gradient information to design principled variations of prioritized replay buffers and exploration strategies (Jiang et al. 2021 2023 The gradient information of the CMDP could also highlight states in the CMDP which are sensitive to context parameters and allow one to focus more training around these sensitive states use adversarial context perturbations (Mehta et al. 2020).

---

We conclude by pointing limitations of our work. While our theory with CEBE applies in general to smooth CMDPs with deterministic rewards, CSE focuses on deterministic transitions which do not always occur in practice. We primarily consider this case because it provides an efficient data augmentation method. Extensions of this may consider taking gradients of appropriately defined transport maps between distributions. Another limitation is that our analysis and experiments focus on fully-observable state and context spaces, This does not always hold practice and future study should examine the sensitivity of CSE with respect to noisy gradients in context space as well as partially observable environments.
Acknowledgments
This project has been supported by NSF DMS-2145630 and NSF CCF-2212520. GM also acknowl- edges support from DARPA AIQ in project HR001 12520014, DFG SPP 2298 project 46410921 109215, 15, and BMFTR in DAAD project 57616814, 14. We also thank Ruibin Lyu for his contributions during an early stage of the project, including writing test cases, verifying model components, and for helpful discussion.
References
Philip J Ball, Cong Lu, Jack Parker-Holder, and Stephen Roberts. Augmented world models facilitate zero-shot dynamics generalization from a single offline environment. Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 619.629. PMLR, 2021. URL https:// httpps::/pproooednggs mlr, prress
v v139/ba111aa hhml
Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. A tutorial on meta-reinforcement learning. Foundations and Trends@ in Machine Learning, 18(2-3):224-884, 2025. ISSN 1935-8237.
doi: 10.1561/22000000800. URL
http://dx. dow.:oooi..og1000/ ww./oo///............ 1561/200000000
Carolin Benjamins, Theresa Eimer, Frederik Schubert, Aditya Mohan, Sebastian Döhler, André Biedenkapp, Bodo Rosenhahn, Frank Hutter, and Marius Lindauer. Contextualize me the case for context in reinforcement learning. Transactions Machine Learning Research, 2023, ISSN 2835-8856. URL htttts:://pperreiiiiee net/forummid /
Michael Beukman, Devon Jarvis, Richard Klein, Steven James, and Benjamin Rosman. Dynamics generalisation in reinforcement learning via adaptive context-aware policies. In Advances in Neural Information Processing Systems, volume 36, pages 40167-40203, Curran Associates. Inc.. 2023. URL https:: https::/pproceedingg neurips cc/ wttpssccccpppppppeeeeeeee htppp/apper___aa files/ile/eeeee (e177781002244333355770707777600 Pwperraoofeeeeeeeeeeeeeeee Conference. wtt.fffeese/pppe////////i/
Djallel Bouneffouf, Irina Rish. and Charu Aggarwal. Survey applications of multi-armed and contextual bandits. In 2020 IEEE congress on evolutionary computation (CEC), pages 1-8. IEEE, 2020.
Runfa Chen, Jiaqi Han, Fuchun Sun, and Wenbing Huang. Subequivariant graph reinforcement learning 3D environments. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 4545-4565. PMLR, 2023. URL (tRpsttts://// /tttps::/proceeednnnng mlr. mlr. press/ press/v202/ w002//2222233../22 chen23i html
Shiyu Chen and Yanjie Li. An overview of robust reinforcement learning. In 2020 ICEE International Conference on Networking, Sensing and Control (ICNSC), pages 1-6. IEEE, 2020.
Xiaoyu Chen, Jiachen Hu, Chi Jin, Lihong Li, and Liwei Wang. Understanding domain randomization for sim-to-real transfer. In International Conference on Learning Representations, 2022. URL
https::/openrevver. peet:forrumidd:T82222
Jung-Hoon Cho, Vindula Jayawardana, Sirui Li, and Cathy Wu. Model-based transfer learning for contextual reinforcement learning. The Thirtt-eighth Annual Conference on Neural Information Processing Systems, 2024. URL
10

---

Karl Cobbe, Chris Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Leurning Research, pages 2048-2056. mlr. wttpeeeeeeeee/bbbeeeeeeeeeeeeeee htmml
PMLR, 2020.
Maria Dimakopoulou, Zhengyuan Zhou, Susan Athey, and Guido Imbens. Estimation considerations in contextual bandits. arXiv preprint arXiv:1711.07077, 2017.
Dibya Ghosh, Jad Rahme, Aviral Kumar, Amy Zhang. Ryan P Adams, and Sergey Levine. Why generalization in RL is difficult: Epistemic POMDPs and implicit partial observability. In Ad- vances in Neural Information Processing Systems, volume 34, pages 25502-25515.. Curran As- sociates, Inc., 2021. URL http///tttpss/pppcceeddggggggee htwp://rwc....n.....n..........p neuriss. htpppapcsp.aeipppppppppppprrrrre files/ www.ff/apeeeee///eeeeeeaaaaa /021/ file/d5ff (55.5555577777777777766600002200
Ralf Gommers, Pauli Virtanen, Matt Haberland, Evgeni Burovski, Tyler Reddy, Warren Weckesser, Travis E Oliphant, David Cournapeau, Andrew Nelson, Pamphile Roy, et al. scipy/scipy: Scipy 1.15. O. Zenodo, 2024. doi: https://doi.org/II ort//0.5228 1 14552608.
Milan Groshev, Gabriele Baldoni, Luca Cominardi, Antonio de la Oliva, and Robert Gazda. Edge robotics: Are we ready? An experimental evaluation of current vision and future directions. Digital Communications and Networks, 91):166-177, 2023.
Shangding Gu, Long Yang, Yali Du, Guang Chen, Florian Walter, Jun Wang, and Alois Knoll. A review of safe reinforcement learning: Methods, theories, and applications. IEEE Transactions on Pattern Analysis and Machine Intelligence, 46(12):11216-1222 2024. doi: 10.1 109//AMI 2024.3457538.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with stochastic actor. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1861-1870. PMLR, 2018. URL https://proceeddggs.ut mlr wtt.a/ais..........ee//aa press/v80/ haarnojal8b. html
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual Markov decision processes, arXiv preprint arXiv: 2015,
Charles R. Harris, K. Jarrod Millman, Stéfan J. van der Walt, Ralf Gommers, Pauli Virtanen, David Cournapeau, Eric Wieser, Julian Taylor, Sebastian Berg, Nathaniel J. Smith, Robert Kern, Matti Picus, Stephan Hoyer, Marten H. van Kerkwijk, Matthew Brett, Allan Haldane, Jaime Fernández del Río, Mark Wiebe, Pearu Peterson, Pierre
Kevin Sheppard, Tyler Reddy, Warren Weckesser, Hameer Abbasi, Christoph Gohlke, and Travis E. Oliphant. Array programming with NumPy. Nature, 585(7825):357-362, 2020. doi: 10. 18600//4116-00000--100 https: org/ 10. 1038/ 00445866-2000 2649-2
URL
James Harrison, Animesh Garg, Boris Ivanovic, Yuke Zhu, Silvio Savarese, Li Fei-Fei, and Marco Pavone. Adapt: zero-shot adaptive policy transfer for stochastic dynamical systems. In Robotics Research: The 18th International Symposium ISRR, pages 437-453. Springer, 2019.
Irina Higgins, Arka Pal, Andrei Rusu, Loic Matthey, Christopher Burgess, Alexander Pritzel, Matthew Botvinick, Charles Blundell, and Alexander Lerchner. DARLA: Improving zero-shot transfer in reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pages 1480-1490. PMLR, 2017. URL https://pptt httpes////oceeeeings.... mlr. www.ppressvvvv caaesss/vvv/hha wwt.///g//ginni./ www.ginnn1a..... html
Sunghoon Hong, Deunsol Yoon, and Kee-Eung Kim. Structure-aware transformer policy for inhomo- geneous multi-task reinforcement learning. In International Conference on Learning Representa- tions, 2022. URL http::://opeereeiieeeeeee http://pprrrrieeee..rrrrrrrrrrr/ net/forumiiddfy X_VVHylX
Wenlong Huang, Igor Mordatch, and Deepak Pathak. One policy to control them all: Shared modular policies for agent-agnostic control. In Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages mlr. wwwppppwsssseeeeeepe wttpas////1///uag2dd... html
PMLR, 2020. http:ttttt/:oceeeeeeee
11

---

J. D. Hunter. Matplodiib: A 2D graphics environment. Computing in Science & Engineering, 9()): 90-95, 2007. doi: 10.1 10.110//MCS222000.555
Nick Jakobi, Phil Husbands, and Inman Harvey. Noise and the reality gap: The use of simulation in evolutionary robotics. In European conference on artificial life, pages 704-720. Springer, 1995.
Minqi Jiang, Michael D Dennis, Jack Parker-Holder, Jakob Nicolaus Foerster, Edward Grefenstette, and Tim Rocktäschel. Replay-guided adversarial environment design. In Advances in Neural Infor- mation Processing Systems, 2021. URL httpps:/
Yiding Jiang, J. Zico Kolter, and Roberta Raileanu. On the importance of ex- ploration for generalization in reinforcement learning. In Advances in Neural In- formation Processing Systems, volume 36. pages 12951-12986. Curran Associates, Inc., 2023, URL htttp:://ppoceeeeeeeeeee wewpeeeddinnnnnnunii.seeee neurips. wwtpcccppppe_eeeeee 2243104444444d333266666666333332 93c.55399933353399-9 Pwwpeeeeeeaeeeeeeeaeeeee--- CoanneCooeeeencce pdf
Steven M Kay. Fundamentals of statistical signal processing: Practical algorithm development,
volume 3. Pearson Education, 2013.
Khimya Khetarpal, Matthew Riemer, Irina Rish, and Doina Precup. Towards continual reinforcement learning: review and perspectives. Journal of Artificial Intelligence Research, 75:1401-1476, 2022.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Research, 76: 201-264, 2023.
Vitaly Kurin, Maximilian Igl. Tim Rocktäschel, Wendelin Boehmer, and Shimon Whiteson. My body is cage: the role of morphology in graph-basede incompatible control. In International Conference on Learning Representations, 2021. URL, http:ttt////peeeeeiieee
Mikko Lauri, David Hsu, and Joni Pajarinen. Partially observable Markov decision processes in robotics: A survey. IEEE Transactions on Robotics, 39(1):21 2023. doi: 10. 10.1109/TR0.202. 3200138.
Suyoung Lee and Yau-Young Chung. Improving generalization in meta-RL with imaginary tasks from latent dynamics mixture. Advances in Neural Information Processing Systems, volume 34, pages 27235. Curran Associates, Inc., 2021. URL https:: Paper pdf
neurips cc/
Ofir Levy and Lior Wolf. Live repetition counting. In Proceedings of the IEEE international
conference on computer vision, pages 3020-3028, 2015.
Orin Levy and Yishay Mansour. Learning efficiently function approximation for contextual MDP. arXiv preprint 2203:2203000099 2022.
Orin Levy, Alon Cohen, Asaf Cassel, and Yishay Mansour. Efficient rate optimal regret for adversarial contextual MDPs using online function approximation. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 19287-19314, PMLR, 2023, URL htctps://proceeeeee mlr. press/ (tress/v202/eev222 html
Boyu Li, Haoran Li, Yuanheng Zhu, and Dongbin Zhao. MAT: Morphological adaptive transformer for universal morphology policy learning. IEEE Transactions on Cognitive and Developmental Systems, 16(4):1611-162 2024. doi: 10. 00.11099TCDD2222433331555
Eric Liang, Richard Liaw, Robert Nishihara, Philipp Moritz, Roy Fox, Ken Goldberg, Joseph Gonzalez, Michael Jordan, and Ion Stoica. RLlib: Abstractions for distributed reinforcement learning. In Proceedings of the 35th International Conference Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 3053-3062, PMLR, 2018. URL httpps //proceedings. mlr, press/ maaass/v80/iie https///o//laaaggbbbb html
12

---

Eric Liang. Zhanghao Wu, Michael Luo, Sven Mika, Joseph E Gonzalez, and Ion Stoica. RLlib Flow: Distributed reinforcement learning is dataflow problem. In Advances in Neural Information Processing Systems, volume 34, pages 5506-5517. Curran Associates, Inc., 2021. URL https::
hew.rediiiinnnnnnisseeeee neurips wttccc//pppp laape__fllss______pppp__ppppppee /ptper_/ileeepppppr///////1/1 /aper/2021/file/
pdf
Richard Liaw, Eric Liang, Robert Nishihara, Philipp Moritz, Joseph E Gonzalez, and Ion Stoica. Tune: A research platform for distributed model selection and training. arXiv preprint arXiv:1807.05118. 2018.
Tyler Lu, David Pal, and Martin Pal. Contextual multi-armed bandits. In Proceedings the the Thirteenth International Conference Artificial Intelligence and Statistics, volume 9 of Proceedings of Machine Learning Research, pages 485-492 Chia Laguna Resort, Sardinia, Italy, 2010. PMLR.
URL https:://ppooeeecc httpps:/pproceedngg... mlr press/ wpwpss/9//lu00000// ht///19911aaammm
Kevin Sebastian Luck, Heni Ben Amor, and Roberto Calandra. Data-efficient co-adaptation of morphology and behaviour with deep reinforcement Icarning. In Proceedings of the Conference on Robot Learning, volume 100 of Proceedings of Machine Learning Research, pages 854-869. PMLR, 2020. htLLhtttps///rrcceddnng mlr presss/v100/luck00a html
Dhruv Malik, Yuanzhi Li, and Pradeep Ravikumar. When is generalizable reinforcement learn- ing tractable? In Advances in Neural Information Processing Systems, volume 34, pages 8032-8045, Curran Associates, Inc., 2021. URL httpps://proceedinngg hww.rdiiinnnnnnnnnneeeeeeeee newrips....... cc/paper_ fi
pdf
Bhairav Mchta, Manfred Diaz, Florian Golemo, Christopher J. Pal, and Liam Paull. Active domain randomization. In Leslie Pack Kaelbling, Danica Kragic, and Komei Sugiura, editors, Proceedings of the Conference on Robot Learning, volume 100 of Proceedings of Machine Learning Research, pages 1162- 1176. PMLR, 2020. URL https::/ .mr...g.m.lr.........m.......mm. prress/vv00/mehha220 html
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing Atari with deep reinforcement learning. arXiv preprint arXiv:13/2.5602 2013.
Aditya Modi, Nan Jiang, Satinder Singh, and Ambuj Tewari. Markov decision processes with continuous side information. In Proceedings of Algorithmic Learning Theory, volume 83 of Proceedings of Machine Learning Research, pages 597-618 PMLR, 07-09 Apr 2018. URL
httops:/pproceeding.. mlr pwtpp/////mmoeeeeeee/pepe htt//////mmmmm.../// html
Aditya Mohan, Amy Zhang, and Marius Lindauer. Structure in deep reinforcement learning: A survey and open problems. Journal of Artificial Intelligence Research, 79:1 197-126,, 2024.
Philipp Moritz, Robert Nishihara, Stephanie Wang, Alexey Tumanov, Richard Liaw, Eric Liang, Melih Elibol, Zongheng Yang, William Paul, Michael Jordan, al. Ray: A distributed framework for emerging Al applications. In 13th USENIX symposium on operating systems design and implementation (OSDI 18), pages 561-577. 2018.
Rhys Newbury, Jack Collins, Kerry He, Jiahe Pan, Ingmar Posner, David Howard, and Akansel Cosgun. A review of differentiable simulators. IEEE Access, 12:97581-97604, 2024. doi: 10.1 /tWp:ACCCESSS.22..44222222
The pandas development team. pandas-dev/pandas: Pandas, February 2020, URL hhttps://do org/100 11111005281zzeoodd 3509134
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary Vito,, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. In Advances in Neural Information Processing Systems. volume 32. Curran As- sociates, Inc., 2019. URL wwwpt:tttttttccccttmm http:/ppppe///ooooceiinn.ees neurips wwwpurrpssccpppppr....e. Fiil://bbe.Btteef/722222ffff2222
pdf
13

---

Sai Prasanna, Karim Farid, Raghu Rajan, and André Biedenkapp. Dreaming of many worlds: Learning contextual world models aids zero-shot generalization. Reinforcement Learning Journal, 3:1317-1350, 2024.
Yi-Ling Qiao, Junbang Liang, Vladlen Koltun, and Ming C Lin. Efficient differentiable simulation of articulated bodies. In Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8661-8671. PMLR, 2021. URL httpps:/prroceedings mlr. www..prees//e... httess/vv39///e v13/////2111//2aaa/
qiao21a. html
Fereshteh Sadeghi and Sergey Levine. CAD2RL: Real single-image flight without single real image. Robolics: Science and Systems, 2016.
Zhenxiong Tan, Kaixin Wang, and Xinchao Wang. C-Procgen: Empowering Procgen with controllable contexts. arXiv preprini arXiv:2311.07312, 2023.
Chen Tang, Ben Abbatematteo, Jiaheng Hu, Rohan Chandra, Roberto Martín,Martín, and Peter Stone.
Deep reinforcement learning for robotics: A survey of real-world successes. Annual Review of Control, Robotics, and Autonomous Systems, 8, 2024.
Vinzenz Thoma, Barna Pasztor, Andreas Krause, Giorgia Ramponi, and Yifan Hu. Con- textual bilevel reinforcement learning for incentive alignment. In Advances in Neural Information Processing Systems, volume 37, pages 127369-127435. Curran Associates, Inc., 2024. URL htttps://proceeddnggg www.riinnnn.n.....e...e........e cc/ppperrcpppppr.reeeeeeepppppee f wileee /ptp///////2///leee #6600aa66bo111101a87888886800 Paper-Conferenoeeeeee pdf
Andrea Tirinzoni, Andrea Sessa, Matteo Pirotta, and Marcello Restelli. Importance weighted transfer of samples in reinforcement learning. In Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 1936-4945 mlr. press//80/tirinzon 1itaooiii1a.... html
PMLR, 2018. httttttt/poooeeeeeeeeee
Josh Tobin, Rachel Fong, Alex Ray. Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/RSJ international conference on intelligent robots and systems (IROS), pages 23-30. IEEE, 2017.
Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model- control. In 2012 IEEE/RSJ International Conference Intelligent Robots and Systems, pages 5026 5033. IEEE, 2012. doi: 10. 109//ROS.200 201..636199,,
Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulão, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: standard interface for reinforcement learning environments. arXiv preprint urXiv:2407.17035 17032, 2024.
Cédric Villani al. Optimal transport: old and new, volume 338. Springer, 2008
Nelson Vithayathil Varghese and Qusay H Mahmoud. survey of multi-task deep reinforcement learning. Electronics, 9393: 2020.
Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler. NerveNet: Learning structured policy with graph neural networks. In International Conference on Learning Representations, 2018. URL
https://openreevex..
Zhechao Wang, Qiming Fu, Jianping Chen, Yunzhe Wang, You Lu, and Hongjie Wu. Reinforce- ment learning in few-shot scenarios: A survey. Journal of Grid Computing, 21(2):30, Jun 2023. ISSN 1572-9184, doi: 10. 10.10077117733-234496666.. Utt.sttts://aoo http://ooo...r///...0.000//
-10723-023 11722-022-09663-0
Michael I.. Waskom. seabor:: statistical data visualization. Journal of Open Source Software 6(60): 3021, 2021. doi: 10.21105/joss.0000 URI, https:://pt.i..... 0rg/10. 21105/joss. 03021
Max Weltevrede, Caroline Horsch, Matthijs TJ Spaan, and Wendelin Böhmer. Training on more reachable tasks for generalisation in reinforcement learning. arXiv preprint 2rXiv:2410.03565 2024.
14

---

Zheng Xiong, Jacob Beck, and Shimon Whiteson. Universal morphology control via contextual modulation. In Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 38286-38300, PMLR, 2023. URL
https://proceedings. mlr wttpreeseeee//eeeeeeeeeeeee
html
Harry Zhao, Safa Alver, Harm van Seijen, Romain Laroche, Doina Precup, and Yoshua Bengio. mnnsciousness-nsspiie spatio-temporal abstractions for better generalization in reinforcement learning. In The Tweifth International Conference on Learning Representations, 2024. URL
Wenshuai Zhao, Jorge Peña Queralta, and Tomi Westerlund. Sim-to-real transfer in deep reinforce- ment learning for robotics: a survey. In 2020 IEEE symposium series on compulational intelligence (SSCI), pages 737-744. IEEE, 2020.
15

---

A Related Works
In this section, we provide more detailed overview of related works and commentary that could not be covered in main text.
Contextual Multi-armed Bandits Multi-armee bandits (MABs) serve as model of reinforcement learning in which the problem horizon is a single step. Contextual Multi-armed bandits (CMABs) are a natural extension of MABs which include contextual information as part of the problem specification. Contextual bandits are a popular modeling tool for a variety of applications including hyperparameter optimization and robotics (Bouneffouf et al. 2020). Lu et al. (2010) introduce clustering algorithm for CMABs over metric spaces with sufficient regularity. Dimakopoulou et al. (2017) propose a method to improve estimation in CMABs and reduce bias for CMDPs with generalized linear reward functions. The assumption of linearity resembles our approach of approximating transition and reward functions with their linearizations.
Challenges in CMDPs CMDPs present new challenges for reinforcement learning as a single CMDP can model potentially infinitely many MDPs. CMDPs can have substantial structure built into the context space that can aid in improving generalization (see the survey by [Moann et al. 2024). In our work, we do not focus specifically on one CMDP's structure, but opt for developing general methods based on first-order information of the context. Some works consider minimal assumptions on the structure of the CMDP by only imposing notion of proximity between different MDPs in the
CMDP. For example, Hallak et al. (2015) consider CMDPs with finite context spaces and introduce a clustering algorithm in context space for training across all contexts. Modi et al. (2018) establish PAC bounds on learning smooth CMDPs. These works primarily consider a zeroth-order approximation and propose covering arguments. This neglects the true changes in dynamics that occur in the CMDP, which we capture through first-order approximations of the CMDP in the context space.
Jiang et al. (2023) examine how better exploration in the training environments can improve gener- alization at test time. Weltevrede et al. (2024) build on this by proposing an additional exploration phase at the start of each episode to broaden the starting state distribution, which can be viewed as an implicit regularization. While these works focus on improving generalization by providing better exploration of the state space, they do not directly account for variations in the context during training.
Other works consider CMDPs where the contexts are limited during training or adversarially chosen. Ghosh et al. (202T) show that contextual MDPs can viewed as Epistemic POMDPs when training on limited number of contexts. Levy et al. 2023) analyze regret in CMDPs with adversarially chosen context, reflecting worst case performance as the context changes. In our work, we consider gradients to generate approximations in all context directions, including potentially adversarially chosen directions.
RL Zero-Shot Generalization Benjamins et al. (0223) identify the need for context in zero-shot generalization and propose a set of environments, CARL, to test generalization in RL. Ball et al. (2021) propose model-based based aata-auummentttion methods for offline RL, which introduce noise into the transition dynamics to reduce overfitting to the observations. This method improves generalization through data augmentation like work, but does not directly account for changes in the context.
Cho et al. 2024) consider related approach which they aim to determine an optimal set of training contexts for generalization. The authors propose an acquisition function that uses a linear function to approximate the generalization gap for performance on out of distribution contexts. Learning this linear model requires access to many contexts, which could be alleviated with access to context gradients. Harrison et al.] (2019) leverage linear approximations of the transition dynamics in the states and actions to improve robustness in model predictive control (MPC) when the agent has access to simulator in the training context and is operating in a different context. Similar to our work, this method relies single training context.
Zhao et al. 2024) use hierarchical planning to design policies which generalize task information across contexts. They consider discrete state, action, and context spaces, evaluating primarily on gridworld environments. They design an architecture which can generalize across discrete contexts, in contrast to our focus the training algorithm and continuous contexts.
16

---

Some works aim to learn representations of the context space. Higgins et al. (2117) ) use B- VAE, and Prasanna et (2024) propose a contextual recurrent state space model, to learn disentangled repre- sentations of the context space. [Beukman et al. (2023) introduce an architecture with hypernetwork which uses the context to generate parameters for other parts of the network.
While several of the aforementioned works use linearity implicitly or gradient information directly, most of these works require access multiple contexts during training. For a more comprehensive survey of zero-shot generalization in RL, we refer the reader to the survey by Kirk et al. 2023
Related areas In multi-task reinforcement learning, trains a policy solve multiple tasks in an environment. This can improve sample complexity since knowledge of the environment can transfer across tasks. For survey of multi-task RI., refer the reader the survey by Vithayathil Varghese and Mahmoud| (2020). Meta-learning seeks to train a policy which can quickly adapt to new contexts. We refer the reader the survey by Beck al. 20255). ). Robust RL aims to improve generalization when model parameters are uncertain and we refer to the survey by Chen and Li (0220).
Another related direction is the simultaneous optimization of the context and policies in CMDPs. optimization methods for optimizing agent morphology in robotics simulations to reduce the number of contexts sampled. Thoma et al. (2024) use bilevel optimization and propose stochastic hypergradients for simultancously optimizing high- and low-level objectives in CMDPs. While these works consider slightly different problems, they may benefit from the first-order methods we consider in our work to reduce the number of contexts sampled.
Luck et al. (2020) consider
Libraries such as Procgen and C-Procgen allow for procedural environment generation (Cobbe et al., 2020; Tan et al. 2023). Not all environments support derivatives in the context variables and even those that do may not support derivatives in the context variables of interest. For more detailed analysis of differentiable simulation, we refer the interested reader to the survey by Newbury et al. (2024).
Due to the inherent partial observability in CMDPs demonstrated by Ghosh et al. (2021) some approaches opt for using methods in POMDPs to study generalization in RL. For applications of POMDPs in robotics, see the survey of Lauri al. (203)..
B Theory
In this section, we prove the results stated in Section
B.1 Preliminaries
First, we introduce some notation and discuss the assumptions make the regularity of the MDP. In the following, p E [1, oo] is constant. Recall from Section that A refers the space of actions available to the agent and W, (.A) denotes the space of probability measures on A equipped with the Wasserstein distance Wp.
Policies: In this article, policy is a Lipschitz map π S + (P() such that for all s E S, x (s) is probability measure. For policy π, we will often write म. in place of (T)).
Q-functions: A Q-function is an element of Lip(S x A). We will often write Q. to denote
the function a + Q(s,
Reward functions: We assume that the reward function R of the MDP is an element of Lip(S x A).
Transition maps: A transition map a Lipschitz mapping T S A + W2(S). For transition map T. we will often write Ts,a in place of ((s.a).
Bellman operator: If S x + (S) is a transition map and π : S + (2(4) is policy, we define the operator AT. Lip(S S x A) - Lip(S x A) by
(A T () := [ / Q (a' (dT) (a' (T.)'TT.a(ss)
(14)
We prove that this operator well-defined in Lemma 7 Observe that the Bellman equation can be written Q R + A (Q where y (0, 1) is the discount factor.
17

---

B.2 Proof of Theorem 1
We need show that the operator A well-defined, in the sense that it maps elements of Lip(S x A) to elements of Lip(S x A), and that is a bounded operator on Lip(S x A). To this end. we prove couple of lemmas establishing the regularity of Lipschitz functions under integration. Lemma 5. Let X be a metric space. Let f Lip(X) and let v1 E (X). Then
1/x Idu - 1xx sdvz S (1,
Proof. First prove the statement for the case p Let >> 0 and let p' be the Hölder conjugate of p. There exists coupling P E [T(1, with
1/8 Urex d(a, P)
< <Wp((1, +c.
Then
1/x fdu - -fx fdvz
1se / (()))T(x.)) J f(u)aT(o.))
J I 15(x) -(())la(x..))
(e 15(=)) - f(()) (T(..))) 1/8 (Je 1 (')a(aa))) 1/8
1/p
v) (W,,((i,,))) + +e,,
where we applied Hölder's inequality in the third line. Since this holds for all > have
| Idv - [dez S L, u2).
For the case co, simply take the limit of the above inequality as - cc.
Lemma 6. Let X and Y be metric spaces. Let f E Lip(X Y), and let : + WY(Y) be Lipschitz with constant Ly- We write ps in place of u(x). Then the function X IR defined by
g(x) = f fodue
Lip(X). Moreover,
sup 19(x)) s sup xex
1f(z, (x,))XXX
and
Lg I L1(11 + Lp)-
Proof. By the triangle inequality,
[g(x)] < sup / Ifeldu xEX
sup
zEX Jy
sup sup (f zeX yEY
Let 01, E X be distinct points. Then
| Jz, dp f Sa duna
(g(x1) r(02)1 =
1 Sz, S2n - fr Say Aundtt
11 fer dpd f Sandtass
<
+
18

---

We bound the terms of the above inequality separately. For the first term,
1 fur dpd - [ feadper
/ If Sex ldper
$ (,d(211222)d)
L,d(211,2) (2).
For the second term, by Lemma 5 we have
|/, far fundues -
- [ fox dutx
- Ly WyyWp
Combining the two terms, we get
(g(x1) - g(x22) & Lyd(21, 02) + 22),
so g has Lipschitz constant L Ln).
Lemma 7 (AT is well-defined). The linear map AT Lip(S x A) + Lip(S x A) given in well-defined. Moreover, for all Q E Lip(S x A),
14
|AT 2l1oo & |1/
and
LAT S LI Lg).
Proof. Let QE E Lip(S A). First, observe that
(AT ()))))(s))))$ S sup
1Q, (dd
sup (a.a)ESAA
(e.0)
I sup
1Q(5'a')),
so |ATT** 11100 < 1101/
Next, let y S IR be defined by
((s) == / Q.dTs.
By Lemma
& E Lip(S) and
Lp I LQ(1 + Ln).
Now
(AT*²)= / ((s') ((s')ddTaa(s))
By Lemma 5
LAT S LighT
S ITQTT(1 T.n).
So A T. maps elements of Lip(S) to Lip(S), and
| . < m1x LT(1 + Ln)).
Next, we show that the regularity of the transition and reward functions implies regularity of the solution to the Bellman equation.
19

---

Lemma 8 (Q is uniquely defined and Lipschitz). Ler y max (1, br(i+t,)) and suppose that R E Lip(S x A). Then there exists unique solution Lip(8 A) the Bellman equation
Q
R R QQ
(15)
Proof. Consider the power series
0C
(AT.r))
S max(1, (T(1 + I-))) Then we have |TAT < 1, so this power series converges absolutely in the operator norm. Then
By Lemma 7
||T
B(I - Σ >
En
n+1 ( (ATTgnnt
n=0
m=0
= I.
The same calculation shows that (I - >--AB 1,
11,00 (I (1 -1- Now Q satisfies the
Bellman equation if and only if
(I (1= R.
Since I - 7/4 is invertible an operator on Lip(S x A), there exists unique Q E Lip(S x A) satisfying this equation.
Lemma (T-stability AT Let ( S Wp(S) transition functiono with
Wp(T(t)s, (2)) 8. (s,a)ESxA
sup
Then
|AT - 1T + S (1 + 1,
Proof. As in the proof of Lemma 7 let Q E Lip(S x A), and let sp S - R be defined by
<()) = / Q.dTa.
By Lemma op E Lip(S) and
Lip S Lo(1 + LR) 11Q1(11 + Ln).
Now by Lemma 5
()), =0)ssa)
1 p(s ((s')d((()) / ((s')d(T!3))
< Li W (T (s, a), (12)
< @Q1((1 + 1.=)
The result follows by taking a supremum over (8, S x A and Q E Lip(S x A).
Proof of Theorem 1
Proof. Let Ai = "T
is and let A = AT(2), - By Lemma
|| A2ll s + Lr)
Since (1) and Q(2) solutions to Bellman equations, have
Q(1) - ((2) = = +)) - +
(R(1)_ (14110)))AA111)))))))))))))))))) + (7)) ) 21
20

---

Then by the triangle inequality,
l!Q Q &cc - || R lcc + +1l - 110 + 7|A4 - -A2Q loc.
< || - (2(2) 10 + +1 A.0
11 + 7114 -2lllQ(() (lip((x))
[R(1) - R(22)
100 + A14 1
Do + x(1 + Iz))rr Lip(SxA)
By assumption, we have || R 1 S 8n. By Lemma
7 we have
|| A1Q 8 1 110 Q 11x0
By the Bellman equation for Q(3),
||Q 2) [1ip(Sx.A)
= |1(I 742) R(21
< |1(T (12) |1|| R2 Lip(sxA)
$ - (1-111A21) (1) R1 [[Ip)(SAA)
By Lemma
||al| S max(1, T-T00) (1 Imm));
(0(00 Lip(Sxa) (1 (1- 111-7/ A2) ( (1) Lip((AA)
(1((2) Lip((xA)
$ y max(1, LT (1 Ln)))
Combining, get
11 -11 + 11 + L R( Lin(SxA)
SR +
1- LITeOs (1 Lx))
and thus
- - 120 I 1->
1
8R + 1- y- max Lx))
711 + L )81 R [ii((SAA)
B.3 Deterministic CEBE approximation error
In this section prove Theorem 2 We first reiterate some of the assumptions implicit in the statement of the theorem and stated before.
The state space and the action space are Banach spaces.
The deterministic transition function T S A x C + WP(S) is of the form T(s, twice continuously differentiable with bounded second partial derivatives. We will write f to denote the function f((, c)
ò f(s,a,c). where f x A x +
The reward function R S A x C -> TR. is twice continuously differentiable with bounded second partial derivatives. We will write R denote the function R( c).
The policy S x C - Wp(A) is Lipschitz.
Lemma 10. Let X be metric space. Then for all 01, X and [1, oo), Wp(6m 832) S d(z1.22).
Proof. First we prove the statement for p < xo. Let I' be the coupling between or, and or, given by the product measure. Then
W Wplden des) S ( d(y1, d(v1,32) dd oz) 1/P
(1 d(y2, d(2r, U2) (11)6 (v2))
1/p
(d(#1, 22)")
d(æ1, v2).
For the case p 00, take the limit + oo in the above bound.
21

---

Proof of Theorem 2 In view of Theorem ! we need to show that the linearized transition and reward functions are good estimates their true values. Let define fce A x C S by
fcr(s, a) = (s, a) + a. (s, a) (c (c-co))
so that (8 a) = 50=$6186ee (s,())
Let 8 E a A, and c C.C. By Taylor's theorem, there exists E |co. c| such that (s, a) (c co)ll = 21l02 (s.a)]lllc coll2
- f
(s,
which we can write as
as) fcu(s, a)ll & ==l
Here 0 denotes the partial derivative with respect to while D denotes the derivative with respect to all parameters. By Lemma
W
a)) & (s. fčn(s, a)ll
s |10 coll?
Hence we have shown that the transition maps are close.
Next show that the reward functions close. Again by Taylor's theorem, there exists o' E [co: c) such that
(s) (s, R. (8, (c (c-ca)l ca)l s \ (s,a)lllee coll?
I | RRloooloe coll?
So
|| - Ro S | Rl coll?
and the reward functions are close.
Next we will show that the transition and reward functions are Lipschitz. By Lemma I0 for all ($1, (1), 02) A, we have
Wp Wy((p(eenaa)))eeeeaa))) S [1(( a1) - - a2)l
< (9. - 82, 01 - -a2]
80 Tore < [D llx Similarly,
W, p (Tég(()), a1). (82 a2))
< [1fEm(1, a1) - fčr 82, a2)l
(1 $2 02)
x (1( f (5,a) 0 0,5 (s) (c co) [((s] 82 a1 -a2)
< || [[s) 52, a2)l || - colll( 82, ai a2)ll
so LT 2 Df + ||= coll-
Since R is twice continuously differentiable, is Lipschitz with constant |R and therefore
| Lip(((x)))))) S max([]/|]lo0)) DR)
For the context-enhanced reward function,
||/eg|lip(xxx) = - + 0.R (c cv) Lip((xx)
< | Lip(Axx) + || (c co)lliip(8xx)
< max() R|] |R ||c coll max(| || RI
Now can apply Theorem Note that
col) Txx) ||D2 1D2
(I lo + 1D2
max(LT- LTG (1+ La)'
22

---

50
|cc - col1
R|loo + 2(1 + [==]] || (
1108 10 x
1- max L Ln))
||e - coll
(1D2 +
R| - || (1D
1 max1 | oo Lx))
B.4 Stochastic CEBE approximation error
In this section we provide proof of Theorem 3
We start by introducing assumptions and notation.
The state space S is a bounded metric space. The action space A is metric space.
For metric space X. let Meas(X) denote the space of finite signed Borel measures equipped with the total variation norm. The transition map is a function T S x A x C MSa and
suppose it is twice continuously differentiable in C, and | |0 x., ||02 < cc. We also assume that T and ²²T are Lipschitz. In this subsection, view T as map into Mcas(S) (S) rather than into Wp, (S). So, for example, the Lipschitz constant LT refers to the Lipschitz constant of T as map S x A x C + Meas(S).
Recall that if X is a metric space and E Meas(X), then admits unique Hahn-Jordee decomposition ب = n² μ Let (X) denote the open subset of Meas( (X) consisting of the signed measures f with μ² (X) > O, and let ((X) denote the probability simplex of Meas X), consisting of the (unsigned) measures μ with ((X) 1. We define the map
P: u(X) ((X) by P(u): T
We suppose the reward function R: S x A + E is twice continuously differentiable with respect C, and || [0, [1] < cc. Moreover, we assume that R is Lipschitz.
A technical difficulty with transition maps in this setting is that the linearization of the parametrization map might produce measures that have negative components and therefore might no longer be probability measures. To fix this issue, we project the linearized transition maps onto the probability simplex (note this is not necessarily on a finite sample space). But first, we must check that it is possible to make this projection. In particular, for projection P of the form (S)) A(S) need to check that the lincarization T a) = T a) O.T (s, a) (c - co) lies in US)..
Lemma 11. Let X be a metric space, and let μ E ((X) be such that ((X) = Then
|| P(x) < 2 111 1|-
Proof. By the triangle inequality,
|| - P(u)ll S 1|
11
P())ll
1++
|| +
||
|| || + ||
||t
|| + 11
-
= (1 + || II) + ||
= 22|u=
In the second last line, used that μ(X)= 1, and therefore that ||
- 11
Lemma 12. For all E S, E A. and co E C, we have
(T, + 0 (c (cc-co)) (S) =
23

---

Proof. Let c > 0 be such that ball of radius el/c cull about is contained in C. Let us define 7: (:(=E,)) Meas(S) by
7t) T
Since T twice continuously differentiable in C, y is twice continuously differentiable. Since (t) E MSas(S) for all t, we have
1 7(1) =1
So for all t + O,
7(1) - -7(0) 0. t
For each N, let pn == 2(1/n)-2(0) 1/n
Since is differentiable, we have un, - (0)) in Meas(S).
So
1 1(0) = [ n n-000
lim pn
n- S Uh
lim
0.
Finally, by the chain rule, we have
(s (s, (c (S) (s.) 1.
The following lemma shows that the linearization 7 a) 0.7% (c- co) lies in (S)).
and therefore that the projection P(T (s, + (s, (c- (c-co) is well-defined.
Lemma 13. Let S. a E A, and c, co E c coll (2| - conditions hold:
1/22, then the following
(T + T (c (c) 2 1 - 1 cull?
(16a)
(s. (c- Eu(S),
(16b)
(P +0. (c- c))) < | /oolllllccc coll2.
(5,
(16c)
Proof. By Taylor's theorem,
(s) - (8, - 0. 0. (c co)l S 100 coll..
Then by the triangle inequality,
a) + (8. (c co))+(
(t))
|T
(c-
(T a)) 110.T ol 1 coll
1 =1 10
a)l=olc (t. coll.
Here we used that (s a) is probability measure on the second line. This establishes 16a In particular, the last line is greater than by the assumption on the norm of ||c coll. so
(T (s. (c - (c-co))+||>o,
0.7 ()
which establishes (6b
24

---

By Lemma 12 we have +.TT))))))00)))))) = 1. We can also bound the negative part of the linearized measure:
(s,a) (a. (c- +
[ - a. (c co)l
102 (| coll
Then by Lemma 11
(s, + 0 (c co) - (
0 (c- c)))
2 (T (s. + 8. (s. (c
cu)) (c-cu))
s | col2
To conclude, apply the triangle inequality:
[ - (e) - 0.7 (s,a) (c- co)l
+ 0,7 (c- co)il
$ (
(c co) P(T (8, + 0. 0,T% (c (c-ca))
ca))
< 10 0 coll
s 7
This establishes (6c)).
Next, show that the projection map sends nearby measures nearby probability measures. Lemma 14. Let X be metric space, and let μt, v E (X). Then
s l|le ml
1
Proof. By repeatedly applying the triangle inequality, get
xc+
/
[P(())
P(v)
|
|
1
|
1
|
|ut|
++
|r+
+
||
|
+
| ||
2 ||
2+
|| ||
||+|||
The following lemma relates the Wasserstein distance to the total variation distance, and is a special case of a result in the monograph of Villani et al. (2008) Theorem 6.12).
Lemma 15. Let X be bounded metric space, and let pr, P2 E Mcas(X). Then
Wi(Hi, (2) s [|m pell Meas(X).
25

---

To use Theorem I we need to show that the context-enhanced transition and reward functions are good estimates of their true values. By Lemma 13 for all (s, S A.
Proof of Theorem 3
(TT a) TTE(s, a)ll M)asss))))) S T -
By Lemma we have
(T (CR(s, a)) - (iaamSS)[7 - Ton(s, (a(ss,))lis a)]lMeas(S) (S) coll?
By Taylor's theorem, there exists d' E [c,coll such that
(s (s, (s. 0 a) co)l S 10 (s, a) a)lllc- coll
1 110 R col?
which rewrite as
[8 a) - C (s, a)l S ||0 |00||oo0100 - col 1²
|| |R - R - |102 R |oo ||c col ?
So the reward functions are close.
Next, we will show that the transition and reward functons are Lipschitz. By Lemma [5 for all (8, a), a') E 8 A, W (T. T' (iam()))[T'ss.)) - (s',',a')) M[M)ss(S)) diam(S)Lrd((sso) a), (s', a')),
so the Lipschitz constant of To: : S A + Wi(S) is at most LS.
Let us define : x - Meas(S) by
(s, = T (8,a) 0. (s, (c- co).
Then ThE PoT. By Lemma for all (s.a) E S A,
as
21 100 |||cc co
Mcas(S)
For all (8) and (s', in S x A. we have
(s, - (ss, a') (llMsss(S)
I T (s', Il + 110T -..T (dd coll
(s)d((s,a)) (8 ,a))
coldd (8),
Then by Lemma 14 for (s, (s', (s') A
2 (T( - -T(s', a')ll d()lMMcss(S)
(((((s)))-P(P)) ((f(s',
Mons(S)
||T a Meaa(S)
2 (LT + ||c - co La. )d( ((,(((ss))))) (s,
0 ~ |c cn 2
$ 4(L + ||c - La. )d( a), (s', a'))
Again by Lemma 15, we have
I (Tén(s,a), Tén(s'a')) ,d)) < diam( (E) TEn(s' (0'6'.'') ll))(MM)s((S) < 4 diam + ||c- coll (s
(), (s',
so the Lipschitz constant THE S x A + (i(S) is at most 4 diam (L + ||e
La.
By assumption, Ro is Lipschitz with constant at most Lx. The Lipschitz constant of RoE Ro0 R. (c-co) is most Lg La_Rllc- coll-
Finally, by Theorem I we have
lI
1%
S 1-7 (119 R|g coll2 By diam + Lz)) |0 coll²
- y max(1, diam(S) (S) + Lz))
26

---

B.5
CEBE policy is approximately optimal
Here we provide proof of Theorem
Proof. Recall that
20 Σ
[ =
7 Rg
So- (s,,
where Rg denotes the reward attained time from following the policy wce context c. Let Let denote the initial state distribution. Since QCe and QBE are close,
JCE c) = 1 / Cer a; Tan)drar(()dd(s)
/ OE (s (m)ddnn((dd((()))
1 (CEE(s, a - Qin a; TBE) (TBB(()dr(()
/A Qin (Bg))
JBE(TBE,C) c) 5.
Similarly,
JBB JEE(TCE, c) / / Qee(s, a; (cn)dren(()dd(ss)
2 s / Cee( (cn))drem(()dd(()
1 1CcE(1 ( Qee a (B)) dT
2 / (or(s, a: (Q)))))
-JCE(TCE, c) 8.
The above inequalities, combined with the (C, -ooptimailyyyy of TBE and सeE their respective MDPs, yield
JBE c) 2 JCE(TCE, c) -
JCn(mnE 6) 8
JBB (THE. c) - 25
sup JBE(p, c) 25 2c.
Hence, TCE is (C, 26 + 2op-mml on the Bellman equation.
B.6 Relation of CSE to regularization
In this section aim provide intuition about what functions learned when using CSE by developing a regularization perspective for the CSE loss. For simplicity we consider the supervised setting. In supervised learning, for an input data distribution D and true function F X x Y we want to learn, we may define loss function for a hypothesis. f X x C
L = (z= c) - F(x, c)
The corresponding CSE loss is then obtained by augmenting the data (by adding perturbed inputs and corresponding labels from the linearized true model) and is given by
L (f) EgEEoD (z, F(z, co) D.F(z,co)El))
27

---

where & is a random vector of our choice, which we will suppose is standard multivariate Gaussian. We then calculate
L (f)
F(x,c)) (z)
- -
+ +1(0ss - 0,F)512 + (5* - ( - E Bn (5 8E, F7 + E(O(1(
-BnDllS - F + Eg
EL.DDI/- F11 BeREyD [0, 0, F) Eg D _ 0%. -E, F)-
The first term of the above expression encourages f to be close to F. The second term is equal to
ExEx~De (0zff 0,F) - E .DD Tracc (0,f 0,F)* (0.f o,F)) D Trace((E(E " [00f 0,,F)"(0,.) 0₂F)) o Ex- Trace (0..f 0, (010)"(0000 (0.f O₂F))
o E 10,5 0 FIF
and therefore encourages the derivative agree with F. Under the assumption that the function is close approximation of F in the training context co. the loss is approximately
(f) R FEE=Dl10,f - 02 FI
Up higher-order terms in the perturbation size, data augmentation with CSE corresponds adding L regularization on the difference between the Jacobian of the trained model f and that of the true function F. other words, this encourages O&f O,F at co.
Additional Experiments
C.1 PendulumGoal
We provide additional plots for the PendulumGoal experiment. Figure we show the evaluation performance of the trained policies when sweeping over the context parameters. We additionally plot the mass and length context parameters. All methods perform similarly as we vary the length. When we vary the mass, the methods perform similarly until the mass is greater than about 1.6 at which point the baseline performs best. Note that CSE still outperforms LDR this case.
C.2 CartGoal
We present our results on the CartGoal environment in Figure 6 This enviromment uses discrete action spaces, so we train the functions and policies using DQN (Mnih al.) 2013). We can see in the figure that all methods perform similarly with respect to the cart mass and pole mass. CSE has slight edge on pole mass when the mass greater than 1.75. Interestingly, baseline performs better than CSE and LDR when changing the gravitational acceleration. We note that this is very out of distribution since the Ac used in training was only 0.1. On pole length, the policies trained with LDR performed best and CSE slightly outperforms baseline for most contexts. When varying the goal state, CSE performs best when the goal state is far from the origin.
28

---

SAC
200
SAC
200
SAC LDR
Mean Return
Mean Return
-600
800
-100
SAC
CSE
-100
Train
-1400
0.5
1.0
15
2.0
2.5
3.0
3.5
4.0
45
3.
0.6
0.8
1.0
1.2
1.4
1.6
1.8
2.0
Context
Context
Gravitational Acceleration
(b) Mass
SAC
-150
-200
-200
SAC
LDR
-250
Meein Return
4000
-300
-000
-150
-400
SAC CSE
-150
-300
Train
-100
0.5
0.5
1.5
1.2
1.4
1.6
1.8
2.0
-1.00
0..5 -..50
-025 0.00 0.25
0.50
0.75
1.00
Context
Context
(c) Length
(d) Goal Torque
Figure 5: Comparison of training methods on PendulumGool as we vary the context parameters.
500
Wean Rec.m
610
Heen
200
220
425
DON ESE
DON
420
DON CSE
DON
DON CSE
DON
100
DON LDR
250
DON
LDR
DON LDR
0
Trah
250
0.25
Trah
3.00
135
1.50 175
0.4
Taan
0.6
0.0
14
2.6
13
Context
Context
CoroeNT
(a) Gravitational Acceleration
(b) Pole Mass
(c) Cart Mass
503
DON DON
506
DON
Them
103
DON
nON CSF
-1022
RON IDR
Candeat
Contecct
(d) Pole Length
(c) Goal State
Figure Comparison of training methods on CartGoal as we vary the context parameters.
C.3 AntGoal
We present the results the AntGoal enviromment Figure 7 least well LDR across all contexts. Both CSE and LDR beat baseline performance across all contexts.
We can see that CSE performs at
29

---

-1000
SAC
SAC CSE
-1500
SAC LDR
Train
Mean Return
2200
-2500
-3000
-3500
4000
Context
Figure Comparison of training methods as we vary angle of the goal state in the AntGoal environment and keep fixed distance of 3. The training context is (300.0.0).
C.4 PendulumGool with Automatic Differentiation
In the other experiments, we use analytical gradients, but this may be limiting in some settings. In this section, we perform an experiment on the PendulumGoal environment and use automatic differ- entiation to obtain gradients. Figure [ shows the evaluation performance of the trained policies when sweeping over context parameters the PenGoal-AD environment, We can see that CSE performs similarly to LDR. In some cases, CSE outperforms LDR (e.g., when gravitational acceleration is at least.3. 3).
-000
200
SAC + CSE
SAC
Mean Return
Mean Return
Train
-600
800
SAC
SAC CSC
-1200
SAC
LDR
-800
1440
Train
15
3.
4.0
1.2
1.6
1.8
Context
Context
(a) Gravitational Acceleration
(b) Mass
SAC SAC + CSE SAC LDR
-200
Mean Return
Train
Moan Return
400-
-300
-350
-800
SAC CSE Train
400
1100
0.5
0.9
1.0
1.2
1.4
1.6
1.8
2.0
-1.00-033 -0.50 -0.25 0.00
0.25
0.50
0.75
1.00
Context
Context
(c) Length
(d) Goal Torque
Figure 8: Comparison of training methods on PendulumGoal-AD as we vary the context parameters.
C.5 Stochastic Transitions
We present the results for the SimpleDirection-Steed -ttochastitt environment in Figure This environment has the same transitions and rewards as in SimpleDirection, but with additive Gaussian noise
30

---

(0, 0.1) applied to both the transitions and rewards at each step. In this environment, the policy trained with CSE has similar performance as with LDR and outperforms the baseline. While this provides empirical evidence that CSE is robust to stochasticity in transitions and rewards, further work should establish theoretical guarantees for this scenario.
SAC
CSE
SAC
LDR
Train
Return
60
40
-1.00 -0.75 -0.50
-0.25
0.00
0.25
0.50
0.75
1.00
Context
Figure 9: Comparison of training methods vary the first context parameter in SimpleDirection-
Stochastic.
C.6 Aggregate Statistics and Context Dimension
In this section, we present a table average performance for each environment. We compute the mean performance from the previous context sweeps and show the results in Table The last column shows the normalized score
CSE Baseline
LDR Baseline
for each environment. Note that CSE is better than the baseline in all cases. CSE performs better than LDR in three of the six environments tested (CheetahVelocits SimpleDirection, and PendulumGoal). The data do not point relationship between the context dimension and the normalized score. This expected since the transitions and rewards vary significantly between environments. We expect that the context dimension have impact on the performance of LDR and CSE and further work should consider testing this when scaling to high dimensional contexts.
Environment
Dimension(C)
Baseline
LDR
CSE
Normalized Score
Cheetah Velocity
681.87
-588222
575.94
1.13
SimpleDirection
2
16.03
37.60
38.05
1.02
AntDirection
2
507.11
560.23
515.98
0.17
AntGoal
2
2790.55
2157.52
-2077.2 24
0.76
PendulumGoal
444.25500 .00
-42.558
394.19
8.83
CartGoal
5
284.65
325.49
315.74
0.76
Table 1: Mean Returns averaged across all context sweeps for each environment using analytic gradients. The last column shows the normalized score for CSE compared to LDR. The highest mean returns for each environment shown in bold.
D Experiment Details
In this section, we outline the details of our experimental setup. We include an overview of the SAC hyperparameters in Table replay buffer hyperparameters in Table For the SAC experiments, all neural networks use three fully connected layers of width 256 with ReLU activations. For the CartGoal experiment, we use DQN and the neural networks have a single fully connected layer of width 256 with Tanh activations. The neural network head is two layer fully connected neural network with width 256 and ReLU activations. We use the default RLLib implementations of SAC and default hyperparameters for each environment are based on the tuned examples from RLLib
31

---

Environment
Epochs
Actor LR
Critic LR
Entropy LR
Initial Entropy
SinpleDirection
500
1. 10-3
2.10-
4. 10
1.0
PendulumGoal
1000
2. 9.10-
8-10
9. 10
1.001
Cheetah
10000
2-10-1
8-10-4
9. 10-1
1.001
AntDirection
5000
3 3-10-
3-
1.10-1
1.001
AntGoal
10000
3 3-10-5
3-10-4
1.10-
0.01
Table 2: SAC Hyperparameters
Environment
Capacity
SimpleDirection
le6
0.6
0.4
PendulumGoal
1e5
1.0
0.0
CartGoal
5e4
0.6
0.4
Cheetah
1e5
0.6
0.4
AntDirection
1e6
0.6
0.4
AntGoal
le6
0.6
0.4
Table 3: Prioritized Episode Replay Buffer Hyperparameters
(iiang et al., 2018 2021). Note that Lee and Chung (2021) use smaller initial entropy penalty for AntGoal, so we use their choice of ox 0.02 as well. We increase the number of epochs in each environment to account for the increased complexity of training with contexts. When performing CSE and LDR, we generate perturbations Ac uniformly from the sphere of radius 0.1. To improve the sampling performance, we vectorize all environments to sample from 8 versions of the environment.
When training with SAC, have the following additional hyperparameters which keep constant across experiments. The target entropy is automatically tuned by RLLib. The policy polyak averaging coefficient is 10-3 We
training batch size of 256 and 4000 environment steps per training epoch. These are the default choices in RLLib, For CartGoal, we also have the following hyperpa- rameters with DQN. The learning rate is 5 10-4 and the train batch size is 32. We additionally enable double DQN and dueling within DQN to improve Q value estimation, implemented within Liang al. 22018 2021 We also use random exploration for 1 timesteps, followed by using e-greedy exploration with = 0.02. These are hyperparameters were chosen based on the default RLLib implementations DQN for the tuned CartPole example from RLLib (Liang et 2018 2021). Note that in both SAC and DQN, we use 1-step returns during training. While the experiments in this paper use SAC and DQN, we note that CSE derived from CEBE which should allow one to apply the method to any algorithm which leverages bootstrapping. Issues may arise when using n-step returns for estimating returns as repeated linear approximations will accumulate error for large n.
Hardware Experiments were run system with Intel(R) Xeon(R) Gold 6152 CPUs @ 2.10GHz and NVIDIA GeForee RTX 2080 Ti GPUs.
Python libraries Torch Paszke et al., 2019); Ray (Moritz et al. 2018); Ray Tune LLaaw et al.
2018); Ray RLlib (Liang et al.) 2018 2021): Seaborn (Waskom) 2021): Matplotlib (Hunter 2007 )
Pandas (pandas development team 2020); Numpy (Harrss et al. 2020); Scipy (Gommers et al. 2024
Mujoco (Todorov et al. 2012); Gymnasium (Towers et al. 2024).
E
Environments
We provide additional details about the environments in this section. A summary of the environment parameters included Table
32

---

Environment
Train Context
7
Horizon
SimpleDirection
(0.0,0.0)
0.9
10
PendulumGoal
1.,
0.99
200
CartGoal
0.99
500
Cheetah Velocih
(2.0)
0.99
1000
AntDirection
(1.0,0.0)
0.99
1000
AntGoal
(100,
0.99
1000
Table 4: Environment Parameters
E.1 Cliffwalking
We begin experiments by first solving the problem in tabular setting. This allows to exactly solve the Bellman equation and the CEBE to avoid noise due to sampling and training in deep RL. This allows us compute the approximation error exactly. We consider the tabular Cliffwalking environment shown in Figure [O] In this environment, the agent must navigate to goal state without first falling off cliff (depicted as black cells in the figure). If the agent walks near the cliff's edge can complete the task faster, but risks falling off the cliff. In this environment, the context is e [0, 1]. which denotes the probability that the agent slips. When the agent slips, the next state picked uniformly at random among the adjacent states. If the slipping probability is high, then the agent is incentivized to first navigate away from the edge of the cliff before navigating the goal state.
STAME
GROAL
Figure 10:
Visualization of the tabular Cliffwalker environment. The agent, represented by *, is taking the right action. The agent is most likely move right, but there is probability of the agent slipping. If the agent slips to moves to an adjacent cell with equal probability.
E.2 SimpleDirection
Recall that the equations for the SimpleDirection environment given by
( a) + + (s a, s') = s
We pick this environment because the transition linear in each variable and rewards are linear in c. The spaces
== R², A=
C=1-1,
and the initial state picked uniformly at random so gradients
Uniform (1-1, This gives the resulting
8T =1 dc
OR
0R 0s
c.
de
E.3 PendulumGoal
The transition function is the same as in the pendulum environment from Towers et al. (2024) and are determined by the solution to the ODE
Ø 3g
sin 0 +
2/
ml?
over time interval dt = 0.02. We modify the reward so that the desired position of the pendulum is not vertical. In particular, we specify be the desired torque applied at the goal state. We choose to specify goal torque instead of goal state since the action space is bounded and an arbitrary goal state may not be stable for any policy under different contexts (e.g., when max lu| < mgl sin(O (oag)))))) Solving for the goal state at the desired torque т, we get
2T
Ogoal arcsin
mgl
33

---

The reward function is then
R n² sin (Q )
0.1 0 0.001 u²
2
where we have also used sin instead of the absolute value function for state error due to periodicity in the problem.
E.4 CartGoal
The transition function is the same as in the cartpole environment from Towers et al. (2024) and are determined by the solution to the system of ODEs
gii 0 + cos 0 -/-mpp
i
Map Mi
my cos my+mm
mpl (02 sin 0 @ cos
mp + mc
over time interval dt = 0.02. The reward function is modified to include a goal state for the cart position
= / 1+ (2-goa))
Note that the environment uses the same termination conditions as CartPole and can terminate early if x &f 1-24,2.4 or 0 of [-0.21, 0.21].
F,5
ODE Environment Simulation
The continuous control environments specified by differential equations and reward functions. Derivatives of T and R are computed symbolically. The differential equations are then converted into first-order systems and compiled into Python functions before solving numerically with Euler's method.
Example with a simple equation equation
Consider the transition function specified by the differential
bF
and reward function
R(a,b.o) (c xtt + At))?
where a, e R are context variables and F is the control. Then the derivatives of the transition function with respect to the context variables are solutions of the following differential equations:
dă
dx da
dă
dx
di dr
+
F,,
0.
du
db
db
de
dc
Similarly, the reward function has derivatives d.(.ab.)) d.((ab.c))
da O, db 0, dR(a.b.))
dR(a.b,)
de -2(c xx),
2(c x).
dx
When solving for the gradients, we use initial conditions
dx (0)=
dr (0) 0,
dr
(0)
da
Z
dc
This is because we care about how the next state would change we change the context. So, the change the next state due the change in the context would be zero we did not move forward in time at all. We include a full derivation for the PendulumGoal system in Appendix F2
34

---

E.6 Goal-based Mujoco Environments
Lee and Chung (2021) introduced goal-based variants the HalfCheetah and Ant from Mujoco
(Todorov et al. 2012). We denote these environment by Cheetah Velciay, AntDirection, and AntGoal. They each follow the same dynamics equations specified by Todorov al. (2012). However, their rewards are modified introduce context dependent tasks.
The original Cheetah VelccthVelocity implementation modifies the forward reward to penalty on the velocity -|u Ugoal
We use smoothed version of this function,
Rellocity
1 (u Ugoul )
We follow the original reward function for the AntDir environment. The velocity component of the Ant reward function is modified to
Direction = C U.
The original AntGoal implementation modifies the forward reward to a goal based penalty -1(x, y) (Fgoul; Bggoaa) )- We smoothed version of this function,
V 1 + ((x, y) (Igoal Vggua)
Roal
Worked Examples
In this section, we provide fully worked out calculations for some of the environments in the paper with the aim of providing more clarity.
F.1 SimpleDirection expected optimal returns
The kth state is given by
8 80
a + Kc.
1=0
Using this, we get the rewards are
c. 8k+1 c. 80 + c.
a: + (k + 1)]|c
0m0
By symmetry on so initial distribution, we have
E(rx) Σ 04 (k: 1)]/cl|.
Due to the constraints on a,, the maximum is obtained when a; = sign(e) (computed element-wise). So, get
B(rx) = (k + 1) (I|ell + [|<ll2)
The expected non-discounted returns thus
#-1
Σ
#-1
Σ (k+ 1) (Iee)) + [|ell2)
H + 1) !) (Iee)) + [|c|(2) 2
Th
0mm
Amma
where H is the horizon.
F.2
PendulumGoal Equations
In this section, we present the full set equations for the PendulumGoal Environment, This includes the transition and reward gradients. Recall that the context parameters are (g. m, l, r), where g is the
35

---

gravitational constant, m is mass, / is the pendulum length, and т is the goal torque. The dynamics are governed by the equation
3g
3
(0
21
ml2
where f is the pendulum angle and u is the control torque. The reward function is
ggual
0 ))
R n² sin
+ + 0.001 u²,
2
where &guaa sin
-27 mgl
Taking gradients of the transition dynamics in each context parameter,
we get
8 2/
(0) + goos(0)gg)
o. 3g cos
21
m2?
3g sin(0)
3g
6
0 2/2
+ cos
21
ml3
3g
0 21
cos(0)
The full transition dynamics with the gradient equations are then converted into the following first order system of differential equations.
à 3g
3
sin
21
ml?
0 Ø
3
6 21
(sin(0) g (0)
0, = Φ9
3g cos (0)(0)
pm 21
22
8m = pm
3g sin(O)
3g
Φi 212
+
cos(0)
21
m/3"
0: Φt
3g
3 = cos(9)),
2]
0 Φr:
where á
d9 2
The reward function and its gradients
8 8goal
R = n² sin
+ 0.1 0² + 0.001 2²
2
0 gnal
0
0 Ogoal
R = x² cos
sin
(0g0 goal +g)++
+ 0.2 12
2
2
Ogual
0
6
ggcal
R = n² cos
sin
(0m O goul
-0y) -0y) 0.2 0.2
2
2
0 gnal
0
0
0goal
Ri 2² cos
sin
(0,0 goal
01) + (
2
2
0gnal
0
0
0goal
RC = 2² cos
sin
(0-0 goal - 02) + 0.2 6-
2
2
36

---

where
2t
Ogggnal
mg21
472
g'=""
goal
m² m²gl 1
4-²
2T
0.0 goal
mgl2 /1
4r2
g-1
-2
a 0 goal
mgl /
172
in g2
G
Broader Impacts
This paper develops foundations for generalization in RL with sufficiently regular CMDPs. This may be applied improve performance out-of-distribution contexts. Improvements in this direction aid in increasing our understanding of CMDPs and can lead to the design of policies which are more robust to perturbations of the context. This work focuses on sufficiently smooth CMDPs, which appear, for instance, in physics problems. The authors do not foresee any direct negative societal impacts this work, but do not have control over how these methods might applied in practice.
37

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Each aspect described in the abstract is represented by section, theoretical, result or experiment in the article.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The paper discusses limitations in the conclusion.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
38

---

Theorem statements and assumptions are included in Section 3 Proofs of the theorems are included in Appendix B We provide proof sketch for our main theorem, Theorem iin Section 3
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We provide all code to reproduce the experiments in GitHub repository linked from the title page of the article.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing verifying the results.
5. Open access to data and code
39

---

Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: We provide all code to reproduce the experiments in GitHub repository linked from the title page of the article.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips... cc/
for more details.
While encourage the release of code and data, we understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: We provide all code to reproduce the experiments GitHub repository linked from the title page of the article (anonymous for double-blind review). We also outline hyperparameters Appendix D
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We report 95% confidence intervals in all experiments involving statistical significance.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
40

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπτог rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
In Appendix D we include include all details of the compute resources.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics httpp::/neuurpp: cc/ppullcc/thiccciiiieee
Answer: [Yes]
The research complies with the NeurIPS Code of Ethics. The research process did not cause harms do the outcomes of the research represent immediate risks create harmful consequences.
Justification:
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g., s if there is special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
We provide a discussion of the broader impacts of this work in the appendix.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
41

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The paper does not release data or models. The code includes some new environments, but these present no risk for misuse.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We properly cite all environments and algorithms used in the paper. We also include citations for python packages in the appendix.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
42

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperssi thood cww datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is not available online, the authors encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
The code base is well documented and includes clear instructions on how to run all experiments done in the paper.
Justification:
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
The paper does not involve human subjects or require IRB approval.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
43

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
The core method development does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy for what should or should not be described.
44


Reviews:
[{"review_id": "UZbaIUZkLq", "reviewer": "Reviewer_xxBc", "summary": "The 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.", "questions": "- 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.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "- 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.", "was_revised": true, "final_justification": "The authors responded to my comments and questions well, along with conducting additional experiments. I will keep my score, as I feel that the weakness still remains as a limitation of this work. However, I'd be happy for this paper to be accepted!", "rebuttal": "We thank the reviewer for their valuable comments, questions and suggestions that will improve the final manuscript.\n\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\nIndeed, other gradient estimation methods can improve the applicability of the method. The theory developed for CSE assumes exact gradient information. However, sufficiently accurate estimates for the gradients should suffice for CSE. Differentiable physics simulators like Brax enable automatic differentiation through the physics engine. Finite difference estimates can be used to obtain gradient estimates as well, but require that one simulates $|\\mathcal{C}| + 1$ contexts. The gradients supplied by these methods can be used with CSE to improve generalization. We ran an additional experiment using automatic differentiation with the PendulumGoal environment and present the mean performance across each context sweep here:\n\n|          | Gravity | Mass    | Length  | Goal Torque |\n|----------|---------|---------|---------|-------------|\n| Baseline | -541.48 | -423.61 | -502.22 | -273.15     |\n| LDR      | -496.67 | -440.14 | -503.97 | -252.42     |\n| CSE      | -388.47 | -398.71 | -477.59 | -245.40     |\n\nIn each sweep across a context variable, CSE performs better than the baseline and LDR. We will add a plot corresponding to this experiment accompanied by the above discussion in the final manuscript.\n\n> What is “Sample-enhanced transition and reward” in line 142? same as Context-enhanced transition and reward?\n\nThank you for pointing this out. This is a typo and we will update it to “context-enhanced transition and reward” in the final manuscript.\n\n> I understand you assume some deterministic transitions for the theory. Can you validate CSE experimentally with the case with stochastic transitions?\n\nThank you for suggesting this. We ran an additional experiment for SimpleDirection with additive Gaussian noise in the transitions and rewards. The results are similar to the SimpleDirection results with CSE performance matching LDR, both of which outperform Baseline. We will include these results in the final manuscript.\n\nIn the theory for CEBE, we do not assume that the transitions are deterministic (Theorems 1 and 3 work both for deterministic and stochastic transitions). For CSE, the deterministic setting confers multiple advantages including an efficient implementation and unique method for augmenting samples. In the stochastic setting, one must directly address how the distribution $\\mathcal{T}^c(s, a)$ changes as the context changes in order to develop a stochastic CSE directly from theory. Due to time and space constraints, we have not included a detailed analysis and experimental validation of CSE in the stochastic setting, but acknowledge its importance and intend to pursue it in future work.  \n\n> Also, Lipschitz assumption may not hold in real-world CMDPs. Do you have any examples with that?\n\nWe assume that CMDPs are Lipschitz to streamline our presentation and to give clearer proofs. This assumption is analogous to that of the Picard-Lindelöf theorem from ODEs. We expect similar results to hold under local Lipschitz assumptions, since we only need to control the Lipschitz constant to the extent that the state space gets traversed. Observe that some amount of smoothness is necessary for context generalization to be possible, since otherwise transitions and rewards could be entirely  different in nearby contexts.\n\nWe have examples with transitions and rewards that are not globally Lipschitz. Consider the PendulumGoal environment, whose gradient equations are shown in Appendix F.2. In the first-order system, we can see that $\\dot \\phi_g$ depends on $\\frac{3}{2l}$, which is not Lipschitz for all $g \\in (0, \\infty)$. From a physical perspective, if the length of the pendulum is low then the pendulum velocity $\\phi = \\frac{d \\theta}{d t}$ is sensitive to the gravitational acceleration. We observe a similar behavior with the reward function for this environment, which has unbounded gradients in the context parameters. Since these gradients are unbounded in the full context space, the function is not globally Lipschitz but is locally Lipschitz. Even so, the experiments demonstrate that the policy can still generalize well in a neighborhood of the training context.\n\n> Your theory gives an $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^*|$ versus context distance or derive an explicit radius $r_{\\text{max}}$ beyond which linearization degrades?\n\nOur theory allows one to provide quantitative bounds. Consider the Cliffwalker Environment with the reward equation specified by Eq 12 right. Note that the transition dynamics for this equation depend linearly on $T$, so the second term in the bound from Theorem 3 navishes. We find that\n$$\\partial_c^2 R^c_{\\text{cliff}} = -20 (1+c)^{-3} \\quad \\partial_c^2 R^c_{\\text{goal}} = 3.75 (1+c)^{-3.5}$$\nThis implies that\n$$\\| \\partial_c^2 R^c \\| \\leq \\max\\left( 20 (1+c)^{-3}, 3.75 (1+c)^{-3.5} \\right) \\leq 20$$\nWith the choice of $\\gamma = 0.1$ used in the experiment, we obtain\n$$\\| Q_{CE}^c - Q_{BE}^c \\|_\\infty \\leq \\| c - c_0 \\|^2 200 / 9 \\approx 22.2 \\| c - c_0 \\|^2$$\n\nThe radius for which the linearization is a good approximation is dependent on the structure of the reward and transition functions. As an illustrative case, consider a CMDP where the transition function is trivial: the CMDP never transitions from one state to another state. Then $DT = 0$ and $D^2T = 0$, so Theorem 2 says that $\\|Q^c_{\\text{CE}} - Q^c_{\\text{BE}}\\| <= \\|c - c_0\\|^2\\|D^2 R\\|_{\\infty}$ which is the second-order Taylor expansion of the reward function. In general, we expect the radius of the approximation around a point (s, a, c) to scale directly with $\\|D^2 R(s, a, c)\\|^{-1}$ and $\\|D^2 T(s, a, c)\\|^{-1}$.\n\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\nWe also perform experiments with DQN in the CartGoal experiment in Appendix C.2. This environment has a discrete action space and continuous state space. We note that CSE requires a continuous state space so that $\\partial_c T$ is well-defined.\n\nWhile we present CSE mainly in the context of SAC and do an experiment with DQN in the appendix, it should work more generally with other algorithms. Importantly, CSE is derived from the CEBE. The CEBE estimates the value of the next state in the perturbed context by bootstrapping on the Q function. Algorithms which leverage bootstrapping should work similarly. Issues may arise when using $n$-step returns for estimating returns as repeated linear approximations will accumulate error for large $n$.\n\nWe will include this discussion about the algorithms in the appendix of the final manuscript."}, {"review_id": "aXBopLRrAE", "reviewer": "Reviewer_dZY9", "summary": "The 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.", "questions": "covered above.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "# 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.", "was_revised": true, "final_justification": "My original review serves as the justification for my score. The authors provided several more results, but I do not see them as substantial enough to raise my score from 4 to 5. I'm still happy for the paper to be accepted.", "rebuttal": "We thank the reviewer for their valuable comments and suggestions that will improve the final manuscript.\n\n> The 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\nIn this work, we conduct experiments across several regimes to demonstrate the benefits and generality of our method. Specifically, we validate Theorem 3 using the Cliffwalker environment, providing a clear link between theory and practice. To assess CSE in discrete control settings, we evaluate it on the CartGoal environment using the DQN algorithm. For continuous control, we test CSE on a diverse set of benchmark environments, including SimpleDirection, PendulumGoal, CheetahVelocity, AntVelocity, and AntDirection. These environments were chosen to cover a range of dynamics, context structures, and action spaces. We ran additional experiments using automatic differentiation and stochastic transitions and rewards, which we will add to the final manuscript.\n\nWe compare against domain randomization, as it is a widely adopted and well-established baseline for out-of-distribution generalization in RL. Notably, domain randomization improves generalization by leveraging exact data from nearby contexts, whereas CSE approximates transitions and rewards in these regions. In this sense, domain randomization can be viewed as an exact counterpart to CSE. Both approaches emphasize the influence of data, rather than other factors such as model architecture. This enables a fair comparison across a diverse set of environments while keeping implementation details consistent between methods.\n\n> A 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\nThe theory developed for CSE assumes exact gradients to simplify the presentation. However, sufficiently accurate estimates for the gradients should suffice for CSE. One may use differentiable simulators like Brax to take gradients through the simulation using automatic differentiation. One can also use finite difference approximations, but this requires that one simulates $|\\mathcal{C}| + 1$ contexts. We ran an additional experiment using automatic differentiation with the PendulumGoal environment and present the mean performance across each context sweep here:\n\n|          | Gravity | Mass    | Length  | Goal Torque |\n|----------|---------|---------|---------|-------------|\n| Baseline | -541.48 | -423.61 | -502.22 | -273.15     |\n| LDR      | -496.67 | -440.14 | -503.97 | -252.42     |\n| CSE      | -388.47 | -398.71 | -477.59 | -245.40     |\n\nIn each sweep across a context variable, CSE performs better than the baseline and LDR. We will add a plot corresponding to this experiment accompanied by the above discussion in the final manuscript."}, {"review_id": "wUQcHyjZ0r", "reviewer": "Reviewer_W69A", "summary": "CLAIMS: \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", "questions": "Zhao 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.", "limitations": "The 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.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "The assumptions of explicit and differentiable context limit the applicability of the method.", "was_revised": true, "final_justification": "The authors corrected my misunderstanding about the dimensionality of the context variables, which addressed my concern regarding the work's scalability measurably.", "rebuttal": "We thank the reviewer for their valuable comments, questions and suggestions that will improve the final manuscript.\n\n> The 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\nThe experiments use higher-dimensional context spaces during training. However, plotting the performance of this against all the dimensions is not possible for a paper. Training is performed with the full dimensionality of the context space, but the plots are generated by performing a sweep over a single dimension. This allows us to view slices of the evaluation performance as we vary the context. In Table 3 from Appendix E, we list details about the training context. We have included the context space dimensions for each environment here for convenience.\n\n| Environment     | Context Space Dimension |\n|-----------------|------------------------:|\n| Cliffwalker     |                       1 |\n| SimpleDirection |                       2 |\n| PendulumGoal    |                       4 |\n| CheetahVelocity |                       1 |\n| AntDirection    |                       2 |\n| CartGoal        |                       5 |\n| AntGoal         |                       2 |\n\n> CSE seems to work well with even stochastic rewards, despite not mathematically proven\n\nOur work focuses primarily on deterministic rewards, but the stochastic reward setting is an important case to consider for future works. We note that the methods used in this paper to address stochastic transitions should naturally extend to stochastic rewards. We ran an additional experiment for SimpleDirection with additive Gaussian noise in the transitions and rewards. The results are similar to the SimpleDirection results with CSE performance matching LDR, both of which outperform Baseline. We will include these results in the final manuscript.\n\nSimilar to the transition function, the structure of the reward function and choice of approximation can influence the overall results. Future works should leverage the structure of the underlying CMDP, domain constraints, and available reward gradient information to motivate these decisions. For example, in simulated environments, one can use the reparameterization trick in combination with automatic differentiation so that the simulator can provide reward gradients. We agree that extending the theory in this direction is a valuable avenue for future work. While we expect the generalization to carry over naturally, we focused on the deterministic case to maintain a clear presentation of the core ideas.\n\n> Zhao 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\nThank you for pointing this out. The concept of domain randomization has existed since the 1990s and can be attributed to multiple sources. We will add the following sources to the final manuscript:\n\n- Nick Jakobi, Phil Husbands, and Inman Harvey. “Noise and the reality gap: The use of simulation in evolutionary robotics”. In: European Conference on Artificial Life. Springer. 1995, pp. 704–720.\n- Ofir Levy and Lior Wolf. “Live repetition counting”. In: Proceedings of the IEEE International Conference on Computer Vision. 2015, pp. 3020–3028.\n- Fereshteh Sadeghi and Sergey Levine. “(CAD)ˆ2 RL: Real Single-Image Flight without a Single Real Image”. In: arXiv preprint arXiv:1611.04201 (2016).\n- Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/RSJ international conference on intelligent robots and systems (IROS), pp. 23–30. IEEE, 2017.\n\nIf there are any other sources that you would recommend, please let us know.\n\n> I think figure 1 can be better designed / scaled. Currently, it takes too much space in the main manuscript and it shows quite little.\n\nThank you for this comment. We added this figure to establish a direct connection between the theory and experiments. We will make this figure smaller in the final manuscript.\n\n> The 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.\n\nThank you for suggesting this source. The work on options and hierarchical planning is very interesting and important, especially in the context of generalization over causal reasoning tasks. This work uses hierarchical planning to generalize task information across variations in gridworld environments. For example, an agent may need to first find a key to unlock a door before navigating to a goal state. The agent must learn spatial abstraction to handle changing key locations, as well as temporal abstraction to learn the causal relationship between the subtasks in the problem. The authors design a hierarchical policy which achieves better generalization on out-of-distribution tasks. It is important to note that this depends on the number of contexts (seeds) used during training. In figure 4 of their paper, we see that out-of-distribution generalization increases significantly when at least 25 training contexts are used.\n\nOur work leverages approximations of the CMDP to improve zero-shot, out-of-distribution generalization. By using gradient information, we can improve generalization with only one training context using data augmentation. The focus of their work is in designing a hierarchical policy while our work is independent of the choice of architecture. Importantly, the flexibility of our method allows one to incorporate CEBE and/or CSE into other frameworks to make improvements. The authors of this work wrote their code only for gridworld environments (discrete state and action spaces) and noted challenges of extending this to continuous problems (e.g., goal identification). As you noted, the context spaces they use are discrete whereas ours are continuous. These factors make it challenging to experimentally compare with our work. An interesting future direction may combine our work with theirs to account for combinations of discrete and continuous context variables.\n\nWe will discuss this in the final version of the manuscript."}]

Author rebuttal:

Meta-review: The reviewers do point out a number of point related to the paper.

1. 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.
2. However the reviewers agree that the paper has some interesting components. The rebuttal was helpful to increase confidence in the reviewers understanding of the paper.
3. Future work should consider more complex environments that will require larger context. And real world environments where gradients and rewards are not easy to access.
Decision: Accept (poster)