Uncertainty-Driven Exploration for Generalization in Reinforcement Learning
Authors: Yiding Jiang
J Zico Kolter
Roberta Raileanu
Venue: Submitted to ICLR 2023 (2023)
Keywords: Deep reinforcement learning
exploration
generalization
procgen
crafter

Abstract: Value-based methods tend to outperform policy optimization methods when trained and tested in single environments; however, they significantly underperform when trained on multiple environments with similar characteristics and tested on new ones from the same distribution. We investigate the potential reasons behind the poor generalization performance of value-based methods and discover that exploration plays a crucial role in these settings. Exploration is helpful not only for finding optimal solutions to the training environments, but also for acquiring knowledge that helps generalization to other unseen environments. We show how to make value-based methods competitive with policy optimization methods in these settings by using uncertainty-driven exploration and distribtutional RL. Our algorithm is the first value-based method to achieve state-of-the-art on both Procgen and Crafter, two challenging benchmarks for generalization in RL. 

Under review as a conference paper at ICLR 2023
UNCERTAINTY-DRIVER EXPLORATION FOR GENERALIZATION IN REINFORCEMENT LEARNING
Anonymous authors
Paper under double-blind review
ABSTRACT
Value-based methods competitive when trained and tested in single environ- ments. However, they fall short when trained on multiple environments with sim- ilar characteristies and tested on new ones from the same family, We investigate the potential reasons behind the poor generalization performance of value-based methods and discover that exploration plays crucial role in these settings. Ex- ploration is helpful not only for finding optimal solutions to the training environ- ments but also for acquiring knowledge that helps generalization to other unseen environments. We show how to make value-based methods competitive in these settings by using uncertainty-drive exploration and distributional RI., Our algo- rithm is the first value-based approach to achieve state-of-the-art on both Procgen and Crafter, two challenging benchmarks for generalization in RL.
INTRODUCTION
Value-based methods (Watkins & Dayan, 1992) (which directly derive a policy from the value func- tions) tend to be competitive on singleton Markov decision processes (MDPs) where agents are trained and tested on the same environment (Mnih et al., 2013; Hessel al., 2018; Badia et al., 2020). However, they fall short in contextual MDPs (CMDPs) (Wang et al., 2020; Ehrenberg et al., 2022), where agents are trained on a number of different environments that share a common struc- ture and tested on unseen environments from the same family (Cobbe et al., 2019; Wang et al., 2020; Mohanty et al., 2021; Ehrenberg et al., 2022). In this work, we aim to understand potential reasons for why value-based approaches work well in singleton MDPs but not in contextual MDPs and how we can make them competitive in CMDPs.
Most of the existing approaches for improving generalization in CMDPs have treated the problem as a pure representation learning problem, applying regularization techniques which are commonly used in supervised deep learning (Farebrother et al., 2018; Cobbe et al., 2018; Igl et al., 2019; Lee et al., 2020; Ye et al., 2020; Laskin et al., 2020; Raileanu et al., 2020). However, these methods neglect the unique structure of reinforcement learning (RL), namely that agents collect their own data by exploring their environments. This suggests that there may be other avenues for improving generalization in RL beyond representation learning.
Here, we identify the agent's exploration strategy as a key factor influencing generalization in contextual MDPs. First, exploration can accelerate training RL, and since neural networks may naturally generalize, better exploration can result in better training performance and consequently better generalization performance. Moreover, in singleton MDPs, exploration can only benefit deci- sions in that environment, while in CMDPs exploration in one environment can also help decisions in other, potentially unseen, environments. This is because learning about other parts of the environ- ment can be useful in other MDPs even if it is not useful for the current MDP. As shown in Figure 1, trajectories that are suboptimal in certain MDPs may turn out to be optimal in other MDPs from the same family, this knowledge can help find the optimal policy more quickly in new MDP encountered during training, and better generalize to new MDPs without additional training.
One goal of exploration is to learn new things about the (knowable parts of the) environment as to asymptotically reduce epistemic uncertainty. To model epistemic uncertainty (which reducible by acquiring more data), need to disentangle it from aleatoric uncertainty (which is irreducible and stems from the inherent stochasticity of the environment). As first observed by Raileanu & Fergus (2021), in CMDPs the same state can have different values depending on the environment,

---

Under review as a conference paper at ICLR 2023
Train
Test
Train
Test
START
START
START
START
Standard
Explorative
Figure I:
Exploration can help agents generalize to new situations at test time, even is not needed for finding the optimal policy on the training environments.
but the agent does not know which environment in so it cannot perfectly predict the value of such states. This is a type of aleatoric uncertainty which can be modeled by learning a distribution over all possible values rather than a single point estimate (Bellemare et al., 2017). Based on these observations, we propose Exploration via Distributional Ensemble (EDE), method that uses an ensemble of Q-value distributions to encourage exploring states with large epistemic uncertainty. We evaluate approach both Procgen (Cobbe al., 2019) and Crafter (Hafner, 2022), two procedurally generated CMDP benchmarks for generalization in deep RL, demonstrating a signif- icant improvement
more naïve exploration strategies. This is the first model-free value-based method to achieve state-of-the-arr performance on these benchmarks, in terms of both sample effi ciency and generalization, surpassing strong policy-optimization baselines (i.e., methods that learn a parameterized policy in addition to a value function) and even a model-based one.
To summarize, this work: (i) identifies exploration as key factor for generalization in CMDPs, (ii) supports this hypothesis using a didactic example in tabular CMDP, (iii) proposes an exploration method based on minimizing the agent's epistemic uncertainty, and (iv) achieves state-of-the-ar performance on two generalization benchmarks for deep RI., Procgen and Crafter.
BACKGROUND
Episodic Reinforcement Learning. A Markov decision process (MDP) is defined by the tuple μ (S, A, RG P. Po). where S is the state space, A the action space, R : S x A - [Rmins Rmax] is the reward function, P: factor, and pn S R20 is the initial state distribution. We further denote the trajectory of an episode to be the sequence T = (8o, ag ro..... 8T, aT, TT, 8T+1) where re = R(sr. a:) and T is the length of the trajectory which can be infinite. If a trajectory generated by probabilistic policy random variable that describes the discounted return the policy achieves. The objective is to find a π- that maximizes the expected discounted return, rr = arg max, Ex-p=(.) [ where p² (r) = (o(s)) IT P(s 8t, (t) a.). For simplicity, we will use E, instead of (Ty)pT()) to denote the expectation over trajectories sampled from the policy æ. With a slight abuse of notation, we use 2 (s, a) to denote the conditional discounted return when starting at s and taking action a (i.e., = s and au = a). Finally, without loss of generality, we assume measures are discrete and their values lie within [0, 1].
S + R2 the transition distribution, y E (0, 1] is the discount
S A + R20. Z = -ET 2rr,
Value-based methods (Watkins & Dayan, 1992) rely on fundamental quantity in RI, the state- action value function, also referred to as the Q-function, QT (s, a) TE, [Z* | 80 = 8, ag = a). The Q-function of policy can be found at the fixed point of the Bellman operator, T (Bell- man, 1957), T=Q(s, a) = E.--P()l,,)) a'>>(-1=') [R(s, a) + QQ('s, a')] Bellemare et al. (2017)
extends the procedure to the distribution of discounted returns, T 7 a) d a(s,a) +
(Z(s', a'), s' ~ P(- 8, a) and a' .(- 8'), where - denotes that two random variables have the same distributions. This extension is referred to as distributional RL (we provide more detailed description of QR-DQN, the distributional RL algorithm we use, in Appendix B. For value-based methods, the policy is directly derived from the Q-function T(a s) 1 arrg max, ((a).s)(a).
Policy optimization approaches (Williams, 1992; Sutton et al., 1999) on the other hand seeks to directly learn the optimal policy via Monte Carlo estimation of the return. Specifically, given
2

---

Under review as a conference paper at ICLR 2023
Reward
Q-learing 4-greedy
Q-leaming UCB
Policy Gradient
-0.04
00.04
-0.30 26
40.04
35
35
00.04
10.04
0.00
subaptimality
25-
suboptinally
255
-0.04
0.00
00.00
0.00
2.0
20
20
-0.04
40.04
40.06
40.00
-0.04
10-
10-
.0.0
40.00
40.00
-0.04
-0.5
250
400 e00
203
-05
2:0
450
800
500
-0.
250
400 602
300
opisode
opisodo
ppioodo
(a) Tabular CMDP
(b) E-greddy
(c) UCB
(d) Policy Gradient
Figure 2: (a) A tabular CMDP that illustrates the importance of exploration for generalization in RL. During training. the agent starts the blue square, while at test time it starts in the orange square. In both cases, the goal is to get to the green square. The other plots show the mean and standard deviation of the train and test suboptimality (difference between optimal return and achieved return) over 100 runs for (b) Q-lcarning with e-greedy exploration, (c) Q-learning with UCB exploration, and (d) policy gradient with entropy bonus.
policy To parameterized by 9, we can compute the gradient of the expected return with respect to & through the policy gradient theorem, a Br, [ = RT, E ((ss, a V log (a st) where (s, a) is the advantage function (Schulman et al., 2015). 1 is updated with crudient-based optimization method such as stochastic gradient descent. The final policy involves directly sampling from Te (at 8t). In general, the ad advantage function T/ (s, a) can also involve learned value func- tions. However, the crucial feature of policy optimization methods is that they have a separate policy function, whereas value- based methods do not.
Generalization in Contextual MDPs. A contextual Markov decision process (CMDP) (Hallak et al., 2015) is special class of partially observable Markov decision procexs (POMDP) consisting of different MDPs, that share state and action spaces but have different R, P, and pe. In addition to the standard assumptions of a CMDP, the MDPs. During training, we are given a (finite or infinite) number of training MDPs, Muain = {uI, /2,,..., fnn ) drawn from qu (Ghosh et al. (2021) refers to this setting as epistemic POMDP). We use (T) to denote the trajectory distribution of rolling out x in μ. The objective is to find a single policy æ that maximizes the expected discounted return over the entire distribution of MDPs,
assume the existence of structured distribution qu (u) over
E 7're Furthermore, we assume the existence of n² (potentially more
than one) that is optimal for all μ E supp(gu), since otherwise zero-shot generalization can be intractable (Malik et al., 2021). If the number of training environments is infinite, the challenge is learning good policies for all of them in a sample-efficiens manner, i.e., optimization; if it is finite, the challenge is also generalization to unseen environments.
3 A DIDACTIC EXAMPLE: GENERALIZATION IN A TABULAR CMDP
Much of the literature treats generalization in deep RL as a pure representation learning prob- lem (Song al., 2020; Zhang et al., 2020a;b; Agarwal et al., 2021a) and aims to improve it by using regularization (Farebrother et al., 2018; Zhang et al., 20;8a; Cobbe et al., 2018; Igl et al., 2019) or data augmentation (Cobbe et al., 2018: Lee al., 2020: Ye et al., 2020; Kostrikov et al., 2020; Laskin et al., 2020; Raileanu et al., 2020; Wang et al., 2020). In this section, we will show that the problem of generalization in RL extends beyond representation learning by considering a tabular CMDP which does not require representation learning. The goal is to provide intuition on the role of exploration for generalization RL using a toy example. More specifically, we show that exploring the training environments can be helpful not only for finding rewards in those environments but also for making good decisions in new environments encountered at test time.
Concretely, we consider generic 5 x 5 grid environment (Figure 2a). During training, the agent always starts at a fixed initial state, (x = 0, y 0), and can move in cardinal directions (i.e., up, down, left, right). The transition function is deterministic and if the agent moves against the boundary, it ends up at the same location. If the agent reaches the terminal state, (x = 4, y = 0), it receives large positive reward, negative reward, r = -0.04. At test time, the agent starts at a different location, (x = 0, y = 4). In other words, the train and test MDPs only differ by their initial state distribution. In addition, each
2 and the episode ends. Otherwise, the agent receives small
3

---

Under review as a conference paper at ICLR 2023
episode is terminated at 250 steps (10 times the size of the state space) to speed up the simulation, but most episodes reach the terminal state before forced termination.
We study 3 classes of algorithms with different exploration strategies: (1) Q-learning with e-greedy (Greedy, Watkins & Dayan (1992)), (2) Q-learning with UCB (JCB, Auer et al. (2002)), and (3) policy gradient with entropy bonus (PG, Williams & Peng (1991)). To avoid any confounding effects of function approximation, we use tabular policy parameterization for Q-values and unnormalized log action probabilities. Both Greedy and UCB use the same base Q-laarning algorithm (Watkins & Dayan, 1992). Greedy explores with s-greedy strategy which takes a random action with proba- bility E and the best action according to the Q-function arg max, Q(s, a) with probability contrast, UCB is uncertainty-drive
E. In so it explores more actions that have been selected fewer times in the past, according to (aa s) = (a = arg max,' Q(s, a') + evlog(t)/N((s a')), where & is the total number of timesteps, e is the exploration coefficient, and (ss, a) is the number of times the agent has taken action in state s, with ties broken randomly'. PG uses the policy gradient method with reward-to-go for variance reduction and entropy bonus for exploration. While Greedy is a naive but widely used exploration strategy (Sutton & Barto, 2018). UCB is an effective algorithm designed for multi- multi-armed bandits (Auer et al., 2002: Audibert et al., 2009). Building on their in- tuition, Chen et al. (2017) showed that uncertainty -baed exploration bonus also performs well in challenging RL. environments. See Appendix A for more details about the experimental setup.
Each method's 's exploration strategy is controlled by single hyperparamter. For each hyperparame- ter, we search over 10 values and run every value for 100 trials. Each trial lasts for 1000 episodes. The results (mean and standard deviation) of hyperparameters with the highest average test returns for each method are shown in Figures 2b, 2c and 2d. We measure the performance of each method by its suboptimality, i.e., the difference between the undiscounted return achieved by the learned policy and the undiscounted return of the optimal policy. While all three methods are able to quickly achieve optimal return for the training MDP, their performances on the test MDP differ drastically. First, we observe that Greedy has the worst generalization performance and the highest variance. On the other hand, UCB can reliably find the optimal policy for the test MDP as the final return has negligible variance. Finally, PG's performance lies between Greedy and UCB. On average, converges to a better solution with lower variance than Greedy but is not as good as UCB. In Appendix A.3, we provide another set of simulations based changing dynamics and make similar observations. These experiments show that more effective exploration of the training environ- ments can result in better generalization to new environments.
is natural to ask whether this example is relevant for realistic deep RL problems. One hypothesis is that this tabular CMDP with two initial state distributions captures a common phenomenon in more challenging CMDPs like Procgen, namely that at test time, the agent can often end up in states that are suboptimal for the training MDPs. Having explored such states during training can help the agent recover from suboptimal states at test time. See Appendix for an illustrative example inspired by one of the Procgen games. This is similar to covariate shift in imitation learning where the agent's 's suboptimality can compound over time. For the purpose of generalization, the effect of efficient exploration is, spirit, similar to that of DAgger (Ross et al., 201 1) it helps the agent learn how to recover from suboptimal situations that are rare during training.
Another potential hypothesis for the superior performance of UCB relative to Greedy that better exploration improves sample efficiency resulting in better training performance, which can in turn lead to better test performance. Indeed, even in the simple tabular MDPs from Figures 2 and 8, UCB converges faster the training MDP. In Appendix E.I, we study this hypothesis in more detail. Note that the above two hypotheses can be simultancously true, but the latter assumes the model learns generalizable representations, whereas the previous hypothesis applies regardless. In both cases, adapting uncertainty-drie exploration to deep RI. on complex CMDPs has additional challenges such representation learning and unobserved contexts. In the next section, we describe these challenges and propose method to mitigate them.
'This UCB is a simple extension of the widely used bandits algorithm to MDP and does not enjoy the same regret guarantee, but we find it to be effective for didactic purpose. Azar et al. (2017) presents formal but much more complicated extension of the UCB algorithm for value iteration on MDPs. The UCB used still represents uncertainty but is more similar a count-hased exploration bonus.
4

---

Under review as a conference paper at ICLR 2023
Note that PG's softmax exploration strategy is not more sophisti- cated than E-greedy so its performance gain over Greedy should be attributed to something else. We hypothesize that this improve- ment is due to the fact that algorithms based on policy gradient use whole-traecctory updates (i.e., Monte Carlo methods) rather than one-step updates like Q-learning (i.e., TD(O)). To test this hypoth- esis, we run Sarsa(A) A) on the same environment. SArsa(A) (Sutton & Barto, 2018, Ch 12.7) is an on-policy method that learns Q- value but also does whole-trajectory updates. We observe that with e == 0.9 and À = 0.9, (Arsa(A) outperforms both Greedy and PG, but still worse than UCB and has a much larger variance (Figure 3). This suggests that although whole-trajectory updates may partially be why policy gradient generalizes better than Q-learning, better exploration can improve generalization further. In the rest of the paper, we will focus on using exploration to improve generalization. Details and further analysis are in Appendix A.2.
Sarsa(A) e.greedy
3.
Suboptimaalty
3.0
2.5-
2.0
1.3
1.0
0.0
opisade
Figure 3: The results Sarsa (A).
EXPLORATION VIA DISTRIBUTIONAL ENSEMBLE
In the previous section, we showed we can improve generalization via exploration in the tabular setting. In this section, we would like to extend this idea to deep RL with function approximation. in
While in the tabular setting shown above there is no intrinsic stochasticity, environments general be stochastic (e.g., random transitions, or random unobserved contexts). At a high level, epistemic uncertainty reflects lack of knowledge which can be addressed by collecting more data, while alcatoric uncertainty reflects the intrinsic noise in the data which cannot be reduced regardless of how much data is collected. One goal of exploration can be gather information about states with high epistemic uncertainty (O' Do'Dnnoguue et al., 2018) since aleatoric uncertainty cannot be reduced, but typical estimates can contain both types of uncertainties (Kendall Gal, 2017). In CMDPs, this is particularly important because a large source of aleatoric uncertainty is not knowing which context the agent is in (Raileanu & Fergus, 2021).
In this section, we introduce Exploration via Distributional Ensemble (EDE), method that en- courages the exploration of states with high epistemic uncertainty, which computed using deep ensembles and distributional RL. While ensembles are useful way of measuring uncertainty in neural networks (Lakshminarayanan et al., 2017), such estimates typically contain both epistemic and aleatoric uncertainty. Here, we build on Clements et al. (2019) who introduced approach for disentangling the epistemic uncertainty from the aleatoric uncertainty of the learned Q-values.
Uncertainty Estimation. Clements et al. (2019) showed that learning the quantiles for QR- DQN (Dabney et al., 2018) can be formulated as Bayesian inference problem, given access to a posterior p(0 D). where 0 is the discretized quantiles of Z(s, a), and D is the dataset of expe- rience on which the quantiles are estimated. j E [N] denote the index for the jth quantile and U11, denote the uniform distribution over integers between 1 and N. The uncertainty of the Q-value, Q(8, a) = EE (0, a)], is the relevant quantity that can inform exploration. The overall uncertainty x² = VaIo~r(O(p) [Q(s, a)] can be decomposed into epistemic uncertainty o and aleatoric uncertainty on such that 0² = o22p + ale where,
d (s, s, a) BBJuEINN [Var (0 a)]] (s a) = VarysueLLN) [Ba= (0, a)]]
(1)
(2)
Ideally, given an infinite or sufficiently large amount of diverse experience, one would expect the posterior to concentrate on the true quantile 0*, so Vare>p(0(p) [0,(s,a)] and consequently ( a) 0. oa would be non-zero if the true quantiles have different values. Intuitively, to improve the sample efficiency, the agent should visit state-action pairs with high epistemic uncer- tainty in order to learn more about the environment (Chen et al., 2017). It should be noted that the majority of the literature on uncertainty estimation focuses supervised learning; in RL, due to various factors such as bootstrapping, non-stationarity, limited model capacity, and approximate sampling, the uncertainty estimation generally contains ertors but empirically even biased epistemic uncertainty is beneficial for exploration. We refer interested readers to Charpentier et al. (2022) for more thorough discussion of this topic.
5

---

Under review as a conference paper at ICLR 2023
Sampling from p(0 D) is computationally intractable for complex MDPs and function approxi- mators such as neural networks. Clements et al. (2019) approximates samples from p(@ (D) with randomized MAP sampling (Pearce al., 2020) which assumes Gaussian prior over the model parameters. However, the unimodal nature of a Gaussian in parameter space may not have enough diversity for effective uncertainty estimation. Many works have demonstrated that deep ensembles tend to outperform regular ensembles and other approximate posterior sampling techniques (Laksh- minarayanan et al., 2017; Fort et al., 2019) in supervised learning. Motivated by these observations, we propose maintain M copies of fully-connected value heads, 9i Rd RIAIX that share single feature extractor / S + R similar to Osband et al. (2016a). However, we train each value head with different mini-batches and random initialization (i.e. deep ensemble) instead of distinct data subsets (i.e.. bootstrapping). This is consistent with Lee et al. (2021) which shows deep en- sembles usually perform better than bootstrapping for estimating uncertainty of Q-values but, unlike EDE, they do not decompose uncertainties.
Concretely, i E [M] is the index for M ensemble heads of the Q-network, and j E [N] is the index of the quantiles. The output of the 1h head for state and action a is (s, E IR where the 5th coordinate, and (1, a) is the jhh quantile of the predicted state-action value distribution for that head. The finite sample estimates of the two uncertainties are:
N
M
N
(5, a) N-M Σ
Σ 1=1
(o, a) -,,0000,sss)) a)) (s,
Σ (0,(s,a) - a(s,a))² (3)
N
where as = M EM 0,y(s, a) and ((s,a) = N (s)
Exploration Policy. There are two natural ways to use this uncertainty. The first one is by using Thompson sampling (Thompson, 1933) where the exploration policy defined by sampling Q- values from the approximate posterior:
Tes a 8) = I arg max, é(a,a') (a), where a(s,aa)= ~ (Q(s,, a'), t0 (s, a'))
(4)
p R2o is exploration coefficient that controls how the agent balances exploration and exploita- tion. Alternatively, we can use the upper-confidence bound (UCB, Chen et al. (2017)):
[a 8) La- (a), where a' arg max Q(s, a') + ip depi a')
a'
which we found to achieve better results in CMDPs than Thompson sampling (used in Clements al. (2019)) on Procgen when we use multiple parallel workers to collect experience, especially when combined with the next technique.
Equalized Exploration. Due function approximation, the model may lose knowledge of some parts of the state space if it does not see them often enough. Even with UCB, this can still happen after the agent learns a good policy on the training environments. To increase the data diversity, we propose to use different exploration coefficients for each copy of the model used to collect experi- ence. Concretely, we have K actors with synchronized weights; the h actor collects experience with the following policy:
(k) Tuch (al 8) = La- (a), where a" = arg max Q(s,a') + ((x @ep((s,
(6)
a'
E (0, 1) and a E R2 are hyperparameters that control the shape of the coefficient distribution. We will refer this technique as temporally equalized exploration (TEE). TEE is similar to Horgan et al. (2018) which uses different values of for the e-greedy exploration for each actor. In prac- tice, the performances are not sensitive to À and a (see Figure 14 in Appendix E). Both learning and experience collection take place on single machine and we do not use prioritized experience replay (Schaul et al., 2015) since prior work found it ineffective in CMDPs (Ehrenberg al., 2022).
To summarize, our agent explores the environment in order to gather information about states with high epistemic uncertainty which measured using ensembles and distributional RI. We build on the algorithm proposed in Clements et al. (2019) for estimating the epistemic uncertainty, but we use deep ensemble instead of MAP (Pearce et al., 2020), and use either UCB or Thompson sampling depending on the task and setting. In addition, for UCB, we propose that each actor uses a different exploration coefficient for more diverse data. While variations of the components of our algorithm have been used in prior works, this particular combination is new (see Appendix C). Our ablation experiments show that each design choice is important and naïvely combining these techniques does not work as well.

---

Under review as a conference paper at ICLR 2023
Median
IQM
Mean
Optimality Gap
QR-DON PPO
IDAAC
EDE (ours)
0.15
0.30
0.45
0.2
0.24 Min-Max Normalized Score
0.3
0.4
0.18
0.30 0.36
0.66 0.72 0.78 0.84
Figure 4: Test performance of different methods on the Procgen benchmark across 5 runs. Our method greatly outperforms all haselines in terms of median and IQM (the more statistically robust metrics), and is competitive with the state-of-thh-arr policy optimization methods in terms of mean and optimality gap. Optimality gap equal 1-mean for the evaluation configuration we chose.
EXPERIMENTS
5.1 PROCGEN
We compare our method with 3 representative baselines on the standard Procgen benchmark (i.e., 25M steps and easy mode) as suggested by Cobbe et al. (2018): (1) QR-DQN which is the prior state-of-the-art value-based method on Procgen (Ehrenberg et al., ;022); (2) PPO which is popular policy optimization baseline on which most competitive methods are built; and (3) IDAAC which is state-of-the the game bi g i sh only and evaluate all algorithms using the metrics proposed in Agarwal et al. (2021b). We run each algorithm on every game for 5 seeds and report the aggregated min-max normalized scores on the full test distribution, and the estimated hootstrap-estimated 95% confidence interval in Figure (simulated with the runs). Our approach significantly outperforms the other baselines in terms of median and interquartile mean (IQM) (which are the more statistically robust metries according to Agarwal et al. (2021b)). In particular, it achieves almost 3 times the median score of QR-DQN and more than 2 times the IQM of QR-DON. In terms of mean and optimality gap, our method is competitive with IDAAC and outperforms all the other baselines. To the best of our knowledge, this is the first value-based method that achieves such strong performance on Proegen. See Appendices G, D, and H for more details about our experiments, results, and hyperparameters.
Procgen and is built on PPO. We tune all hyperparameters of our method on
Ablations. We aim to better understand how each component of our algorithm contributes to the final performance by running a number of ablations. addition, we compare with other popular ex- ploration techniques for value-based algorithms. Since many of the existing approaches are designed for DQN, we also adapt a subset of them to QR-DQN for a complete comparison. The points of comparison we are: (1) Bootstrapped DQN (Osband & Van Roy, 2015) which uses bootstrapping to train several copies of models that have distinct exploration behaviors, (2) UCB (Chen et al., 2017) which uses an ensemble of models to do uncertainty estimation, (3) ez-greedy exploration (Dabney et al., 2021) which repeats the same random action (following zeta distribution) to achieve tempo- rally extended exploration, (4) UA-DQN (Clements al., 2019), and (5) NoisyNet (Fortunato al., 2017) which adds trainable noise to the linear layers. When UCB is combined with QR-DQN, we use the epistemic uncertainty unless specified otherwise. Details can be found in Appendix G.
First, note that both using the epistemic uncertainty via UCB and training on diverse data from TEE are crucial for the strong performance of EDE, with QR-DQN+TEE being worse than QR- DQN+UCB which is itself worse than EDE. Without using the epistemic uncertainty, the agent cannot do very well even if it trains on diverse data, Le., EDE is better than QR-DQN+TEF. Sim- ilarly, even if the agent uses epistemic uncertainty, it can still further improve its performance by training on diverse data, i.e., EDE is better than QR-DQN+UCB
Both Bootstrapped DQN and DQN+UCB, which minimize the total uncertainty rather than only the epistemic one, perform worse than DQN, although both of them are competitive exploration methods on Atari. This highlights the importance of using distributional RL in CMDPs in order to measure epistemic uncertainty. QR-DQN-UCB on the other hand does outperform QR-DQN and DQN because it seeks to reduce only the epistemic uncertainty. In Figure 12c from Appendix E, show that indeed exploring with the total uncertainty o² performs significantly worse at test time than exploring with only the epistemic uncertainty o² UA-DQN also performs worse than QR- DQN+UCB suggesting that deep ensemble may have better uncertainty estimation (Appendix F).
QR-DQN with e2 ez-greedy exploration marginally improves upon the base QR-DQN, but remains significantly worse than approach. This may be due to the fact that ez-greedy exploration can
7

---

Under review as a conference paper at ICLR 2023
Median
IQM
Mean
DON
QR-DQN
UA-DON
Bootstrapped DQN
DQN+ UCB
UR-DQN+UCB
DON+ TEE
QR-DQN+TEE
DON+ e-eeeedyyy
QR-DQN 1e0++-e--rreee ez-greedy
DON+ NoisyNett
OR-DON+NOisy/e
QR-DQN+UCB+TTE (EDE)
0.00
0.30
0.45
0.00
0.15
0.30
0.45
0.00
0.30
Min-Max Normalized Score
Figure 5: Test performance of different exploration methods on the Proegen henchmark across 5 runs.
induce temporally extended exploration but is not aware of the agent's 's uncertainty. Not accounting for uncertainty can be detrimental since CMDPs can have a much larger number of effective states than singleton MDPs. If the models have already gathered enough information about state, further exploring that state can hurt sample efficiency, regardless of whether the exploration temporally extended or not.
NoisyNet performs better than the other points of comparison we consider, but remains worse than EDE. Intuitively, the exploration behaviors of NoisyNet naturally adaptive the agents will take into account what they have already learned. While direct theoretical comparison between NoisyNet and our method is hard establish, believe adaptivity is common thread for methods that perform well on CMDPs. Nonetheless, if we consider IQM, none these methods significantly outperforms one another whereas our method achieves much higher IQM. Note that our TEE cannot be applied to NoisyNets and ez-greedy which already use an implicit "schedule".
5.2 CRAFTER
To test the generality our method beyond the Procgen bench- mark, conduct experiments the Crafter environment (Hafner, 2022). Making progress on Crafter requires a wide range of capabil- ities such strong generalization, deep exploration, and long-term reasoning. Crafter evaluates agents using a score that summarizes the agent's abilities into a single number. Each episode is proce- durally generated, so the number of training environments is prac- tically infinite. While Crafter does not test generalization to new environments, it still requires generalization across the training en- vironments in order to efficiently train on all of them. We build method on top of the Rainbow implementation (Hessel et al., 2018) provided in the open- open-sourced code of Hafner (2022). We use Thompson sampling instead of UCB+TEE since only one environment is used. As seen in Table 5, algorithm achieves signif- icantly higher scores compared to all the baselines presented Hafner (2022), including Dream- erV2 (Hafner al., 2021) which is a state-of-the-art model-based RI. algorithm. The significant improvement over Rainbow, which is competitive value-based approach, suggests that the explo- ration strategy is crucial for improving performance on such CMDPs.
Method
Score (%)
EDE (ours)
11.7 + 1.0
Rainbow
4.3 - 0.2
PPO
4.6 + 0.3
DreamerV2
10.0 1.2
Table 1: Results on Crafter ter 1M steps and over 10 runs.
RELATED WORK
Generalization in RL. large body work has emphasized the challenges of training RL agents that can generalize to new environments and tasks (Rajeswaran et al., 2017: Machado et al., 2018; Justesen et al., 2018; Packer et al., 2018; Zhang et al., 201 8a;b; Nichol et al., 2018; Cobbe et al., 2018; 2019; Juliani et al., 2019; Kuttler et al., 2020; Grigsby & Qi, 2020; Chen, 2020; Bengio al.,
2020: Bertrán et al., 2020; Ghosh et al., 2021; Kirk et al., 2021; Ajay et al., 2021; Ehrenberg et al.,
2022; Lyle et al., 2022). This form of generalization is different from generalization in singleton
8

---

Under review as a conference paper at ICLR 2023
MDP which refers to function approximators generalizing to different states within the same MDP. A natural way to alleviate overfitting is to apply widely-used regularization techniques such as implicit regularization (Song et al., 2020), dropout (Igl et al., 2019), batch normalization (Farebrother et al., 2018), or data augmentation (Ye et al., 2020; Lee et al., 2020; Laskin et al., 2020; Raileanu et al., 2020; Wang et al., 2020; Yarats et al., 2021; Hansen & Wang, 2021; Hansen et al., 2021; Ko & Ok, 2022). Another family of methods aims to learn better state representations via bisimulation metries (Zhang et al., 2020b;a; Agarwal et al., 2021a), information bottlenecks (Igl al., 2019; Fan & Li, 2022), attention mechanisms (Carvalho et al., 2021), contrastive learning (Mazoure et al,, 2021), adversarial learning (Roy & Konidaris, 2020; Fu et al., 2021; Rahman & Xue, 2021), or decoupling representation learning from decision making (Stooke et al., 2020; Sonar et al., 2020). Other approaches use information-theoretti approaches (Chen, 2020; Mazoure et al., 2020), non- stationarity reduction (Igl et al., 2020; Nikishin et al., 2022), curriculum learning (Jiang et al., 2020;
Team et al., 2021: Jiang et al., 2021; Palder-Holder et al., 2022), planning (Anand et al., 2021), forward-back backward representations (Touati & Ollivier, 2021), diverse policies (Kumar et al., 2020). More similar to our work, Raileanu& & Fergus (2021) show that the value function can overfit when trained on CMDPs and propose decouple the policy from the value optimization train more robust policies. However, this approach cannot be applied to value-based methods since the policy is directly defined by the Q-function. Most of the above works focus on policy optimization methods, and none emphasizes the key role exploration plays in training more general agents. In contrast, our goal is to understand why value-based methods are significantly worse on CMDPs.
Ex Exploration. Exploration is a fundamental aspect of RL (Kolter & Ng. 2009; Fruit & Lazaric, 2017; Tarbouricch et al., 2020). Common approaches include s-greedy (Sutton & Barto, 2018), count- based exploration (Strehl & Littman, 2008; Bellemare et al., 2016; Machado et al., 2020), curiosity- based exploration (Schmidhuber, 1991; Stanton Clune, 2016; Pathak et al., 2017), or novelty- based methods specifically designed for exploring sparse reward CMDPs (Raileanu & Rocktäschel, 2020; Zha et al., 2021; Flet-Berliac et al., 2021; Zhang et al., 2021). These methods are based on policy optimization and focus on training agents in sparse reward CMDPs. In contrast, we are in- terested in leveraging exploration as a way of improving the generalization of value-based methods to new MDPs (including dense reward ones). Some of the most popular methods for improving exploration in value-based algorithms use noise (Osband et al., 201 Fortunato et al., 2017), boot- strapping (Osband & Van Roy, 2015; Osband et al., 6a), ensembles (Chen et al., 2017; Lee al., 2020; Liang et al., 2022), uncertainty estimation (Osband et al., 2013; O' O'Donoghue et al., 2018; Clements et al., 2019), or distributional RL (Mavrin al., 2019). The main goal of these works was to use exploration to improve sample efficiency on singleton MDPs that require temporally ex- tended exploration. Another related area is agsk-agnostic RL (Pong et al., 2019; Zhang et al., 2020c) where the agent explores the environment without reward and tries to learn a down-stream task with reward, but to our knowledge these methods have not been successfully adapted to standard bench- marks like Procgen where the MDPs can have large variations. Our work is the first one to highlight the role of exploration for faster training on contextual MDPs and better generalization to unseen MDPs. Our work also builds on the distributional RL perspective (Bellemare et al., 2017), which is useful in CMDPs for avoiding value overfitting (Raileanu & Fergus, 2021).
DISCUSSION
In this work, aim understand why value-based methods are significantly worse than policy optimization methods in contextual MDPs, while being much better in singleton MDPs with similar characteristics. Our tabular experiments indicate that effective exploration is crucial for generaliza- tion to new environments. In CMDPs, exploring an environment is not only useful for finding the optimal policy in that environment, but also for acquiring knowledge that can be useful in other envi- ronments the agent may encounter at test time. In order to make value-based competitive with policy optimization in CDMPs, we propose an approach based on uncertainty-diit -diveen exploration and dis- tributional RL. This is the first value-based based method to achieve Wate--ofthhe-eee performance on both Procgen and Crafter, two challenging benchmarks for generalization in RL. Our experiments suggest that exploration is important for all RL algorithms trained and tested on CMDPs. While here we focus on value-based methods, similar ideas could be applied in policy optimization to further improve their generalization abilities. One limitation of our approach is that it is more computa- tionally expensive due to the ensemble (Appendix F). Thus, we expect it could benefit from future advances in more efficient ways of accurately estimating uncertainty in neural networks.
9

---

Under review as a conference paper at ICLR 2023
8 REPRODUCIBILITY STATEMENT
We include the source code of EDE in the supplementary materials. This work builds on several existing open-source repositories and we have provided a detailed description of the algorithm, hyperparameters, and architecture in Appendix C and G. The codebase for Procgen is built on com- ponents from Jiang et al. (2020) and Ehrenberg et al. (2022); the codebase for Crafter is built on baseline provided by Hafner (2022). Access to these codebases can all be found in these papers.
REFERENCES
Rishabh Agarwal, Marlos C. Machado, Castro, and Mare G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. 2021a.
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural informa- tion processing systems, 14:29304-293200 2021b.
Anurag Ajay, Ge Yang, Ofir Nachum, and Pulkit Agrawal. Understanding the generalization gap in visual reinforcement learning. 2021.
Ankesh Anand, Jacob Walker, Yazhe Li, Eszter Vértes, Julian Schrittwiese,, Sherjil Ozair, Théophane Weber, and Jessica B Hamrick. Procedural generalization by planning with self- supervised world models. arXiv preprint arXiv:2111.0158) 2021.
Jean- Yves Audibert, Rémi Munos, and Csaba Szepesvári. Exploration tradcoff using variance estimates in multi-armed bandits. Theoretical Computer Science, 410(19):1876-1902, 2009.
Peter Auer, Nicolo Cesa-Bianchi, and Paul Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47(2):235-256, 2002.
Mohammad Gheshlaghi Azar, Ian Osband, and Rémi Munos. Minimax regret bounds for reinforce- ment learning. In International Conference on Machine Learning, pp. 263-272. PMLR, 2017.
Adrià Puigdomenech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark. In International Conference on Machine Learning, pp. 507-517. PMLR, 2020.
Marc Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. Advances neural information pro- cessing systems, 29, 2016.
Mare G Bellemare, Will Dabney, and Rémi Munos. distributional perspective on reinforcement learning. In International Conference Machine Learning, pp. 449-458. PMLR, 2017.
Richard Bellman. markovian decision process. Journal of mathematics and mechanics, pp. 679
684, 1957.
Emmanuel Bengio, Joelle Pineau, and Doina Precup. Interference and generalization temporal difference learning. ArXiv, abs/2003.06350, 2020.
Martín Bertrán, Martínez, Mariano Phielipp, and G. Sapiro. Instance based generalization in reinforcement learning. ArXiv, abs/2011.0100 2020.
Wilka Carvalho, Andrew Lampinen, Kyriacos Nikiforou, Felix Hill, and Murray Shanahan. Feature-attending recurrent modules for generalization in reinforcement learning. arXiv preprint arXiv:2112.08866, 2021.
Bertrand Charpentier, Ransalu Senanayake, Mykel Kochenderfer, and Stephan Günnemann. Dis- entangling epistemic and aleatoric uncertainty in reinforcement learning. arXiv preprint arXiv:2206.01155. 2022.
10

---

Under review as a conference paper at ICLR 2023
Jerry Zikun Chen.
Reinforcement learning generalization with surprise minimization. 0bss2204.11399) 2020.
ArXiv,
Richard Y Chen, Szymon Sidor, Pieter Abbeel, and John Schulman. Ucb exploration via qr ensembles. arXiv preprint arXiv:1706.01500 2017.
William R Clements, Bastien Van Delft, Benoît-Marie Robaglia, Reda Bahi Slaoui, and Sébastien Toth. Estimating risk and uncertainty in deep reinforcement learning. arXiv preprini (rxxiv:905.09000 2019.
Karl Cobbe, Oleg Klimov, Chris Hesse, Taehoon Kim, and John Schulman. Quantifying generaliza- tion in reinforcement learning. arXiv preprint arXiv: (cXivv/8/2.0234/ 2018.
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural genera- tion benchmark reinforcement learning. arXiv preprint drXiv:1912.01588, 2019.
Will Dabney, Mark Rowland, Mare Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelli- gence, volume 32, 2018
Will Dabney, Georg Ostrovski, and Andre Barreto. Temporally-extenie E-greedy exploration. In International Conference on Learning Representations, 2021. URL. https: //tp/////ee.././...////eee. net/ forum? 1d=ONBPHE272G..
Andy Ehrenberg, Robert Kirk, Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. A study of off-policy learning in environments with procedural content generation. ICLR Workshop on Agent Learning in Open-Endedness, 2022.
Lasse Espcholt, Hubert Soyer, Remi Munos, Karen Simonyan, Volodymir Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. arXiv preprint arXiv: 180.00/00....,,, 2018.
Jaamnng Fan and Wenchao Li. Dribo: Robust deep reinforcement learning via multi-view infor- mation bottleneck. In International Conference on Machine Learning, pp. 6074-6102. PMLR, 2022.
Jesse Farebrother, Marlos C. Machado, and Michael H. Bowling. Generalization and regularization in dqn. ArXiv, abs/1810.00123, 2018.
Yannis Flet-Berliac, Johan Ferret, Olivier Pietquin, Philippe Preux, and Matthieu Geist. Adversar- ially guided actor-critic. CoRR, abs/2102.04376, 2021. URL https://arxxv.crrrreee/// abss 2102. 04376.
Stanislav Fort, Huiyi Hu, and Balaji Lakshminarayanan. Deep ensembles: A loss landscape per- spective. arXiv preprint 1xXiv:191..02757, 2019.
Meire Fortunato, Mohammad Gheshlaghi Azar, Bilal Piot, Jacob Menick, Ian Osband, Alex Graves, Vlad Mnih, Remi Munos, Demis Hassabis, Olivier Pietquin, et al. Noisy networks for exploration. arXiv preprint arXiv: 7066. 10295, 2017.
Ronan Fruit and Alessandro Lazaric. Exploration-exxloiiation in mdps with options. In Artificial intelligence and statistics, pp. 576-584. PMLR, 2017.
Xiang Fu, Ge Yang. Pulkit Agrawal, and Tommi Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning, pp. 3480-3491. PMLR, 2021.
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. Advances in Neural Information Processing Systems, 34:25502-25515, 2021.
Jake Grigsby and Yanjun Qi. Measuring visual generalization in continuous control from pixels. ArXiv, Abs/2010.06744 2020.
11

---

Under review as a conference paper at ICLR 2023
Danijar Hafner. Benchmarking the spectrum of agent capabilities. In International Confer- ence on Learning Representations, 2022. URL https: /oppr///eee/eeeeieeeeoeee net forum?ide UNCZ96MFEo1
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In International Conference Learning Representations, 2021. URL https:/ wwpp/ppeeeviiieeeeeeeeeeeeeeeeee net/ Forum?id=0oabwy:
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv preprint arXiv: 1502.02259 2015.
Nicklas Hansen and Xiaolong Wang. Generalization reinforcement learning by soft data aug- mentation. In 2021 IEEE International Conference Robotics and Automation (ICRA), pp. 13611-13617. IEEE, 2021.
Nicklas Hansen, Hao Su, and Xiaolong Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. Advances in Neural Information Processing Systems, 34: 3680-3693, 2021.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp- 770-778, 2016.
Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver. Rainbow: Combining improvements in deep reinforcement learning. In Thirty-second AAAI conference on artificial intelligence, 2018.
Dan Horgan, John Quan, David Budden, Gabriel Barth-Maron, Matteo Hessel, Hado van Hasselt, and David Silver. Distributed prioritized experiecce replay. In International Conference on Learn- ing Representations, 2018. URL https: opwprevieeeeeeeeeeeeeeeeeeeeeeee net/ forum -id--mmd-------------
Maximilian Igl, Kamil Ciosek, Yingzhen Li, Sebastian Tschiatschek, Cheng Zhang, Sam Devlin, and Katja Hofmann. Generalization reinforcement learning with selective noise injection and information bottleneck. In Advances in Neural Information Processing Systems, pp. 13956- 13968, 2019.
Maximilian Igl. Gregory Farquhar, Jelena Luketina, Wendelin Böhmer, and Shimon White- son. The impact of non-stationarity on generalisation deep reinforcement learning. ArXiv, abs/2006.05826, 2020.
Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. Prioritized level replay. ArXiv, abs/2010.03934, 2020.
Minqi Jiang, Michael Dennis, Jack Parker-Holdee Jakob Foerster, Edward Grefenstette, and Tim Rocktäschel. Replay-guided adversarial environment design. Advances in Neural Information Processing Systems, 34:188441881 1897, 2021.
Arthur Juliani, Ahmed Khalifa, Vincent-Pierre Berges, Harper, Hunter Henry, Adam Crespi, J. To- gelius, and D. Lange. Obstacle tower: A generalization challenge in vision, control, and planning. ArXiv, abs/1902.01378. 2019.
Niels Justesen, Ruben Rodriguez Torrado, Philip Bontrager, Ahmed Khalifa, Julian Togelius, and Sebastian Risi. Illminatiin generalization in deep reinforcement learning through procedural level generation. arXiv: Learning, 2018.
Alex Kendall and Yarin Gal. What uncertainties do need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel. A survey of generalisation in deep reinforcement learning. arXiv preprint 1rXiv:211100979) 2021.
Byungchan Ko and Jungseul Ok. Efficient scheduling of data augmentation for deep reinforcement learning. arXiv preprint 2rX1v:2206.005188 2022.
12

---

Under review as a conference paper at ICLR 2023
J Zico Kolter and Andrew Y Ng. Near-bayesian exploration in polynomial time. In Proceedings of the 26th annual international conference on machine learning, pp. 513-520, 2009.
Ilya Kostrikov, Denis Yarats, and Rob Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
Saurabh Kumar, Aviral Kumar, Sergey Levine, and Chelsea Finn. One solution is not all you need: Few-shot extrapolation via structured maxent rl. Advances in Neural Information Processing Systems, 33:81 98 8210, 2020.
Heinrich Kuttler, Nantas Nardelli, Alexander H. Miller, Roberta Raileanu, Marco Selvatici, Edward Grefenstette, and Tim Rocktäschel. The nethack learning environment. ArXiv, abs/2006.13760, 2020.
Balaji Lakshminarayanan. Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances neural information processing systems, 30, 2017.
Michael Laskin, Kimin Lee, Adam Stooke, Lerrel Pinto, Pieter Abbeel, and Aravind Srinivas. Re- inforcement learning with augmented data. arXiv preprint 20t:::2004.14990 2020.
Kimin Lee, Kibok Lee, Jinwoo Shin, and Honglak Lee. Network randomization: A simple tech-
nique for generalization in deep reinforcement learning. In International Conference on Learning Representations. https://openreview. net/forum, 2020.
Kimin Lee, Michael Laskin, Aravind Srinivas, and Pieter Abbeel. Sunrise: A simple unified frame- work for ensemble learning in deep reinforcement learning. In International Conference on Ma- chine Learning, pp. 6131-6141. PMLR, 2021.
Litian Liang, Yaosheng Xu, Stephen McAleer, Dailin Hu, Alexander Ihler, Pieter Abbeel, and Roy value estimation via ensemble of deep networks.
Fox. Reducing variance in In International Conference on Machine Learning, pp.
PMLR, 2022.
Evan Z Liu, Aditi Raghunathan, Percy Liang, and Chelsea Finn. Decoupling exploration and ex- ploitation for meta-reinforcemmet learning without sacrifices. In International conference on ma- chine learning, pp.
PMLR, 2021.
Clare Lyle, Mark Rowland, Will Dabney, Marta Kwiatkowska, and Yarin Gal. Learning dynamics and generalization in reinforcement learning. arXiv preprint arXiv:220.002126 2022.
Marlos C. Machado, Mare G. Bellemare, Erik Talvitie, Joel Veness, Matthew Hausknecht, and Michael H. Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. In IJCAI, 2018.
Marlos C Machado, Mare Bellemare, and Michael Bowling. Count-based exploration with the successor representation. In Proceedings of the AAAI Conference on Artificial Intelligence, vol- ume 34, pp. 5125-5133, 2020.
Dhruv Malik, Yuanzhi Li, and Pradeep Ravikumar. When generalizable reinforcement learning tractable? Advances in Neural Information Processing Systems, 34:8032-8045, 2021.
Borislav Mavrin, Hengshuai Yao, Linglong Kong, Kaiwen Wu, and Yaoliang Yu. Distributional reinforcement learning for efficient exploration. In International conference on machine learning, pp. 4424-4433 PMLR, 2019.
Bogdan Mazoure, R' R'emi Tachet des Combes, Thang Doan, Philip Bachman, and R. Devon Hjelm. Deep reinforcement and infomax learning. ArXiv, fbs/2206.00211,, 2020.
Bogdan Mazoure, Ilya Kostrikov, Ofir Nachum, and Jonathan Tompson. Improving zero-shot gener- alization in offline reinforcement learning using generalized similarity functions. arXiv preprint arXiv:2/ /1. /4629, 2021.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wier- stra, and Martin Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv: 1312.5603 2013.
13

---

Under review as a conference paper at ICLR 2023
Jonas Mockus. The bayesian approach to global optimization. In System Modeling and Optimiza- tion, pp. 473 481.. Springer. 1982.
Sharada Mohanty, Jyotish Poonganam, Adrien Gaidon, Andrey Kolobov, Blake Wulfe, Dipam Chakraborty, Gražvydas Semetulskis, João Schapke, Jonas Kubilius, Jurgis Pašukonis, et al. Mea- suring sample efficiency and generalization in reinforcement learning benchmarks: Neurips 2020 procgen benchmark. arXiv preprint arXiv:2103.15332 2021.
Jun Morimoto and Kenji Doya. Robust reinforreeet learning. Neural computation, 17(2):335-35,, 2005.
Alex Nichol, V. Pfau, Christopher Hesse, O. Klimov, and John Schulman. Gotta learn fast: A new benchmark for generalization rl. ArXiv, abs/1804.03720, 2018.
Evgenii Nikishin, Max Schwarzer, Pierluca D'Oro, Pierre-Luc Bacon, and Aaron Courville. The primacy bias in deep reinforcement learning. In International Conference on Machine Learning, pp. 16828-16847. PMLR, 2022.
lan Osband and Benjamin Van Roy. Bootstrapped thompson sampling and deep exploration. arXiv preprint arXiv15577000300 2015.
lan Osband, Daniel Russo, and Benjamin Van Roy. (more) efficient reinforcement learning via posterior sampling. Advances in Neural Information Processing Systems, 26, 2013.
Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems, 29, 2016a.
Ian Osband, Benjamin Van Roy, and Zheng Wen. Generalization and exploration via random- ized value functions. In International Conference on Machine Learning, pp. 2377-2386. PMLR, 2016b.
Brendan O' O'Donoghue, Ian Osband, Remi Munos, and Volodymyr Mnih. The uncertainty bellman equation and exploration. In International Conference on Machine Learning, pp. 3836-3845, 2018.
Charles Packer, Katelyn Gao, Jernej Kos, Philipp Krähenbühl, Vladlen Koltun, and Dawn Xiaodong Song. Assessing generalization in deep reinforcement learning. ArXiv, abs/1810.12282, 2018.
Jack Polker-Holder, Minqi Jiang, Michael Dennis, Mikayel Samvelyan, Jakob Foerster, Edward Grefenstette, and Tim Rocktäschel. Evolving curricula with regret-based environment design. arXiv preprint arXiv:2203.013333 2022.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning, pp. 2778-2787. PMLR, 2017.
Tim Pearce, Felix Leibfried, and Alexandra Brintrup. Uncertainty in neural networks: Approxi- mately bayesian ensembling. In International conference on artificial intelligence and statisties, pp. 234-244 PMLR, 2020.
Vitchyr H Pong, Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine. Skew- fit: State-covering self-supervised reinforcement learning. arXiv preprint arXiv: arXiv:1903.00 1903.0369,, 2019.
Md Masudur Rahman and Yexiang Xue. Adversarial style transfer for robust policy optimization in reinforcement learning. 2021.
Roberta Raileanu and Rob Fergus. Decoupling value and policy for generalization in reinforcement learning. In International Conference on Machine Learning, pp. 8787-8798, PMLR, 2021.
Roberta Raileanu and Tim Rocktäschel. Ride: Rewarding impact-driven exploration for procedurally- environments. In International Conference on Learning Representations, 2020. URL https: /pwnrevieeeeeeeeeeee net fww.eoridee.---tmaeee TJRFPR.
14

---

Under review as a conference paper at ICLR 2023
Roberta Raileanu, M. Goldstein, Denis Yarats, llya Kostrikov, and R. Fergus. Automatic data aug- mentation for generalization in deep reinforcement learning. ArXiv, abs/2006.12862, 2020.
Aravind Rajeswaran, Kendall Lowrey, Emanuel Todorov, and Sham M. Kakade, Towards general- ization and simplicity in continuous control. ArXiv, abs/1703.00266) 2017.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. reduction of imitation learning and struc- tured prediction to no-regret online learning. In Proceedings of the fourteenth international con- ference on artificial intelligence and statistics, pp. 627-635. JMLR Workshop and Conference Proceedings, 2011.
Josh Roy and George Konidaris. Visual transfer for reinforcement learning via wasserstein domain confusion. arXiv preprint arXiv:2006.03466, 2020.
Tom Schaul, John Quan, Ioannis Antonoglou, and David Silver. Prioritized experience replay, arXiv preprint arXiv: 00..5511.005552 05552, 2015.
Jürgen Schmidhuber. A possibility for implementing curiosity and boredom in model-building neu- ral controllers. In Proc. of the international conference simulation of adaptive behavior: From animals to animats, pp. 222-227, 1991.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High- dimensional continuous control using generalized advantage estimation. arXiv preprint arXiv:1506.02438, 2015.
Burr Settles. Active learning literature survey. 2009.
Anoopkumar Sonar, Vincent Pacelli, and Anirudha Majumdar. Invariant policy optimization: To- wards stronger generalization in reinforcement learning. ArXiv, abs/2006.01096, 2020.
Xingyou Song, Yiding Jiang, Stephen Tu, Yilun Du, and Behnam Neyshabur. Observational overfit- ting in reinforcement learning. In International Conference Learning Representations, 2020.
URL /tttps://opente
net forum?
Christopher Stanton and Jeff Clune. Curiosity search: Producing generalists by encouraging indi- viduals to continually explore and acquire skills throughout their lifetime. PLOS ONE, 11(9): 1-20, 09 2016. doi: 10. 11..7//joumal..ooe..0066.. journal, pone. 01 62235.
URL
1371/
Adam Stooke, Kimin Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforcement learning. ArXiv, abs/2009.083 08319, 2020.
Alexander I. Strehl and Michael Littman. An analysis of model-based interval estimation for markov decision processes. Journal of Computer and System Sciences, 130-)11309-1331, 2008.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
Richard Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient meth- ods for reinforcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
Jean Tarbouriech, Evrard Garcelon, Michal Valko, Mattco Pirotta, and Alessandro Lazaric. No- regret exploration in goal-oriented reinforcement learning. In International Conference Ma- chine Learning, pp. 9428 9437. PMLR, 2020.
Open Ended Learning Team, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, Max Jaderberg, Michael Mathieu, et al. Open-ended learning leads to generally capable agents. arXiv preprint arXiv:2107.128000 2021.
William R Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-2944 1933.
Ahmed Touati and Yann Ollivier. Learning one representation to optimize all rewards. Advances in Neural Information P Processing Systems, 14:13-23, 2021.
15

---

Under review as a conference paper at ICLR 2023
K. Wang, Bingyi Kang, Jie Shao, and Jiashi Feng. Improving generalization in reinforcement leam- ing with mixture regularization. ArXiv, abs/2010. 108.00..0144 2020.
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8(3):279-292, 1992.
Ronald J Williams. Simple statistical gradient-followeee algorithms for connectionist reinforcement learning. Machine learning, 8(3):229-2) 22622, 1992.
Ronald J Williams and Jing Peng. Function optimization using comectioonistrriifffcceeeee learning algorithms. Connection Science, 3(3):241-268, 1991.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous trol: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645 2021.
Chang Ye, Ahmed Khalifa, Philip Bontrager, and Julian Togelius. Rotation, translation, and crop- ping for zero-shot generalization. arXiv preprint arXiv:2001.09900 2020.
Daochen Zha, Wenye Ma, Lei Yuan, Xia Hu, and Liu. Rank the episodes: simple approach for
exploration in procedurally-genendee environments. arXiv preprint arXiv:2101.08152, 2021.
Amy Zhang, Nicolas Ballas, and Joelle Pineau. A dissection of overfitting and generalization in continuous reinforcement learning. ArXiv, abs/ 10060997,7797 2018 8a.
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. arXiv preprint arXiv:2003.06016 2020a.
Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, and Sergey Levine. Learning invariant representations reinforcement learning without reconstruction. arXiv preprint 1rX1::2006.10742, 2020b.
Chiyuan Zhang, Oriol Vinyals, Rémi Munos, and Samy Bengio. A study overfitting in deep reinforcement learning. ArXiv, abs/ 0ts/1804.00/ 006.06 2018b.
Tianjun Zhang, Huazhe Xu, Xiaolong Wang, Yi Wu, Kurt Keutzer, Joseph E. Gonzalez, and Yuan- dong Tian. Noveld: A simple yet effective exploration criterion. In Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https: opwnrevieeeeeeeeeooeeeeeeee.oeee net/ forum?ideeee (crrmm?id=CCUzppnOEFF
Xuezhou Zhang, Yuzhe Ma, and Adish Singla. Task-agnostic exploration in reinforcement learning. Advances in Neural Information Processing Systems, 33:11734 1743, 2020c.
16

---

Under review as a conference paper at ICLR 2023
A
TABULAR EXPERIMENTS
A.1
GENERALIZATION TO DIFFERENT INITIAL STATES
05
05
-2
0.6
0.0
12.
-05
02.
opionde
opisode
apacdd
(a) e-greedy
(b) UCB
(c) Policy Gradient
Figure 6: Mean test return for each method with different exploration hyperparameters.
Both the training MDP, Murann and the test MDP, juest, share the same state space S, [5] [5]. and action space A, [4], corresponding to {left, right, up, down). From each state, the transition, P(s') a). to the next state corresponding to action is 100%. 7 is 0.9. The two MDPs differ by only their initial state distribution: (train() ((0.0) (s) and Prest(s) 6) = (4.0) (s).
For both policy gradient and Q-learning, training progresses. Specifically, at time step f, a₁ = vi -0-
base learning rate of 0.05 that decays as
For Q-learning, the Q-function is parameterized as E RRSSX|A| where each entry a state-action value. The update is:
Q(t+1) (s,)ss, at (re (re+ + max max (s+11, (st,
where Q(s, a) Ds.a- For Q-learning with UCB, the agent keeps N(ss a) that keeps track of how many times the agent has taken action a from s over the course of training and explore according to:
logt
Tubh (al () 1a. (a) where a" arg max Q(s, a) + +CV N(s,a)
(8)
For Q-learning with e-greedy, the exploration policy is:
Regreedy (a 8) = (1 s1 aar max,, Q(a,a)) (a) te
Σ Ea
la (a)
(9)
\a'Ed
Ties in arg max are broken randomly which acts the source of randomness for UCB. For policy gradient, the policy is parameterized as & E [R( exp(d>.a) The overall gradient is:
Al where each entry is the logit for the
distribution taking action a from 8, i.e., Ta 8) Eved exp(D....
T
VO [J(T)] Eng
Σ
v (ag st) Σ k rk
(10)
VoHl, = V VolET
T sx) st) log (s)
(11)
t=0 aeA
The update is:
g(++1) = ((() + a. (V.E., [J(T) + 3V (1167.91100 )1==000
(12)
where the expection is estimated with single trajectory.
For each method, repeat the experiment for the following 10 hyperparameter values:
17

---

Under review as a conference paper at ICLR 2023
{0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
(5, 10, 15, 20, 25, 30, 35, 40, 45, 50)
3: {0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
Each experiment involves training the algorithm for 100 trials. Each trial consists of 1000 episodes and each episode ends when the agent reaches a terminal state or at 250 steps. At test time, all ap- proaches are deterministic: Q-learning follows the action with highest Q-value and policy gradient follows the action with the highest probability. The mean test returns for each experiment are shown in Figure 6. We observe that both UCB and policy gradient are relatively robust to the choice of hy- perparameters, whereas e-greedy's generalization performance varies greatly with e. Furthermore, even with extremely large E, the generalization performance is still far worse than UCB and policy gradient. The best performing values are respectively ex = 0.9, c* = 45 and 3* = 0.1.
A.2 SARSA(A)
Sarsa(A) is an on-policy value-based method that learns Q-value with shole-trajectorn update via eligibility traces. Since the algorithmic details are more involved than the other methods we con- sider, we refer the reader to Chapter 12.7 of Sutton & Barto (2018) for the details. 'he Sarsa( (X) with e-greedy experiments inherit all the hyperparameters of Q-learning e-greedy experiments, but have one more hyperparameter À E (0, 1] that interpolates the method between Monte Carlo method and one-step TD update. We use fixed À = 0.9 for all experiments. We sweep e across 10 values (100 trials each):
* E: {0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9}
It can be seen from Figure 7a that even with whole-episode updates, the average return is still higher for larger exploration coefficients, which means exploration can still help. Furthermore, larger exploration coefficients also tend to have smaller variances in the test performance (Figure 7b).
Sarsa(A) different
Standard Deviation Sarsa(A) different
2.0
0.5-
0.5
0-
-0.3-
-0.3
200
400
600
200
400
episode
episode
(a) Average of return.
(b) Standard deviation of return. Figure 7: Tet performance of Sarsa(X) across different E.
A.3 GENERALIZATION TO DIFFERENT DYNAMICS
In previous section, we presented simple scenario that exploration helps generalization to different initial state distribution. Here, we demonstrate another scenario where exploration helps general- ization to new dynamics. We use 7 x grid as the state space (Figure 8a). The agent receives a small negative reward for every location except for the green square (4, 2). where the agent receives a reward of 2 and the episode ends. The difference between this environment and the previous one that the agent always starts at (0, 3) but at test time. gust of wind can blow the agent off course. Specifically, with a probability 40% the agent will be blown to +2 unit in the y-direction. This environment an example of the windy world from Sutton & Barto (2018). The results are shown
18

---

Under review as a conference paper at ICLR 2023
Reward
Qulearning E-groody
O-leaming UCB
Palicy Gradient
suboptimallt
2
6.54
230
400 600 episode
800
400 episode
episode
(a) Tabular CMDP
(b) E-greedy
(c) UCB
(d) Policy Gradient
Figure 8: (a) During both training and test time, the agent starts in the blue square, but test time, with 40% probability, wind can blow the agent down by two unit. both cases, the goal is to get to the green square. The other plots show the mean and standard deviation of the train and test suboptimality (difference between optimal return and achieved return) over 100 runs for (b) Q-learning with E- e-greedy exploration, (c) Q-lcarning with UCB exploration, and (d) policy gradient with entropy bonus.
different opsilion
UCB different exploration
Folicy gradient different
-6
-12
200
200
4:01
4.81
610
opisode
cp/oode
cp.cdd
(a) e-greedy
(b) UCB
(c) Policy Gradient
Figure 9: Mean test return for each method with different exploration hyperparameters when the transition
dynamic changes at test time.
in Figure 8 and 9, where we see similar results as the previous section, but at test time the optimal expected return is generally worse and has higher variance for all methods since there is stochasticity (suboptimality is computed with respect to the return of an oracle that consistently moves towards the goal). This setup is related robust reinforcement learning (Morimoto Doya, 2005), but exploration makes no assumptions about the type or magnitude of perturbation (thus making more complicated perturbation, e.g., different random initializations of the same game, possible) and does not require access to simulator.
B ADDITIONAL BACKGROUND
B.1 QUANTILE REGRESSION DQN
Quantile regression DQN (QR-DQN) is a variant of distributional RL that tries to learn the quantiles of the state-action value distribution with quantile regression rather than produce the actual distri- bution of Q-values. Compared to the original C51 (Bellemare et al., 2017), QR-DQN has many favorable theoretical properties. Here, we provide the necessary background for understanding QR- DQN and refer the readers to Dabney et al. (2018) for a more in-depth discussion of QR-DQN.
Given a random variable Y with associated CDF Fy (v) IPY S v)), the inverse CDF of Y is
(T' = inf (y E R T S Fy (v))
To approximate Y, we may discretize the CDF with qunatile values:
(T), (T)), F7 (Tx)} where
N
Using the same notation as main text, we use 0 S x A R to denote a parametric function that maps a state-action pair, (s, a), to the estimated quantiles (0, (s, a) N The approximated CDF
19

---

Under review as a conference paper at ICLR 2023
of (ss, a) is thus:
N Σ 8 (0, N
1
2(8, a)
where 5 (0, a)) is the Dirac delta distribution centered at (s, a). To learn the quantiles, we use quantile regression which allows for unbiased estimate of the true quantile values. Concretely, for larget distribution Z and quantile т, the value of Fz (T) is the minimizer of the quantile regression loss:
2 (T) (T)
arg min Cox = Eznz (0- 0)] where p+(u) u L(u 0)).
Due to discretization, for each quantile Ty, the minimizer of the 1-Wasserstein distance within that quantile is at 0 that satisfies (0) = [ + Ty) ig (see Lemma 2 of Dabney et al. (2018)). Consequentily, to approximate Z, we can simultaneously optimize for all of the quantiles, {0, 1 by minimizing the following loss:
N Σ
N Σ ==1
(0,)
Banz (et, (=-)0= 0,]]
QR
To stablize optimization, QR-DQN uses modified version of quantile loss called quantile Huber loss which utilizes the Huber loss with hyperparameter k:
(u) 1 1 S k) (Iul 1 こz 1(/u] xt).
The quantile Huber loss is defined as:
(=) = (u) < 0)] (u)
In QR-DQN, both the approximating distribution and the target distribution are discretized with quantiles. The target distribution is computed with bootstrapping through the distributional Bellman operator. Concretely, given target network, furget and transition tuple (s, a, r, 8'), the target distribution is approximated as:
N
1 N
Q(s',a')
0 larget
(s', a')
a arg max Q(s', a')
Tθ = r 0 arget (s', a²).
(13)
The TD target Q-value distribution's quantiles and distribution are:
(O,,,T0,,... ,ON), TZ d Σ (TT,))
N
j=1
Treating the target distribution as an oracle (i.e., not differentiable), we update the current quantile estimates minimize the following quantile regression loss:
N E E p2,(2 ()())a))
N
N
Σ
Σ
Ai (sq
[T0] as)
(14)
sg denotes the stop gradient operation as we do not optimize the target network (standard DQN practice).
20

---

Under review as a conference paper at ICLR 2023
C ALGORITHMIC DETAILS
In this section, we describe EDE in detail and highlight its differences from previous approaches. In Algorithm 1, we describe the main training loop of EDE. For simplicity, the pseudocode omits the initial data collection steps during which the agent uses uniformly random policy to collect data and does not update the parameters. We also write the parallel experience collection as for loop for easier illustration, but in the actual implementation, all loops over K are parallelized. The main algorithmic loop largely resembles the standard DQN training loop. Most algorithmic difference happen in how exploration is being performance i.e., how the action is chosen and how the ensemble members are updated.
In Algorithm 2, we describe how EDE chooses exploratory actions during training. For each state, we first extract the features with the shared feature extractor f and then compute the state-action value distributions for each ensemble head. With the state-action value distributions, compute the estimate for the epistemic uncertainty as well as the expected Q-values. Using these quantities, choose the action based on either Thompson sampling or UCB.
In Algorithm 3, we describe how EDE updates the parameters of the feature extractor and ensemble heads. This is where our algorithm deviates significantly from prior approaches. For each ensemble head, we sample an independent minibatch from the replay buffer and compute the loss like deep ensemble. Inside the loop, the loss is computed if each ensemble member QR-DQN. How- ever, doing so naïvely means that the feature extractor f will be updated at a much faster rate. To ensure that all paramters are updated the same rate, for each update, we randomly choose a sin- gle ensemble member that will be responsible for updating the feature extractor and the gradient of other ensemble heads are not propagated to the feature extractor. This also saves some compute. In practice, this turns out to be crucial for performing good uncertainty estimation and generalization. We show the effect of doing so in Figure 13c.
Attentive readers would notice EDE is more expensive than other methods since it needs to process M minibatches for each update. our experiments, the speed of EDE with 5 ensemble member approximately 2.4 times slower than the base QR-DQN (see Appendix F). On the other hand, we empirically observed that the deep ensemble is crucial for performing accurate uncertainty estima- tion in CMDPs. UA-DQN (Clements et al., 2019), which uses MAP sampling, performs much worse than EDE, even when both use the same number of ensemble members e.g., 3 (see Figure 5). this work, do not investigate other altrenatives of uncertainty estimation but believe future works in improving uncertainty estimation in deep learning could significantly reduce the computation overhead of EDE.
Algorithm EDE
Initialize feature extractor f and ensemble heads M
M target M ===1
2: Initialize target feature extractor flagget and ensemble heads {g,
3: Initialize replay buffer Buffer
4: {8(8), done (k) KK ]k=1 reset all environment
5: for from to do
A Standard DQN training loop
6: for k from 1 to K do
7: a + CHOOSEAATTIO((!!! f. (o.
8:
s' r,, done (A) +- Take in the environment add (s(k) a, (1)) to Buffer
9:
10:
(kk) s'
if done (k) then
s s(k) d(n), &- reset the 1 environment
12:
13:
end end for UPDATE(Buffer, J, (g fuarget {g,
14:
sarget M
Ji=1
16:
if t mod update apdate..reeqencc = then
17:
faagett {g, jarget M
J==1 f f, 1M
18: end if
19:
end for
21

---

Under review as a conference paper at ICLR 2023
Algorithm CHOOSEACTION( /, (9i)
1: feature f- f(s)
Only one forward pass on the feature extractor
2: for from to M do
3: {0ir(s, a)) j=1
9 (feature)
Compute individual ensemble prediction
4: end for
5: Compute ap (8, a) using Equation
Estimate epistemic uncertainty with ensemble Explore with estimated uncertainty
6: Compute a' using Equation 6 or Equation 4
7: return a.
Algorithm 3 UPDATE(Buffer, to flarge {9; mrgeet AM i==l
L
2: indd_ ~ (1,
D Randomly select ensemble member
3: for from to M
Train each ensemble member independently
4: {s,a,r, done) Sample from Buffer
5: feature(s') f- (
6: (0 (8', 1N +- 9,
targer
(feature(s"))
7: Compute ,
TO.N) using Equation 13 for the 7th ensemble member
8: feature(s) f- [(s)
9:
if & grad_index then
Prevent over-training the feature extractor
10:
feature(s) + STOPGRADIENT (feature(s))) end if
11:
12:
(s, a)) + g, (feature(s))
13:
L. <- compute QR loss using Equation 14 for the 5th ensemble member
C
15:
end for
16: Compute gradient of w.r.t the parameters of f,
M JM
17: Update the parameters with Adam
22

---

Under review as a conference paper at ICLR 2023
D PROCGEN RESULTS
In Table 3 and Table 2, we show the mean and standard deviation of the final unnormalized train and test scores (at the end of 25M environment steps) of EDE along side with other methods in the literature. We adapt the tables from Raileanu & Fergus (2021). In Figure 10, we show the curves of min-max normalized We see that for a subset of games, EDE is significantly more sample-efficien For other games, EDE either on par with the policy optimization methods fails train, indicating there are other issues behind its poor performance on these games. Note that IDAAC results tune the hyperparameters for individual games whereas we only tune the hyperparameters on bi bigfish.
score that we reproduced using the code from Raileanu & Fergus (2021), I).
Game
PPO
MixRey
PLR
UCB-DRAC
PPG
DAAC
IDAAC
EDe
bigtish
3.7
7.1
10.9
18.5
22.1
SurPilot
24.9 1.0
32.4 = 1.5
27.9 : 4.4
30.0 1.3
47.2 2 1.6
36.4 2.8
37.0222
49.6 2.4
PriiBot
26.2 1.2
27.3 = 0.8
28.0
1.4
27.6
27.8 0.6
28.6
27.9
25.7 +
Ninja
0.5
7.240.
0.6
6.6
6.4
6.4
6.1
Plunder
6.2 5.2 0.6
0.5
8.7
8.3
14.3
20.7 333
2.4
4.9 0.6
Cavellyer
5.1 =
6.1
6.3 0.5
5.0
7.0
4.8
5.0
7.9 0.4
CoinRun
6.6
9.4
0.5
Jumper
0.5
0.3
0.1
6.6
6.3
5.7 0.3
Chaser
3.5 0.9
5.8
6.9
0.6
0.5
6.6
68
1.6
Climber
5.6 0.5
0.7
6.3
0.6
2.4
T.2
8.3
3.7 1.1
Dodgebull
2.6
2.4
2.9
3.20A
2.8
3.2
0.2
1.5 0.3
Lenger
4.9 4.2
1.3
6.8
4.8
1.0
1.3
7.7 1.0
6.4 0.3
Mase
0.3
0.2
5.5 0.8
0.1
0.1
0.2
5.6 0.3
3.4 0.5
Miner
9.200.6
levels after training 25M environment steps. The mean and standard deviation are computed using 5 runs with different seeds.
Table 2: Procgen
Game
PPO
MisRez
PLR
UCB DRAC
PPG
DAAC
IDAAC
EDE
bigtish
+ 2.7
1.3
19.9
20.1
StacPloo
29.0 d 1.1
28.7 + 1.1
2.6
0.3
33.1 1.3
49.6
2.1
38.0
2.6
38.6
2.2
46.9 + U.T
FruaRat
28.8 = 0.6
29.9 0.5
15.9 1.3
0.5
31.1 0.5
31.1
29.7 0.4
29.1
0.7
27.7 + 1.0
BassFight
0.4
0.7
0.4
11.1
10.0
Niga
0.2
Plunder
6.1 + 0.8
6.2
4.1 4.1.11 1.3
10.221.70
16.4 1.9
22.5
24.6 1.6
$.2 0.8
CavePlyer
6.7 + D.6
6.2 0.7
6.4
0.9
9.6
5.8 0.4
6.2
10.6 0.1
Coinkun
5.4
9.4
0.4
Janger
-
8.6
0.6
8.2
0.7
8.6
8.7
9.0
Chaser
4.1 0.3
3.4
6.3
7.0
10.7
6.212
7.8
2.2 0.1
Climber
6.9
7.5
0.2
0.6
1..2
1..
10.2
10.0 - 0.3
Dodgshail
7.3
5.6
0.2
0.3
7.1
1.2044
0.6
7.4
0.2
4.5
7.2
Leaper
5.5 0.4
3.2
6.4
5.0
9.3
8.0
6.3
9.0 0.3
Maze
9.1
8.7
0.1
8.3
9.0
0.4
5.7 11.6
6.6
1.3
11.3
11.
1.1
Table 3: Procgen scores train levels after training 25M environment steps. The mean and standard deviation are computed using 5 runs with different seeds.
23

---

Under review as a conference paper at ICLR 2023
bigfish
bossfight
caveflyer
chaser
crainn
66-
Gricn
201
PPV
MAAAA Mo
0.2-
33
0.1
41
33
0.0
05
22
0.5
15
1.5
0.0
2.0
climber
coinrun
dodgeball
fruitbot
2.0
0.7
0.0
0.4
64
1.
1:2
heist
Jumper
leaper
maze
EDC
COC
0.5
COC
3.2
0.0
0.4
0.4
-0.6
Merull VW Map
800
00-
miner
ninja
plunder
starpilot
EDC
202
0.35
0.5
626
3.20
3.85
322
01-
03
ortin
110
0.
05
3.0
25
6.0
2.5
3.5
2.5
2.0
0.0
0.5
3.5
2.5
0.0
05
2.0
23
Figure 10: Min-max normalized test performance of few representative methods on individual Procgen games. Mean and standard deviation are computed over 5 seeds. can see for games that value-based methods can make meaningful progress, EDE almost improve up QR-DQN. In many cases, EDE significantly more sample efficient compared to the other methods. There are notably games (chaser, maze, heist, miner) where value-based methods (including EDE) still perform much worse than policy optimization. These games all require generalization with long horizon planning which appear to remain challenging for value-based methods, Getting value-basse methods to generalize on these games would be an interesting future direction.
E ADDITIONAL ANALYSIS
E. OTHER SOURCES OF POOR GENERALIZATION
While our method improves train and test performance on two different CMDP benchmarks, it still does not fully close the generalization gap. In this section, we study other potential reasons for the poor generalization of RI. algorithms, particularly value-based ones.
Optimization vs. Generalization. In supervised learning, deep neural networks have good gener- alization, meaning that they perform similarly well on both train and test data. However, RL, the main bottleneck remains optimization rather than generalization. For example in Procgen, the state- of-the-art art training performance is far below the theoretical optimum (Cobbe et al., 2019; Raileanu & Fergus, 2021), suggesting that current RL algorithms do not yet "overfit" the training distribu- tion. To verify this hypothesis, we first conduct meta-analysis existing algorithms. We choose representative algorithms (Appendix D) and measure whether better performance on training MDPs
24

---

Under review as a conference paper at ICLR 2023
Task ID ablation
Comelaatio between train test performance
1.0
Singleton Training Performance
0.6
0.8
0.6
0.4
0.4
return
0.2
0.0
-0.2
0.0
-0.4
miner
dudirtial
Cal
bigrinn
sharpling
caveltye
0.0
PPO
DON
-02-
w/o task ID
task ID
0.0
0.5
1.0
1.5
steps
166
steps
1e7
(a) Meta Analysis Figure 11: Investigations of other potential sources of poor generalization in RL.
(b) Singleton Experiments
(c) DQN with Task ID
translates into better performance on test MDPs. More specifically, we compute the Spearman's rank correlation, Tas between the training and test performance of the algorithms, for each Procgen environment (Figure 1la). For many Procgen environments, training performance is strongly corre- lated with test performance, suggesting that further improvements in optimization could also result in better generalization.
In Figure 12a and Figure 12b, we show the scatter plot of the training performace and test perfor- mance of each algorithm on games with positive correlation and games with negative correlation (each color represents a different game). For games with positive correlation, we see that the cor- relation is almost linear with slope of 1 (i.e., almost no generalization gap), which suggests that for these games, we can expect the good generalization if the training performance is good. On the games with negative correlation, it is less clear that improving training performance will necessarily lead to better generalization or more severe overfitting. Most of the games with negative correlation either saturated in terms of performance or too challenging for current methods to make much progress. This suggests that despite their apparent similarity, there are actually two "categories" of games based on the algorithm's 's tendency to overfit. Note that this dichotomy likely depends on the algorithm: here, all algorithms are PPO-based and we hypothesize the picture would be different for value-based methods. It would be interesting to understand what structural properties of these CMDPs are responsible for this difference in overfitting, but this is out of scope for the current work.
Another interesting thing to notice is that most of the outliers in the two plots are PLR (Jiang al., 2021) which is an adversarial method that actively tries to reduce training performance. This may explain why the trend does not hold for PLR as all other methods use a uniform distribution over the training environments.
Sample Efficiency. One hypothesis for the poor performance of value-based methods relative to policy optimization ones in CMDPs is that the former is less sample efficient than the latter in each individual MDP μ E .Mimnin If that is the case, the suboptimality can accumulate across all the training MDPs, resulting in a large gap between these types of methods when trained in all the environments in Musin To test this hypothesis, we run both DQN and PPO on a single environment from each of the Procgen games. As shown in Figure 1 11,, we that there is not a significant gap between the average performances of DQN and PPO on single environments (across 3 seeds for 10 environments on which both methods reach scores above 0.5). In fact, DQN usually slightly more sample-efficient. Thus, we can rule this out as a major factor behind the observed discrepancy.
Partial Observability. Another sensible hypothesis for why DQN underperforms PPO in CMDPs is the partial observability due to not knowing which environment the agent is interacting with at any given time (Ghosh al., 2021). Policy optimization methods that use trajectory-wiset Monte Carlo update are less susceptible to partial observability whereas value- based methods that use temporal difference updates can suffer more since they rely heavily on the Markovian assumption. To test this hypothesis, we provide the task ID to the Q-network in addition to the observation similar to Liu et al. (2021). The access to the task ID means the agent knows exactly which environment it is in (even though the environment itself may still be partially observable like Atari). In Figure 11b, we see that both methods do well on the singleton environments, so with task IDs, the algorithms should in theory be able to do well singleton environments even if there is transfer between different MDPs because the model can learn them separately. In practice, we embed the discrete task IDs into R and add them input the final layers of the Q-network. Since there is semantic
25

---

Under review as a conference paper at ICLR 2023
relationship between discrete task IDs, do not expect this to improve generalization performance, but, surprisingly, we found that it does not improve training performance either (Figure 1 Icc.. This suggests that partial observability may not be the main problem in such environments as far as training is concerned. Note that this issue is related to but not the same as the aforementioned value- overfitting issue. Having task ID means that the agent can have accurate point estimate for each MDP (as far as representation is concerned), but optimization remains challenging without proper exploration.
Value Overfitting. Most deep RL algorithms model value functions as point estimates of the ex- pected return at each state. As discussed in Raileanu & Fergus (2021), this is problematic in CMDPs because the same state can have different values in different environments. Hence, the only way for the values to be accurate is to rely on spurious features which are irrelevant for finding the optimal action. For policy optimization algorithms, the policy can be protected from this type of overfitting by using separate networks to train the policy and value, as proposed in Raileanu & Fergus (2021). However, this approach cannot be applied to value-based methods since the policy is directly defined by the Q-function. An alternative solution is distributional RL (Bellemare et al., 2017) which learns distribution (rather than point estimate) over all possible Q-values for given state. This mod- els the aleatoric uncertainty resulting from the unobserved contexts in CMDPs (i.e., not knowing which MDP the agent is interacting with), thus being able to account for one state having different potential values depending on the MDP. Our method uses distributional RI. to mitigate the problem of value overfitting. As seen in Figure 5, while learning value distribution leads better results than learning only a point estimate (i.e., QR-DQN > DQN), the largest gains are due to using more sophisticated exploration method (i.e. OR-DQN- + UCB + TNN DQN).
Our analysis indicates that sample efficiency and partial observability are not the main reasons be- hind the poor generalization of value- based methods. In contrast, value overfitting is indeed a prob- lem, but it can be alleviated with distributional RL. Nevertheless, effective exploration proves to be a key factor improving the train and test performance of value-based approaches in CMDPs. Our results also suggest that better optimization on the training environment can lead to better general- ization to new environments, making it a promising research direction.
Train vs. return for games with negative correlation
Train vs. retumn
a² vs. oz test performance
1.0
1.0
0.25
0.8
MM M
0.20
0.6-
test
test
0.15
0.4
0.10-
0.2
0.2
0.05
Anmm
0.00
-0.2
0.2
0.0
0.2
0.4
0.6
0.8
1.0
-0.05
train
0.5
1.0
1.5
2.0
Retin
Cuuaer
Purre
107
nulbe:
(c) Comparing different uncer- tainties (mean test min max nor- malized scores).
(a) Games with positive train-test correlation.
(b) Games with negative train-test correlation.
Figure 12: Additional empirical results.
F UNCERTAINTY
Central to our method is the estimation of uncertainty in Q-values. We found that what kind of un- certainty to model and how we model it are crucial to the performance of EDE. The first important consideration is what kind of uncertainty to model. Specifically, we argue that only epistemic uncer- tainty is useful for exploration since they represent the reducible uncertainty. This is a longstanding view shared by many different sub-field of machine learning (Mockus, 1982; Settles, 2009). The aleatoric uncertainty on the other hand is intrinsic to the environment (e.g.. label noise) and cannot be reduced no matter how many samples we gather. In standard regression, even with ensembles, the estimated uncertainty (i.e., variance) contains both epistemic and aleatoric uncertainty. In Atari, the main source of aleatoric uncertainty comes from the sticky action which has relatively simple
26

---

Under review as a conference paper at ICLR 2023
structure. However, in CMDPs such Procgen, the aleatoric uncertainty coming from having differ- ent environments is much more complicated and possibly multimodal, making it important to model them explicitly. Empirically, using alcatoric uncertainty to explore is detrimental to the test per- formance which may partially explain DQN+UCB's extremely poor performance (Figure 5), even compared D (i.e., uncertainty estimation based ensemble of non-distribational DQNs contains both types of uncertainties). To further test this hypothesis, we run QR- BR-DQN+UCB with both aleatoric and epistemic uncertainty and version with only epistemic uncertainty (Fig- ure 12c). The results show that using only epistemic uncertainty performs significantly better than using both uncertainties in terms of both training and test performance.
Some additional important detail about uncertainty estimation with an ensemble are the number of ensembles and how the members of the ensemble are updated. These details have implications for both the computational complexity of inference and the quality of estimated uncertainty. For the main experiments, we use ensemble members which, without specific optimization, results in about 2.5 times more time per algorithm step. The fol lowing are the time complexity of our algorithm on Tesla V100:
EDE (5 ensemble members):
0.3411 1 seconds per algorithm step
EDE (3 ensemble members):
0.2331 seconds per algorithm step
QR-DQN:
0.1356 seconds per algorithm step
In Figure 13b, we found that using ensemble heads achieves comparable performance. On the other hand, using deep ensemble seems to be crucial for good estimation of epistemic uncertainty as UA- DQN (Clements et al., 2019), which uses MAP sampling with 3 ensemble members, performs much worse than QR-DQN with 3 ensemble members trained with different minibatch and initialization (which are presumably more independent from each other). We believe further understanding this phenomenon would be an exciting future research direction.
G ARCHITECTURE AND HYPERPARAMETERS
G.1 PROCGEN
Architecture. We use standard ResNet (He al., 2016) feature extractor (f) from IMPALA (Es- peholt et al., 2018), which was adopted by Ehrenberg et al. (2022) for value-based methods. The feature maps are flattened (d 2048) and processed by M copies of 2-layer MLPs with the same architecture to output M x |A| x N values (for other models, this shape is changed accordingly for the appropriate output dimension). |A| = 15 for all Proegen games. Each MLP conisists of liner layers (both of dimension 512) separated by a ReLU activation; the first layer is g Rdd R5522 and the second layer is g2 51 + K The overall archetecture is
g o f (2) ReLU o ()) Flatten o f.
The hyperparameters for the main method is shown in Table 4, and are directly adapted from Ehren- berg et al. (2022) other than those of new algorithmic components.
For hyperparameters specific to EDE, search over the following values the game bigfish
p (10, 20, 30,50) . A: (0.5, 0.6, 0.7} a [6, 7, 8,9)
and picked the best combination.
Points of comparisons. We try to keep the hyperaparameters the same as much as possible and only make necessary changes to keep the comparison fair:
For base DQN and QR-DQN, we use e-greedy exploration with decaying e. Specifically, at every algorithmic step, the epsilon is computed with the following equation:
(-t- (-t-2000) 2000
(t)) = min 0.1 + 0.9
8000
27

---

Under review as a conference paper at ICLR 2023
Name
Value
number of training envs proegen distribution mode minibatch size
200
casy
512
number of parallel workers (K)
64
replay buffer size
10
discount factor (+)
0.99
n-step return
3
frame stack
1
dueling network
target network update frequency initial random experience collection steps
32000 (algo step)
2000 x K (env step)
total number of steps update algo step env step per algo step
25 x 10 (env step)
1
K
Adam learning rate Adam epsilon
2.5 10
1.5 x 10-4
prioritized experience replay
no
number of quantiles (N)
200
number of ensemble heads (M)
5
Huber loss
1
gradient clip norm
9 (UCB)
30
À (TEE)
0.6
& (TEE)
7
Table 4: Hyperparameters for the main Procgen experiments. One algorithm step (algo step) can have multiple environment steps (env step) due to the distributed experience collection.
For Bootstrapped DQN (Osband & Van Roy, 2015), we use bootstrapping to trained the model and uniformly sample a value head for each actor at the beginning of each episode that used for the entire episode. We use ensemble members just like the other methods.
For UCB,
the estimated epistemic uncertainty when combined with QR-DQN and use the standard variance when combined with DQN (Chen et al., 2017 and use 30.
For TEE without UCB, we use s-greedy exploration without decaying for both QR-DQN and DQN, and EDE with the same À and
For ez-greedy, we use the same & decay schedule in combination with n = 10000 and μ 2.0 which are the best performing hyperparameter values from Dabney et (2021). Each parallel worker has its own ez-greedy exploration process.
For NoisyNet, we modify the base fully-connected layer with the noisy implementa- tion (Fortunato al., 2017) with 0.5.
G.2 CRAFTER
For Crafter, we use the default architecture and hyperparameters that are used for the Rainbow experiments in Hafner (2022) with the following changes:
Change the distributional RL component from C51 to QR-DQN with appropriate architec- ture change and use EDE for exploration
Change the minibatch size from 32 64 speed up training
* Remove prioritized experience replay
Once again, we emphasize that everything else stays the same, so the performance gain can be largely attributed to the proposed method. For EDE, since the default configuration from Hafner (2022) is single thread, we use Thompson sampling instead of UCB and do not use TEE.
For 42, searched over (0.5, 1.0, 2.0) and picked the best value.
28

---

Under review as a conference paper at ICLR 2023
Name
Value
minibatch size
64
number of parallel workers (K)
replay buffer size
10
discount factor (7)
0.99
n-step return
3
frame stack
4
dueling network
no
target network update frequency
8000 (algo step)
initial random experience collection steps
20000 (env steps)
total number of steps
10 (env steps)
algo step per update env step per algo step
4
1
Adam learning rate Adam epsilon
6,25 x 10
1.5 10
prioritized experience replay
no
number of quantiles (N)
200
number of ensemble heads (M)
5
Huber loss k gradient clip norm
1
10
* (Thompson sampling)
0.5
Table 5: Hyperparameters for the final Crafter experiments.
H SENSITIVITY STUDY
In this section, we study the sensitivity of the test performance (i.e., aggregated min-max normalized scores) various hyperparameters on subset of games. First, without TEE, we study the sensitivity to different (UCB exploration coefficients). different M (number of ensemble members), and whether the feature extractor trained with gradients from each ensemble member.
For to and M, we run the algorithm on: bossfight, ninja, plunder, starpi lot, caveflyer, jumper, bigfish, leaper for I seed and report the aggregated min-max normalized scores in Figure 13a and Figure 13b. We observe that the algorithm is slightly sensitive to the choice of yp, but not sensitive at all to M. This encouraging since a large amount compu- tational overhead comes from the ensemble. Note that while smaller yp performs the best here, use a larger value (4 = 30) combination with TEE.
In Figure 13c, show the effect of only training the feature extractor using the gradient from one member the ensemble at every iteration. The results are computed on: ninja, plunder, jumper, caveflyer, bigfishh training the feature extractor leads to lower performance, corroborating our intuition that the feature extractor should be trained at the same speed the individual ensemble members.
leaper, cl immer for seed. We observe that always
In Figure 14, we study the performance under different hyperparameter values TEE. We use fixed M 5 and yp 30 and vary the values of either a À while holding the other one fixed. We observe no significant difference across these, suggesting that the algorithm is robust to the values of a and A.
29

---

Under review as a conference paper at ICLR 2023
Varying
Varying
Always
0.45
0.40
0.30
0.40
0.35
0.35
0.25
0.30
0.30
0.25
0.25
return
0.20
0.20
0.15
0.15
0.10
0.10
=
30
0.10
M
True
0.05
= 50
0.05
M =
False
0.00
0.5
1.5
2.0
0.0
1.0
1.5
2.0
0.5
1.5 steps
2.0
steps
1e7
steps
1e7
1e7
(a) ablation.
(b) M ablation.
(c) Feature extractor ablation.
Figure 13: Aggregated min-mat normalized test scores for (for fixed M 3, and training feature extractor for all value heads), M (for fixed 50 and training feature extractor for all value heads), and whether feature extractor trained with all value head (for fixed yp 50 and M 3).
Varying
Varying
0.5
0.4
reture
0.3
0.2
0.2
0.1
A = 0.5
A
0.0
0.7
0.0
0.5
1.0
1.5 steps
2.0
0.0
0.5
1.0
1.5
2.0
steps
1e7
Figure 14:
Aggregated min-max normalized test scores for (for fixed 7) and a (for fixed À = 0.6) on games: bossf ight, climber, paseerlpuunne plunder, starpilet.
I
CASE STUDY: SIMPLIFIED BIGFISH
To further illustrate the intuition that exploration on the training MDPs can help zero-shot general- ization to other MDPs, will use slightly simplified example inspired by the game bigfist fish (Figure 15) from Procgen. As illustrated in Figure 16, in bigfish the agent (green circle) needs eat smaller fish (small red circle) and avoid bigger fish (large red circle). The images with solid borders are the observed frames and the images with dashed borders are the transitions.
If the agent collides with the small fish, it eats the small fish and gains +1 reward; if the agent collides with the big fish, it dies and the episode ends. The blue rectangle (top) represents the training environment and the red rectangle (bottom) represents the test environment. In the training MDP, the agent always starts in the state shown in the T = 0 of the greedy trajectory (top row). Using random exploration strategy (e.g., e-greedy), the agent should be able quickly identify that going to the top will be able to achieve the +1 reward. However, there is an alternative trajectory (middle) where the agent can go down and right to eat the small fish on the right (orange border). This trajectory is longer (therefore harder to achieve via uniformly random exploration than the first one) and has lower discounted return.
From the perspective of solving the training MDP, the trajectory is suboptimal. If this trajectory is sufficiently hard to sample, the agent with naive exploration will most likely keep repeating the greedy trajectory. On the other hand, once the agent has sampled the greedy trajectory sufficiently, uncertainty-awe aware exploration (e.g., UCB or count-based exploration) will more likely sample these rarer trajectories since they have been visited less and thus have higher uncertainty. This has
30

---

Under review as a conference paper at ICLR 2023
Figure 15: Example frames bigfish. The goal for the agent (green fish) avoid fish bigger than itself and eat smaller fish. The spawning of enemy fish is different across different MDPs and the background can change.
reward
Training
Greedy
Trajectory
rawrd
Explorative Trajoctory
Test MDP
Test Trajectory
Figure 16: Simplified version of bigfish.
impact the performance in the training MDP because the agent has learned the optimal policy regardless of which exploration strategy is used.2
However, the test MDP, the initial state is shown in T' of the test trajectory (bottom row). There is no guarantee that the agent knows how to behave correctly in this state because it has not seen it during training. One could hope that the neural network has learned a good representation so the agent knows what to do, but this does not necessarily have to be the case the objective only cares about solving the training MDP and does not explicitly encourage learning representations that help generalization. Suppose that the agent keeps moving up due to randomness in the initialization of the weight, it will run into the big fish and die (bottom row, = 2). Notice that even though the two environments are distinct, the test trajectory (bottom row) and the explorative trajectory (middle row) share the same state at
1 (pink arrow). Due to this structure, agents that have learned about the explorative trajectory during training time will be able to recognize that the better action from = 1 to move to the right which will avoid death and ultimately result in +1 reward.
Once again, emphasize that this is simplified example for illustrative purpose. In reality, the sequences of states do not need to be exactly the same and the neural network can learn to map
Using function approximators with insufficient capacity may induce effect the optimal policy.
31

---

Under review as a conference paper at ICLR 2023
similar states similar representations. In principle, the neural network can also learn avoid the big fish but at T = 0 of the training MDP, the behavior "moving up" is indistinguishable from the behavior "avoiding the big fish". Good exploration during training can lead the agent run into the big fish from all directions which is much closer to "avoid the big fish", the intuitive generalizable behavior. Clearly, not all CMDPs will have this structure which allows via exploration, but it should not hurt. Like most problems in deep RL, it is difficult to prove this hypothesis analytically but empirically it is indeed the case that EDE has better or the same performance as QR-DQN in 11 out of 16 Procgen environments and perform approximately the same in the remaining environments (Figure 10).
improve generalization
This perspective allows us to tackle the problem with exploration that is in principle orthogonal to representation learning; however, since exploration naturally collects more diverse data, it may also be beneficial for representation learning as side effect.
32


Reviews:
[{"review_id": "y4bM7lElJ9", "reviewer": "Reviewer_ziZP", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper argues that effective exploration strategies can improve generalization in contextual MDPs (CMDPs). The paper proposes a method called exploration via distributional ensemble (EDE), which uses a deep ensemble to (1) disentangle epistemic uncertainty from aleatoric uncertainty and (2) encourage exploration by trying to reduce the epistemic uncertainty. EDE is shown to achieve SOTA for value-based methods in two popular benchmarks.", "summary_of_the_review": "As mentioned before, this paper makes two separate contributions: (1) claiming that generalization can be improved with better exploration, and (2) proposing an exploration method. I don’t think it is fair to asses the paper as a whole since the contributions are somewhat orthogonal. I believe that the first contribution, although certainly interesting, is insufficient for acceptance in its current form. The second contribution does not seem very novel since the method is just a combination of previous solutions. Moreover, in order to determine its value, the method should be compared against other exploration methods.  \n\nMy suggestion to the authors is to focus on one of the two. Either try to develop the first idea, providing further analysis and experiments (using different exploration strategies) to show that bad exploration is indeed the cause for poor generalization, or work on the exploration method and show that it can outperform other exploration strategies. This last direction is probably riskier since there are many works that focus on this problem.", "strength_and_weaknesses": "### Strengths\n\nI really like the idea of casting the generalization problem as an exploration problem. I agree that the standard RL formulation may be too task-centered, which makes it difficult for the agent to learn about the environment. The paper shows that this narrow focus may negatively affect generalization.\n\n### Weaknesses\n\nI am quite confused about what the true objective of this paper really is. The paper starts by claiming that any effective exploration strategy can improve generalization. Then, it proposes a method that targets exploration. I am not certain whether the proposed method is somehow addressing generalization directly (didn’t see any evidence for this in the paper) or if it is just another exploration method. If the latter, then the paper makes two separate contributions: \n\n (1) The hypothesis that exploration can improve generalization in CMDPs, which is interesting and seems to be partly supported by the experiments. \n\n (2) A method for encouraging exploration that does not seem very novel since it is based on existing ideas, and whose performance should be compared with other exploration methods.\n\n\n### Questions/Suggestions\n\nThe abstract and introduction state that policy-based methods perform better than value-based methods in CMDPs. The authors identify exploration as the main reason for value-based methods not being able to generalize well in CMDPs. Is this suggesting that the reason policy-based methods outperform value-based methods is mainly that the former are more effective at exploring?\n\nIn order to be able to properly assess the novelty of the method, Section 4 should list only the ideas that are original and move those that belong to previous work to the background section. From what I understand, only equalized exploration is original here. I believe this idea is interesting and could be highlighted and further evaluated in the experiments. \n\nHave you investigated the reasons why EDE performs better in terms of median/IQM but falls behind IDAAC when looking at the mean? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "lNdSJkDgq15", "reviewer": "Reviewer_WM6p", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes a value-based RL architecture that is aimed at achieving good generalization, in the settings of \"contextualized\" MDPs (where the agent is trained on one set of MDPs and tested on a different set sampled from the same distribution). This is achieved by encouraging a more explorative behavior during training. The architecture itself builds on and combines concepts from distributional RL, ensemble training, and classic exploration approaches (UCB-like and Thompson sampling). The method is evaluated on two benchmarks of procedurally generated RL environments. ", "summary_of_the_review": "The architecture proposed in the paper seems effective and the results are overall promising. However at its current form the paper is missing a coherent (and supported) main message. It offers some general high-level intuitions about the importance of exploration, but how and if these are actually related to the proposed performance of the algorithm is unclear, and probably cannot be deduced based on performance metrics alone (without considering the behavior, and/or some more direct measure for its \"explorativitiy\"). Taken together with the technical issues, the paper seems not ready for publication at its current stage.\n\n\n**EDIT**\nThe revised manuscript is an improvement, and the authors have addressed most of the technical issues raised in my review.\nI've changed my recommendation to reflect these improvements.\n", "strength_and_weaknesses": "The paper tackles a challenging benchmark and has good evaluation methodology, including an ablation study for the different components of the architecture, and comparison with strong baselines. The empirical results are strong, or at least competitive with the state-of-the-art on an aggregate level. \n\nOn the other hand, the main conceptual message of the paper seems confused, and some claims are not well supported by neither theory nor empirical evaluation (see main issues below). Moreover the large heterogeneity of performance on individual task basis (Appendix D) makes the empirical results much less conclusive.\n\nMain issues:\n\n* The distinction between \"Value-based\" and  \"policy optimization\" (I will refer to these as PG, for policy gradient) which is emphasized repeatedly in the paper seems a bit arbitrary. First, most (if not all) PG methods discussed and compared to are some kind of an actor-critic (\"hybrid\") architecture which does, in fact, involve Value estimation/learning.\n\n* The tabular example does not seem to support the later claims at all. \n\t* First, the claim that (in Procgen) \"the agent can often end up in states that  are suboptimal for the training MDPs.\" is far from being trivial and it requires at least a convincing demonstration by examples. It seems that the main issue in Procgen is generalization in terms of good \"semantic\" feature extraction for the state structure (identifying platforms from background, enemies, etc.) rather than learning to adapt to a \"new task\" in the same environment (e.g., different target or initial state).\n\t* The exploration of the PG agent is by no means more sophisticated than that of the $\\epsilon$-greedy agent, both rely on purely random (and uniform) exploration (moreover, the entropy regularization have no real effect here, as is seen clearly in Fig. 6c). The main reason limiting e-greedy compared to PG here is most likely the reliance on 1-step updates versus trajectory-based updates. In PG, after each rewarding episode the policy for states along the entire trajectory will be updated. By contrast, the simple Q-learning has to \"propagate\" values back from the rewarding state. But, and related to the previous point, this is not an essential limitation of *value*-based (for example, TD($\\lambda$) will not suffer from this). \n\n* Some strong claims are made about exploration and its components but the supporting evidence is too weak. Importantly, this kind of claims are very hard to support based on some aggerated performance measure alone, as they are really claims about the behavior and learning of the agents (in specific tasks). Most conclusions about why Alg1 > Alg2 are not much more than interpretations, which may or may not be a key reason.\n  Particular examples: \n\t- The mentioned \"importance of Distributional RL\" (which seems not inline with NoisyNet performance, for example)\n\t- The ez-greedy \"not accounting for what the model already knows\" as opposed to the proposed architecture.\n\t- Even the fact that the learning is successful in reducing the epistemic uncertainty is not exemplified.\n More generally speaking, the fact that \"exploration as a key factor for good sample efficiency and generalization\" is almost a triviality. And surely it doesn't provide an answer for the question \"why value-based approaches work well in singleton MDPs but not in contextual MDPs\". (see previous points as well).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "dPUroDdGj-", "reviewer": "Reviewer_iatu", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a value-based RL method (termed EDE) to improve generalization in procedurally generated environments such as Procgen. In particular, the authors introduce two novel modifications to QR-DQN: UCB exploration with learned epistemic uncertainty (UCB) and temporally equalized exploration (TEE), which greatly boosts the zero-shot test performance on Procgen.", "summary_of_the_review": "In summary, I think the paper makes good contributions to the field, but requires some modifications before acceptance.", "strength_and_weaknesses": "**Strength:**\n\n* Solid experiments\n* Significant empirical improvements\n* EDE would serve as a strong baseline for future works and be of importance to the RL generalization community.\n\n**Weaknesses:**\n\nThe connection between \"more effective exploration during training\" and \"better generalization in testing\" is weak. I understand Sec.3 is meant to build this connection. However, since the tabular example is quite far from the realistic setting in Procgen, it is hard to say this is the reason why EDE works well in Procgen. In fact, from Tab.2&3, it seems that the improvements of EDE on testing levels are highly correlated to the improvements on training levels. So a more plausible logic flow is: better exploration -> better performance on training levels -> better zero-shot generalization on testing levels.\n\nDespite the weaknesses, I do think the empirical contributions of this paper are good on their own. So I would suggest the authors reposition the paper from a different perspective.\n\n**Minor issues:**\n\n* In ablation results in Sec.5.1, please add reference to Fig.4.\n* Some statements are not grounded and may mislead readers.\n  * First paragraph in introduction section: \"In this work, we aim to understand why value-based approaches work well in singleton MDPs but not in contextual MDPs\". The statement that value-based approaches does not work well in contextual MDPs is unsupported. The only seeming evidence in the context is that most existing SOTA methods on Procgen are based on PPO. But this is probably because the baseline code released from Procgen authors is PPO.\n  * Second paragraph in introduction section: \"the main goal of exploration is to learn an optimal policy for that environment\". This is not correct. Exploration is also useful in the absence of reward. From the context, I assume the authors want to make a contrast that exploration in singleton env can only benefits that environment while in CMDP exploration in one env may also help other envs. However, I see this more of the difference between singleton env and multiple envs. In CMDP, anything we change in one env will likely affect the behaviour in another env. In short, this is not specific to exploration.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "GpaZN6vma4Q", "reviewer": "Reviewer_R7TS", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper studies the task of generalization in RL tasks defined by contextual MDPs (CMDPs). In contrast to meta-RL tasks, this paper focuses on tasks where the agent should learn a policy that performs well on all MDPs instantiated by a CMDP. The paper argues that exploration is key to improving the performance of RL algorithms in such tasks. Here the goal of exploration is not to find a more rewarding policy for a single task, but to better explore the state space to perform well on test tasks. The proposed EDE algorithm estimates the epistemic uncertainty of the Q function to guide exploration. EDE is evaluated on the Procgen and the Crafter benchmarks.", "summary_of_the_review": "Although the adopted uncertainty estimation techniques are not new, the paper demonstrates that epistemic uncertainty is useful to guide exploration to improve generalization in RL. The effect of uncertainty estimation under a bootstrapped target (Q function learned by Q-learning) is not discussed in the paper, which in my opinion, is a weakness of the paper. ", "strength_and_weaknesses": "Strength:\n\n- Although the idea of separating aleatoric and epistemic uncertainty is not new, this paper shows that this idea can be used to improve generalization in RL environments. \n\n- The experiments conducted in the paper are comprehensive: the diversity of the adopted tasks and baselines is large enough to demonstrate the superiority of the proposed method. \n\nWeaknesses and questions:\n\n- It is unclear to me how the bootstrap learning procedure of the Q function affects the estimation of epistemic uncertainty. The uncertainty estimation process adopted in the paper is well-studied on tasks such as classification and environment model estimation. However, it is unclear to me whether the estimation is still accurate when Q-learning style updates are applied. Therefore, I believe this paper should provide at least some intuition into this. \n\n- Prior work such as SAC and TD3 maintain two value heads to mitigate the value overestimation problem of Q-learning. Since EDE also uses multiple value heads, is it possible that it also implicitly mitigates the value overestimation problem? It would be nice to do an additional ablation study to verify this.\n\n- In the abstract and the introduction, the authors mentioned that there is a gap between policy-based methods and value-based methods. However, the paper does not discuss why policy-based methods work better than value-based methods. And the paper is mainly about the importance of exploration + uncertainty estimation, which can also be applied to policy-based methods. It would be nice to either adjust the story or provide more discussion on why the proposed technique works for value-based methods in general.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject