SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yiding Jiang* Carnegie Mellon University yidingji@cs.cmu.edu
J. Zico Kolter  
Carnegie Mellon University  
zkolter@cs.cmu.edu
Roberta Raileanu Meta AI Research raileanu@meta.com
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 Progen 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.
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
*Work done while interning at Meta AI Research.
Figure 1: Exploration can help agents learn about parts of the environment which may be useful at test time, even if they are not needed for the optimal policy on the training environments. Note that this picture only illustrates one of many scenarios where exploration helps.
benchmarks for RL. In this work, we show that the agent's 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 Progen [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, Progen and Crafter.
Episodic Reinforcement Learning. A Markov decision process (MDP) is defined by the tuple  $\mathcal{M} = (\mathcal{S},\mathcal{A},R,P,\gamma ,\mu)$ , where  $\mathcal{S}$  is the state space,  $\mathcal{A}$  is the action space,  $R:S\times \mathcal{A}\to [R_{\min},R_{\max}]$  is the reward function,  $P:S\times \mathcal{A}\times \mathcal{S}\rightarrow \mathbb{R}_{\geq 0}$  is the transition distribution,  $\gamma \in (0,1]$  is the discount factor, and  $\mu :S\to \mathbb{R}_{\geq 0}$  is the initial state distribution. We further denote the trajectory of an episode to be the sequence  $\tau = (s_0,a_0,r_0,\dots ,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  $\pi :S\times \mathcal{A}\to \mathbb{R}_{\geq 0}$ ,  $Z^{\pi} = \sum_{t = 0}^{T}\gamma^{t}r_{t}$  is a random variable that describes the discounted return the policy achieves. The objective is to find a  $\pi^{\star}$  that maximizes the expected discounted return,  $\pi^{\star} = \arg \max_{\pi}\mathbb{E}_{\tau \sim p^{\pi}(\cdot)}[Z^{\pi}]$ , where  $p^\pi (\tau) = \mu (s_0)\prod_{t = 0}^T P(s_{t + 1}\mid s_t,a_t)\pi (s_t\mid a_t)$ . For simplicity, we will use  $\mathbb{E}_{\pi}$  instead of  $\mathbb{E}_{\tau \sim p^{\pi}(\cdot)}$  to denote the expectation over trajectories sampled
from the policy  $\pi$ . With a slight abuse of notation, we use  $Z^{\pi}(s, \pmb{a})$  to denote the conditional discounted return when starting at  $\pmb{s}$  and taking action  $\pmb{a}$  (i.e.,  $s_0 = s$  and  $\pmb{a}_0 = \pmb{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^{\pi}(s,a) = \mathbb{E}_{\pi}[Z^{\pi} \mid s_{0} = s, a_{0} = a]$ . The  $Q$ -function of a policy can be found at the fixed point of the Bellman operator,  $\mathcal{T}^{\pi}$  [11],  $\mathcal{T}^{\pi}Q(s,a) = \mathbb{E}_{s' \sim P(\cdot|s,a),a' \sim \pi(\cdot|s')} [R(s,a) + \gamma Q(s',a')]$ . The value function for a state  $s$  is defined as  $V^{\pi}(s) = \mathbb{E}_{a \sim \pi(\cdot|s)}[Q^{\pi}(s,a)]$ . Bellemare et al. [10] extends the procedure to the distribution of discounted returns,  $\mathcal{T}^{\pi}Z(s,a) \stackrel{d}{=} R(s,a) + \gamma Z(s',a')$ ,  $s' \sim P(\cdot \mid s,a)$  and  $a' \sim \pi(\cdot \mid s')$ , where  $\underline{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  $\pi(a \mid s) = \mathbb{1}_{\arg \max_{a'} Q(a',s)}(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  $\mu$ . In addition to the standard assumptions of a CMDP, we assume the existence of a structured distribution  $\Psi$  over the MDPs. During training, we are given a (finite or infinite) number of training MDPs,  $\tilde{\mathcal{M}}_{\mathrm{train}} = \{\mathcal{M}_1,\mathcal{M}_2,\dots ,\mathcal{M}_n\}$ , drawn from  $\Psi$  (Ghosh et al. [35] refers to this setting as epistemic POMDP). For each  $\mathcal{M}$ , we use  $Q_{\mathcal{M}}^{\pi}(s,a)$  and  $V_{\mathcal{M}}^{\pi}(s)$  to denote the Q-function and value function of that  $\mathcal{M}$ . We use  $p^{\pi ,\mathcal{M}}(\tau)$  to denote the trajectory distribution of rolling out  $\pi$  in  $\mathcal{M}$ . The objective is to find a single policy  $\pi$  that maximizes the expected discounted return over the entire distribution of MDPs,  $\mathbb{E}_{\tau \sim p^{\pi ,\mathcal{M}(\cdot)},\mathcal{M}\sim \Psi (\cdot)}\left[\sum_{t = 0}^{T}\gamma^{t}r_{t}\right]$  without retraining on the unseen  $\mathcal{M}$  (i.e., zero-shot generalization). We will refer to this quantity as the test return<sup>2</sup>.
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  $\mathcal{M}_{\mathrm{train}}$ , we can say that algorithm 1 generalizes better than algorithm 2. Furthermore, we assume the existence of  $\pi^{\star}$  (potentially more than one) that is  $\alpha$ -optimal for all  $\mathcal{M} \in \operatorname{supp}(\Psi)$  and all  $s_0$  in  $\operatorname{supp}(\mu_{\mathcal{M}})$ , i.e.,  $V_{\mathcal{M}}^{\pi^{*}}(s_{0}) \geq \max_{\pi} V_{\mathcal{M}}^{\pi}(s_{0}) - \alpha$  where  $\alpha$  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 \times 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
It may be useful to contrast this with the notion of generalization gap, the difference between training and test returns. Without good training performance, the generalization gap is not a meaningful measure: a random policy would have no generalization gap at all. The generalization gap is important for supervised learning because we can get nearly perfect training performance, but this is not the case for RL.
(a) Tabular CMDP
(b)  $\varepsilon$ -greedy
Figure 2: (a) A tabular CMDP that illustrates the importance of exploration for generalization in RL. During training, the agent starts in the blue square, while at test time it starts in the orange square. In both cases, the goal is to get to the green square. The other plots show the mean and standard deviation of the train and test suboptimality (difference between optimal return and achieved return) over 100 runs for (b) Q-learning with  $\varepsilon$ -greedy exploration, (c) Q-learning with UCB exploration.
(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  $\varepsilon$ -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  $\varepsilon$ -greedy strategy which takes a random action with probability  $\varepsilon$  and the best action according to the Q-function  $\arg \max_{a} Q(s, a)$  with probability  $1 - \varepsilon$ . In contrast, UCB is uncertainty-driven so it explores according to  $\pi(a \mid s) = \mathbb{1}(a = \arg \max_{a'} Q(s, a') + c \sqrt{\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<sup>3</sup>. 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 hyperparameter. 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<sup>4</sup>. 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
This UCB is a simple extension of the widely used bandits algorithm to MDP and does not enjoy the same regret guarantee, but we find it to be effective for the didactic purpose. Azar et al. [7] presents a formal but much more complicated extension of the UCB algorithm for value iteration on MDPs. The UCB used is more similar to a count-based exploration bonus.
4While the training and test MDPs are seemingly drawn from different distributions, the setting can be made IID by taking a mixture of the train and test MDP. Since there is a significant discrepancy between train and test performance, all the observations would remain valid.
Figure 3: A simplified version of bigfish in Procgen (Figure 6) that captures the spirit of the grid world example from Figure 2a. The green dot represents the agent and the red dots represent the enemies (see Appendix B for more details). The goal is to eat smaller dots while avoiding larger dots. An agent that explores both trajectories in the training MDP could recover the optimal behavior on the test MDP whereas an agent that only focuses on solving the training MDP would fail.
suboptimal states that are rarely visited by simple exploration on the training MDPs. Having explored such states during training can help the agent recover from suboptimal states at test time. See Figure 3 and Appendix B for an illustrative example inspired by one of the Progen 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(\pmb{\theta} \mid \mathcal{D})$ , where  $\pmb{\theta}$  is the discretized quantiles of  $Z(s, a)$ , and  $\mathcal{D}$  is the dataset of experience on which the quantiles are esti-
mated. Let  $j \in [N]$  denote the index for the  $j^{\text{th}}$  quantile and  $\mathcal{U}\{1, N\}$  denote the uniform distribution over integers between 1 and  $N$ . The uncertainty of the Q-value,  $Q(\boldsymbol{s}, \boldsymbol{a}) = \mathbb{E}_{j \sim \mathcal{U}\{1, N\}}[\theta_j(\boldsymbol{s}, \boldsymbol{a})]$ , is the relevant quantity that can inform exploration. The overall uncertainty  $\sigma^2 = \operatorname{Var}_{\boldsymbol{\theta} \sim p(\boldsymbol{\theta}|\mathcal{D})}[Q(\boldsymbol{s}, \boldsymbol{a})]$  can be decomposed into epistemic uncertainty  $\sigma_{\mathrm{epi}}^2(\boldsymbol{s}, \boldsymbol{a})$  and aleatoric uncertainty  $\sigma_{\mathrm{ale}}^2(\boldsymbol{s}, \boldsymbol{a})$  such that  $\sigma^2(\boldsymbol{s}, \boldsymbol{a}) = \sigma_{\mathrm{epi}}^2(\boldsymbol{s}, \boldsymbol{a}) + \sigma_{\mathrm{ale}}^2(\boldsymbol{s}, \boldsymbol{a})$ , where,
$$
\sigma_ {\mathrm {e p i}} ^ {2} (\boldsymbol {s}, \boldsymbol {a}) = \mathbb {E} _ {j \sim \mathcal {U} \{1, N \}} \left[ \operatorname {V a r} _ {\boldsymbol {\theta} \sim p (\boldsymbol {\theta} | \mathcal {D})} \left[ \boldsymbol {\theta} _ {j} (\boldsymbol {s}, \boldsymbol {a}) \right] \right], \sigma_ {\mathrm {a l e}} ^ {2} (\boldsymbol {s}, \boldsymbol {a}) = \operatorname {V a r} _ {j \sim \mathcal {U} \{1, N \}} \left[ \mathbb {E} _ {\boldsymbol {\theta} \sim p (\boldsymbol {\theta} | \mathcal {D})} \left[ \boldsymbol {\theta} _ {j} (\boldsymbol {s}, \boldsymbol {a}) \right] \right]. \tag {1}
$$
Ideally, given an infinite or sufficiently large amount of diverse experience, one would expect the posterior to concentrate on the true quantile  $\theta^{\star}$ , so  $\mathrm{Var}_{\theta \sim p(\theta | \mathcal{D})}[\theta_j(s, a)]$  and consequently  $\sigma_{\mathrm{epi}}^2(s, a) = 0$ .  $\sigma_{\mathrm{ale}}^2$  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<sup>5</sup> is beneficial for exploration. We refer interested readers to Charpentier et al. [17] for a more thorough discussion.