The Benefits of Being Categorical Distributional: Uncertainty-aware Regularized Exploration in Reinforcement Learning

We would like to express our gratitude for your time and effort in reviewing our paper. We really appreciate the insightful and detailed comments and suggestions you have provided, and we would like to address each comment in the section on Weaknesses and Questions you raised in your review.

Weakness

1. The paper is not easy to follow. Many words are vague. This is especially noticeable in section 4.2. For instance, the phrase "characterizing the impact of action-state return distribution beyond its expectation" is vague to me and may need explanation (Line 196)

Thank you for highlighting this clarity issue. Based on your suggestion, we have revised this statement in the updated paper to better explain the impact of the induced histogram function on the behavior difference between distributional and classical RL. We have also thoroughly proofread the entire paper to improve clarity, and we would appreciate any further feedback if there are still vague statements or terms.

2. There have been prior works on the theory of distributional RL, including those referenced around Line 50 (Rowland et al, 2018, 2023, 2024). Therefore, the statement in the abstract is not very accurate: "Despite the remarkable empirical performance of distributional reinforcement learning (RL), its theoretical advantages over classical RL remain elusive."

We apologize for the statement in the abstract that may have been inaccurate. Based on your suggestion, we have revised it to 'Despite the remarkable empirical performance of distributional reinforcement learning (RL), its theoretical advantages over classical RL are still not fully understood.'

3. When reading Line 55, I found the authors mentioning "Yet, their findings are not directly based on typical distributional RL algorithms commonly used in practice, such as C51 or QR-DQN". This implies that this authors is probably going to analyze a typical distribution RL algo like C51. However, the paper instead focuses on DERPI and develops DERAC in section 5.2. Both are different from C51 (although they have some similarities). Thus, I feel (1) Line 55 is not precise enough and (2) the paper's scope is somewhat narrow, given that categorical distributional RL is already a specific subset within the field.

We want to clarify that we indeed study C51 and investigate its behavior differences from DQN using return density decomposition, particularly in Section 4. Section 5 further explores the decomposed regularization effect within the policy gradient framework, and, as a result, we introduce DERPI and DERAC as byproduct algorithms to guarantee the convergence when incorporating the uncertain-aware regularization term. As also mentioned in Section 6.2, the primary purpose of introducing DERAC is to interpret the benefits of CDRL from the perspective of uncertain-aware regularized exploration rather than to pursue empirical superiority. Therefore, while the two follow-up algorithms are proposed following our theoretical analysis, our main goal remains to interpret the benefits of the practical CDRL algorithm.

Although CDRL is a specific class of distributional RL, it is recognized as the first successful implementation in distributional RL, with broad recognition in the community. Unfortunately, the benefits of CDRL are not yet fully understood by the current research community. Therefore, it is valuable to focus on CDRL for a rigorous analysis as the first step before extending our analysis to other algorithmic classes. To further enhance our contribution, we have provided a discussion about how to decompose quantile-based distributional RL in Appendix M of the revised paper. We invite the reviewer to refer to this discussion, and we welcome any further feedback.

Thank you for taking the time to review our paper. We appreciate your positive assessment and insightful feedback, and we would like to address the concerns you raised in your review.

Weakness

1. Histogram function estimator is sensitive to bin widths, yet it is unclear how the bin widths in the experiments were selected. I suggest that an ablation study varying the bin widths be carried and report how it impacts the results - this would be informative.

The C51[1] is a typical CDRL algorithm, which uses 51 atoms (equivalent to 51 bins). We thus also use 51 bins in our algorithm implementation to maintain the original favorable performance of C51. To further address your concern, we have conducted an ablation study by varying the bin sizes (number of atoms) in two Atari games, with the results presented in Figure 10 of Appendix L.4. These results suggest that the performance of the new ablation study is consistent with our reported results in Section 6.1. Therefore, our conclusion regarding the regularization effect from the return density decomposition is robust against different bin sizes.

2. Please provide intermediate steps in equation (29) for proof of proposition 2 that shows how the constant disappeared. It currently unclear and would be great if it is explicitly shown.

We would like to clarify that in Eq. 29 (Eq 28 in the updated version), $1(x\in\Delta^i_E) / \Delta \log q_\theta^{s_i, a_i}(\Delta_j) / \Delta$ represents $\frac{1(x\in\Delta^i_E)}{\Delta} \log \frac{q_\theta^{s_i, a_i}(\Delta_j)}{\Delta}$ and we apologize for this potential confusion. To eliminate this ambiguity, we updated this formula in the revised paper. Regarding how $\Delta^{-2}$ disappears, the first $\frac{1}{\Delta}$ disappears as we take the integral $\int_{l_i}^{l_{i+1}}$. The second one is within the $\log$, and serves as a constant in the end, which does not affect the optimization.

Questions

1. How can the p_E (coefficient on bin \Delta_E) be computed?

Given the target return distribution of $Y_i^k$, which is fixed in each phase of Neural FZI, we can evaluate the expectation of the histogram density function by $E[Y_i^k] = \sum_{i=0}^{N-1} \frac{l_l + l_{i+1}}{2} p_i.$ Then we index the bin $\Delta_E$ that contains $E[Y_i^k]$. Consequently, the probability (the height times the bin size) on that bin is $p_E$. This computation process is also consistent with what we have done in our implementation.

2. What are the effects of various bin sizes on the results?

Please refer to the response to Weakness 1.

In summary, thank you again for pointing to these potential areas of improvement. We very much appreciate your suggestions. Please let us know if you have any further comment or feedback.

Reference

[1] Bellemare, Marc G., Will Dabney, and Rémi Munos. "A distributional perspective on reinforcement learning." International conference on machine learning. PMLR, 2017.

Thank you very much for taking the time to review our paper. We appreciate your positive assessment and insightful feedback, and we would like to address the concerns you raised in your review.

Question 1: Some theoretical results are achieved when $\Delta$ (the bin size) tends to zero. As $\Delta$ approaches zero, $P_E$ typically approaches zero. What happens to the hyperparameter $\epsilon$ in Proposition 1 and $\alpha$ in Equation (4) of Propositions 2 and 3 in the limit case? It is stated on page 4 that “Proposition 1 demonstrates that the return density distribution is valid when the hyperparameter $\epsilon$ is specified as $\epsilon \geq 1 - P_E$. Does this imply that $\epsilon$ needs to be close to 1 when $P_E$ is very small?

We acknowledge that as $\Delta \rightarrow 0$, $P_E \rightarrow 0$ and $\epsilon \rightarrow 1$ ($\alpha \rightarrow \infty$) to ensure a valid histogram function, satisfying $\epsilon \geq 1 - p_E$ in the return density decomposition. Note that this proposed return density decomposition is valid for any non-trivial $\epsilon$ as long as $\epsilon \geq 1 - p_E$, which allows for developing practical algorithms if needed. In contrast, the results presented in Proposition 2 and 3 provide an asymptotically approximated connection between optimizing the term (a) in the decomposed distribution loss in Eq.3 (the updated paper) and optimizing the classical RL objective function.

We understand the review may be concerned about the gap between the asymptotic connection in our analysis with the practical algorithm implementation. In fact, we have considered this potential concern in the submitted paper and have provided the rate of convergence $o(\Delta)$ in terms of the absolute difference of cumulative distribution function (CDF) between the converging CDF of term (a) and the step function (the CDF of a Dirac Delta function) in Appendix G. In the process of $\Delta \rightarrow 0$, the minimizer of optimizing term (a) gradually tends to that of optimizing classical RL, with the error shrinking at the rate of $o(\Delta)$. Please review Appendix G for more details.

Question 2: In Section 7, the authors mention the potential to extend their work to other distributional RL methods, such as QR-DQN. Given that distributional dynamic programming with quantile representation is not mean-preserving, the extension to QR-DQN may be problematic. The authors need to address why they believe their approach could extend to QR-DQN and whether they have considered using a normal representation instead. Distributional dynamic programming with a normal representation is mean-preserving, making it a more suitable candidate for extending this work.

Thank you for raising this interesting and technical point regarding our future work. We still believe that our analysis can be extended quantile-based distributional RL at least partially, though the techniques involved may be different. We acknowledge that distributional dynamic programming with the quantile representation is not $exactly$ mean-preserving as quantiles are non-linear functionals of distribution and focus on capturing the shape of distributions. However, we emphasize the fact that quantile dynamic programming can be asymptotically mean-preserving owing to the property that the expected quantiles can converge to the true expectation of the distribution as the number of quantiles approaches infinity. For a given random variable $X$, we have $\lim_{N\rightarrow+\infty} \frac{1}{N} \sum_{i=1}^{N} F^{-1}(\frac{i}{N+1}) = \int_{0}^{1} F^{-1}(\tau) d \tau = \int_{-\infty}^{+\infty} x d F(x) = \mathbb{E}\left[X\right].$ Additionally, we provide a more detailed explanation, including two potential decomposition methods on quantile distributional loss in Appendix M of the revised paper. We recommend the reviewer to refer to this discussion, where we extend our decomposition method to quantile-based distributional RL. The revised paper thus has enhanced contribution and mitigated the limitation present in the original submission.

Thank you also for bringing the normal representation to our attention. We fully agree with you that using normal representation (Section 5.4 in [1]) can be (theoretically) convenient. However, there still remains a potential gap between the well-behaved normal representation and more practical fixed-size representation (Section 5.5 in [1]), e.g., quantile, categorical distribution, and particles.

We have fixed the typo mentioned in the minor comment in the revised paper.

Overall, thank you again for pointing out these potential areas for improvement. We appreciate your suggestions. Please let us know if you have any further comments or feedback.

Reference

[1] Bellemare, Marc G., Will Dabney, and Mark Rowland. Distributional reinforcement learning. MIT Press, 2023.

We would like to sincerely express our gratitude for your time and effort spent reviewing our paper. We really appreciate the insightful and detailed comments and suggestions that have been provided. We would like to address each concern you raised in your review and are very happy to discuss them with you further.

Weakness: 1. Quality of writing

1. Section 3 and 4.1 only covers the superficial backgrounds to understand how return density decomposition can be conducted. For example, there is no detailed explanation on C51, which helps to understand why $\mu$ can be defined.

Section 3 covers the basic background and definition of distributional RL, while the Neural FZI in Section 4.1 characterizes the key part of distributional RL, particularly from the perspective of distributional loss. Notably, we decompose the distributional loss function directly based on Neural FZI in Section 4.2, when we choose $d_p$ as KL and histogram function to represent $Z_\theta$, which leads to the induced histogram function $\widehat{\mu}$.

We did not introduce too many details of CDRL to save space, as CDRL is the first successful algorithm with broad recognition with the research community. We agree that adding more details of CDRL would be better as these details could cater to more interested readers with more diverse backgrounds beyond distributional RL. To this end, we have added one paragraph at the beginning of Section 4.2 before we employ the return distribution decomposition, and we have added more details, particularly the CDRL update algorithm, in Appendix A of the revised paper. We believe these revisions can provide readers with sufficient background on CDRL and C51 before we dive deeper into the distributional loss decomposition in Section 4.2.

2. The notations and definitions are not well presented to make readers understand the main contribution. For example, the definition of is followed by Equation (3) without enough statements.

The notation $\widehat{\mu}$ following Eq.3 (Eq. 2 in the revised paper) represents the induced histogram density function after applying the return density decomposition for a given histogram density function $\widehat{p}$. To enhance clarity, we have rewritten the notations and definitions in Proposition 2 and Eq.3 in the revised paper, which we believe significantly improves the readability.

We apologize if any other definitions or notations were not well presented. We have carefully proofread our presentation and updated all of them in the revised paper. We sincerely appreciate any further feedback the reviewer can provide regarding specific areas that may still need enhancement.

3. Experiment section is presented without any descripton of baseliens and environemnts.

We would like to clarify that we employ both Atari games and MuJoCo environments, which are widely used in RL research. In the submitted version, we have stated the environment in both the text and captions of figures in the experimental section.

We have also explained the baselines in the text of each subsection of the experimental section. For example, we have stated how to derive $\mathcal{H}(\mu, q_\theta)$ algorithm in the first paragraph of Section 6.1, and what are the definitions of AC, AC+UE, AC+VE, and AC+UE+VE referring to in the caption of Figure 5.

To further emphasize and differentiate our algorithms, we provided another subsection in Appendix L.1 of the updated paper, where more detailed implementation and comparison of all algorithms considered are provided. Although we appreciate this potential clarity concern about baseline algorithms, we believe the updated version has provided a sufficient description of all baselines. Please feel free to let us know if you are still concerned about any considered algorithm in our experiments.

Thank you for the detailed response and clarification.

As far as I understand, the induced histogram density function $\mu$ is for the target distribution, and $q$ is the distribution for learning. Then, it makes sense that the difference of two distributions represents the uncertainty. If it is true, I have more following questions and concerns.

1. The authors can apply this concept to other distributiononal method, such as QRDQN. Is there any specific reason why the categorical distribution enjoys the proposed concept?
2. Bootstrapped DQN can be a important baseline, because it predicts the return distribtuion by bootstrapping multiple heads and shows the deep exploration concept by the uncertainty concept. Can the authors provide some brief explanation on the difference with the proposed method.
3. Using the uncertainty of value distribution function for exploration, [A] has already shown that the distribution aware exploration method, PQR outperforms the baseline exploration methods and other distributional RL methods. What is the strength of the proposed method compared to PQR?
4. It seems weired that AC (without vanilla entropy) works in mujoco enviornment in Figure 5. It means that the agent performs well without any exploration scheme. Is it correct?
5. To consider the manuscript as the deep interpretation of categorial distribution RL, it is natural to adjust the softmax property with temperature to demonstrate why the categorial distribution is benefical than the naive method. However, this paper introduces the additional uncertainty term into the value function for exploration like SAC, which cannot prove the strength of categorical distribution itself.

[A] Cho, Taehyun, et al. "Pitfall of optimism: distributional reinforcement learning by randomizing risk criterion." Advances in Neural Information Processing Systems 36 (2024).

2. Bootstrapped DQN can be a important baseline, because it predicts the return distribtuion by bootstrapping multiple heads and shows the deep exploration concept by the uncertainty concept. Can the authors provide some brief explanation on the difference with the proposed method.

We acknowledge that Bootstrapped DQN [4] maintains several independent Q-estimators and randomly samples one of them, enabling the agent to perform temporally extended exploration. The exploration mechanism of Bootstrapped DQN is similar to that of DLTV [6], as both aim to enhance exploration by leveraging the stochasticity of the return (as a random variable). Specifically, Bootstrapped DQN uses bootstrapping on the value function in classical RL, while DLTV utilizes the variance from the intrinsically learned return distribution in distributional RL algorithms.

In contrast, our approach decomposes the distributional loss into two parts, with the regularization term being interpreted as a form of exploration that is intrinsically promoted when we optimize the distributional loss in distributional RL. This interpretation aligns well the one important exploration strategy in the exploration literature [5], called uncertain-based exploration. The key difference, however, is that while both Bootstrapped DQN and DLTV explicitly use the ``post-''learned uncertainty to explore--a relatively common strategy--we argue that distributional RL intrinsically promotes uncertainty-aware exploration as a natural consequence of optimizing the distributional loss. This is a novel and interesting connection that bridges the optimization of distributional loss and exploration. Moreover, our explanation offers an intuitive explanation grounded in the exploration literature.

3. Using the uncertainty of value distribution function for exploration, [A] has already shown that the distribution aware exploration method, PQR outperforms the baseline exploration methods and other distributional RL methods. What is the strength of the proposed method compared to PQR?

We would like to clarify that, similar to DLTV, PQR is also an improved exploration strategy by using the estimated uncertainty from distributional RL effectively. As previously explained in our response to Question 2, the purpose of our paper is not to develop advanced exploration based on the learned return distribution from distributional RL. Instead, our decomposed regularization is used to demonstrate that the benefits of using distributional loss in RL are inherently promoting an uncertain-aware exploration. Albeit also involving the exploration concept, the primary objective of our paper is fundamentally different from the subsequent exploration strategies based on distributional RL, such as PQR and DLTV.

4. It seems weired that AC (without vanilla entropy) works in mujoco enviornment in Figure 5. It means that the agent performs well without any exploration scheme. Is it correct?

For a fair comparison, AC is directly adapted from the implementation of SAC but excludes the entropy term used in SAC. With the remaining architecture, such as the use of target networks, unchanged, this ablation explains why AC still generally performs well. In Figure 5, SAC (AC+VE) outperforms AC in most environments, demonstrating the widely recognized performance improvement (though not in all cases) of incorporating the vanilla entropy term in SAC.

Our implementation is based on RLKit (https://github.com/rail-berkeley/rlkit), a widely recognized reinforcement learning framework in the RL community.

5. To consider the manuscript as the deep interpretation of categorial distribution RL, it is natural to adjust the softmax property with temperature to demonstrate why the categorial distribution is benefical than the naive method. However, this paper introduces the additional uncertainty term into the value function for exploration like SAC, which cannot prove the strength of categorical distribution itself.

We acknowledge that well-known property of the softmax function, which interpolates between a point mass on the value with the largest output probability and a concrete uniform distribution, but it is challenging to establish a meaningful connection between this interpolation with the classical RL loss (the naive method in your statement), particularly regarding the largest probability. This contrasts with our proposed return density decomposition, which focuses on the probability associated with the bin containing the expectation of the return.

Instead, incorporating the decomposed uncertainty term into the policy gradient framework, as presented in Section 5, is a natural extension immediately following the formed analysis in value-based RL in Section 4. If we do not introduce the additional uncertainty term into the value function, we, alternatively, can also directly analyze the distributional critic loss and further apply the return density decomposition again to derive the uncertainty term. However, the subsequent analysis and conclusion after applying this redundant analysis remains the same as the current version, while the current version with a direct comparison to SAC is more readable.

In summary, thank you again for your review and feedback. We hope our response will help eliminate your concerns. Should this rebuttal address your concerns, we would be grateful for a revised score. Of course, we remain at your disposal for any further clarifications.

Reference

[1] Bellemare, Marc G., Will Dabney, and Rémi Munos. "A distributional perspective on reinforcement learning." International conference on machine learning. PMLR, 2017.

[2] Lyle, Clare, Marc G. Bellemare, and Pablo Samuel Castro. "A comparative analysis of expected and distributional reinforcement learning." Proceedings of the AAAI Conference on Artificial Intelligence. Vol. 33. No. 01. 2019.

[3] Rowland, Mark, et al. "The statistical benefits of quantile temporal-difference learning for value estimation." International Conference on Machine Learning. PMLR, 2023.

[4] Osband, Ian, et al. "Deep exploration via bootstrapped DQN." Advances in neural information processing systems 29 (2016).

[5] Hao, Jianye, et al. "Exploration in deep reinforcement learning: From single-agent to multiagent domain." IEEE Transactions on Neural Networks and Learning Systems (2023).

[6] Mavrin, Borislav, et al. "Distributional reinforcement learning for efficient exploration." International conference on machine learning. PMLR, 2019.