Interpreting Categorical Distributional Reinforcement Learning: An Implicit Risk-Sensitive Regularization Effect

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.

how the epsilon in Eq. 3 is pre-specified? Should it not be dependent on the actual return distribution which cannot be known in advance? Am I correct in the understanding that you are actually not learning the full return distribution, but still its expectation albeit allowing some uncertainty around it?

A: Eq.3 is to show that the return density can be decomposed via a pre-specified $\epsilon$, which controls the proportion of the expectation and the regularization part. Eq.3 does not imply that we are going to learn only based on the expectation, but to focus on the validity of the density decomposition. For example, in the sensitivity analysis in Figure 2, we are using the regularization term with different $\varepsilon$ to optimize the neural network in order to verify the impact of it.

if so, it seems to me you're matching the expaction whilst expanding it according to a pre-defined level of uncertainty retained from the target distribution. In this case, how would you ensure that the so expanded distribution still has the same expectation/the correct expectation?

A: Yes. The expectation can be inherently maintained by simply taking the expectation over both sides in Eq. 3. Since the first part maintains the same/correct expectation with the weight $1-\epsilon$, the remaining term accordingly has the same expectation with the weight $\epsilon$.

if you are minimizing KL divergence, how would you ensure the return distribution model is and remains as a histogram which is count-based?

A: The validity of the induced distribution is guaranteed in Proposition 1 for a certain range of $\epsilon$. The counted-based histogram is inherently maintained via the decomposition in Eq. 3.

in Eq. 5, you're regularzing the cross-entropy between consecutive 'return distributions'. How is it connected to maximum-entropy RL which encourages 'policy' entropy or cross-entropy between the policy and a reference policy? There seems no bridge to me to be built between uncertainty in action selection and that in return estimate, the only connection perhaps being both can be approached by augmenting the reward function with the entropy term. In fact, I failed to understand the claim that your regularization is favouring target return distributions with large dispersion. It seems to me only to push the dispersion of the return distribution model to catch up with the dispersion of the target return distribution.

A: Let us look at Eq. 10 to compare with MaxEnt RL in Eq. 7. Since we know $a_i = \pi(\cdot | s_i)$, optimizing the regularization term implicitly updates the policy $\pi$. By contrast, MaxEnt RL explicitly update the policy via maximizing the entropy in Eq. 7

In terms of Eq. 5, the regularization term in the value-based RL includes $q_\theta$, and optimizing Eq. 5 will directly affect the update of $q_\theta$. Note that $a_i =\arg\max_a \mathbb{E} [q_\theta^{s_i, a}]$, which therefore connects the return estimate with the action selection.

We would like to clarify that we did not claim the regularization would certainly favor target return distribution with a large dispersion. We state that it tends to do that as the original $q_\theta$ only captures the expectation and it tends to expand the distribution with a larger variance to match the distribution. However, we agree your statement is safer and more accurate, which is that regularization in general helps to catch up with the dispersion of the target distribution. We have revised them appropriately in the updated paper.

(minor) how and why a policy can be learned to reduce the intrinsic uncertainty of the environment, shouldn't it by definition be intrinsic and thus independent from the policy?

A: The intrinsic uncertainty has been discussed in [1, 2], which is defined as the stochasticity of the environment, including the transition dynamics and stochastic reward function. Since distributional RL makes full use of the distribution information of return, it is deemed to reduce the intrinsic uncertainty of the environment for better performance.

We thank the reviewer once again for the time and effort in reviewing our work! We are more than happy to answer any other questions you have.

[1] Implicit Quantile Networks for Distributional Reinforcement Learning (ICML 2018)

[2] Distributional Reinforcement Learning for Efficient Exploration (ICML 2019)

Thank you for taking the time to review our paper. We appreciate your comments and we have revised the paper accordingly. Here we would like to address the concerns you raised in your review.

Eq. (3) and the subsequent remarks confused me regarding what $\widehat{\mu}$ is? Does mean the coefficient of the corresponding bin?

A: We apologize for this confusion and we thus have provided the detailed definition of $\widehat{\mu}$ in the revised version. $p^\mu_{i}$ is indeed the coefficient of the corresponding bin with the bin height $p^\mu_{i}/\Delta$.

Moreover, I failed to understand eq. (6) and what Proposition 3 tries to convey, specifically regarding the integral on the right.

A: Generally speaking, Eq. 6 and Proposition 3 try to convey that minimizing the term (a) in the decomposed Neural FZI is equivalent to Neural FQI in the limit case, allowing us to use the regularization term to interpret the benefits of distributional RL over classical RL analyzed later.

The integral part additionally shows the uniform convergence rate of distribution difference in terms of $\Delta$, which turns out to be an order of $\mathcal{O}(\Delta)$.

Additionally, Theorem 2 asserts the convergence of the learned policy, yet it does not specify under what notion it converges, that is, it lacks clarity on the measure under which it converges.

A: The optimal policy is defined based on the optimal Q function. In our proof, following SAC paper, we show the algorithm converges to the optimal Q function~(the maximum), and thus the corresponding policy/policies would be optimal. Please refer to Appendix H.3 for more details.

I found most theoretical results rely on the assumption that the bounded entropy....

A: The bounded assumption is natural and mild as the convergence of the Bellman operator typically requires a bounded reward assumption. Note that our cross-entropy term serves as the augmented rewards, for which it is natural to assume to be bounded to allow the convergence accordingly.

I think this paper may need some necessary polish to improve clarity before I can make an informed evaluation.

A: Thanks for this suggestion and we have substantially improved the clarity based on all the reviewers' suggestions in the rebuttal phase. We appreciate an informed evaluation of our paper.

The formulation of neural FZI is similar to [2]. The latter focuses on maximizing the log-likelihood, which is equivalent to minimizing the KL divergence or cross-entropy. The authors also proposed to specify the divergence in eq.(2) to KL divergence, so the formulation looks almost the same. Hence, I am wondering about the potential connection here.

A: [2] can be viewed as a distributional version of Fitted Q evaluation in the offline setting, whose target is to provide the prediction error guarantee in the off-policy evaluation. Based on the relationship of MLE they used and KL divergence, [2] can approximately serve as the prediction error guarantee for offline CDRL. However, our target is to interpret the benefits of distribution RL starting from CDRL, which is very different from [2].

We would also like to clarify that [2] used TV and Wasserstein distance instead of KL in their theoretical analysis to derive the prediction error guarantee in an offline setting, which we focus on interpreting the KL-based practical CDRL algorithms in online cases.

In summary, the main contribution and research target between [2] and us are different and based on different divergences, although the algorithm part shares some similarities with CDRL algorithm.

We thank the reviewer once again for the time and effort in reviewing our work! We would greatly appreciate it if the reviewer could check our responses and let us know whether they address the raised concerns. We are more than happy to provide further clarification if you have any additional concerns. Should this rebuttal address your concerns, we would be grateful for a revised score.

Q6: I don't understand the claim of equation (6). It says that with probability.....

A: $Z^{k+1}$ is a random variable and the convergence in probability 1 indicates that $Z^{k+1}$ would equal to a constant $\mathcal{T} Q_{\theta^*}$. The $\mathcal{T} Q_{\theta^*}$ is a constant in each Neural FZI as we use the target network trick, where the target is fixed in each iteration phase of Neural FZI.

Q7: The proof of proposition 3 is sloppy. Firstly, equation (29) is incorrect.....

A: Thanks for pointing out some inaccurate parts in our proof and we have revised this proof accordingly, also including the limit of the minimizer issue in the Appendix.

It is correct that proposition 3 is about optimizing the term (a) and thus the target distribution does not include the $\widehat{\mu}$ term. Regarding the relationship between [3], An Analysis of Categorical Distributional Reinforcement Learning" by Rowland et. al, note that [3] and other related works showed that most distributional RL algorithms can enjoy the mean preserving properties, but our focus is to interpret the extra benefits of distributional RL over only this expectation part. To this end, it is necessary to connect the term (a) with the optimization in terms of the expectation part, thus allowing us to attribute the benefits to the regularization. As far as we are concerned, proposition 3 is a crucial step towards interpreting the regularization effect. The results in proposition 3 actually coincide with the mean-preservation properties in distributional RL, which can help to substantiate the technical soundness of our decomposition.

We thank the reviewer once again for the time and effort in reviewing our work! We would greatly appreciate it if the reviewer could check our responses and let us know whether they address the raised concerns. We are more than happy to provide further clarification if you have any additional concerns. Should this rebuttal address your concerns, we would be grateful for a revised score.

[1] Bellmare et al. Distributional Reinforcement Learning (MIT press)

[2] Dabney et al. Implicit quantile network for distributional reinforcement learning (ICML 2018)

[3] Rowland et. al. An Analysis of Categorical Distributional Reinforcement Learning" (AISTATS 2018)

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

Q1: In the "Bellman Operators vs Distributional Bellman Operators" section, it says tha random variable definition ....

A: This statement can refer to Chapter 4 of the book [1]. In particular, the random-variable operator may not be rigorous enough as it specifies a return distribution without identifying a mapping from the sample space to the real number, i.e., definite the probability measure. However, the theory of contraction mappings needs a clear definition of the space on which an operator is defined, i.e., the probability space of the return distribution defined in advance.

We clearly state this difference in case any theory reviewers favor the more rigorous definition. However, we agree with you the difference is minor.

Q2: I doubt Theorem 1 is novel, it is essentially a concentration inequality on histogram estimates....

A: The big O definition is a commonly used convergence process in probability theory, and here we provide the wiki reference: https://en.wikipedia.org/wiki/Big_O_in_probability_notation. In particular, it characterizes a stochastic boundedness in the convergence in probability.

The main similarity is both Theorem 1 and DKW are derived by leveraging the commonly used concentration inequality, and thus some parts share a similar form. However, the main difference is that DKW focuses on the (sample-based) empirical distribution, while Theorem 1 provides results about the histogram-based estimator in the context of CDRL.

Q3: The objection $\widehat{\mu}^{s, a}$ is not defined clearly ....

A: Thank you for pointing out this clarity issue and we have fixed it according. It is true that $\widehat{\mu}$ is the induced histogram estimator, which is used to characterize the regularization effect / the benefits of distributional RL regardless of the expectation of return distributions.

Q4: Right after Proposition 1, what does "this decomposition exactly keeps the vanilla (categorical) distributional RL...

A: We have paraphrased this sentence as ''Under this valid return density decomposition condition, this return density decomposition approach precisely maintains the standard categorical distribution framework in distributional RL''. This statement is made to clarify the potential concerns that this decomposition may change the objective function of CDRL. However, as stated in the new sentence, the original CDRL framework is carefully maintained.

Q5: I don't agree that the second term of equation (5) is doing "risk-sensitive entropy regularization"....

A: We would like to clarify this potential misunderstanding. According to the literature [2], risk refers to uncertainty over possible outcomes, and risk-sensitive policies are those that depend upon more than the mean of the outcomes. This implies that our risk-sensitive policies, which are indeed risk-neutral, typically refer to the capability to reduce the intrinsic uncertainty of the environment. We know distributional RL provides a framework for risk-sensitive RL, including the risk-seeking or risk-averse policies, and therefore, it can be achieved by changing the convexity or concavity of the utility function [2] on the return distribution in the interaction phase of the environment, i.e., the behavior policy side. For example, IQN [2] changes different weights to average the quantiles to select the actions interacting with the environment. This risk-seeking or risk-averse framework can be very naturally incorporated into our algorithm or Neural FZI, by changing the greedy rule $\pi_Z(s^\prime_i)= \arg\max_{a^\prime} \mathbb{E}\left[Z_{\theta^*}^{k}(s_i^\prime, a^\prime)\right]$ in Eq. 2 with some specified utility function, $\pi_Z(s^\prime_i)= \arg\max_{a^\prime} U(Z_{\theta^*}^{k}(s_i^\prime, a^\prime))$.

Regarding the regularization mechanism in Eq. 5, note that $a_i$ is selected by the policy, which is also determined by $q_\theta$. Thus, minimizing the term (a) and the regularization term will optimize the policy via $q_\theta$.

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

General Question: The paper is highly obscure and filled with ambiguous statements and typos that prevented me from validating most of its claims. It also appear to support its theory using arguments that lack rigour or purely contradictory. I finally disagree with the claim on page 4 that distributional RL produces risk sensitive policies:

A: We apologize for the confusion generated by the clarity issue. Thus, we have improved our claim following your suggestions throughout the paper as appropriate, and sincerely hope the current version is easier to follow. We will also be happy to edit the paper further based on any further helpful comments from the reviewers. The specific update corresponds to each clarity issue you raised as follows.

Examples of lack of rigour: the use of the concept of "ideal case" mentioned at the bottom of page 2 to support the idea that Neural FQI converges to optimal Q-function.

A: Thanks for pointing out this issue. We have provided more specific explanations, respectively, instead of using ``the ideal case''. This is a commonly used assumption about the richness of function class in statistical learning theory. In Neural FQI, we explain as $\\{Q_\theta: \theta \in \Theta\\}$ is sufficiently large such that it contains $\mathcal{T}^{opt} Q_{\theta^*}^{k}$, while in Neural FZI and proposition 3, we explain as $\\{Z_\theta: \theta \in \Theta\\}$ is sufficiently large such that it contains $\{Y_i\}_{i=1}^n$. We also revised it accordingly to the updated version.

Examples of lack of rigour: a mismatch of assumption between the statement of proposition 3 and its proof ...

A: We would like to clarify that $\mathbb{E}\left[Y_i\right]$ can be representable is a corollary if $\\{Y_i\\}$ is representable as the $\mathbb{E}\left[Y_i\right]$ is (the limit of ) the first component of $\{Y_i\}$'s pdf via our decomposition.

Example of contradiction: Proposition 3 is contradicting itself in equation (6). The left equation claims that the minimizer over converges to the k+1 iterate of the Q function in FQI while on the right claiming that it converges to the actual Q-value function.

A: We apologize for this typo, where the action Q-value function should be replaced with the $k+1$ iterate of the Q function, which is thus consistent with the left equation. We have fixed the typo in the revised paper, where, specially, the correct form of the right equation is updated as:

$\int_{-\infty}^{+\infty}\left|F_{q_\theta}(x)-F_{\delta_{ \mathcal{T}^{\text{opt}} Q^k_{\theta^*}(s, a)}}(x)\right| d x =o(\Delta)$

Q: Typos and inclarity include: ...

$A:$ Thank you for noting these typos and we have fixed them appropriately. For example, we acknowledge that p.4 $\Delta_E^i$ represents the interval of $\mathbb{E}[Y_i]$ instead of the current return distribution and added a notation warning about the low case random variables. We also gave a more accurate definition of $\widehat{\mu}$.

Please refer to the updated paper for more details. Please let us know if there are any other typos you had in mind for correction.

We thank the reviewer once again for the time and effort in reviewing our work! As suggested, in the rebuttal period, we have finished the throughout notation revisions to allow an informed evaluation. We would greatly appreciate it if the reviewer could check our responses and are more than happy to provide further clarification if you have any additional technical concerns about our paper.