On the Importance of Exploration for Generalization in Reinforcement Learning
Authors: Yiding Jiang
J Zico Kolter
Roberta Raileanu
Venue: NeurIPS 2023 poster (2023)
Keywords: reinforcement learning
generalization
procgen
crafter
TL;DR: We highlight the importance of exploration for generalization in reinforcement learning and propose an exploration method that achieves sota on procgen.

Abstract: Existing approaches for improving generalization in deep reinforcement learning (RL) have mostly focused on representation learning, neglecting RL-specific aspects such as exploration. We hypothesize that the agent's exploration strategy plays a key role in its ability to generalize to new environments.
Through a series of experiments in a tabular contextual MDP, 
we show that exploration is helpful not only for efficiently finding the optimal policy for the training environments but also for acquiring knowledge that helps decision making in unseen environments. Based on these observations, we propose EDE: Exploration via Distributional Ensemble, a method that encourages the exploration of states with high epistemic uncertainty through an ensemble of Q-value distributions. 
The proposed algorithm is the first value-based approach to achieve strong performance on both Procgen and Crafter, two benchmarks for generalization in RL with high-dimensional observations. The open-sourced implementation can be found at https://github.com/facebookresearch/ede.

On the Importance of Exploration for Generalization in Reinforcement Learning
Yiding Jiang"
Zico Kolter Carnegie Mellon University zkolterOcs. cmu. edu
Roberta Raileanu Meta Al Research rail wwwlleanummtt.ccmm com
Carnegie Mellon University yidingji@cs. cmu. edu
Abstract
Existing approaches for improving generalization deep reinforcement learn- ing (RL) have mostly focused on representation learning, neglecting RL-specific aspects such as exploration. We hypothesize that the agent's exploration stral- egy plays key role in its ability generalize a series of experiments in a tabular contextual MDP, we show that exploration is helpful not only for efficiently finding the optimal policy for the training envi- ronments but also for acquiring knowledge that helps decision making in unseen environments. Based these observations, propose EDE: Exploration via Distributional Ensemble, a method that encourages the exploration of states with high epistemic uncertainty through ensemble of Q-value distributions. The proposed algorithm is the first value-based approach to achieve strong performance both Procgen and Crafter, benchmarks for generalization in RL with high- dimensional observations. The open-sourced implementation can be found at https://gi thu.. com/facebookreserrcchceeoo
environments. Through
Introduction
Current deep reinforcement learning (RL) algorithms struggle to generalize in contextual MDPs (CMDPs) where agents are trained on a number of different environments that share a common structure and tested unseen environments from the same family [22, 112, 75], despite being competitive in singleton Markov decision processes (MDPs) where agents are trained and tested the same environment [73, 43, 8]. This is particularly true for value-based methods [113] (i.e., methods that directly derive a policy from the value functions), where there has been little progress on generalization relative poliey-optimization methods (i.e., methods that learn parameterized policy in addition to a value function) [26]. Most existing approaches for improving generalization in CMDPs have treated this challenge pure representation learning problem, applying regularization techniques which are commonly used in supervised deep learning [29, 21, 45, 61, 116, 60, 92]. 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 RL beyond representation learning.
Exploration can help an agent gather more information about its environment, which can improve its ability generalize new tasks environments.
While this statement seemingly intuitive, a formal and explicit connection has surprisingly not been made outside of more niche sub-areas of RL (e.g., task-agnostic RL [123] meta-RL. RL [65]). has there been a convincing empirical demonstration of this statement on common generalization
Work done while interning at Meta Al Research.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Train
Train
Test
SONTT
Standard
Explorative
Figure 1: Exploration help agents learn about parts of the enviromment which may be useful test time, even they are not needed for the optimal policy on the training environments. Note that this picture only illustrates one of many scenarios where exploration helps.
benchmarks for RL. In this work, we show that the agent's 's exploration strategy is a key factor influencing generalization in contextual MDPs. First, exploration can accelerate training in deep RI, and since neural networks tend to naturally generalize, better exploration can result in better training performance and consequently better test performance. More interestingly, in singleton MDPs, exploration can only benefit decisions in that environment, while in CMDPs exploration in one environment can also help decisions in other, potentially unseen, environments. This is because learning about different parts of the environment can be useful in other MDPs even if it is not useful for the current MDP. As shown in Figure I, trajectories that are suboptimal in certain MDPs may turn out to be optimal in other MDPs from the same family, this knowledge help find the optimal policy more quickly in other MDPs encountered during training, and better generalize to new MDPs without additional training.
One goal of exploration as to reduce epistemic uncertainty. To model epistemic uncertainty (which is reducible by acquiring more data), we need to disentangle from aleatoric uncertainty (which irreducible and stems from the inherent stochasticity of the environment). As first observed by Raileanu and Fergus [90], in CMDPs the same state can have different values depending on the environment, but the agent does not know which environment it is in so it cannot perfectly predict the value of such states. This is type of aleatoric uncertainty that can be modeled by learning a distribution over possible values rather than single point estimate [10]. Based on these observations, we propose Exploration via Distributional Ensemble (EDE), a method that uses an ensemble of Q-value distributions to encourage exploring states with large epistemic uncertainty. We evaluate EDE on both Procgen [22] and Crafter [37], two procedurally generated CMDP benchmarks for generalization in deep RI, demonstrating a significant improvement over more naive exploration strategies. This is the first model-free value-based method to achieve state-of-theear performance on these benchmarks, in terms of both sample efficiency and generalization, surpassing strong policy-optimization baselines and even a model-based one. Crucially, EDE only targets exploration so it can serve as a strong starting point for future works.
learn new information about the (knowable parts of the) environment so
To summarize, our work makes the following contributions: (i) identifies exploration as a key factor for generalization in CMDPs and supports this hypothesis using a didactic example in a tabular CMDP, (ii) proposes exploration method based on minimizing the agent's 's epistemic uncertainty in high-dimensionnl CMDPs, and (ii) achieves sfate-of-the-art performance on two generalization benchmarks for deep RL, Procgen and Crafter.
2 Background
Episodic Reinforcement Learning. A Markov decision process (MDP) is defined by the tuple M = A, R, P.y. H1), where S is the state space, A is the action space, R : S x AA- -> [Rmin Runax] is the reward function P S A x S + Rg0 is the transition distribution, discount factor, and μ : S + R2 is the initial state distribution. We further denote the trajectory of episode to be the sequence T is the length of the trajectory which can be infinite. If a trajectory is generated by a probabílistic policy S x A + R20. Z= = "TT 7' is a random variable that describes the discounted return the policy achieves. The objective is find nh" that maximizes the expected discounted return, x' arg max, & (yppp=()) [2*], where (r) p(()) Ie= P(s | 81, (() a.)- For simplicity, we will use E, instead of Er-p)() ) to denote the expectation over trajectories sampled
E (0, 1] is the
PT, (Su, ag ro..... ST, aT,rT, ,T, ST+1) where rt R(s:, a.) and
2

---

from the policy л. With slight abuse of notation, we use (s, a) to denote the conditional discounted return when starting at s and taking action (i.e., Sq and ap = a). Finally, without loss of generality, we assume all measures are discrete and their values lie within [0. 1].
Value-based methods [113] rely tion, also referred to as the O-fumction, Q= a) = 1E, | 80 8, ao a). The Q-function of policy can be found at the fixed point of the Bellman operator, T [11], T=Q(s,
fundamental quantity in RI., the state-action value func-
aaaaaaa'-"""""" a0) [Rs8,00 a) + +Q((s, (o(ss'aa)) a')]. The value function for state 8 is defined as (*) = Baan(-1l)) [Q (1
a)]- Bellemare et al. [10] extends the procedure to the distribution
of discounted returns, T Z(s, a) - R(s,a) + ( a'), s' P'- s, and a' ~ -((-) |s),
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 D.I). For balue-based methods, the greedy policy is directly derived from the Q-function a(a 8) = Larg max, Q(a'.s)(a).
Generalization in Contextual MDPs. A contextual Markov decision process (CMDP) [39] is special class of partially observable Markov decision process (POMDP) consisting of differ- ent MDPs, that share state and action spaces but have different R, P. and μ.. In addition to the standard assumptions of CMDP, assume the existence of structured distribution V over the MDPs. During training, we are given a (finite or infinite) number of training MDPs, drawn from Vi (Ghosh et al. [35] refers to this setting as epistemic POMDP). For each M,, we use QM a.) and VM (s) to denote the Q-function and value function of that M. We use (T) denote the trajectory distribution of rolling out in M. The objective is to find a single policy x that maximizes the expected discounted return over the entire distribu- tion of MDPs, (...p..))....... Σ 7'r1] without retraining on the unseen M (i.e., zero-shot generalization). We will refer this quantity as the test return?.
Mimin = MM,...,......
Since in RL the algorithm collects its own data, the appropriate notion of generalization is the performance learning algorithm can achieve given a finite number of interactions. Crucially, if algorithm achieves a better test return than algorithm 2 given the same number of interactions with Musin. we can say that algorithm 1 generalizes better than algorithm 2. Furthermore, we assume the existence of π* (potentially more than one) that is a-optimal for all M E supp( ( and all so in s) i.e., V (sp) 2 max, V (sp) - a where a is small. This assumption only ensures that zero-shot generalization is intractable [69] but it usually not used explicitly in algorithm design. If the number of training environments is infinite, the challenge is learning good policies for all of them in sample-efficient manner, i.e., optimization; if is finite, the challenge also generalization to unseen environments.
3 Generalization in a Tabular CMDP
Much of the literature treats generalization in deep RL as pure representation Icarning problem [100, 119, 120, 1] and aims to improve by using regularization [29, 118, 21, 45] or data augmentation [21, 61, 116, 56, 60, 92, 1 112]. In this section, we will first show that the problem of generalization in RL extends beyond representation learning by considering tabular CMDP, which does not require any representation learning. The goal is to provide intuition on the role of exploration for generalization in 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 a generic 5 x 5 grid environment (Figure 2a). During training, the agent always starts at a fixed initial state, (z 0, y = 0), and can move in 4 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 a large positive reward, 2 and the episode ends. Otherwise, the agent receives a small negative
"It may be useful to contrast this with the notion of generalization gap, the difference between training and test returns. Without good training performance, the generalization gap is not a meaningful measure: a random policy would have no generalization gap at all. The generalization gap is important for supervised learning because we get nearly perfect training performance, but this is not the case for RL.
3

---

Reward
a')), where is the O-leaming E-ggreedy
Q-learming UCB
4.0
4.0
+2.04
-2.04
.0.04
-0.04
3.5-
3.0
3.0
-0.04
0..0
-3.04 -0.00
2.5-
Subopttnaalty
2.5-
2.0-
M
2.0
1.5-
1.5-
-0.04
-0.04
.0.04
-3.04
-0.0
1.0
1.0
0.5-
0.
-2.06
-0.04
0.04
-0.04
-0.04
-0.5
200
400
800
-0.
200
400
episode
episode
(a) Tabular CMDP
(b) e-greedy
(c) UCB
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 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-learning with e-greedy exploration, (c) Q-learning with UCB exploration.
reward, r = -0.00000 0.04. At test time, the agent starts at different location, (a = 0, y = 4). In other words, the train and test MDPs only differ by their initial state distribution. In addition, each 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 two classes of algorithms with different exploration strategies: (1) Q-learning with e-greedy (Greedy, Watkins and Dayan [113), (2) Q-learning with UCB (UCB, Auer et al. [6]). To avoid any confounding effects of function approximation, we use tabular policy parameterization for Q-values. Both Greedy and UCB use the same base Q-learning algorithm [1 13]. Greedy explores with s-greedy strategy which takes random action with probability and the best action according to the Q-function arg max, Q(s, a) with probability 1 1- E. E. In contrast, UCB is uncertainty-driver so it explores according to T s) = (a = arg max, Q(s, a') + ey (l/oo(t)/N(ss, total number of timesteps, is the exploration coefficient, and N a) is the number of times the agent has taken action a in state 8, with ties broken randomly². While Greedy is a naive but widely used exploration strategy [105], UCB is effective algorithm designed for multi-armed bandits [6, 5). Chen et al. [19] showed that uncertainty-basee exploration bonus also performs well in challenging RL environments. See Appendix C for more details about the experimental setup.
Each method's exploration strategy is controlled by a single hyperparamter. For each hyperparameter, 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 and 2c. We measure the performance of each method by its suboptimality, Le., 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 quickly achieve an 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 a negligible variance In Appendix C.4, we provide another set of simulations based on changing dynamics and make similar observations. These experiments show that more effective exploration of the training environments can result in better generalization to new environments.
is natural to ask whether this example is relevant for realistic deep RL problems. Analytically, we can motivate the importance of exploration for generalization using sub-MDPs [51] and the sample complexity of Q-learning [63]. Due to space constraints, we develop this argument in Appendix A. Conceptually, the tabular CMDP with two initial state distributions captures common phenomenon in more challenging CMDPs like Procgen, namely that at test time, the agent can often end up in
"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 the didactic purpose. Azar et al. [7] presents a formal but much more complicated extension of the UCB algorithm for value iteration on MDPs. The UCB used is more similar to count-basee exploration bonus.
'While the training and test MDPs are seemingly drawn from different distributions, the setting can be made IID by taking mixture of the train and test MDP. Since there is a significant discrepancy between train and test performance, all the observations would remain valid.

---

To)
THI
TR2 Te3
Training MDP
Oreedy Trajectrry
Exgloratvve Trajectoy
Trjjetory
Figure simplified version of bigfish in Procgen (Figure 6) that captures the spirit of the grid world example from Figure 2a. The green dot represents the agent and the red dots represent the enemies (see Appendix for more details). The goal is to cat smaller dots while avoiding larger dots. An agent that explores both trajectories the training MDP could recover the optimal behavior on the test MDP whereas an agent that only focuses on solving the training MDP would fail.
suboptimal states that are rarely visited by simple exploration on the training MDPs, Having explored such states during training can help the agent recover from suboptimal states at time. See Figure and Appendix B 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. The effect of efficient exploration is, in spirit, similar to that of DAgger [94] - it helps the agent learn how to recover from situations that are rare during training.
In Appendix G. we explore some other interpretations for why exploration improves generalization (which are also related to the sample complexity of Q-learning). In addition, also explore other potential reasons for why current RL methods generalize poorly in Appendix G.2. Note that we do claim exploration be the only way of improving generalization RL different environments may have very different properties that require distinct generalization strategies. In Appendix C.2 and C.3, we study other potential avenues for improving generalization. The benefits of exploration are complementary to these approaches, and they may be combined to further improve performance. In this paper, we focus on using exploration to improve generalization contextual MDPs.
Exploration via Distributional Ensemble
In the previous section, we showed we can improve generalization via exploration in the tabular setting. In this section, would like to extend this idea deep RL with function approximation. While in the tabular setting shown above there is no intrinsic stochasticity, environments can in general be stochastic (e.g., random transitions, epistemic uncertainty reflects a lack of knowledge which can be addressed by collecting more data, while aleatoric uncertainty reflects the intrinsic noise in the data which cannot be reduced regardless of how much data is collected. One goal of exploration is to gather information about states with high epistemic uncertainty [83] since aleatoric uncertainty cannot be reduced, but typical estimates can contain both types of uncertainties [52]. In CMDPs, this is particularly important because large source of aleatoric uncertainty is not knowing which context the agent is in [90].
random unobserved contexts). At high level,
In this section, we introduce Exploration via Distributional Ensemble (EDE), method that encourages the exploration of states with high epistemic uncertainty. Our method builds on several important ideas from prior works, the most important which are deep ensembles and distributional RL. While ensembles are useful way of measuring uncertainty in neural networks [59], such estimates typically contain both epistemic and aleatoric uncertainty. To address this problem, we build on Clements et al. [20] which introduced approach for disentangling the epistemic uncertainty from the aleatoric uncertainty of the learned Q-values.
Uncertainty Estimation. Clements et al. [20] showed that learning the quantiles for QR-DQN [23] can be formulated as Bayesian inference problem, given access posterior p(0) DD,, where the discretized quantiles of (ss, a), and D is the dataset of experience which the quantiles are esti-
5

---

mated. Let j E [V] denote the index for the nth quantile and ( ] denote the uniform distribution over integers between 1 and N. The uncertainty of the Q-value, Q(s, a) = B)-M((,N) [0,(s, a)}. 0)].1. is the relevant quantity that can inform exploration. The overall uncertainty o² = Varo=p(o(p) [Q(s, a)] can be decomposed into epistemic uncertainty d (s, a) and aleatoric uncertainty dae (s, a) such that (8 a. = oégi (s, a) + oie a), where,
(s.(s, a) = (2. [Vrre-p(orp) (0, a)]] (z a) Var,-ur(1.N) (Ee-scePD) 10, (s. a)]] (1)
Ideally, given infinite or sufficiently large amount of diverse experience, one would expect the posterior to concentrate on the true quantile 0*, so Varg=p(0 [0, a)] and consequently (e a) = 0. oae 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 uncertainty in order to learn more about the environment [19]. It should be noted that the majority of the literature on uncertainty estimation focuses on supervised learning; in RL, due to various factors such as bootstrapping, non-stationarity. limited model capacity, and approximate sampling, the uncertainty estimation generally contains errors but empirically even biased epistemic uncertainty is beneficial for exploration. We refer interested readers to Charpentier et al. [17] for more thorough discussion.
Sampling from p(0 | (D) computationally intractable for complex MDPs and function approx- imators such as neural networks. Clements et al. [20] approximates samples from p(0 'D) with randomized MAP sampling [87] which assumes a Gaussian prior over the model parameters. How- ever, the unimodal nature of a Gaussian in parameter space may not have enough diversity for effective uncertainty estimation [31]. Many works have demonstrated that deep ensembles tend to outperform other approximate posterior sampling techniques [59, 31] in supervised learning. Motivated by these observations, we propose to maintain M copies of fully- colleceondeeeed value heads, 9; F + I |R| that share a single feature extractor f -> TRd, similar to Oshand et al. [81]. 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 al. [62] which shows deep ensembles usually perform better than bootstrapping for estimating the uncertainty of Q-values but, unlike EDE, they do not decompose uncertainties.
Concretely, 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 7th head for state s and action (, a) E NN,, where the jth coordinate, and 0,y a.) is the jth quantile of the predieted acttinnttion value distribution for that head. The finite sample estimates of the two uncertainties are:
(8 a) = N M Σ
1
Σ (s, - a)) ( a) = N E (0) - Q(+, a))² (2)
where (8) = गo EM 0,y(s, a) and Q(8 a) = + Σリ (8.
Exploration Policy. There are two natural ways to use this uncertainty. The first one is by using Thompson sampling [110] where the exploration policy is defined by sampling Q-values from the posterior:
(a s) 1 arg max, ((a.a') (a), where {(s, ~ (Q(s,a'), & (s, (3)
E IRzo is a coefficient that controls how the agent balances exploration and exploitation. Alterna- tively, we can use the upper-confidence bound (UCB, Chen et al. [19]):
(a s) = la (a), where a" = arg max Q(s, a') + & (s,
a'
which we found to achieve better results in CMDPs than Thompson sampling [20] on Procgen when use multiple parallel workers to collect experience, especially when combined with the next technique.
Equalized Exploration. Due to 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
't may still contain some aleatoric uncertainty due the heteroskedasticity of the problem.

---

Median
Mean
Optimality Gap
QR-DQN
PPO
IDAAC
EDE (ours)
0.15
0.30 0,45
0.2
0.3
0,4
0.18 0.24 0.30 0.36
0.66 0.72 0.78 0.84
Min-Max Normalized Score
Figure 4: Test performance of different methods on the Procgen benchmark across 5 runs. Our method greatl outperforms all baselines in terms of median and IQM (the more statistically robust metrics) and is competitive with the sate-of-the-an policy optimization methods in terms of mean and optimality gap. The optimality gap equal 1-mean for the evaluation configuration we chose.
propose to use different exploration coefficients for each copy of the model used to collect experience. 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,
Ōepi(s, a'). (5)
a'
À (0, 1) and Rxo are hyperparameters that control the shape of the coefficient distribution. We will refer to this technique as temporally equalized exploration (TEE). TEE is inspired by Horgan et al. [44] which uses different values of for the e-greedy exploration for each actor. In practice, the performances are not sensitive to À and a (see Figure 18 in Appendix G). Both learning and experience collection take place single machine and prioritized experience replay [96] is used since prior work found ineffective in CMDPs [26].
To summarize, our agent explores the environment in order to gather information about states with high epistemic uncertainty which is measured using ensembles and distributional RI. We build on the algorithm proposed in Clements et al. [20] for estimating the epistemic uncertainty, but ensemble instead of MAP [87], and use either UCB or Thompson sampling depending on the task and setting". In addition, for UCB, we propose that each actor uses 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 E). Our ablation experiments show that each design choice important and applying these techniques individually or naïvely combining them performs significantly worse.
deep
5 Experiments
5.1
Procgen
We compare our method with representative baselines on the standard Procgen benchmark (i.e., 25M steps and easy mode) as suggested by Cobbe et al. [21]: (1) QR-DQN which is the prior state-of-the-ar value-based method on Procgen [26]; (2) PPO which a popular policy optimization baseline on which most competitive methods are built; and (3) IDAAC which is state-of-the-art on Procgen and is built on PPO. We tune all hyperparameters of our method on the game bigfish iss only and evaluate all algorithms using the 4 metrics proposed in Agarwal et al. [2]. We run each algorithm game for seeds and report the aggregated min-max normalized scores on the full test distribution, and the estimated nootstrap-estimatee 95% confidence interval in Figure 4 (simulated with the runs). Our approach significantly outperforms the other baselines terms of median and interquartile mean (IQM) (which are the more statistically robust metries according to Agarwal et al. [2]). In particular, it achieves almost times the median score of QR-DQN and more than 2 times the IQM of QR-DQN. In terms of mean and optimality gap, method is competitive with IDAAC and outperforms all the other baselines. To the best of our knowledge, this the first value-based method that achieves such strong performance on Procgen. See Appendices the performance on individual games, hyperparameters, and sensitivity analysis.
every
and for more details about the experiments,
Ablations and Exploration Baselines. We aim to better understand how each component of our algorithm contributes to the final performance by running a number of ablations. In addition, we
'UCB and Thompson sampling have similar regret guarantees but their actual performance can be problem- dependent [16].
7

---

Median
IQM
Mean
DQN
QR-DQN
UA-DQN
Bootstrapped DQN
DQN+UCB
QR-DQN+UCB
DQN+ TEE
QR-DQN+TEE
DON+ NON+ee--reedj
QR-DQN WE-DON+e--guue +2-ggeedy
NNisNNiyy
OR-DON+NOisyNe
QR-DQN+UCB++TE (EDE)
0.00 0.15 0.30
0.45
0.00
0.15
0.30
0.45
0.00
0.15
0.30
Min-Max Normalized Score
Figure 5: Test performance of different exploration methods on the Proegen benchmark across 5 runs.
compare with other popular exploration 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 complete comparison. The points of comparison bootstrapping to train several copies of models that have distinct exploration behaviors, (2) UCB [19] which uses ensemble of models
are: (1) Bootstrapped DQN [79] which uses
do uncertainty estimation, (3) ez-greedy exploration [24] which repeats the same random action (following a zeta distribution) to achieve temporally extended exploration, (4) UA-DQN [20], and (5) NoisyNet [32] which adds trainable noise the linear layers' ?
When UCB is combined with QR-DQN, we use the epistemic uncertainty unless specified otherwise. The results are shown in Figure 5. Details can be found in Appendix I.
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+UCD which is itself worse than EDE. Without epistemic uncertainty, the agent cannot do very well even if it trains diverse data, i.e. EDE is better than QR-DQN+TEE Similarly, 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 are competitive exploration methods on Atari. This highlights the importance of using distributional RL in CMDPs in order to disentangle the epistemic and aleatoric uncertainty. QR-DQN+UCB the other hand outperforms QR-DQN and DQN because it only targets the (biased) epistemic uncertainty. In Figure 16c from Appendix G, show that indeed exploring with the total uncertainty o² performs significantly worse 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 H).
QR-DQN with ez-greedy exploration marginally improves upon the base QR-DQN, but remains significantly worse than EDE. This may be due the fact that ez-greedy exploration can induce temporally extended exploration but it is not aware of the agent's 's uncertainty. Not accounting for uncertainty can detrimental since CMDPs can have much larger number of effective states than singleton MDPs. If the models have gathered enough information about state, further exploring that state can hurt sample efficiency, regardless of whether the exploration is temporally extended.
NoisyNet performs better than the other points of comparison we consider but it is still worse than EDE. The exploration behaviors of NoisyNet Net are naturally adaptive the agents will take into account what they have already learned. While direct theoretical comparison between NoisyNet and our method hard to establish, we believe adaptivity is a common thread for methods that perform well on CMDPs. Nonetheless, if we consider IQM, none of these methods significantly outperforms one another whereas our method achieves a much higher IQM. Note that TEE cannot be easily applied to NoisyNets and e2- ez-greedy which implicitly use an schedule.
'Taiaa et al. [107] found that, on the whole Atari suite, NoisyNet outperforms more specialized exploration strategies such as ICM [86] and RND [14]. This suggests that these exploration methods may have overfitted to sparse reward environments.

---

5.2 Crafter
To test the generality of our method beyond the Procgen benchmark, we conduct experiments on the Crafter environment [37]. Making progress on Crafter requires wide range of capabilities such as strong generalization, deep exploration, and long-term reasoning. Crafter evaluates agents using a scalar score that summarizes the agent's s abilities. Each episode is procedurally generated, so the number of training environments is practically infinite. While Crafter does not test generalization to new environments, it still requires generalization across the training environments in order to efficiently train on all of them. We build our method on top of the Rainbow implementation [43] provided in the open-soureed code of Hafner [37] and only modified the exploration.
We use Thompson sampling instead of UCB+TEE since only one environment is used and Thompson sampling outperforms UCB by itself in this environment. As seen in Table 1, our algorithm achieves significantly higher scores compared all the baselines presented in Hafner [37], including Dreamer V/ [38] which is a state-of-the-art model-based RL algorithm. It is also competitive with LSTM-SPCNN [101] which uses a specialized architecture that does not have spatial pooling with two orders of magnitude more parameters and extensive hyperparameter tuning. The significant improvement over Rainbow, which is a competitive value-based ap- proach, suggests that the exploration strategy is crucial for improving performance on such CMDPs.
Table 1: Results on Crafter af- ter IM steps and over 10 runs.
Method
Score (%)
EDE (ours)
11.7 1.0
Rainbow PPO
4.3 0.2
4.6
DreamerV2 LSTM-SPCNN
10.0 1.2
12.1 + 0.8
6 Related Works
Generalization in RL. A large body of work has emphasized the challenges of training RL agents that can generalize to new environments and tasks [93, 67, 50, 84, 118, 121, 77, 21, 22, 49, 58, 36, 18, 12, 13, 35, 53, 3, 26, 66]. This form of generalization is different from generalization in singleton MDP which refers to function approximators generalizing to different states within the same MDP. A natural way alleviate overfitting is to apply widely-used regularization techniques such as implicit regularization [100]. dropout [45], batch normalization [29], or data augmentation [116, 61, 60, 92, 112. 115, 40, 41, 54]. Another family of methods aims to learn better state representations via bisimulation metrics [120, 119, 11, information bottlenecks [45, 28). attention mechanisms [15], contrastive learning [72], adversarial learning [95, 34, 89], or decoupling representation learning from decision making [103, 99]. Other approaches use
approaches [18, 71],
non-stationarity reduction [46, 78], curriculum learning [47, 109, 48, 85], planning [4], forward- backward representations [1 [1]]], or diverse policies [57]. More similar to our work, Raileanu and Fergus [90] show that the value function can overfit when trained on CMDPs and propose to decouple the policy from the value optimization to 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 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-basee methods significantly worse on CMDPs.
Exploration. Exploration is a fundamental aspect of RL [55, 33, 108]. Common approaches include e-greedy [105), -ouut- exploration [104, 9, 68), curiosity-based exploration [97, 102, 86), or novelty- methods specifically designed for exploring sparse reward CMDPs [91, 117, 30, 122]. These methods are based on policy optimization and focus on training agents in sparse reward CMDPs. In contrast, we are interested 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 [82, 32), bootstrapping [79, 81], ensembles [19, 61, 64], uncertainty estimation [80, 83, 20], or distributional RL [70]. This class of method implements the principle of "optimism the face of uncertainty". EDE also falls under this broad class of algorithms. The main goal of these works is to balance return maximization (i.e., exploitation) and exploration to improve sample efficiency on sin- gleton MDPs that require temporally extended exploration. For generalization, both over-exploration and over-exploitation would result in poor generalization in addition to poor training performance, so methods that balance exploration and exploitation would likely be preferred.

---

Another related area is task-agnostic RL [88, 123] 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 benchmarks like Procgen. 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 [10], which is useful in CMDPs for avoiding value overfitting [90].
7 Conclusion
In this work, we study how exploration affects an agent's ability to generalize to new environments. Our tabular experiments indicate that effective exploration of the training environments crucial for generalization 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 be useful in other environments the agent may encounter at test time. To this end, we propose to encourage exploration of states with high epistemic uncertainty and employ deep ensembles and distributional RL to disentangle the agent's 's epistemic and aleatoric uncertainties. This results in the first value-basee based method to achieve ttateoff-thear performance on both Procgen and Crafter, two benchmarks for generalization in RL with high dimensional observations. Our results suggest that exploration is important for all RL algorithms trained and tested on CMDPs. While here we focus on value-bare based methods, similar ideas could be applied to policy optimization to further improve their generalization abilities. In a broader context, perhaps important to emphasize the exploration is not the only piece of puzzle for generalization in RL. There are still environments in Procgen where EDE does not significantly improve the performance even compared QR-DQN (Figure 13), which indicates that there other bottlenecks beyond poor exploration. Another limitation of our approach is that is more computationally expensive due to the ensemble (Appendix H). Thus, we expect it could benefit from future advances in more efficient ways of accurately estimating uncertainty in neural networks.
Acknowledgement
We would like thank Alessandro Lazaric, Matteo Pirotta, Samuel Sokota, and Yann Ollivier for the helpful discussion during the early phase of this project. We would also like to thank the members of FAIR London for their feedback this work and the members of Locus Lab for their feedback on the draft.
References
[1] R. Agarwal, M. C. Machado, P. S. Castro, and M. G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. 2021.
[2] R. Agarwal, M. Schwarzer, P. Castro, A. Courville, and M. Bellemare. Deep rein- forcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 34:29304 24422,04-293220, 2021.
[3] A. Ajay, G. Yang, O. Nachum, and P. Agrawal. Understanding the generalization gap in visual reinforcement learning. 2021.
[4] A. Anand, J. Walker, Y. Li, E. Vértes, J. Schrittwieser, S. Ozair, T. Weber, and J. Hamrick. Procedural generalization by planning with self-supervised world models. arXiv preprint arXiv:2111.015887 2021.
[5] J.-Y. Audibert, R. Munos, and C. Szepesvári. Exploration exxllootonn tradcoff using variance estimates in multi-armed bandits. Theoretical Computer Science, 410(19):1876-199) 1902, 2009.
[6] P. Auer, N. Cesa-Bianchi, and P. Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47(2):235-22 2566 2002.
[7] M. G. Azar, I. Osband, and R. Munos. Minimax regret bounds for reinforcement learning. In
International Conference on Machine Learning, pages 263-272. PMLR, 2017.
10

---

[8] A. P. Badia, B. Piot, S. Kapturowski, P. Sprechmann, A. Vitvitskyi, D.D. Guo, and C. Blundell. Agent57: Outperforming the atari human benchmark. In International Conference on Machine Learning, pages 507-517. PMLR, 2020.
[9] M. Bellemare, S. Srinivasan, G. Ostrovski, T. Schaul, D. Saxton, and R. Munos. Unifying count-based exploration and intrinsic motivation. Advances in neural information processing systems, 29, 2016.
[10] M. G. Bellemare, W. Dabney, and R. Munos. A distributional perspective reinforcement learning. In International Conference on Machine Learning, pages 449-458. PMLR, 2017.
[11] R. Bellman. A markovian decision process. Journal of mathematics and mechanics, pages 679-684, 1957.
[12] E. Bengio, J. Pineau, and D. Precup. Interference and generalization in temporal difference learning. ArXiv, 0t//2003.063550 2020.
[13] M. Bertrán, N. Martínez, M. Phielipp. and G. Sapiro. Instance based generalization in reinforcement learning. learning ArXiv, aabs2201.00108 2020.
[14] Y. Burda, H. Edwards, A. Storkey, and O. Klimov. Exploration by random network distillation. arXiv preprint arXiv:1810. 12894, 2018.
[15] W. Carvalho, A. Lampinen, K. Nikiforou, F. Hill, and M. Shanahan. Feature-attending recurrent modules for generalization in reinforcement learning. arXiv preprint arXiv:2112.0836 2021.
[16] O. Chapelle and L. Li. An empirical evaluation of thompson sampling. Advances in neural information processing systems, 24, 2011.
[17] B. Charpentier, R. Senanayake, M. Kochenderfer, and S. Günnemann. Disentangling epistemic and aleatoric uncertainty in reinforcement learning. arXiv preprint arXiv:2206.01558. 2022.
[18] J. Z.. Chen. Reinforcement learning generalization with surprise minimization. ArXiv, abs/2004. 12399, 2020.
[19] R. Y. Chen, S. Sidor, P. Abbeel, and J. Schulman. Ucb exploration via q-ensembles. arXiv preprint arXiv: 502, 2017.
[20] W. R. Clements, B. Van Delft, B.-M. Robaglia, R. B. Slaoui, and S. Toth. Estimating risk and uncertainty in deep reinforcement learning, arXiv preprint arXiv:1905.09638 2019.
[21] K. Cobbe, O. Klimov, C. Hesse, T. Kim, and J. Schulman. Quantifying generalization in reinforcement learning. arXiv preprint 1812.02344 2018.
[22] K. Cobbe, C. Hesse, J. Hilton, and Schulman. Leveraging procedural generation to bench- mark reinforcement learning. arXiv preprint
2019.
[23] W. Dabney, M. Rowland, M. Bellemare, and R. Munos. Distributional reinforcement learning with quantile regression. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 32, 2018
[24] W. Dabney, G. Ostrovski, and A. Barreto. Temporally-extended s-greedy exploration. In International Conference on Learning Representations, 2021. URL net/f oorumrmmm id 2NBPHHF77224..
[25] C. Dann and E. Brunskill. Sample complexity of episodic fixed-horizon reinforcement learning. Advances in Neural Information Processing Systems, 28, 2015.
[26] A. Ehrenberg, R. Kirk, M. Jiang, E. Grefenstette, and T. Rocktäschel. A study of off-policy learning in environments with procedural content generation. ICLR Workshop on Agent Learning in Open- Endednes,, 2022.
[27] L. Espeholt, H. Soyer, R. Munos, K. Simonyan, V. Mnih T. Ward, Y. Doron, V. Firoiu, T. Harley, Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. arXiv preprint arXiv:1802.01 2018.

---

[28] J. Fan and W. Li. Dribo: Robust deep reinforeement learning via milti-view information bottleneck. In International Conference on Machine Learning, pages 6074-6102, PMLR. 2022.
[29] J. Farebrother, M. C. Machado, and M. H. Bowling. Generalization and regularization in dqn. ArXiv, abs/1810.001233, 2018.
[30] Y. Flet-Berliac, J. Ferret, O. Pietquin, P. Preux, and M. Geist. Adversarially guided actor-critic. CoRR, abs/2102.04376, 2021. URL https://arxiv.. httgg/ab///22...////
[31] S. Fort, H. Hu, and B. Lakshminarayanan. Deep ensembles: A loss landscape perspective. arXiv preprint vrXiv:1912.027555 2019.
[32] M. Fortunato, M. G. Azar, B. J. Menick, I. Osband, A. Graves, V. Mnih, R. Munos, D. Has- sabis, O. Pietquin, et Noisy networks for exploration. arXiv preprint arXiv: I 7706. 10295, 2017.
[33] R. Fruit and A. Lazaric. Exploration-cxploiaaiii in mdps with options. In Artificcaa intelligence and statistics, pages 576-584. PMLR, 2017.
[34] X. Fu, G. Yang. P. Agrawal, and T. Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning, pages 3480-3491. PMLR, 2021.
[35] Ghosh, J. Rahme, Kumar, A. Zhang, P. Adams, and S. Levine. Why generalization in rl is difficult: Epistemic pomdps and implicit partial observability. Advances in Neural Information Processing Systems, 34:25502 25,55502-2555,, 2021.
[36] J. Grigsby and Y. Qi. Measuring visual generalization in continuous control from pixels. ArXiv, abs/2010.06740, 2020.
[37] D. Hafner. Benchmarking the spectrum of agent capabilities. In International Confer- ence on Learning Representations, 2022. URL. httpt:t/t//ppeee http::////peeeeeeeeeeeeee wwt/ffrumm/ 1W0z96MFEoH.
[38] D. Hafner, T. P. Lillicrap, M. Norouzi, and Ba. Mastering atari with discrete world models. In International Conference on Learning Representations, 2021. URL /
net/f form/foum?? Od=
[39] Hallak, D. Di Castro, and S. Mannor, Contextual markov decision processes. arXiv preprint arXiv:/502.088
2015.
[40] N. Hansen and X. Wang. Generalization in reinforcement learning by soft data augmentation. In 2021 IEEE International Conference on Robotics and Automation (ICRA), pages 13611- 13617. IEEE, 2021.
[41] N. Hansen, H. Su, and X. Wang. Stabilizing deep q-learning with convnets and vision transformers under data augmentation. Advances in Neural Information Processing Systems, 34:3680- 0698000-63,00000 2021.
[42] He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770- 778, 2016.
[43] M. Hessel, J. Modayil, H. Van Hasselt, T. Schaul, G. Ostrovski, W. Dabney, D. Horgan, B. Piot, M. Azar, and D. Silver. Rainbow: Combining improvements in deep reinforcement learning. In shirty-second AAAI conference on artificial intelligence, 2018.
[44] D. Horgan, J. Quan, D. Budden, G. Barth-Maron, M. Hessel, H. van Hasselt, and D. Sil- ver. Distributed prioritized experience replay. In International Conference on Learning Representations, 2018. URL https:// wwwpeerrrviee..ne//oormmm wwt..e//oooomiii-tyyy-------
[45] M. Igl, K. Ciosek, Y. Li, S. Tschiatschek, C. Zhang, S. Devlin, and K. Hofmann. Generaliza- tion in reinforcement learning with selective noise injection and information bottleneck. In Advances in Neural Information Processing Systems, pages 13956-13966 2019.
12

---

[46] M. Igl. G. Farquhar. J. Luketina, W. Böhmer, and S. Whiteson. The impact of non-stationarity on generalisation in deep reinforcement learning. ArXiv, abs/2006. 2020.
[47] M. Jiang, E. Grefenstette, and T. Rocktäschel. Prioritized level replay. ArXiv, abs/2010.03934, 2020.
[48] M. Jiang, M. Dennis, J. Parker-Holder, J. Foerster, E. Grefenstette, and T. Rocktäschel. Replay- guided adversarial environment design. Advances in Neural Information Processing Systems, 14:1884-1888 1897, 2021.
[49] A. Juliani, A. Khalifa, V.P. Berges, J. Harper, H. Henry, A. Crespi, J. Togelius, and Lange. Obstacle tower: A generalization challenge vision, control, and planning. ArXiv, abs/1902.01378, 2019.
[50] N. Justesen, R. R. Torrado, P. Bontrager, A. Khalifa, J. Togelius, and S. Risi. Illuminating generalization in deep reinforcement learning through procedural level generation. arXiv: Learning, 2018
[51] M. Kearns and S. Singh. Finite-sample convergence rates for q-learning and indirect algorithms. Advances in neural information processing systems, 11, 1998.
[52] Kendall and Gal. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems, 30, 2017.
[53] R. Kirk, A. Zhang, E. Grefenstette, and T. Rocktäschel. A survey of generalisation in deep reinforcement learning. arXiv preprint arXiv:2 (tt::2//1.090 1.09994, 2021.
[54] B. Ko and J. Ok. Efficient scheduling of data augmentation for deep reinforcement learning. arXiv preprint 1rXiv22206.00518, 2022.
[55] J. Z. Kolter and A. Y. Ng. Near-bayesian exploration in polynomial time. In Proceedings of the 26th annual international conference on machine learning, pages 513-520, 2009.
[56] I. Kostrikov, D. Yarats, and Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004. 13649, 2020.
[57] S. Kumar, A. Kumar, S. Levine, and C. Finn. One solution is not all you need: Few-shot extrapolation via structured maxent rl. Advances in Neural Information Processing Systems, 33:8 1119-99811,,, 2020.
[58] H. Kuttler, N. Nardelli, A. H. Miller, R. Raileanu, M. Selvatici, E. Grefenstette, and T. Rock- täschel. The nethack learning environment. ArXiv, aws/aab6..376666 13760, 2020.
[59] Lakshminarayannan A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
[60] M. Laskin, K. Lee, A. Stooke, L. Pinto, Abbeel, and Srinivas. Reinforcement learning with augmented data. arXiv preprint arXiv:2004. /4990, 2020.
[61] K. Lee, K. Lee, J. Shin, and H. Lee. Network randomization: A simple technique for generalization in deep reinforcement learning. In International Conference on Learning Representations. https://openreview. nev/forum, 2020.
[62] K. Lee, M. Laskin, A. Srinivas, and P. Abbeel. Sunrise: A simple unified framework for ensemble learning in deep reinforcement learning. In International Conference on Machine Learning, pages 6111-6141. PMLR, 2021.
[63] G. Li, Y. Wei, Y. Chi, Y. Gu, and Y. Chen. Sample complexity of asynchronous q-learning: Sharper analysis and variance reduction. Advances neural information processing systems, 33:7031-7043, 2020.
[64] L. Liang, Y. Xu, S. McAleer, D. Hu, A. Ihler, P. Abbeel, and R. Fox. Reducing variance in demporel value estimation via ensemble of deep networks. In International Conference on Machine Learning, pages 13285-13301. PMLR, 2022.
13

---

[65] E. Z. Liu, A. Raghunathan, P. Liang, and C. Finn. Decoupling exploration and exploitation for metu-reinforeement learning without sacrifices. In International conference on machine learning, pages 6925 6935. PMLR, 2021.
[66] C. Lyle, M. Rowland, W. Dabney, M. Kwiatkowska, and Y. Gal. Learning dynamics and generalization in reinforcement learning, arXiv preprint arXiv:2206.02126. 2022.
[67] M. C. Machado, M. G. Bellemare, E. Talvitie, J. Veness, M. J. Hausknecht, and M. H. Bowling. Revisiting the arcade learning environment: Evaluation protocols and open problems for general agents. In IJCAI, 2018.
[68] M. C. Machado, M. G. Bellemare, and M. Bowling. Count-based exploration with the successor representation. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pages 5125-5133, 2020.
[69] D. Malik, Y. Li, and P. Ravikumar. When is generalizable reinforcement learning tractable? Advances in Neural Information Processing Systems, 34:8032-8044) 2021.
[70] B. Mavrin, H. Yao, L. Kong. K. Wu, and Y. Yu. Distributional reinforcement learning for efficient exploration. In International conference on machine learning, pages 4424-4434, PMLR, 2019.
[71] B. Mazoure, R. T. des Combes, T. Doan, P. Bachman, and R. D. Hjelm. Deep reinforcement and infomax learning. ArXiv, abs/2006.07217 2020.
[72] B. Mazoure, I. Kostrikov, O. Nachum, and J. Tompson. Improving zero-shot generaliza- tion in offline reinforcement learning using generalized similarity functions. arXiv preprini 1rXiv:2/11.14600 2021.
[73] V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. Riedmiller. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
[74] J. Mockus. The bayesian approach to global optimization. In System Modeling and Optimiza- tion, pages 473-481. Springer, 1982.
[75] S. Mohanty, J. Poonganam, A. Gaidon, A. Kolobov, B. Wulfe, D. Chakraborty, G. Semetulskis, J. Schapke, J. Kubilius, J. Pašukonis, et al. Measuring sample efficiency and generalization in reinforcement learning benchmarks: Neurips 2020 procgen benchmark. arXiv preprint 7xXi:22033/55322 2021.
[76] J. Morimoto and K. Doya. Robust reinforcement learning. Neural computation, 17(2):335-359 2005.
[77] A. Nichol, V. Pfau, C. Hesse, O. Klimov, and J. Schulman. Gotta learn fast: A new benchmark for generalization in rl. ArXiv, abs/1804.023 804.03720, 2018.
[78] E. Nikishin, M. Schwarzer, P. D'Oro, P.-L. Bacon, and A. Courville. The primacy bias in deep reinforcement learning. In International Conference on Machine Learning, pages 16828-16847. PMLR, 2022.
[79] I. Osband and B. Van Roy. Bootstrapped thompson sampling and deep exploration. arXiv preprint arXiv:1507.00300 2015.
[80] I. Osband, D. Russo, and B. Van Roy. (more) efficient reinforcement learning via posterior sampling. Advances in Neural Information Processing Systems, 26, 2013.
[81] I. Osband, C. Blundell, A. Pritzel, and B. Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems, 29, 2016.
[82] I. Osband, B. Van Roy, and Z. Wen. Generalization and exploration via randomized value functions. In International Conference on Machine Learning, pages 2377-2386. PMLR, 2016.
[83] O'Donoghue, Osband, R. Munos, and Mnih. The uncertainty bellman equation and exploration. In International Conference on Machine Learning, pages 3836-3845, 2018.
14

---

[84] C. Packer, K. Gao. J. Kos, P. Krähenbühl, V. Koltun, and D. X. Song. Assessing generalization 2018.
in deep reinforcement learning. ArXiv,
[85] J. Parker-Holder, M. Jiang, M. Dennis, M. Samvelyan, J. Foerster, E. Grefenstette, and T. Rocktäschel. Evolving curricula with regret-based environment design. arXiv preprint arXiv:2203.01302. 2022.
[86] D. Pathak, P. Agrawal, A. A. Efros, and T. Darrell. Curiosity-driven exploration by self- supervised prediction. In International conference on machine learning, pages 2778-2787. PMLR, 2017.
[87] T. Pearce, Leibfried, and A. Brintrup. Uncertainty in neural networks: Approximately bayesian ensembling. In International conference on artificial intelligence and statistics, pages 234-244, PMLR, 2020.
[88] H. Pong, M. Dalal, S. Lin, A. Nair, Bahl, and S. Levine. Skew-fit: State-covering self-supervised reinforcement learning, arXiv preprint arXiv:1903.03690 2019.
[89] M. M. Rahman and Y. Xue. Adversarial style transfer for robust policy optimization in reinforcement learning. 2021.
[90] R. Raileanu and R. Fergus. Decoupling value and policy for generalization in reinforcement learning. In International Conference on Machine Learning, pages 8787-8798. PMLR, 2021.
[91] Raileanu and Rocktäschel. Ride: Rewarding impact-driven exploration for procedurally- generated environments. In International Conference on Learning Representations, 2020. URL httptt/////peeeeeeeeeee.eo www.eeeneeettoommmddtdeesee TJBFPB.
[92] R. Raileanu, M. Goldstein, D. Yarats, I. Kostrikov, and R. Fergus, Automatic data augmentation for generalization in deep reinforcement learning. ArXiv, abs/2006. 12862, 2020.
[93] A. Rajeswaran, K. Lowrey, E. Todorov, and S. M. Kakade. Towards generalization and simplicity in continuous control. ArXiv, abs/ 17011703.00 7226600 2017.
[94] S. Ross, G. Gordon, and D. Bagnell. A reduction of imitation learning and structured predic- tion to no-regret online learning. In Proceedings of the fourteenth international conference on artificial intelligence and statistics, pages 627 635. JMLR Workshop and Conference Proceedings, 2011.
[95] J. Roy and G. Konidaris. Visual transfer for reinforcement learning via wasserstein domain confusion, arXiv preprint arXiv:2006.00 033665, 2020.
[96] T. Schaul, Quan, I. Antonoglou, and D. Silver. Prioritized experience replay. arXiv preprint 1rXiv:/5/1.05990 2015.
[97] J. Schmidhuber. A possibility for implementing curiosity and boredom in model-building neural controllers. In Proc. of the international conference on simulation of adaptive behavior: From animals to animats, pages 222-222222222222222220 222, 1991.
[98] B. Settles. Active learning literature survey. 2009.
[99] A. Sonar, V. Pacelli, and A. Majumdar. Invariant policy optimization: Towards stronger
generalization in reinforcement learning. ArXiv, abs/2006.01096, 2020.
[100] X. Song, Y. Jiang, S. Tu, Y. Du, and B. Neyshabur. Observational overfitting in reinforcement learning. In International Conference on Learning Representutions, 2020. URL https: // //openrevier net
[101] A. Stanić, Y. Tang, D. Ha, and J. Schmidhuber. An investigation into the open world survival game crafter. In Decision Awareness in Reinforcement Learning Workshop at ICMI. 2022, 2022. URL https:/ htt.r/ereee...eeerermmreoecsteoe
15

---

[102] C. Stanton and J. Clune. Curiosity search: Producing generalists by encouraging individuals continually explore and acquire skills throughout their lifetime. PLOS ONE, 11(9):1-20, 09 2016. doi: 10. 10.137 1aj/////orlll.ppp.. pone. 62235.. URI. https::tt//... 1ttr//0..371/jurnnaa pone. 0162235.
[103] A. Stooke, K. Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforcement learning. ArXiv, abs/2009.087 089319, 2020.
[104] A. L. Strehl and M. L. Littman. An analysis of model-based interval estimation for markov decision processes. Journal of Computer and System Sciences, 74(8):1309-1331, 2008.
[105] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
[106] R. S. Sutton, D. McAllester, S. Singh, and Y. Mansour. Policy gradient methods for rein- forcement learning with function approximation. Advances in neural information processing systems, 12, 1999.
[107] A. A. Taiga, W. Fedus, M. C. Machado, A. Courville, and M. G. Bellemare. On bonus based ex- ploration methods in the arcade learning environment. In International Conference Learning Representations, 2020. URL https: /estps://openrrree/ openrevieneeeeeeeeoe net meet/fooummiiee Bet.ruunnndddJeesssslsl
[108] J. Tarbouriech, E. Garcelon, M. Valko, M. Pirotta, and A. Lazaric. No-regret exploration in goal-oriented reinforcement learning. International Conference on Muchine Learning, pages 9428 9437. PMLR, 2020.
[109] E. Team, A. Stooke, A. Mahajan, C. Barros, Deck, Bauer, J. Sygnowski, M. Trebacz, M. Jaderberg, M. Mathieu, et al. Open-ended learning leads to generally capable agents. arXiv preprint arXiv:2107. 12808, 2021.
[110] W. R. Thompson. On the likelihood that one unknown probability exceeds another in view of the evidence of two samples. Biometrika, 25(3-4):285-22 1933.
[111] A. Touati and Y. Ollivier. Learning one representation to optimize all rewards. Advances in Neural Information Processing Systems, 34:13-23, 2021.
[112] K. Wang, B. Kang, J. Shao, and J. Feng. Improving generalization in reinforcement learning with mixture regularization. ArXiv, 0tts/2010.1081... 2020.
[113] J. Watkins and Dayan. Q-learning Machine learning, 8(3):799-22 22222, 1992.
[114] R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement reineooreeessreiiiffrreeeee learning. Machine learning, 8(3):29--266, 1992.
[115] D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645, 2021.
[116] C. Ye, A. Khalifa, P. Bontrager, and J. Togelius. Rotation, translation, and cropping for zero-shot generalization. arXiv preprint arXiv:2001.0990 009008 2020.
[117] D. Zha, W. Ma. L. Yuan, X. Hu, and J. Liu. Rank the episodes: A simple approach for exploration procedurally-
environments. arXiv preprint arXiv:2101.08152, 2021.
[118] A. Zhang, N. Ballas, and J. Pineau. A dissection of overfitting and generalization continuous reinforcement learning. ArXiv, abs//806.. 0806707377 2018.
[119] A. Zhang, C. Lyle, S. Sodhani, A. Filos, M. Kwiatkowska, J. Pineau, Y. Gal, and D. Precup. Invariant causal prediction for block mdps. arXiv preprint arXiv:2003.060166 2020.
[120] A. Zhang, R. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006. 2020.
[121] C. Zhang, O. Vinyals, R. Munos, and Bengio. A study on overfitting in deep reinforcement learning. ArXiv, abs/ 18 804, 06893, 2018.
16

---

[122] T. Zhang, H. Xu, X. Wang. Y. Wu, K. Keutzer, J. E. Gonzalez, and Y. Tian. Noveld: A simple yet effective exploration criterion. In A. Beygelzimer, Y. Dauphin, Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL http:/ouumiid-C/pppdDDa0
[123] X. Zhang, Ma, and A. Singla. aask-agnostic exploration in Neural Information Processing Systems, 33:1 11773-111444 2020.
explorrrttiiellllttttooo in reinforcement learning. Advances
17

---

A Sample Complexity of Q-Learning
While the statement that "exploration helps generalization" should be fairly intuitive, the theoretical argument is surprisingly more nuanced. Instead of proving new results, we will rely several existing results in the literature. To understand why exploration could help the performance of the learned policy under different initial state distributions (i.e., generalization to different initial state distributions), we can study the asymptotic sample complexity of Q-learning in the tabular case, which measures how many steps are required to get good estimate for all Q(8, a). If the amount of experience is less than the sample complexity, the Q-value can be inaccurate, thus leading to poor performance.
In the theoretical analysis of asynchronous Q-learning, it is common to assume access to a fixed behavioral policy, Tb, which collects data by interacting with the environment and that the environ- ment is ergodic (see all assumptions in Li et al. [63]). The update to the value function is done in an off- off-policy manner with the data generated by Tr. Intuitively, the behavioral policy is responsible for exploration. We will first define some useful quantities that relate to the inherent properties of Ty and the MDP. Recall that j is an initial state distribution and let (r s) be the trajectory distribution induced by и starting from 8, we can define the state marginal distribution at time t given that the initial state is drawn from as (s(8; n. (4) = Eropr (-1=0) [1 (s. 8}}. the state-action marginal distribution Pe(s, 0, T. p.) = x(a Ps; J.), the stationary state occupancy distribu- tion, dxe (b) = lim T4I Et (s (4) and finally the stationary state-action distribution (s, a) = (aa (s) 8. Using Drv to denote the total variation distance, we define:
(1min =
min (s,,),, (x.a)E5xx
(6)
11 Div (P((s, a; n, 5(80, au)), d. = 4
(mi((n) min
max (
(7)
Intuitively, (mia(()) is the occupancy of the bottleneck state that is hardest to reach by æ, and mmix captures how fast the distribution generated by decorrelates with any initial state distribution. In this setting, Li et al. [63] showed that with Ô [(m((T)[11)))
1
Menin (aai(T))) (To)(1 (1--)) steps", one can expect
fais Emia((a)) (STa)
the Q-value to be e-close to the optimal Q-value, i.e., max A (Q((s a) (s, a)l S c with high probability. The second term of the complexity concerns the rate at which the state-action marginal converges to steady state, and the first term characterizes the difficulty of learning once the empirical distribution has converged. The crucial point of this analysis is that the sample complexity is inversely proportional to the visitation probability of the hardest state to reach by Te. so better exploration policy that aims to explore the state space more uniformly (such as UCB) would improve the sample complexity, which in turn leads to better generalization (i.e., lower Q-value error for all state-action pairs).
The argument above shows that good exploration could ensure that the Q-values are accurate for any initial state distribution (which guarantees generalization in a uniform convergence sense), but it does not immediately explain why the Q-values are accurate around the training initial state distribution even though the training initial states further away have inaccurate Q-values, i.e., convergence at these states further away is slower. It turns out that this is related to resetting. It is perhaps worth noting that the analysis so far deals with the continuous RL (i.e., not episodic and does not have an end state). To account for the effect of reset distribution, we may assume that after the episode ends (e.g., reaching the goal in Figure 2), the agent is reset to a state drawn from the initial state distribution or that the agent may be reset to a state drawn from the initial state distribution at every timestep. We use this the continuous MDP as a proxy for the original MDP 10, this setting, dx (s) and consequently dz a) will naturally have a higher density around the initial state distribution, so the Q-values at states further away from the initial state would converge slower, leading to these states having poor return even though the policy has high return around the states in the training initial state
'Notcee that the stationary state-action distribution is independent of since it is the stationary distribution of Markov chain; however, it could still he close to particular /reee if the environment resets Prese and m. does not explore efficiently (L.e,, the Markov chain has high prohability of visiting the initial state distribution).
9Ô hides logarihgaii factors.
1fff the original episodic MDP is ergodic excluding the end state, then the continuous MDP will also be ergodic; however, this modified continuous MDP, in general, does not have the same optimal policy as the original episodic MDP, but it may be close for many common MDPs.
18

---

distribution. In other words, we can roughly say that the more state-action pair is visited, the more likely that its Q-value has a low error for many MDPs.
To further formalize this intuition, we will leverage the concept of sub-MDP [51], where we prune (s, a) pairs from MM based on the behavioral policy. Kearns and Singh [51] prunes M based on whether dx (s, a) is larger than a, a hyperparameter that is larger than Amin (mon) and argues that one can obtain better sample complexity pairs that are not likely visited by the behavioral policy in the analysis. Concretely, the sub-MDP M₁ (a) an MDP where the state space
the sub-MDP. This can be seen as throwing away state-action
(8) If further assume that all the state-action pairs visited by the optimal policy, x², starting from l are contained within G(a)). then the optimal policy of the Ma, (ax recovers that of the M starting from μ. This assumption generally holds for many reasonable MDPs (e.g. those with deterministic transition or bounded reward) and exploration policies тT - in the grid world MDP (Figure 2a), the optimal policy starting from the top left only moves right and never visits any other states, and these states are well covered by E-greedy exploration; for the training initial state, one can disregard all states that are not on the first row and still retain the optimal policy. By definition > Mmin (Tb) on M. (a) and thus Q-values would converge to the optimal value on Mx, (a) faster than the full MDP, M. a is sufficiently larger than ( the sample complexity on M (ax) could be much better than on M, explaining the policy's superior performance near the initial state when the Q(, a) further away from the initial states have not converged. Note that Mz, (ax) only a construct for making the theoretical argument. The Q-values of the "pruned" state-action pairs as usual algorithmically. Even if the assumption is not true, the optimal policy on Ma, (a) may still be good enough. If the optimal policy Ma, (a) or has to be small for the assumption to approximately hold, then one would need better behavioral policy to solve the training environment efficiently, much less having good generalization performance on the test environment.
(a) = ((s,a) E E x A ds, > at.
still updated
Finally, Li al. [63] showed that with high probability min,,a (s, a)/t S Mmin (b)) which means that having behavior policy that encourages visiting states with low visitation count, (s,, a). should result increasing mins,a 3 (s, at/t and thus increasing /mmin (T6). This indicates that UCB-based exploration should have a better sample complexity than e-greedy which does not explicitly target the visitation count and can be general pretty bad at covering large state space. Indeed, for episodic RL, the state-of-the-ard algorithms generally rely UCB-based exploration [25] to ensure good sample complexity, which is more similar to what we do in the tabular tasks. Nonetheless, we focus on the continuous case because it highlights the importance of exploration via Mmin.
B Case Study: Simplified Bigfish
"1 rEward
Training
Greedy Trajectory
muarr
Explorative Trajoctory
Test MDP
Trajectory
Figure 7: Simplified version of bigfish.
19

---

Figure 6: Example frames of bigfish. The goal for the agent (green fish) is to avoid fish bigger than itself and eat smaller fish. The spawning of enemy fish is different across different MDPs and the background can change.
To further illustrate the intuition that exploration on the training MDPs can help zero-shot gener- alization to other MDPs, we will use a slightly simplified example inspired by the game bigfish (Figure 6) from Procgen. As illustrated in Figure 7, in bigfish, the agent (green circle) needs to 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 al always starts in the state shown in the 7 = 0 of the greedy trajectory (top row). Using random exploration strategy (e.g., e-greedy), the agent should be able to 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- -wware 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 no impact on the performance in the training MDP because the agent has learned the optimal policy regardless of which exploration strategy is used.!!
However, on the test MDP, the initial state 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 the agent has not seen it during training. One could hope that the neural network has learned 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
(pink arrow). Due to this structure, agents that have learned about
Using function approximators with insufficient capacity may have some effects the optimal policy.
20

---

the explorative trajectory during training time will be able to recognize that the better action from 1 move to the right which will avoid death and ultimately result in +1 reward.
Once again, we emphasize that this simplified example for illustrative purposes. In reality, the sequences of states do not need be exactly the same and the neural network can learn to map similar states to similar representations. In principle, the neural network can also learn to 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 to 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 us to improve generalization via exploration, but should not hurt. Like most problems deep RL, is difficult to prove this hypothesis analytically but empirically it is indeed the case that EDE as better or the same performance as QR-DQN in out of 16 Procgen environments and perform approximately the same in the remaining environments (Figure 13).
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.
C Tabular Experiments
C.1 Generalization to Different Initial States
greeddy wi different opsilio
UCB wi different exploration bonus
Edily gradient ai different entrey borus
13-
15
10
22.2
213
-05
-05
0.2
200
800 20
6.0
500
-10
(a) e-greddy
(b) UCB
(c) Policy Gradient
Figure 8 Mean test return for each method with different exploration hyperparameters.
Both the training MDP, Musin and the test MDP, Musts 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' 8, a), to the next state corresponding an action is 100%. 7 is 0.9. The two MDPs differ by only their initial state distribution: (Ptrai = (0.0) (s) and Prest (s) = 1 (1.0) (s).
For both policy gradient and Q-learning, we use base learning rate of oq 0.05 that decays as training progresses. Specifically, at time step t, a₁ = Vt. ag.
For Q-learning, the Q-function parameterized as 9 E giiiiAA where each entry is a state-action value. The update is:
a.) (s), + a. (re max Q11) QQ (s).
where ((s, a) & For Q-learning with UCB, the agent keeps N(s, a) that keeps track of how many times the agent has taken action a from over the course of training and explores according to:
logt
(a 8) 1a+ (a) where a" arg max Q(5,a) + +oV N(s,a)
(10)
For Q-learning with e-greedy, the exploration policy is:
Regreedy (a = (1 - -E) larg max, Q(a,a)) (a) +e
Σ
|A| 1ar (a)
(a((a))
(11)
Va'EA
21

---

Ties in arg max are broken randomly which acts as the source of randomness for UCB.
C.2 Policy gradient
For policy gradient [114, 106], the policy is parameterized A IR Six Al where each entry is the
logit for the distribution of taking action a from 8, i.e., Ta 8) Eo'ed explu,,.') gradient is:
exp(d...) The overall
T Σ
V LO [J(T)] = Exo
V log To( 86)
Σ Tx k=t
(12)
1=
VoHoo
VOlET
To(a 8t) log log
(13)
1=0
The update is:
g(++1) g(t) + (VoEn, [J(T]] + BV,H. H )
(14)
where the expectation estimated with single trajectory.
For each method, we repeat the experiment for the following 10 hyperparameter values:
[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9)
* c: [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 approaches 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 8. We observe that both UCB and policy gradient hyperparameters, whereas e-greedy's 's generalization performance varies greatly with E. Furthermore, with extremely large e, the generalization performance is still far gradient. The best-performing values are respectively &* = 0.9, c* = 45 and 3 = 0.1 and the train/test performance of this configuration are shown in Figure 9.PG's performance lies between Greedy and UCB. On average, it converges to a better solution with lower variance than Greedy but is not as good as UCB.
relatively robust to the choice of
than UCB and policy
Policy Gradient
4.0
3.5
train
suboptimalisy
2.0
15-
1.0
0.5
0.0
0.5
episode
Figure 9: Suboptimality of the best hyperparameters of Policy Gradient
C.3 Full trajectory update and Sarsa X)
Exploration is evidently not the only approach one can take improve generalization in RL. In fact, we see that policy gradient with rather naive softmax exploration can also outperform e-greedy exploration (Appendix C.2). We hypothesize that this improvement is due to the fact that these algorithms use whole-trajectory updates (i.e., Monte Carlo methods) rather than one-step updates like Q-learning. To study this, we also consider Sarsa(A) [105, Ch 12.7] which is an on on-policy RI. method that learns vavaue with whole-trajectory updates. The results are shown in Figure 10a.
22

---

While both on-policy methods use fairly naive exploration methods (softmax exploration for PG and e-greedy for Sarsa(A), we observe that both are able to outperform Greedy. This observation suggests that their improved performance could be attributed to something else. We believe this is due to the fact that when we use whole-trajectory updates (or n-step returns), the learning signal is propagated to states with low density as soon good trajectory is sampled. On the other hand, when we use off-policy approaches, the signal is only propagated when the state action pair is sampled multiple times by the behavioral policy, which may be exponential in the horizon. Nonetheless, there is still a gap between these two methods and UCB, suggesting that the generalization benefits of whole- trajectory updates may be distinct from that of exploration. Hence, better exploration can further improve generalization.
C.3.1 Details of Sarsa())
(Arsa(A) is an on-policy value-based method that learns Q-value with whole-trajectory update via eligibility traces. Since the algorithmic details are more involved than the other methods we consider, refer the reader to Chapter 12.7 of Sutton and Barto [105] for the details. The Sarsa(A) (A) with E-greedy experiments inherit all the hyperparameters of Q-lcarning + 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 a fixed À = 0.9 for all experiments. We sweep & 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 10a 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 10b). With the best hyperparameter configuration (Figure 10c), Sarsa(A) outperforms both Greedy and PC, but is still worse than UCB and has a much larger variance.
Sarsa(A) + E-greedy
SamallA w/ d
Standard Destation Sartalái dfferent
4.0
3.5-
suboptimalit
2.5-
2.0
750m
1.5-
0.00
1.0
0.0
0.
0.5
0.0
-05
-0.5
200
400
600
800
200
402
610
600
-10
200
400
610
500
episode
opisade
<pisode
(c) Best hyperparameter. and = 0.9,
0.9
(a) Average of return.
(b) Standard deviation return.
Figure 10: Test performance of Sarsa A) across different s. The results of Sarsa (X). SA)sa(A) is value-hased method, but it is able achieve competitive performance because it uses whole-trajectory similar to policy gradient.
C.4 Generalization to Different Dynamies
In previous section, we presented a 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 a 7 x 7 grid as the state space (Figure 1 1la). 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 is that the agent always starts at (0, 3) but test time, a gust of wind can blow the agent off course. Specifically, with probability of 40%, the agent will be blown to +2 unit in the y-direction. This environment is an example of the windy world from Sutton and Barto [105]. The results are shown in Figure 11 and 12, 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 oracle that consistently moves towards
23

---

Reward
Q-learning E-greedy
Q-leaming UCB
Policy Gradient
0.14
train
suboptimallty
Subaptimall
4.C4 0.04
episode
episode
cpisode
(a) Tabular CMDP
(b) e-greedy
(c) UCB
(d) Policy Gradient
Figure 11: (a) During both training and test time, the agent starts in the blue square, but at test time, with 40% probability, wind can blow the agent down by two unit. In both cases, the goal is 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-leaming with e-greedy exploration, (c) Q-lcarning with UCB exploration, and (d) policy gradient with entropy bonus.
>igeeelyy different wpuilnn
UCB wi eeploration bonux
Pilicy gradient wt diffeent entropy borus
14:27
18.
2.4
-10
200
4.0
episdde
epionde
(a) e-greedy
(b) UCB
(c) Policy Gradient
Figure 12: Mean test return for each method with different exploration hyperparameters when transition dynamic changes at test time.
the goal). This setup is related to robust reinforcement learning [76], but exploration makes no assumptions about the type magnitude of perturbation (thus making complicated perturbation, e.g. different random initializations of the same game, possible) and does not require access to a simulator.
D
Additional Background
D.1 Quantile Regression DQN
Quantile regression DQN (QR-DQN) is variant of distributional RL that tries learn the quantiles of the state-action value distribution with quantile regression rather than produce the actual distribution of Q-values. Compared the original C51 [10], QR-DQN has many favorable theoretical properties. Here, we provide the necessary background for understanding QR-DQN and refer the readers to Dabney al. [23] for more in-depth discussion of QR-DQN.
Given a random variable Y with associated CDF Fy (y) = PYY & y), the inverse CDF Y is
Fy'
inf {y
S Fy (v))
To approximate Y, we may discretize the CDF with qunatile values:
(T), (TO), (Tw)) where n
N
Using the same notation as main text, we use 0 S A - R.N to denote a parametric function that maps ante-action pair, (s, a), to the estimated quantiles {0) a) The approximated CDF of
24

---

(s,, a) is thus:
N N Σ ((0)(s,)))
Z(s, a
where 8 (0, a)) is the Dirac delta distribution centered at ( a). To learn the quantiles, quantile regression which allows for unbiased estimate of the true quantile values. Concretely, for target distribution and quantile т, the value of Fz (7) is the minimizer of the quantile regression loss:
Fz (T) arg min LōR (0) Ez (p- 0)] where p+(u) - < 0)).
Due to discretization, for each quantile 7j, the minimizer of the 1-Wasserstein distance within that quantile is at that satisfies Fz (0) (T + T2) fÚ (see Lemma 2 of Dabney et al. [23]). Consequentily, to approximate Z, we can simultaneously optimize for all of the quantiles, {0, ) by minimizing the following loss:
FOR (0,)
Σ Renz [pt,(2 (2,((=-0))) 0,)]
To stablize optimization, QR-DQN uses an modified version of quantile loss called quantile Huber loss which utilizes the Huber loss with hyperparameter i::
(x(u) = 2 S x))+ + x [Iu]
1 =x 1([u] > x).
The quantile Huber loss is defined as:
(u) = =T - 1(u 0)) (w).
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, (arget and transition tuple (8, a, r, 8'), the target distribution is approximated as:
Q(s',a')
1 N
Σ
0 glarge(( a')
a' arg max Q(s', a')
T0, r 70 0 (8' a²). The TD target Q-value distribution's quantiles and distribution are:
(15)
(T01, TON).
TZ
8 (TO,).
Treating the target distribution oracle (i.e., not differentiable), update the current quantile estimates to minimize the following quantile regression loss:
Σ E 0f,(2 (s.
p. (sg [O, (,(s, a))
(16)
sg denotes the stop gradient operation as we not optimize the target network (standard DQN practice).
E 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 pscudocode omits the
25

---

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 differences happen in how exploration is being performed ie., 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, we compute the estimate for the epistemic uncertainty well the expected Q-values. Using these quantities, choose the action based on either Thompson sampling or UCB.
In Algorithm 3, 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 as if each ensemble member is a QR-DQN. However, doing so naïvely means that the feature extractor will be updated at much faster rate. To ensure that all parameters are updated at the same rate, for each update, we randomly choose a single ensemble member that will be responsible for updating the feature extractor and the gradient of other ensemble heads is not propagated to the feature extractor. This also saves some compute, In practice, this turns out be crucial for performing good uncertainty estimation and generalization. We show the effect of doing so in Figure 17c.
Attentive readers would notice that EDE is expensive than other methods since it needs to process M minibatches for each update. In our experiments, the speed of EDE with 5 ensemble members is approximately 2.4 times slower than the base QR-DQN (see Appendix H). On the other hand, we empirically observed that the deep ensemble is crucial for performing accurate uncertainty estimation in CMDPs. UA-DQN [20], 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). In this work. we do investigate other alternatives of uncertainty estimation but we believe future works in improving uncertainty estimation in deep learning could significantly reduce the computation overhead of EDE
Algorithm EDE
1: Initialize feature extractor f and ensemble heads {g
L.M
2: Initialize target feature extractor flagget and ensemble heads {g.
target M T==
3: Initialize replay buffer Buffer
4: {s(8), done (k)14 K reset all environment
JR=1
5: for from to do
> Standard DQN training loop
6: for & from 1 to K do
7: a a (HOOSEACTION(S!*) f. (o.)
8: s' r, (k) t Take in the environment
add (M( a, r, done (8) to Buffer
9:
(k)
11:
if done then
((k) , done <- reset the th environment
13:
end if end for
14:
15:
UPDATE(Buffer, f, M4 flrget
target M {g.
Ji=1)
16:
if t mod update_freetee frequency = 0 then
flarget {g, jarger j + f, M
18:
end if
19:
end for
26

---

Algorithm 2 CHOOSEACTION(Ss f. (9ih-1,))
1: feature - -f(s)
D Only one forward pass on the feature extractor
2: for from to M do
3: (0,,(8,a)))) f (feature)
Compute individual ensemble prediction
4: end for
5: Compute (s, a) using Equation 2
Estimate epistemic uncertainty with ensemble
6: Compute a' using Equation 5 or Equation 3
b Explore with estimated uncertainty
7: return a*
Algorithm 3 UPDATe,Buffer, J, {gi) flaget
(g, tageet
1: C C+0 0
2: grad_indes ~ M}
D Randomly select an ensemble member
3: for from 1 to M do
>> Train each ensemble member independently
4: {s,a, s', done) +- Sample from Buffer
5: feature +- f (s')
6:
(s', a. ] (0
9,
target
(feature(s'))
7: Compute (70.1: T0 ...,, TO.N) using Equation 15 for the 7th ensemble member
8: feature(s) (8) f f(s)
9:
i # grad_inder index then
b Prevent over-training the feature extractor (FOGGRA
10:
feature(s)
11:
end
12:
{0., (s, a)} y=1 (feature(s))
C. 4- compute QR loss using Equation 16 for the 7th ensemble member
14:
C Li
end for
16: Compute gradient of L w.r.t the parameters of f, b
17: Update the parameters with Adam
27

---

Procgen Results
In Table 3 and Table 2, we show the mean and standard deviation the final unnormalized train and test scores (at the end of 25M environment steps) of EDE long side with other methods in the literature. We adapt the tables from Raileanu and Fergus [90]. In Figure 13, we show the curves of min-max normalized test score that reproduced using the code from Raileanu and Fergus [90]. We see that for a subset of games, EDE significantly more sample-efficient For other games, EDE either on par with the policy optimization methods or fails train, indicating there other issues behind its poor performance on these games. Note that IDAAC results tune the hyperparameters for individual games whereas only tune the hyperparameters on bigfish.
Game
PPO
MixReg
FLR
UCB DRAC
PPG
DAAC
IDAAC
CDE
biglish
StrPiltt
24.9
32.4
27.9
30.0 1.3
47.2 1.6
36.4 2.8
37.0
2.3
49.6
FuiiRo
26.2 2 1.2
27.3 0.8
28.0
1.4
27.6 0.4
27.8
0.6
28.6 1
27.9 27.9 0.5
25.7
BossFight
7.4.:0. 0.4
6.2 0.7
8.9 0.4
7.8
1.2
0.6
9.8
10.0
0.2
5.9 6.5
2.2
832
14.3
20.7 3.3
23.3
Caveflyer
5.1
6.1
6.3 0.5
0.8
7.0
0.2
5.0
7.9
CoinRun
0.2
0.3
8.8 0.5
0.2
0.0
0.2
9.4
6.7 2 0.5
Jumper
Chaser
5.8 1.111
6.9
63206
0.8
8.2
6
1.100. 0.1
Climber
5.6 0.5
0.7
0.8
0.6
2.80044
T.2
8.4
5.7 1.1
Dodgeball
0.1
1.4
1.8
0.9
2.3
3.3
13.3
3.400.
Lenger
2.2
5.3
0.8
0.9
1.0
7.3
7.7
6.4 0.3
Mase
5.5
5.2
5.3
0.8
0.1
0.1
0.2
5.6 0.3
3.4 0.5
Miner
0.4
0.4
9.6 0.6
9.220.6
7.2
0.8
0.4
0.7 2 0.1
Table 2: Procgen scores on test levels after training on 25M environment steps. The mean and
standard deviation are computed using runs with different seeds.
PPO
MisRag
UCB-DRAC
PPG
bigfish
9.2 2 2.7
15.0 + 1.3
7.8 4 1.0
12,8
1.8
19.9 1.7
20.1 $ 1.6
21.8 $ 1.8
27.5 2.0
StaPPloe
29.0 1.1
28.7
1.1
2.6
0.3
34.1 1.3
49.6 2.1
38
2.6
38.6
22
46.9 2 U.T
FraRRt
0..
BoatFight
11.1
Niga
7.3 +: 0.3
8.2 8.2 0.4
3.4 0.5
8.0 8.0 0.4
8.9
0.2
8.8
022
8.9 0.9 0.3
9.1 0.4
Plunder
6.1 + 0.8
6.2 0.3
4.1 1.3
10.2 1.76
16.41.. 1.9
22.6
24.6
8.2
CaveFiyer
0.7
6.1
5.8
9.6
0.4
0.6
10.6 0.1
CoinRun
- 0.3
0.5
0.4
9.2
9.9
0.8
0.8
6.6
Janper
$.3 2 0.2
8.6
0.6
0.2
8.7
0.6
8.7
9.0
Chaser
4.1 0.3
0.9
0.7
7.0 0.6
1..7
69
7.8
2.2 0.1
Climber
1..0
10.2
Dodgeball
0.1
20
7.3
5.6
5.240.4
4.3
0.3
Heist
T.1 0.5
4.4 0.3
1.4
0.6
7.4
0.7
4.5
7.2 D.1
Leaper
5.5 0.4
3.2
0.4
5.0
9.3
8.0
0.3
8.0 0.3
Made
8.6
9.0
6.6
6.4
0.5
8.9.0.
9.7
1..0
11.3 1.0
11.3..0.9
11.5 0.5
1.1 0.2
Table 3: Procgen scores on train levels after training on 25M environment steps. The mean and standard deviation are computed using 5 runs with different seeds.
28

---

bigfish
bosstight
caveflyer
chaser
C21
626
60-
33-
86-
0.2
064
4.5
00-
0.3
65
3.
1.3
0.5
1-4
0.0
2.
climber
coinrun
dodgeball
fruitbot
3.4
0.5
044
66
0.1
3.
helst
Jumper
leaper
imaee
0.3
0.2
0.0
-0.2
Vwn Mype
miner
ninja
3x7
plunder
626 star starpilot
307
606
205
3.25
0.14
64
11-
-3.5
3.
-322
0.0
Figure 13: Min max normalized test performance of few representative methods individual Proegen games. Mean and standard deviation are computed over seeds. We can see that games on which value-bast methods can make meaningful progress, EDE almost improve upon QR FR-DQN. In many cases, EDE is significantly more sample efficient compared to the other methods. There are notably 4 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 appears remain challenging for value-based methods. Getting value-based methods generalize on these games would an interesting future direction.
F.1 Additional Ablation
Median
IQM
Mean
Optimality Gap
DON+UCB+TEE
0.00
0.02
0.04
0.06
0.015 0.030 0.045
0.015 0.030 0.045
0.960 0.975 0.990
Min-Max Normalized Score
Figure 14: The results on Procgen for DQN The performance is much worse than the other ablations.
This is not surprising. Since DQN+UCB performed unfavorably, it is unlikely that TEE would improve its performance.
29

---

Task ID ablation
Correlation train performance
1.0
Singleton Training Performance
0.6-
0.8
0.6
0.4-
0.6
return
0.2
0.0
-02
0.0
-0.4
0.2
mere
niter
heist
materr
ded genaal
conrert
chatto
Eiprati
plinder
storpili cowttyer
0.0-
PPO
DON
-02-
weo task ID
0.0
game
steps
1e6
steps
(a) Meta Analysis
(b) Singleton Experiments
(c) DQN with Task ID
Figure 15: Investigations of other potential sources of poor generalization in RI..
G Additional Analysis
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 RL algorithms, particularly value-based ones.
G.1 Alternative hypothesis
The second hypothesis for the superior performance of UCB relative Greedy is that better explo- ration improves sample efficiency resulting in better training performance, which can, in turn, lead to better test performance. Indeed, even the simple tabular MDPs from Figures 2 and 11, UCB converges faster on the training MDP. Here, we study this hypothesis in more detail. Note that this hypothesis and the one proposed in Section here assumes the model readily learns generalizable representations. whereas the first hypothesis applies even if the agent has perfect state representation (e.g., tabular). In both cases, adapting uncertainty-drive exploration to deep RL on complex CMDPs has additional challenges such as representation learning and unobserved contexts.
be simulhaneously true, but the hypothesis
Optimization vs. Generalization. supervised learning, deep neural networks have good general- ization, meaning that they perform similarly well both train and test data. However, in RL, the main bottleneck remains optimization rather than generalization (which also boils down to sample complexity). For example in Proegen, the state-of-the-art training performance is far below the theoretical optimum [22, 90], suggesting that current RL algorithms do not yet "overfit" the train- ing distribution. To verify this hypothesis, we first conduct a aatllass of existing algorithms. We choose 7 representative algorithms (Appendix F) and measure whether better performance on training MDPs translates into better performance on test MDPs. More specifically, we compute the Spearman's 's rank correlation, T., between the training and test performance of the algorithms, for each Proegen environment (Figure 15a). For many Procgen environments, training performance is strongly correlated with test performance, suggesting that further improvements in optimization could also result in better generalization.
In Figure 16a and Figure 16b, we show the scatter plot of the training performace and test performance of each algorithm games with positive correlation and games with negative correlation (each color represents a different game). For games with positive correlation, we see that the correlation almost linear with slope of 1 (i.e., almost no generalization gap), which suggests that for these games, we can expect good generalization if the training performance is good. On the games with negative correlation, it less clear that improving training performance will necessarily lead better generalization or more severe overfitting. Most of the games with negative correlations are either saturated terms of performance 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 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. would interesting to understand what structural properties of these CMDPs are responsible for this difference in overfitting, but this is out of the scope of the current work.
30

---

Another interesting thing to notice is that most of the outliers in the two plots are PLR [48] which 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 uniform distribution over the training environments.
G.2 Other Sources of Poor Generalization
Sample efficiency on a single environment. 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 MMainn If that is the case, the suboptimality can accumulate across all the training MDPs, resulting in large gap between these types of methods when trained in all the environments in Munin To test this hypothesis, we run both DQN and PPO on a single environment from each of the Proegen games. As shown in Figure 1 15,, we that there is not significant gap between the average performances of DQN and PPO on single environments (across seeds for 10 environments which both methods reach scores above 0.5). In fact, DQN is 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 [35]. Policy optimization methods that use trajectory-wise 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. [65]. 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 15b, we see that both methods do well on the singleton environments, so with task IDs, the algorithms should in theory be able to do as well as 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 as input the final layers of the Q-network. Since there no semantic relationship between discrete task IDs, we do not expect this to improve generalization performance, but, surprisingly, we found that does not improve training performance either (Figure 15c). This suggests that partial observability may not be the main problem in such environments as far as training is concerned. Note that this issue related to but not the same as the aforementioned oivrffiinnni issue. Having task ID means that the agent can have an 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 expected return at each state. As discussed in [90], 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 [90]. However, this approach cannot be applied to aalue-based methods since the policy is directly defined by the Q-function. An alternative solution distributional RL [10] which learns a distribution (rather than a point estimate) over all possible Q-values for a given state. This models the aleutoric uncertainty resulting from the unobserved contexts CMDPs (i.e., not knowing which MDP the agent interacting with), thus being able to account for one state having different potential values depending on the MDP. Our method uses distributional RL to mitigate the problem of value overfitting. As seen in Figure 5, while Icarning a value distribution leads to better results than learning only point estimate (i.e., QR-DQN > DQN), the largest gains are due to using a more sophisticated exploration method (i.e., QR-DQN + UCB + TNN ((a) DQN).
Our analysis indicates that sample efficiency and partial observability are not the main reasons behind the poor generalization of value-basee methods. In contrast, value overfitting is indeed problem, but it can be alleviated with distributional RL. Nevertheless, effective exploration proves to be key factor in improving the train and test performance of value-based based approaches in CMDPs. Our results also suggest that better optimization on the training environment can lead to better generalization to new environments, making it promising research direction.
31

---

Train return for with negalive correlation
Train
o² vs. Cggi test performance
0.25
0.8
0.20
M
0.6
tost
0.4
0e5t
0.4
0.15
0.10
0.05
Anmm
0.0
0.00
0.2
0.0
0.2
train
0.4
066
0.8
1.0
-0.2
train
0.4
0.6
0.8
1.0
-0.05
0.0
0.5
3.0
1.5
2.0
1e7
bosright
despetall
heid
niver
(9
mubber
maae
(c) Comparing different uncertain- ties (mean test min-max normal- ized scores).
(a) Games with positive train-test correlation.
(b) Games with negative train-test correlation.
Figure 16: Additional empirical results.
H
Uncertainty
Central to our method is the estimation of uncertainty in Q-values. We found that what kind of uncertainty model and how 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 uncertainty is useful for exploration since they represent the reducible uncertainty. This is a longstanding view shared by many different sub-field of machine learning [74, 98]. The aleatoric uncertainty on the other hand 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 Atari, the main source
uncertainty (i.e., variance) contains both epistemic and aleatoric uncertainty. of aleatoric uncertainty comes from the sticky action which has a relatively simple structure. However, in CMDPs such Procgen, the aleatoric uncertainty coming from having different environments is much more complicated and possibly multimodal, making it important to model them explicitly. Empirically, using aleatoric uncertainty to explore detrimental to the test performance which may partially explain DQN+UCB's extremely poor performance (Figure 5), even compared to DQN+s- greedy (i.e., uncertainty estimation based on ensemble of non-distributional DQNs contains both types of uncertainties). To further test this hypothesis, we run QR-DQN+UCB with both aleatoric and epistemic uncertainty and version with only epistemic uncertainty (Figure 16c). 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 5 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 Tesla V100 (note that this only occurs at the training time since test time all heads share the same feature extractor):
EDE ensemble members):
0.3411 seconds per algorithm step
EDE ensemble members):
0.2331 seconds per algorithm step
QR-DQN:
0.1356 seconds per algorithm step
In Figure 17b, we found that using 3 ensemble heads achieves comparable performance. On the other hand, using deep ensemble seems to be crucial for a good estimation of epistemic uncertainty as UA-DQN [20], which uses MAP sampling with 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 exciting future research direction.
32

---

I Architecture and Hyperparameters
1.1 Procgen
Architecture. We use standard ResNet [42] feature extractor ( from IMPALA [27], which was adopted by Ehrenberg et [26] for value-based methods. The feature maps are flattened = 2048)
and processed by M copies of 2-layer MLPs with the same architecture to output M [A]] x N values (for other models, this shape is changed accordingly for the appropriate output dimension). |A| = 15 for all Procgen games. Each MLP conisists of two liner layers (both of dimension 512) separated by a ReLU activation; the first layer is g()) R 7 1R61 and the second layer is g RR512 IR x The overall archetecture
g ʃ - ( ReLU (() Flatten
The hyperparameters for the main method is shown in Table and are directly adapted from Ehrenberg et al. [26] other than those new algorithmic components.
For hyperparameters specific to EDE, we search over the following values on 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.
Name
Value
number of training cnvs proegen distribution mode minibatch size
200
easy
512
number of parallel workers (K)
64
replay buffer size discount factor (n)
10
0.99
n-step return frame stack
3
dueling network
no
target network update frequency initial random experience collection steps
32000 (algo step)
2000 x K (env step)
total number of steps update per algo step env step per algo step
25 10 (env step)
Adam learning rate Adam epsilon
2.5 x 10
1.5 10-4
prioritized experience replay
no
number of quantiles (N)
number of ensemble heads (M)
5
Huber loss & gradient clip norm
1
10
(UCB)
30
X (TEE)
0.6
a (TEE)
Table 4: Hyperparameters for the main Procgen experiments. One algorithm step (algo step) can have multiple environment steps (env step) due the distributed experience collection.
Points of comparisons. We try to keep the hyperparameters the same as much as possible and only make necessary changes to keep the comparison fair:
33

---

For base DQN and QR-DQN, we use E-greedy exploration with decaying Specifically, at every algorithmic step, the epsilon is computed with the following equation:
2000
<tt) min 0.1 + 0.9
8000
For Bootstrapped DQN [79], we use bootstrapping to train the model and uniformly sample value head for each actor at the beginning of each episode that used for the entire episode.
We use 5 ensemble members just like the other methods.
For UCB, we use the estimated epistemic uncertainty when combined with QR-DQN and the standard variance when combined with DQN [19] and +p 30.
For TEE without UCB, we use s-greedy exploration without decaying for both QR-DQN and DQN, and EDE ith the same and a.
For ez-greedy, we use the same decay schedule in combination with 10000 and J = 2.0 which are the best performing hyperparameter values from Dabney et al. [24]. Each parallel worker has its own ez-greedy exploration process.
For NoisyNet, modify the base fully-connected layer with the noisy implementation [32] with = 0.5.
I.2 Crafter
For Crafter, the default architecture and hyperparameters that used for the Rainbow experiments in Hafner [37] with the following changes:
Change the distributional RL component from C51 QR-DQN with appropriate architecture
change and use EDE for exploration
Change the minibatch size from 32 to 64 to 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 the proposed method. For EDE since the default configuration from Hafner [37] is single-thread, we use Thompson sampling instead of UCB and do not use TEE.
For sp, searched over (0.5, 1.0, 2.0} and picked the best value.
34

---

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
ducling 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
Adam learning rate
6.25 X 10
Adam epsilon
1.5 10
prioritized experience replay
no
number of quantiles (N)
200
number ensemble heads (M)
5
Huber loss x gradient clip norm
1
10
(Thompson sampling)
0.5
Table 5: Hyperparameters for the final Crafter experiments.
J Sensitivity Study
In this section, we study the sensitivity of the test performance (i.e., aggregated min-max normalized scores) to various hyperparameters on subset of games. First, without TEE, study the sensitivity to different o (UCB exploration coefficients). different M (number of ensemble members), and whether the feature extractor is trained with gradients from each ensemble member.
For P and M, we run the algorithm on: bossfight, ninja, plunder, starpilot, caveflyer, jumper, bigfish, leaper for 1 seed and report the aggregated min-max normal- ized scores in Figure 1 and Figure 17b. We observe that the algorithm slightly sensitive the choice of but not sensitive at all to M. This is encouraging since a large amount of computational overhead comes from the ensemble. Note that while smaller ;p performs the best here, we use larger value (yp 30) in combination with TEE.
In Figure 17c, show the effect of only training the feature extractor using the gradient from one member of the ensemble at every iteration. The results are computed on: ninja, plunder, jumper, caveflyer, bigfish, leaper, climber for seed. We observe that always training the feature extractor leads to lower performance, corroborating our intuition that the feature extractor should be trained the same speed as the individual ensemble members.
In Figure 18, we study the performance under different hyperparameter values TEE. We use fixed M = 5 and y = 30 and vary the values of either or while holding the other fixed. We observe no significant difference across these, suggesting that the algorithm is robust to the values of and X.
K Hardware
The experiments are conducted on 2080 and V100 and take approximately 250 GPU days.
35

---

Varying
Varying
Always
0.45
0.40
0.30
0.40
0.35
0.35
0.25
Maky
0.30
0.25
0.25
return
0.20
0.20
0.15
0.15
10
0.10
0.10
=
30
0.10
M
True
0.05
50
0.05
M
5
Fabse
0.00
1.5
1.0 steps
1.5
steps
1e7
1e7
steps
1e7
(a) ablation.
(b) M ablation.
(c) Feature extractor ablation.
Figure 17: Aggregated min-max normalized test scores y (for fixed M training feature extractor for all value heads), M (for fixed y = 50 and training feature extractor for all value heads), and whether feature extractor is trained with all value head (for fixed yp = 50 and M = 3).
Varying
Varying
0.4
0.3
0.3
0.2
0.14
À 0.5
7
À - 0.6
0.0
=
0.0
A= 0.7
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1'5
2.0
steps
1e7
steps
1e7
Figure 18: Aggregated min-max normalized test scores for À (for fixed a = 7) and a (for fixed À = 0.6) on games: bossfight, climber, plunder, starpilot.
36


Reviews:
[{"review_id": "wGGtcrrpuU", "reviewer": "Reviewer_tLzS", "summary": "The paper focuses on the importance of exploration in the generalizability in contextual MDPs. The proposed method is built on QR-DQN where the epistemic uncertainty can be separated from aleatoric uncertainty via ensemble. The epistemic uncertainty is then used in a UCB manner to promote exploration. The resulting algorithm is tested in two benchmarks and shown significant improvement over previous baselines.", "questions": "The authors emphasize separating epistemic uncertainty from the overall uncertainty, which motivates them to build their framework on top of QR-DQN, from my understanding. This decision makes sense theoretically. However, there is a lack of empirical study to show the actual significance. If this separation is unimportant, this method can potentially be applied to a wider range of algorithms.\n\nMinor typos / mistakes:\n- No legend in Figure 2(b) and 2\\(c\\).", "limitations": "The authors adequately addressed the limitations and potential negative societal impact of their work.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The paper is overall well-organized and easy to follow.\n- The proposed method is novel, well-motivated and empirically strong.\n", "weaknesses": "- The motivating environment introduced in Sec 3 doesn't seem exactly appropriate for the following reasons:\n    - According to the definition of CMDP, the distribution of MDPs during training and testing should be the same, which is not the case here.\n    - Since only the starting distribution is different, the value function is the same in both MDPs. Combined with above mentioned different distribution during training and testing, this example seems too obviously engineered towards promoting visiting more states during training.\n    - Claiming UCB is more effective in exploration compared to $\\epsilon$-greedy seems dry. It would be helpful to at least compare the training stage state-action visiting frequency for both methods.\n- Same as the motivating example, it would be helpful to show either qualitatively or quantitatively the difference in the extent of exploration among different methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the thoughtful review and support for the paper! We were glad to hear that you found our paper \"well-organized and easy to follow\" and our proposed method \"novel, well-motivated, and empirically strong\". We hope our answers below will address your remaining questions.\n\n\n>The motivating environment introduced in Sec 3\n>\nWe agree that the claim that exploration can help generalization to new environments (with different initial states or dynamics) may seem obvious *after* we have put them into the context. In fact, we think that it's good that it's so intuitive since the toy example is meant to provide intuition. But of course, the motivating example is not meant to and cannot capture everything that happens in deep RL. \n\nAs far as we know, the observation that exploration affects generalization in RL, has not been explicitly discussed in any prior work, so we believe it is still a valuable contribution to explicitly emphasize this, provide intuition using simplified settings and controlled experiments, and validate this hypothesis in more challenging environments like Procgen and Crafter, which our paper does.  \n\n> The distribution of MDPs during training and testing should be the same\n> \nRegarding the distribution of MDPs, the distributions can be easily made the same: the full distribution contains the two starting states with equal probability and the training set contains a single sample from the distribution which is the environment starting from the top left corner. In this case, the population performance would be just the average of the two curves (blue and orange) and UCB would still perform much better since the performance on training MDP (blue) is about the same for both methods. We'd be happy to add a discussion on this point in the appendix. \n\n\n>Importance of separation of different uncertainties \n>\nOur paper contains multiple experiments that support the importance of separating the two types of uncertainties. First, in Figure 16c from the appendix, we directly compare aggregated (both aleatoric and epistemic) with only epistemic uncertainties showing that using only the epistemic uncertainty to guide exploration is significantly better. Second, Figure 5 from the main paper shows that the ensemble of DQN (which contains both uncertainties) performs significantly worse. \n\n\n> Extent of exploration\n>\nIt is difficult to visualize exploration in a procedurally generated environment because there could be a potentially infinite number of distinct states. Moreover, we probably do not want the policy to visit all of these states since that would be extremely inefficient.\n"}, {"review_id": "YGcwd0nOGE", "reviewer": "Reviewer_PttU", "summary": "The paper proposes that effective exploration is important for generalization and shares a value-based method which gets good generalization performance on procgen and outperforms Rainbow on crafter.", "questions": "You say this is the first value-based method to get good performance on procgen. The Q-learning baseline from [1] simply adds learning a value-equivalent model as an auxiliary task (but does not use it for planning) and gets pretty good performance. Can you discuss the difference and how your method does better than theres? Is it still true that yours is the first value-based method to get this type of performance on procgen? MuZero also does quite well and is value-based.\n\nThe authors compare to LSTM-SPCNN, note that LSTM-SPCNN is a \"specialized architecture\"  but LSTM-SPCNN is just a visual-transformer like architecture with PPO? So it seems fairly general. Why is the comparison is there?\n\nOf the *many* exploration algorithms that are available to compare against, why did you choose these particular ones?\n\nWhile I like the idea of improving generalization by improving exploration, do you provide evidence that bad exploration is indeed the reason for poor generalization? \n\nIt's hard to judge this paper without answers to the above questions.\n\nWhy did you choose crafter over other domains? I could imagine using the mini-grid environments which seem like more relevant environments since you can explicitly setup experiments that test both exploration and generalization.\n\n[1]: Procedural Generalization by Planning with Self-Supervised World Models", "limitations": "There is a single line on limitations (expensive due to ensembling) which is insufficient in my opinion. \n\nWill this method improve exploration & generalization in 3D environments like Mine Craft? How about mazes like mini-grid? Why or why not?\n\nAre there different conditions under which you expect the epistemic uncertainty estimate to be more-or-less informative?", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "originality: The idea to leverage improved exploration to improve generalization is novel/seldomly explicitly tackled.\nquality: The evaluation is very high-quality.\nclarity: The paper is fairly clear. I like starting with a tabular MDP to motivate that exploration improves generalization. The methods section is dense but clear.", "weaknesses": "Figure 1 could be improved to show some abstraction for how the type of exploration they propose can improve generalization. \n\nThe method is only evaluated on two domains. It could be improved if evaluated on more domains.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you for taking the time to review our paper and the detailed questions! We were glad to hear you found the idea novel and the evaluation of very high quality. We hope our answers below will address your remaining concerns. \n\n> Figure 1 could be improved to show some abstraction \n\nFigure 1 is meant as an accessible illustration of one possible intuition/interpretation (which other reviewers have found effective), rather than a full description of our method. We opted for simplicity as we were concerned that illustrating uncertainties would result in an overly-complicated diagram which may confuse readers.\n\n> more domains\n> \nRespectfully, we believe this is quite an extensive evaluation of the algorithm. The procgen benchmark already contains **16 distinct games**. Since this is the standard benchmark for studying generalization in RL, we believe this choice is appropriate. The experiments and ablations for these games alone take more than 250 GPU days which are quite extensive in our opinion. On top of that, Crafter, which was designed to resemble Minecraft, adds another domain that poses very different challenges from Procgen. We believe our experiments on both Procgen and Crafter offer sufficient evidence to support our main claims and emphasize the generality of our approach.\n\n> Q1: QL agent from [1]\n\nThank you for going through the reference and pointing out this one. The QL agent from [1] performs on par with PLR which is worse than IDAAC [2]. In our experiments, our method outperforms IDAAC which suggests that we should outperform [1] as well. Further, we do not claim that exploration is the only way to improve generalization. The main difference is that the QL agent used in [1] uses additional self-supervised learning objectives. This is orthogonal to our contribution, so combining these methods may lead to further improvements. Also note that [1] uses a different setup than we do. Our work adopts the standard setup with 25M training steps on 200 levels from the ``Easy`` mode which allows for direct comparisons to existing methods, whereas [1] uses 30M training steps on 500 levels from the ``Hard`` mode which is not commonly used in the literature. \n\n> Q2: LSTM-SPCNN \n\nLSTM-SPCNN is not a vision transformer (the transformer-based architecture does worse than LSTM-SPCNN). It is a CNN that does not reduce the spatial dimension via striding (hence the ~100x more parameters than QRDQN’s architecture) because striding would lose information about the backpack in the bottom of the frame in crafter which is important for doing well, so LSTM-SPCNN is specialized for crafter.\n\n> Q3: the choice of baselines\n\nPlease see our general response at the top.\n\n> Q4 evidence that bad exploration is indeed the reason for poor generalization\n\nIn Figure 16c of the appendix, we showed that if one targets both aleatoric and epistemic uncertainty in EDE, the performance is significantly worse, suggesting that the type of exploration indeed affects performance. Nonetheless, there can be many reasons for poor generalization so even a good exploration policy will not always result in better performance if these other components are not appropriate (e.g., poor optimization).\n\n> Q5 Crafter\n\nCrafter is designed to be an extremely challenging environment that requires, among other things, strong generalization and exploration, so we believe it is very relevant for our study. On the other hand, there is no standard benchmark for generalization in mini-grid that everyone uses. In comparison, Crafter is more challenging and has a standardized benchmark to allow for easy comparison and shows how far we can get just by doing better exploration.\n\n> Will this method improve exploration & generalization in 3D environments like Mine Craft? How about mazes like mini-grid? Why or why not?\n\nWe believe that it would help in environments like Minecraft since the structure of Crafter is designed to resemble Minecraft (e.g., crafting tree, resource collection and survivial). The 3D part of Minecraft is more challenging for reasons that are less related to exploration (e.g., representation learning). For mazes, it is less clear whether exploration during training is the main bottleneck (it could still be useful during test time). In procgen, there is a maze game. In this setting, we observed that EDE did not improve QR-DQN significantly. We hypothesize that this is due to the fact that mazes require good long-term planning which requires better representation rather than exploration. We commented on this in the caption of Figure 13 but we would be happy to add more discussion.\n\n> Are there different conditions under which you expect the epistemic uncertainty estimate to be more or less informative?\n\nCMDPs are a large and underexplored class of problems so we do not have an exact theoretical characterization right now about when we can expect epistemic uncertainty to help, but it should not hurt in general (it improves over QR-DQN in almost every game). We would be happy to provide some speculations if you think that's helpful. \n\nFor different initial states, good exploration should generally help. We argue this analytically in Appendix A. For different dynamics, the story is more complicated. “The change in transition is not highly correlated with the action” could be a good condition, but we do not have a precise theoretical characterization. Finally, for different rewards, very strict assumptions are needed to ensure that generalization is possible. Please see [1] for more analysis.\n\n**References**\n\n[2] Decoupling Value and Policy for Generalization in Reinforcement Learning. Raileanu et al.\n\n\n**Conclusion**\n\nWe hope the above clarifications and proposed revision have adequately addressed all your questions regarding the paper. If there are no outstanding concerns, will you consider raising your score to provide strong support for our paper? \n"}, {"review_id": "Eclx7eX9d8", "reviewer": "Reviewer_saZD", "summary": "This paper introduces a method called EDE (Ensemble Distributional Exploration) that promotes the exploration of states with high epistemic uncertainty through an ensemble of Q-value distributions. The authors evaluate EDE and compare it to several baselines on Procgen and Crafter.", "questions": "1. Could you provide some insight into the decision to use Thompson Sampling in the Crafter experiment? Specifically, why was this approach chosen over UCB combined with TEE? Understanding the rationale behind this choice would be beneficial.", "limitations": "The authors adequately addressed the limitations", "rating": 5, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "**Originality**\nThe paper's originality lies in its approach to exploration in reinforcement learning. The authors introduce the concept of using ensembles for uncertainty estimation to guide exploration.\n\n**Quality**\nSee Weaknesses Section.\n\n**Clarity**\nThe paper is well-written and organized, making it easy for readers to follow the authors' thought process and understand the methodology and results.\n\n**Significance**\nBy demonstrating that ensembles for uncertainty estimation can effectively guide exploration, the authors have opened up new possibilities for exploration in reinforcement learning. The paper's focus on improving an existing method for better exploration is particularly relevant in today's context, where there is a growing emphasis on improving the efficiency and effectiveness of reinforcement learning.", "weaknesses": "One potential weakness is the increased computational cost associated with the method. The use of an ensemble to guide exploration, while useful, also makes the method more expensive to run. This could limit its applicability in scenarios where computational resources are constrained.\n\nIn Figure 5, it would be beneficial to add QR-DQN with Thompson Sampling in the ablation.\nFurthermore, in Table 1 it would be more informative to evaluate QR-DQN, and EDE (with UCB and TEE instead of Thompson Sampling).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper and for your valuable feedback! We were glad to hear you think our paper \"opens up new possibilities\" and is \"particularly relevant in today's context\". \n\n> computational cost\n>\nWe agree that this is an important topic. Estimating uncertainty in an efficient manner is an active area of research in machine learning. In supervised learning, methods such as epinet [1] could significantly reduce the computational cost while providing good uncertainty estimates, so we believe it could also be useful in RL. For this work, we chose to use deep ensembles because it’s a commonly used technique that is robust and easy to implement.\n\nFurthermore, the cost of ensembles occurs mostly at training time, and at test time, the additional cost is much smaller due to the shared feature extractor. Since ensembles can achieve much better results, this cost may be justified depending on the application as it only adds a very small additional cost at inference time.\n\n> Thompson sampling vs UCB\n> \nAs mentioned in the paper, for Crafter, the environment is not naturally parallelized so we cannot easily use TEE. In this setting without TEE, we observed that TS outperforms UCB. Since our computational budget was limited, we chose to focus on the method that does better in each setting. In the provided source code, both settings can be used for those interested. For procgen, in our preliminary experiment, QR-DQN+TS+TEE did not perform as well as QR-DQN+UCB+TEE. Since evaluating the entire Procgen suite is expensive (~250 GPU days for all the experiments), we did not conduct the full experiment for TS and chose to conduct other ablation and baselines. For bandits, UCB and TS are actually very similar algorithms (in spirit at least), and they have similar regret but their actual performance can be problem-dependent [2] so it's hard to say one will always outperform the other. We will make this more clear in the revision.\n\n**References**\n\n[1] Epistemic Neural Networks. Osband et al. \n\n[2] An Empirical Evaluation of Thompson Sampling. Chapelle et al.\n\n\n**Conclusion**\n\nWe hope our answers above have addressed your questions sufficiently well to alleviate your concerns regarding the paper. Please let us know if there is anything standing between us and a strong recommendation of acceptance."}, {"review_id": "vwH9BvbHZw", "reviewer": "Reviewer_zePN", "summary": "This paper proposes an exploration method for value-based RL for contextual MDPs (CMDP) motivated by the idea that good generalization in RL requires attention to RL specific problems such as exploration. The method uses an ensemble of quantized Q-functions to estimate the epistemic uncertainty about the value function. Then uncertainty estimate is used to conduct upper confidence bound (UCB) exploration based on the estimated distribution of the Q-values. Additionally, the algorithm samples more diverse data by using different exploration weights for the different actors. The proposed method achieves better performance than previous value-based methods on procgen, where it is also competitive with policy optimization-based methods. On crafter, the proposed method achieves performance on par with the state of the art method with much smaller network.\n\n## Acknowledgment\n\nI have read the rebuttal and the following discussion and updated my review accordingly. Concerns raised by other reviewers convinced me to lower my rating by one grade.", "questions": "- In CMDP settings, where the training MDP distribution is the same as the testing one, would exploration still be the most important driver of generalization performance?", "limitations": "The paper discusses relevant limitations to the proposed algorithm.", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "## Originality\nThis paper presents a well executed study on generalization for value-based RL algorithms in CMDPs. Other studies of generalization in CMDPs have been presented in previous work, but the focus on value-based RL and exploration seems like a relatively unexplored area. The proposed method is a novel combination of existing ideas from previous work. It is well motivated by the presented analysis and achieves state-of-the-art results for value-based RL in challenging CMDPs, so presenting this combination seems like a potentially important contribution to the RL community.\n\n## Quality\nThe study of exploration as a driver for generalization is of high quality. It starts from illustrative examples and applies that intuition to more complicated cases. The paper hits a good balance in what level of detail to cover the background information. The proposed method is motivated by the analysis and each of the design decisions are carefully analysed in the empirical section. The empirical results are exceptionally thorough and well presented.\n\n## Clarity\nThe paper is well written throughout. It explains its ideas at a helpful level of detail. The illustrations and figures are clear and well designed. The results for procgen are presented using a style that has been promoted as the best practice in communicating RL results.\n\n## Significance\nValue-based RL is an interesting topic and for various tasks, value-based RL algorithms are found to be state-of-the-art. Therefore, contributions that improve the generalization ability of value-based RL algorithms for CMDPs can be significant contributions to the RL community.", "weaknesses": "## Discussion on exploration-exploitation tradeoff\n- The paper proposes exploration as a critical component for achieving generalization for CMDPs. Choosing a more exploratory policy always introduces an exploration-exploitation tradeoff, yet such tradeoff is not discussed in the paper.\n- The illustrative examples assume a distribution shift between training and test time MDPs. If there was no distribution shift, using a more exploratory policy would result in an exploration-exploitation trade-off that needs to be balanced. It would be good to comment on how the presented intuition is still relevant when there is no distribution shift.\n\n## Other weaknesses\n- The results for the Crafter are presented with less detail than the Procgen ones. I would have appreciated learning curves and the median/IQM/... plots for Crafter as well.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the generous review and strong support of our paper! We were excited to hear you found our paper to be an \"important contribution to the RL community\" and the empirical results to be \"exceptionally thorough and well presented\".\n\n>exploration-exploitation tradeoff \n>\nIndeed this is a fundamental problem in RL. In our setting, we don’t think the tradeoff is significantly different from other RL settings. Given a fixed time budget, one would need to choose the extent to which you want to explore before settling on a reasonable solution. Both over-exploration and over-exploitation would result in poor generalization in addition to poor training performance. This is why we chose optimism-based exploration which tries to balance between the two automatically. In EDE, this is controlled by the hyperparameter $\\varphi$ (w/o TEE) or $\\alpha$ and $\\lambda$. That being said, in practice, we observed that EDE is not particularly sensitive to these hyperparameters. You can find this ablation in Figures 17 and 18 in the appendix. We will also add more discussion about this to the paper.\n\n> IQM for crafter\n>\nThe baseline results are taken from the original Crafter and other papers that use Crafter which do not use IQM so we wouldn't be able to compare to them.\n\n> Distribution of MDPs (Q1)\n>\nSuppose the full distribution contains the two starting states with equal probability and the training MDP is a single sample from the environment starting from the top left corner, the population performance would be just the average of the two curves (blue and orange) and UCB would still perform much better because the training performance (blue curve) is similar for both methods. In this case, there is no distribution shift and the conclusion still stays the same.\n\n**Conclusion**\n\nWe thank you again for your feedback and strong support of the paper. We hope our answers have addressed all your questions but let us know if you have any outstanding concerns. \n"}, {"review_id": "aByHm1EwNP", "reviewer": "Reviewer_dHNU", "summary": "The paper motivates the importance of exploration for generalization in contextual MDP with a tabular example, where the context is either the starting state or an uncontrollable random deviations in the transition model. The authors introduce an exploration method for distributional DQN, called EDE, which estimates the epistemic uncertainty of the network as the variance of quantiles between ensemble members. They use this uncertainty in an UCB-style exploration algorithm with varying exploration coefficients, called TEE. EDE/TEE is evaluated on the ProcGen baselines and on the Crafter environment. Results indicate that the average median over test MDP of all environments is higher than baselines with epsilon greedy exploration.  \n", "questions": "1. Can you specify precisely under which conditions exploration is useful for tabular CMDP?\n2. What is the connection between the first and the second part of the paper? Wouldn't every exploration algorithm help this type of CMDP generalization?\n3. Which type of exploration is the right one for CMDP generalization (i.e. for your answer of Q1)? \n4. Which type of exploration did the PPO and DAAC baselines use? \n5. How does the generalization gaps in Figure 13 look like?\n", "limitations": "It is unclear which exploration is best for CMDPs and whether the performance of EDE is significantly better than the baselines.\n", "rating": 4, "confidence": 5, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "The paper is very well written and contains an enormous amount of references, for which the authors are to be commended! The insight that exploration has a larger role for generalization in CMDP than just finding the path to the largest reward (as it is in single MDP) is very interesting and well made. The presented exploration method makes sense and outperforms epsilon-greedy baselines in some of the ProcGen baselines and probably also in Crafter. ", "weaknesses": "The biggest issue with the paper is that its two parts (exploration for CMDP and EDE) do not really have much to do with each other. EDE/TEE is a general exploration algorithm, which does not use any of the insights from the first part to explore specifically for CMDP. Furthermore, the first part could be more precise for which CMDP exploration is helpful (in the presented analogy), and the second part is not all that novel. Finally, the experiments are missing important baselines and the results of EDE do not seem to be significantly better. To clarify these criticisms:\n- The insight given in Figure 1 is very nice, but *only* applies to CMDPs where the context is different initial states. The example given in Figure 11 goes beyond this restriction, but it remains unclear what exactly the class of CMDP that can be solved by exploration is. It appears that either the context is fully observable in the state, in which case the tabular analogy breaks down, or the assumption that all $\\mathcal M \\sim q_{\\mathcal M}$ have the same optimal policy $\\pi^*$ does not apply to most CMDP, including the ProcGen benchmarks. It also does not help that the authors phrase uncertainty about an unobserved context as aleatoric uncertainty (which is a reasonable classification), as the existence of $\\pi^*$ would make the distinction effectively meaningless.\n- EDE is based on the ensemble disagreement (variance) of the output of a distributional RL algorithm (QR-DQN). While the reviewer could not find a paper that does this *exact* combination, it is a very straight forward combination of many papers that pair ensemble-disagreement for exploration (many cited by the authors) with quantile regression DQN. While not in the context of exploration, other papers have already made this combination for risk-avoiding RL (Eriksson et al., 2022; Hoel et al., 2023). In its defense, EDE/TEE uses some non-standard components, like every head being trained by their own minibatch, instead of subsets of the same minibatch as in [78, 80, 81], and different exploration parameters for different actors. These are not new, though, and the overall novelty of EDE/TEE seems low.\n- What is the connection between the first and the second part of the paper? It seems that every exploration algorithm would help this type of CMDP generalization, so why did the author specifically propose EDE? This seems to be a missed opportunity, as the first part could be strengthened by a clearer definition of exploration, too. For example, most exploration methods (including TEE) optimize a trade-off between value and epistemic uncertainty, whereas others like task-agnostic exploration [87, 122] only maximize entropy or minimize uncertainty. Which type of exploration is the right one for CMDP generalization? \n- Does EDE ever use the aleatoric uncertainty $\\sigma_\\text{ale}$ in any meaningful way? Is there any significant difference between EDE and a DQN ensemble trained in the same manner? The reviewer doubts that the variance of a DQN ensemble is more tainted by aleatoric uncertainty than one made of QR-DQN. Figure 5 seems to suggest this, but is DQN-UCB/TEE also been trained in the exact same way as EDE? \n- It is commendable that the authors ran a range of comparisons in Figure 5, but it is surprising that in Figure 5 EDE seems to have a significant advantage in all quantities, but against PPO/IDAAC there seems to be no advantage in mean and optimality gap. This could be accurate, but the reviewer still recommends the authors to check their code again. Which type of exploration did the PPO and DAAC baselines use here exactly? The reviewer also would have liked a comparison with intrinsic reward methods or UBE [82], as TEE only takes local uncertainty into account, which has proven to be suboptimal for exploration. \n- Finally, the conclusions drawn from the presented results appear questionable. Improving over PPO (assumingly without advanced exploration) w.r.t. median but not mean is suspicious and looking at the learning curves in Figure 13 shows that EDE significantly beats PPO in 5 environments, but significantly looses against PPO in another 5 environments (there seems to be no significant difference in the other 6 environments). This does not look like a significant improvement by EDE! Moreover, the story of the paper is that exploration helps generalization. Can you also plot the generalization gaps for all algorithms/environments in Figure 13?\n\n**Additional References**\n- H. Eriksson, D. Basu, M. Alibeigi, and C. Dimitrakakis. Sentinel: taming uncertainty with ensemble based distributional reinforcement learning. In Uncertainty in Artificial Intelligence, pages 631-640. PMLR, 2022.\n- C.-J. Hoel, K. Wolff, and L. Laine. Ensemble quantile networks: Uncertainty-aware reinforcement learning with applications in autonomous driving. IEEE Transactions on Intelligent Transportation Systems, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for the detailed review and insightful questions! We were glad you found our insight interesting and well made and our literature review extensive. \n\n> Connection between two parts (Q2)\n\nAs we explain at the beginning of Sec 4, the first part is meant to show that good exploration can help generalization even in tabular MDPs w/o deep learning. The second part shows that applying this insight to deep RL requires a more careful design of the exploration strategy since empirically not all exploration methods help in more complex environments.\n\n> novelty\n\nWe believe that we have been straightforward with the fact that many components of EDE have been studied in prior works. The main claim we are making is that they can be used to improve generalization significantly. We believe this claim is novel and substantiated by our empirical results. Please see the common response for more details on the core contribution of our paper.\n\n> Intrinsic motivation and UBE\n\nPlease see our general response. We believe that the baselines are representative of general-purpose exploration methods. For example, UBE belongs to the family of Thomspon sampling algorithms following Bootstrapped DQN.\n\n> Fig 1 and conditions on CMDPs (Q1)\n\nThe goal of Fig 1 is a high-level intuition on how exploration can help generalization in RL. Characterizing CMDPs (which is a very large class of problems) is an open area of research. Given the complexity of the problem, we believe it is beyond the scope of our paper. For different initial states, good exploration should generally help. We argue this analytically in Appendix A. For different dynamics, the story is more complicated. “The change in transition is not highly correlated with the action” could be a good condition, but we do not have a precise theoretical characterization. Finally, for different rewards, very strict assumptions are needed to ensure that generalization is possible. Please see [1] for more analysis.\n\n> existence of $\\pi^\\star$\n\nThank you for pointing this out. The precise statement [1] is that $\\pi^\\star$ achieves nearly optimal performance (up to $\\alpha$) on all MDP: $V^{s_0}\\_{\\mathcal{M}}(\\pi^\\star) \\geq \\max\\_\\pi  V\\_\\mathcal{M}^{s_0}(\\pi) - \\alpha, \\forall \\mathcal{M}$ i.e., we assume there exists a policy that does reasonably well on all MDPs of interest without re-training on the unseen MDP. This assumption only ensures that the problem is solvable, but the RL algorithm does not rely on this assumption in any way (i.e., it will find a bad solution if no good solution exists). We will clarify this in the revision.\n\n> Use of aleatoric uncertainty\n\nEDE uses the aleatoric uncertainty by actively ignoring it. This is distinct from risk-sensitive RL which aims to avoid undesirable outcomes by avoiding actions with high aleatoric uncertainty. This is an important topic but not one that we focus on. Future works could combine both (e.g., avoiding high aleatoric uncertainty during test time for better test performance in environments where bad actions can have large penalties).\n\n> EDE vs DQN ensemble\n\nBoth DQN ensemble and QRDQN ensemble are not trained on separate minibatches, so DQN+UCB and DQN+TEE are directly comparable to QRDQN+UCB and QRDQN+TEE. We see that QRDQN variants improve upon DQN variants, highlighting the need for aleatoric uncertainty. Without this separation, the variance of DQN ensemble will contain both uncertainties. In Fig 16c, we show an ablation of QRDQN+UCB with both uncertainties and it does significantly worse which means the separation is important.\n\n> Comparison against PPO and checking code\n\nWe use the standard rliable [2] for the results so it's unlikely that there is a bug.\n\n> Difference in mean and median\n\nIt is well-known in statistics that the median and mean can have different ordering. We have not claimed that EDE outperforms PPO in every game and, the lower mean is because, as you have noticed, there are some games where all value-based (VB) methods perform much worse than policy-based (PB) methods (due to reasons unrelated to exploration) and mean is sensitive to extreme values.\n\nWhen there are many games, it is customary to report robust summary [2]. Research on Atari shows that rarely does a single algorithm outperform prior works in all games. We agree that it is valuable to understand why VB approaches fail in particualr games, but we don't think this contradicts our claim that the EDE outperforms PPO on Procgen overall, in the sense that the standard RL literature uses the term \"*outperform*\". We already comment on this in the caption of Fig 14 but would be happy to emphasize it.\n\n> Exploration policy of PPO/DAAC (Q4)\n\nThey use the standard Boltzmann policy to explore.\n\n> What kind of exploration (Q3)\n\nThe type of exploration we support in this paper is the classical optimism-based that balances value and uncertainty since there is a well-defined task. Please see the common response for more details. \n\n> Generalization gap (Q5)\n\nAs pointed out in the footnote of page 3, the notion of generalization we focus on is **test return after a fixed number of interactions** [3] because the generalization gap (gap) can be misleading (e.g., a random policy has 0 gap). In supervised learning, the gap is meaningful because almost all models can fit the training data perfectly, but this is not the case for RL. In terms of gap, EDE is on par with other methods.\n\n**References**\n\n[1] When Is Generalizable Reinforcement Learning Tractable? Malik et al.\n\n[2] Deep Reinforcement Learning at the Edge of the Statistical Precipice. Agarwal et al.\n\n[3] Leveraging Procedural Generation to Benchmark Reinforcement Learning. Cobbe et al.\n\n**Conclusion**\n\nWe hope our answers and proposed revision have adequately addressed your questions and concerns and that you will consider raising your score. Please let us know if there is anything preventing you from recommending acceptance and we will respond to any outstanding questions."}]

Author rebuttal:
We thank all the reviewers for taking the time to provide valuable feedback on our work. Overall, the reviewers found our paper to be clear and easy to follow, the idea of using exploration to improve generalization to be novel and well-motivated, and the experimental results to be thorough and strong. Below we clarify the main contribution of the paper and motivate our choice of baselines since multiple reviewers had questions about these. 


> Core contribution of the paper

We'd like to emphasize that **the main message of this paper is the insight that appropriate exploration can help generalization to new environments in RL**. This message is quite intuitive on the tabular MDP we constructed, but it is only relevant if it helps empirically on more complex problems such as Procgen or Crafter, which we use to validate this hypothesis. Our experiments also show that many exploration methods which were designed to improve performance on a single training environment (e.g., NoisyNet, ez-greedy, bootstrapped DQN etc.) are insufficient when applied in such settings. This motivated the design of EDE, a new exploration method that significantly improves performance in new environments. While it is true that EDE has strong empirical performance (especially for a value-based method), we think the algorithmic contribution is secondary to the conceptual message as algorithms can always be improved. Since most existing methods for Procgen do not target exploration explicitly, we think they would benefit from better exploration. We hope our work is a first step towards designing exploration methods that improve generalization in RL and that it will inspire more research on better understanding the relationship between exploration and generalization. 


> Choice of baselines (reviewers PttU and dHNU)

Exploration is a longstanding problem in RL and it is infeasible for us to evaluate all possible methods in the literature. We believe that the baselines and ablations we use in the paper are extensive and representative of general-purpose exploration methods in the literature. The estimated time for all Procgen experiments (without considering hyperparameter tuning) is more than **250 GPU days** and covers **5 popular baselines + 8 ablations**. We believe this ensures that the conclusions from our experiments are reliable.

Most of the baselines we chose are general-purpose exploration methods that aim to accelerate learning rather than solve an environment with extremely sparse reward (which is not the setting we consider here and would require additional techniques). They are widely cited and distinct from each other (e.g., bootstrapped DQN is cited 1222 times and NoisyNet is cited 568 times). They follow the principle of "optimism in the face of uncertainty", trading off return maximization with uncertainty. We will refer to them as optimism-based exploration although the exact implementation can vary. Since EDE follows the same principle to explore, we believe these are the most relevant and informative baselines to compare against. 

Several reviewers mentioned intrinsic reward methods which we assume refer to methods such as ICM [2] or RND [3]. As we explained in footnote 4, [1] finds that *existing intrinsic reward methods are no better and much less sample-efficient than optimism-based exploration methods such as NoisyNet on Atari*, except for the environments with extremely sparse reward. Note that Procgen and Crafter don't fall into this category since they generally have dense rewards and not-so-long episodes. Thus, we think it is unlikely that intrinsic reward methods would outperform our baselines (e.g., NoisyNet), so we did not include them due to the limited computation budget. However, we believe that a better understanding of how intrinsic rewards influence generalization in RL is an important and interesting research question that would be better addressed in future work. 


**References**

[1] On bonus based exploration methods in the arcade learning environment. Taiga et al. 2021.

[2] Curiosity-driven Exploration by Self-supervised Prediction. Pathak et al. 2017.

[3] Exploration by Random Network Distillation. Burda et al. 2018.

**Conclusion**

We thank you again for your feedback and support of the paper. We hope our clarifications have adequately addressed all your concerns but let us know if you have any remaining questions. 


Meta-review: The paper addresses the role of exploration in generalization for cMDP solution strategies. The paper proposes a ensemble-based method using quantized Q-funcitons to quantify epistemic and aleatoric uncertainty. The inferred epistemic uncertainty is then used to guide exploration. The method is tested on two cMDP enviornments that serve as benchmarks for generalization: Crafter and ProcGen, and outperforms the considered baselines in these environment.

There was a extended discussion on these papers, both between the authors and the reviewers and amongst the reviewers themselves.

The reviewers mostly agreed that the paper was well written, provided insight on a significant problem, and contained original ideas. Several reviewers also noted that from the results it is hard to understand the differences in exploration behavior.

There was discussion on the significance of the contribution of the paper. The contribution was rated as “2, fair” by 3 of the five reviewers. The novelty of the paper seems to be in two main aspects:

1. Introducing a new methods, composed of elements of two known approaches (quantile learning and ensemble methods) to improve generalization
2. making the conceptual connection between exploration and generalization.

The results presented by the authors on two domains are quite convincing compared to the considered benchmarks. However, the experiments could have been more convincing in several aspects, such as (1) more directly measuring exploration (2) considering more environments (3) more directly measuring generalization. As is, the experiments measure policy performance, which makes it hard to exactly understand whether (1) the performance is explained by other factors such as tuning (2) the performance is explained by improved exploration alone (3) the performance improvement is due to improved generalization.

Another point of discussion was the relationship between the example in section 3 and the method development in section 4. The authors clarified in the discussion and rebuttal that the example in section 3 is a motivating example and not part of the core contribution, and most reviewers found the example helpful if not perhaps ideally suited.
Decision: Accept (poster)