Title: On the Importance of Exploration for Generalization in Reinforcement Learning
Venue: NeurIPS 2023
Accepted: yes
Keywords: reinforcement learning

generalization

procgen

crafter
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

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 highdimensional observations. The open-sourced implementation can be found at https://github.com/facebookresearch/ede .

1 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 on unseen environments from the same family [22, 112, 75], despite being competitive in singleton Markov decision processes (MDPs) where agents are trained and tested on 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 to policy-optimization methods ( i.e., methods that learn a parameterized policy in addition to a value function) [26]. Most existing approaches for improving generalization in CMDPs have treated this challenge as a 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 in RL beyond representation learning.

Exploration can help an agent gather more information about its environment, which can improve its ability to generalize to new tasks or environments.

While this statement is 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] or meta-RL [65]), nor has there been a convincing empirical demonstration of this statement on common generalization

???

Train

Train

Test

Test

Standard

Explorative

benchmarks for RL. In this work, we show that the agent's exploration strategy is a key factor influencing generalization in contextual MDPs. First, exploration can accelerate training in deep RL, 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 1, trajectories that are suboptimal in certain MDPs may turn out to be optimal in other MDPs from the same family, so this knowledge can 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 is to learn new information about the (knowable parts of the) environment so as to reduce epistemic uncertainty . To model epistemic uncertainty (which is reducible by acquiring more data), we need to disentangle it from aleatoric uncertainty (which is 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 a type of aleatoric uncertainty that can be modeled by learning a distribution over possible values rather than a 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 RL, demonstrating a significant improvement over more naive exploration strategies. This is the first model-free value-based method to achieve state-of-the-art 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.

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 an exploration method based on minimizing the agent's epistemic uncertainty in high-dimensional CMDPs, and (iii) achieves state-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 = ( S , A , R, P, γ, µ ) , where S is the state space, A is the action space, R : S × A → [ R min , R max ] is the reward function , P : S × A × S → R ≥ 0 is the transition distribution, γ ∈ (0 , 1] is the discount factor, and µ : S → R ≥ 0 is the initial state distribution. We further denote the trajectory of an episode to be the sequence τ = ( s 0 , a 0 , r 0 , . . . , s T , a T , r T , s T +1 ) where r t = R ( s t , a t ) and T is the length of the trajectory which can be infinite. If a trajectory is generated by a probabilistic policy π : S × A → R ≥ 0 , Z π = ∑ T t =0 γ t r t is a random variable that describes the discounted return the policy achieves. The objective is to find a π ⋆ that maximizes the expected discounted return, π ⋆ = arg max π E τ ∼ p π ( · ) [ Z π ] , where p π ( τ ) = µ ( s 0 ) ∏ T t =0 P ( s t +1 | s t , a t ) π ( s t | a t ) . For simplicity, we will use E π instead of E τ ∼ p π ( · ) to denote the expectation over trajectories sampled

from the policy π . With a slight abuse of notation, we use Z π ( s , a ) to denote the conditional discounted return when starting at s and taking action a ( i.e., s 0 = s and a 0 = a ). Finally, without loss of generality, we assume all measures are discrete and their values lie within [0 , 1] .

Value-based methods [113] rely on a fundamental quantity in RL, the state-action value function, also referred to as the Q-function , Q π ( s , a ) = E π [ Z π | s 0 = s , a 0 = a ] . The Q-function of a policy can be found at the fixed point of the Bellman operator, T π [11], T π Q ( s , a ) = E s ′ ∼ P ( ·| s , a ) , a ′ ∼ π ( ·| s ′ ) [ R ( s , a ) + γQ ( s ′ , a ′ )] . The value function for a state s is defined as V π ( s ) = E a ∼ π ( ·| s ) [ Q π ( s , a )] . Bellemare et al. [10] extends the procedure to the distribution of discounted returns, T π Z ( s , a ) d = R ( s , a ) + γZ ( s ′ , a ′ ) , s ′ ∼ P ( · | s , a ) and a ′ ∼ π ( · | s ′ ) , where d = denotes that two random variables have the same distributions. This extension is referred to as distributional RL (we provide a more detailed description of QR-DQN, the distributional RL algorithm we use, in Appendix D.1). For value-based methods, the greedy policy is directly derived from the Q-function as π ( a | s ) = arg max ′ Q ( a ′ , s ) ( a ) .

✶ a Generalization in Contextual MDPs. A contextual Markov decision process (CMDP) [39] is a special class of partially observable Markov decision process (POMDP) consisting of different MDPs, that share state and action spaces but have different R , P , and µ . In addition to the standard assumptions of a CMDP, we assume the existence of a structured distribution Ψ over the MDPs. During training, we are given a (finite or infinite) number of training MDPs, ˜ M train = {M 1 , M 2 , . . . , M n } , drawn from Ψ (Ghosh et al. [35] refers to this setting as epistemic POMDP). For each M , we use Q π M ( s , a ) and V π M ( s ) to denote the Q-function and value function of that M . We use p π, M ( τ ) to denote the trajectory distribution of rolling out π in M . The objective is to find a single policy π that maximizes the expected discounted return over the entire distribution of MDPs, E τ ∼ p π, M ( · ) , M∼ Ψ( · ) [ ∑ T t =0 γ t r t ] without retraining on the unseen M ( i.e., zero-shot generalization). We will refer to this quantity as the test return 2 .

Since in RL the algorithm collects its own data, the appropriate notion of generalization is the performance a learning algorithm can achieve given a finite number of interactions. Crucially, if algorithm 1 achieves a better test return than algorithm 2 given the same number of interactions with M train, we can say that algorithm 1 generalizes better than algorithm 2. Furthermore, we assume the existence of π ⋆ (potentially more than one) that is α -optimal for all M∈ supp (Ψ) and all s 0 in supp ( µ M ) , i.e., V π ∗ M ( s 0 ) ≥ max π V π M ( s 0 ) -α where α is small. This assumption only ensures that zero-shot generalization is intractable [69] but it is 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 a sample-efficient manner, i.e., optimization; if it is finite, the challenge is also generalization to unseen environments.

3 Generalization in a Tabular CMDP

Much of the literature treats generalization in deep RL as a pure representation learning problem [100, 119, 120, 1] and aims to improve it by using regularization [29, 118, 21, 45] or data augmentation [21, 61, 116, 56, 60, 92, 112]. In this section, we will first show that the problem of generalization in RL extends beyond representation learning by considering a 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 × 5 grid environment (Figure 2a). During training, the agent always starts at a fixed initial state, ( x = 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, r = 2 and the episode ends. Otherwise, the agent receives a small negative

0

1

2

3

4

x

0

1

2

3

4

y

-0.04

-0.04

-0.04

-0.04

2

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

Reward

(a) Tabular CMDP

0

200

400

600

800

episode

0.5

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

suboptimality

Q-learning +  -greedy

(b)

ε

-greedy

0

200

400

600

800

episode

0.5

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

suboptimality

Q-learning + UCB

(c) UCB

reward, r = -0 . 04 . At test time, the agent starts at a different location, ( x = 0 , y = 4) . In other words, the train and test MDPs only differ by their initial state distribution. In addition, each 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 ε -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 [113]. Greedy explores with ε -greedy strategy which takes a random action with probability ε and the best action according to the Q-function arg max a Q ( s, a ) with probability 1 -ε . In contrast, UCB is uncertainty-driven so it explores according to π ( a | s ) = ✶ ( a = arg max a ′ Q ( s, a ′ ) + c √ log( t ) /N ( s, a ′ )) , where t is the total number of timesteps, c is the exploration coefficient, and N ( s, a ) is the number of times the agent has taken action a in state s , with ties broken randomly 3 . While Greedy is a naive but widely used exploration strategy [105], UCB is an effective algorithm designed for multi-armed bandits [6, 5]. Chen et al. [19] showed that uncertainty-based 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 , i.e., the difference between the undiscounted return achieved by the learned policy and the undiscounted return of the optimal policy. While all three methods are able to quickly achieve 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 4 . 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 .

It 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 a common phenomenon in more challenging CMDPs like Procgen, namely that at test time, the agent can often end up in

T=0

T=1

T=2

T=3

+1 reward

TrainingMDP

Greedy

Trajectory

+1 reward

Explorative

Trajectory

episode ends

TestMDP

Test

Trajectory

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 test time. See Figure 3 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 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.1, we explore some other interpretations for why exploration improves generalization (which are also related to the sample complexity of Q-learning). In addition, we also explore other potential reasons for why current RL methods generalize poorly in Appendix G.2. Note that we do not claim exploration to be the only way of improving generalization in RL as 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 in contextual MDPs.

4 Exploration via Distributional Ensemble

In the previous section, we showed we can improve generalization via exploration in the tabular setting. In this section, we would like to extend this idea to deep RL with function approximation. While in the tabular setting shown above there is no intrinsic stochasticity, environments can in general be stochastic ( e.g., random transitions, or random unobserved contexts). At a high level, 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 a large source of aleatoric uncertainty is not knowing which context the agent is in [90].

In this section, we introduce Exploration via Distributional Ensemble (EDE), a method that encourages the exploration of states with high epistemic uncertainty. Our method builds on several important ideas from prior works, the most important of which are deep ensembles and distributional RL. While ensembles are a 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 an 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 a Bayesian inference problem, given access to a posterior p ( θ | D ) , where θ is the discretized quantiles of Z ( s , a ) , and D is the dataset of experience on which the quantiles are esti-

mated. Let j ∈ [ N ] denote the index for the j th quantile and U{ 1 , N } denote the uniform distribution over integers between 1 and N . The uncertainty of the Q-value, Q ( s , a ) = E j ∼U{ 1 ,N } [ θ j ( s , a )] , is the relevant quantity that can inform exploration. The overall uncertainty σ 2 = Var θ ∼ p ( θ |D ) [ Q ( s , a )] can be decomposed into epistemic uncertainty σ 2 epi ( s , a ) and aleatoric uncertainty σ 2 ale ( s , a ) such that σ 2 ( s , a ) = σ 2 epi ( s , a ) + σ 2 ale ( s , a ) , where,

Ideally, given an infinite or sufficiently large amount of diverse experience, one would expect the posterior to concentrate on the true quantile θ ⋆ , so Var θ ∼ p ( θ |D ) [ θ j ( s , a )] and consequently σ 2 epi ( s , a ) = 0 . σ 2 ale 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 5 is beneficial for exploration. We refer interested readers to Charpentier et al. [17] for a more thorough discussion.

Sampling from p ( θ | D ) is computationally intractable for complex MDPs and function approximators such as neural networks. Clements et al. [20] approximates samples from p ( θ | D ) with randomized MAP sampling [87] which assumes a Gaussian prior over the model parameters. However, the unimodal nature of a Gaussian in parameter space may not have enough diversity for effective uncertainty estimation [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-connected value heads, g i : R d → R |A|× N that share a single feature extractor f : S → R d , similar to Osband 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 et 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, i ∈ [ M ] is the index for M ensemble heads of the Q-network, and j ∈ [ N ] is the index of the quantiles. The output of the i th head for state s and action a is θ i ( s , a ) ∈ R N , where the j th coordinate, and θ ij ( s , a ) is the j th quantile of the predicted state-action value distribution for that head. The finite sample estimates of the two uncertainties are:

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:

π ts ( a | s ) = ✶ arg max a ′ ξ ( s , a ′ ) ( a ) , where ξ ( s , a ′ ) ∼ N ( Q ( s , a ′ ) , φ ˆ σ epi ( s , a ′ )) . (3) φ ∈ R ≥ 0 is a coefficient that controls how the agent balances exploration and exploitation. Alternatively, we can use the upper-confidence bound (UCB, Chen et al. [19]):

which we found to achieve better results in CMDPs than Thompson sampling [20] on Procgen when we 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

0.15

0.30

0.45

EDE (ours)

IDAAC

PPO

QR-DQN

Median

0.2

0.3

0.4

IQM

0.18

0.24

0.30

0.36

Mean

0.66

0.72

0.78

0.84

Optimality Gap

Min-Max Normalized Score

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 k th actor collects experience with the following policy:

λ ∈ (0 , 1) and α ∈ R > 0 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 ε -greedy exploration for each actor. In practice, the performances are not sensitive to λ and α (see Figure 18 in Appendix G). Both learning and experience collection take place on a single machine and no prioritized experience replay [96] is used since prior work found it 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 RL. We build on the algorithm proposed in Clements et al. [20] for estimating the epistemic uncertainty, but we use deep ensemble instead of MAP [87], and use either UCB or Thompson sampling depending on the task and setting 6 . In addition, for UCB, we propose that each actor uses a different exploration coefficient for more diverse data. While variations of the components of our algorithm have been used in prior works, this particular combination is new (see Appendix E). Our ablation experiments show that each design choice is important and applying these techniques individually or naïvely combining them performs significantly worse.

5 Experiments

5.1 Procgen

Wecompare our method with 3 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-art value-based method on Procgen [26]; (2) PPO which is 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 only and evaluate all algorithms using the 4 metrics proposed in Agarwal et al. [2]. We run each algorithm on every game for 5 seeds and report the aggregated min-max normalized scores on the full test distribution, and the estimated bootstrap-estimated 95% confidence interval in Figure 4 (simulated with the runs). Our approach significantly outperforms the other baselines in terms of median and interquartile mean (IQM) (which are the more statistically robust metrics according to Agarwal et al. [2]). In particular, it achieves almost 3 times the median score of QR-DQN and more than 2 times the IQM of QR-DQN . In terms of mean and optimality gap, our method is competitive with IDAAC and outperforms all the other baselines. To the best of our knowledge, this is the first value-based method that achieves such strong performance on Procgen. See Appendices I, F, and J for more details about the experiments, the performance on individual games, hyperparameters, and sensitivity analysis.

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

0.00

0.15

0.30

0.45

QR-DQN+UCB+TEE (EDE)

QR-DQN+NoisyNet

DQN+NoisyNet

QR-DQN+ez-greedy

DQN+ez-greedy

QR-DQN+TEE

DQN+TEE

QR-DQN+UCB

DQN+UCB

Bootstrapped DQN

UA-DQN

QR-DQN

DQN

Median

0.00

0.15

0.30

0.45

IQM

0.00

0.15

0.30

Mean

Min-Max Normalized Score

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 a complete comparison. The points of comparison we use are: (1) Bootstrapped DQN [79] which uses bootstrapping to train several copies of models that have distinct exploration behaviors, (2) UCB [19] which uses an ensemble of models to do uncertainty estimation, (3) ϵz -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 to the linear layers 7 . 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+UCB which is itself worse than EDE. Without epistemic uncertainty, the agent cannot do very well even if it trains on 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 on the other hand outperforms QR-DQN and DQN because it only targets the (biased) epistemic uncertainty. In Figure 16c from Appendix G, we show that indeed exploring with the total uncertainty σ 2 performs significantly worse at test time than exploring with only the epistemic uncertainty σ 2 epi . UA-DQN also performs worse than QR-DQN+UCB suggesting that deep ensemble may have better uncertainty estimation (Appendix H).

QR-DQN with ϵz -greedy exploration marginally improves upon the base QR-DQN, but remains significantly worse than EDE. This may be due to the fact that ϵz -greedy exploration can induce temporally extended exploration but it is not aware of the agent's uncertainty. Not accounting for uncertainty can be detrimental since CMDPs can have a much larger number of effective states than singleton MDPs. If the models have gathered enough information about a 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 are naturally adaptive - the agents will take into account what they have already learned. While a direct theoretical comparison between NoisyNet and our method is 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 ϵz -greedy which implicitly use an schedule.

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 a 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 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-sourced 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 to all the baselines presented in Hafner [37], including DreamerV2 [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 approach, suggests that the exploration strategy is crucial for improving performance on such CMDPs.

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 to 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, 1], 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 information-theoretic approaches [18, 71], non-stationarity reduction [46, 78], curriculum learning [47, 109, 48, 85], planning [4], forwardbackward representations [111], 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 on policy optimization methods, and none emphasizes the key role exploration plays in training more general agents. In contrast, our goal is to understand why value-based methods are significantly worse on CMDPs.

Exploration. Exploration is a fundamental aspect of RL [55, 33, 108]. Common approaches include ε -greedy [105], count-based exploration [104, 9, 68], curiosity-based exploration [97, 102, 86], or novelty-based 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 in 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 singleton 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 is 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 can 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 epistemic and aleatoric uncertainties. This results in the first value-based based method to achieve state-of-the-art 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-based methods, similar ideas could be applied to policy optimization to further improve their generalization abilities. In a broader context, it is 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 to QR-DQN (Figure 13), which indicates that there are other bottlenecks beyond poor exploration. Another limitation of our approach is that it 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.

References

R. Agarwal, M. C. Machado, P. S. Castro, and M. G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. 2021.

R. Agarwal, M. Schwarzer, P. S. Castro, A. C. Courville, and M. Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems , 34:29304-29320, 2021.

A. Ajay, G. Yang, O. Nachum, and P. Agrawal. Understanding the generalization gap in visual reinforcement learning. 2021.

A. Anand, J. Walker, Y. Li, E. Vértes, J. Schrittwieser, S. Ozair, T. Weber, and J. B. Hamrick. Procedural generalization by planning with self-supervised world models. arXiv preprint arXiv:2111.01587 , 2021.

J.-Y. Audibert, R. Munos, and C. Szepesvári. Exploration-exploitation tradeoff using variance estimates in multi-armed bandits. Theoretical Computer Science , 410(19):1876-1902, 2009.

P. Auer, N. Cesa-Bianchi, and P. Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning , 47(2):235-256, 2002.

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.

A. P. Badia, B. Piot, S. Kapturowski, P. Sprechmann, A. Vitvitskyi, Z. D. Guo, and C. Blundell. Agent57: Outperforming the atari human benchmark. In International Conference on Machine Learning , pages 507-517. PMLR, 2020.

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.

M. G. Bellemare, W. Dabney, and R. Munos. A distributional perspective on reinforcement learning. In International Conference on Machine Learning , pages 449-458. PMLR, 2017.

R. Bellman. A markovian decision process. Journal of mathematics and mechanics , pages 679-684, 1957.

E. Bengio, J. Pineau, and D. Precup. Interference and generalization in temporal difference learning. ArXiv , abs/2003.06350, 2020.

M. Bertrán, N. Martínez, M. Phielipp, and G. Sapiro. Instance based generalization in reinforcement learning. ArXiv , abs/2011.01089, 2020.

Y. Burda, H. Edwards, A. Storkey, and O. Klimov. Exploration by random network distillation. arXiv preprint arXiv:1810.12894 , 2018.

W. Carvalho, A. Lampinen, K. Nikiforou, F. Hill, and M. Shanahan. Feature-attending recurrent modules for generalization in reinforcement learning. arXiv preprint arXiv:2112.08369 , 2021.

O. Chapelle and L. Li. An empirical evaluation of thompson sampling. Advances in neural information processing systems , 24, 2011.

B. Charpentier, R. Senanayake, M. Kochenderfer, and S. Günnemann. Disentangling epistemic and aleatoric uncertainty in reinforcement learning. arXiv preprint arXiv:2206.01558 , 2022.

J. Z. Chen. Reinforcement learning generalization with surprise minimization. ArXiv , abs/2004.12399, 2020.

R. Y. Chen, S. Sidor, P. Abbeel, and J. Schulman. Ucb exploration via q-ensembles. arXiv preprint arXiv:1706.01502 , 2017.

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.

K. Cobbe, O. Klimov, C. Hesse, T. Kim, and J. Schulman. Quantifying generalization in reinforcement learning. arXiv preprint arXiv:1812.02341 , 2018.

K. Cobbe, C. Hesse, J. Hilton, and J. Schulman. Leveraging procedural generation to benchmark reinforcement learning. arXiv preprint arXiv:1912.01588 , 2019.

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.

W. Dabney, G. Ostrovski, and A. Barreto. Temporally-extended ε -greedy exploration. In International Conference on Learning Representations , 2021. URL https://openreview. net/forum?id=ONBPHFZ7zG4 .

C. Dann and E. Brunskill. Sample complexity of episodic fixed-horizon reinforcement learning. Advances in Neural Information Processing Systems , 28, 2015.

A. Ehrenberg, R. Kirk, M. Jiang, E. Grefenstette, and T. Rocktäschel. A study of off-policy learning in environments with procedural content generation. In ICLR Workshop on Agent Learning in Open-Endedness , 2022.

L. Espeholt, H. Soyer, R. Munos, K. Simonyan, V. Mnih, T. Ward, Y. Doron, V. Firoiu, T. Harley, I. Dunning, et al. Impala: Scalable distributed deep-rl with importance weighted actor-learner architectures. arXiv preprint arXiv:1802.01561 , 2018.

J. Fan and W. Li. Dribo: Robust deep reinforcement learning via multi-view information bottleneck. In International Conference on Machine Learning , pages 6074-6102. PMLR, 2022.

J. Farebrother, M. C. Machado, and M. H. Bowling. Generalization and regularization in dqn. ArXiv , abs/1810.00123, 2018.

Y. Flet-Berliac, J. Ferret, O. Pietquin, P. Preux, and M. Geist. Adversarially guided actor-critic. CoRR , abs/2102.04376, 2021. URL https://arxiv.org/abs/2102.04376 .

S. Fort, H. Hu, and B. Lakshminarayanan. Deep ensembles: A loss landscape perspective. arXiv preprint arXiv:1912.02757 , 2019.

M. Fortunato, M. G. Azar, B. Piot, J. Menick, I. Osband, A. Graves, V. Mnih, R. Munos, D. Hassabis, O. Pietquin, et al. Noisy networks for exploration. arXiv preprint arXiv:1706.10295 , 2017.

R. Fruit and A. Lazaric. Exploration-exploitation in mdps with options. In Artificial intelligence and statistics , pages 576-584. PMLR, 2017.

X. Fu, G. Yang, P. Agrawal, and T. Jaakkola. Learning task informed abstractions. In International Conference on Machine Learning , pages 3480-3491. PMLR, 2021.

D. Ghosh, J. Rahme, A. Kumar, A. Zhang, R. 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-25515, 2021.

J. Grigsby and Y. Qi. Measuring visual generalization in continuous control from pixels. ArXiv , abs/2010.06740, 2020.

D. Hafner. Benchmarking the spectrum of agent capabilities. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id= 1W0z96MFEoH .

D. Hafner, T. P. Lillicrap, M. Norouzi, and J. Ba. Mastering atari with discrete world models. In International Conference on Learning Representations , 2021. URL https://openreview. net/forum?id=0oabwyZbOu .

A. Hallak, D. Di Castro, and S. Mannor. Contextual markov decision processes. arXiv preprint arXiv:1502.02259 , 2015.

N. Hansen and X. Wang. Generalization in reinforcement learning by soft data augmentation. In 2021 IEEE International Conference on Robotics and Automation (ICRA) , pages 1361113617. IEEE, 2021.

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-3693, 2021.

K. 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 770778, 2016.

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 Thirty-second AAAI conference on artificial intelligence , 2018.

D. Horgan, J. Quan, D. Budden, G. Barth-Maron, M. Hessel, H. van Hasselt, and D. Silver. Distributed prioritized experience replay. In International Conference on Learning Representations , 2018. URL https://openreview.net/forum?id=H1Dy---0Z .

M. Igl, K. Ciosek, Y. Li, S. Tschiatschek, C. Zhang, S. Devlin, and K. Hofmann. Generalization in reinforcement learning with selective noise injection and information bottleneck. In Advances in Neural Information Processing Systems , pages 13956-13968, 2019.

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.05826, 2020.

M. Jiang, E. Grefenstette, and T. Rocktäschel. Prioritized level replay. ArXiv , abs/2010.03934, 2020.

M. Jiang, M. Dennis, J. Parker-Holder, J. Foerster, E. Grefenstette, and T. Rocktäschel. Replayguided adversarial environment design. Advances in Neural Information Processing Systems , 34:1884-1897, 2021.

A. Juliani, A. Khalifa, V.-P. Berges, J. Harper, H. Henry, A. Crespi, J. Togelius, and D. Lange. Obstacle tower: A generalization challenge in vision, control, and planning. ArXiv , abs/1902.01378, 2019.

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.

M. Kearns and S. Singh. Finite-sample convergence rates for q-learning and indirect algorithms. Advances in neural information processing systems , 11, 1998.

A. Kendall and Y. Gal. What uncertainties do we need in bayesian deep learning for computer vision? Advances in neural information processing systems , 30, 2017.

R. Kirk, A. Zhang, E. Grefenstette, and T. Rocktäschel. A survey of generalisation in deep reinforcement learning. arXiv preprint arXiv:2111.09794 , 2021.

B. Ko and J. Ok. Efficient scheduling of data augmentation for deep reinforcement learning. arXiv preprint arXiv:2206.00518 , 2022.

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.

I. Kostrikov, D. Yarats, and R. Fergus. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. arXiv preprint arXiv:2004.13649 , 2020.

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:8198-8210, 2020.

H. Kuttler, N. Nardelli, A. H. Miller, R. Raileanu, M. Selvatici, E. Grefenstette, and T. Rocktäschel. The nethack learning environment. ArXiv , abs/2006.13760, 2020.

B. Lakshminarayanan, A. Pritzel, and C. Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems , 30, 2017.

M. Laskin, K. Lee, A. Stooke, L. Pinto, P. Abbeel, and A. Srinivas. Reinforcement learning with augmented data. arXiv preprint arXiv:2004.14990 , 2020.

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. net/forum , 2020.

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 6131-6141. PMLR, 2021.

G. Li, Y. Wei, Y. Chi, Y. Gu, and Y. Chen. Sample complexity of asynchronous q-learning: Sharper analysis and variance reduction. Advances in neural information processing systems , 33:7031-7043, 2020.

L. Liang, Y. Xu, S. McAleer, D. Hu, A. Ihler, P. Abbeel, and R. Fox. Reducing variance in temporal-difference value estimation via ensemble of deep networks. In International Conference on Machine Learning , pages 13285-13301. PMLR, 2022.

E. Z. Liu, A. Raghunathan, P. Liang, and C. Finn. Decoupling exploration and exploitation for meta-reinforcement learning without sacrifices. In International conference on machine learning , pages 6925-6935. PMLR, 2021.

C. Lyle, M. Rowland, W. Dabney, M. Kwiatkowska, and Y. Gal. Learning dynamics and generalization in reinforcement learning. arXiv preprint arXiv:2206.02126 , 2022.

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.

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.

D. Malik, Y. Li, and P. Ravikumar. When is generalizable reinforcement learning tractable? Advances in Neural Information Processing Systems , 34:8032-8045, 2021.

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.

B. Mazoure, R. T. des Combes, T. Doan, P. Bachman, and R. D. Hjelm. Deep reinforcement and infomax learning. ArXiv , abs/2006.07217, 2020.

B. Mazoure, I. Kostrikov, O. Nachum, and J. Tompson. Improving zero-shot generalization in offline reinforcement learning using generalized similarity functions. arXiv preprint arXiv:2111.14629 , 2021.

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.

J. Mockus. The bayesian approach to global optimization. In System Modeling and Optimization , pages 473-481. Springer, 1982.

S. Mohanty, J. Poonganam, A. Gaidon, A. Kolobov, B. Wulfe, D. Chakraborty, G. Šemetulskis, J. Schapke, J. Kubilius, J. Pašukonis, et al. Measuring sample efficiency and generalization in reinforcement learning benchmarks: Neurips 2020 procgen benchmark. arXiv preprint arXiv:2103.15332 , 2021.

J. Morimoto and K. Doya. Robust reinforcement learning. Neural computation , 17(2):335-359, 2005.

A. Nichol, V. Pfau, C. Hesse, O. Klimov, and J. Schulman. Gotta learn fast: A new benchmark for generalization in rl. ArXiv , abs/1804.03720, 2018.

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.

I. Osband and B. Van Roy. Bootstrapped thompson sampling and deep exploration. arXiv preprint arXiv:1507.00300 , 2015.

I. Osband, D. Russo, and B. Van Roy. (more) efficient reinforcement learning via posterior sampling. Advances in Neural Information Processing Systems , 26, 2013.

I. Osband, C. Blundell, A. Pritzel, and B. Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems , 29, 2016.

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.

B. O'Donoghue, I. Osband, R. Munos, and V. Mnih. The uncertainty bellman equation and exploration. In International Conference on Machine Learning , pages 3836-3845, 2018.

C. Packer, K. Gao, J. Kos, P. Krähenbühl, V. Koltun, and D. X. Song. Assessing generalization in deep reinforcement learning. ArXiv , abs/1810.12282, 2018.

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.

D. Pathak, P. Agrawal, A. A. Efros, and T. Darrell. Curiosity-driven exploration by selfsupervised prediction. In International conference on machine learning , pages 2778-2787. PMLR, 2017.

T. Pearce, F. Leibfried, and A. Brintrup. Uncertainty in neural networks: Approximately bayesian ensembling. In International conference on artificial intelligence and statistics , pages 234-244. PMLR, 2020.

V. H. Pong, M. Dalal, S. Lin, A. Nair, S. Bahl, and S. Levine. Skew-fit: State-covering self-supervised reinforcement learning. arXiv preprint arXiv:1903.03698 , 2019.

M. M. Rahman and Y. Xue. Adversarial style transfer for robust policy optimization in reinforcement learning. 2021.

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.

R. Raileanu and T. Rocktäschel. Ride: Rewarding impact-driven exploration for procedurallygenerated environments. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id=rkg-TJBFPB .

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.

A. Rajeswaran, K. Lowrey, E. Todorov, and S. M. Kakade. Towards generalization and simplicity in continuous control. ArXiv , abs/1703.02660, 2017.

S. Ross, G. Gordon, and D. Bagnell. A reduction of imitation learning and structured prediction 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.

J. Roy and G. Konidaris. Visual transfer for reinforcement learning via wasserstein domain confusion. arXiv preprint arXiv:2006.03465 , 2020.

T. Schaul, J. Quan, I. Antonoglou, and D. Silver. Prioritized experience replay. arXiv preprint arXiv:1511.05952 , 2015.

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-227, 1991.

B. Settles. Active learning literature survey. 2009.

A. Sonar, V. Pacelli, and A. Majumdar. Invariant policy optimization: Towards stronger generalization in reinforcement learning. ArXiv , abs/2006.01096, 2020.

X. Song, Y. Jiang, S. Tu, Y. Du, and B. Neyshabur. Observational overfitting in reinforcement learning. In International Conference on Learning Representations , 2020. URL https: //openreview.net/forum?id=HJli2hNKDH .

A. Stani´ c, Y . Tang, D. Ha, and J. Schmidhuber. An investigation into the open world survival game crafter. In Decision Awareness in Reinforcement Learning Workshop at ICML 2022 , 2022. URL https://openreview.net/forum?id=C5q00n9dBxC .

C. Stanton and J. Clune. Curiosity search: Producing generalists by encouraging individuals to continually explore and acquire skills throughout their lifetime. PLOS ONE , 11(9):1-20, 09 2016. doi: 10.1371/journal.pone.0162235. URL https://doi.org/10.1371/journal. pone.0162235 .

A. Stooke, K. Lee, P. Abbeel, and M. Laskin. Decoupling representation learning from reinforcement learning. ArXiv , abs/2009.08319, 2020.

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.

R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction . MIT press, 2018.

R. S. Sutton, D. McAllester, S. Singh, and Y. Mansour. Policy gradient methods for reinforcement learning with function approximation. Advances in neural information processing systems , 12, 1999.

A. A. Taiga, W. Fedus, M. C. Machado, A. Courville, and M. G. Bellemare. On bonus based exploration methods in the arcade learning environment. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id=BJewlyStDr .

J. Tarbouriech, E. Garcelon, M. Valko, M. Pirotta, and A. Lazaric. No-regret exploration in goal-oriented reinforcement learning. In International Conference on Machine Learning , pages 9428-9437. PMLR, 2020.

O. E. L. Team, A. Stooke, A. Mahajan, C. Barros, C. Deck, J. 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.

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-294, 1933.

A. Touati and Y. Ollivier. Learning one representation to optimize all rewards. Advances in Neural Information Processing Systems , 34:13-23, 2021.

K. Wang, B. Kang, J. Shao, and J. Feng. Improving generalization in reinforcement learning with mixture regularization. ArXiv , abs/2010.10814, 2020.

C. J. Watkins and P. Dayan. Q-learning. Machine learning , 8(3):279-292, 1992.

R. J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine learning , 8(3):229-256, 1992.

D. Yarats, R. Fergus, A. Lazaric, and L. Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. arXiv preprint arXiv:2107.09645 , 2021.

C. Ye, A. Khalifa, P. Bontrager, and J. Togelius. Rotation, translation, and cropping for zero-shot generalization. arXiv preprint arXiv:2001.09908 , 2020.

D. Zha, W. Ma, L. Yuan, X. Hu, and J. Liu. Rank the episodes: A simple approach for exploration in procedurally-generated environments. arXiv preprint arXiv:2101.08152 , 2021.

A. Zhang, N. Ballas, and J. Pineau. A dissection of overfitting and generalization in continuous reinforcement learning. ArXiv , abs/1806.07937, 2018.

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.06016 , 2020.

A. Zhang, R. McAllister, R. Calandra, Y. Gal, and S. Levine. Learning invariant representations for reinforcement learning without reconstruction. arXiv preprint arXiv:2006.10742 , 2020.

C. Zhang, O. Vinyals, R. Munos, and S. Bengio. A study on overfitting in deep reinforcement learning. ArXiv , abs/1804.06893, 2018.

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, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems , 2021. URL https://openreview.net/forum?id=CYUzpnOkFJp .

X. Zhang, Y. Ma, and A. Singla. Task-agnostic exploration in reinforcement learning. Advances in Neural Information Processing Systems , 33:11734-11743, 2020.

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 on 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 a good estimate for all Q ( s, 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 , π b , which collects data by interacting with the environment and that the environment is ergodic (see all assumptions in Li et al. [63]). The update to the value function is done in an off-policy manner with the data generated by π b . Intuitively, the behavioral policy is responsible for exploration. We will first define some useful quantities that relate to the inherent properties of π b and the MDP. Recall that µ is an initial state distribution and let p π ( τ | s ) be the trajectory distribution induced by π starting from s , we can define the state marginal distribution at time t given that the initial state is drawn from µ as P t ( s ; π, µ ) = E τ ∼ p π ( ·| s 0 ) ,s 0 ∼ µ [ ✶ { s t = s } ] , the state-action marginal distribution P t ( s, a ; π, µ ) = π ( a | s ) P t ( s ; π, µ ) , the stationary state occupancy distribution, d π ( s ) = lim T →∞ 1 T +1 ∑ T t =0 P t ( s ; π, µ ) and finally the stationary state-action distribution d π ( s, a ) = π ( a | s ) d π ( s ) 8 . Using D TV to denote the total variation distance, we define:

Intuitively, µ min ( π ) is the occupancy of the bottleneck state that is hardest to reach by π , and t mix captures how fast the distribution generated by π decorrelates with any initial state distribution. In this setting, Li et al. [63] showed that with ˜ O ( 1 µ min ( π b )(1 -γ ) 5 ϵ 2 + t mix ( π b ) µ min ( π b )(1 -γ ) ) steps 9 , one can expect the Q-value to be ϵ -close to the optimal Q-value, i.e., max ( s,a ) ∈S×A | Q t ( s, a ) -Q ⋆ ( s, a ) | ≤ ϵ with high probability. The second term of the complexity concerns the rate at which the state-action marginal converges to a 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 π b , so a 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 . In this setting, d π ( s ) and consequently d π ( s, 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

distribution. In other words, we can roughly say that the more a 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 M based on the behavioral policy. Kearns and Singh [51] prunes M based on whether d π ( s, a ) is larger than α , a hyperparameter that is larger than µ min ( π b ) , and argues that one can obtain better sample complexity on the sub-MDP. This can be seen as throwing away state-action pairs that are not likely visited by the behavioral policy in the analysis. Concretely, the sub-MDP M π ( α ) is an MDP where the state space is:

If we further assume that all the state-action pairs visited by the optimal policy, π ⋆ , starting from µ are contained within G ( α ) , then the optimal policy of the M π b ( α ) 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 π b -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 ε -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 α > µ min ( π b ) on M π b ( α ) and thus Q-values would converge to the optimal value on M π b ( α ) faster than the full MDP, M . If α is sufficiently larger than µ min ( π b ) , the sample complexity on M π b ( α ) could be much better than on M , explaining the policy's superior performance near the initial state when the Q ( s, a ) further away from the initial states have not converged. Note that M π b ( α ) is only a construct for making the theoretical argument. The Q-values of the 'pruned' state-action pairs are still updated as usual algorithmically. Even if the assumption is not true, the optimal policy on M π b ( α ) may still be good enough. If the optimal policy on M π b ( α ) or α has to be small for the assumption to approximately hold, then one would need a better behavioral policy to solve the training environment efficiently, much less having a good generalization performance on the test environment.

Finally, Li et al. [63] showed that with high probability min s,a 2 3 N ( s, a ) /t ≤ µ min ( π b ) which means that having a behavior policy that encourages visiting states with low visitation count, N ( s, a ) , should result in increasing min s,a 2 3 N ( s, a ) /t and thus increasing µ min ( π b ) . This indicates that UCB-based exploration should have a better sample complexity than ϵ -greedy which does not explicitly target the visitation count and can be in general pretty bad at covering large state space. Indeed, for episodic RL, the state-of-the-art algorithms generally rely on 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 µ min.

B Case Study: Simplified Bigfish

0=1

T=1

T=2

T=3

Training MDP

+1 reward

Greedy

Trajectory

+1reward

Explorative

Trajectory

episode ends

Test MDP

Test

Trajectory

To further illustrate the intuition that exploration on the training MDPs can help zero-shot generalization 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 always starts in the state shown in the T = 0 of the greedy trajectory (top row). Using random exploration strategy ( e.g., ϵ -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 a 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 a naive exploration will most likely keep repeating the greedy trajectory. On the other hand, once the agent has sampled the greedy trajectory sufficiently, uncertainty-aware exploration ( e.g., UCB or count-based exploration) will more likely sample these rarer trajectories since they have been visited less and thus have higher uncertainty. This has no impact on the performance in the training MDP because the agent has learned the optimal policy regardless of which exploration strategy is used. 11

However, on the test MDP, the initial state is shown in T = 0 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 a good representation so the agent knows what to do, but this does not necessarily have to be the case - the objective only cares about solving the training MDP and does not explicitly encourage learning representations that help generalization. Suppose that the agent keeps moving up due to randomness in the initialization of the weight, it will run into the big fish and die (bottom row, T = 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 T = 1 (pink arrow). Due to this structure, agents that have learned about

the explorative trajectory during training time will be able to recognize that the better action from T = 1 is to move to the right which will avoid death and ultimately result in +1 reward.

Once again, we emphasize that this is a simplified example for illustrative purposes. In reality, the sequences of states do not need to 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 it should not hurt. Like most problems in deep RL, it is difficult to prove this hypothesis analytically but empirically it is indeed the case that EDE as better or the same performance as QR-DQN in 11 out of 16 Procgen environments and perform approximately the same in the remaining environments (Figure 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 a side effect.

C Tabular Experiments

C.1 Generalization to Different Initial States

0

200

400

600

800

episode

4

3

2

1

0

1

2

return

-greedy w/ different epsilon

= 0.0

= 0.1

= 0.2

= 0.3

= 0.4

= 0.5

= 0.6

= 0.7

= 0.8

= 0.9

(a)

ε

-greedy

0

200

400

600

800

episode

1.0

0.5

0.0

0.5

1.0

1.5

2.0

return

UCB w/ different exploration bonus

c

c

= 5.0

c

= 10.0

c

= 15.0

c

= 20.0

c

= 25.0

c

= 30.0

c

= 35.0

c

= 40.0

c

= 45.0

c

= 50.0

(b) UCB

0

200

400

600

800

episode

1.0

0.5

0.0

0.5

1.0

1.5

2.0

return

Policy gradient w/ different entropy bonus

= 0.0

= 0.1

= 0.2

= 0.3

= 0.4

= 0.5

= 0.6

= 0.7

= 0.8

= 0.9

(c) Policy Gradient

Both the training MDP, µ train, and the test MDP, µ test, 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 ′ | s, a ) , to the next state corresponding to an action is 100% . γ is 0 . 9 . The two MDPs differ by only their initial state distribution: ρ train ( s ) = (0 , 0) ( s ) and ρ test ( s ) = (4 , 0) ( s ) .

✶ ✶ For both policy gradient and Q-learning, we use a base learning rate of α 0 = 0 . 05 that decays as training progresses. Specifically, at time step t , α t = 1 √ t α 0 .

For Q-learning, the Q-function is parameterized as ϑ ∈ R |S|×|A| where each entry is a state-action value. The update is:

where Q ( s, a ) = ϑ 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 s over the course of training and explores according to:

For Q-learning with ε -greedy, the exploration policy is:

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 as ϑ ∈ R |S|×|A| where each entry is the logit for the distribution of taking action a from s , i.e., π ( a | s ) = exp( ϑ s,a ) ∑ a ′ ∈A exp( ϑ s,a ′ ) . The overall gradient is:

The update is:

where the expectation is estimated with a 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}

β : {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 are relatively robust to the choice of hyperparameters, whereas ε -greedy's generalization performance varies greatly with ε . Furthermore, even with extremely large ε , the generalization performance is still far worse than UCB and policy gradient. The best-performing values are respectively ε ⋆ = 0 . 9 , c ⋆ = 45 and β ⋆ = 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 it is not as good as UCB .

0

200

400

600

800

episode

0.5

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

suboptimality

Policy Gradient

train

test

C.3 Full trajectory update and Sarsa ( λ )

Exploration is evidently not the only approach one can take to improve generalization in RL. In fact, we see that policy gradient with rather naive softmax exploration can also outperform ε -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 ( λ ) [105, Ch 12.7] which is an on-policy RL method that learns a Q-value with whole-trajectory updates. The results are shown in Figure 10a.

While both on-policy methods use fairly naive exploration methods (softmax exploration for PG and ε -greedy for Sarsa ( λ ) , 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 as a 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 ( λ )

Sarsa ( λ ) 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, we refer the reader to Chapter 12.7 of Sutton and Barto [105] for the details. The Sarsa ( λ ) with ε -greedy experiments inherit all the hyperparameters of Q-learning + ε -greedy experiments, but have one more hyperparameter λ ∈ (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):

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 ( λ ) outperforms both Greedy and PG , but is still worse than UCB and has a much larger variance.

9

0

200

400

600

800

episode

1.0

0.5

0.0

0.5

1.0

1.5

2.0

return

) w/ different

= 0.0

= 0.1

= 0.2

= 0.3

= 0.4

= 0.5

= 0.6

= 0.7

= 0.8

= 0.9

(a) Average of return.

0

200

400

600

800

episode

1.0

0.5

0.0

0.5

1.0

1.5

2.0

return

Standard Deviation of Sarsa(

) w/ different

= 0.0

= 0.1

= 0.2

= 0.3

= 0.4

= 0.5

= 0.6

= 0.7

= 0.8

= 0.9

(b) Standard deviation of return.

0

200

400

600

800

episode

0.5

0.0

0.5

1.0

1.5

2.0

2.5

3.0

3.5

4.0

suboptimality

) +  -greedy

(c) Best hyperparameter.

ε

= 0

.

and

λ

= 0

.

9

,

C.4 Generalization to Different Dynamics

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 generalization to new dynamics. We use a 7 × 7 grid as the state space (Figure 11a). 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 at test time, a gust of wind can blow the agent off course. Specifically, with a 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 an oracle that consistently moves towards

0

1

2

3

4

5

6

x

0

1

2

3

4

5

6

y

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

2

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

-0.04

Reward

(a) Tabular CMDP

0

200

400

600

800

episode

0

1

2

3

4

5

6

7

suboptimality

Q-learning +  -greedy

(b)

ε

-greedy

0

200

400

600

800

episode

0

1

2

3

4

5

6

7

suboptimality

Q-learning + UCB

(c) UCB

0

200

400

600

800

episode

0

1

2

3

4

5

6

7

suboptimality

Policy Gradient

train

test

(d) Policy Gradient

0

200

400

600

800

episode

12

10

8

6

4

2

0

return

-greedy w/ different epsilon

= 0.0

= 0.1

= 0.2

= 0.3

= 0.4

= 0.5

= 0.6

= 0.7

= 0.8

= 0.9

(a)

ε

-greedy

0

200

400

600

800

episode

4

3

2

1

0

1

return

UCB w/ different exploration bonus

c

c

= 5.0

c

= 10.0

c

= 15.0

c

= 20.0

c

= 25.0

c

= 30.0

c

= 35.0

c

= 40.0

c

= 45.0

c

= 50.0

(b) UCB

0

200

400

600

800

episode

4

3

2

1

0

1

return

Policy gradient w/ different entropy bonus

= 0.0

= 0.1

= 0.2

= 0.3

= 0.4

= 0.5

= 0.6

= 0.7

= 0.8

= 0.9

(c) Policy Gradient

the goal). This setup is related to robust reinforcement learning [76], but exploration makes no assumptions about the type or magnitude of perturbation (thus making more complicated perturbation, e.g., different random initializations of the same game, possible) and does not require access to a simulator.

D Additional Background

D.1 Quantile Regression DQN

Quantile regression DQN (QR-DQN) is a variant of distributional RL that tries to learn the quantiles of the state-action value distribution with quantile regression rather than produce the actual distribution of Q-values. Compared to 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 et al. [23] for a more in-depth discussion of QR-DQN.

Given a random variable Y with associated CDF F Y ( y ) = P ( Y ≤ y ) , the inverse CDF of Y is

To approximate Y , we may discretize the CDF with N qunatile values:

Using the same notation as main text, we use θ : S × A → R N to denote a parametric function that maps a state-action pair, ( s , a ) , to the estimated quantiles { θ j ( s , a ) } N j =1 . The approximated CDF of

Z ( s , a ) is thus:

where δ ( θ j ( s , a )) is the Dirac delta distribution centered at θ j ( s , a ) . To learn the quantiles, we use quantile regression which allows for unbiased estimate of the true quantile values. Concretely, for a target distribution Z and quantile τ , the value of F -1 Z ( τ ) is the minimizer of the quantile regression loss:

Due to discretization, for each quantile τ j , the minimizer of the 1-Wasserstein distance within that quantile is at θ that satisfies F Z ( θ ) = 1 2 ( τ j -1 + τ j ) = ˆ τ j (see Lemma 2 of Dabney et al. [23]). Consequentily, to approximate Z , we can simultaneously optimize for all of the quantiles, { θ j } N j =1 , by minimizing the following loss:

To stablize optimization, QR-DQN uses an modified version of quantile loss called quantile Huber loss which utilizes the Huber loss with hyperparameter κ :

The quantile Huber loss is defined as:

ρ κ τ ( u ) = | τ -✶ ( u < 0) | L κ ( u ) . In QR-DQN, both the approximating distribution and the target distribution are discretized with quantiles. The target distribution is computed with bootstrapping through the distributional Bellman operator. Concretely, given a target network, θ target , and transition tuple ( s , a , r, s ′ ) , the target distribution is approximated as:

The TD target Q-value distribution's quantiles and distribution are:

Treating the target distribution as an oracle ( i.e., not differentiable), we update the current quantile estimates to minimize the following quantile regression loss:

sg denotes the stop gradient operation as we do 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 pseudocode omits the

initial data collection steps during which the agent uses a uniformly random policy to collect data and does not update the parameters. We also write the parallel experience collection as a 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 i.e., how the action is chosen and how the ensemble members are updated.

In Algorithm 2, we describe how EDE chooses exploratory actions during training. For each state, we first extract the features with the shared feature extractor f and then compute the state-action value distributions for each ensemble head. With the state-action value distributions, we compute the estimate for the epistemic uncertainty as well as the expected Q-values. Using these quantities, we choose the action based on either Thompson sampling or UCB.

In Algorithm 3, we describe how EDE updates the parameters of the feature extractor and ensemble heads. This is where our algorithm deviates significantly from prior approaches. For each ensemble head, we sample an independent minibatch from the replay buffer and compute the loss like deep ensemble . Inside the loop, the loss is computed as if each ensemble member is a QR-DQN. However, doing so naïvely means that the feature extractor f will be updated at a 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 to 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 more 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 not 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 1 EDE

Algorithm 2 CHOOSEACTION ( s , f, { g i } , k

̸

F Procgen Results

In Table 3 and Table 2, we show the mean and standard deviation of the final unnormalized train and test scores (at the end of 25M environment steps) of EDE 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 we reproduced using the code from Raileanu and Fergus [90]. We see that for a subset of games, EDE s significantly more sample-efficient. For other games, EDE s either on par with the policy optimization methods or fails to train, indicating there are other issues behind its poor performance on these games. Note that IDAAC results tune the hyperparameters for individual games whereas we only tune the hyperparameters on bigfish .

0.0

0.5

1.0

1.5

2.0

1e7

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

bigfish

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.0

0.2

0.4

0.6

0.8

bossfight

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.2

0.0

0.2

0.4

0.6

caveflyer

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.0

0.1

0.2

0.3

0.4

chaser

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.2

0.0

0.2

0.4

0.6

climber

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.4

0.2

0.0

0.2

0.4

0.6

0.8

1.0

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.0

0.1

0.2

0.3

0.4

0.5

0.6

0.7

dodgeball

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.0

0.2

0.4

0.6

0.8

fruitbot

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.5

0.4

0.3

0.2

0.1

0.0

0.1

0.2

0.3

heist

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.2

0.0

0.2

0.4

0.6

jumper

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.2

0.0

0.2

0.4

0.6

0.8

1.0

leaper

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.6

0.4

0.2

0.0

0.2

0.4

0.6

maze

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.1

0.0

0.1

0.2

0.3

0.4

0.5

0.6

miner

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.2

0.0

0.2

0.4

0.6

ninja

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.10

0.05

0.00

0.05

0.10

0.15

plunder

EDE

qrdqn

ppo

daac

0.0

0.5

1.0

1.5

2.0

1e7

0.0

0.2

0.4

0.6

0.8

starpilot

EDE

qrdqn

ppo

daac

F.1 Additional Ablation

Median

IQM

Mean

OptimalityGap

DQN+UCB+TEE

0.00

0.02

0.04

0.06

0.015

0.030

0.045

0.015

0.0300.045

0.9600.975

0.990

Min-MaxNormalizedScore

maze

miner

heist

fruitbot

ninja

dodgeball

climber

jumper

chaser

bigfish

plunder

starpilot

caveflyer

bossfight

leaper

game

0.4

0.2

0.0

0.2

0.4

0.6

0.8

r

s

Correlation between train & test performance

(a) Meta Analysis

0

5

steps

1e6

0.0

0.2

0.4

0.6

0.8

1.0

return

Singleton Training Performance

DQN

PPO

(b) Singleton Experiments

0.0

0.5

1.0

1.5

steps

1e7

0.2

0.0

0.2

0.4

0.6

return

Task ID ablation

w/ task ID

w/o task ID

(c) DQN with Task ID

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 to Greedy is that better exploration improves sample efficiency resulting in better training performance, which can, in turn, lead to better test performance. Indeed, even in the simple tabular MDPs from Figures 2 and 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 3 can be simultaneously true , but the hypothesis 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-driven exploration to deep RL on complex CMDPs has additional challenges such as representation learning and unobserved contexts.

Optimization vs. Generalization. In supervised learning, deep neural networks have good generalization, meaning that they perform similarly well on 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 Procgen, 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 training distribution. To verify this hypothesis, we first conduct a meta-analysis 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 rank correlation, r s , between the training and test performance of the algorithms, for each Procgen 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 on games with positive correlation and games with negative correlation (each color represents a different game). For games with positive correlation, we see that the correlation is almost linear with a 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 is less clear that improving training performance will necessarily lead to better generalization or more severe overfitting. Most of the games with negative correlations are either saturated in terms of performance or too challenging for current methods to make much progress. This suggests that despite their apparent similarity, there are actually two 'categories' of games based on the algorithm's tendency to overfit. Note that this dichotomy likely depends on the algorithm; here, all algorithms are PPO-based and we hypothesize the picture would be different for value-based methods. It would be interesting to understand what structural properties of these CMDPs are responsible for this difference in overfitting, but this is out of the scope of the current work.

Another interesting thing to notice is that most of the outliers in the two plots are PLR [48] which is an adversarial method that actively tries to reduce training performance. This may explain why the trend does not hold for PLR as all other methods use a uniform distribution over the training environments.

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 µ ∈ M train. If that is the case, the suboptimality can accumulate across all the training MDPs, resulting in a large gap between these types of methods when trained in all the environments in M train. To test this hypothesis, we run both DQN and PPO on a single environment from each of the Procgen games. As shown in Figure 15b, we see that there is not a significant gap between the average performances of DQN and PPO on single environments (across 3 seeds for 10 environments on which both methods reach scores above 0.5). In fact, DQN 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 no transfer between different MDPs because the model can learn them separately. In practice, we embed the discrete task IDs into R 64 and add them as input to the final layers of the Q-network. Since there is no semantic relationship between discrete task IDs, we do not expect this to improve generalization performance, but, surprisingly, we found that it 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 is related to but not the same as the aforementioned value-overfitting 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 value-based methods since the policy is directly defined by the Q-function. An alternative solution is distributional RL [10] which learns a distribution (rather than a point estimate) over all possible Q-values for a given state. This models the aleatoric uncertainty resulting from the unobserved contexts in CMDPs ( i.e., not knowing which MDP the agent is interacting with), thus being able to account for one state having different potential values depending on the MDP. Our method uses distributional RL to mitigate the problem of value overfitting. As seen in Figure 5, while learning a value distribution leads to better results than learning only a 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 G ( α ) DQN).

Our analysis indicates that sample efficiency and partial observability are not the main reasons behind the poor generalization of value-based methods. In contrast, value overfitting is indeed a problem, but it can be alleviated with distributional RL. Nevertheless, effective exploration proves to be a key factor in improving the train and test performance of value-based approaches in CMDPs. Our results also suggest that better optimization on the training environment can lead to better generalization to new environments, making it a promising research direction.

0.2

0.0

0.2

0.4

0.6

0.8

1.0

train

0.2

0.0

0.2

0.4

0.6

0.8

1.0

test

Train vs. test return

bigfish

bossfight

caveflyer

chaser

climber

jumper

leaper

plunder

starpilot

H Uncertainty

Central to our method is the estimation of uncertainty in Q-values. We found that what kind of uncertainty to model and how we model it are crucial to the performance of EDE The first important consideration is what kind of uncertainty to model. Specifically, we argue that only epistemic 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 is intrinsic to the environment ( e.g., label noise) and cannot be reduced no matter how many samples we gather. In standard regression, even with ensembles, the estimated uncertainty ( i.e., variance) contains both epistemic and aleatoric uncertainty. In Atari, the main source of aleatoric uncertainty comes from the sticky action which has a relatively simple structure. However, in CMDPs such as 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 is detrimental to the test performance which may partially explain DQN+UCB's extremely poor performance (Figure 5), even compared to DQN+ ε -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 a 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 on a Tesla V100 (note that this only occurs at the training time since at test time all heads share the same feature extractor):

EDE 5 ensemble members):

0.3411 seconds per algorithm step

EDE 3 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 3 ensemble members, performs much worse than QR-DQN with 3 ensemble members trained with different minibatch and initialization (which are presumably more independent from each other). We believe further understanding this phenomenon would be an exciting future research direction.

0.2

0.0

0.2

0.4

0.6

0.8

1.0

train

0.2

0.0

0.2

0.4

0.6

0.8

1.0

test

Train vs. test return for

games with negative correlation

dodgeball

fruitbot

heist

maze

miner

ninja

(b) Games with negative train-test correlation.

0.0

0.5

1.0

1.5

2.0

1e7

0.05

0.00

0.05

0.10

0.15

0.20

0.25

2

vs.

2

epi

test performance

2

2

epi

(c) Comparing different uncertainties (mean test min-max normalized scores).

I Architecture and Hyperparameters

I.1 Procgen

Architecture. We use a standard ResNet [42] feature extractor ( f ) from IMPALA [27], which was adopted by Ehrenberg et al. [26] for value-based methods. The feature maps are flattened ( d = 2048 ) and processed by M copies of 2-layer MLPs with the same architecture to output M ×|A|× 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 (1) : R d → R 512 and the second layer is g (2) : R 512 → R |A|× N . The overall archetecture is

The hyperparameters for the main method is shown in Table 4, and are directly adapted from Ehrenberg et al. [26] other than those of new algorithmic components.

For hyperparameters specific to EDE, we search over the following values on the game bigfish :

φ : { 10 , 20 , 30 , 50 }

λ : { 0 . 5 , 0 . 6 , 0 . 7 }

α : { 6 , 7 , 8 , 9 }

and picked the best combination.

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:

For base DQN and QR-DQN, we use ε -greedy exploration with decaying ε . Specifically, at every algorithmic step, the epsilon is computed with the following equation:

For Bootstrapped DQN [79], we use bootstrapping to train the model and uniformly sample a value head for each actor at the beginning of each episode that is 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 use the standard variance when combined with DQN [19] and use φ = 30 .

For TEE without UCB, we use ε -greedy exploration without decaying ε for both QR-DQN and DQN, and EDE ith the same λ and α .

For ϵz -greedy, we use the same ε decay schedule in combination with n = 10000 and µ = 2 . 0 which are the best performing hyperparameter values from Dabney et al. [24]. Each parallel worker has its own ϵz -greedy exploration process.

For NoisyNet, we modify the base fully-connected layer with the noisy implementation [32] with σ = 0 . 5 .

I.2 Crafter

For Crafter, we use the default architecture and hyperparameters that are used for the Rainbow experiments in Hafner [37] with the following changes:

Change the distributional RL component from C51 to 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 to 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 φ , we searched over { 0 . 5 , 1 . 0 , 2 . 0 } and picked the best value.

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 a subset of games. First, without TEE , we study the sensitivity to different φ (UCB exploration coefficients), different M (number of ensemble members), and whether the feature extractor is trained with gradients from each ensemble member.

For φ and M , we run the algorithm on: bossfight, ninja, plunder, starpilot, caveflyer, jumper, bigfish, leaper for 1 seed and report the aggregated min-max normalized scores in Figure 17a and Figure 17b. We observe that the algorithm is slightly sensitive to 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 φ performs the best here, we use a larger value ( φ = 30 ) in combination with TEE.

In Figure 17c, we 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 1 seed. We observe that always training the feature extractor leads to lower performance, corroborating our intuition that the feature extractor should be trained at the same speed as the individual ensemble members.

In Figure 18, we study the performance under different hyperparameter values of TEE. We use fixed M = 5 and φ = 30 and vary the values of either α or λ while holding the other one fixed. We observe no significant difference across these, suggesting that the algorithm is robust to the values of α and λ .

K Hardware

The experiments are conducted on 2080 and V100 and take approximately 250 GPU days.

0.0

0.5

1.0

1.5

2.0

steps

1e7

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

0.45

return

Varying

= 10

= 30

= 50

(a)

φ

ablation.

0.0

0.5

1.0

1.5

2.0

steps

1e7

0.05

0.10

0.15

0.20

0.25

0.30

0.35

0.40

return

Varying

M

M

= 3

M

= 4

M

= 5

(b)

M

ablation.

0.0

0.5

1.0

1.5

2.0

steps

1e7

0.00

0.05

0.10

0.15

0.20

0.25

0.30

return

Always train feature extractor

True

False

(c) Feature extractor ablation.

0.0

0.5

1.0

1.5

2.0

steps

1e7

0.0

0.1

0.2

0.3

0.4

0.5

return

Varying

= 6

= 7

= 8

0.0

0.5

1.0

1.5

2.0

steps

1e7

0.0

0.1

0.2

0.3

0.4

0.5

return

Varying

= 0.5

= 0.6

= 0.7

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nThe 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.\n\nQUESTIONS:\n1. 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\nLIMITATIONS:\nIt is unclear which exploration is best for CMDPs and whether the performance of EDE is significantly better than the baselines.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n## 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.\n\nWEAKNESSES:\n## 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nThe paper discusses relevant limitations to the proposed algorithm.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-28"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n**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.\n\nWEAKNESSES:\nOne 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).\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nThe authors adequately addressed the limitations", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\noriginality: 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.\n\nWEAKNESSES:\nFigure 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\nQUESTIONS:\nYou 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\n\nLIMITATIONS:\nThere 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?", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- The paper is overall well-organized and easy to follow.\n- The proposed method is novel, well-motivated and empirically strong.\n\nWEAKNESSES:\n- 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.\n\nQUESTIONS:\nThe 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\\).\n\nLIMITATIONS:\nThe authors adequately addressed the limitations and potential negative societal impact of their work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}