Uncertainty Prioritized Experience Replay

We thank the reviewer for their constructive review. We are glad you find the method insightful. We hope to address some of your concerns around clarity, phrasing, and evaluations below.

W1 [also addresses Q2]: Our prioritization scheme is not limited to use with quantile representations of the distribution. As noted in your review the method is in principle independent of the underlying algorithm. However, the quality of our epistemic and aleatoric uncertainty estimates do matter. The specific algorithms you mention (C51 and rainbow) are categorical distributional algorithms, which are well known to overestimate the variance when used in discounted bootstrapping (see for instance Fig. 2 of Rowland et. al 2019). This overestimate will significantly impact the quality of our information gain criterion and so we do not expect UPER to perform well with C51 or rainbow. We would fully expect a quantile regression version of rainbow to work well with UPER, but developing this agent was beyond the scope of this paper.

W2: We appreciate the possible confusion associated with Table 1. It is worth clarifying that the computational costs associated with the methods we are discussing can be decomposed into two parts: first is the cost associated with performing prioritized sampling. These costs are discussed in the original PER paper, and amount to approximately “2-4% increase in running time and negligible additional memory usage”. This part of the cost is identical in our method, so we do not include our own analysis. The second component of the cost pertains to the computation of the prioritization variables themselves. In PER there is no additional cost in this component compared to vanilla DQN since the TD error is already computed. In our method however we require an ensemble of distributional estimators to compute our information gain prioritization criterion. This cost is the focus of the comparisons shown in Table 1. Since the majority of our network is shared across distribution and ensemble heads, the additional cost is insignificant, particularly when leveraging parallelization of GPU operations. We have added a sentence to the manuscript to clarify the implications of Table 1.

Q1: We appreciate the feedback on the readability of the introduction. We have updated the introduction to be more comprehensive and better sign-post our contributions. In particular, we give intuitions for the use of epistemic and aleatoric uncertainty under the information gain criterion in prioritising replay. In addition to the noisy TV thought experiment borrowed from the exploration literature to explain the potential shortcomings of PER, we discuss a toy game involving estimating two distributions to help motivate the intuition behind the information gain. We also mention more explicitly and sooner in the narrative the particulars of our methodology i.e. ensembles and distributions. Hopefully, this helps to address your concerns regarding the depth of the introduction.

Q3: This is a fair distinction to raise. We have amended the phrasing to ‘Prioritized experience replay is a crucial component of contemporary value-based deep reinforcement learning models’. On the broader point about PER in policy-gradient methods: it has generally been difficult to use PER in policy gradient methods, especially when training the actor (Saglam et. al 2024). However approaches such as phasic policy gradient methods (Cobbe et. al 2020), which splits training into two phases, may allow prioritizing when training the critic without high variance updates from large error transitions harming training of the policy network.

Q4: Thank you for catching these oversights in the style of some of our citations. We have gone through the manuscript and updated the citation commands to match the ICLR guidelines (as outlined here: https://arxiv.org/html/2409.06957v1). We hope this improves the readability of the paper–for instance in the introductory passage you highlighted.

Refs

Statistics and Samples in Distributional Reinforcement Learning (Rowland et. al 2019)

Actor Prioritized Experience Replay (Saglam et. al 2024)

Phasic Policy Gradient (Cobbe et. al 2020)

Thank you for providing meticulous revisions and responses to my requests.

However, it is still crucial to demonstrate that the proposed method is effective across various RL distributional methods, even if some results show only marginal improvement. Otherwise, the method will be viewed as being overly specialized for QR-DQN, thus, its impact may be regarded as too limited by the research community.

Thanks for your response. We are glad to have been able to address the majority of your concerns. The extended rebuttal period did give us the necessary time to perform some additional experiments. Please see our general response that shows significant improvements to C51 with UPER. Thanks again for your engagement and feedback. We hope these experiments address the last of your concerns.

Thank you for your detailed review and questions. We appreciate that you found our work novel and supported by comprehensive experiments. However, we realize we could have been clearer about the goal of our proposed method.

UPER modifies the prioritization variable to account for uncertainty when sampling transitions from the memory buffer for updating the agent’s Q-value estimates. It does not explicitly enhance exploration or alter the agent's interaction with its environment, though prioritizing uncertainty-based transitions might have such effects indirectly. As a result, while UPER is related to exploration literature, we do not consider it an exploration algorithm, as we do not promote exploration directly through methods like rewarding unvisited, uncertain, or high-error states. Instead, our focus is on how transitions are sampled from the buffer for action-value function updates, borrowing ideas from the exploration-exploitation trade-off literature to estimate the quantities needed for our method, epistemic and aleatoric uncertainty. We extended the introduction and discussion section to clarify this. We provide detailed answers to your comments below and hope this clarification helps highlight the contribution of our method. Please let us know any further concerns.

Comparison with other baselines

In [3], the authors use pseudo-counts to quantify state uncertainty. This measure does not distinguish between epistemic and aleatoric uncertainty but is instead used to reward unvisited states with an exploration bonus (as described in the first equation of [3]). Unlike this approach, UPER does not modify state rewards based on novelty or uncertainty. Regarding the application of our method to C51, BBF, or Rainbow, methods built on categorical distribution estimation, we do not expect UPER to perform well. This is because these agents rely on fixed support for categorical distribution estimation, and when bootstrapping with discounting, the target distribution needs to be projected onto the support of the estimated distribution; here the variance is inherited and tends to be overestimated. This is a well documented phenomenon (see e.g. Rowland et. al 2019). It may have been possible to try UPER on a quantile regression version of Rainbow, this would have required developing a new agent entirely, which we felt was beyond the scope of this paper. Meanwhile, QR-DQN is generally a far stronger algorithm than C51 (see original QR-DQN paper).

Q1: Thank you for noticing this. The MSE shown in Figure 1b is the average across all arms, calculated between the estimated mean and the true mean (i.e., the mean used to sample rewards from the arm). By definition, this represents the expected generalization error and is not subject to overfitting, and it was used to show the performance for each method exactly for this reason. We included this explanation in the text. Additionally, we revised Figure 1 and its explanations to improve clarity, as well as the other suggested changes in the writing style.

Q2: We compare our method with the one you refer to as “naive prioritized DQN” [1]. In our work, we call it Prioritized Experience Replay (PER), as shown in Figure 2 and mentioned in the first paragraph of Section 5. Regarding Atari games that require extensive exploration, as we stated earlier, our method does not explicitly promote exploration. It only alters the prioritization variable to sample transitions from the replay buffer. Therefore, we do not expect our method to improve performance on these games, as it was not designed to incentivize exploration explicitly. One example of such a game is Montezuma’s Revenge, as mentioned by reviewer C2gj. This game has been solved by explicitly incentivizing exploration, such as by rewarding unvisited, high-uncertainty, or high-error states, or through imitation learning. For a review of these methods and ideas, please refer to our response to reviewer C2gj, W2.

Q3: Yes, you are correct. The total uncertainty $\hat{\mathcal{U}}{\delta}$ in Equation 9 depends on the error $\delta{\Theta}$, the ensemble disagreement (epistemic uncertainty) $\hat{\mathcal{E}}$, and the aleatoric uncertainty $\hat{A}$. Here, $\delta_{\Theta}$ represents the TD error between the true target and the combined estimation from all quantiles and ensembles. Considering the expected error across quantiles and ensembles as total uncertainty introduces an extra term, yielding the target epistemic uncertainty $\mathcal{E}{\delta}$, which accounts for both the ensemble disagreement and the distance to the estimation quantity. Compared to previous literature, which typically considers only the decomposition in Equation 7 (taking total uncertainty as the variance of the prediction), empirical results (see App E) demonstrate that using the target epistemic uncertainty $\mathcal{E}{\delta}$ is more effective than relying solely on the original epistemic uncertainty $\mathcal{E}$ from Equation 8.

Theoretical Analysis

We agree that theoretical analysis could help substantiate and guarantee improvements to our method. While we explored this during development, it remains challenging due to the coupling between value estimation and buffer sampling. Existing work on replay buffer dynamics ([1] from reviewer C2gf) would require incorporating epistemic and aleatoric uncertainty, which is beyond our proposal's scope. Nonetheless, we demonstrate our method's effectiveness through toy models, simulations, and ablation studies. For instance, Appendix G.1 shows that performance improvements on the Atari-57 benchmark stem solely from the prioritization variable. Additionally, Appendix D provides insights into the effects of prioritization based on uncertainty.

UPER is not an ad hoc solution as it was derived in a principled way. The epistemic and aleatoric uncertainty decomposition, derived from the expected error in Equation 9, along with the information gain, are both mathematically derived (see Appendices B and D.1), hence UPER is a general prioritization scheme, and its quality will depend on the method used to estimate the distribution.

Refs

Statistics and Samples in Distributional Reinforcement Learning (Rowland et. al 2019)

Thank you for your review. We are pleased that you found the experiments to clearly demonstrate UPER's advantages and that you consider the paper well-written. You raise valid points regarding the comparison with [1], the performance of Montezuma’s Revenge, and the applicability of our proposed method to vanilla DQN. We are committed to addressing each of your concerns in detail, and we will incorporate these responses into the revised manuscript. Please find the details below:

W1: Comparison with Langevin Dynamics Sampling method [1]

In the referenced study, the authors address limitations of PER that are unrelated to uncertainty. Specifically, they tackle two issues: (1) outdated priorities due to batch-like priority updates in the replay buffer, and (2) insufficient sample space coverage, especially in high-dimensional problems. To address these, they train a model based on a Langevin dynamical system to estimate state transitions. This model generates "ideal transitions" with updated priorities and better sample space coverage. We argue that the strengths of their method are complementary to uncertainty estimation. While our approach may face issues with outdated priorities and sample space coverage, their method will be subject to aleatoric uncertainty. Although integrating their advantages into our method seems like a promising research direction, this falls beyond the scope of our paper and could merit an entirely new study.

W2: Montezuma’s Revenge.

As you correctly noted, Montezuma’s Revenge is a challenging game for RL agents due to its sparse rewards and the need for long-term strategies, requiring deep exploration of possible trajectories. Several approaches have been proposed to address this, such as imitation learning in “Playing Hard Exploration Games by Watching YouTube” (Aytar et al., 2018), promoting exploration in “Go-Explore” (Ecoffet et al., 2021), and assigning intrinsic rewards to uncertain states in “Flipping Coins to Estimate Pseudocounts for Exploration in Reinforcement Learning” (Lobel et al., 2023). These strategies generally involve updating policies with highly exploratory trajectories. While UPER affects the resulting policy after training, it does not explicitly promote exploration. Therefore, we do not expect our method to improve performance in Montezuma’s Revenge, as its primary goal is prioritizing memory buffer transitions based on uncertainty measures, not fostering exploration. However, as we mentioned earlier, our approach could complement exploration-focused methods, making this a promising direction for future research. We will include this important discussion in the revised manuscript.

W3: Computational cost of PER, UPER, and uniform sampling.

PER and UPER are indeed more computationally expensive than uniform sampling, but the additional cost is minimal. For PER, the main extra calculations involve computing priority probabilities (Equation 2) and maintaining priorities for all transitions in the buffer. According to the original Prioritized Experience Replay paper (Schaul et al., 2016, Appendix B.2.1), the runtime overhead compared to uniform sampling is only 2%-4%, with negligible memory impact. Since we use DeepMind’s DQN Zoo implementation, our method incurs the same small overhead.

Q1: Can UPER be used in a Vanilla DQN?

This is an excellent question, and the short answer is "probably not." Uncertainty measures are typically derived from higher-order moments of the underlying and estimated probability distribution. A standard DQN agent estimates only the mean and lacks a concept of how volatile this estimate might be. To capture such volatility, the agent must estimate moments beyond the mean.

This is why many works extend the vanilla DQN architecture. For example:

• Direct Epistemic Uncertainty Prediction (Section 2.3.3) trains a network to predict sample-specific errors.
• Lobel et al. (2023) add a network component to estimate counts as a proxy for uncertainty.
• Direct Variance Estimation (Appendix A.1) introduces a separate network to estimate return variance.
• Bootstrapped DQN (Section 2.3.1) uses ensembles to approximate the posterior distribution and capture higher-order moments.
• Quantile regression (Section 2.3.2) includes a parameter for estimating different parts of the distribution through quantiles. These examples highlight the need for additional mechanisms to estimate uncertainty in Q-learning. Thus, a standard DQN alone cannot provide precise measures of epistemic and aleatoric uncertainty.