SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sampling from  $p(\pmb{\theta} \mid \mathcal{D})$  is computationally intractable for complex MDPs and function approximators such as neural networks. Clements et al. [20] approximates samples from  $p(\pmb{\theta} \mid \mathcal{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}: \mathbb{R}^{d} \to \mathbb{R}^{|A| \times N}$  that share a single feature extractor  $f: S \to \mathbb{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 \in [M]$  is the index for  $M$  ensemble heads of the Q-network, and  $j \in [N]$  is the index of the quantiles. The output of the  $i^{\text{th}}$  head for state  $s$  and action  $a$  is  $\theta_{i}(s, a) \in \mathbb{R}^{N}$ , where the  $j^{\text{th}}$  coordinate, and  $\theta_{ij}(s, a)$  is the  $j^{\text{th}}$  quantile of the predicted state-action value distribution for that head. The finite sample estimates of the two uncertainties are:
$$
\hat {\sigma} _ {\mathrm {e p i}} ^ {2} (\boldsymbol {s}, \boldsymbol {a}) = \frac {1}{N \cdot M} \sum_ {j = 1} ^ {N} \sum_ {i = 1} ^ {M} \left(\boldsymbol {\theta} _ {i j} (\boldsymbol {s}, \boldsymbol {a}) - \bar {\boldsymbol {\theta}} _ {j} (\boldsymbol {s}, \boldsymbol {a})\right) ^ {2}, \hat {\sigma} _ {\mathrm {a l e}} ^ {2} (\boldsymbol {s}, \boldsymbol {a}) = \frac {1}{N} \sum_ {j = 1} ^ {N} \left(\bar {\boldsymbol {\theta}} _ {j} (\boldsymbol {s}, \boldsymbol {a}) - Q (\boldsymbol {s}, \boldsymbol {a})\right) ^ {2}, \tag {2}
$$
where  $\bar{\theta}_j(\pmb{s},\pmb{a}) = \frac{1}{M}\sum_{i = 1}^{M}\pmb{\theta}_{ij}(\pmb{s},\pmb{a})$  and  $Q(\pmb {s},\pmb {a}) = \frac{1}{N}\sum_{j = 1}^{N}\bar{\theta}_j(\pmb {s},\pmb {a})$
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:
$$
\pi_ {\mathrm {t s}} (\boldsymbol {a} \mid \boldsymbol {s}) = \mathbb {1} _ {\arg \max  _ {\boldsymbol {a} ^ {\prime}} \xi (\boldsymbol {s}, \boldsymbol {a} ^ {\prime})} (\boldsymbol {a}), \quad \text {w h e r e} \quad \xi (\boldsymbol {s}, \boldsymbol {a} ^ {\prime}) \sim \mathcal {N} (Q (\boldsymbol {s}, \boldsymbol {a} ^ {\prime}), \varphi \hat {\sigma} _ {\mathrm {e p i}} (\boldsymbol {s}, \boldsymbol {a} ^ {\prime})). \tag {3}
$$
$\varphi \in \mathbb{R}_{\geq 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]):
$$
\pi_ {\mathrm {u c b}} (\boldsymbol {a} \mid \boldsymbol {s}) = \mathbb {1} _ {\boldsymbol {a} ^ {\star}} (\boldsymbol {a}), \quad \text {w h e r e} \boldsymbol {a} ^ {\star} = \underset {\boldsymbol {a} ^ {\prime}} {\arg \max } Q (\boldsymbol {s}, \boldsymbol {a} ^ {\prime}) + \varphi \hat {\sigma} _ {\mathrm {e p i}} (\boldsymbol {s}, \boldsymbol {a} ^ {\prime}), \tag {4}
$$
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
It may still contain some aleatoric uncertainty due to the heteroskedasticity of the problem.
Figure 4: Test performance of different methods on the Progen benchmark across 5 runs. Our method greatly outperforms all baselines in terms of median and IQM (the more statistically robust metrics) and is competitive with the state-of-the-art policy optimization methods in terms of mean and optimality gap. The optimality gap is equal to 1-mean for the evaluation configuration we chose.
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^{\mathrm{th}}$  actor collects experience with the following policy:
$$
\pi_ {\mathrm {u c b}} ^ {(k)} (\boldsymbol {a} \mid \boldsymbol {s}) = \mathbb {1} _ {\boldsymbol {a} ^ {\star}} (\boldsymbol {a}), \quad \text {w h e r e} \boldsymbol {a} ^ {\star} = \underset {\boldsymbol {a} ^ {\prime}} {\arg \max } Q (\boldsymbol {s}, \boldsymbol {a} ^ {\prime}) + \left(\varphi \lambda^ {1 + \frac {k}{K - 1} \alpha}\right) \hat {\sigma} _ {\mathrm {e p i}} (\boldsymbol {s}, \boldsymbol {a} ^ {\prime}). \tag {5}
$$
$\lambda \in (0,1)$  and  $\alpha \in \mathbb{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  $\varepsilon$  for the  $\varepsilon$ -greedy exploration for each actor. In practice, the performances are not sensitive to  $\lambda$  and  $\alpha$  (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<sup>6</sup>. 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 naively combining them performs significantly worse.
5.1 Procgen
We compare our method with 3 representative baselines on the standard Procogen 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 Procogen [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 Procogen 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 Procogen. 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
UCB and Thompson sampling have similar regret guarantees but their actual performance can be problem-dependent [16].
Figure 5: Test performance of different exploration methods on the Progen benchmark across 5 runs.
compare with other popular exploration techniques for value-based algorithms. Since many of the existing approaches are designed for DQN, we also adapt a subset of them to QR-DQN for 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)  $\epsilon 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<sup>7</sup>. 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  $\sigma^2$  performs significantly worse at test time than exploring with only the epistemic uncertainty  $\sigma_{\mathrm{epi}}^2$ . UA-DQN also performs worse than QR-DQN+UCB suggesting that deep ensemble may have better uncertainty estimation (Appendix H).
QR-DQN with  $\epsilon z$ -greedy exploration marginally improves upon the base QR-DQN, but remains significantly worse than EDE. This may be due to the fact that  $\epsilon 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  $\epsilon z$ -greedy which implicitly use an schedule.
Taiga et al. [107] found that, on the whole Atari suite, NoisyNet outperforms more specialized exploration strategies such as ICM [86] and RND [14]. This suggests that these exploration methods may have overfitted to sparse reward environments.
5.2 Crafter
To test the generality of our method beyond the Progen 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.
Table 1: Results on Crafter after 1M steps and over 10 runs.
Method	Score (%)
EDE (ours)	11.7 ± 1.0
Rainbow	4.3 ± 0.2
PPO	4.6 ± 0.3
DreamerV2	10.0 ± 1.2
LSTM-SPCNN	12.1 ± 0.8
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], forward-backward 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  $\varepsilon$ -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-exploitation 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].
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 Progen 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 Progen 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.
We would like to thank Alessandro Lazaric, Matteo Pirotta, Samuel Sokota, and Yann Ollivier for the helpful discussion during the early phase of this project. We would also like to thank the members of FAIR London for their feedback on this work and the members of Locus Lab for their feedback on the draft.
[1] R. Agarwal, M. C. Machado, P. S. Castro, and M. G. Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. 2021.
[2] R. Agarwal, M. Schwarzer, P. 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.
[3] A. Ajay, G. Yang, O. Nachum, and P. Agrawal. Understanding the generalization gap in visual reinforcement learning. 2021.
[4] A. Anand, J. Walker, Y. Li, E. Vétres, 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.
[5] 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.
[6] P. Auer, N. Cesa-Bianchi, and P. Fischer. Finite-time analysis of the multiarmed bandit problem. Machine learning, 47(2):235-256, 2002.
[7] M. G. Azar, I. Osband, and R. Munos. Minimax regret bounds for reinforcement learning. In International Conference on Machine Learning, pages 263–272. PMLR, 2017.
[122] T. Zhang, H. Xu, X. Wang, Y. Wu, K. Keutzer, J. E. Gonzalez, and Y. Tian. Noveld: A simple yet effective exploration criterion. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https://openreview.net/forum?id=CYUzpn0kFJp.
[123] 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,  $\pi_{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  $\pi_{b}$ . Intuitively, the behavioral policy is responsible for exploration. We will first define some useful quantities that relate to the inherent properties of  $\pi_{b}$  and the MDP. Recall that  $\mu$  is an initial state distribution and let  $p^{\pi}(\tau \mid s)$  be the trajectory distribution induced by  $\pi$  starting from  $s$ , we can define the state marginal distribution at time  $t$  given that the initial state is drawn from  $\mu$  as  $P_{t}(s;\pi ,\mu) = \mathbb{E}_{\tau \sim p^{\pi (\cdot |s_{0}),s_{0}\sim \mu}}[\mathbb{1}\{s_{t} = s\} ]$ , the state-action marginal distribution  $P_{t}(s,a;\pi ,\mu) = \pi (a\mid s)P_{t}(s;\pi ,\mu)$ , the stationary state occupancy distribution,  $d_{\pi}(s) = \lim_{T\to \infty}\frac{1}{T + 1}\sum_{t = 0}^{T}P_{t}(s;\pi ,\mu)$  and finally the stationary state-action distribution  $d_{\pi}(s,a) = \pi (a\mid s)d_{\pi}(s)^{8}$ . Using  $D_{\mathrm{TV}}$  to denote the total variation distance, we define:
$$
\mu_ {\min } (\pi) = \min  _ {(s, a) \in \mathcal {S} \times \mathcal {A}} d _ {\pi} (s, a), \tag {6}
$$
$$
t _ {\operatorname {m i x}} (\pi) = \min  \left\{t \left| \max  _ {(s _ {0}, a _ {0}) \in \mathcal {S} \times \mathcal {A}} D _ {\mathrm {T V}} \left(P _ {t} (s, a; \pi , \delta \left(s _ {0}, a _ {0}\right)), d _ {\pi} (s, a)\right) \leq \frac {1}{4} \right. \right\}. \tag {7}
$$