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

Authors: Ke Sun, Yingnan Zhao, Enze Shi, Yafei Wang, Xiaodong Yan, Bei Jiang, Linglong Kong
Primary area: reinforcement learning
Keywords: distributional learning, reinforcement learning, exploration

Abstract:
Despite the remarkable empirical performance of distributional reinforcement learning (RL), its theoretical advantages over classical RL are not fully understood. Starting with Categorical Distributional RL (CDRL), we propose that the potential superiority of distributional RL can be attributed to a derived distribution-matching regularization by applying a return density function decomposition technique. This less-studied regularization in the distributional RL context aims to capture additional knowledge of return distribution beyond only its expectation, contributing to an augmented reward signal in policy optimization. In contrast to the standard entropy regularization in MaxEnt RL, which explicitly encourages exploration by promoting diverse actions, the regularization derived from CDRL implicitly updates policies to align the learned policy with environmental uncertainty. Finally, extensive experiments substantiate the significance of this uncertainty-aware regularization derived from distributional RL on the empirical benefits over classical RL. Our study offers a new perspective from the exploration to explain the benefits of adopting distributional learning in RL.

Decision: reject
Reviewer scores: 3, 6, 8, 3

Reviews and discussion:
<thread>
creation date:
1730339981944

thread_type:
official_review

<thread_object>
creator: reviewer_hctq

<content>
summary:
This paper examines the success of categorical distributional RL algorithms by formulating the distributional component of these algorithms as a form of regularization. The authors support this perspective through theoretical analysis. It is also followed by experiments to validate their findings.

strengths:
- This paper provides theoretical explanations on distributional RL (concretely, on categorial distributional RL algorithms). It is an important problem.

- This paper explores both the theoretical aspects (considering the regularization effect) and empirical validation, making it relatively comprehensive overall.

weaknesses:
- 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)
- 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."
- 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.
- I have more concerns regarding novelty that I listed in the questions section.

questions:
The theoretical contributions appear to have limited novelty in light of prior work. For instance, proposition 3 looks pretty close to prior results on the convergence of categorial algorithm (such as in [1]) ? We already know such convergence results when $\Delta$ goes to zero. Could you clarify the novelty here?

[1] Bellemare, Marc G., Will Dabney, and Mark Rowland. *Distributional reinforcement learning*. MIT Press, 2023.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730577139594

thread_type:
official_review

<thread_object>
creator: reviewer_u83m

<content>
summary:
This paper provides an explanation of how the Categorical Distributional RL (CDRL) outperforms classical RL. Essentially, CDRL contains an implicit entropy regularization that encourages exploration through visiting states where the environment uncertainty is largely underestimated. This uncertainty-aware regularization provides context about the superiority of distributional RL.

strengths:
Authors neatly derived the uncertainty-aware regularization term in distributional RL that is responsible for their superiority. Furthermore, an explicit demonstration of the link between distributional and classic RL through regulating the effects of the uncertainty-aware regularization term is provided.

weaknesses:
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. 

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

questions:
1) How can the $p_{E}$ (coefficient on bin $\Delta_{E}$) be computed? 
2) What are the effects of various bin sizes on the results?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730666456921

thread_type:
official_review

<thread_object>
creator: reviewer_bvvf

<content>
summary:
This paper presents a novel decomposition of CDRL’s objective function into an expectation-based term and a distribution-matching regularization. The authors argue that the latter term provide an augmented reward signal, contributing to implicit regularized exploration. Unlike the explicit entropy regularization in MaxEnt RL, this approach encourages exploration aligned with environmental uncertainty. The paper further introduces the Distribution-Entropy-Regularized Actor Critic (DERAC) algorithm, which balances between expectation-based and distributional learning, supported by experiments in Atari and MuJoCo environments.

strengths:
This paper provides a new perspective on interpreting distributional RL by leveraging the decomposition of the action-value distributions. The derivation of uncertainty-aware regularization from the decomposition of CDRL’s return distribution is an innovative contribution that highlights the advantages of distributional RL over traditional methods.

The comparison between the proposed uncertainty-aware regularization and entropy regularization in MaxEnt RL offers valuable insights into the added benefits of distributional RL.

The experiments are solid and provide strong empirical evidence supporting the paper’s theoretical claims.

weaknesses:
The analysis in this paper focuses solely on categorical distributional reinforcement learning (C51). Given that other distributional learning methods, such as QR-DQN and IQN, are becoming more popular, this appears to be a weak point.  It would be helpful for the authors to discuss how their analysis might extend to or differ from other popular distributional RL methods like QR-DQN and IQN.

questions:
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 \ge 1-P_E$.”  Does this imply that $\epsilon$ needs to be close to $1$ when $P_E$ is very small?

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.

Minor comment: On page 7, the first sentence after Lemma 1 likely should use "retain" instead of "remain”.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730739652288

thread_type:
official_review

<thread_object>
creator: reviewer_x6vg

<content>
summary:
This paper studies one of the advantages that distributional reinforcement learning, especially the categorical version, has in executing uncertainty-aware regularized exploration.
The authors claim that categorical distributional reinforcement learning (CDRL) implicitly updates the return distribution and uncertainty simultaneously when we learn the return distribution by value (return distribution) iteration.
The induced uncertainty term can be considered as an entropy regularization which is analogous to maximum entropy RL.
By the proposed analysis, CDRL has a strength in terms of exploration compared to standard RL algorithms.

strengths:
In the research of distributional RL, there are a few analyses of what strengths distributional RL has compared to standard RL.

This paper provides an interesting view that uncertainty of return distribution is additionally learned and provided as the entropy for encouraging exploration.

In addition, this paper provides the experimental results of environments with both continuous and discrete action spaces.

weaknesses:
Despite the above strengths, this paper has the following critical weaknesses.

If I misunderstood and was worng about the following comments, I’ll be happy to discuss with the authors.

1. **Quality of writing**: the main weakness of this paper is the quality of writing, which makes readers hard to understand the main contribution.
    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  $\hat{\mu}$ can be defined.
    2. The notations and definitions are not well presented to make readers understand the main contribution. For example, the definition of $\hat{\mu}$ is followed by Equation (3) without enough statements.
    3. Experiment section is presented without any descripton of baseliens and environemnts.
2. **Not rigorous proof**: the theortical proofs to show the property of categorical distributional RL ar not rigorous to validate the main concept.
    1. There is no explicit proof how the decomposed term in Equation (3) is related to categorical representation. Especially, there is no explanation on the relation between $\mu$ and categorical distribution.
    2. Proposition 2 has too strong approximation that KL divergence is proportional to the cross entropy. This can underestimate the uncertainty of ground-truth distribution. 
3. **Weak presentation of experiment section**: Since the author claims that the proposed method is related to exploration, but there is no discussion on exploration schedule and comparison with other exploration schem, such as DLTV. In addition, the proposed method cannot ourperform C51, which only uses the default exploration method, epsilon-greedy.

questions:
What is the main role and exact definition of $\mu$? I can partially agree that the categorical representation can be considered as a histogram density representation, but I cannot understand how $\mu$ can be introduced.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732120533615

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Author Response

comment:
We thank all reviewers for their in-depth and constructive comments. We have revised and updated our manuscript by following all reviewers' suggestions. We provide the major revision points in the revised paper:

* **Discussion about decomposing quantile-based distributional loss in Appendix M**. Our discussion examines the possibility of extending our decomposition-style analysis to quantile-based distributional loss. This new discussion broadens our research scope and can further enhance our contributions.
* **Enhanced clarity.** We improved the clarity that may not have been accurate or clear in the original submission, including the explanation of $\widehat{\mu}$, comparison with prior work, the equivalence between our work and C51, some steps in the proof, and more details of environments and baselines.
* **More experiments.** We also added more experiments to corroborate our conclusions further, including the ablation study by varying the bin sizes and the equivalence between KL and cross-entropy in our algorithms. 

We hope the revised version addresses all the reviewers' concerns. We would appreciate it if the reviewers could check our response to their comments and look through our revised paper further. We are consistently committed to enhancing the quality of our work and remain at your disposal for any further clarifications or questions you might have.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732580982664

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers, 

The rebuttal phase is coming to a close. Thank you again for the effort you put into reviewing our paper. We are eager to learn if our rebuttal adequately addresses your concerns. We are happy to provide further clarification and engage in any discussion. Thank you! 

Sincerely,
The Authors
</content>
</thread_object>

</thread>
<thread>
creation date:
1733090611113

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
Dear Reviewers,

The author-reviewer discussion is nearing its final day. We believe we have made our best efforts to address all your concerns, and we remain eager to hear your feedback, even in the final moments of this discussion period. For reviewers to whom we have responded but have not yet had the chance to reply, we sincerely hope our efforts have addressed your concerns and kindly request a revised score if appropriate. Thank you once again for your thorough and thoughtful review of our paper.

Sincerely,  
Authors
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper explores the role of uncertainty-aware exploration in Categorical Distributional RL (CDRL) algorithms, presenting the Distribution-Entropy-Regularized Actor-Critic (DERAC) method. The authors argue that the distributional component serves as a form of regularization, which promotes exploration. The theoretical framework is supported by experimental validation in Atari and MuJoCo environments.

While the paper provides a novel perspective on the benefits of distributional RL, there are still concerns about its clarity, novelty, and scope. Reviewers appreciated the theoretical contributions and the solid experimental validation, but many noted that the presentation could be improved. Some sections, particularly in the theoretical analysis, remain unclear, and there are concerns about vague terminology. Additionally, the focus on DERPI and DERAC, which are not traditional categorical RL algorithms, feels narrow, and there is a need for a broader discussion of more established methods like C51, QR-DQN, and IQN. The novelty of the theoretical contributions, particularly Proposition 3, was also questioned, as it appears similar to existing convergence results in the literature.

The paper would benefit from a clearer explanation of the novelty of the proposed methods, a broader discussion of traditional distributional RL methods, and further clarification of technical terms. Additionally, an ablation study on the impact of bin size would help demonstrate the robustness of the approach.

Overall, the paper has potential but requires revisions to improve its clarity, address novelty concerns, and expand its scope to fully strengthen its contribution to the field.