Intrinsic Benefits of Categorical Distributional Loss: Uncertainty-aware Regularized Exploration in Reinforcement Learning
Authors: Ke Sun
Yingnan Zhao
Enze Shi
Yafei Wang
Xiaodong Yan
Bei Jiang
Linglong Kong
Venue: NeurIPS 2025 poster ()
Keywords: distributional learning
reinforcement learning
exploration
uncertainty
Primary area: reinforcement_learning

Abstract: The remarkable empirical performance of distributional reinforcement learning~(RL) has garnered increasing attention to understanding its theoretical advantages over classical RL. By decomposing the categorical distributional loss commonly employed in distributional RL, we find that the potential superiority of distributional RL can be attributed to a derived distribution-matching entropy regularization. This less-studied entropy regularization 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 vanilla entropy regularization in MaxEnt RL, which explicitly encourages exploration by promoting diverse actions, the novel entropy regularization derived from categorical distributional loss implicitly updates policies to align the learned policy with (estimated) environmental uncertainty. Finally, extensive experiments verify the significance of this uncertainty-aware regularization from distributional RL on the empirical benefits over classical RL. Our study offers an innovative exploration perspective to explain the intrinsic benefits of distributional learning in RL.

Intrinsic Benefits of Categorical Distributional Loss: Uncertainty-aware Regularized Exploration in Reinforcement Learning
Ke Sun', Yingnan Zhao?, Enze Shi', Yafei Wang', Xiaodong Yan", Bei Jiang", Linglong Kong"*
UUniversity Alberta, Canada
Harbin Engineering University, China
"Xi'an an Jiaotong University, China
(ksun6, eshi, yafei2, beil, .xoong zhaoyingnan@hrreu. edw.rrbbbuuuuuuue.e.........eee. cn yanxlaodongoxjja tu. edu.
Abstract
The remarkable empirical performance of distributional reinforcement learn- ing (RL) has garnered increasing attention to understanding its theoretical ad- vantages classical RL. By decomposing the categorical distributional loss commonly employed in distributional RI,, we find that the potential superiority of distributional RL can be attributed to derived regularization. This less-studied entropy regularization aims to capture additional knowledge of return distribution beyond only its expectation, contributing an augmented reward signal in policy optimization. In contrast to the vanilla en- tropy regularization in MaxEnt RL, which explicitly encourages exploration by promoting diverse actions, the novel entropy regularization derived from categori- cal distributional loss implicitly updates policies align the learned policy with (estimated) environmental uncertainty. Finally, extensive experiments verify the significance this uncertainty-awant regularization from distributional RL the empirical benefits over classical RL. Our study offers an innovative exploration perspective explain the intrinsic benefits of distributional learning in RI.
cerrributionmmatchiin entropy
Introduction
The fundamental characteristics classical reinforcement learning (RL) [57, such as Q-learning 62 rely on estimating the expectation of discounted cumulative rewards that an agent observes while in- teracting with the environment. In contrast to the expectation-basee RL, novel branch of algorithms. termed distributional RL, seeks to estimate the entire distribution of total returns and has achieved state-of-the of-the--rtttttt performance across a diverse array of environments 5
10
63 56 52
Meanwhile, discussions of distributional RL have increasingly extended into a broader range of fields, such as risk-sensitive control [1 8) offline learning [35 67 policy exploration [36
49 25 ro-
bustness [55 53 11, optimization [51] 28 54 statistical inference [69] multivariate rewards 770 66]
and continuous-time setting 65
Motivation:
Understanding the Benefits of Employing (Categorical) Distributional Loss in RL.
Despite the impressive empirical success of distributional RL algorithmss our comprehension their advantages over classical RL remains incomplete, especially for the general function approximation setting and practical implementations. Early work [33] demonstrated that in many realizations of tabular and linear approximation settings, distributional RL behaves similarly to classic RL, suggesting
Corresponding author
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

that its benefits are mainly realized in the non-linear approximation setting. Although their findings offer profound insights, their analysis, based on a coupled update method, overlooks several factors, such as the optimization effect under various losses. The statistical benefits of quantile temporal difference (QTD), employed in quantile distributional RI., e.g., QR-DQN [T], were highlighted in [51, 50) which posited that the robust estimation of QTD fosters the benefits in stochastic environments. The foundational theoretical aspects of Categorical Distributional RL (CDRL), e.g., C51 [5], were discussed in 448 27 however, explaining the advantages of categorical distributional learning remains under-explored. Recent studies [6]] elucidate the benefits of distributional RL by introducing the small-loss and second-order PAC bounds, revealing the enhanced sample efficiency, particularly in specific cases with small achievable costs. Yet, their findings are not directly based on practical distributional RL algorithms, such C51 or QR-DQN. Therefore, it is imperative to close this gap between understanding their theoretical advantages and practical deployment in complex environments for distributional RL algorithms.
Contributions. In this study, we interpret the potential superiority of distributional learning in RL. over classical RL, specifically focusing Categorical Distributional RL (CDRL), the pioncering family within distributional RL. We examine the benefits through the lens of a regularized exploration effect, offering distinct perspective relative to existing literature. Our investigation begins by decomposing the categorical distributional loss into a mean-related term and a distribution-matee regularization term, facilitated by our proposed return density decomposition technique. The resulting regularization acts as an augmented reward in the actor critic framework, encouraging policies to explore states whose current return distribution estimates lag far behind the (estimated) environmental uncertainty in the target return. This derived regularization from the categorical distributional loss in CDRL promotes uncertainty-awar exploration effect, which diverges from the exploration for diverse actions commonly used in MaxEnt RL 664 17 18]. We also provide the convergence foundations when leveraging the decomposed uncertain-aware regularization in the actor critic. Empirical evidence underscores the pivotal role of the uncertainty-an empirical success of adopting categorical distributional loss in RL over classical RL on both Atari games and MuJoCo tasks. We further elucidate the distinct roles that the uncertainty-awan in distributional RL and the vanilla entropy in MaxEnt RL play by exploring their mutual impacts on learning performance, providing consitent evidence for []. This opens new avenues for future research in this domain. Our contributions are summarized as follows:
entropy regularization in the
entropy
1. By applying return density decomposition on the categorical distributional loss, we derive distribution-matte mattiinng regularization. This regularization promotes uncertainty-aware exploration, interpreting the benefits of categorical distributional learning in RI.
2. We extend the benefit interpretation of the categorical distributional loss to policy gradient methods. We compare the different exploration effects of decomposed uncertainty- regularization from distributional RL and the vanilla entropy regularization MaxEnt RI.
Empirically, we verify the uncertainty-awwe regularization effect on the performance improvement of distributional RL and investigate the mutual impacts of regularizations in the learning.
Outline. We provide the related work and background knowledge in Sections and respectively. We begin by interpreting the benefits of categorical distributional learning as uncertainty-awar regularized exploration in value-based RL in Section 4 We further probe this exploration benefit in the basicy-based RL,, especially the actor critic framework, in Section []] where we directly compare it with the vanilla entropy regularization in MaxEnt RL. Extensive experiments demonstrate the benefits of regularized exploration in distributional RL and its mutual impact with entropy regularization in MaxEnt RL in Section 6
2 Related Work
Distributional Learning via Categorical Representation. Categorical learning has been widely
employed, with advantages representation [46] 26 and optimization (24 [54] Recently, the empirical superiority of categorical distribution learning has been further investigated in various RL tasks 14. A pressing need exists to examine the theoretical foundations of categorical distributional learning, particularly in RL. The perspective of uncertainty-aware regularized exploration that our study introduces provides significant insights into understanding the benefits of employing categorical distribution loss in the RL context.

---

Uncertainty-orienter Exploration. envertainty-orientee exploration plays an integral part in exist- ing exploration methods (20), which leverages uncertainty either the (posterior) estimation of the
value function, as seen in Bayesian framework
45
4
Bootstrap [44], and Ensemble methods [0)
or in the entire distribution of returns [58]
36
For example, Decaying Left Truncated Vari-
ance (DLTV) [36] and Perturbed Quantile Regression (PQR) 9 exploit the variability of the learned return distribution to promote an optimistic exploration in distributional RL. In contrast, the primary aim of this study is to demonstrate that distributional learning in RL entails an intrinsic exploration effect against environmental uncertainty, contributing to the outperformance of distributional RL over classical RL. Our study goal is independent of designing advanced exploration strategies top of distributional RL. Similarly, MaxEnt RL [64], which includes soft Q-learning [[6], Soft Actor Critic (SAC) [17] and their variants [191, also promotes uncertainty orientte exploration by relying on the stochasticity of the learned policy. A more detailed discussion of related work is provided in
Appendix A
3 Preliminaries
Markov Decision Process (MDP) and Classical RL. An environment is modeled via an Markov Decision Process A, P.
with a set of states S and actions A. the bounded reward function R ::S S x A P(immin Rmox)), the transition kernel P S x A PSS), and discounted factor 7 E [0, IJ. We denote the reward the agent receives at time i as r3 ~ R(s,, az). Given a policy æ, the key quantity of interest is the return 2", which is the total cumulative rewards over the course of trajectory defined by Z Za a) = Ei=o 7'r1l 80 = 8, 20000 a. Classical RL focuses on estimating the
expectation of the return, i.e., Q (8, a) E, Et 70 8,00 = a. We also define Bellman evaluation operator T=Q(s, a) = E[R(s, a) + P.spPar' (s', a')], and Bellman optimality operator Qs, a) [[ a y max, E.' (Q a')].
Distributional RI. and CDRL.. Instead of only learning the expectation in classical RI, distributional RL models the full distribution of the return z". The return distribution 1y" : S x A - P(R) defined as 17t (s,a) a) D(2" a)), where DD extracts the distribution of random variable.
17" (s, a) is updated via the distributional Bellman operator I², defined by Is, _ R(s,
yZ a'), where 2 implies that random variables of both sides are equal in distribution. Categorical Distributional RL (CDRL) is the first successful distributional RL family that approximates the return distribution by a discrete categorical distribution n = Σn p:o.,, where { is set of fixed supports and { are learnable probabilities. The leverage of heuristic projection operator TTc (see Appendix for more details) and the Kullback-Leibler (KL) divergence guarantee the theoretical convergence of CDRL under Cramér distance or Wasserstein distance in the tabular setting [48).
4 Regularization Benefits in Value-based Distribution RL
In this section, we simplify value-based distributional RL to Neural Fitted Z-Iteration (Neural FZI) process in Section 4.1 within which the distributional loss used in distributional RL can be further rewritten as entropy- reguuaiereg form as shown in Section 4.2 Finally, we characterize the role of the derived entropy- based regularization as uncertain- aware regularized exploration in Section 4.3
4.1 Distributional RL: Neural FZI
offers statistical explanation of DQN [39, capturing its key features, including experience replay and the target network Qo-. In Neural FQI, update parameterized Qo in each iteration & of an iterative regression framework:
Classical RL: Neural Fitted Q-Iteration (Neural FQI). Neural FQI 13 47
1 n
2 argming,
Σ n
10t Q0 (s))
(I)
where the target v¹ r(s;, ai) 7 maxaeA Q6- (5%, a) is fixed within every Tiager steps to update target network Qa- by letting Q2- Q+. The experience buffer induces independent samples
3

---

((s;, de, Ti, st) ne(m If (Qo 0 E )) is sufficiently large such that it contains 7 Q/- i.e,, the
realizable assumption in learning theory [40), Neural FQI has the solution Q - 7 Q/ which is exactly the updating rule under Bellman optimality operator [3]
Distributional RI: Neural Fitted Z-Iteration (Neural FZI). Analogous to Neural FQI, we simplify value-based distributional RL algorithms with the parameterized Zo as Neural FZI:
zz
argmin 2o
Σ 1m1
dp(Y, Zo (s1, ai)),
where we denote the target return as Y = R(s;, az) +7 R() Zz-- (8/; (z(()) with the policy тz following the greedy rule (s{) argmax, E Z (5',, a')] The target Y₁* is fixed within every Target steps to update target network Zg.. dp is distribution divergence between two distributions. While our analysis is not intended to involve properties of deep neural networks, we interpret distributional RL as Neural FZI, as it is by far the closest to the practical algorithms. More details about the motivation of Neural FZI are provided in Appendix C
4.2 Distributional RL: Entropy- Neural FQI
As mentioned previously in preliminary knowledge in Section 3 CDRL employs neural networks to learn the probabilities { in a discrete categorical distribution to represent Zg. and choose KL divergence as dp in Eq. 2 of Neural FZI. We next decompose the KL-based distributional loss d, in CDRL by utilizing an equivalent histogram density estimator p in representing Zo.
1.
Return Density Decompasition. To characterize the impact of additional knowledge from the return distribu- tion beyond its expectation, we use a variant of gress error model from robust statistics [23], which was also similarly applied to analyze Label Smoothing [42] and Knowledge Dis-
ps,a
I(x E 4A)
A
1
A
+
tillation 221) Akin to the categorical parameterization in CDRL, we utilize histogram function estimator (zx'(x)) with N bins to approximate an arbitrary continuous density p" p=x of 2 (s, a), given a state s and action In contrast to categorical parameterizatton defined on set of fixed supports, the histogram estimator operates over a continuous interval, enabling more nuanced analysis within continuous functions. Given fixed set of supports < 4 S ... S Ly with the equal bin size A, each bin is thus denoted as A; = [4, i 1, density estimator is formulated by (>)(x) = Σ p; I (ær A.)/A with p; as the coefficient in the i-th bin A, Denote Âg as the interval that E [7 (s,a) falls into, i.e., [2" a)] E A g- Putting all together, we apply an action-state return density decomposition over the histogram density estimator p
Ag
AF
Figure 1: Return Density Decomposition on Histograms.
with AN = [NN-1, As such, the histogram
==
Th (1- E AE) + 20 (x),
where pr. is decomposed into a single-bin histogram 1(r AB) A with all mass on Ag and an induced histogram density function u evaluated by p (x) EN 1 1(x e A with rf as the coefficient of the i-th bin A, e is a hyper-parameter pre-specified before the decomposition, controlling the proportion between 1 [x E AB)/A A and (z*"(2)) See Figure[ the illustration of the decomposition. More specifically, the induced histogram density function p. in the second term of Fq. (4 repeessetts the difference between the full histogram function 7 and a single-hin histogram (x E C A where ( E AE)/ A only captures the mean. This difference indicates that fi captures the additional distribution information of 7 7* a) beyond its expectation E. [2" a)), incorporating higher- the performance of distributional RL. The additional leverage of pa:9 in the distributional loss explains the behavior differences between classical and distribution RL algorithms. We next demonstrate that a is valid probability density under certain in Proposition
information. This reflects the influence of using a full distribution on

---

Proposition 1. (Decomposition Validity) Denote pr E AF) = PF _ where pr is the coefficient on the bin AE. (x)) = - E A.)/A is a valid density if and only if é e21-1 PE.
The proof can be found in Appendix [ Proposition demonstrates that the return density decompo- sition is valid when the analysis maintains the standard categorical distributional learning in distributional RL.
is well specified 2 1 Pg- Under this condition,
Equivalence between Histogram Parameterization and Categorical Representation. The his- togram function is continuous estimator contrast to the discrete nature of categorical parameteri- zation. Although the underlying connection is relatively straightforward, equivalence representing a density function in Appendix for completeness. As a supplementary analysis, with attribution to [6]]. we also discuss the necessary theoretical underpinnings of the histogram density estimator in the context of distributional RL in Appendix F
still demonstrate their
Distributional RL: Entropy-regularized Neural FQI. We apply the decomposition in Eq. on the histogram density function, denoted as (l(o)(e)) of the target return Y** R (8; (4) Z (5), 2 (s()) in Eq. 2 of Neural FZI. Consequently,
have (z(==((0) A///A (1))) (ar). where A represents the interval that the expectation of the target return E A E. and ((()))(()) the induced histogram density function, similar to the role of 7a in Eq. Let HHU,V) be the cross-entropy between two probability measures U' and V, ie.. HUU V) - - U (xx) log (zz) dz. Immediately, we derive the following nuuropy-reullrrzze loss function form of Neural FZI for distributional RL in Proposition 2 The proof is provided in Appendix H
(1 c)
Y, falls into, i... E [Y]]]
Proposition 2. (Decomposed Neural FZI) Denote To as the histogram density estimator of Z (s. a) in Neural FZI. Based on the decomposition Eq. 3and the KL divergence as dgp the Neural FZI process in Eq. ii simplified as
Z argmin Z
log L8 : (A'z) aH(To ,0g )1,,
9o
T1
Mean Mean-Related Related Term
Regudarization Term
where a = 5/(1 - > O and the mean-related term is negative log- nigeliilldoodd centered on A
Connection between Neural FQI and FZI. A crucial bridge between classical RL and distributional RL is established in Proposition where we demonstrate that minimizing the mean-related term in Eq. 29 Neural FZI is asymptotically equivalent to minimizing Neural FQI in terms of the minimizers A + 0. As such, with this equivalence in the objective function, the remaining regularization term (zH(lllazzz)) 40 :..a)) ) in Eq. 29 thus interprets the potential benefits of CDRL over classical RL. For the uniformity of notation, we still use 8, a in the following analysis instead of Sc, dg.
Proposition 3. (Equivalence between the Mean-Related term in Decomposed Neural FZI and Neural FQI) In Eq. 29 assume the function class {Zo : 0 E θ is sufficiently large such that it contains the for all k, when A -> O, minimizing the mean-related term Eq. 29 implies
target [Y*)) for
P(Z
T a)) =,
(5)
where Q8 (8,0) is the scalar-valued target in the k-th phase of Neural FQI.
Proposition demonstrates that as A - O, the random variable 2s, a) with the limiting distribution in Neural FZI (distributional RL) will degrade to constant T (s, the minimizer (scalar-valued target) in Neural FQI (classical RL). That being said, minimizing the mean-related term in Neural FZI asymptotically equivalent to minimizing Neural FQI with the same limiting minimizer. A formal proof for convergence in distribution with the convergence rate o(A) is given in Appendix [
The realizable assumption that ZZg 0 E } is sufficiently large such that it contains {Y²")² implies good in-distribution generalization performance in each phase of Neural FZI, which is also adopted in (77] This connection is also consistent with the mean hean-preserving property of distributional
RL the tabular setting 48), but we extend this conclusion the arbitrary function approximation with histogram density estimator. Proposition especially focuses on the asymptotic property of the
5

---

mean-related term, which is different from existing convergence results based on the entire categorical distribution [48 6]. Given the connection between optimizing the mean-related term of Neural FZI with Neural FQI in Proposition we can leverage the regularization term afH( 9o (::))) to
explain the behavior difference between CDRL and classical RL, as analyzed later.
4.3 Uncertainty-awart Regularized Exploration
Regularization Effect. turns out that minimizing the regularization term (((( 96
in Neural FZI pushes 9 for the current return density estimator to catch up with the target return density function of (2f (s)) Importantly, (5)) encompasses the uncertainty of the entire return distribution the learning course beyond only its expectation, given that ( is the induced histogram density after applying the return density decomposition in Eq. 3 Since it is prevalent notion that distributional RL can significantly reduce intrinsic uncertainty of the environment [36 [0],
the derived Destribution-matcingg regularization term (p742(22))) 9p (i:))) helps to capture more
uncertainty of the environment by modeling higher moments of the whole return distribution beyond
the expectation. In Section [ we further demonstrate that this derived regularization contributes to an
uncertainty-awont regularized exploration effect in the policy optimization or actor critic.
Remark: Approximation of (z)), In practical distributional RL algorithms, we typically use temporal-different (TD) learning to attain the target probability density estimate 7it' (#2("') based on Eq. provided E (s, a)] exists and E 2 - PE in Proposition [ The approximation error of a (3(") is fundamentally determined by the TD learning nature. A desirable approximation of n (z((") intuitively leads to performance improvement in distributional RI.. As KI. divergence is used in CDRL, we also discuss the usage of KL divergence in distributional RL in Appendix G
5 Regularization Benefits in Actor Critic
Notations. In this section, we use boldface notations to represent random variables, such as (St. ar). at time t for clarity in the learning process of the actor critic.
5.1 Connection with MaxEnt RL
Motivation for the Connection. The maximum entropy regularization is commonly used in RL,
which has various conceptual and practical advantages. Firstly, the learned policy is encouraged to visit states with high entropy in the future, promoting the exploration of diverse actions [19 17 64] also considerably improves the learning speed (37] and therefore widely employed in otate-of- the-art algorithms, e.g., Soft Actor-Critt (SAC) [17. Similar empirical benefits of both distributional RL and MaxEnt RL motivate us to probe their underlying connection, especially by comparing their exploration effects.
Explicit Entropy Regularization in MaxEnt RL. MaxEnt RL explicitly encourages exploration by optimizing for policies to reach states with higher entropy in the future:
T
J(T) = Σ (( [r (st, ax) + (l(ss))s
(6)
1=0
where H (mo (-Is:)) Ea no (als;) log my (als;) and Par is the generated distribution following т. The temperature parameter |3 determines the relative importance of the entropy term against the cumulative rewards and thus controls the action diversity of the optimal policy learned via Eq. 6
Implicit Entropy Regularization in Distributional RL. For a direct comparison with MaxEnt RL, it is required to specifically analyze the impact of the regularization term in Eq. 29 Therefore, we directly incorporate the iiriibutionmmacchinn
regularization of distributional RL in Eq. 29 into the soft Q-vulue. The new Q function can be computed iteratively by applying a modified Bellman operator denoted as TT,, called Distribution- Entropy- Regullrzed Bellman Operator. Given fixed qo. T is defined
Actor Critic (AC) framework, enabling us to consider
Q (st. ax) & (St, a.) + 7Ex,, 110 [V (S1+1)],
where a new soft value function V (st) is defined by V (st) = Em,-, ~n IQ (st, ax) + I(H (1 (10001, ())))))
(8)

---

where f is continuous increasing func- tion over the cross-entropy H. 1 is the induced true target return histogram density function via the decomposition Eq. []] which excludes its expectation. Note that p can be approximated via boot- strap TD estimate f 1: 1) similar Eq. 29 In this specific tabular setting re- garding St; at, we particularly use do to approximate the true density function of (st,ar)) The f transformation over the cross-entropy H between je and gg "0"""(2) serves as the uncertainty-awar entropy regularization that we implicitly derive from value-based distributional RL in Section 4.2 By optimizing qe that is involved in the value-based critic component in actor critic, this regularization reduces the mismatch between the target return distribution and current estimate, aligning with the regularization effect 98 6.0 is optimized to catch up with the uncertainty involved in the target return distribution of p= iteratively expanding the agent's 's knowledge about the environment uncertainty to contribute more informative decisions. Next, we elaborate on its additional impact on policy learning in the actor critic compared to MaxFnt RI..
(x)(x)
A x*)(x))
xx)
qo q0=(x)
Figure 2: do is optimized disperse (left) or concen- trate (right) to align with the uncertainty of target return distributions.
analyzed in Section 4.3 As illustrated in Figure 2
Reward Augmentation for Policy Learning. As opposed to the vanilla entropy regularization in MaxEnt RL that explicitly encourages the policy to explore, our derived regularization term in the distributional loss of RL plays the role of reward augmentation for policy learning. Compared with classical RL, the augmented reward from the distributional loss incorporates additional knowledge of the return distribution in the learning process. As we will show later, the augmented reward encourages policies to reach states s; with actions a, ~ (.|sst), whose current action-state return distribution qe.at lags far behind the (estimated) environmental uncertainty from the target returns.
For a detailed comparison with MaxEnt RI., we now focus on the properties of our decomposed distribution matching regularization in the actor critic. In Lemma 1 we demonstrate that Distribution- Entropy- Regularized Bellman operator TH inherits the convergence property in the policy evaluation phase with a cumulative augmented reward function as the new objective function J'(T).
Lemma 1. (Distribution Fnttrrop - Policy Evaluation) Consider the distribution-entrupp regularized Bellman operator Tí in Eq. Aand assume H( 9g (..a)) is bounded for all (se, as) e
S x A. We define Q'+1 = T* Given qo, Q++1 will converge to corrected Q-value of as k + 20 with the new objective function J' J'(z)) defined as
J'(x)= =
Σ TE (e [r (St, a1) (H (yot. T6 (...1))]]
(9)
The updating rule is Tnew = arg EE En.wa [Q ac) + [(H (fCt: ,g ".."))]] in phase of policy optimization. Next, derive new policy iteration algorithm, called Aitribation-Entropp Regularized Policy Iteration (DERPI), alternating between policy evaluation and policy improvement. It provably converges to a policy regularized by the iestribution-matchinn term
(Distribution--nnnr) Eetrooonn--tttt---oe Regularize Policy Iteration) Repeatedly applying distribution- policy evaluation in Eq. 7 and the policy improvement, the policy converges to an optimal policy x²" such that Q (st, ax) 2 Q" (s), for all E IT.
Theorem 1.
Please refer to Appendix for the proof of Lemma and Theorem 1 Theorem I demonstrates that we incorporate the decomposed regularization into the actor critic in Eq. 9 we can design a variant of "soft policy iteration" [17] that can guarantee the convergence to an optimal policy given any fixed qe- In summary, our theoretical investigation is variant of the standard analytical framework in MaxEnt RL that allows a comparable analysis. Importantly, we next recognize a fundamental difference between our decomposed entropy regularization and the vanilla entropy regularization in MaxEnt RL.
Uncertainty-aware Regularized Exploration in CDRL Compared with MaxEnt RL. For the objective function (TT) in Eq. 6 of MaxEnt RL, the state-wise entropy ((T(-(s:) is maximized explicitly w.rl. æ for policies with higher entropy in terms of diverse actions to encourage an explicit exploration. For the objective function J'(xr) in Eq. 9 of distributional RL, the policy is implicitly optimized through the action selection process a, ~ π (-[ss) guided by an augmented reward signal from the distribution distribution-matciing regularization fHH (1 .66 ...).)). Concretely, the
7

---

learned policy is encouraged to visit state S, along with the policy-determined action via ~ (T|ss)), whose current action-state return distributions of .i lag far behind the target return distributions with a large discrepancy. This discrepancy is measured by the magnitude of the cross entropy between two return distributions of d0 si. .P: and f large discrepancy indicates that the uncertainty of the current return distribution is considerably misestimated for the considered states, enabling an uncertainty-awee exploration against these states in the policy optimization phase. This also indicates that the policy learning in CDRL is additionally driven by the uncertainty difference between the current and the target estimates, leading to a distinct exploration strategy of distributional RL relative to MaxEnt RL.
Interplay of Uncertainty-aware Regularization in Distributional Actor Critic. Putting the critic and actor learning together in distributional RL, we reveal their interplay impact pertinent to the uncertainty-awer regularized exploration. For the actor component, the policy learning seeks states and actions whose current return distribution estimate lags far behind the environmental uncertainty of the target returns. For the critic component, the critic learning reduces the return distribution mismatch on the states and actions explored by the policy, with two situations illustrated in Figure 2 This uncertainty-aware exploration effect arises from the decomposed regularization via the return density decomposition, interpreting the benefits of CDRL over classical RL.
6 Experiments
We comprehensively demonstrate our theoretical analysis using both Atari games and MuJoCo tasks. In Section 6.1] we validate that the uncertainty-awe aware regularization is crucial to the outperformance of CDRL over classical RL by varying e in the return density decomposition. We also investigate the mutual impacts between the vanilla entropy regularization in MaxEnt RL and the uncertainty-awan entropy regularization from CDRL in Section 6.2 More implementation details, including the description of basclines, are provided in Appendix
6.1 Regularization Effect in Performance by Varying
H(H, (e) = 0.8/0.5/0.1). For the categorical distributional loss in C51 or the distributional critic loss in the actor critic, we employ f instead of p.. as the target return distribution, leading to the decomposed algorithms, denoted by H(p, qo). This decomposed algorithm enables to assess the uncertainty-awan regularization effect of distributional RI. by directly comparing its performance with the classical RI. and CDRL.
Baseline Algorithm:
Experimental Details. We substantiate that the decomposed uncertainty entropy regularization, derived in Eq. [2 through the return density function decomposition, plays a crucial role in the empirical superiority of CDRL over classical RL. We compare CDRL with the decomposed bascline algorithm Hi, do) under different based on Eq. 3 To ensure a pre-specified e that guarantees a valid decomposition analyzed in Proposition
employ new notation E, which is proportional to
Alien
BattleZone
Breakout
Frostbite
Avergge Retum
Averag Retur
Avange Ratum
20:0
Corl
5
15:0
1RL-
1000
1::0
Millons Frames
Millons
Millona Framos
Gopher
Hero
Seaquest
StarGunner
1200
20403
Average Retum
1200
10.0
nafewritt to
15900
Has wale
Avarage Raturn
Avaroge Ralaro
2009
2010
4000
2003
Millons of Frames (M)
Millons of Frames (M)
Millons of Frames (M)
Milins of Frames (M)
Figure 3: Learning curves of value-based CDRL (C51) and the decomposed algorithm H(H,, go)(E 0.8/0.5/0.1) after applying the return distribution decomposition with different& & on eight Atari games. Results are averaged over three seeds, and the shade represents the standard deviation.
8

---

é but is more convenient in the implementation. See Appendix K.2 for more explanation, including the transformation equation between e and E, and the details of the baseline algorithm H(p.g))
Figure 3 showcases that as é gradually decreases from 0.8 to 0.1, learning curves of (.8/0.5/0.1) )), tend to degrade from C51 to DQN across most Atari games. The sensitivity of the decomposed algorithm ((,, Qa) regarding depends on the environment. Similar results MuJoCo environments can be found in Appendix [.] Overall, our empirical result corroborates that the decomposed uncertainty- entropy regularization from the categorical distributional loss is pivotal to the empirical advantage of CDRL over classical RL.
Results.
decomposed C51, i.e., H(n,
6.2 Mutual Impacts of the Two Entropy Regularization
Baseline Algorithms. For a detailed comparison of the mutual impacts between Vanilla Entropy (VE) in MaxEnt RL and Uncertainty-awart Entropy (UE) in CDRL, we conduct ablation study across several related baseline algorithms. We denote SAC with/without vanilla entropy as AC+VE and AC. We denote Distributional SAC (DSAC)
B4 with/without vanilla entropy as AC+UE+V+ VE and AC+UE. AC+UE also denoted as DAC. The implementation details can found in Appendix K
Experimental Details. We demonstrate that the two types of regularized exploration in MaxEnt RL and CDRL play distinct roles in policy learning when employed simultaneously. including mutual improvement or potential interference. We perform our experiments for both DSAC (C51) in Figure and DSAC (IQN) Figure Appendix [.2] where the latter used to examine the mutual impacts in quantile-hased distributional RI. heuristically.
walkerr2d
humanoidsstandup
reacher
bipedalwalkerraaddoore
Average Retur
+000
3000
closve
-150
8.5
halfcheetah
humanoid
swimmer
Average Return
1300
XG
AG
180
622 04
65
24 Time Steps (1e6)
62
Time Stepe (1e6)
Time Sleps (1e6)
Time Steps (1e6)
Figure 4: Learning curves AC, AC+VE (SAC), AC+UE (DAC) and AC+UB+VE (DSAC) over five seeds across seven MuJoCo environments where the distributional RL part is based on C51. (First Row): Mutual Improvement. (Second Row): Potential Interference.
Results. In the first row of Figure simultaneously employing uncertainty-aware and vanilla entropy regularization renders a mutual improvement. Conversely, the two kinds of regularizations, when adopted together, can also lead to performance degradation, exhibited in the second row in Figure For instance, AC+UE+VE outperforms both AC+ VE (SAC) and AC+UE (DAC) on humanoidstandup. while suffering from performance degradation on Ant and Swimmer. We posit that the potential interference may result from distinct exploration directions in the policy learning for the two types of regularizations. SAC optimizes the policy to visit states with high entropy, while distributional RL updates the policy to explore states and the associated actions whose current return distribution estimate lags far behind the environment uncertainty in target returns.
7 Extension to Quantile Distributional Loss
As an extension, we consider decomposing the quantile distributional loss, which is also commonly used in distributional RL, such as QR-DQN and IQN. Due to space limitations, more detailed description deferred to Appendix M The quantile distributional loss can be viewed as variant of composite quantile loss (72). We commit decomposing into mean-related term and residual

---

term, where the mean reeanerelate term is related to the expected quantile values. We demonstrate that minimizing the decomposed mean-related term asymptotically of quantiles approaches infinity. The induced residual term, therefore, captures the information from the return distribution that excludes expectation, serving as the benefit to explain the superiority of quantile-based based distributional RI.
pressernng 49 as the number
8 Conclusion and Discussion
In this study, we interpret the benefits of CDRL over classical RL as uncertainty-awarr regularization via return density decomposition. In contrast to the exploration to encourage diverse actions in MaxEnt RI, the uncertainty-awarc regularization in CDRL promotes exploring states where the environmental uncertainty is largely underestimated. Our study offers a novel exploration perspective to analyze the benefits of (categorical) distributional learning in RL.
Limitation and Future Work. The uncertainty-awarc regularized exploration from distributional loss is mainly founded on CDRL. Although briefly examined in Section 7 remains interesting yet challenging to extend our conclusion to general distributional RL, given that the analytical techniques, such as those in QR-DQN, are largely different from CDRL. We leave this extension for future work.
Acknowledgements
This work was primarily conducted while Ke Sun was at the University of Alberta. Some revi- sions were completed during his postdoctoral research at Harvard University, where he was sup- ported by NIH/NIDA P50DA054039, NIH/NIA 5P30AG073 107-03 GY3 Pilots, and NIH/NIDCR UH3DE028723. Bei Jiang and Linglong Kong were partially supported by grants from the Canada CIFAR AI Chairs program, the Alberta Machine Intelligence Institute (AMII,, and Natural Sciences and Engineering Council of Canada (NSERC), and Linglong Kong was also partially supported by grants from the Canada Research Chair program from NSERC. Xiaodong Yan was supported by the National Key R&DD Program of China (No. 2023YFA 1008701) and the National Natural Science Foundation of China (No. 12371292) equally. The authors express their gratitude for the insightful feedback provided by all reviewers and the area chairs, which significantly enhanced the initial version of this paper.
References
[1] Rohit Agrawal and Thibaut Horel. Optimal bounds between f-divergences and integral proba- bility metries. Journal of Machine Learning Research, 22(128):1-59) 2021.
[2] Martin Arjovsky and Léon Bottou. Towards principled methods for training generative adver- sarial networks. International Conference on Learning Representations, 2017.
[3] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pages 214-223. PMLR, 2017.
[4] Kamyar Azizzadenesheli, Emma Brunskill, and Animashree Anandkumar. Efficient exploration through bayesian deep genttworks. In 2018 Information Theory and Applications Workshop (ITA), pages 1-9. IEEE, 2018
[5] Mare G Bellemare, Will Dabney, and Rémi Munos. distributional perspective on reinforce- ment learning. International Conference on Machine Learning (ICML), 2017.
[6] Marc G Bellemare, Will Dabney, and Mark Rowland. Distributional reinforcement learning.
MIT Press, 2023.
[7] Mare G Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan. Stephan Hoyer, and Rémi Munos. The cramer distance as a solution to biased wasserstein gradients. arXiv preprint arXiv: 100. 2017.
[8] Yu Chen, Xiangcheng Zhang, Siwei Wang, and Longbo Huang. Provable risk-sensitive distribu- tional reinforcement learning with general function approximation. International Conference on Machine Learning, 2024.
10

---

[9] Tachyun Cho, Seungyub Han, Heesoo Lee, Kyungjae Lee, and Jungwoo Lee. Pitfall of optimism: Distributional reinforcement learning by randomizing risk criterion. Advances in Neural Information Processing Systems, 2023.
[10] Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. International Conference on Machine Learning (ICML), 2018.
[11] Will Dabney, Mark Rowland, Mare G Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. Association for the Advancement of Artificial Intelligence (AAAT), 2018.
[12] Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time-iil -ii.. Communications on pure and applied Mathematics, 29(4):389 461, 1976.
[13] Jianqing Fan, Zhaoran Wang, Yuchen Xic, and Zhuoran Yang. A theoretical analysis of deep q-learning. In Learning for Dynamics and Control, pages 486-489, PMLR, 2020.
[14] Jesse Farebrother, Jordi Orbay, Quan Vuong, Adrien Ali Taiga, Yevgen Chebotar, Ted Xiao, Alex Irpan, Sergey Levine, Pablo Samuel Castro, Aleksandra Faust, et al. Stop regressing: Training value functions via classification for scalable deep rl. arXiv preprint arXiv:2403.039500 2024.
[15] Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error actor-critic methods. International conference on machine learning, pages 1587-1596. PMLR, 2018.
[16] Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International Conference on Machine Learning, pages 1352-1361. PMLR, 2017.
[17] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018
[18] Tuomas Haarojja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al. Soft actor-critic algorithms and upplications. arXiv preprint arXiv: 1812,0:1 2018.
[19] Seungyul Han and Youngehul Sung. A max-min entropy framework for reinforcement learning. Advances in neural information processing systems (NeurIPS), 2021.
[20] Jianye Hao, Tianpei Yang, Hongyao Tang, Chenjia Bai, Jinyi Liu, Zhaopeng Meng, Peng Liu, and Zhen Wang. Exploration in deep reinforcement learning: From single-agent to multiagent domain. IEEE Transactions Neural Networks and Learning Systems, 2023.
[21] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network. NIPS Deep Learning Workshop, 2015.
[22] Peter J Huber. Robust estimation of location parameter. In Breakthroughs in slatistics: Methodology and distribution, pages 492-518. Springer, 1992.
[23] Peter Huber. Robust Statistics, volume 523. John Wiley & Sons, 2004.
[24]
Ehsan Imani and Martha White. Improving regression performance with distributional losses. In International Conference Machine Learning, pages 2157-2166. PMLR, 2018.
[25] Haque Ishfaq, Guangyuan Wang. Sami Nur Islam, and Doina Precup. Langevin soft actor-critie: Efficient exploration through uncertainty-drit critic learning. International Conference on Learning Representations,, 2015.
(26) Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. International Conference on Learning Representations, 2016.

---

[27] Tyler Kastner, Mark Rowland, Yunhao Tang, Murat A Erdogdu, and Amir- masir- Farahmand. Categorical distributional reinforcement learning with kullback-leibler divergence; Convergence and asymptotics. In Forty-second International Conference on Machine Learning, 2025.
[28] Qi Kuang, Zhoufan Zhu, Liwen Zhang, and Fan Zhou. Variance control for distributional reinforcement learning. International Conference on Machine Learning. 2023.
[29] Pawel Ladosz, Lilian Weng, Minwoo Kim, and Hyondong Oh. Exploration in deep reinforce- ment learning: A survey. Information Fusion, 85:1-22, 2022.
[30] Kimin Lee, Michael Laskin, Aravind Srinivas, and Picter Abbeel. Sunrise: simple unified framework for ensemble learning in deep reinforcement learning. In International Conference on Machine Learning, pages 6131-6141. PMLR, 2021.
[31] Shiau Hong Lim and Hyas Malik. Distributional reinforcement learning for risk-sensitive policies. Advances in Neural Information Processing Systems, 35:30977-30989, 2022.
[32] Owen Lockwood and Mei Si. A review of uncertainty for deep reinforcement learning. In Pro- ceedings of the AAAI Conference on Artificial Intelligence and Interactive Digital Entertainment, volume 18, pages 155-162, 2022.
[33] Clare Lyle, Marc G Bellemare, and Pablo Samuel Castro. A comparative analysis of expected and distributional reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 33, pages 4504 111444511220 2019.
[34] Xiaoteng Ma, Li Xia, Zhengyuan Zhou, Jun Yang, and Qianchuan Zhao. Dsac: Distributional soft actor critic for risk-sensitive reinforcement learning, arXiv preprint arXiv:2004. 14547, 2020.
[35] Yecheng Ma, Dinesh Jayaraman, and Osbert Bastani. Conservative offline distributional reinforcement learning, Advances in neural information processing systems, 34:19235-1 2021.
[36] Borislav Mavrin, Shangtong Zhang, Hengshuai Yao, Linglong Kong, Kaiwen Wu, and Yaoliang Yu. Distributional reinforcement learning for efficient exploration. International Conference on Machine Learning (ICML), 2019.
[37] Jincheng Mei, Chenjun Xiao, Csaba Szepesvari, and Dale Schuurmans. On the global con- vergence rates of softmax policy gradient methods. In International Conference on Machine Learning, pages 6820-6829. PMLR, 2020.
[38] Alberto Maria Metelli, Amarildo Likmeta, and Marcello Restelli. Propagating uncertainty in reinforcement learning via wasserstein barycenters. Advances in Neural Information Processing Systems, 32, 2019.
[39] Volodymyr Mnih, Koray Kavukeuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland. Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51897540: 2015.
[40] Mehryar Mohri. Foundations of machine learning, 2018.
[41] Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki Tanaka. Parametric return density estimation for reinforcement learning. Proceedings of the Twenty-Sixth Conference on Uncertainty in Artificial Intelligence (UAI), 2011.
[42] Rafael Müller, Simon Kornblith, and Geoffrey Hinton. When does label smoothing help?
Neural Information Processing Systems (NeurIPS), 2019.
[43] Thanh Tang Nguyen, Sunil Gupta, and Svetha Venkatesh. Distributional reinforcement learning with maximum mean discrepancy. Association for the Advancement of Artificial Intelligence (AAAI), 2020.
[44] Ian Osband, Charles Blundell, Alexander Pritzel, and Benjamin Van Roy. Deep exploration via bootstrapped dqn. Advances in neural information processing systems, 29, 2016.
12

---

[45] Ian Osband, Benjamin Van Roy. and Zheng Wen. Generalization and exploration via randomized value functions. In International Conference on Machine Learning, pages 2377-2386. PMLR, 2016.
[46] Yangchen Pan, Kirby Banman, and Martha White. Fuzzy tiling activations: A simple approach to learning sparse representations online. International Conference on Learning Representations, 2019.
[47] Martin Riedmiller. Neural fitted q cerration-fistt experiences with data efficient neural reinforcement learning method. In European conference on machine learning, pages 317-328. Springer, 2005.
[48] Mark Rowland, Mare Bellemare, Will Dabney, Rémi Munos, and Yee Whye Teh. An analysis of categorical distributional reinforcement learning. In International Conference on Artificial Intelligence and Statistics, pages 29-37. PMLR, 2018.
[49] Mark Rowland. Robert Dadashi, Saurabh Kumar, Rémi Munos, Mare G Bellemare, and Will Dabney, Statistics and samples in distributional reinforcement learning. International Confer- ence on Machine Learning (ICML), 2019.
[50] Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Harutyunyan, Karl Tuyls, Mare G Bellemare, and Will Dabney. An analysis of quantile temporal-differenn learning. Journal of Machine Learning Research (JMLR), 2024.
[51] Mark Rowland, Yunhao Tang, Clare Lyle, Rémi Munos, Mare Bellemare, and Will Dab- ney. The statistical benefits of quantile temporal-difference learning for value estimation. International Conference on Machine Learning, 2023.
[52] Mark Rowland, Li Kevin Wenliang, Rémi Munos, Clare Lyle, Yunhao Tang, and Will Dabney. lear-minimax-option optimal distributional reinforcement learning with generative model. Advances in Neural Information Processing Systems (NeurIPS), 2024.
[53] Yang Sui, Yukun Huang, Hongtu Zhu, and Fan Zhou. Adversarial learning of distributional reinforcement learning. In International Conference on Machine Learning, pages 32783-32796. PMLR, 2023.
[54] Ke Sun, Bei Jiang, and Linglong Kong. How does return distribution in distributional reinforce- ment learning help optimization? In ICML 2024 Workshop: Aligning Reinforcement Learning Experimentalists and Theorists, 2024.
[55] Ke Sun, Yi Liu, Yingnan Zhao, Hengshuai Yao, Shangling Jui, and Linglong Kong. Exploring the training robustness of distributional reinforcement learning against noisy state observa- tions. European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases (ECML-PKDD), 2023.
[56] Ke Sun, Yingnan Zhao, Yi Liu, Bei Jiang, and Linglong Kong. Distributional reinforcement learning with regularized wasserstein distance. Advances in Neural Information Processing Systems, 2024.
[57] Richard S Sutton and Andrew G Barto. Reinforcement learning: An Introduction. MIT press, 2018.
[58] Yunhao Tang and Shipra Agrawal. Exploration by distributional reinforcement learning. Inter- national Joint Conference on Artificial Intelligence (IJCAI), 2018.
[59] Kaiwen Wang, Owen Oertell, Alekh Agarwal, Nathan Kallus, and Wen Sun. More benefits of being distributional: Second-order bounds for reinforcement learning. International Conference on Machine Learning, 2024.
[60] Kaiwen Wang, Kevin Zhou, Runzhe Wu, Nathan Kallus, and Wen Sun. The benefits of being distributional: Small-loss bounds for reinforcement learning. Advances in neural information processing systems, 2023.
[61] Larry Wasserman. All of nonparametric statistics. Springer Science & Business Media, 2006.
13

---

[62] Christopher JCH Watkins and Peter Dayan. Q-learning Machine learning, 8(3-4):279-29) 1992.
[63] Li Kevin Wenliang, Grégoire Déletang, Matthew Aitchison, Marcus Hutter, Anian Ruoss, Arthur Gretton, and Mark Rowland. Distributional bellman operators over mean embeddings. International Conference on Machine Learning, 2024.
[64] Ronald Williams and Jing Peng. Function optimization using connectionist reinforcement learning algorithms. Connection Science, 3(3):241 1991.
[65] Harley Wiltzer, Mare G Bellemare, David Meger. Patrick Shafto, and Yash Jhaveri. Action gaps and advantages in continuous-time distributional reinforcement learning, Advances in Neural Information Processing Systems (NeurIPS), 2024.
[66] Harley Wiltzer, Jesse Farebrother, Arthur Gretton, and Mark Rowland. Foundations of mul- tivariate distributional reinforcement learning. Advances in Neural Information Processing Systems (NeurIPS), 2024.
[67] Runzhe Wu, Masatoshi Uchara, and Wen Sun. Distributional offline policy evaluation with predictive emrrr guarantees. International Conference Machine Learning, 2023.
[68] Derek Yang. Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tie-Yan Liu. Fully parameterized quantile function for distributional reinforcement learning. Advances in neural information processing systems, 32:6193-6200 2019.
[69] Liangyu Zhang, Yang Peng, Jiadong Liang. Wenhao Yang, and Zhihua Zhang. Estimation and inference in distributional reinforcement learning. Annals of Statistics, 2025.
[70] Pushi Zhang, Xiaoyu Chen, Li Zhao, Wei Xiong, Tao Qin, and Tie- Yan Liu. Distributional reinforcement learning for multi-dimensional reward functions. Advances in Neural Information Processing Systems, 34: 34:1519-1529, 2021.
[71] Fan Zhou, Jianing Wang, and Xingdong Feng. Non-crossing quantile regression for distri- butional reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020.
[72] Hui Zou and Ming Yuan, Composite quantile regression and the oracle model selection theory. The Annals of Statistics. 36 (3) 1108 1126, 2008
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
This study focuses on revealing the intrinsic benefits of adopting distributional learning in RI., particually under categorical distributional loss, from novel perspective of exploration.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results be expected to generalize other settings.
is fine include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: We have discussed the potential limitation of our work in Section 8 and valuable to extend the conclusion find to other types of distributional learning RL.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
15

---

Answer: [Yes]
Justification: We have made rigorous assumptions toward complete proof. The full proof is detailed in Appendix.
Guidelines:
The answer NA means that the paper does not include theoretical results. All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We have provided the implementation details in our proof, such as in Ap- pendix K
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, other means that appropriate to the research performed.
also be provided via detailed
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reprodueibility may be tricky in some cases, in which authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path reproducing or verifying the results.
5. Open access to data and code
16

---

The authors should provide scripts to reproduce all experimental results for the new reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
We are going to release our code once the paper accepted facilitate reproductivity.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ https/cccppblligggddpeeeeeeeeeee Coaee//ddeeSebmmmcia ssi onPolicy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We have provided sufficient experimental details in Appendix L
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All learning curves are reported across a certain number of random seeds and the shade represents the standard deviation.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
17

---

The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative rates).
show in tables or
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The experiments are conducted by using multiple RTX 3090 GPUs.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:///wwurrr...cpppaaaccccaaaa
Answer: [Yes]
Justification: We fully adhere to NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: As this paper focuses on understanding existing RL algorithms, we believe there is no negative societal impacts involved.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
18

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Question: Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Answer:[NA]
Justification:
The paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer:[NA)
Justification:
The paper does not use existing assets.
Guidelines:
The answer NA means that the paper does use existing assets. The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
19

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moddl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer[[NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsoureing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
20

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification:
We are not using LLM to propose our methodology.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.LM for what should or should not be described.
21

---

Appendix
Table of Contents
A Related Work: More Discussions about Unvertainty-oriented Exploration in RI.
23
More Details about Categorical Distributional RI. and Algorithm Description of C51
23
Explanation about the Efficacy Neural FZI Framework
D) Proof of Proposition 1
25
E Equivalence between Categorical Parameterization and Histogram Density Estimation
in Distributional RL
25
F Convergence Guarantee of Histogram Density Estimator in Distributional RL
26
Discussion about KT, Divergence in Distributional RL.
27
Properties of KL divergence in Distributional RL
27
Equivalence between Cross-Entropy Loss and KL Divergence in Neural
H
Proof Proposition 2
29
Proof of Proposition! 3
30
Convergence Proof of DERPI in Theorem
32
J.1
Proof of Distribution Entropy Regularized Policy Evaluation in Lemmal
!
32
J.2
Policy Improvement with Proof
J.3
Proof DERPI Theorem
32
J.4
Discussion about DERPI with Varying Go
33
K
iplemennaaion Details
33
K.
1
More Descriptions of Basclines Algorithms
33
K.
Replacing e with the ratio for Visualization
33
Hyner-parameters and Network structure
34
L
Experiments Results
35
L.1
Uncertainty-aware Regularization Effect by Varying Actor Critic
35
2.2
Mutual Impacts DSAC (IQN)
36
L.3 Ablation Study across Different Bin Sizes (Number of Atoms)
36
Discussion on Decomposing Quantile-based Distributional Loss
36
22

---

A Related Work: More Discussions about Uncertainty-orientel Exploration in RL
Uncertainty in RL. Uncertainty is ubiquitous in RL and sequential Decision-making and therefore harnessing uncertainty is always crucial in designing efficient algorithms [32]. In the literature of uncertainty quantification, uncertainty is often decomposed into two sources: aleatoric uncertainty and epistemic uncertainty.
Aleatoric uncertainty, also called intrinsic or environmental uncertainty, originates from the stochas- tic or probabilistic nature of the environment, encompassing three main sources: stochastic transi- tion dynamics, stochastic policy, and stochastic reward function. Aleatoric uncertainty determined by the environment, which is thus irreducible. However, we can design more efficient algorithms by capturing more environmental uncertainty in the learning process, e.g.. via distributional RL.
Epistemic uncertainty, also called parametric uncertainty, often originates from the stochasticity in statistical estimation the presence of limited data or incomplete knowledge. As opposed to aleatoric uncertainty, epistemic uncertainty is reducible and should decrease over more data, which contributes to a more reliable statistical estimation.
Uncertainty-oriented Exploration. There are few survey papers that comprehensively summarize existing exploration approaches 20). Following (20), we classify the exploration strategies into two main categories: uncertainty-oriinte exploration and intrinsic motivation-orienter exploration. The latter is inspired by psychology, which is not the focus of our study. Importantly, according the two categories of uncertainty in RL, uncertainty-orienti exploration, which often applies Optimism in the Face of Uncertainty (OFU) principle, involves aleatoric and epistemic uncertainty.
Epistemic uncertainty-orienter exploration takes advantage of the uncertainty in the (posterior) estimation value functions. The typical exploration methods include Bayesian framework 45 38) Bootstrap 44, and Ensemble methods (30). For instance, Bootstrapped DQN 44 maintains several independent Q-estimators and randomly samples one them, enabling the agent to perform temporally extended exploration.
Aleatoric incertainty-ooienter exploration aims to capture more environmental uncertainty from three sources of stochastic transition dynamics, stochastic policies. and stochastic reward function, all of which can be comprehensively integrated into return distribution. [9] employs Perturbed Quantile Regression (PQR) to promote the optimistic exploration within the distributional RL framework, while Decaying Left Truncated Variance (DLTV) [6 utilizes the variance from the learned return distributions. (58) investigates the approximate posterior sampling in distributional RL to encourage the exploration. By contrast, our primal goal in this study is to attribute the benefits of distributional RL to its intrinsic uncertainty-awaat
exploration we derived via return density decomposition instead of harnessing the learned return distribution to develop subsequent alcatoric uncertainty-oriente exploration strategies in (58] 36). On the other hand, MaxEnt RL [16] 17] utilizes the stochasticity of learned policy, one of the three sources in environmental uncertainty, to encourage diverse actions. Therefore, MaxEnt RL can also be categorized into the aleatoric uncertainty-orient oriented exploration, and it is thus intuitive and interesting make detailed comparison of the exploration effects between distributional RL and MaxEnt RL, conducted in Section 5.1 of our study.
B More Details about Categorical Distributional RL and Algorithm Description of C51
Distributional Loss and Projection in CDRL. Categorical Distributional RL uses the heuristic projection operator Iec which was defined as
S 21
Ic (8,,
5t,
82411
2; < y S 2
(10)
2,41
8=N
y > 2N
After applying the distributional Bellman operator TH on the current return distribution 17" (s, a) in each update, the resulting new distribution, which we denote as 77 (s, a), typically no longer lies
23

---

in the same (discrete) support with the original one on { To maintain the same support, the underpinning the KL divergence, CDRL additionally applies the projection operator IIc on the new distribution 1j" (s, a). This projection rule distributes the weight of 6, across the original support points { based on the linear interpolation. For example, if y lies in between two support points 2; and 2141, the probability mass on y is split between 2; and 2++1 with the weight inversely proportional to its distance ratio to 2; and 1. Therefore, the projection extends affinely to finite mixtures
of Dirac measures, such that for a mixture of Diracs Σ p; we have I[ç (Ei=1 P: =
Ei= P; (. The Cramér distance was recently studied as an alternative the Wasserstein distances in the context of generative models [7]. Recall the definition of Cramér distance in the following.
Definition 1. (Definition 3 (48]) The Cramér distance Eg between two distributtions V1, V2 9 (R). with cumulative distribution functions F2,, Fv, respectively, is defined by:
(1 (Frn Foz (x)) dr
1/2
l2 (v2,
Further, the supremum-Cramed metric Eg defined between two distribution functions E PP ((B)XxA by
lz (1p(01)) ple.e)).
sup (x,
Thus, the contraction of categorical distributional RL can be guaranteed under Cramér distance: Proposition 4. (Proposition 2 (88) The operator lcT is a V-contraction in Eg.
An insight behind this conclusion is that Cramér distance endows a particular subset with notion of orthogonal projection, and the orthogonal projection onto the subset is exactly the heuristic projection He (Proposition in [88). distance.
also states that the operator Iee is contractive under Wasserstein
Description of CDRL Algorithm: C51. With N 51, C51 instantiates the CDRL algorithm. To elaborate the algorithm, we first introduce the pushforward measure f E ((R) from Definition in [48]. This pushforward measure shifts the support of the probability measure l according the map f. which is commonly used in distributional RL literature. In particular, we consider an affine shift map fr,> : + R. defined by fr. = r yx. As Algorithm displays, we first apply the pushforward measure on the target return distribution As'' a²) by affinely shifting its support points, leading to a new distribution 77ss a). Next, we project the support points of ij(s, a) by employing He onto the original support, allowing us to compute the KL divergence in the end. Notably, we decompose the distributional objective function on the KL loss KL (Targa (s, a) a)))))((s. a)).
Algorithm CDRL Update (Adapted from Algorithm in (48)
Require: Number of atoms N, c.g., N = 51 in C51, the categorical distribution M( a) = Em P; 80, for the current return distribution.
Input: Sample transition (s, a, r,
1: if Policy evaluation: then
2:
TT-ss")
3: else if Control: then
4: a" +- arg maX EA ER-a(s)aa)) [R]
5: end if
6: 77s. + (frg) ns''s
((sssaa)) a²) # Distributional Bellmen update by applying दभ 7: Marzet (8, + Meñss, a) # Project target support points and then distribute the probabilities
Output: Compute the distributional loss KL(uungee (s, (ss,)) a)) # Choose KL divergence dp
C Explanation about the Efficacy of Neural FZI Framework
The main reason for adopting the Neural FZI framework in Section 4.T to interpret the behavior of distributional RL is that it is by far the closest framework to the practical algorithms, to the best of
24

---

our knowledge. Neural FZI or FQI was initially proposed in the context of offline RL with fixed dataset of transitions (s, a, s',r). ) the leverage of the Neural FZI framework does not assume a static dataset and accommodates an evolving dataset when using exploration mechanism. Specifically, our theoretical analysis focuses on the decomposition of the categorical distributional loss, but not on formal convergence proof of Neural FZI under an evolving dataset. The related convergence proof might be complicated beyond classical assumptions on Neural FQI or FZI in an offline setting. This distinction is important: the goal of our analysis is to reveal how the categorical loss can be interpreted as form of distributional entropy regularization in Proposition 2 regardless of whether the underlying data distribution changes over time. During each phase of Neural FZI, the return distribution can be updated based on newly collected transitions, allowing the exploration of new samples based on the decomposed distribution-matee matching regularization term
D Proof of Proposition
Proposition 1 (Decomposition Validity) Denote fr AR) - pP/A where pr is the coefficient on the bin Ag. (>) = Σn -. 1/1(1 A.) is a valid density if and only if 2 1 - Pg.
Proof. Recap valid probability density function requires non-negative and one-bounded probability in each bin and all probabilities should sum 1. We start prove all probabilities should sum 1, which is straightforward by taking the integral of both sides of Eq
f (x)(x = (1 -1- / (xr E AE)
dx + e f (o))(
(1- 1 (1 6) + (x)de,
A
(11)
1=
which directly implies (z)dz = 1. Next, we show necessity and sufficiency of non-negative and one-bounded probability in each bin.
can simplified as PE/4 A = (1 (1- -)- + ////A,, where P = fa( E Ag). Thus, P pR -e 2 e 2 - Pg. Obviously.
Necessity. (1) When x E Ag. Eq.
pe 1-6
1-6
(12)
which is guaranteed by the validity of PE (2) When x & A we have p./ A = ///"/A, i.e., When
x E Ag, We immediately have 1/, = P. I 1-pe < when e 2 1 PE- Also, 1t = P: 2 0.
Sufficiency, (1) When x E Ag, let D = 22 - 1=s 2 O. we have E 2 - PE. 1 = pe 1= <
in nature, (2) When & A Ag, p P. 2 nature. Let pf = 2h S 1, we have pj S e. We need to
take the intersection set of (1) and (2), and we find that
c 2 1 - PE = c21 PE 2 Pi:
(13)
which satisfies the condition in (2). Thus, the intersection set of (1) and (2) would be e 2 PE.
In summary, as e 2 1 - Pg is both the necessary and sufficient condition, we have the conclusion that ((x:)) is valid probability density function > 221 1 PE-
E Equivalence between Categorical Parameterization and Histogram Density Estimation in Distributional RL
Proposition 5. Suppose the target categorical distribution Ei' and the target histogram
function l((x)= = = pt (x E A,) updating the parameterized categorical distribution co under KL divergence is equivalent to updating the parameterized histogram function he.
25

---

Proof. For the histogram density estimator hg and the true target density function p(z), we can simplify the KL divergence as follows.
(z)
P=(2)
DRI ha)
Σ
log
dr
h
N
(z(
Pz(x)
N
(x(x)
hg
log
dr
log
dx
(14)
A
44-1
(0)
(x(x)
log
de
(b)
Pt log
hgg
(c)
Σ pi log hig
A
where hg is determined by and 0, which is independent of (a) is true because the target distribution with all p. is fixed. (b) follows because pı(x) remains constant for æ E [1, 4+1]. Finally, (c) holds as the remaining term involving p., and A is also constant.
On the other hand, we consider the KL-based objective function in learning categorical distribution estimator. Given the target categorical distribution e = Σx= PrOg, where the probability p; is fixed for each atom 24, we aim updating the current categorical estimator co. Then, we have:
N
N
N
N
(K co) = Σ log 1=1
pi C
logp; - Σ 1=1
pr log co x -EP log cj. 1=1
(15)
where co -E= c/ is the current categorical estimator and oj is the learnable probability. By comparing the final loss function forms in Eq. and Eq. [1]] it turns out that they are equivalent as both cg and hg are the learnable probabilities, which are parameterized by the same neural network.
Remark. In CDRL, we use a discrete categorical distribution with probabilities centered on the fixed atoms { In contrast, the histogram density estimator in our analysis is a continuous function defined on [20, 2, enabling more nuanced analysis within continuous functions. Proposition indicates that minimizing the KL divergence with the categorical distribution in Eq. I5 amounts to the cross-entropy loss with the parameterized histogram function in Eq.
F Convergence Guarantee of Histogram Density Estimator in Distributional RL
Histogram Function Parameterization Error: Uniform Convergence in Probability. The previous discrete categorical parameterization error bound [8 (Proposition 3) is derived between the true return distribution and the limiting return distribution denoted as iteratively updated via the Bellman operator Ie in expectation, without considering an asymptotic analysis when the number of sampled {s,, 0; a pairs goes to infinity. As a complementary result, we provide a uniform convergence rate for the histogram density estimator in the context of distributional RL. In this -e.0 particular analysis within this subsection, we denote Pe as the density function estimator for the we show that the sample-based histogram estimator Pc can approximate any arbitrary continuous limiting density function p under a mild condition. This ensures the use of a histogram density estimator in the implementation of our subsequent algorithm adapted from CDRL.
25.4 true limiting return distribution Iee IT with its true density s.a In Theorem 2
Theorem 2. (Uniform Convergence Rate in Probability) Suppose Pe (z) is Lipschitz continuous, and the support of random variable is partitioned by N bins with bin size A. Then
log N n
sup IP (x) PC (x) (x)l (A))
(16)
Proof. Our proof is mainly based on the non-parametric statistics analysis 61 In particular, the difference of Pe (x: PC (1r) can be written
DC (x) Pc 5.a (x)
E (PC
(x))))))))))))))))))
PC
8.0
(a)
+
Pc
25.0
(x) E
(17)
stochastic variation
26

---

(1) The first bias term. Without loss of generality, we consider I' E Ak, we have
P(X= Ak)
E (PC
1h (x))
A
MOTRA
1)A ply)dy
A
(18)
F(lo + (kc 1)4) - F(lo + (K - 1)4)
10 + K (lo + (k: 1)4)
Pe s.a (z'),
where the last equality is based on the mean value theorem. According to the L-Lipschitz continuity property, we have
E (Pc th (x)) (x)) (x') Pe (x)] S L|r' x S LA
(19)
(2) The second stochastic variation term. If we a Ak. then PC 23.0 we thus have
Pk E 1(X; AE).
P
sup
Pc (x) - E (Pé (x))
>e)
max
j= NN
7
max
Σ
>D<
(20)
j=1,,, NN
N
E > >
j=1
S exp (-2n (by Hoeffding's inequality),
where in the last inequality we know that the indicator function is bounded in [0, 1]. We then let the last term be a constant independent of N, n, A and simplify the order of c. Then, we have:
(V123) log N
sup x
Pe
(x) E (P (x)) = Or
n
(21)
In summary, the above inequality holds for each thus have the uniform convergence rate of histogram density estimator
sup " (x) Ve (x:) < sup JE (x:))))) Pe ( + sup PC (x) E (Pé (x))
log (V AA
(22)
= 0 +0p
G Discussion about KL Divergence in Distributional RL
G.1 Properties of KL divergence in Distributional RL
Remark on KL Divergence. As stated in Section 3 of CDRL 5 when the categorical parameter- ization is applied after the projection operator IIc, the distributional Bellman operator T has the contraction guarantee under Cramér distance or Wasserstein distance [48], albeit the direct use of a non-expansive KI. divergence [T.. Similarly, our histogram density parameterization with the projection He and KI. divergence also enjoys a contraction property due to the equivalence between optimizing histogram function and categorical distribution analyzed in Appendix
We summarize
some properties of KL divergence in distributional RL in Proposition 6
27

---

Proposition 6. Given two probability measures and v. we define the supreme DEL as a functional P(X)SXA x P(X)SxA R. i.e. (R..H.2) v) (A DxI. (u(, v(s,) we have:
(1) is non-expaniive distributional Bellman operator under D i.e.
(R 2, T Z2) S DRL(Z, Z2),
(23)
(2) (2 2) - implies the Wasserstein distance W (Z,, 2) -
Proof. We first assume Zo is absolutely continuous and the supports of two distributions in KL divergence have a negligible intersection 22]. under which the KL divergence is well-defined.
(1) The contraction analysis of distributional Bellman operator Ix under a distribution divergence dp depends on its scale sensitive (S) and sum invariant (I) properties 77 [5]. We say dp is scale sensitive (of order T) if there exists T > O, such that for all random variables X, Y and real value a > 0, d,(aX, aY) S da Y). dp has the sum invariant property whenever a random variable A is independent from X, y, we have d, + X, + Y) S dx, Y). We first prove that the D is sum-invariant, which is based on the dual form of KL divergence via the variational representation [2
(X,) = sup [Ex[f(x)] log (By (e) sec
(24)
where L is the space of bounded measurable functions. Consequently, have
Dxt(A X X+Y)) sup fec
{Ez, x+ x+x(((=1)] - log [Ezz=AtY (e
(EA Ex [f + +a]]] log (BA By (eflutto]])
sup
SEC
(6)
sup [EAEx [x + a)] BA log (E (eflotai))
SEC
sup (EA Ex [f + a)] - log (Ry (eflotta])
(25)
fEL
(c)
EA sup (xx [f + a)] - log (y (cllota]) SEL
S
(d)
EA sup {xx [x(x)]] - log Ey (coto)))
gee
(X,,),
where (a) results from the independence between A and X (Y), (b) and (c) rely on the Jensen inequality for the function log and the operator sup. (d) is because the translation is still within the same bounded functional space. Next, we show that DKL is not seale-sensitive, where we denote the probability density function of X and Y as p and q.
__ P log
1 1 (
p(y) log p(y) dy (x,
DKL(aX,
dr dz=
(26)
((v)
Putting the two properties together and given two return distributions (i a) and (z a), we have the non-expansive contraction property of the supremal form of DxL as follows.
DRL (I (1.5%2) (2) sup D (T
sup (R(s, + ( M(s,a) + (2((''a)))
8.0
(c) DKL (s))))))
(27)
( DRL (s',
S sup DKL (2)(
= DR (Z1, (Z),
28

---

where (e) relies on the sum invariant property of DKL and (f) holds due to the non-scale sensitive property of DKI. By applying the well-known Banach fixed point theorem, we have unique return distribution when convergence of distributional dynamic programming under DKL.
(2) By the definition of Dki, have sups.a DKL (Z (s, a), [(s a() 0 implies DRL. (Z,, Z) - O. Dxi (Zn Z) - 0 implies the total variation distance 8(Zn Z) + according to straightforward application of Pinsker's inequality
1 (Zn, 2) S V
V DKL (Zn,2)
2z, S V 22DKe (m, ~0
1
(28)
Based on Theorem 2 in WGAN [3. 8(2,, 2) + 0 implies (Z Z) + O. This is trivial by recalling the fact that 6 and W give the strong and weak topologies on the dual of (C(2), 1l when restricted to Prob(X').
G.2 Equivalence between Cross-Entropy Loss and KL Divergence in Neural FZI
If the target density function in evaluating the KL divergence is not fixed, using cross-entropy loss instead of the KL divergence may underestimate the uncertainty of return since this simplification may fail to capture the exact shape or uncertainty spread of the true target return distribution. However, this underestimation issue does occur in our analysis. Particularly, the leverage of the target network in Neural FZI, which is fixed in the updating of each phase, guarantees that the KI. divergence is exactly proportional to the cross-entropy loss. Figure [ suggests that C51 I with cross-entropy loss (DSAC_ AB) behaves similarly to the vanilla C51 equipped with KL divergence (DSAC) both three Atari games and MuJoCo environments with continuous action space.
Breakout
Pong
Seaquest
550
DON
DON
1000
DON
Average Return
Average Retur
Average Retum
C51
10
C51
OE
6000
CS1
4000
250
100
2000
0.4
0.0
0.2
044
0.5
Time Sleps (1c7)
Time Steps (1e7)
Time Steps (1e7)
ant
swimmer
bipedalwalkerhardcorr
6000
120
Averuge Return
5.00
Average Renum
100
Average Return
4003
100
5003
2003
1000
-1000
-1m0
Stape
Figure 5: (First row) Learning curves of C51 under cross- entropy loss on Atari games over seeds. (Second row) Learning curves of DSAC with C51 under cross- cross-entropy loss MuJoCo environments over five seeds.
H Proof of Proposition
3.4 as the histogram density estimator of Z (s. a) Proposition 2 Decomposed Neural FZI) Denote dg in Neural FZI. Based on the decomposition in Eq. and the KL divergence as dp, the Neural FZI process in Eq. 2 is simplified as
z x++ argmin
log $1.0, ".... (A + HllTrmmae))) 96 g6
(29)
qu
12
Mean- Related Term
Regularization Term
where =/(1 - > 0 and the mean-related term is negative likelihod centered on A
29

---

Proof. Firstly, given fixed p(z) we know that minimizing DEL (p, qo) is equivalent to minimizing Hipp q) by following
Pı(
A
(P,,P, 90) =
Σ
log
dr
A
%
4-11
1m1
Σ
(z
log
9 dr
((( E/e
log
dx
(30)
A
A
A A
H(pp 90) H((p)
x H(pp90) 90)
where = E 1Z E / and go = E q/A. Based on H((p. go). we use (x))=(()(())
N
to denote the target probability density function of the random variable (ss,, 0s) + (s', TT (s())) Then, derive the objective function within each Neural FZI as
H(p 140 (
1
-11 k) Σ
1(x ((xE)))
log
4g
81,01 (Aj)
dx
Σ
N
Fj
log
4o
de
A
A
A
A
f=1
j=1
1--1
a. (A)) + eH(TI)mmaa)) 26 (6.11)) + +(1 (1-6) E1- (1-c)(- T
n
1
log 90 s. (A²) at (( .0....).. where a 1
n
(31)
where recall that (o( = Ei 1 (z) E A.)/A En ///A A for conciseness and denote 90 qg.": شك Go $1.0 (A,)/A. The cross-entropy H(Z)((z(66)) 9 (aa) is based on the discrete distribution when i = 1,.... N. A/e represent the interval that E [R(s; a.) + 22 (8) (zx2) falls into, i..., E (R(s; + ZZg. (5%, 72(())) E A
I Proof of Proposition 3
Proposition
Equivalence between the Mean-Related term in Decomposed Neural FZI and Neural
FQI) In Eq. 29 assume the function class (Zo : 0 E θ} is sufficiently large such that it contains the when A - O, for all k. minimizing the mean-related term in Eq. 29 implics
target (Y: L
P(Z (1,4++((,,0)) - (s, a))= 1, and
(x) F8 Fo (nay (x) do o(A),
(32)
where ToRQ (s, a) is the scalar-valued target in the k-th phase Neural FQI, and 8ywey (s.a) the Dirac delta function defined the scalar T (s,a).
Proof. Limiting Case. Firstly, we define the distributional Bellman optimality operator Tap as follows:
TTh ITH a) 2 R(s, a) +S"
(33)
where S' P(- 8, a, and a" arginaxlE [2 a')) If (Zo : 0 [) is sufficiently large
a'
enough such that it contains Zg- ((Y,))m J 1) then optimizing Neural FZI in Eq. 2 leads to 2 - Top
Secondly, we apply the return density decomposition on the target histogram function (x)) Consider the parameterized histogram density function he and denote h / A as the bin height in the
30

---

bin Ag. under the KL divergence between the first histogram function 1(x E Ag) with hz)),, the objective function simplified as
DKL. (1 E AF)/A,hor1)
log
A I
dx = log hg
(34)
JaEAg
A
Since [Ze : 0 E } is sufficiently large enough that represent the pdf of {Y," implics that {Zo : 0 E } can represent the mcan-related term part in its pdf via the return density decomposition. The KL minimizer would be h2 = 1(r E AF)/ A in expectation. Then, lima-so arg minn, ( E Ar)/A,hh(rr) = 6g [.]],, where (2 is Dirac Delta function centered at E [Z a]] and can be viewed
Ji=1+ it also
generalized probability density function. That being said, the limiting probability density function (pdf) converges to Dirac delta function E [Zurgetss (s, a)). The limit behavior from histogram function p to a continuous one for Zuaget is guaranteed by Theorem 2 and this also applies from hg to Zo. In Neural FZI, we have Zuarer
I Zg.. Here, we use Z a) as the random variable whose edf is the limiting distribution. According to the definition of the Dirac function, in the limiting case where A - 0, we attain that
P(Z a) = E [ Zz- (s.a))) =1
(35)
This is because the pdf of the limiting return random variable ' a) is a Dirac delta function, which implies that the random variable takes this constant value with probability one. Due to the linearity of expectation in Lemma 4 of [5], we have
E I Zg- (8. a)] = TME [Z (s, a)] T " (8,
(36)
Finally, we obtain the convergence in probability one in the limiting case:
F(ZØ+1
TT (s) = as A 0
(37)
Convergence in Distribution. The connection established above is in the limiting case. Alternatively, we can provide more formal proof by using the language of convergence in distribution. Here, we use z to replace 21+1 to explicitly consider its asymptotic behavior. According to the fact that fx E A is the optimizer when minimizing the mean-related term in Eq. 29 given a fixed A, the convergence in distribution is:
lim 2( lim D(1 E AR)/A) = o0 (s.a)) A
(38)
A-30
where &oo (s.a) is the Dirac Delta function centered at Topt Q as. D6 (s.a)) is the corresponding step function, where ((67m*g)) (x)))))))) = 1 if 2 Top Q/. (s, a). and equals otherwise. Note that the convergence in distribution in terms of the Dirac delta function implies that P(Z (s.a)= = TU (s, a)) = 1 as A + in Eq 37
Convergence Rate. In order to characterize how the difference varies when A + O, we further define Ag = [Te, le41) and we have:
1 ((T'T (s,a) -to)² T
Pao (x) Formag--an (ao)) (z)) dx
2A
(a² + (A - o))
2A
S A/2
o(A))
(39)
where Qb- a) E [I (s, a)] E Ag and we denote 7 (s, (s, a) - 4. The first equality holds go (z:), the KL minimizer while minimizing the related term, will follow uniform distribution on A F. i.e.. 9o 1(r E A Thus, the integral of LHS would be the area of two centralized triangles accordingly. The inequality holds as the maximizer is obtained when
a = A or O. The result implies that the convergence rate in distribution difference is o(A).
31

---

J Convergence Proof of DERPI in Theorem I
J.1 Proof of DistribttonnFattoy---egggrrrr Policy Evaluation in Lemmal 1
Policy Evaluation) Consider the distributiinneettroyy- regularized Bellman operator Tá in Eq. and assume H(p :00 is bounded for all (st, ar) A. Define Q*+1 "T" Q*. then Q4+1 will converge corrected Q-value of π as k with the new objective function J' (²) defined as
J'(T)
Σ E (o Ir (s., ax) (f (porae (a)
1=
Proof. Firstly, we plug in (8 into RHS of the iteration in Eq. then we obtain
T T7Q((s (8)
=r(s;, a,) + 1
(S++1)]
= r (s,, az) + y1En, 11 P( (8).))))))))10)
11 + ( IQ (S++1, at-1]]
ra (st, ar) + 7
1Q (st+1, a(+1)]
(40)
where 'n (st. a,) - (S,, a.) + P a.---- [((((=*+))00+)) 9 - ++++))))))
))] the en-
tropy augmented reward. Applying the standard convergence results for policy evaluation [57], can attain that this Bellman updating under T is convergent under the assumption of |.A| < oc and bounded entropy augmented rewards Tx-
J.2 Policy Improvement with Proof
Lemma (Distribution Entropy-Regulln
Policy Improvement) Let E II and new pol- icy Tnew be updated via the policy improvement step in the policy optimization: Tnew
arg ma Exrr [Q (st, at) + f(H (Jrt." 40 ())). Then Q Tee (st, ax) 2 awu (st, (s,,a)) for all (st, a.) E S x A with |A| < xc.
Proof. The policy improvement in Lemma 2 implies that Ear- Tame (Qaa (s:, at) f(H
(H % ((.3+)) Em.~ 1QTH a,) f(H (p .60 ()).))))
We consider the Bellman equation via the distibution--ntooyyyrrgaaaicii Bellman operator Tzk (st, ac)
Q TTu
r (Se, az) + 2I. 1 TTH (8c+1)]
r -st, a.) + 7 BaI
[f(H(1 0::11111110:1 01:1
(1+1:9+11) + QTiaa (s1+1. a++1)] 20
S (Sc, a.) + 71E 1 Ea BaLLINMeN ((
90
Bc+1 (l41)))11))
+ Q (St+1 ac+1)]
(St, at) + yE r P.ar mo [f [H((p=--1) .90
(14100411))
yE P. ~om [Quau (s 2(41)]
Tame (St, at) + 7B Paa P Tew
Ta (841),
< TM (8141, 1t+1),
(41)
where Ria (St+1. at+1) indicates that the future actions are taking following Told, given S.41 and a,+1- We have repeated expanded aTu on the RHS by applying the distribution-enrooyy -ntropy-reuullrizee distributional Bellman operator. Each following step will then incorporate the actions following the new policy. Convergence to Fiew follows from Lemma
JJ3 Proof of DERPI in Theorem i
Theorem (Distribution-Entree -eaee-ooooeouolrrriieee Regularized Policy Iteration) Repeatedly applying distribution- autropy-regularicg policy evaluation in Eq. 7 and the policy improvement, the policy converges to an optimal policy x² such that Q" (st, ax) 2 Q (s).a) ax) for all II.
32

---

Proof. The proof is similar to soft policy iteration [7]. For completeness, we provide the proof here. By LLLmmaa 2 as the number of iteration increases, the sequence Q at i-th iteration monotonically increasing. Since we assume the uncertainty-awal entropy is bounded, the Q" is thus bounded as the rewards are bounded. Hence, the sequence will converge to some π². Further, we prove that x in fact optimal. At the convergence point, for all E IT, it must be case that:
Ear-s TQTHM (Se; ax) 2 Ear The (St, ar]]
TQTH
According to the proof in Lemma 2 can attain Q" (st, ar) Q" (sr, a.) for (st, az). That is
to say, the "corrected" value function of any other policy II is lower than the converged policy, indicating that x² is optimal.
J.4 Discussion about DERPI with Varying
In the tabular setting, have shown that the convergence of DERPI holds given fixed go. The primary goal for us to derive this convergence result is to demonstrate the uncertainty- aware regular- ized exploration promoted by the decomposed regularization from (categorical) distributional loss. we hope to develop further algorithm in the function approximation, we need to consider how to interplay parameterized function, policy, and q. For example, we may leverage separate neural network works for each component. Alternatively, we can use one single neural network to represent the whole return distribution (go) and then take the expectation to evaluate the Q function.
K
Implementation Details
K.1 More Descriptions of Baselines Algorithms
Algorithms in Section 6.1
DQN 39 and C51
H(P, (E X): variant of C51 algorithm, where we replace the original target histogram function fraa with the induced fa for each (s, a) pair in the update. By varying X, H(u go) relies on the distributional loss to different extents the RL learning. For examples, when 1, H(,,, go) = X) degenerates to the vanilla C51 I algorithm. On the contrary, decreasing & in H.l,s, ga) will reduce the leverage of knowledge from the distributional loss, leading to performance degradation in distributional learning context.
Algorithms in Section 6.2
AC: This implementation is the same as AC.
AC+VE: This is exactly the standard SAC algorithm.
AC+UE: This implementation is also the same as DAC (C51), where distributional critic loss in the AC algorithm.
AC+UE+VE: Based on the SAC algorithm, i.e., AC+VE, we additionally use the distribution objective in C51 the critic loss.
K.2 Replacing e with the ratio for Visualization
shares the same utility as e, but it is more convenient in implementation. is defined as the mass proportion centered at the bin that contains the expectation when transporting the mass to other bins. A large proportion probability e, which transports less mass other bins, corresponds large in Eq. 3 Increasing indicates that the decomposed algorithm performs more similarly to a pure CDRL algorithm. As Proposition [ elucidates, the return density decomposition requires that exceed certain thresholds to ensure the resultant decomposed fi qualifies valid density function. In practice, pinpointing this lower boundary for in each iteration to regulate its range could be prohibitively time-intensive. A more pragmatic approach involves redistributing the mass from the bin that contains the expectation to other bins in specified ratios, thereby introducing the corresponding ratio term By varying e from 0 to 1, it invariably meets the validity condition outlined in Proposition T thereby streamlining the process for conducting ablation studies concerning a as demonstrated in Figure
33

---

To delineate the relationship between the ratio and the coefficient é in constructing fr after some calculations we establish their equivalence as follows:
PE (1-c)
(42)
PEΕ
where Pg represents the weighting assigned to the bin Ag as specified in Proposition | The resulting E E [0, 1] has a monotonically increasing relationship with e. addition, = I implies = 1. These properties facilitate the visualization without undermining our conclusion.
Decomposition Details. By varying E, can evaluate via the transformation equation in Eq. 42 which guarantees the validity of return density decomposition. Next, under different e, we compute the induced histogram density fr via the return density decomposition in Eq.
N
(x: & AB)/E + ( AB)E,
(43)
where combines Eq.
3 and Eq. 42
Importantly, by summing all the probabilities of p/ in μt, we have:
PE
PE
(1
=
(44)
This substantiates the validity of our decomposition by using instead of E for visualization. Next, we replace 5=.- with fi in C51 or the critic loss in Distributional AC (C51) the decomposed algorithm M(H, ge) and compare the performance of all considered algorithms. Please refer to the code the implementation for more details.
K.3 Hypperppsamettes and Network structure
Our implementation is adapted from the popular RLKit platform. For Distributional SAC with C51, we use 51 atoms similar to the C51 [5]. For distributional SAC with quantile regression, instead of using fixed quantiles in QR-DQN, we leverage the quantile fraction generation based on IQN [10] that uniformly samples quantile fractions in order to approximate the full quantile function. In particular, we fix the number of quantile fractions as N and keep them in ascending order. Besides, we adapt
Table 1: Slyee--araameters Sheet.
Hyperparameter
Value
Shared
Policy network learning
3c-4
(Quantile) Value network learning rate
3e-4
Optimization
Adam
Discount factor
0.99
Target smoothing
5e-3
Batch size
256
Replay buffer size
le6
Minimum steps hefore training
le4
DSAC with C5}
Number of Atoms (N)
51
DSAC with IQN
Number of quantile fractions (N)
32
Quantile fraction embedding size
Huber regression threshold
Hyperparameter
Temperature Parameter /
Max episode lenght
Walker2d-v2
0.2
1000
Swimmer-v2
0.2
1000
Reacher-v2
0.2
1000
Ant-v2
1000
HalfChectah-v2
0.2
1000
Humanoid-v2
0.05
1000
HumanoidSSanndup-222
1000
BipedalWalkerHHardonne22
0.002
2000
34

---

the sampling as O, = = ee/E N-1 hyper-parameters, which are listed in Table I and network structure as in the original distributional SAC paper (4
i=0 e. where e; U0, 11, N.
N. We adopt the same
L Experiments Results
L.1 Uncertainty-aware Regularization Effect by Varying é in Actor Critic
humaroldstanduu
reacher
swimmer
walker2d
1000
Average Relurn
bioon
kina
CAC
EA5
SAC
SAC HL1. IIAh miti
-00
Ma
Nia Ai
Ha 1.Xp-da
Ma
Nia 1o
Time Steps (1e6)
Steps (1e6)
Time (1e6)
Time Steps (166)
Figure 6: Learning curves DSAC (C51) with the return distribution decomposition H(H, qo) under different 8.
We study the uncertainty-aware regularization effect from being categorical distributional in the actor- critic framework, where decompose the C51 critic loss in Distributional SAC (DSAC) according to Eq. 3 We denote the decomposed DSAC (C51) with different & as H(p, qe (E = 0.9 /..9/..5/0.1.. As suggested in Figure 6 the performance of ((, tends vary from the vanilla DSAC (C51) to SAC with the decreasing of on four MuloCo environments. In some environments, the difference of H(,,,0o) across various may not be pronounced between DSAC (C51) and SAC. We hypothesize that the algorithm performance is not sufficiently sensitive when e changes within this restricted range. Although (0, 1) is designed to guarantee valid density decomposition, it does not guarantee that in Eq. 3 can flexibly vary from 0 to 1. It is worth noting that our return density decomposition valid only when 2 Pe as shown in Proposition] [ and therefore can not strictly go to where H(p, 9a) would degenerate to SAC ideally. Therefore, compared with the ablation study in Figure the trend varying from DSAC to SAC in Figure decreasing e may not be as pronounced as that in value- basue RL evaluated on Atari games. One crucial reason behind is that the actor-cris critic architecture generally perceived be more prone to instability compared value-bare in RI. As outlined in 15), this instability stems from the policy updates, which likely introduces additional bias or variance from the critic learning process.
learning
halfcheetah
humanoid
walker2d
humanoidstandup
Average Return
1203
6300
3030
4000
2000
2593
200
AC-US-VE
1999
ACHLE--V
50000
ADHUC-VV
ant
reacher
swimmer
bipeddalwalkerhaddoor
Average Return
*00
AC
Ac AOME waiDe ADHLE-VV
ADVVE
AD-VE
NCVE ACHLE ACHUE-VE
AD-UE
AC-UZ
AD-UEHEE
AD-IE-VE
-100
02 044
0,5
02
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Figure 7: Learning curves of AC, AC,AC+VV VE (SAC), AC+UE (DAC) and AC+UE+ VE (DSAC) over five seeds across eight MuJoCo environments where DAC and DSAC are based on IQN. (First Row): Mutual improvement. (Second Row): Potential interference.
35

---

L.2 Mutual Impacts on DSAC (IQN)
To extend the mutual impact of the two types of regularization to broader distributional RL algorithm, we investigate the learning behavior of distributional RI. based IQN. The conclusion when using IQN is similar to that when using categorical distributional learning in Figure In particular, in the first row of Figure 7 simultaneously employing uncertainty-aware and vanilla entropy regularization renders a mutual improvement. Conversely, the two kinds of regularizations, when adopted together, can also lead performance degradation, as exhibited in the second row in Figure 7 For instance, on Swimmer and Reacher, AC+ UE+ VE is significantly inferior to AC+UE or AC+ VE. These results about potential interference also serve as the emprical evidence to reveal distinct exploration directions in the policy learning for the two types of regularizations.
L.3 Ablation Study across Different Bin Sizes (Number of Atoms)
To further demonstrate our regularization effect based on the return density decomposition, we conducted an additional ablation study by varying the number of bins/atoms (equivalent to adjusting the bin sizes) of both C51 and our decompose algorithm H(u, go). Consistent with the tendency shown in Figure 3 in Section 6.1 Figure 8 also suggests that decreasing & implies that H(,,, go) degrades from CST with the same bin size to DQN, Another interesting observation is that, as shown in Breakout (the first row in Figure [). increasing the number of atoms (reducing the bin size) restricts the range of for a valid return density decomposition in Proposition[] [] Consequently, a small number of atoms or large bin size can allow broader variation of H(N, qo) from C51 to DQN, facilitating the demonstration of our regularization effect empirically.
Breakout (Atoms: 20)
Breakout (Atoms: 80)
Breakout (Atoms: 100)
600
Average Raturn
400
MM
Average Return
500
Average Retum
900
nelle
400
Mit
Mui qutte 0.5)
402
qulte
Viui 6..)) 00.5)
300
Mi Malle 2.8)
2824 dale - 0.5)
202
200
200
102
2000
100
10
15
20
25
30
35
40
10
15
20
25
30
35
40
5
10
15
25
25
30
35
Millions of Frames (M)
Millions of Frames (M)
Millions of Frames (M)
StarGunner (Atoms: 20)
StarGunner (Atoms: 80)
StarGunner (Atoms: 100)
DGn
Average Return
Avarage Return
40000
Average Retum
-0.00
Iapu-uute-222 2:2 MULUAA/--0 +.00
CS1
MELL 4u2tc=
.EA quitc 0.5)
50006
mia quite 0.51
MMAA qulle 0.21
30.000
30000
NEJ mote In 0.8)
PADA quite D.81
20000
20000
10000
19830
1000
10 15
20
25 30
35 40
10 15
20
25 30
35
40
19
Millions of Frames (M)
Millions of Frames (M)
Milions of Fromes (M)
Figure 8: Learning curves of value-based CDRL, i.e., C51 algorithm, and the decomposed algorithm H(M, qo) across different numbers of atoms (various bin sizes) on two Atari games. Results are averaged over three seeds, and the shade represents the standard deviation.
M Discussion on Decomposing Quantile-based Distributional Loss
In order to extend our analysis and conclusion to broader distributional RL algorithm classes, we need to discuss another commonly-use usee algorithm based on quantile regression loss [0]. Although it may be possible to discuss both categorical and quantile representation based on the particle representation (Definition 5.13 in ()) committing either fixed atoms in categorical representation or fixed quantiles in quantile representation can simplify the algorithm analysis. In this section, we discuss how to decompose the quantile-based distributional loss in quantile regression distributional RL.
36

---

Quantile-basee Distributional Loss. In each phase of Neural FZI, we know that the return distribu- tion, parameterized by quantiles, is fixed. This, therefore, leads to a composite quantile loss [72]. which is initially developed to capture the full conditional distribution of the response variable by predicting or estimating its multiple quantiles:
l Equantile
Σ BynPy lPT. (y - Zz:)]s )], N
1
(45)
where pr, is the quantile (pinball) loss defined by Pr. (u)=u(T [((e))) Yu E R. We use Py to denote the fixed target return distribution. In quantile-bass distributional RL, can directly sample y from the quantile function Fy' of the fixed target return given that both the current and target return distributions are parameterized by the quantiles. Zg represents the estimated q--qunttle value of the current return distribution. Alternatively, Pr; can be the quantile Huber loss (22], a smooth version of vanilla quantile loss at zero, by additionally introducing a hyper-parameter K. We thus denote the quantile Huber loss as pf. which is defined as:
L (u)
Pr, (u) |T: 1 (mcc))
(46)
where
{{ 1 (1u) - 1x),
if |u| S k
(u)
(47)
otherwise
As K + 0, it is easy to show that the quantile Huber loss reverts to the vanilla quantile loss. To simplify the notation, we consider the inner-level loss for fixed y:
N Pr. (y )
1
Lquantile
(48)
N
Quantile Representation and Asymptotic Mean-Preserving Property. The normal representation and categorical representation with the categorical projection Ilc in Eq. [0] satisfy the mean- preserving property (Section 4.3 in [4)), as seen in (8)18) in Section 5.4 for normal representation [6] and in Lemma 4.8 for categorical representation in [49]. By contrast, quantile distributional dynamic programming is generally not mean- preserrinn (Lemma 4.8 in [49]), as the quantiles are non-linear functionals of distribution. However, we show that the quantile representation has asymptotic mean-preservve property as the mean of quantiles is asymptotically equivalent to the expectation of the considered distribution when the number of quantiles tends to infinity. Particularly, assume that i 1.N
we have N evenly spaced quantiles { [++ 1i=1- we approximate the expectation by the mean of all quantiles values defined by
N
1
EEr'
)
(49)
N
1=1
Consequently, given a random variable X with its quantile function F-1, we have the following property of quantile function:
1 ( N 1+11== = f.' (F
1 +oo adF = E (X),
lim
(50)
where the first equation results from the relationship between the limit of Riemann Sum and its integral, and the second equation holds by changing the variable 7 = F(x). Note that this asymptotic regime
is similar to that in our histogram function analysis for CDRL, where A + 0 + N + +oo.
According to this equivalence regarding the mean quantiles and the expectation of a random variable, we consider the two decomposition ways as follows.
Decomposition Method 1. We denote 2 + ΣN 2J- as the mean of the quantiles for the current return. Consequently, we have a straightforward composition as follows:
Pr, (y 23 = PT ((V Ž) + (Z -2,-)) = p*, (v 2) + dr.
(51)
where
8, = p.. (1v - 2) + (2 (2-z;)) p:. 2).
(52)
37

---

Therefore, we have the decomposed composite quantile loss as
N 8,
Lquantile
2)
(53)
1=1
Mcan-Related Term
Residual Term
The first term is a mean-related one, which we will elaborate on later, while the induced 8, in the residual term is aimed at capturing the distribution information beyond only the expectation. Particularly, minimizing P* ((v - 2) + (Z 25 )) in 6, will push the deviations Z - Z from
the current return estimator to capture the deviations from the target return distribution of Z. This regularization term contributes preserving the richness of the quantile representation for distributional information, especially the dispersion, from the return.
In terms of the mean-related term, let consider the approximation. As the quantile Huber loss is typically used in quantile-based distributional RL, when x is large, the mean-related term can be simplified
N Σ
N Σ
1
1
1 1(1 -2)²
2)
T5 1 2207 2 2)2 2
(54)
where the first approximation holds because (u) = 1=²² with high probability. The second approximation holds because [T; ([--2<o) [ - 2)² is just the quantile value scaled version of least squared loss. Since Z is the expectation of all estimated quantiles, it can be approximately symmetric to E [Y]. Suppose P(y - < P(y 2 2 0) = 2. we have
T; 1 (v- æ<0) ||
2 (IT: - 11 + 2
(55)
Therefore, this approximation in the mean-related term holds, as shown in Eq. 54 This implies that
the mean quantile estimator ZB- arg minz Ey-Py y - & Σm= 23¹ captures the expectation of
the target return distribution from y ~ Py. Recap the asymptotic equivalence between the expected quantiles and the true expectation of random variable in Eq. 50, the limiting estimator of Zg by minimizing the mean-related term in Equamite satisfies:
E M
Σ
E
Zol
N - +00.
(56)
8
This implics that the learned expected return IE Lg is asymptotically mean-preserving when mini- mizing the mean-related term in the quantile-based distributional loss.
In summary, the first decomposition method decomposes the quantile-base distributional loss into the mean-related and residual terms. After mild approximation, the mean-related term can be simplified as a suus-- loss equipped with an expected quantiles estimator. Combining the equivalence regarding the limiting behavior of the expected quantiles, the mean-related term is thus approximately equivalent to the standard least-squared loss used in classical RI., thus asymptotically satisfying the mean-pre preaan-rreeviigggg property. Moreover, the residual term is able to capture the return distribution information beyond its expectation. In the context of uncertain -ecertai--wware aware regularized exploration in paper, the residual term plays a similar role to the cross-entropy-basee -basdd regularization derived in Proposition 2 of CDRL.
Decomposition Method 2. Another decomposition method can directly follow the return density decomposition proposed in Eq. 3 but we apply the decomposition on the quantile function F- (r) for T E [0, 1]. We expect that this decomposition also leads to two parts, where the first part can involve the quantile defined the bin A, that contains the expected quantiles F-1 ΣE Fr;(T;) and the second term relates to the distribution part. However, this detailed decomposition is largely beyond the scope of this paper, and takes more effort to think about it carefully. We leave this decomposition regarding the quantile- distributional loss as future work.
38


Reviews:
[{"review_id": "MisTBSoMcH", "reviewer": "Reviewer_5szr", "summary": "This paper seeks to deepen the understanding of the benefits of the Categorical Distributional Loss. It does so by decomposing the return distribution estimated in distributional RL into the Dirac delta distribution of the mean and the remaining distribution capturing higher-moments information, and then applying the distributional analogue of the neural fitted Q iteration to analyze the objective that categorical distributional RL is minimizing.\n\nThis decomposition and analysis uncovers a connection between distributional RL, whose objective seeks to reach states where the cross-entropy between the estimation of the distribution at the current state and the distribution at the target state is more, and maximum entropy (MaxEnt) RL, whose objective aims to reach states where the entropy of the policy is maximized.\n\nExperiments in Atari and MuJoCo seek to validate this analysis. First, by showing that the performance of Categorical DQN can be smoothly reduced to DQN by varying how much the algorithm focuses on minimizing the cross-entropy term. Second, by comparing exploration caused by the cross-entropy minimization versus the action entropy maximization.", "questions": "The last two subsections of Section 5.1 indicate that the policy seeks states that have a higher discrepancy between the distribution estimate at that state and the return distribution from the next states. However it is not made clear how the policy is affected over predecessor states where this discrepancy might not have propagated yet, and also how that discrepancy might propagate back. It is not simply a question of comparing to MaxEnt RL since there the entropy regularization is not an error term.", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 3, "originality": 4, "strengths_and_weaknesses": "## Strengths\n* The analysis of the paper seems solid\n* The appendix includes a lot of details and is well constructed.\n* The experiments are chosen well to validate the hypothesis\n\n## Weaknesses\n* The way the idea is communicated can be improved further. The core result of the paper is hard to parse unless the reader is familiar with a lot of the related literature. But given the deep analysis of this paper, perhaps that is the audience they are tailoring it for.\n* The experiments on the Atari domain should be replicated in MuJoCo if possible to showcase that the idea holds in discrete and continuous action spaces as well as with value based and actor critic based algorithms.\n* The results in the bottom row of Figure 4 are mixed, but the paper merely presents a hypothesis for the mixed results instead of validating the hypothesis.\n* Minor point: Evaluations with 3 seeds in Atari and 5 seeds in MuJoCo might not be statistically sound, but consistency across the different tasks in these two domains points to some validity.\n* Minor point, another paper that has used a mixture of distributional RL and MaxEnt RL to great success is [1]. Perhaps that paper validates or disproves some of the hypotheses in this one.\n\n### References\n[1] Outracing champion Gran Turismo drivers with deep reinforcement learning., Wurman et al., Nature, 2022", "was_revised": true, "final_justification": "Authors addressed my minor concerns through rebuttal. I do not feel a need to update my rating", "rebuttal": "We thank the reviewer for the insightful feedback and positive assessment of our work, especially on solid analysis, detailed appendix, and well chosen experiments. We would like to address the concerns you raised in your review.\n\n\n> **Q1:[Weakness 1]** The way the idea is communicated can be improved further. The core result of the paper is hard to parse unless the reader is familiar with a lot of the related literature. But given the deep analysis of this paper, perhaps that is the audience they are tailoring it for.\n\n**A1:** Thank you for the helpful feedback. In future revisions, we will improve the clarity and presentation of the core ideas to make them more accessible to a broader audience while still retaining the depth of our analysis.\n\n>  **Q2:[Weakness 2]** The experiments on the Atari domain should be replicated in MuJoCo if possible to showcase that the idea holds in discrete and continuous action spaces as well as with value based and actor critic based algorithms.\n\n**A2:** We have conducted additional experiments in MuJoCo environments, and similar results can be found in **Appendix M.1**.  This empirical evidence consistently demonstrates the key effect of the decomposed regularization term from the categorical distribution loss on the superiority of and distributional RL over classical RL.\n\n\n>  **Q3:[Weakness 3]** The results in the bottom row of Figure 4 are mixed, but the paper merely presents a hypothesis for the mixed results instead of validating the hypothesis.\n\n**A3:** Yes. Results in **the first row of Figure 4** show that simultaneously employing uncertainty-aware and vanilla entropy regularization renders a mutual improvement. By contrast, the two kinds of regularization can also lead to potential interference in policy optimization through the results in **the second row of Figure 4**, revealing an interesting phenomenon to be explored in the future. \n\n\n\n>  **Q4:[Weakness 4]** Minor point: Evaluations with 3 seeds in Atari and 5 seeds in MuJoCo might not be statistically sound, but consistency across the different tasks in these two domains points to some validity.\n\n**A4:** Thank you for this suggestion, and we will run more seeds in the final version of the paper.\n\n\n>  **Q5:[Weakness 5]** Minor point, another paper that has used a mixture of distributional RL and MaxEnt RL to great success is [1]. Perhaps that paper validates or disproves some of the hypotheses in this one.\n\n**A5:** Thanks for sharing this reference paper, which indeed validates that combining distribution learning and MaxEnt RL can lead to mutual improvement in some scenarios.\n\n\n> **Q6:[Questions]** The last two subsections of Section 5.1 indicate that the policy seeks states that have a higher discrepancy between the distribution estimate at that state and the return distribution from the next states. However, it is not made clear how the policy is affected over predecessor states where this discrepancy might not have propagated yet, and also how that discrepancy might propagate back. It is not simply a question of comparing to MaxEnt RL since there the entropy regularization is not an error term. \n\n**A6:** Thank you for the insightful question. It is correct that the policy is encouraged to explore states with high discrepancy between the current estimated return distribution and the target return distribution. Since the policy operates in an MDP, this exploratory behavior is primarily driven by the current state rather than directly by predecessor states.\n\nHowever, we agree that the influence of discrepancies in predecessor states can indirectly propagate backward through the Bellman updates during training. If the discrepancy in the predecessor states may not propagate yet, the policy may not be immediately simulated to visit those earlier states, which could limit the full exploitation of distributional uncertainty in long trajectories. This subtle effect on the trajectory-based exploration is beyond the scope of this paper, but it is interesting to have a deeper theoretical and empirical investigation. We will consider this in future work. \n\nThank you again for pointing out these potential areas of improvement. We appreciate your suggestions. Please let us know if you have any further comments or feedback."}, {"review_id": "U195qNBvfs", "reviewer": "Reviewer_oCWL", "summary": "The paper attempts to shed light on the improved performance seen in deep RL when using distributional RL, and in particular categorical distributional losses, by decomposing the loss into a mean-related and entropy regularization term. The argument is that this (distribution-matching) entropy regularization term implicitly, but also explicitly in this paper, encourages the policy to visit states/actions where the distribution is poorly estimated.\n\nThe primary contributions include this decomposition, theoretical results supporting this decomposition and the correctness/convergence of an explicitly regularized policy evaluation method, empirical evaluation of the explicit form compared with C51 and DQN, and a study in MuJoCo environments on the combination of this form of regularization with the standard form of entropy regularization.", "questions": "* On line 239 the text says “we use uppercase notation…” and then proceeds to give an example “(s_t, a_t)” which is lowercase. Uppercase for random variables is ideal, but this juxtaposition was confusing and I thought maybe a typo.\n* There are other methods that use categorical losses without attempting to estimate the distribution of returns (at least not fully) such as so called TwoHot encodings and HL-Gauss. Could your approach be adapted to analyze these methods as well? Or in these cases do you think that there would not be a related regularization effect?\n* Can you say more about why this form of regularization *should* be expected to be useful, or when it might be counter productive?", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths:\n\n* There is actual novel insight here with rigorous theory and informative experiments to back it up.\n* The experimental work asks exactly the right set of questions (for the scope of a conference paper). In Figure 3 we quite clearly see how this explicit form the regularization leads to performance which largely interpolates between C51 and DQN as the strength of regularization is varied. Figure 4 studies a very natural and interesting question of combining with policy entropy regularization, although I feel that this result did not get quite enough discussion for it to fully land.\n* The paper is generally well written and the flow of the paper introduces us to this idea very effectively.\n\nWeaknesses:\n\n* The notation is a little rough. It is precise, but there is a fair bit to keep track of and I found it a bit less intuitive than it could be. In particular, I keep wanting to treat mu as the mean-related term and forgetting it is the opposite. Lots of super- and sub-scripts, there’s a lot to keep track of.\n* An intuitive explanation of the effects of this form of regularization could be improved. The phrase “lags far behind” seemed like it kept getting used, but I really wanted more investigation and discussion of how to understand the effects here. Overall, I came away from reading this able to parrot that concept but lacking a more complete intuition for what that would really mean in terms of affecting behavior.", "was_revised": true, "final_justification": "Still find this to be strong work worthy of acceptance and author rebuttal did not lower my evaluation. I would be willing to argue in favor of this work.", "rebuttal": "We thank the reviewer for the insightful feedback and positive assessment of our work, especially on the rigorous theory and experiments to back up our novel insight and effective flow of our paper. We would like to address the concerns you raised in your review.\n\n> **Q1:** The notation is a little rough. It is precise, but there is a fair bit to keep track of and I found it a bit less intuitive than it could be. In particular, I keep wanting to treat mu as the mean-related term and forgetting it is theo pposite. Lots of super- and sub-scripts, there’s a lot to keep track of.\n\n**A1:** Thank you for pointing out this notation issue. We will simplify notations and improve clarity to enhance the readability in the final version of the paper.\n\n> **Q2:** An intuitive explanation of the effects of this form of regularization could be improved. The phrase “lags far behind” seemed like it kept getting used, but I really wanted more investigation and discussion of how to understand the effects here. Overall, I came away from reading this able to parrot that concept but lacking a more complete intuition for what that would really mean in terms of affecting behavior\n\n**A2:** When revisiting the objective function in **Eq. 8 of Lemma 1**, optimizing the policy implies an increase of the regularization term $f(\\mathcal{H} ( \\mu^{\\mathbf{s}\\_t, \\mathbf{a}\\_t},  q_{\\theta}^{\\mathbf{s}\\_t, \\mathbf{a}\\_t} ) )$. Given a fixed $q_\\theta^{s,a}$, the increasing of $f(\\mathcal{H} ( \\mu^{\\mathbf{s}\\_t, \\mathbf{a}\\_t},  q_{\\theta}^{\\mathbf{s}\\_t, \\mathbf{a}\\_t} ))$ indicates that the future trajectory with the states and actions is expected to have a large distribution discrepancy. Given the current state, the policy is optimized to explore states and actions whose current return distribution $q_\\theta^{s, a}$ is dramatically different from the estimated target return distribution $\\mu^{s, a}$. Intuitively, the exploration is promoted to diminish the agent's current understanding of the environmental uncertainty (through the return distribution) and the (estimated) real environmental uncertainty. We will add a more detailed explanation in the revised version.\n\n\n\n> **Q3:** On line 239 the text says “we use uppercase notation…” and then proceeds to give an example “(s_t, a_t)” which is lowercase. Uppercase for random variables is ideal, but this juxtaposition was confusing and I thought maybe a typo.\n\n**A3:** Thank you for catching this typo. The notation should use boldface to represent random variables, not uppercase letters. This follows the convention used in much of the RL literature, such as in Soft Actor-Critic. We will correct the text for clarity in the final version.\n\n\n> **Q4:** There are other methods that use categorical losses without attempting to estimate the distribution of returns (at least not fully) such as so called TwoHot encodings and HL-Gauss. Could your approach be adapted to analyze these methods as well? Or in these cases do you think that there would not be a related regularization effect?\n\n**A4:** This is a great question. It is indeed possible to extend our analysis to methods like HL-Gauss [1, 2], which also utilize histogram-like representations, albeit under a Gaussian assumption. Our histogram-based return distribution decomposition can be potentially applied to HL-Gauss and we expect that a similar regularization effect may emerge from their simplified formulations. We consider this an interesting direction for future work.\n\n\n> **Q5:** Can you say more about why this form of regularization should be expected to be useful, or when it might be counter productive?\n\n**A5:** This regularization promotes exploration by encouraging the agent to reduce the environmental uncertainty between the estimated and target return distributions by assuming that **the target distribution is well-estimated**. However, if the TD target is poorly estimated (e.g., due to limited distributional representation or the limitations of one-step TD as opposed to multi-step TD), the policy may be misled toward uninformative regions of the state space. In such cases, the distributional loss can be counterproductive, and classical RL may perform better.\n\n\nThank you again for pointing out these potential areas of improvement. We appreciate your suggestions. Please let us know if you have any further comments or feedback.\n\n### Reference \n\n[1] Stop Regressing: Training Value Functions via Classification for Scalable Deep RL   \n[2] Improving regression performance with distributional losses. (ICML 2018)"}, {"review_id": "FY8nqLawLD", "reviewer": "Reviewer_WN6r", "summary": "This work revisits the benefits of distributional RL, focusing on categorical distributional RL (CDRL), which uses categorical losses. In Section 4.2, they show that CDRL uses additional information by learning the entire return distribution, compared to standard RL. They also show that the “mean term” in the distributional objective asymptotically behaves like the standard RL objective.\n\nThe paper further presents a convergence analysis for distribution-entropy-regularized policy iteration (DERPI), showing that it enjoys similar guarantees as both standard RL and MaxEnt RL. However, they emphasize a key difference: unlike MaxEnt RL, where exploration is explicitly encouraged, DERPI promotes exploration implicitly through the additional reward signal coming from modelling the return distribution, and encouraging the policy to pursue a \"target\".\n\nThe theoretical insights are supported by experiments on both Atari and MuJoCo environments. Also, this work shows that combining both types of regularized exploration (distributional and MatEnt) can sometimes reach better performances.", "questions": "1. There are additional parameters involved, e.g. the parameter $\\\\beta$ in MaxEnt RL. I may have overlooked it, but how is this parameter set in your experiments?\n\n2. Also, when combining the two types of exploration in DSAC, how do you choose the corresponding parameters to balance their contributions? From Figure 4, it seems this combination introduces larger fluctuations in the learning curves. Do you have any insights on whether there are ways to stabilize it?", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "## Strengths\n\nThis work revisits the benefits of a specific form of distributional RL, i.e., CDRL, and provides a rigorous comparison to standard RL. It clearly shows, by formula, the additional informational used by CDRL. Furthermore, due to the implicit exploration behavior in distributional RL, the paper also theoretical compares to MaxEnt RL. This work makes such comparison clear and also quite elegant.\n\nOverall speaking, this work provides nice insights to the relationships among standard RL, distributional RL, and MaxEnt RL. \n\nAlso, the presentation is smooth and well-structured, with the logical flow of ideas.\n\n## Weaknesses\n\nI don’t have much to complain about. However, \n\n1. The analysis is limited to CDRL, rather than general distributional RL frameworks. While this is explicitly acknowledged in the paper, the readers should be careful when intending to generalize the conclusions beyond CDRL.\n\n2. When using UE, the method appears to introduce some instability during training. It could be worth discussing this behavior in more detail.", "was_revised": true, "final_justification": "Most reviewers share similar opinions to the work as I do, and the authors have addressed most of my questions.", "rebuttal": "We thank the reviewer for the insightful feedback and positive assessment of our work, especially the clear formula to show the additional information about CDRL and the theoretical comparison with MaxEnt RL. We would like to address the concerns you raised in your review.\n\n\n> **Q1:** The analysis is limited to CDRL, rather than general distributional RL frameworks. While this is explicitly acknowledged in the paper, the readers should be careful when intending to generalize the conclusions beyond CDRL.\n\n**A1:** The main analysis is intended for CDRL, but we hope such a similar conclusion also applies to general distributional RL. An extension to quantile distributional loss is also given in Section 7.\n\n> **Q2:** When using UE, the method appears to introduce some instability during training. It could be worth discussing this behavior in more detail.\n\n**A2:** Thank you for the observation. The instability—such as in the HumanoidStandup task in Figure 4—may result from divergence under certain random seeds. We plan to run additional seeds to better quantify this effect. However, in most environments, introducing UE in algorithms remains stable and does not introduce significant training instability.\n\n\n> **Q3:** There are additional parameters involved, e.g. the parameter $\\beta$ in MaxEnt RL. I may have overlooked it, but how is this parameter set in your experiments?\n\n**A3:** Following the original SAC paper, we directly set $\\beta=1$ across all experiments.\n\n\n> **Q4:** Also, when combining the two types of exploration in DSAC, how do you choose the corresponding parameters to balance their contributions? From Figure 4, it seems this combination introduces larger fluctuations in the learning curves. Do you have any insights on whether there are ways to stabilize it?\n\n**A4:** We choose the default hyperparameters of distributional RL and MaxEnt RL without tuning them. A key conclusion we make in Section 6.2 is that the combination of the two kinds of entropy can have different mutual impacts, either a mutual improvement (the first row in Figure 4) or potential intervention (the second row in Figure 4). We posit that the potential interference may result from distinct exploration directions in the policy learning for the two types of regularizations, thus potentially leading to a large fluctuation. It is hard to predetermine which exploration is more useful given an environment, and we thus leave it as future work. \n\n\nThank you again for pointing out these potential areas of improvement. We appreciate your suggestions. Please let us know if you have any further comments or feedback."}, {"review_id": "SkEgQoNRVg", "reviewer": "Reviewer_m2h8", "summary": "The paper aims to explain the benefits of (categorical) distributional RL by relating it to entropy regularization, which encourages exploration. Theoretical results are widely provided, and experiments follow.", "questions": "I listed my questions above. My understanding might be inaccurate so please correct me if I was wrong.", "limitations": "The authors discussed the limitations at the end of the paper. I don't see any potential negative societal impact.", "rating": 2, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 1, "clarity": 1, "significance": 2, "originality": 2, "strengths_and_weaknesses": "This paper presents an interesting perspective on distributional RL by drawing connections to entropy regularization (and thus maxent RL). Both theoretical and experimental components are provided. However, I found the theoretical section somewhat difficult to follow. Some terminology is unclear. For instance, the \"Return Density Decomposition\" paragraph uses heavy notation and probably need clearer explanations. Terms like $p^\\mu_i$, which is the coefficient of the i-th bin, it needs to be strictly defined. The use of the word \"induced\" in line 161, although it is bold, is kind of clear.\n\nBased on my current understanding, the theoretical contributions may be limited in general. Regarding Proposition 1, I may be misunderstanding it, but it appears to be narrow and not referred to in other theoretical results. Proposition 2 seems fairly standard, as it directly applies the KL divergence to histogram representations. Proposition 3 also appears to be straightforward: when $\\Delta$ is approaching 0, the function should reduce to the desired target.\n\nOn the empirical side, the paper is verifying that the \"uncertainty-aware regularization is crucial to the outperformance of CDRL over classical RL...\" (line 335) However I found it difficult to draw this conclusion from the results. According to Figure 3, runs with verifying $\\epsilon$ is worse than C51 while being better than DQN. This could suggest that the gains might come from the general benefits of distributional RL rather than the proposed uncertainty-aware regularization itself. Additional analysis or ablation might help clarify the specific contribution of the regularization component.", "was_revised": true, "final_justification": "i am keeping my score since my concern about the clarity persists after rebuttal.", "rebuttal": "We would like to sincerely express our gratitude to the reviewer for their time and effort in reviewing our paper. We appreciate the reviewer's positive recognition of our work on an interesting perspective on distributional RL, the connection with entropy-regularized RL, and both theoretical and empirical demonstrations. We address each concern you raised and are happy to answer any further questions.\n\n> **Q1:** I found the theoretical section somewhat difficult to follow. Some terminology is unclear. For instance, the \"Return Density Decomposition\"paragraph uses heavy notation and probably need clearer explanations. Terms like $p_i^\\mu$ , which is the coefficient of the i-th bin, it needs to be strictly defined. The use of the word \"induced\" in line 161, although it is bold, is kind of clear.\n\n**A1:** When it comes to a distribution decomposition, it requires rigorous descriptions and definitions about which random variable of this distribution is, how the support is partitioned, and what the coefficients represent in histogram functions. Although we acknowledge this may increase the cognitive load, we choose this way for the presentation to ensure rigor. For instance, for the induced density function $\\hat{\\mu}^{s, a}(x)=\\sum_{i=1}^N p_i^\\mu \\mathbb{1}\\left(x \\in \\Delta_i\\right) / \\Delta$, it rigorously represents a histogram function with each coefficient of this histogram defined on each bin $\\Delta_i$. This return distribution decomposition is also illustrated in **Figure 1**, where the coefficient $p_i^\\mu$ represents the height on the histogram of the bin $\\Delta_i$. **``Induced''** is used intentionally to indicate that given the hyper-parameter $\\epsilon$ (under the condition of **Proposition 1**) and a histogram density function $\\hat{p}^{s, a}$, $\\mu^{s, a}$ results from decomposing the original density function $\\hat{p}^{s,a}$, rather than being predefined.\n\nWe will revise the text to reduce unnecessary notation and improve explanations of key terms for better clarity and accessibility.\n\n\n> **Q2:** Based on my current understanding, the theoretical contributions may be limited in general. Regarding Proposition 1, I may be misunderstanding it, but it appears to be narrow and not referred to in other theoretical results.\n\n\n**A2:** Thank you for the comment. Firstly, the return distribution decomposition is both novel and non-trivial. It is not generally true but rigorously holds when $\\epsilon > 1 -p_E$ as shown in Proposition 1. Proposition 1 serves as the theoretical foundation for the following regularization-based analysis, as the return distribution decomposition will be invalid if the conditions in Proposition 1 are not satisfied. Although Proposition 1 is not directly referenced in subsequent theoretical results, it underpins our analysis with a solid foundation.\n\n\n> **Q3:** Proposition 2 seems fairly standard, as it directly applies the KL divergence to histogram representations.\n\n**A3:** We need to highlight that the tricky part in deriving Proposition 2 is introducing the **return distribution decomposition** to represent the target distribution when we apply the KL divergence, with the detailed proof in **Eq. 28 in Appendix H**. It is also crucial to guarantee a fixed target distribution (ensured by the Neural FZI framework) such that the KL divergence is equivalent to the cross entropy $\\mathcal{H}(p, q)$ as shown in **Eq. 27 in Appendix H**. Putting them together, we eventually derive the decomposed Neural FZI objective function in Proposition 2. Importantly, this decomposed Neural FZI objective function is novel, as we first introduce the return distribution decomposition technique that we propose initially in this paper. The derivation avoids complicated assumptions but provides a general and principled formulation when the conditions in Proposition 1 hold.\n\n\n> **Q4:** Proposition 3 also appears to be straightforward: when $\\Delta$ is approaching 0, the function should reduce to the desired target.\n\n\n**A4:** Given the novel decomposed objective function of Neural FZI in Proposition 2, it is both novel and non-trivial to analyze the relationship between the minimizer of its first term (i.e., the mean-related component) and the minimizer of the classical RL objective function (i.e., that of Neural FQI). This motivates Proposition 3, whose result was previously unknown before our formulation.\n\nThe derivation of Proposition 3 involves a rigorous asymptotic analysis of the Dirac delta function (as a generalized probability density function), an equivalence analysis between distinct minimizers, and properties of the distributional Bellman optimality operator. We emphasize that this derivation is not straightforward, nor can it be directly inferred from existing literature. The conclusion in Proposition 3 is intuitive in hindsight, but this should not be mistaken for simplicity of derivation—it is the result of careful and non-trivial theoretical analysis. \n\n\n\n> **Q5:** On the empirical side, the paper is verifying that the \"uncertainty-aware regularization is crucial to the outperformance of CDRL over classical RL...\" (line 335) However, I found it difficult to draw this conclusion from the results. According to Figure 3, runs with verifying $\\epsilon$ is worse than C51 while being better than DQN. This could suggest that the gains might come from the general benefits of distributional RL rather than the proposed uncertainty-aware regularization itself. Additional analysis or ablation might help clarify the specific contribution of the regularization component.\n\n\n**A5:** We would like to emphasize that the proposed uncertainty-aware regularization is directly derived from the decomposition of the distributional loss function used in CDRL. As shown in **Figure 3**, decreasing $\\varepsilon$ from 0.8 to 0.1 reduces the contribution of the distribution component in the baseline algorithm $\\mathcal{H}(\\mu, q_\\theta)$, leading to a learned target $\\hat{\\mu}^{s,a}$ that captures less high-order moments knowledge of the original return distribution $\\hat{p}^{s,a}$ (as defined in Eq. 1). This reduction results in notable performance degradation—from C51 (a distributional RL method) to DQN (a classical RL method)—demonstrating the crucial role of modeling the full return distribution, beyond just its expectation, in enhancing learning performance.\n\nThis observed performance gain is rigorously captured by the regularization term in **Eq. 3 of Proposition 2**, which we interpret as uncertainty-aware exploration grounded in the theoretical insights established by our propositions. Accordingly, the ablation study in **Figure 3** precisely illustrates that the decomposed term from the categorical distributional loss plays a pivotal role in explaining the performance gap between classical and distributional RL methods. We welcome any follow-up questions for clarification on this explanation.\n\nThanks again for the reviewer's consistent dedication to reviewing our work, and we are happy to answer any further questions or concerns. Should this rebuttal address your concerns, we would be grateful for an increased score."}]

Author rebuttal:

Meta-review: This paper is divisive among reviewers. Everyone appreciates the interesting insightful theoretical perspective on an old empirical observation that mystifies researchers for almost a decade now. However, the paper has two critical flaws: 1. the analysis is complex, hard to understand without extensive prior knowledge and uses overly complex notation; and 2. the empirical results show the effect in some experiments, but remain ambiguities in others.

A clear majority of reviewers, who state they are also more confident in their assessment, accepts the paper is for a more niece audience that is fluent enough in literature and willing to disentangle the complex notation. Only m2h8 maintains that the paper is too hard to read and the interpretation of the results questionable. This could be easily dismissed (m2h8 does so as well) as an individual opinion, but I need to admit that I somewhat agree. Complexity is relative, but the interpretation of the results is not as clean cut as the authors make it seem. For example, Figure 8 in the appendix ablates the influence of the proposed regularizer for Mujoco environments. However, here the results do not "interpolate" as they do for Atari environments in Figure 3. These kinds of ambiguities are common in deep RL research, but they do cast doubt on some of the empirical conclusions.

I nonetheless recommend to accept this paper. The theory, while hard to digest, sheds light on an old unresolved question, and the results provide enough evidence to make the authors' interpretation reasonable, if not unquestionable.
Decision: Accept (poster)