How Does Value Distribution in Distributional Reinforcement Learning Help Optimization?
Authors: Ke Sun
Bei Jiang
Linglong Kong
Venue: Submitted to ICLR 2023 (2023)
Keywords: distributional reinforcement learning
optimization

Abstract: We consider the problem of learning a set of probability distributions from the Bellman dynamics in distributional reinforcement learning~(RL) that learns the whole return distribution compared with only its expectation in classical RL. Despite its success to obtain superior performance, we still have a poor understanding of how the value distribution in distributional RL works. In this study, we analyze the optimization benefits of distributional RL by leverage of additional value distribution information over classical RL in the Neural Fitted Z-Iteration~(Neural FZI) framework. To begin with, we demonstrate that the distribution loss of distributional RL has desirable smoothness characteristics and hence enjoys stable gradients, which is in line with its tendency to promote optimization stability. Furthermore, the acceleration effect of distributional RL is revealed by decomposing the return distribution. It turns out that distributional RL can perform favorably if the value distribution approximation is appropriate, measured by the variance of gradient estimates in each environment for any specific distributional RL algorithm. Rigorous experiments validate the stable optimization behaviors of distributional RL, contributing to its acceleration effects compared to classical RL. The findings of our research illuminate how the value distribution in distributional RL algorithms helps the optimization.

Under review as a conference paper at ICLR 2023
How DOES VALUE DISTRIBUTION IN DISTRIBU- TIONAL REINFORCEMENT LEARNING HELP OPTI- MIZATION?
Anonymous authors Paper under double-blind review
ABSTRACT
We consider the problem of learning set of probability distributions from the Bellman dynamics in distributional reinforcement learning (RL) that learns the whole return distribution compared with only its expectation in classical RL. De- spite its success to obtain superior performance, we still have poor understanding of how the value distribution in distributional RL works. In this study, analyze the optimization benefits of distributional RL by leveraging its additional value distribution information over classical RL in the Neural Fitted 7-Iteration (Neu- ral FZI) framework. To begin with, we demonstrate that the distribution loss distributional RL has desirable smoothness characteristics and hence enjoys stable gradients, which is in line with its tendency to promote optimization stability. Fur- thermore, the acceleration effect of distributional RL is revealed by decomposing the return distribution. It turns out that distributional RL can perform favorably if the value distribution approximation appropriate, measured by the variance of gradient estimates in each environment for any specific distributional RI. algo- rithm. Rigorous experiments validate the stable optimization behaviors of distri- butional RL, contributing its acceleration effects compared to classical RL. The findings our research illuminate how the value distribution in distributional RL algorithms helps the optimization.
INTRODUCTION
Distributional reinforcement learning (Bellemare et al.) 2017a Dabney al. 2018h a Yang al. 2019 Zhou et al.) 2020; Nguyen et al. 20200] Luo et al. 2021 Sun et al.) (022) characterizes the intrinsic randomness of returns within the framework
Reinforcement Learning (RL). When the agent interacts with the environment, the intrinsic uncertainty of the environment seeps in the the stochasticity of rewards the agent receives and the inherently chaotic state and action dynamics of physical interaction, increasing the difficulty of the RL algorithm design. Distributional RL is aimed at representing the entire distribution of returns in order to capture more intrinsic uncertainty of the environment, and therefore in stark contrast to the classical RI. that only focuses on the expectation the return distribu- tions, such as temporal-diiferenc (TD) learning (Sutton Barto 2018 and Q-learning (Watkins & Dayan 1992).
these value distributions evaluate and optimize the policy. This
As a promising branch of RL algorithms, distributional RL has demonstrated the state-of-the-ar performance in a wide range of environments, e.g., Atari games, in which the representation of re- turn distributions and the distribution divergence between the current and target return distributions within each Bellman update are pivotal to its empirical success (Dabney et al.] 201 Sa, (Sun et al.) 2021b 2022). Specifically, categorical distributional RI., e.g., CS1 (Bellemare al.) [01 Row- land et al. 2018). integrates a categorical distribution by approximating the density probabilities in pre-specified bins with bounded range and Kullback-Leible (KL) divergence, serving as the first successful distributional RL family in recent years. Quantile Regression (QR) distributional RL, e.g., QR-DQN (Dabney et al. 22188b) approximates Wasserstein distance by the quantile regression loss and leverages quantiles represent the whole return distribution. Other variants of QR-DQN, including Implicit Quantile Networks (IQN) (Dbneyy et al. (0188) and Fully parameterized Quan- achieve significantly better performance across
tile Function (FQF) (Yang al. 2019),

---

Under review as a conference paper at ICLR 2023
plenty of Atari games. Moment Matching distributional RL Nguuyen et al. 2020) learns determin- istic samples to evaluate the distribution distance based on Maximum Mean Discrepancy, while a more recent work called Sinkhorn distributional RL (Sun et al., 2022) interpolates Maximum Mean Discrepancy and Wasserstein distance via Sinkhorn divergence (Sinkhorn 1967). Meanwhile, dis- tributional RI also inherits other benefits in risk-sensitive control (Dabney et al. [0188a), policy exploration settings (Mavrin et al., 2019; Rowland et al. 2019) and robustness (Sun et al. 2021a).
Despite the remarkable empirical success of distributional RL, the illumination on its theoretical advantages is still less studied. A distributional regularization effect (Sun et al. 2021b) stemming from the additional value distribution knowledge has been characterized to explain the superiority of distributional RL over classical RL, but the benefit of the proposed regularization on the optimization of algorithms has not been investigated as the optimization plays a key role RI. algorithms. In the literature of strategies that can help the learning in RL, recent progresses mainly focus on the policy gradient methods (Sutton & Barto 2018). Mei et al.] (2020) show that the policy gradient with softmax parameterization converges at a (11/)) rate, with constants depending on the problem and initialization, which significantly expands the existing asymptotic convergence results. Entropy regularization (Haarnoja al.) 2017: 2018) has gained increasing attention it can significantly speed up the policy optimization with a faster linear convergence rate (Mei et al.) (020). [Ahmed] al. (2019) provide fine-grained understanding on the impact of entropy on policy optimization, and emphasize that any strategy, such as entropy regularization, can only affect learning in one of two ways: either it reduces the noise in the gradient estimates or it changes the optimization landscape. These commonly optimization impact of distributional RL arising from the exploitation of return distributions.
strategies that accelerate RL learning inspire us to further investigate the
In this paper, we study the theoretical superiority of distributional RL over classical RL from the op- timization standpoint. We begin by analyzing the optimization impact of different strategies within the Neural Fitted Z-Iteration (Neural FZI) framework and point out two crucial factors that contribute to the optimization of distributional RI., including the distribution divergence and the distribution parameterization error. The smoothness property of distributional RI. loss function has also been revealed leveraging the categorical parameterization, yielding its stable optimization behavior. The uniform stability in the optimization process can thus be more easily achieved for distributional RL in contrast to classical RL. In addition to the optimization stability. we also elaborate the accelera- tion effect of distributional RL algorithms based the value distribution decomposition technique proposed recently. It turns out that distributional RL can be shown to speed up the convergence and perform favorably if the value distribution is approximated appropriately, which is measured by the variance of gradient estimates. Empirical results corroborate that distributional RL indeed enjoys stable gradient behavior by observing smaller gradient norms terms of the observations the agent encounters in the learning process. Besides, the variance reduction of gradient estimates for distributional RL algorithms with respect to network parameters also provides strong evidence to demonstrate the smoothness property and acceleration effects of distributional RL. Our study opens up many exciting research pathways in this domain through the lens of optimization, paving the way for future investigations to reveal more advantages of distributional RL.
PRELIMINARY KNOWLEDGE
Classical RL. In standard RL setting, the interaction between an agent and the environment is modeled as a Markov Decision Process (MDP) (S.AA, R, P, r), where S and A denote state and action spaces.
the transition kernel dynamics, discount factor. For a fixed policy π, the return, -T - >² RR, is random variable representing the sum of discounted rewards observed along one trajectory of states while following the policy π. Classical RL focuses on the value function and action-value function, the expectation of returns 2". The action-value function Q (s. a) is defined as Q= (s. a) = E [Z (s, a)] = EE 7 R R(t, a,), where su- s, ag a, S1+1 P(- St, a.). and ag (-1ss)).
the reward measure and 7 E (0, 1) is the
Distributional RL. Distributional RL, on the other hand, focuses on the action-value distribution, the full distribution of (s,ss, a) rather than only its expectation, i.e., Q= a). Leveraging knowl- edge on the entire value distribution can better capture the uncertainty of returns and thus can be advantageous to explore the intrinsic uncertainty of the environment (Dabney et al.] [2018a] MMvrin ee al. 2019). The scalar-based classical Bellman updated is therefore extended to distributional
2

---

Under review as a conference paper at ICLR 2023
Bellman update, which allows fluryy of distributional RL algorithms, mainly including Categori- cal distributional RL (Bellemare et al. (0177) and Quantie Regression Distributional RL. (Dabney et al.
a).
Categorical Distributional RL. As the first successful distributional RL family, Categorical dis- tributional RL (Bellemare a 2017a) approximates the action-valle distribution r) by a categorical distribution n Σi=1 fiór, where 11,12,-- 12......4 44 is a set of fixed supports and {ft] are learnable probabilities, normally parameterized by a neural network. A projection is also introduced to have the joint support with newly distributed target probabilities, equipped by a KI. divergence to com- pute the distribution distance between the current and target value distribution within each Bellman update. In practice, C51 (Bellemare et al. 2017a), an instance of Categorical Distributional RL with k = 51, performs favorably wide range of environments.
Quantile Regression (QR) Distributional RL. QR Distributional RL (Dabney et al. 201 8b a) ap- proximates the value distribution by a mixture of Dirac f N Σi drg where T; P 22_1 are the learnable quantile values at the fixed quantiles 2i-11 J and F=1 is the inverse cumulative distribution function of Since the quantile regression loss proposed in QR distributional RL can be used to approximate the Wasserstein distance, it gains favorable performance on Atari games. Moreover, the performance has been further improved by series of variants based on quantile re- gression loss (Dabney al., (0,,2018: Sa| Yang al.) 2019 Zhou al., 2020). For example, Implicit Quantile Network (IQN) (Dabney et al. (2018a) Sa) utilizes a continuous mapping for the quantile func- tion F (2/=1) rather than a fixed set of quantiles, which expands the expressiveness power of function approximators to represent the value distribution.
2N
3 OPTIMIZATION EFFECT OF DISTRIBUTIONAL RL
We consider the function approximation setting to analyze the optimization benefit of distributional RL. In Section we begin by showing the different roles of components in distributional RL on the entire optimization of RL algorithms within the Neural FZI framework. Further, in Section 3.2 we reveal the desirable smoothness properties of distributional RL loss function opposed classical RL, contributing to the stable optimization. Finally, the acceleration effect of distributional RL stemming from the additional value distribution is analyzed in Section 3.3 which is characterized by the variance of gradient estimates.
3.1 How TO OPTIMIZE NEURAL FITTED Z-ITERATION FOR DISTRIBUTIONAL RL?
In classical RL, Neural Fitted Q Itteration (Neural FQI) (Fan al. 2020 Riedmiller 2005 provides
a statistical interpretation of DQN (Mnih et al.) 2015) while capturing its two key features, i.e., the leverage target network and the experience replay:
Qt+1 arggminn Qo Ti Σ
1
[v Q3 (si, a:))²
where the target y: rr(s; 0:) + maxae Q/- (s),a) is fixed within every Turget steps to up- date target network Q by letting {(8,, di, &)), + heen) and ideally without the optimization and TD approximation errors, Neural FQI is exactly the update under Bellman optimality operator (Fan et al.) 2020). Similarly, in distribu- tional RL, Sun et al. (2021b); Ma et al. (2021) proposed Neural Fitted Z-Iteration (Neural FZI), a distributional version of Neural FQI based the parameterization of Zo:
A. The experience buffer induces independent samples
1 Σ
2 argmin Za
doi, 23 (si,
(2)
where the target Y; = R(s,, a.) + 7Zz. (5/, (s())) is a random variable, whose distribution is also fixed within every T Tiage steps. The target follows greedy policy rule, where T2(3() argmax, E [zp- (s(,a') and dp is the choice of distribution distance. Within the Neural FZI process, we can easily perceive that there are mainly two crucial components that determine the comprehensive optimization of distributional RL algorithms.
3

---

Under review as a conference paper at ICLR 2023
The choice of dz. dp in fact has two-fold impacts on the optimization of the whole Neu- ral FZI. Firstly, dp determines the convergence rate of distributional Bellman update. For instance, distributional Bellman operator under Crámer distance is " (Belle- mare et al.] 2017b), and ?-contraction when dp is Wasserstein distance (Bellemare et al. (2)77). Apart from the impact on the distributional Bellman update speed, d, also largely affects the continuous optimization problem to estimate parameter in Zg within each iteration of Neural FZI, including the convergence speed and the bad or good local minima issues.
The parameterization manner of Ze. The distribution representation way of dp plays an integral part of the optimization for deep RL algorithms. For example, with more ex- pressiveness power on quantile functions, IQN outperforms QR-DQN on a wider range of Atari games, which is intuitive as a more informative representation way can approximate the true value distribution more reasonably. A smaller value distribution parameteriation error is also potential to help the optimization albeit in an indirect avenue.
Owing to the fact that convergence rates of distributional Bellman update under typical dp are basi- cally known, our optimization analysis mainly focuses on the impact of d, and the paramterization error of Ze on the continuous optimization within Neural FZI of distributional RL by comparing Neural FQI of classical RL. In Sections 3.2 and 3.3 we attribute the optimization benefits of distri- butional RL to its distribution objective function, consisting of the aforementioned two factors, as opposed to the vanilla least squared loss in classical RL.
3.2 STABLE OPTIMIZATION ANALYSIS BASED ON CATEGORICAL PARAMETERIZATION
To allow for a theoretical analysis, we resort to the categorical parameterization equipped with KL. divergence in categorical distributional RL (Bellemare al. 2017a), e.g., C51, in order to investigate the stable optimization properties within each iteration in Neural FZI. Concretely, we assume Za is absolutely continuous and the current and target value distributions under KL divergence within bounded range have joint supports (Arjovsky & Bottou| 2017). under which the KL divergence is well-defined. Note that this analysis strategy slightly different from vanilla Categorical distribu- tional RI, which also introduces a projection to redistribute probabilities of target value distribution by the neighboring smoothing without the joint support assumption. We slightly simplify Categori- cal distributional RL by assuming that the target distribution is still within the pre-specified support, which is still easy to satisfy in practice given a relative large bounded range (4o 4e] in advance.
To approximate the categorical distribution, we leverage the histogram Sunction fa, with k uni- form partitions on the support to parameterize the approximated probability density function of Z(s, a). With KL divergence as dy. we can eventually derive the distribution objective function to be optimized within each update in Neural FZI, which is similar to the histogram distributional loss proposed in (Imani & White 2018).
In particular, we denote x(s) as the state feature on each state s, and we let the support of Z(s, a) be uniformly partitioned into bins. The output dimension of f can be |A] x k. where we use the in- dex to focus on the function fe.e_ Hence, the function f=.a_ : -> [0, 1]. provides a k-dimensional vector f (x()) of the coefficients, indicating the probability that the target is in this bin given the state feature x(s) and action a. Next, we use softmax based on the linear approximation x(s) 0, use to express fa.", i.e., f..a.00 (x(s)) exp (x(s) TB.) Σ/=1 exp (x(s) For simplicity,
I (x())) to replace f. (x(s)). Note that the form of fa.. is similar to that in Softmax policy gra- dient optimization (Mei et al. 2020 Sutton & Barto 2018), but here we focus on the value-based
RL rather than the policy gradient RL. Our prediction probability ft is redefined the probability in the i-th bin over the support of Z(s, a). thus eventually serving as a density function. While the linear approximator is clearly limited, this is the setting where far the cleanest results have been achieved and understanding this setting is a necessary first step towards the bigger problem of under- standing distributional RL algorithms. Under this categorical parameterization equipped with KI. divergence, the resulting distributional objective function La (s, a) for the continuous optimization for each s, pair in each iteration of Neural FZI (Eq.
2)
can be expressed as:
ao(s,a)== = Σ
&
|
A.t:
(g) log
f (x(s)) dy x
Σ p. log (x(s))
(3)
Wi
4

---

Under review as a conference paper at ICLR 2023
where 0 = (80..... 0x) and pi is the probability in the i-th bin of the true density function p (**(() for Z(s, a) defined in 6 w; is the width for the i-th bin (4, 4.1]. The derivation of the categorical distributional loss under the categorical parameterization is given in Appendix A To attain the stable optimization property of distributional RI., we firstly derive the appealing properties of the new categorical distributional loss in Eq. 3 as shown in Proposition []
Proposition 1. (Properties of Categorical Distributional Loss) Assume the state features ||x())]| - I for each state s. then Lo is kl-Lipschitz continuous, -mooohh and convex w.E.t the parameter 8.
Please refer to Appendix B for the proof. The derived smoothness properties of d, under the categor- ical distributional loss plays an integral role in the stable optimization for distributional RL. In stark contrast, classical RL optimizes a least squared loss function (Sutton & Barto 2018) in Neural FQI. It is known that the least squared estimator has bounded Lipschitz constant in general and is only Amax -mooth, where Amax is the largest singular value of the design or data matrix. More specifically, for the categorical distributional loss in distributional RI. we have V Coll || I kl, while the gradient norm in classical RL is ly:- - 2 (s, a) x (+)]l, where ( a) = Σ (4: + 11+1)58 (x(*)) /2w; un- der the same categorical parameterization for a fair comparison. Clearly, Q a) can be sufficiently large if the support |4o, [x] is specified to be large, which is common in environments where the agent is able to attain a high level of expected returns (Bellemare et al. 2017a). As such, 1// - Q a)l can vary significantly more than l and therefore classical RL with the potentially larger upper bound of gradient norms is prone to the instability optimization issue.
After providing the intuitive comparison in terms of gradient norms above, we next show that distri- butional RL loss can induce an uniform stability property under the desirable smoothness properties analyzed in Proposition [] We recap the definition of uniform stability for an algorithm while run- ning Stochastic Gradient Descent (SGD) in Definition [
Definition 1. (Uniform Stability) (Hardt al. 2016) Consider loss function gw (z) parameterized by w encountered on the example 2, a randomized algorithm M is uniformly stable if for all data sets D, DD' such that D, 'D' differ in most one example, we have
sup B. [=M - 9MMD') (=)] S Estab
(4)
In Theorem T.
we show that while running SGD to solve the categorical distributional loss within each Neural FZI, the continuous optimization process in each iteration is 6aab Ssibuunifrm stable.
Theorem 1. (Stable Optimization for Distributional RL) Suppose that we run SGD under Co in Eq. 3 with step sizes x S 2/kl² for T steps. Assume [x(s]] $ &forr each state and action a, then we have La satisfies the uniform stability in Definition 7 with Estab < AKT i.e.
4KT
E Co a) - Loy, (s, a) S
(5)
where 0T and Gr are the minimizers after T steps under the dataset D and D', respectively.
Please refer to the proof of Theorem in Appendix C The stable optimization has multiple advan- tages. In deep learning optimization literature (Hardt et al. 20166). an uniform stability can guarantee Guah -bounded generalization gap. In reinforcement learning, algorithms with more stability tend to achieve a better final performance (Bjorck et al. 2021 Li & Pathak [021] Ahmed et af. 2019).
In summary, under the categorical parameterization equipped with KL divergence, the continu- ous optimization objective function within each update of Neural FZI for distributional RL is uni- formly stable with the stability errors shrinking at the rate of O(n -1). and the immediately obtained bounded generalization gap also guarantees a desirable local minima. This advantage can be owing to the desirable smoothness property of categorical distributional loss with a potentially smaller up- per bound of gradient norms compared with classical RL. Empirically, in Section 4 we validate the stable gradient behaviors of categorical distributional RI, and similar results are also observed in Quantile Regression distributional RL. By contrast, without these smooth properties, classical RL. may not yield the stable optimization property directly. For example, Amax -sooth may be of less help for the optimization given bad conditional number of the design matrix where Amax could be sufficiently large. The potential optimization instability for classical RL can be used to explain its inferiority to distributional RL in most environments, although it may not explain why distri- butional RL could not perform favorably in certain games (Ceron & Castro 202I). We leave the comprehensive explanation future works.
5

---

Under review as a conference paper at ICLR 2023
Remark Non-linear Categorical Parameterization. Aithough the aforementioned stability optimization conclusions are established on the linear categorical parameterization on the value distribution of Z*. Similar conclusions can be extended in the non-convex optimization case with a non-linear categorical parameterization by techniques proposed in (Hardt et al. (016). We also empirically validate our theoretical conclusions in the experiments by directly applying practical neural network parameterized distributional RL algorithms.
3.3 ACCELERATION EFFECT OF DISTRIBUTIONAL RL
To characterize the acceleration effect of distributional RL,, we additionally leverage the recently proposed value distribution decomposition (Sun et al. 2021b) to decompose the target pa1..
Value Distribution Decomposition. In order to decompose the optimization impact of value dis- tribution into its expectation and the remaining distribution part, adopt the wisdom from robust statistics via a variant of gross error model (Huber 2004). Value distribution decomposition (Sun) et al., 2021b) was successfully applied derive the distributional regularization effect of distri- butional RL. We utilize Fe.a to express the distribution function of Z" a) and we consider the function class of 1e that satisfies the following expectation decomposition:
(x) = (1 - <) (x23(2 (s.a))) (x) + (
(6)
where the distribution function F is determined by F*.a and e measure the impact of remaining distribution independent of its expectation E [1 a)]. controls the proportion of " (az) and the indicator function 1
[ a)], otherwise 0. Although the function class of F*. is restricted to satisfy this decomposition equality, it is still rich with the rationale rigorously demonstrated in (Sun al.] 2021b). To reveal the speeding up effect of distributional RL loss, we consider the density function form of Eq. 6 i.e., p² (x) (1 <) () (z) + (xx), where 6 = (s.aa) is a Dirac function centered at E IZ (s, a)] to characterize the expectation im- pact and n= is the density function of F to measure the addition value distribution information.
[=23]2* (s.ai]} 1 if
Within Neural FZI, our goal is to minimize 1
We rewrite Ca(s, a) as
La La""",, Jz""), where the target density function g can be p" 1:a0a or ((0=812"( and fs:a.8 is rewritten as fo for conciseness. We denote G*(0) (0) = E [[((0(==) [e and use G(O) for (0) for simplicity. Based on the categorical parameterization in Section 3.2. the convex and smooth properties with respect to the parameter in fe as shown in Proposition sstll hold for G(0). As the KI divergence enjoys the property of unbiased gradient estimates, let the variance of its stochastic gradient over the expectation 6 ((==B(==ssa))) be bounded, i.e.,
(s..))=p= [[|V Lo ((( )0
VG(O) =
Next, following the similar label smoothing analysis (Xu al. 2020), we further characterize the approximation degree of fer to the target value distribution 1eaaaa by measuring its variance as ko?:
E (s,a)~p" [ ff VG(O)] 0² := =
Notably, can be used to measure the approximation error between f and μ.² and we do not assume 8² to be bounded as can be arbitrarily large. This expression kod for 82 allows us to utilize to characterize different acceleration effects for distributional RL given different k. Concretely, a favorable approximation of fo to fr* would lead to small x that contributes to the acceleration effect of distributional RL as shown in Theorem 2
Proposition 2. Based on the value distribution decomposition in Eq. 6 and Eq. 8 we have:
(s,a)~pp
[ ff - VG(O)) (1- + eno?
Based on Eq. g. we immediately have Proposition 2 with proof in Appendix for the proof. Before comparing the sample complexity in the optimization process of both classical and distributional RL, we provide the definition of the first-order r-stationary point, which preferred in the optimization of deep learning rather than the simple stationary point in order to guarantee the generalization.
Definition 2. (First-order r-Stationary Point) While solving ming G(B). the updated parameters Or after T steps is first-order T-stationary point V < T, where the small is in (0, 1).

---

Under review as a conference paper at ICLR 2023
Based on Definition 2 we formally characterize the acceleration effects for distributional RL in Theorem that depends upon approximation errors between μ²= and fgg measured by k.
Theorem 2. (Sample Complexity and Acceleration Effects of Distributional RL) While running SGD to minimize La in Eg. within Neural FZI, we assume the step size À = 1/kl", = 1/(1+k) across (2) and (3), and the sample is uniformly drawn from T samples, then:
(1) (Classical RL) When minimizing (00000=0 to 7-stationary point in expectation.
= 0( such that Lg converges
(2) (Distributional RL with S 1/4 When minimizing La Ca(p" (000), let T converges to a T-stationary point in expectation.
4G100)
= 0((2), La
xx -xx
(3) (Distributional RL with >> t/) When minimizing C :6 let = G(0o) Xucoo
( La
does not converge to T-stationary point, but can guarantee a (1K2) ((*=))ssationnary point.
The proof is provided in Appendix E Theorem iis is inspired by the intuitive connection between the value distribution in distributional RL and the label distribution in label smoothing technique (Xu demonstrates that solving categorical distributional loss of distributional RI can speed up the convergence if distribution approximation error is favorable. Otherwise, the convergence point, albeit stationary, may not guarantee a desirable performance un- der an agnostic Ki which may be very large on certain environments.
al.) 2020). Importantly, Theorem 2
In Classical RL scenario, we provide an equivalence between Lo [o(0 Sg f2 and mean squared error loss (Eq. [1 in Neural FQI in Appendix H In the first seenario ((2) in Theorem 2 there is only a small approximation or paramterization error between fg and p²* (or 2²* corre- sponding small with K S 2/4 In this case, solving Lo based on the categorical parameteriza- tion can reduce the sample complexity from ((+) to ( compared with classical RL in (I) of Theorem 2 and meanwhile guarantees a 7-stationary point. In the second scenario ((3) in Theo- rem 2 especially for some challenging environments with much intrinsic uncertainty, we can also attain a relatively large approximation error or parameterization error of Ze with a large k > 2n as the distributional TD approximation error could be potentially large in practice. Under this circum- stance, distributional RI. algorithms may fail to speed up the convergence or achieve the superior performance compared with classical RI. as (xx)) could be potentially large on some complex en- vironments. If ((x2) is proper, distributional RL can still potentially perform reasonably due the to point guarantee.
Theses theoretical results also coincide with past empirical observations (Dabney et al. 201 8b Ceron & Castro 202I), where distributional RL algorithms outperform classical RL. in most cases, but are inferior in certain environments. Based certain environments have much intrinsic uncertainty, the distribution parameterization error be- tween Za and the true value distribution under the distributional TD approximation is still too large (k > 20 to guarantee a favorable convergence point for distributional RL algorithms with different dp, which is intuitive.
results in Theorem [ we contend that these
4 EXPERIMENTS
We perform extensive experiments on eight continuous control MuJoCo games to validate the the- oretical optimization advantage of distributional RL algorithms analyzed in Section 3 including the stable gradient behaviors of distributional RL to achieve the uniform stability as well as the acceleration effects determined by the distribution parameterization error.
Implementation. Our implementation is based Soft Actor Critic (SAC) (Haarnoja et al. 2018)
and distributional Soft Actor Critic (Ma et al. 2020). We eliminate the optimization impact of entropy regularization in these algorithm implementations, and thus we denote the resulting algo- rithms as Actor Critic (AC) and Distributional Actor Critic (DAC) for the conciseness. For DAC, we firstly perform the C51 algorithm to the critic to extend the classical critic loss to the distribu- tional version denoted by DAC (C5I) as our theoretical analysis in Sections 3.2 and 3.3 are mainly Quantile
based categorical parameterization. We further apply empirical demonstration Regression distributional RL heuristically, i..., Implicit Quantile Network (IQN), which is denoted as DAC (IQN). Hyper-parameter and more implementation details provided in Appendix F
7

---

Under review as a conference paper at ICLR 2023
ant
humanoid
walker2d
bipedalwalkkerardcore
6in
Average Return
5000
5000
250
250
3000
4000
1000
150
100
2010
3000
AC
AC
1000
DAC(C51)) DAC(IQN)
NABICST CS1)
DACIGE1
1000
DAC(IQN)
1000
AC(ION)
" DAC([ON)
-100
-1000
6.6
halfcheetah
reacher
swimmer
humanoidstandup
180000
Average Return
-1.00
40000
600
60
4000
AC
-300
40
DAC(C51)
DAC(C51)
(C5(C
DAC(C51)
2000
50000
DAC(ION)
-100
DAC(IQN)
DAC((ON)
60000
DAC(ION)
-500
02
04
86
n4
1.0
0.2
#4
1.
02
04
66
6a
10
02
n.
dn
1.0
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Figure 1: Performance. Learning curve of AC, DAC (C51) and DAC (IQN) over seeds with smooth size 5 across eight MuloCo games.
4.1 PERFORMANCE AND UNIFORM STABILITY IN DISTRIBUTIONAL RL OPTIMIZATION
Figure Suggests that DAC (IQN) orange lines outperforms its classical version AC (black lines) across all environments, while DAC (C51) in red lines is inferior to AC on humanoid, walker2d and reacher. This could be explained by a more flexible parameterization of IQN over C51.
We then demonstrate the advantage of uniform optimization stability for distributional RL over classical RL. According to Theorem [ the stable optimization of distribution loss with Neural FZI is described as a bounded loss difference for a neighboring dataset in terms of each state s and action a. In other words, the error bound holds by taking the supreme over each state the agent encounters. To measure this algorithm stability, while far from perfect, we consider to leverage the average gradient norms with respect to the state feature x(s) in the whole optimization process as the proxy due to the fact that the gradient could measure the sensitivity of loss function regarding each state the agent observes. From Figure 2 turns out that both DAC (C51) and DAC (IQN) entail much smaller gradient norm magnitude opposed their classical version AC (black
ant
humanoid
walker2d
bipedalwalkerhardcore
7.5
7.5
5.0
5.0
2.5
2.5
2.5
0.0
0.0
0.0
2.5
2.5
2.5
DAGIG51)
-5.0
DACIOS1;
DAC(O51)
DACTION)
-5.0
DIACTOO)
-7.5
DACIION)
DACTION)
-75
7.5
0.2 0.4
0.6 0.8
1.0
0.2 0.4 0.6 0.8
1.0
0.2 0.4
0.6 0.8 swimmer
1.0
0.2 0.4 0.6 0.8 1.0
halfcheetah
reacher
humanoidstandup
10
log logll x
100
DAC(O51)
-8
DAC(C51)
5
DACIO51)
DACTION)
-10
DACIION)
DACTIO
-10
DACTION)
-15
0.2
0.4 0.6 0.8
1.0
12
0.4
0.6 0.8
1.0
0.6 0.8 1.0
0.6 0.8 1.0
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Figure 2: Stable Optimization. The critic gradient norms in the logarithmic scale regarding the state during the training of AC, DAC (C51), DAC (IQN) over seeds eight MuloCo environments.
8

---

Under review as a conference paper at ICLR 2023
ant
humanoid
walker2d
bipedalwalkerhardcore
10
T1 6ol
ACC DAC(C51)
AG
DACICSO
DAC(C51)
DACIC51)
DACION]
RAGIONI
DACIGO)
0.2 0.4 0.6 0.8 1.0
0.2
0.4 0.6 0.8 1.0
0.2 0.4
0.6
0.8
1.0
0.2 0.4 0.6 0.8 1.0
halfcheetah
reacher
swimmer
humanoidstandup
Vel
5.0
10
log I
2.5
7'
0.0
-2.5
DACIC51)
DACICS
DACIC51) DACTION]
5.0
DACI(ON)
DACIONI)
0.2 0.4 0.6 0.8 Time Steps (1e6)
0.4 0.6 0.8 1.0
0.8 1.0
0.2 0.4 0.6 0.8
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Figure 3: Acceleration Effect.. The critic gradient norms the logarithmic scale regarding network parameters in the training of AC, DAC (C51), DAC (IQN) over 5 seeds on MuJoCo environments.
lines) across all eight MuJoCo environments, which corroborates the theoretical advantage of the uniform optimization stability for distributional RL analyzed in Theorem I
4.2 SMOOTHNESS PROPERTY AND ACCELERATION EFFECT OF DISTRIBUTIONAL RL
Theorem 2 demonstrates that distributional RL can speed up the convergence if the distribution parameterization is appropriate. characterized by the variance of the gradient estimates with a small (case (2) in Theorem 2 To demonstrate it, we use the proxy by evaluating the &y-norms of gradients with respect network parameters of the critic for AC and DAC. We mainly focus on direct comparison between vanilla AC and DAC algorithm, although their network architectures are slightly different. Similar results under the same architecture and via the value distribution decomposition of Eq. are provided in Appendix G
Figure showcases that both DAC (C51) and DAC (IQN) have smaller gradient norms terms of network parameters 0 compared with AC the whole optimization process, which directly validates that distributional RL loss is more likely to enjoy smoothness properties in Proposition [ In terms of acceleration effects, the property of stationary points, albeit being different, in cases (2) and (3) of Theorem 2 guarantees bounded gradient norms, but the precise evaluation of x is tricky in order discriminate either case (2) or (3) for each algorithm in specific environment. Nevertheless, by considering the fact that DAC (IQN) outperforms DAC (C51) in most environments in Figure [ we hypothesize that the inferiority of DAC (C51) on humanoid, walker2d and reacher could be owing to its larger parameterization eтrоrs k in these environments. This results in the worse performance of DAC (C51) compared with DAC (IQN) that is more likely accord with the case (3) in Theorem due to its richer distribution expressiveness power than C51.
DISCUSSIONS AND CONCLUSION
Our optimization analysis of distributional RL is based on the categorical parameterization, and the alternative analysis on Wasserstein distance can be integral complementary for our conclusions. Acceleration effects could be further investigated explain whether a typical distributional RI. algorithm can perform favorably in specific environment. We leave them future works.
In our paper, we answer the question: how does value distribution in distributional RL help the opti- mization from two perspectives, including the stable optimization analysis based on the smoothness property categorical distributional loss, as well as the acceleration effects determined by the vari- ance of gradient estimates. We theoretically and empirically show that distributional RL embraces stable gradient behaviors and could speed up the convergence if the distribution approximation is desirable or the parameterization error is sufficiently small.

---

Under review as a conference paper at ICLR 2023
Ethics Statement. Due to the fact that study is about the theoretical properties of distributional RL algorithms, we do not think our research is involved with any ethics issues.
Reproducibility Statement. As stated in Section our implementation is based on the public code of SAC (Haarnoja et al.) 2018) and Distributional SAC (Ma et al. 2020). We also provide implementation details in Appendix F for reproducibility. For the theoretical results, rigorous proof is also given in Appendix from A to
REFERENCES
Zafarali Ahmed, Nicolas Le Roux, Mohammad Norouzi, and Dale Schuurmans. Understanding the impact of entropy policy optimization. In International conference machine learning, pp- 151-160. PMLR, 2019.
Martin Arjovsky and Léon Bottou. Towards principled methods for training generative adversarial networks. International Conference on Learning Representations, 2017, 2017.
Mare G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective reinforcement
learning. International Conference on Machine Learning (ICML), 2017a.
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 gradi- ents. arXiv preprint arXiv: 10743, 2017b.
Johan Bjorck, Carla P Gomes, and Kilian Q Weinberger. Towards deeper deep reinforcement learn- ing. Advances in neural information processing systems (NeurIPS), 2021.
Johan Samir Obando Ceron and Pablo Samuel Castro. Revisiting rainbow: Promoting more insight- ful and inclusive deep reinforcement learning research. In International Conference on Machine Learning, pp. 1333-1383. PMLR, 2021.
Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. International Conference on Machine Learning (ICML), 2018a.
Will Dabney, Mark Rowland, Mare G Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. Association for the Advancement of Artificial Intelligence (AAAI), 2018b.
Jianqing Fan, Zhaoran Wang. Yuchen Xie, and Zhuoran Yang. theoretical analysis of deep q- learning. In Learning for Dynamics and Control, pp. 486-489, PMLR, 2020.
Tuomas Haarooja, Haoran Tang. Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-based policies. In International Conference on Machine Learning, pp. 1352-1361. PMLR, 2017.
Tuomas Haarnoja, 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 appli- cations. arXiv preprint arXiv:1812.05900 2018.
Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability stochastic gradient descent. In International Conference on Machine Learning, pp. 1234-1234. PMLR, 2016.
Peter Huber. Robust Statistics, volume 523. John Wiley & Sons, 2004.
Ehsan Imani and Martha White. Improving regression performance with distributional losses. In International Conference Machine Learning, pp. 2157-2166. PMLR, 2018.
Alexander Li and Deepak Pathak. Functional regularization for reinforcement learning via learned fourier features. Advances in Neural Information Processing Systems, 34, 2021.
10

---

Under review as a conference paper at ICLR 2023
Yudong Luo, Guiliang Liu, Haonan Duan, Oliver Schulte, and Pascal Poupart. Distributional rein- forcement learning with monotonic splines. In Interaational Conference on Learning Represen- tations, 2021.
Xiaoteng Ma, Li Xia, Zhengyuan Zhou, Jun Yang, and Qianchuan Zhao. Dsac: Distributional soft actor critic for risk-sensitive reinforcement learning. arXiv preprint 10Xiv:2004.15547, 2020.
Yecheng Jason Ma, Dinesh Jayaraman, and Osbert Bastani. Conservative offline distributional rein- forcement learning. arXiv preprint arXiv:2107 10..:11077666666 006606,666000 2021.
Borislav Mavrin, Shangtong Zhang, Hengshuai Yao. Linglong Kong, Kaiwen Wu, and Yaoliang
Yu. Distributional reinforcement learning for efficient exploration. International Conference on Machine Learning (ICMI.), 2019.
Jincheng Mei, Chenjun Xiao, Csaba Szepesvari, and Dale Schuurmans. On the global convergence rates of softmax policy gradient methods. International Conference on Machine Learning, pp- 6820-6829. PMLR, 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 1177540):529-333 2015.
2015.
Thanh Tang Nguyen, Sunil Gupta, and Svetha Venkatesh. Distributional reinforcement learning with maximum mean discrepancy. Association for the Advancement of Artificial Intelligence (AAAI), 2020.
Martin Riedmiller. Neural fitted q iteration-first experiences with data efficient neural reinforce- ment learning method. In European conference on machine learning, pp. 317-328. Springer, 2005.
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, pp. 29-37. PMLR, 2018.
Mark Rowland, Robert Dadashi, Saurabh Kumar, Rémi Munos, Marc G Bellemare, and Will Dab- ney. Statistics and samples in distributional reinforcement learning. International Conference on Machine Learning (ICML), 2019.
Richard Sinkhorn. Diagonal equivalence matrices with prescribed row and column sums. The American Mathematical Monthly, 74(4):402-405, 1967.
Ke Sun, Yi Liu, Yingnan Zhao, Hengshuai Yao, Shangling Jui, and Linglong Kong. Exploring the robustness of distributional reinforcement learning against noisy state observations. arXiv preprint arXiv:2109.0877 2021a.
Ke Sun, Yingnan Zhao, Yi Liu, Shi Enze, Wang Yafei, Yan Xiaodong, Bei Jiang, and Linglong Kong. Interpreting distributional reinforcement learning: A regularization perspective. arXiv preprint 10..::221000115.. 2021b.
Ke Sun, Yingnan Zhao, Yi Liu, Bei Jiang, and Linglong Kong. Distributional reinforcement learning via sinkhorn iterations. arXiv preprint arXiv:2202.00 00769, 2022.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An Introduction. MIT press, 2018.
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8(3-4):279-292, 1992.
Yi Xu, Yuanhong Xu, Qi Qian, Hao Li, and Rong Jin. Towards understanding label smoothing. arXiv preprint arXiv:2006. / 1653, 2020.
Derek Yang, Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tic- Yan Liu. Fully parameterized quan- tile function for distributional reinforcement learning. Advances in neural information processing systems, 32:6193-6200 2019.
Fan Zhou, Jianing Wang, and Xingdong Feng. Non-crossing quantile regression for distributional reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020.

---

Under review as a conference paper at ICLR 2023
A DERIVATION OF CATEGORICAL DISTRIBUTIONAL LOss
We show the derivation details of the Categorical distribution loss starting from KL divergence between p and go- pi is the cumulative probability increment of target distribution {Yi within
the i-th bin, and qo corresponds to (normalized) histogram, and has density values (x(()) per bin. Thus, we have:
DKL (p"0 990 s.a) = /
b
(1) logg"'"""" (y)dd - J." p²' log to (v)dy
paia log 40 " (y)dy
sutt (v) log
f (x(s)) dy
wi
(10)
f. fx(((s)) (F (1, + wes F (4))
log
Wi
Pr log S (x(s))
where the first xx results from the fixed target pa in the Neural FZI framework. The second equality based on the categorical parameterization for the density function % a The last holds because the width parameter ut; can be ignored for this minimization problem.
B PROOF OF PROPOSITION 1
Proof. For the Categorical distributional loss below,
exp (x((s T0:)
Co(s,a)= = Σ
3.a log 1 (x(s)), where (x(s)) E/=1 exp (x(s) TØ,)
p
(1) Convexity. Note that - log Σソ_ exp(x(s) = log Σラ= exp (o)) - xs) T 0. the first term is Log-sum-exp. which is convex (see Convex optimization by Boyd and Vandenberghe), and the second term is affine function. Thus, Lo a) is convex.
(xp(x(>>) 8.)
(2) (s a) is Al-Lipschitz continuous. We compute the gradient of the Histogram distributional loss regarding 0,:
a
sa log (x({))((s)
Pj j=1
09,
Σ Pj
8.0
1
Vo.ff(x(() (x(s)
19 (x(s)
==l
k
=
Σ
Pj 5.0
89 (x(s))
1
1. (x(s)) - (x(s))))x((
(11)
1=1
(1 (x(ss))) - " -pf")(((())) x(s)
Pi
j#i
(p p. 34 (x(s)) (1 - p. /. )o (x(s)) x(s)
(p, (x(())) x(s)
12

---

Under review as a conference paper at ICLR 2023
where 6/ 1 i j, otherwise O. Then, have ||x(s)] s L, we bound the norm of its gradient
log 1j (x(+))]
00
j=1
0
Pj log 13 (x()))I
00,
1=1
/=1
(12)
(v/
s.a
1 (x(()) x(s)
$
Σ
s.a - (xx))))][x())
i=1
S kl
The last equality satisfies because pp. (x(()) is less than 1 and even smaller. Therefore, obtain that Lo ki-Lipschitz.
(3) Le is - smooth. A lemma is that log(1 + exp(x)) is -smooth its second- order gradient is bounded by 1. and if (w)) is B-smooth w.r.t. w, then g((r, w)) is |/|=|||-smoons -moothh Based on this knowledge, we firstly focus on the 1-dimensional case of function log f (z), where expa;
(z) Dis exp
As we have derived, we know that 80, log (2,) dij f (2). Then
1 (=)(E) - f2())) (=)
the second-order gradient is
if k. log (=) S 1, which implies that log ff(2) is 1-smooth. Thus, otherwise (=) Clearly, 00,86x
80.80%
log ((x, 01)) is ||0|| smooth, or 12 smooth. Further, Σド= log ff (x(s)) is also -mmomo as we have
IVa. 8, Σ
3. log (x(s)) Vo, Σ Pj log 1 (x())
pj
k
5.0
Vo, log S (x(x) Vo, log 1z (x()))
Pj
(13)
k
<
>
Pj 8,a
112 vll
j=l
1 f||p-~1) vll
for each parameter y and v. Therefore, further have
V 8
log (x(s) - V
log ff((((())) (x(s)))
1=1
1=1
0, Σ
Pj
5.0 log 1/t (x(s))
Ve,
Σ
p
sa
log fogfy (x(s))
(14)
j=1
j=1
(2 111u vll
1=1
k2 || μ
Finally, we conclude that (o a) is -mooth.
13

---

Under review as a conference paper at ICLR 2023
C PROOF OF THEOREM 1
Proof. Consider the stochastic gradient descent rule as Gx,c(0) = 0 - \VgCa. Firstly, we provide Iwo definitions about Le for the following proof.
Definition 3. (a-bounded) An update rule is a o-bounded if supa |10 XV S
Definition 4. (n-expansivve) An update rule is n-expansive if sup,, p.w.
ا|س=a| [ S n.
| (())
Lemma 1. (Grow Recursion, Lemma 2.5 Hard et al. 2016 Fix arbitrary sequence of updates
Gi..., GT and another sequence Gi,,.., GT. Let 0g = 8 be the starting point and define ó = 110, Bell, where 0t and f are defined recursively through
0c+1 Gx,c(0,,,, 044 (())
Then we have the recurrence relation:
G: = Gú is n-expansive
8c+1 I
min(n, 1)6, + 20,
Gt and G2, are a-bounded Ge is expansive
Lemma 2. (Lipschitz Continuity) Assume Lo is L-Lipschitz, the gradient update Gx,c is (AL)- bounded.
Proof. |10 -
| |AVoCo AL
Assume Lemma 3. (Lipschitz Smoothness) Assume Lo is 3-smooth, then for any 3. the gradient update Gx,c is I-expansive.
Proof. Please refer to Lemma 3.7 in (Hardt al. 2016) for the proof.
Based on all the results above, start to prove Theorem 1 Our proof is largely based (Hardt et al. 2016), but it is applicable in distributional RL setting as well as considering desirable prop- attain that Lo is kl-Lipschitz as well as kl=smooth, and thus based on Lemma 2 and Lemma 3] we have Gx,C is -bounddudddd, and 1-expansive if X S RTZ- In the step f, SGD selects samples that are both in and DD, with prob- ability In this case, G2 = G¹,, and thus 5(41 I 6₁ as Ge is 1-expansive based on Lemma The other case is that samples selected are different with probability 4, where de+1 S 6, + 2 based Lemma 1 Thus, if A, F7.. for each state and action a, we have:
erties of histogram distributional loss. According to Proposition [
Lor ()) Loy (s, a)l S kIE [ōr]. where by = 1|0 8711
2AT kl
< kd
(1
IE (8T
E (ōr +
2AT
kd
E 18T
+
2
(15)
kl
E
[ō0]
Σ
t=
22
T-1
2
71
R:2
4kT
Since this bounds hold for all D, and 8, a, we attain the uniform stability in Definition for our categorical distributional loss applied in distributional RL.
Define the population risk as:
R 00 (s,
14

---

Under review as a conference paper at ICLR 2023
and the empirical risk as:
Rs [
Σ
Co S; 0:)
n
According to Theorem 2.2 in (Hardt et al. 2016). if an algorithm M is Estab Estab -atu-rni stable, then the generalization gap 6aab bounded, i.c.,
|Es.A [Rs[MDD)] R[M(D)II S
stabb
D PROOF OF PROPOSITION 2
(s,a)~-p" [ - VG
s ( (1- (1-c) + Eno2
(16)
Proof. As we know that (r)) = (1- (1-c)0(m= (2) cp. and we use KL divergence in Cg. then we have:
" 100 (1 10- + + LV _1²0)
Therefore,
((.a)ppp V (p :10 f((a) [G(+))1
", -
= (1 -)- + exxo²
(17)
where the first inequality uses the triangle inequality of norm, i.e., I1(1 1a11-l((---aaee c)a-cbl22 S (1-6)3 @|Ibil?, and the last equality uses the definition of the variance of Lo Co (==0 S0 and
2+
Lo Co (aaa).
E PROOF OF THEOREM 2
Proof. (1) If we only consider the expectation of Z* a), we use the information 8 ((=F(( to construct the loss function. As La Lo(0)u=8/2 "0) is -mooth,, we have
G(O++1) G(O+)
ki2 11041 0e1
< <VVG(0)). - 01) +
(18)
= =-V(VG())) VLo( 1" + x2
2 (1 f2-))
where the last first equation is according to the definition of Lipschitz-smoothene and the last second one based on the updating rule of f. Next, we take the expectation on both sides,
E (G((++1) G(.))
x2² E
AE [V
1/ VG(A.) VG(A)))
x1
VG(A))] + 112
S -EE
E [IVG(()))
2
A(KP2 2)
x
2 E
[IV
+ 2
0²
- -AA [IVG
x2 0
(19)
where the first two equation hold because VG(0) E. [VLg] and the last inequality comes from RTS Through the summation, we obtain that
XP
BlG(Or) - G(O)) < - t=o
E [IVG((0))]]]]
0
2
15

---

Under review as a conference paper at ICLR 2023
We let [G(+r)] 0, we have
2G((0)) XT
T Σ E [INGOO)
+ =12200"
By setting À that the degenerated loss function based on the expectation ó (s can achieve r-stationary point if the sample complexity T 0
2k and 4G16.)
Xr² we can have - [IV < 73, implying
(2) and (3) We are still based on the -moothnesss of C(p f8
G(O++1) G(O+)
1.2 10 Bell
$ (VG(()), 0(+1 01) +
= - (VG(O)). VColpre ,10²)) + x2 [VLolp ,18"))1
(20)
2 A IVG(())) A (VG(()) - VColp" 11a)) A(kl2> 1)
| Cp "(")
2
X (VG(()) - VCo (0000 ,100011²
S - --1/VG(())1
where the second equation based on (a, -b) 2 (Iaa bll |||22 [[b| and the last in- equality is according to $ 1 After taking the expectation, we have
E [G(0)+1) - G(
s E [[VG(()) + [[V - (
(21)
s - |[G(()) + e)=2 + x²n²²)
where the last inequality is based on Proposition We take the summation, and therefore,
TA ((1 - =)²0² - +e²no?)
E [G((T) G(0o))
E [IVG[())
2
t=
We let R[G(Or)) and 1 then,
14x
E [IVGG())1
1=0
2G((o)
(1 c)2 + 2kk
XI
(22)
2G(Oo)
2x²
(1+ k}+ k))
XI
2G((u)
2
XT
F S 20 and let T 1G(0,)
Ar this leads to $ Et [ [[V((() S 72, i.c. r-stationary point, with the sample complexity as ( Thus, (2) has been proved. the other hand, x 2
z- This implies that & -I- E [IVGG(O)] 4n² O Therefore, the degree of stationary point is determined the degree of distribution approximation measured by k. Thus, obtain (3).
set 7 G(O,)
F IMPLEMENTATION DETAILS
Our implementation is directly adapted from the source code in (Ma et al. 2020). For DAC (IQN), we consider the quantile regression for the distribution estimation on the critic loss. Instead of us- ing fixed quantiles in QR-DQN (Dabney al. 2018b), we leverage the quantile fraction generation
16

---

Under review as a conference paper at ICLR 2023
Table I: Hyper-parametter Sheet.
Hyperparameter
Value
Shared
Policy network learning rate
3e-4
(Quantile / Categorical) 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 before training
le4
DAC (IQN)
Number quantile fractions (N)
32
Quantile fraction embedding size
Huber regression threshold
DAC (C51)
Number Atoms (k)
51
Hyperparameter
I for C51
Max episode lenght
Walker2d-v2
500
1000
Swimmer-v2
1000
Reacher-v2
500
1000
Ant-v2
500
1000
HalfChectah-v2
10,000
1000
Humanoid-v2
5,000
1000
HumanoidStandup--2
15,000
1000
BipedalWalkerrHarddooeez
50
2000
based IQN (Dabney al.) (018a) that uniformly samples quantile fractions in order approxi-
mate the full quantile function. In particular, we fix the number of quantile fractions as and keep
them in an ascending order. Besides, we adapt the sampling as To = 0, = = E/Eno N-1 where G E [0, 1), N.
PAPEE-PARAMETERS AND NETWORK STRUCTURE
We adopt the same hyper-parameters, which is listed Table T and network structure as in the original distributional SAC paper (Ma et al. 2020).
F.2 BEST L FOR DAC (C51)
As suggested in Table 1 after line search for the hyperparameter tuning, we select 4, as 500, 10,000, 15,000, 160, 50, 5,000, 500, 500 for ant, halfcheetah, humanoidstand, swimmer, bipedal- walkerhardeore, humanoid, walker2d and reacher, respectively.
G EXPERIMENTAL RESULTS ON ACCELERATION EFFECTS OF DISTRIBUTIONAL RL
Same Architecture. For a fair comparison, we keep the same DAC network architecture and eval- uate the gradient norms of DAC (C51) and a variant of AC, which optimized based on the ex- pectation of the represented value distribution within the DAC implementation framework. Figure 4
17

---

Under review as a conference paper at ICLR 2023
suggests DAC (C51) still enjoys smaller gradient norms compared with AC in this fair comparison setting.
ant
humanoid
walker2d
bipedalwalkerhardcore
log 1
20
30
20
15
15
10
10
5
AC
AC
AC
BAC(CST)
0
DACIC511
0
DARISED
-5
DAC(C51)
5
0.6 halfcheetah
0.2
0.4
0.8
1.0
0.2
0.4 0,6 0.8
1.0
0.2
0.4
0.6
0.8
1.0
0.2 0.4 0.6
0.8
1.0
reacher
swimmer
humanoidstandup
30
log 1
15
10
20
20
5
10
10
0
AC
AC
-5
AC
AC
-10
-10
-10
-10
LAC
0.2 0.4 0.6
0.8
1.0
0.2 0.4 0.6 0.8
1.0
0.2 0.4 0.6
0.8 1.0
0.2 0.4 0.6 0.8 1.0
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Figure 4: The critic gradient norms in the logarithmic scale during the training of AC and DAC (C51) over seeds three MuJoCo games. We keep the same DAC network architecture and evaluate based on the expectation of the represented value distribution.
Results under Value Distribution Decomposition We also provide gradient norms of both ex- pectation and distribution based on the value distribution decomposition Eq. 6 Similar results can be still observed in Figure
ant
humanoid
walker2d
bipedalwalkkeraardoore
log [e
10
15
10.0
7.5
5.0
7.5
5
10
2.5
0.0
0
2.5
-2.5
0.0
-5.0
DAC(C51)
-2.5
(AC(OST)
DAC
-5
-5.0
7.5
0.2 0.4 0.6 1.0 halfcheetah
0.6 0.8 reacher
0.2
1.0
0.4 0.6 0.8
swimmer
humanoidstandup
log eeI12
15
10
10
5
0
-5
AC
-10
(AC
-10
-10
-10
DAC(C51)
0.2 0.4 0.6 0.8 1.0
0.2 0.4 0.6 0.8 1.0
0.2 0.4 0.6 0.8 1.0
0.2 0.4 0.6 0.8 1.0
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Time Steps (1e6)
Figure 5: The critic gradient norms in the logarithmic scale during the training of AC and DAC(C5)) over 5 seeds on three MuJoCo games. Results of AC is the expectation part calculated via the Value Distribution Decomposition.
H EQUIVALENCE BETWEEN THE LOSS FUNCTION IN THEOREM AND MEAN SQUARED LOSS IN NEURAL FOI
Proposition (Connection between Theorem and Mean Squared Loss in Neural FQI) In Eq. 2 of Neural FZI, if the function class (Zo- E ②) is sufficiently large such that it contains the
18

---

Under review as a conference paper at ICLR 2023
target (Y, 1- where Y; R(s,, a,) 7Zz (2())= Minimizing (o _6 in
Theorem 1J. as max; Wi - 0 implies
((Zo(s,a) 7 (s- 1.
(23)
Proof. Bellman Optimality Operator is ( a) = E[R(s, a)] + 7 maxx Eg' 1Q (s', [')] We
also define the distributional Bellman optimality operator Iup as follows:
T Z(5, a) _ R(s + 77 (S', (5', "15',''' aa)
(24)
S' P(. s,a,, a" = argmaxE 2 (S',
For the uniform notation, we ignore k in Neural FZI. If {Zg 0 such that contains Zp- then optimizing Neural FZI in Eq. leads to Zg = TOp 7g. Similarly, optimizing Neural FQI yields Qo T Qg- ideally.
- } is sufficiently large enough
We denote we as the interval that [targe(ss.a] and Zurger Io in expectation as shown in Neural FZI. fāra q000 as Wmax max; W/ O, where ge is the continuous target probability density function. Then we have:
Lo (a RiZ (s,a))) f
JEEE (1=10/2400ssaa1 log f2 (w""wg))
)x
(25)
+ log 20 8.0 (FE [ (s,a)) as WE + 0
where we know [((r)d 1. Since {2 0 θ) is sufficiently large enough, the KL min- imizer would be Ug = (s.2))((..)))- where (8(20)"(s.a)) is a Dirae Delta function centered at
[s a)] and can be viewed generalized probability density function. According to the definition of Dirac Delta function, Wmmx O, we attain
P P(Zo a) E 7uugel (s,a))
E T Zg. (s, a)])
(26)
Due to the linearity of expectation analyzed in Lemma 4 of (Bellemare al. 2017a), have (s,
E T Zo.
" [Ze. (s,
(27)
-o- (s,
Finally, for each iteration in Neural FZI, the following equation always holds:
P(žo(s, TT (s, a)) as max w.
(28)
This indicates that the optimal value distribution Zo may take other values instead of the expectation,
but the probabilities of these events happen are O. This implies that minimizing the KL divergence in terms of the dirac Delta function "almost" equivalent the minimizer of mean squared loss. O
19


Reviews:
[{"review_id": "oW5WOxyfFn", "reviewer": "Reviewer_bwzS", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper analyzes the optimization dynamics induced by distributional losses. It does so using two main classes of theoretical tools: first, it studies the smoothness of distributional losses and the consequent stability of gradient descent on these losses. Second, it studies the variance of the gradients obtained by distributional losses and characterizes the convergence rate of gradient descent. The qualitative predictions of the theoretical results are evaluated on two distributional algorithms.", "summary_of_the_review": "This paper addresses an interesting problem, but suffers from a lack of clarity and technical depth. Of its main contributions, the result on stability and smoothness of distributional RL losses are straightforward corollaries of previously known results. The analysis of the convergence rate of distributional vs expected updates is not clear, and it is difficult to interpret how restrictive the necessary assumptions are for Theorem 2 (in particular case (2)) to remain valid. Finally, there is a fairly large gap between the setting of the theoretical analysis and the empirical experiments, which limits the insight these experiments are able to provide. \n", "strength_and_weaknesses": "### Strengths\n- The paper combines a set of results from fairly disparate regions of the literature into a unified story to explain the benefits of distributional RL. \n\n- The strategy of decomposing the distributional RL objective based on the mean value of the state and a perturbation distribution is appealing.\n\n- The figures in section 4 are easy to read and interpret. \n\n\n### Weaknesses\n\n- Section 3.2: \t\n  - The proof of smoothness of cross-entropy/KL is widely known and similar analysis is already provided by Imani and White. \n  - The stability result is a straightforward corollary of the previous observations on stability of SGD and smoothness of the CE loss. These bounds are also not very informative of generalization in practice because they are vacuous for optimization trajectories that don’t converge to a low loss in less than one epoch.\n\n- Section 3.3: \t\n  - The notation in section 3.3 is unclear and makes it difficult to follow precisely what the claim is that is stated in Proposition 2. \n  - The notation $G^k(\\theta)$ is not obvious: what is the interpretation of $k$? Is it the loss with respect to the current targets at iteration $k$ of FZI? Why not include this dependence elsewhere in that case?\n    - Is f_\\theta^{s,a} meant to be the density of the true return, or of the bellman targets?\n    - The overloading of $\\mathcal{L}_\\theta$ is ambiguous as the paper does not specify which argument corresponds to the prediction and which to the target distribution, and hence which order these appear in the KL divergence.\n\t-I am assuming that the notation $\\nabla G$ corresponds to the expectation of the left hand side gradient in equations 7 and 8 from context, but it is defined as the expectation of $L(s,a)$ in the paragraph previously, which was previously defined to be shorthand for the gradient of the current function approximator predictions with respect to the true FZI targets. This inconsistency makes it difficult to interpret the proposition 2 and theorem 2. \n\t- The result on distributional regularization would benefit from clarification on the role of $\\kappa$ and $\\sigma^2$. In the case of a deterministic MDP it seems that case (2) should be identical to case (1), so I assume that instead the assumptions required by theorem 2 preclude this type of prediction target, in particular due to the lower bound on the value of $\\kappa$. In this case it seems that the assumptions of theorem 2 are relatively restrictive — what amount of noise in practice would be needed to fall into this sweet spot where convergence is faster than deterministic targets but still guarantees a $\\tau$-stationary point?\n\n- Section 4\n\n- The use of actor-critic algorithms in section 4 is not well-motivated. Prior work (Ilyas et al.) has shown that the critic value estimates in these architectures are often wildly inaccurate and yet still provide significant benefits to performance. Thus it cannot be concluded that any performance improvement obtained by altering the critic architecture is caused by improved value prediction accuracy or stability of the critic.\n- The IQN baseline isn’t informative as it follows a markedly different training procedure from that studied in previous sections. I would prefer to have seen a QR-DQN baseline with a similar parameterization as the C51 agent for a more reasonable comparison.\n- The empirical results exhibit a similar ranking as is predicted by the theoretical results, though the convergence rates don’t seem to be consistent.\n- Section 3 focuses on a single iteration of FZI whereas SAC follows an entirely different learning algorithm. It is not clear to me why we should expect the findings on convergence to a fixed set of targets to hold for the SAC algorithm, which updates the targets more frequently.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "VK_jlZUAzfU", "reviewer": "Reviewer_d87u", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies the convergence rate of neural fitted Z-iteration with a discrete distribution approximation of the Q distribution. In particular, the authors demonstrate that under the linear categorical parameterization assumption, neural fitted Z-iteration converges faster with the distributional RL objective compared to the classical RL objective (where the expectation is used as the fitting target rather than a distribution). The paper also shows that the gradient norm of the critics with respect to the network parameters for practical distributional RL algorithms is generally lower than the gradient norm for practical non-distributional RL algorithms, corroborating the theoretical results.  ", "summary_of_the_review": "While the theoretical contributions are interesting on their own, my biggest concern of the paper at its current state is that both the theoretical results and the empirical results do not demonstrate why distributional RL is more stable than classical RL (which from my understanding is the main point of the paper). Specifically, the theoretical results do not analyze the correct objective for the classical RL (see above) and because of it, the empirical results that aim to corroborate the theoretical insights are also confounded by the objective mismatch. Therefore, I would not recommend acceptance of the paper.", "strength_and_weaknesses": "*Strength*\n\n- The theoretical results in the paper look reasonable and likely to be correct. Though admittedly I did not check all the proof details in the appendix. \n- The paper is doing a great job of positioning itself in the literature.   \n\n*Weaknesses*\n\n- The paper argues that there are acceleration effects of distributional RL compared to classical RL which optimizes the mean square error $(y\\_i - Q\\_\\theta^k(s_i, a_i))^2$. However, this is not well supported by the theorem as the main result that demonstrates the acceleration effects do not use the mean square error. Rather it uses KL with the target density being the dirac delta of expectation of the value target (see Theorem 2 (1) -- $\\mathcal{L}\\_\\theta(\\delta\\_{x=\\mathbb{E}\\left[Z^\\pi(s, a)\\right]}, f\\_\\theta^{s, a})$). It is not obvious to me how the results under this KL objective may be translated to the mean square error objective. \n\n- The empirical results do not support the theory well. While the paper shows that the critic gradient norm with respect to the parameters for distributional RL is smaller, this comparison is confounded by the difference in the objective (e.g., AC uses mean square error, C51 uses KL, and IQN uses Wasserstein distance).", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "Bnw7qNPqSMF", "reviewer": "Reviewer_poww", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper proposes to understand the benefits of distributional RL through an optimization perspective. In fitted Z iteration and under categorical distributional parameterization, the paper analyzes the effect of distributional parameterization on the optimization problem, concretely a certain notion of stability (Thm 1) and guarantee to stationary point under finite sample (Thm 2). They also show some experiments to support the claim and showcase improvements of distributional RL over value-based RL in the control case.", "summary_of_the_review": "A few questions.\n\n=== **Goal of the paper and theory vs. practice** ===\n\nA high level concern I have for the paper, is that the proposed theoretical explanation to justify the benefits of C51 over value-based RL, is a bit detached from practice. As a concrete example, Thm 1 shows that softmax parameterization entails certain uniform stability for the optimization problem. The logic here is that as a result of the stability, C51 should work better than just optimizing the expected value. This is not a satisfying explanation.\n\nTo see why, consider an algorithm that \"represents values with a categorical distribution\", but is not doing distributional RL. This algorithm effectively just does value-based RL, but represents the mean value as a softmax distribution instead of a scalar. This technique is commonly applied in a number of deep RL agents, which have proved useful empirically [1,2]. Note that such an algorithm does not use distributional back-up, so does not learn a proper distribution. \n\nIn other words, the arguments in Thm 1 explain why softmax parameterization may work better from an optimization perspective, but does not explain why learning a distribution is useful. However, learning distributions lies at the core of distributional RL, and failing to provide an explanation for the benefits of actually \"learning a distribution\" defeats the purpose of the paper.\n\n[1] MuZero: Mastering Go, chess, shogi and Atari without rules, Schrittwieser et al, 2019\n[2] Muesli: Combining Improvements in Policy Optimization, Hessel et al, 2020\n\nEmpirically, it is worth trying out the above algorithm to see if its performance is in between value-based RL and full C51. If it is indeed the case, then Thm 1 provides a partial explanation why the value-based categorical parameterized algorithm works better than value-based RL, but there is still a gap from full distributional RL or C51.\n\n=== **Results in Thm 2** ===\n\nResults in Thm 2 suffer from the same issue as Thm 1 -- the characterization of the stationary point comes purely from the softmax parameterization, but not from the fact that dist RL actually learns an approximation to the target return distribution. All the characterizations in Thm 2 and Thm 1 are based on softmax parameterization's advantage over scalar parameterization, but not why learning a dist is better than just learning the mean.\n\n=== **Title and scope of the work** ===\n\nGiven that only softmax parameterization is analyzed in the work, I think it is better to modify the title of the paper as reflecting this. Quantile representations and QR-DQN type of algorithms are not included in the analysis. \n\n=== **Experiments** ===\n\nIn the experiment, the authors have compared IQN variants of distributional actor-critic with other baselines, and have showed improvements. This is a bit confusing, because IQN uses quantile representations which are drastically different from C51 in terms of optimization landscape. Do we expect quantile representations to enjoy similar uniform stability and stationary point guarantee as C51? There is no proof in the paper. If not, where should we expect the performance improvements come from?\n\n", "strength_and_weaknesses": "=== Strength ===\n\nThe paper proposes an interesting conceptual question that may have been overlooked in prior literature, i.e., distributional RL may just help with optimizing the mean value by enjoying a better optimization landscape. While much of the conventional wisdom is that distributional RL helps with the side-effect of learning a better representation for value-based agents, this work posits a relatively new and orthogonal perspective.\n\n=== Weakness ===\n\nThe theoretical results are a bit weak and cannot satisfactorily explain what happens in practice. I believe this gap renders the paper not very useful in obtaining insightful understanding of how and why distributional RL works better than value-based RL. I will expand more in the detailed comments.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject