Towards Understanding Distributional Reinforcement Learning: Regularization, Optimization, Acceleration and Sinkhorn Algorithm
Authors: Ke Sun
Yingnan Zhao
Yi Liu
Enze Shi
Yafei Wang
Aref Sadeghi
Xiaodong Yan
Bei Jiang
Linglong Kong
Venue: ICLR 2022 Submitted (2022)
Keywords: distributional reinforcement learning

Abstract: Distributional reinforcement learning~(RL) is a class of state-of-the-art algorithms that estimate the whole distribution of the total return rather than only its expectation. Despite the remarkable performance of distributional RL, a theoretical understanding of its advantages over expectation-based RL remains elusive. In this paper, we interpret distributional RL as entropy-regularized maximum likelihood estimation in the \textit{neural Z-fitted iteration} framework and establish the connection of the resulting risk-aware regularization with maximum entropy RL. In addition, We shed light on the stability-promoting distributional loss with desirable smoothness properties in distributional RL, which can yield stable optimization and guaranteed generalization. We also analyze the acceleration behavior while optimizing distributional RL algorithms and show that an appropriate approximation to the true target distribution can speed up the convergence. From the perspective of representation, we find that distributional RL encourages state representation from the same action class classified by the policy in tighter clusters. Finally, we propose a class of \textit{Sinkhorn distributional RL} algorithm that interpolates between the Wasserstein distance and maximum mean discrepancy~(MMD). Experiments on a suite of Atari games reveal the competitive performance of our algorithm relative to existing state-of-the-art distributional RL algorithms.

Under review as a conference paper at ICLR 2022
TOWARDS UNDERSTANDING DISTRIBUTIONAL REIN- FORCEMENT LEARNING: REGULARIZATION, OPTI- MIZATION, ACCELERATION AND SINKHORN ALGO- RITHM
Anonymous authors Paper under double-blind review
ABSTRACT
Distributional reinforcement learning (RL) is class of state-of-the-ans algorithms that estimate the whole distribution of the total return rather than only its expec- tation. Despite the remarkable performance of distributional RI,, theoretical un- derstanding its advantages over expectation- paper, interpret distributional RL intropy-regularize maximum likelihood estimation the neural Z-fited iteration framework, and establish the connection of the resulting risk-awart regularization with maximum entropy RL. In addi- tion, We shed light on the stability prooootin distributional loss with desirable smoothness properties in distributional RL, which can yield stable optimization and guaranteed generalization. We also analyze the acceleration behavior while optimizing distributional RL algorithms and show that an appropriate approxi- mation to the true target distribution can speed up the convergence. Finally, we propose class of Sinkhorn distributional RL algorithm that interpolates between the Wasserstein distance and maximum mean discrepancy (MMD). Experiments on a suite of Atari games reveal the competitive performance of our algorithm relative to existing state-of-the-art distributional RL algorithms.
RI. remains elusive. In this
INTRODUCTION
The intrinsic characteristics of classical reinforcement learning (RL.) algorithms, such as temporal- difference (TD) learning (Sutton & Barto 2018) and Q-learning (Watkins & Dayan 1992). are based on the expectation of discounted cumulative rewards that an agent observes while interacting with the environment. In stark contrast to the classical expectation-basee RL, new branch of algorithms called distributional RL estimates the full distribution of total returns and has demonstrated state- of of-the-art art performance in wide range of environments (Bellemare al.] [017a Dabney et al.) 2018b] by Yang et al., 2019) Zhou et al. 2020 Nguyen et al. 2020). Meanwhile, distributional RL has also enjoyed further benefits in risk-sensitive control, policy exploration settings (Mavrin al., 2019 Rowland et al. 2019) and robsutness (Sun et al. 20211)
Despite the existence of numerous algorithmic variants of distributional RL with remarkable empir- ical success, theoretical studies of advantages of distributional RL over expectation-baase RL are less established. Existing works include (Lyle et al. 2019) that proved in many tabular and linear approximation settings, distributional RL behaves exactly the same as expectation-
RL. Lyle al. (20IT) investigated the impact of distributional RL from the perspective of representation dy- namics. Martin et al.) (2020) recently mapped distributional RL problems to a Wasserstein gradient flow problem, treating the distributional Bellman residual potential energy functional. Offline distributional RL (Ma et al. 2021) has also been proposed to investigate the efficacy of distributional RL in both rist-neutral and lisk-vverse domains. Recent works have tended towards closing the gap between theory and practice in distributional RL.
From an algorithmic perspective, the Sinkhorn loss (Sinkhom, 1967) can be used to tractably ap- proximate the Wasserstein distance and has been successfully applied in numerous crucial machine learning developments, including the Sinkhorn-GAN (Genevay et al., 2018) and Sinkhorn-bast RI. literature, Martin et al.
adversarial training (Wong al. 2019). Inspired by the distributional

---

Under review as a conference paper at ICLR 2022
2020) argues for the use of second-order stochastic dominance relation to select among mul- tipiiiity of competing solutions via Sinkhorn iteration (Sinkhorn (967). which can be useful to manage stochastic uncertainty in RL paradigms. However, a Sinkhorn- based distributional RL al- gorithm has not yet to be formally proposed and investigated.
In this paper, theoretically illuminate the superiority of distributional RL expectation-basel RL from the perspectives of regularization, optimization, acceleration, and representation. Specif- ically, we simplify distributional RL as neural Z-fitted iteration, within which we establish an equivalence between distributional RL and a form of entropy-regularized maximum likelihood esti- mation (MLE). We also demonstrate that the resulting novel cross entropy regularization correlates strongly with the behavior of maximum entropy RL. By incorporating histogram distributional loss, we further achieve stable optimization and guaranteed generalization of distributional RI., at- tributable to desirable smoothness properties of the distribution loss. We further characterize the effect of acceleration on distributional RL and discuss when distributional RI algorithms are effec- tive in various environments. After gaining insights from the theoretical advantages of distributional RL, we further propose a novel distributional RL algorithm based on Sinkhorn loss that interpolates between the Wasserstein distance and maximum mean discrepancy (MMD). Our approach allows us to find a trade-off that simultaneously leverages the geometry of the Wasserstein distance and the favorable high-dimensional sample complexity and unbiased gradient estimates of MMD. In sum- mary, our analysis opens the door to deeper understanding of theoretical advantage of distributional RL.
2 PRELIMINARY KNOWLEDGE
In the standard RL setting, agent interacts with environment via Markov decision pro- cess (MDP), a 5-tuple (S, R, P) where and A are the state and action spaces, respectively. P is the environment transition dynamics, R is the reward function and y E (0, is the discount factor.
State-value function vs. state-value distribution. Given a policy æ, the discounted sum of future rewards is random variable Z² (s(saa = Er= 7 R[s;, az), where so = S, ag a, 8141 ~ P(-ss, dz), and ag T-) In the control setting, expectation-basee RL focuses on the state-value function (s,saa), which is the expectation of (s, a). i.e., (*(ssa) E [** (s, Distributional RL, on the other hand, focuses on the state-value distribution, the full distribution of (s, a). Leveraging knowledge on the entire distribution can better capture the uncertainty of returns in the MDP beyond the expectation of return (Dabney al. 2018a Mavrin et al. 2019).
Bellman operators
distributional Bellman operators. For the policy evaluation in expectation-based RL, the value function is updated via the Bellman operator T"Q(s,)) = E[R(s, a]] + >m,
- IQ d')] In distributional RL, the state-value distribution of 2 27 a) is updated via the distributional Bellman operator Ix
X I² R(s, 77Z(6',a))
(I)
where S' P(-|s, a) and d' (-1s²).
From a theoretical perspective, both the Bellman operator T in the policy evaluation setting and the Bellman optimality operator T in the control setting are contractive in the stationary policy case. In contrast, the distributional Bellman operator I is contractive under certain distribution divergence metrics, but the distributional Bellman optimality operator I can only converge to a set of optimal non-stationary value distributions in weak sense (Elie & Arthur 2020).
EFFECT OF REGULARIZATION ON DISTRIBUTIONAL RL
Although the theoretical framework of distributional RL in the tabular setting has been basically established mentioned in Section 2 the theoretical understanding of its advantages over expectation- based RL has been less studied. this section, we attribute the superiority distributional RL into its regularization effect.
2

---

Under review as a conference paper at ICLR 2022
3.1 DISTRIBUTIONAL RL: NEURAL Z-FITTED ITERATION
Neural Q-Fitted Iteration. In the function approximation setting, Deep Q Learning (Maih et al., 2015) can be simplified into Neural Q-Fitted Iteration (Fan al. 2020) under tricks of experience replay and the target network Qø., where we update parameterized Qo(s, a) in each iteration k:
1 = argginn - Qo n
yi - Q6 (st, as), Σ
(2)
where the target Y: r(s;, 04) + m Ø3 (5/ a) is fixed within every Turget steps to up- date target network Qg by letting = 9 and the experience buffer induces independent samples (8i, , Oe, Fi,, ie(n)- In ideal case that neglects the non-convexity and TD approximation errors, we have Q/+11 TQ.. which is exactly the updating under Bellman optimality operator. Under the two conditions, the optimization problem in Eq. _ can be viewed as Least Square Estimation (LSE) in neural network parametric regression problem between the updating of target network Q....
Neural Z-Fitted Iteration. Analogous to neural Q-fitted iteration, we can also simplify value-based distributional RL methods based parameterized Zo into Neural Z-fitted Iteration as
2 zk+1 argmin Zo
dYi,, Zg a:)),
where the target Y; R(s;, aq) + ZZz (5) (5))) with (z')') argmax, E [Z]- (s', a') is fixed within every Turget steps to update target network Za. and dp is a divergence metric between two distributions. Notably, the options of representation manner on Zg and the metrie dp are pivotal for the empirical success of distributional RL algorithms. For instance. QR-DQN (Dabney et al. 20186) approximates Wasserstein distance W,. which leverages quantiles to represent the distri- bution of Zg. C51 (Bellemare et al. 2017a) represents Zp via categorical distribution under the convergence of Cramér distance (Bellemare et al., 2017 Rowland al. 2018). a special case with p 2 of the lp distance (Elie & Arthur 2020), while Moment Matching (Nguyen al. 2020) learns deterministic samples to represent the distribution of La based on Maximum Mean Discrep- ancy (MMD). Contractive properties under typical metries dp can be summarized follows
T is 7-contractive under the supreme form of Wassertein distance Wp.
T is // contractive under the supreme form of lp distance.
T is /:t/:///////o/p contractive under MMD with the kernel ka y) = -|| ull Va E R.
For the completeness, the definition of mentioned distances and the proof of contraction are pro- vided in Appendix [ Although the widely used Kullback-Leible (KL) divergence is not contrac- tion (Morimura et al. 2012), we show Proposition that the KL. divergence still enjoys desirable properties in distributional RL context, which can be reasonable for the theoretical analysis. We assume Zo is absolutely continuous and has joint supports, under which the KL divergence is well- defined. Proof of Proposition and the definition of supreme DKL are provided in Appendix B
Proposition 1. Denote the supreme of DRL D
have: (1) _ is a non-expansive operator
under D&g i.e., DRI(I" ( Z1, I" Z2) < DRS Z1, (2,, Z2), ). (2) D (Zm. Z) - 0 implies Wp(Z,, Z) + O, (3) the expectation of 2 is still y-contructive, i.e., E Z E Zellov S |BZ EZZ2. 1
3.2 DISTRIBUTIONAL RL: NOVEL ENTROPY-REGGLARRIZE MLE
The reasonable properties of KL divergence in Proposition I allows to leverage it to conduct the theoretical analysis. To separate the impact of additional distribution information from the expec- tation of Z*, we leverage the variant technique of gross error model from robust statisties (Huber) 2004), similar to the technique to analyze Label Smoothing (Müller et af.) 2019) and Knowledge Distillation (Hinton et al., 20155). Specifically, we denote the one-dimensional full distribution of 2" as F, and the distribution on the remaining support getting rid of E [ as Ey. Hence, we can obtain the distribution decomposition for 2 (s, a)
F*((x) (1 --) (02331222( (I) + FF
3

---

Under review as a conference paper at ICLR 2022
where e controls the proportion of (x) and the indicator function 1 (a27]2 (s.a])) if
E (s, a)], otherwise 0. After taking derivatives on both sides, we attain the relationship of their density functions as p*.((z) = (1 =])(===((=(ssa))) (œ) + epari(a (z), where (x)(()) is the density function related to (s, a) on remaining supports removing E [2"(s,, a)]. It is worth noting that the existence of (x) can be simply guaranteed by directly computing 1=* (x) = p='(e//
(1 (1-e6 =x0 (s.a]))/ e as long as p²*(() and the expectation of 2= a) exist. Next, we use (2) and 96 *.00 (x) to denote the density distributions behind (ei and Z (s,a) in neural Z- fitted iteration via Eq. 3 respectively. Therefore, we can derive the following result in Proposition
Proposition 2. Let H(P,, Q) as the cross entropy, i.e., H(P, Q) = =-Saex P(a:) log Q(x) dr. Let be a positive constant, and based on the decomposition in Eq. 4 and Dxz. as dp, Neural Z-fitted iteration Eg. Jean be reformulated as
zzz1
argmin Za
H(0 (s.as))) (s.a))))))) "h) + aH(p .00 000))
(5)
71
We provide the proof in Appendix C
For the uniformity of notation, we still use s, a in the following Importantly, the first term in Eq. [5] can be further simplified as -Txxx log 1 5.0 (E (s, a)))- Minimizing this first term can be viewed as a variant of Maximum
analysis instead of 8t, a, in Eq. [
Likelihood Estimation (MLE) on the expectation E [2, a)] rather than the traditional MLE directly on observed samples. The cross entropy regularization the second term pushes q3 approximate the distribution yr in order to fully utilize the additional distributional information while learning, serving as the key to the superiority of distributional RL. This novel cross entropy regularization regarding 1a= and g0 is different from the classical entropy regularization used in RL, which we further analyze their connection and discrepancy Section [.3 In summary, distributional RL can be simplified as a novel entropy-regullrizee MLE within neural Z.-itted iteration framework in stark contrast to the Least- SaasttSquare estimation of expectation-bast RL in the neural Q-fitted iteration.
3.3 CONNECTION WITH MAXIMUM ENTROPY RL
We establish the connection between the derived novel cross entropy regularization in Eq. in distributional RL with the classical maximum entropy RL (Williams & Peng. 1991). Maximum entropy RI., including Soft Q-Learning (Haarnoja al. 2017). greedily maximizes the entropy of the policy π (-1s) in each state:
((T) = Σ 1=0
Bo (s2, at) + (-(((-())))
where H (To (-Ist)) -. Ta (alse) log no (alst) 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 stochasticity of the optimal policy. This maximum tropy regularization has various conceptual and practical advantages. Firstly, the policy encourages exploration, avoiding situations in which the agent might fall into local optimum. Secondly, it considerably improves learning speed over classical methods and therefore are widely used in state- of-the-art algorithms, c.g., Soft Actor-Critic (SAC) (Haarnoja al. 2018) and PPO (Schulman al.) 2017). Lastly, maximum entropy can yield more robustness to abnormal or rare events while devel- oping task. Similar robustness superiority of distributional RL against noisy state observations has been investigated in (Sun et al. )
Similar benefits of both distributional RL and maximum entropy RL motivate us to explore their mathematical connection, and we refocus on the Bellman updating based on these two kinds of entropy-based regularization. In particular, Soft Policy Evaluation proposed has shown that by alternately iterating the following equations. the algorithm is equivalent to maximize the maximum entropy RL objective function in Eq. 6 T (st, az) 2 r (st, ax) + ylE,, [V (81+1)), where V (8r4-1 Ea, ~ [Q (S14 ae+1) 1] - log ग (a++1 (844 )] is the soft value function (Haarnoja al. 2018). Interestingly, we still have the similar convergence result
SAC
related to our novel cross entropy regularization in Eq. [5] which is shown in Theorem I
Theorem 1. Consider the soft Bellman operator T in Eq. and mapping Q" S A
with [A] & oc, and define Q²-1 TT Q². Given the true distribution prese for each t, then the
4

---

Under review as a conference paper at ICLR 2022
sequence Q" will converge to soft Q-value of I as k 00 with the new entropy objective function as J' (0) Ei= ())ra))> |r (s), 0z) 7/H (feraa (g))))
T (5t, a1) _ r ($t, az) + + IV ($1+1)}
where V ($1+1) = Bai www, -0 [Q ($141, a++1) + log 90 ( and the policy a fol- lows the classical greedy policy rule, i.e. ([-(s) = arg max, E [Z s, a)}, determined by 0.
Please refer to Appendix for the proof. The new en- tropy augmented objective function (0) in Theorem [ reveals that different from adding the entropy regarding the policy π ((11s)) in maximum entropy RL, distributional RL can be viewed as subtracting a cross entropy term be- tween μ and qo. Moreover, the entropy in maximum en- tropy RL state-wise, while our cross entropy regular- ization is state-action-
90 (x)
μ sa (x)
which is more fine-grained characterization on the variability of the full distribution of return Z.
Notably, we highlight that our derived cross entropy reg- ularization is risk-aware. If the true aate-action distri- bution Fμ. has higher degree of dispersion, e.g., a larger variance, than the current qa. minimizing our entropy au- tomatically encourages the policy to explore the uncer- tainty of the environment through changing Ze. A sketch map in Figure iss provided to illustrate the impact of our risk-aware entropy regularization. We remark that this risk-aware regularization may not necessarily encourage exploration under the greedy action rule Tz (s) argmax, EE [Zo a)] even though qo changes towards lr as exhibited in Figure 1 A risk-aware action rule on the current greedy version might be preferred in the future.
Figure 1: Impact of the risk- reg- ularization in distributional RL. The entropy-based regularization will push 4g s. to match the true density func- tion p=.aa determined by the environ- ment, and thus let the learned policy be fully aware of the uncertainty of envi- ronment.
4 STABLE OPTIMIZATION AND GUARANTEED GENERALIZATION
In this section, we further characterize q6 5.a in order to attain stable optimization and guaranteed generalization properties of distributional RL. Concretely, we leverage the histogram function f*.e to parameterize the density function of Zo(s, a), i.e., Ug a.4 still based on the KI. divergence as dp. yielding the histogram distributional loss (Imani & White| 2018) within each update in the neural Z-fitted iteration. Denote x(s) as the state feature each state s, and we let the support of x(s) be uniformly partitioned into & bins. We let the function fa. : X - [0, 11ah provide a k-dimensional vector fa.0 (x()) of the coefficients indicating the probability the target is in that bin given the state 8 and action a pair, and x(s). We use softmax based on the linear approximation x(s) 0, to express f i.e., to replace f..a.00 (x(s)). Therefore, the resulting histogram distributional loss Lo is formulated as:
= exp (o)) / Σj= exp (x(s) 0g). T0,). For simplicity, we use 1 (x()))
*
Lo( = DKL(p" " 3.0 / (x) (x) log (p (z) 2 (z)) dr Σ xEX
p 3,0 log 1 (x(s))
(8)
where 0 0= (01, ......,, } and ,, is the histogram probability of pz'((x) defined in Eq. 4 The deriva- tion of the histogram distributional loss is given in Appendix E To attain the stable optimization
property of distributional RL, we firstly derive Lemma as follows
Lemma 1. (Properties of Histogram Distributional Loss) Assume each state feature is bounded, i.e., [|x(s)] < 1, then Lo is af2-smooth, convex and kl kl-Lipschitz continuous w.r.t.
Please refer to Appendix F for the proof. The derived Lipschitz properties of dp under histogram distributional loss plays an integral part in the stable optimization of distributional RL, but the clas- iteration for expectation- based RL does not enjoy
sical Least-Squared estimation in neural-Q-fitte these smoothness properties. In particular, for histogram distributional loss we have || V/ < kl, while the counterpart in expectation- is 1y (s, a)][[x(+)||, which might be arbitrarily large
5

---

Under review as a conference paper at ICLR 2022
due to the unbounded 1y (s, a)). To further derive the uniform stability of distributional RL while running stochastic gradient descend (SGD), we introduce its definition in Definition I
Definition 1. (Uniform Stubiity) (Hrdd et al. 2016, Consider a function g, a randomized algorithm M is uniformly stuble for all data sets S' such that S,S' differ in at most one exumple, we have
RM (M(((M); 9 (M (5)) ;x)] S Csuub
Given the definition of uniform stability, we show that distributional RL under histogram distribu- tional loss while running SGD is Estab- sca--unifrm stable and generalization guaranteed in Theorem 2
Theorem 2. (Stable Optimization and Guaranteed Generalization) Suppose that we run SGD under Co in Eq.
8 with step sizes > S 2/ for T steps, and [|x())]| S 1 for each state s, we have:
(1) SGD under Le satisfies the uniform stability in Definition ! with Estab S
4KT
(2) Suppose that the data distribution is stable within each neural Z-fitted iteration, the generaliza- tion gap for solving the entropy regularized MLE in Eq. 5 is Ssabb -bounded.
Please refer to the proof of Theorem in Appendix G In summary, the optimization solve the entropy regularized MLE of distributional RL within the update of neural Z-fitted iteration is stable with the stability errors shrinking the rate of Onn This stability optimization properties is ascribed to the smooth and convex properties from histogram distributional loss (non-convex case still holds in (Hardt al., 2016)) and can further control the generalization gap. By contrast, the least-estimation in neural Q- pitttted iteration for expectation based RL, which is without these smooth properties, can not yield the stable optimization and guaranteed generalization directly.
5 ACCELERATION EFFECT
Based on the setting in Section 4 we further incorporate the decomposition of (x)((x) proposed in Eq. 4 into the analysis, where p p=*(x)) (1 =)
+ exxx'(x)) in order to derive the acceleration effect of distributional RL. Within each update in the neural Z-fitted iter- ation, the target is to minimize n Ei= Co(ss, (s; a:). We denote (#) as the expectation of Co, i.e., G* (sa))~>'" [Ca(s, a)}. As such, the convex and smooth properties in Lemma | still hold for G* (0). We use G()) for (0) for simplicity. As the KL divergence enjoys the property of unbiased gradient estimates, we let the variance of its stochastic gradient over the expectation be bounded, i.e., E(s.a)<p" Co (==0(())) () [0)]] = =². Next, we characterize
the approximation degree of % to the ground -tunh distribution uaa defined in Eq. 5 by measuring its variance as ko²;
(s,a)pp [IV[o((e .66 ((-))-G(()133 = no².
(10)
Notably, a favorable approximation of q0 to P1 5. would lead to a small k. in which case an ac- celeration effect of distributional RI. can be derive as shown in Theorem Based Eq. TO immediately have the following lemma.
Lemma 2. As (**(())) (I) (1 (20 (x) + (x)) we have:
E(e,a)=pp [I 46 [= s (1 e)o² + enO"
(11)
Please refer to Appendix H for the proof. To measure the convergence of the entropy regularized
MLE in distributional RL or least-square estimation expectation-bastt RL, we need the following definition of the first-order r-stationary point.
Definition 2. (First-order F --Stationary Point) While solving ming G(Ø), the updated parameters Or
after steps is a first-ordrr r-stationary point ([ S t, where the small T in (0, 1).
We formally show our Theorem to characterize the acceleration effect of distributional RL deter- mined by the variance magnitude T.
Theorem 3. (Acceleration Effect) While running SGD to solve the entropy- MLE within neural Z-ftted iteration in Eq. with the step size X 1/kl2 and e 1/(1 + T), we have:

---

Under review as a conference paper at ICLR 2022
(1) The sample complexity is ((2.) we only consider the expectation in p* Le., [[==
(2) When K S =2 and let T = -x--
4010)) the regularized MLE of distributional RL within each neural
T
Z-fitted iteration converges t0 a r-stationary point in expectation with sample complexity O
(3) When K foz and let = A
G(fa)
the regularized MLE of distributional RL does not converge
to T-stationary point, but we have E [[]] S O(k).
The proof is provided in Appendix ! Theorem 3 is inspired by (Xu al. 2020) to analyze the convergence of label smoothing, which can be similarly utilized to analyze the acceleration effect of distributional RL.. Importantly, we need to emphasize that Theorem why distributional RI. algorithms can achieve inconsistent superiority across different Atari games. This phenomenon is mainly owing to the different degree of TD target approximation to pa1 es- pecially μ² for various environments. In particular, small approximation error of qo to p (or μ)) corresponds to a small K, This can be normally S r² 22/40², yielding better sample efficiency and accelerating the algorithm relative to expectation-based RL. Conversely, an unsatisfactory approxi- mation may not lead to the acceleration effect, but the bounded gradient normally corresponds to a reasonable performance, coinciding with previous empirical observations on various environments.
3 in fact reveals the reason
Extension of Representation Effect. Apart from the acceleration effect, we also conduct some empirical analysis of distributional RL from the perspective of representation. particular, distri- butional RL encourages state representation from the same action class classified by the policy in tighter clusters. Please refer to Appendix K for more details.
ALGORITHM: SINKHORN DISTRIBUTIONAL RL
Through the theoretical analysis, we can gain insights into the advantages of distribution RL over aopectation-based RL, attributing to cntropy-basee regularization, stable optimization and accelera- tion effect determined by the distributional approximation error (measured by k). From the algorith- mic perspective, the choice of representation manner on Zo and the distributional divergence metric dp mentioned in Section 3.] are pivotal for the final performance of distributional RL algorithms.
we further design Sinkhorn distributional RL algorithm. Sinkhorn loss (Sinkhorn 1967) is tractable loss to approximate optimal transport problem by leveraging an entropic regularization to turn the original Wasserstein distance into differentiable and more robust quantity. The resulting loss can be computed using Sinkhorn fixed point iterations, which naturally suitable for modern deep learning frameworks. In partieular, the entropic smoothing generates family of losses inter- polating between Wasserstein distance and Maximum Mean Discrepancy (MMD). Thus it allows us to find sweet trade-off that simultaneously leverages the geometry of Wasserstein distance the one hand, and the favorable high-dimensional sample complexity and unbiased gradient estimates of MMD. We introduce the entropic regularized Wassertein distance
I(I, c(z, ((z.))d y) + /log (du())do(y)
min (EIl(
dII(2, v),
(12)
where e is the cost function in optimal transport. This objective function associated with cost func- tion e can be rewritten as We,,(,, = -f [( y)dlI, (2, y). Therefore, the sinkhorn loss between two measures u and v is defined as
We.(u,,) = 2We.cuu,
We,e u) We,e(u,
(13)
Theorem 4. If we leverage Sinkhorn loss the metric in distributional RL and choose c as unrec- tified kernel kan iee, Kin (x, y) = - || - "|la, for Vo E R. we have:
A1) O. Wec(u,)) - 2W, v), and thus T is a "-contraction.
(2) A2) 0o, War(u, MMD (u, v), and thus IT is 20/00/222o0 conttractit.nonn
Proof is provided in [ppendix]] inforreminndicttee that if we choose as the unrectified kernel, the limiting behaviors of distributional Bellman operator I are both contractive under Sinkhorn loss. Note that without the limitation, the entropic regularization restricts the search space the optimal
7

---

Under review as a conference paper at ICLR 2022
transport that eventually corresponds to a Gibbs kernel, based on which I may not be contraction. However, similar to MMD method, we show that our approach can still achieve empirical success and is very competitive across wide range of Atari games in Section
In the algorithm design, similar to MMD distri- butional RL (Nguyen al. 2020), we apply par- ticle representation to represent Za(s, a) by di- rectly generating multiple deterministic samples. Finally, we minimize the Sinkhorn loss between the approximate distribution via multiple sam- ples and its distributional Bellman target. A de- tailed description of algorithm is provided in Al- gorithm |
Algorithm 1 Generic Sinkhorn Algorithm
Require: Number of particles N, number of
Sinkhorn iteration L and hyperparameter Input: Sample transition (s, a, r's'')
1: if Policy evaluation then
2: a" ~ TT-Is').
3: else
Remark. MMD enjoys convergence rate of Onn-1/2) regardless of the underlying dimen- sion while I-Wasserstein distance has a conver- gence rate of O(n (//))/ if d > which is slower for large d. Thus, Sinkhorn loss has the potential to enjoy the faster convergence of MMD.
4: a" + arg max 1 N
Ew 7o (8,
5: end if
SisN Output: Waze ((Zo(s,a)))== (32,)
6: IZ + r 7 -=0-20-05,,"))
EXPERIMENTS
We demonstrate the effectiveness of Sinkhorn distributinal RL (SinkhornDRL) as described in Al- gorithm [ on a wide range of Atari 2600 games. Specifically, we leverage the same architecture as QR-DQN [Dabney al.] 20188) for fair comparison. More advanced techniques that can ex- pand the model expressiveness from IQN (Dabney et al. 2018a) and FQF (Yang et al.) 2019) can be naturally incorporated into our framework.
Baselines. Due to interpolation characteristic of SinkhornDRI, we choose 3 typical distributional RL algorithms as classic baselines, including QR-DQN (Dabney et al., (018h), C51 (Bellemare eet al. 2017a) and MMD (Nguyen et al., 2020). as well as DQN (Mnih et al. 20155). MMD algorithm implemented with the same architecture QRDQN, and leverages Gaussian kernels ky (2, y)
Breakout
Enduro
Pong
DSN Mell
3583
DSN
Average Retum
3000
ORTER
vsit
Gracion
Sinenom
Average
300
1200
0.3
0.2
0.4
8.6
0.8
Time Steps (1e7)
Time Steps (167)
Time Steps (167)
Inacers
Berzerk
BattleZone
CB1
1600
19800
2014
Average Rolum
NVDD
GRIONN
Average Retum
290
GROON
Average Return
20000
MMC
1250
204ce
1030
7550
433
12000
.00
10000
5150
Time Steps (107)
6.2
Time Steps (107)
8.4
0.5
0.0
0.2
Time
044
Steps (1G77
0.6
0.5
1.0
Figure 2: Performance of SinkhorDDRL on six Atari games.
8

---

Under review as a conference paper at ICLR 2022
exp(-(x /h) with the kernel mixture trick covering range of bandwidths h, which is same as the basic setting in the original MMD paper (Nguyen et al. 2020).
Hyperparameter settings. For fair comparison with QR-DQN, C51 and MMD, we used the same hyperparamters: the number of generated samples N 200, Adam optimizer with Ir 0.00005, AAdam = /30 We used a target network to compute the distributional Bellman target, which fits well in the neural Z-fitted iteration framework. In addition, we choose number of Sinkhorn iterations = 10 and smoothing hyperparameter E
10.0 in Section 7.1 as they are not sensitive
within proper interval demonstrated in Section 7.2 We choose
2 in Ka.
7.1 PERFORMANCE OF SINKHORNDRL
Figure 2 illustrates that SinkhornDR can achieve the state-of-the-art competitive performance on typical Atari games compared with various baseline algorithms (DQN, QR-DQN, C51, MMD) in different metrics dp and representation manners on Ze. Even though C51 outperforms others breakout game, it is significantly inferior on other games, e.g., Enduro, Pong and Berzerk. Moreover, SinkhornDRL significantly outperforms MMD breakout and Spacelnvader, and this superionity can be owing to the theoretical advantage of Sinkhorn loss over MMD. This coincides with the theoretical results as demonstrated in Theorem Sinkhorn loss interpolates between Wasserstein distance and MMD, which can simultaneously make full use of the data geometry from Wasserstein distance and the faster convergence, unbiased gradient estimates from Maximum Mean Discrepancy. We provide the competitive performance of SinkhornDRL on other Atari games in Appendix
7.2 SENSITIVITY ANALYSIS
We further examine the im- pact of different hyperpa- rameters in SinkhornDRL on the final performance, including the smoothing
Breakout
Breakout
1=1
210
Average Retum
400
280
hyperparameter e in Eq. [2
the number of determin- istic samples N in Algo- rithm [ From the left di- agram in Figure 33 we can observe that our algorithm not sensitive to the mag- nitude of long as within an appropriate inter- val, e.g., [, 100). Mean- while, it turns out that an overly large number of samples N can even slightly worsen the per- formance of SinkhorDRRL.
200
100
69
62
8.C Time Steps (167)
04
09
04
6.5 08
Time Steps (1o7)
(a) Hyper-parameter
(b) Number of Samples
Figure 3: Sensitivity analysis of SinkhornDRL on Breakout.
DISCUSSIONS AND CONCLUSION
Implicit generative models can be further incorporated into SinkhomDDL, including parameterizing the cost function in Sinkhorn loss, which be naturally expected to achieve more promising performance in the future. Moreover, a direct analysis on Wasserstein or l, distance can be more effective than KL divergence presented in this paper, albeit being theoretically tricky.
In this paper, we illuminate the superiority of distributional RL over expectation-based RL from the perspectives of regularization, optimization, acceleration and representation. In addition, novel family of distributional RL algorithms based on Sinkhorn loss is designed that accomplishes the promising performance on Atari games. Our analysis paves the way towards deeper understanding of distributional RL, and further promote its deployment in real applications.
Ethics Statement. Revealing the advantage of distributional RL would promote the application of distributional RI. algorithms in real scenarios. As distributional RL. enjoys the robustness against

---

Under review as a conference paper at ICLR 2022
abnormal events, e.g., noisy state observations, it can also be beneficial for the privacy of algorithms. Besides, the deeper insights into distributional RL plays a key role into the research integrity issue. Based on our knowledge, it is not related to harmful applications or faimess issues.
Reproducibility Statement. For the theoretical part, we clearly state the related assumption and detailed proof process in the appendix. terms of the algorithm implementation, our Sinkhorn algorithm is directly adapted from the public distributional RL algorithms, such as MMD.
REFERENCES
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference machine learning, pp. 211-2212222 PMLR, 2017.
Mare G Bellemare, Will Dabney, and Rémi Munos. distributional perspective on 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 solution to biased wasserstein gradi- ents. arXiv preprint arXiv:1705. 1074: 2017b.
Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. International Conference on Machine Learning (ICML), 8a.
Will Dabney, Mark Rowland, Marc G Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. Association for the Advancement of Artificial Intelligence (AAAI), 2018b.
Odin Elie and Charpentier Arthur. Dynamic Programming in Distributional Reinforcement Learn- ing. PhD thesis, Université du Québec à Montréal, 2020.
Jianqing Fan, Zhaoran Wang, Yuchen Xie, and Zhuoran Yang. A theoretical analysis of deep q- learning. In Learning for Dynamics and Control, pp. 486-489, PMLR, 2020.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn di- vergences. In International Conference on Artificial Intelligence and Statistics, pp. (608-1617) PMLR, 2018.
Tuomas Haarnoja, Haoran Tang. Pieter Abbeel, and Sergey Levine. Reinforcement learning with deep energy-basee basedd policies. 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: /0/2.008// 2018.
Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In International Conference on Machine Learning, pp. 1225-1234. PMLR, 2016.
Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge neural network. arXiv preprint 1rX3v:1503.02531, 2015.
Peter Huber. Robust Statistics, volume 523. John Wiley Sons, 2004.
Ehsan Imani and Martha White. Improving regression performance with distributional losses. In International Conference on Machine Learning, pp. 2167-2166. PMLR, 2018.
Clare Lyle, Mare 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, pp. 4511, 2019.
10

---

Under review as a conference paper at ICLR 2022
Clare Lyle, Mark Rowland, Georg Ostrovski, and Will Dabney. On the effect of auxiliary tasks representation dynamics. In International Conference on Artificial Intelligence and Statistics, pp- 1-9. PMLR, 2021.
Yecheng Jason Ma, Dinesh Jayaraman, and Osbert Bastani. Conservative offline distributional rein- forcement learning. arXiv preprint rrXiv:2107.06166 2021.
John Martin, Michal Lyskawinski, Xiaohu Li, and Brendan Englot. Stochastically dominant distribu- tional reinforcement learning. In International Conference on Machine Learning, pp. 6745-6754.
PMLR, 2020.
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.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare Belle- mare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 5187550))5529-33. 2015.
Tetsuro Morimura, Masashi Sugiyama, Hisashi Kashima, Hirotaka Hachiya, and Toshiyuki Tanaka. Parametric return density estimation for reinforcement learning. arXiv preprint arXiv: arXiv:/203.36: /20..3997 2012.
Rafael Müller, Simon Kormblith, and Geoffrey Hinton. When does label smoothing help? arXiv preprint arXiv:/906.0266 2019.
Thanh Tang Nguyen, Sunil Gupta, and Svetha Venkatesh. Distributional reinforcement learning with maximum mean discrepancy. Association for the Advancement of Artificial Intelligence (AAAI), 2020.
Mark Rowland, Marc 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 distributional reinforcement learning. International Conference on Machine Learning (ICML), 2019.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Richard Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. The American Mathematical Monthly, 74(4):402-400 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.00 08776, 2021.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An Introduction. MIT press, 2018.
Gábor J Székely. E-statistics: The energy of statistical samples. Bowling Green State University, Department of Mathematics and Statistics Technical Report, 3(05):1-18. 2003.
Christopher JCH Watkins and Peter Dayan. Q-learning. Machine learning, 8(3-4):279 292, 1992.
Ronald J Williams and Jing Peng. Function optimization using connectionist reinforcement Ieareeennnnnttnnnnnnnnn
algorithms. Connection Science, 3(3):241-268, 1991.
Eric Wong, Frank Schmidt, and Zico Kolter. Wasserstein adversarial examples via projected sinkhorn iterations. In International Conference on Machine Learning, pp. 6808-6817. PMLR, 2019.
Yi Xu, Yuanhong Xu, Qi Qian, Hao Li, and Rong Jin. Towards understanding label smoothing. arXiv preprint a006v:200.. /653, 2020.
11

---

Under review as a conference paper at ICLR 2022
Derek Yang, Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tie- Yan Liu. Fully parameterized quan- tile function for distributional reinforcement learning. Advances in neural information processing systems,
2019.
Fan Zhou, Jianing Wang, and Xingdong Feng. Non-crossing quantile regression for distributional reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020.
Florian Ziel. The energy distance for ensemble and scenario reduction. arXiv:2005.14670 2020.
arXiv preprint
12

---

dr Under review as a conference paper at ICLR 2022
A DEFINITION OF DISTANCES AND CONTRACTION
Definition of distances. Given two random variables X and Y, p- Wasserstein metric Wp between the distributions of X and Y is defined
1/D
W (X Y) - (o -1 dw) = =11Ex Fy'lls 1las
(14)
which F-1 is the inverse cumulative distribution function of random variable with the cumulative distribution function as F. Further, Ep distance (Elie Arthur 2020) is defined as
(x, = ( (Fx(W) - ( dw = Fx Fylle
1/p
(15)
The ép distance and Wassertein metric are identical at p = 1, but are otherwise distinct. Note that when p 2, Ep distance is also called Cramér distance (Bellemare al. 2017b) (x (Y,Y). Also, the Cramér distance has a different representation given by
1 1 -
dc dc(X,Y) =| Y
x'I E 1 -Yl,,,,
(16)
where X' and Y' the i.i.d. copies of X and Y. Energy distance (Székely 2003 Ziel 2020 is natural extension of Cramér distance to the multivariate case, which is defined as
1 Y Y'Il, dx(X, B - Y 1 2 - X
(17)
2
where X and Y are multivariate. Moreover, the energy distance is special case of the maximum mean discrepancy (MMD), which is formulated as
MMD(X, Y; (E (X, + - Y)
(18)
where k(---) is continuous kernel on X. In particular, if l is trivial kernel, MMD degener- ates to energy distance, Additionally, we further define the supreme MMD, which is functional P(X))xX P(X)SxA R defined as
(1,
sup (MD (x,a)eSxA
(19)
Proof of Contraction.
Contraction under supreme form of Wasserstein diatance is provided in Lemma (Belle- mare et al.) 2017a).
Contraction under supreme form of 4p distance can refer to Theorem 3.4 Elie & Arthur 2020
Contraction under MMD do is provided in Lemma 6 (Nguyen al. 2020
B PROOF OF PROPOSITION 1
Proof. (1) We recap three crucial properties of a divergence metric. The first is scale sensitive (S) (of order 8, beta > 0), i.e., (cX, (f.,,cY)) S (x, The second property is shift invariant (I), i.e, (A X,A + Y) < (), The last is unbiased gradient (U). We denote the density function of two random variables X and Y, and thus (X,
p and q to (X,Y) Y) is defined as DKL (X, Y) f p(x) p(z) dr. Firstly, we show that (XL(X, Y) is NOT scale sensitive: ((z)
1
de(())
Dxa(XX,
p
a L9(
p(y) p(y)
(20)
dy
((v)
DXL Y), with 3 =
13

---

Under review as a conference paper at ICLR 2022
We further show that (xL(X, Y) is shift invariant:
p(e AA p(z A) A)
9( A) dx
p(u) p(y) dy
(21)
g(y)
(X,)
Moreover, it is well-known that KL divergence has unbiased sample gradients (Bellemare et al. 2017b). The supreme DKL a functional ((X)SxA P(X)SxAA E defined
sup (R, (x.a)
Drr (u(,, a), v(2
(22)
Therefore, we prove I" is at best a non-expansive operator under the supreme form of DKL:
Zz) =sup DKL (I" (1 a), Z2(s,a))
si
DKL(R(ss a) + 7721(', R(s.a) + 7Z2('',
(23)
DK(( a) + (2(()))
sup DKL (R(s (35, stad
DEL (Z,, Zz)
23 There we have DRL ( Z2 S DRi (2), implying that I non-expansive operator under DRL.
(2) By the definition of DEL, we have sup,,, DKL (Z,
D (s, a), ((s, a)) 0 implies (Zn, 2)
0. DRr. 7) 0 implies the total variation distance 5(Zn, Z) according straightfor- ward application of Pinsker's s inequality
(Zn; 2) S V DKL (Zn
(24)
8 2,, S DKL (7, Zn,) +
Based on Theorem 2 in WGAN (Arjovsky al. 2017). 8(Zn= Z) 0 implies Wp (Z,, Z) - 0. This is trivial by recalling the fact that ò and W give the strong an weak topologies on the dual of (C(X), lloe) when restricted to Prob(X').
(3) The conclusion holds because the I" degenerates to T regardless of the metric dy Specifically, due to the linearity of expectation, obtain that
E 21 E Z2llac = -| EZ, This implies that the expectation of Z under DxL exponentially converges to the expectation of Z" ie., c-nntraccion..
T EZZl S EZ E22 zo.
(25)
C PROOF PROPOSITION 2
Proof. Firstly, given fixed p(x) we know that minimizing DKL qo) is equivalent to minimizing H(p, q) by following
(p.(p.90 go) = +oc
p(z) log p(x) dr (x(x)
p(x) log go (x) dr (- +00
p(x) logp(z) dx)
(26)
H(p, go) H(p)
x H(p, go)
14

---

Under review as a conference paper at ICLR 2022
Based on H(p. q0), can derive the objective function within each neural Z-fitted iteration
H(p (pprii00 (z), giaai (x))
7
p0t.0e(2) (og))oog@0 1o 40 (x) dr
-/ + (e(it((())) log up simi (xx) dar
-00
1000
11==3|2<0,00111 log "i (x) dæ + eHl(eria, "0, (araa))
(1 c)
(27)
[11 (000011122 ( + eH(p (e (..0)
((((==8|2=(( )..01) + +0 07H(11 fne, .00 _) where a 1 - > 0
x
log 90 n.. (B|Z"(ss a))) dr + aH((ptma "000) (0
12
D PROOF OF THEOREM
Proof. Firstly, have:
V (st+1) Bortina, 19111 [Q (S1+1, at+1) + log 4o
10-110 (x))
=E [Q (8141, ] + Bat 24
[log 90 10gg9/110000 11-1.0411 (x)
-Ea,11- ~~ 1Q (0441,0441)) - Eag1
10111101111 (x) log 10 au (x) dæ
-Roci 1Q (81+1, - Buyiow H (1o011-011 (1)))++1]]
(28)
7 then obtain Further, plug in V($+41) into RHS of the iteration in Eq.
TTQ a
=r(s;, a)) + 7E [V (8141))
r (st, ax)) (1B(o0410044)>)) [H (n1 (1*+41101-1 :40 ]111 + 7 (9441 1Q (8 1)]
A rx (st, + () 1Q ($141, 0411),
(29)
a.) where Tn (st, az) A r at) - 7I 15 ) H (1 110 (141)14331 the entropy aug- mented reward we redefine. Thus, repeatedly updating via the above equations is equivalent to
maximizing the new objective function J'()) = Er (ss,))))=p [r (st,,
H (μ Se
96 (.."))].
)]
Applying the standard convergence results for policy evaluation (Sutton & Barto 2018). we can attain that this Bellman updating under T convergent under the assumption bounded entropy augmented rewards r=-
[ 0o and
E DERIVATION OF HISTOGRAM DISTRIBUTIONAL Loss
We show the derivation details of the Histogram distribution loss starting from KL, divergence be- tween and yoo is the cumulative probability increment of target distribution (Y))e(m) within the
i-th bin, and go corresponds (normalized) histogram, and has density values ff
per bin.
ws
15

---

Under review as a conference paper at ICLR 2022
Thus, have:
DKL (p* .40 = -0
p (y) log To 2. (y)ddy
pl
f. (x(s)) dy
(1) log
1=1
(30)
f (x((s))
log
(F (4, + we) Fss0 (1,)
i=1
Pc
Pi 8.0 log f (x(s))
where the last equality holds because the width parameter wy can be ignored for this minimization problem.
F PROOF OF LEMMA 1
Proof. For the histogram distributional loss below,
exp (x(s) TO:)
Lo( = Σpd pr va log 12 (x(s)), where (x(s))
Σラ exp (x(s)))))
exp(x(s)TT.) = log Ei= exp (8;) we firstly prove its convexity. Note that - log Ej=n exp(x(*) 0,) x(s) f., 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.
Secondly, show that Ca(s, a) is kl- kl-Lipschitz. We compute the gradient of the Histogram distri- butional loss regarding 0;:
s. log (x((s))
Pj
09,
j=1
Σ Pj
3.0
1
(x.f((((3))
19 (x(s)
==l
k
=
Σ
Pj 5.0
89 (x(s))
1
1. (x(s)) - (x(s))))x((
(31)
1=1
Pi
(1 (x(ss))) - -PT" "
x(s)
j#i
(p p. J4 (x(s)) (1- p. /. )o (x(s)) x(s)
(p, (x(())) x(s)
16

---

Under review as a conference paper at ICLR 2022
where 6/ 1 i j, otherwise O. Then, have ||x(s)] || s L, we bound the norm of its gradient
a
pj log fl (x()))]
00
j=1
k
k
Pj log f3 (x()))]
00;
1=1
j=1
(32)
|| (0, 9. (x(s))) x(+)Il
a.a (xx))))]l
1=1
< kl
The last equality satisfies because [p; (x(s)) is less than I and even smaller. Therefore, we obtain that La kl-Lipschitz.
Lastly, we show that Lo kl2-Lipschitz smooth. A lemma that log(1 + exp(x)) I-smooth as its second-order gradient is bounded by 1. and if g(w) is 3-smooth w.r.t. then g (2, w)) is @||z:|2-smoooth Based on this knowledge, we firstly focus on the 1-dimensional case of function log Jf (2). As we have derived, we know that à log (z) dij - ff(). Then the second- 80;
8 log -8,00E
(Ō.R [[(=)) - (=) if = k. otherwise [2()). Clearly, S 1, which implies that log (z) is 1-smooth. Thus, log ((x 0.)) |smoosho, or -moohhh. Further, Σラ= D; log 19 (x(s)) also -mmsooh as we have
order gradient is
8 log (=)
00,30,
IV 0,
Pj
sa log 1 (n) Vo, sa log 13 (v))
Pj 3.4 VV log 1 (1) Vo, log 1og (v)
(33)
/=1
< Σ pj 12 - wll
/=1
= e [ill-vll w|l
for each y and v. Therefore, we further have
V
6
log (M) Vo
Σ
pj
log 1 (v)ll
/=1
1=1
k
k
k
Vo. 0,
Pj
log /9 (1)
Vo.
Dj
k.a log J (v)
(34)
1=1
/=1
==1
11 1| vll
kl2 ||
wil
Finally, we conclude that (o a) is k.l smooth.
17

---

Under review as a conference paper at ICLR 2022
G PROOF OF THEOREM
Proof. Consider the stochastic gradient descent rule as Gx.c(0) = - Firstly, we provide Iwo definitions about Lo for the following proof.
Definition 3. (a-bounded) An update rule is a-bounded if supa 110 XV a.
[G,,c()) |س-=i|
Definition 4. (1-expansive) An update rule is s--xpanssive if supe,w
S
Lemma 3. (Grow Recursion, Lemma 2.5 (Hardi et al. 2016)) Fix an arbitrary sequence of updates
G1 G..... GT and another sequence C',...,C', Grp. Let 90 0g be the starting point and define d, |10, Bell. where dt and fl are defined recursively through
Oct1 GGx,(00),,, 02+1 - (0))
Then we have the recurrence relation:
nne [min(n, 1)6, 20,
G: G2 is n-expansive
det1
Gt and G¹, are o-bounded GG is n expansive
Lemma 4. (Lipschitz Continuity) Assume La is L-Lipschitz, the gradient update Gx,c is (AL)- bounded.
Proof. |10 - Gx,c(0) (AVoC(0)] S AL
Lemma 5. (Lipschitz Smoothhsss)))))) Assume Ce is B-smooth, then for any X < 2. the gradient update Gx,c is I-expansive.
Assume
Proof. Please refer Lemma 3.7 in (Hardt al. 2016) for the proof.
Based on all the results above, we start to prove Theorem 2 Our proof is largely based on (Hard)
al. 2016), but is applicable in distributional RL setting well considering desirable proper- ties of histogram distributional loss. According Lemma [ we attain that Lo kl-Lipschitz as well as kl*-smooth, and thus based on Lemma 4 Gx,c is (AKl)-bounded, and 1-expansive if À S R7.. In the step f, SGD selects samples that are both in S and 5", with probability 1 n. In this case, G = -G,, and thus 6++1 I 6, as G2 is l-expansive based on Lemma 3 The other case is that sam- ples selected
different with probability ب where d,+1 S 6, + 27, based on Lemma Thus,
d S 2
we have:
E7TT
E [|((0r;)) 2(6))=2) < kE [8r]. where &y = 10 - 6/
2 -kk
< kl
(1
E 18r + E 18
+ 2Ar
kl
E E
2dekl
(35)
kl
E
[80]
+ Σ
n1
t=
2
T-1
2
71
h.2
4kT
Since this bounds hold for all S, S' and x, we attain the uniform stability in Definition for our histogram distributional loss applied in distributional RL.
Define the population risk as:
R [0] Fr C(O;
18

---

Under review as a conference paper at ICLR 2022
and the empirical risk as:
n [(; x4)
Rs [0]
Σ 72
According to Theorem 2.2 in (Hardt et al. 2016). if an algorithm M is Estab -aab-uniformly stable, then the generalization gap is 6sab -ounded, i.e.,
[Rs, [Rs(M[S)] R[M(S)] & stab
H PROOF OF LEMMA
((eaa)=pp V "0000) ()) VG(0)]²] s (1 (1-c)o² + EnO
(36)
Proof. As know that pao (1 (0=== (2) e.1 (x), then have:
Co 20 -00)) = (100==(((ooo(000-e -1 + cV Lo :10
Therefore,
B(ea)>pp [I]V Lo (p):a (0²)) - [G
I ((e.))=p= [11
Lo Co00 1) VG(0)]1+ VlV ,)0²)) -V((0)
(1 - c) CRO?
(37)
where the first inequality uses the triangle inequality of norm, i.e., II(1 t)a + ebll2 < (1
l + and the last equality uses the definition of the variance of Lo (0(6 Fs2* 98 2.0
+ e|bb/
and Ca (xx 46
5.a
I PROOF OF THEOREM 3
Proof. (1) If we only consider the expectation of (*,s a), the entropyy-regularieed MLE would degenerate to the pure MLE regarding 0[==2]2 (sail) As Lo (=(00)=)127(.00)) (00)) ) is kmo-smooth, we have
G(0++1) G(O,)
kt 1 0,1
(VG(()), 0c) +
2
(38)
([ =-A((GG00,VVo00=====z (s,a)) )) x2
(0 (x.o))) 4 112
where the last first equation is according to the definition of Lipschitz-smoothhe smoothness, and the last second one is based the updating rule of f. Next, we take the expectation both sides,
[G((0++1 G(O,)]
AE ([V(0))))
EE [[
96 ") VG(O.)
VG(()++V(())]
x2
VG(0,) + kt2x
-E -A
[[G
2
A(KPA - 2)
xl2X2
E
|
+ 2
0²
2
A [IVG(()11
x
2
(39)
where the first two equation hold because VG(0) E [VCp] and the last incquality comes from A < 1/2 Through the summation, we obtain that
B[G(Or) - G(O)) < - Σ
XPXTT
[IVGG(0))]
0
2
t=o
19

---

Under review as a conference paper at ICLR 2022
We let [G(rr)] 0, we have
2G((0) E [IVGG()))
XI + =kl2
By setting S 2
2kPo and 7 40100)
have E - 72, implying that
Ar
the degenerated MLE can achieve s-station point the sample complexity T O..).
(2) and (3) We are still based on the smoothhthsssss of L(p"' "400 %
G(((+1)- G(O,)
122
I (VG(0.), 0.+1 - 0,) + 2 -100+1 - 0ell
2t
- VLolD"M -)00))
\V .90 "g00)")) )11
(40)
2
2 1VVC(0))1 A -2
A(kI2>
VLo(p"n .90 -012)
V Lo . )|
A -- - 1 IVG(())) VLo(p"10 m12
where the second equation based on (a, -b) = 1/2 (Iaa 1 bl1² |a|2 2 [bb| and the last in- equality is according to À S में After taking the expectation, we have
E [G(00+1) G(O,)
A = [[V))()) - (Co (0
5 hーク [G((0))) +
(41)
X
S - [|VG(
c)= + eno?)
where the last inequality is based on Lemma
We take the summation, and therefore,
E [G((r) - G(6o)) A T-1
E [[G()) + T>
((1 (1=c)o²² - ekO?)
2
We let [G(rr) = 0 and 1 then,
1+x*
T-1
E [I
T t=
2G((o)
+1 - c)o² + eRo²
XT
(42)
2G(Ou)
2k 0²
+ 1 K
XI
2G((o)
2na
XT
E1= D-A [IVVG(0))]]]] < ie, r-stationary point, with the sample complexity as 0 Thus, (2) has been proved. On the other hand, if > |Z. we set T G(%,
If +2 T0z and let T 4G(0g) this leads
Axoi This implies that T T-1 E [| [2| s 4ko² = 0 Therefore, the degree of stationary point is determined the degree of distribution approximation measured by k. Thus, obtain (3).
J PROOF OF THEOREM 4
Proof. 1. As E - 0, it is obvious to observe that Sinkhomn loss degenerates to the wasserstein distance.
20

---

Under review as a conference paper at ICLR 2022
dil(r,)) y). We can consider its dual problem by introducing Lagrange multipliers and By additionally considering the primal-dual relation, we can solve the dual problem, which gives v = 0, and thus the optimal coupling is simply the product of the marginals ie., πT
As - 0o, Sinkhorn loss turns to be f log du(z))v(z)
T(a.y)
Please refer more detailed proof in Sinkhorn-GAI (Genevay et al. 2018). Since the cost function is the rectified kernel ka, under with MMD (guyen et al. 2020) can be -onttcontr
K
REPRESENTATION EFFECT OF DISTRIBUTION RL
From the perspective of representation, we find that distributional RL encourages state representa- tion from the same action class classified by the policy in tighter clusters.
The intrinsic characteristic of distri- butional RI. is that it enables to learn a richer and more faithful represen- tation of the environment, leading to more stable and efficient learning. To investigate the more informative rep- resentation resulting from distribu- tional RL, we visualize how distribu- tional RI changes the representation learned in the penultimate layer the value network. We collect state fea- tures the penultimate layer classified by the policy π different action classes, and perform t-SNE to reduce them to the two-dimensional space. In Figure points in different colors represent state features classified into different action classes. It illustrates that both QR- QR-DQN and C51 encourage the representation of state observations from the same action class group tighter clusters rela- tive to expectation-basel DQN. Interestingly, this phenomenon is similar to label smoothing (Müller al. 2019). which leverages additional distributional knowledge in soft labels relative to hard la- bels. Therefore, we conclude that similar to the benefit of label smoothing, distributional RL also encourages the activations of the penultimate layer be close the template of the correct action class and distant to the templates of the incorrect action classes.
(a) DQN
(b) C51
(c) QRDQN
Figure 4: Visualization on the penultimate layer of the value network Breakout, Each color denotes one action class.
L MORE EXPERIMENTAL RESULTS
21

---

Under review as a conference paper at ICLR 2022
Revenge
Alien
DQN
2250
DON
Average Retur
20000
QRDQN
Average Return
CS1
1750
QRDON
Sinkhom
1500
MMD
1250
Sinkhom
10000
1000
500
0.0
0.2
0.4
0.6
0.8
0.4
1.0
Time Steps (1e7)
Time Steps (1e7)
UpNDown
Qbert
25000
16000
CS1
Average Return
QRDON
Average Return
CS1
20000
MMD
12000
QRDON
Strkhom
15000
Sinkhom
8000
4000
0.0
0.2
0.4
0.6
0.8
1.0
Time Steps (1e7)
Time Steps (1e7)
Figure 5: Performance of SinkhomDDL on YarRevenge, Aline, Qbert and UpNDown.
22


Reviews:
[{"review_id": "yKgkOtWusk", "reviewer": "Reviewer_Q4X1", "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": "Pros:\t\n1. I think the general question this paper studies is important and also underexplored. A deeper understanding of the advantages of DRL will be beneficial to both methodological and applied research related to DRL. There have been few previous studies trying to formalize these advantages.\n2. This paper proves some theoretical advantages of distributional RL from some new perspectives, such as optimization and sample efficiency. These, to my knowledge, are quite novel compared with previous work.\n\n\nConcerns:\n1. Section 4 and 5 show the theoretical advantages of distributional RL in terms of optimization stability, generalization, and sample efficiency. However, most results are based on one or more assumptions not used in popular distribution RL algorithms, such as assuming the optimizer as SGD with a certain learning rate, using the KL divergence as the distance measure. These assumptions are fine for theoretical analysis. However, to convince people that these theoretical results are significant, the paper should include some empirical evidence that these advantages do exist in practical distributional RL algorithms without these assumptions. Otherwise, it is not obvious whether these advantages come more from the distributional RL algorithm or these artificial assumptions.  \n2. The conclusion in Section 6, that distributional RL learns a better state representation, is weak and not well-supported. Figure 2 cannot be used to show that distributional RL does a better job in clustering representations than expectation-based RL. In fact, from humans’ eyes, all subfigures a, b, and c are quite similar in terms of clustering performance. Moreover, it is not sufficient to draw the strong conclusion at the end of Section 6 from just one visualization on one environment. More empirical and theoretical analysis are needed to make this claim.\n3. Section 5 argues that distributional RL is more sample efficient than expectation RL. I find this result a bit counter-intuitive. Why does learning a distribution require fewer samples than its expectation? Also, Section 3 shows that distributional RL encourages exploration, which should lead to more samples required. Some intuitive clarifications here would help.\n4. I cannot follow the proof of Theorem 4. It is difficult to see how the second case is related to MMD. \n5. The main argument in Section 7, “Sinkhorn loss leverages the advantages of MMD and Wasserstein distance”, is not convincing. The paper shows two extreme cases (epsilon = 0, \\infty) of Sinkhorn loss lead to MMD and Wasserstein distance. However, that does not automatically mean that the in-between cases (when epsilon is a positive number) will inherit the advantages of the two methods. In fact, they may get the disadvantages of MMD and Wasserstein distance. More empirical and theoretical analysis are needed to make this claim.\n6. The paper only shows the results of 10 Atari games including the appendix. However, the Atari benchmark consists of 57 games, and each of them has very different environments and difficulties. The convention in the RL community is to do evaluations on all 57 Atari games, which is also the case for the previous work in distributional RL, such as, C51 and MMD. Also, there is no explanation for how these 10 games are selected from the benchmark. This makes these empirical results less trustworthy.\n7. There seems to be a discrepancy between the results presented in the main paper (Figure 3) and the appendix (Figure 5). All games where the proposed method performs poorly are included in the appendix rather than the main paper. This selective presentation could mislead readers who haven’t read the appendix.\n\nMinor suggestions\n1. Overall, the paper can feel like some disconnected pieces glued together. The beginning of Section 7 says “Based on the theoretical analysis, we further design Sinkhorn distributional RL algorithm”. However, I could not see the connection between the proposed loss and the previous theoretical analysis. It would be good to point out these connections explicitly in writing, which will make the paper seem better-motivated and more pleasant to read. \n2. I found some figures a bit confusing and found it difficult to know their purpose, such as Figure 1. I think it would help a bit to make the caption of each figure more detailed and self-contained.\n", "summary_of_the_paper": "Distributional reinforcement learning (DRL) is a family of RL algorithms that estimates distributions of the value function rather than the expectation. The first part of the paper interprets distributional RL as adding a cross-entropy regularizer to the traditional objective. The second part discusses several advantages of DRL over expectation-based RL, including more stable optimization, better sample efficiency, and better state representations. The last part proposes a new distributional RL algorithm based on Sinkhorn loss, aiming to leverage the nice geometric interpretation of Wasserstein distance and sample complexity of Maximum Mean Discrepancy.", "summary_of_the_review": "Overall, the questions that this paper wants to address are important, but the answers provided are not very satisfactory. The significance of most theoretical results in the paper is questionable due to the discrepancy with practical distributional RL algorithms (for example, most distributional RL do not use KL divergence to measure distance) and lack of empirical evidence to back up (such as Section 3, 4, 5). Some major claims are not well-supported by the rigorous evaluations (such as Section 6) and clear explanations (such as Section 4 and 5). The results section (Section 8) for the proposed method is not very trustworthy due to cherry-picking the results and lack of important details. \n", "strength_and_weaknesses": "", "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": "Not applicable"}, "extra_text": {}}, {"review_id": "uh7FnXHrBQd", "reviewer": "Reviewer_Qz4z", "summary": "", "questions": "", "limitations": "", "rating": 1, "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": "I don’t find any clear points or guarantees for distributional reinforcement learning. The authors provide some claims, but unfortunately, all of these claims do not directly relate to the benefits of distributional reinforcement learning. \n\nDetailed Comments:\n\n* On the regularization part: \n  * Please note that, the decomposition of (4) is non-sense. I don’t find any interpretation of such results. Specifically why shall we separate the event $x = \\mathbb{E}[Z^\\pi(s, a)]$? After this decomposition, $F^{s, a}(x)$ is not the distribution of the value function any more, how should I parse the remaining results? Please also notice that when the value support on finite set that does not contain $\\mathbb{E}[Z^\\pi(s, a)]$, to make $F$ is still a valid probability measure, we can only let $\\epsilon = 1$, which makes the remaining argument nonsense, as $\\alpha = \\frac{\\epsilon}{1-\\epsilon}$. Even the value support on a continuous set, after you make the claimed decomposition, the distribution significantly changes, as all of the events other than $x = \\mathbb{E}[Z^\\pi(s, a)]$ is down-scaled, which I don’t find any clear interpretation on that. Hence, I don’t think any of the claim for the regularization part makes any sense.\n  * For Theorem 1, how can we know the exact $\\mu(s^{t+1}, a^{t+1})$? If not, what’s the purpose of this theorem? Also, what’s the role of $\\theta$ here? Are we aiming at the policy induced by $\\theta$? Even the $Q$ sequence converge to the solution of that objective, what’s the theoretical property of the solution? When will it be used? What’s the theoretical benefit of using that? I have no idea on any of these questions.\n\n* On the optimization and generalization part:\n  * How do we partite the support of the state feature? Will it affect the accuracy? As each partition will lead to a totally different output probability, if I understand correctly?\n  * The parameterization of $f$ doesn’t seem to partite the feature of $x$. It seems like perform a classification on the state, which means the authors partite the output rather than the input. Also, where are the actions take the role?\n  * I don’t think properly match the density of each bin will provide a good estimation on the distribution, and hence I think if we minimize (8), we are not guaranteed to find the distribution of value.\n  * I don’t check the proof of Theorem 2, but even it’s correct, it doesn’t mean anything. As I said, even we can perfectly match the probability on each bin, we don’t have any guarantees on the original distributional reinforcement learning problem. Can the authors state the guarantee of the original performance guarantee? For fitted Q-iteration, at least we can guarantee we can find the exact Q function under proper regularity condition. Also, the proof highly depends on the specific parameterization, I don’t think the authors can say such results can be held with neural network approximation. It’s never a ``neural’’algorithm.\n\n* On the acceleration part:\n  * As I said in the regularization part, we cannot set $\\epsilon$, or otherwise we are not really doing distributional reinforcement learning. \n  * Why should we parameterize $q$ in this way, if we only want the expectation? Theorem 3 (1) doesn’t make any sense to me.\n  * As the authors don’t provide the bound for $\\kappa$, and results of Theorem 3(3) is not satisfactory, I don’t get any information the authors want to argue.\n\n* On the Sinkhorn part:\n  * There are so many works for entropy regularized optimal transport, so I don’t think the loss is novel. Meanwhile, I don’t see any motivation on use this.\n\nIn summary, the authors never provide clear problem setup for each individual points and never provide rigorous theory on the benefits of distributional reinforcement learning compared with expectation based reinforcement learning. Instead, the authors only list several unrelated thing from other work and try to connect them with distributional reinforcement learning.", "summary_of_the_paper": "The authors want to provide some new perspectives on distributional RL from the view of regularization, optimization, acceleration and eventually propose a class of Sinkhorn algorithm for distributional RL.", "summary_of_the_review": "I don’t get any points the authors want to show, and the authors don’t provide any meaningful theoretical guarantees on the distributional reinforcement learning. The authors only list several unrelated thing and try to make them connected with distributional reinforcement learning, without going any real problems of distributional reinforcement learning.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "1: The main claims of the paper are incorrect or not at all supported by theory or empirical results.", "technical_novelty_and_significance": "1: The contributions are neither significant nor novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "kHGt4JTx5_", "reviewer": "Reviewer_r4R2", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "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": "**Novelty**: The paper presents several interesting insights including the interpretation of distributional RL as entropy-regularized MLE and the finding of the representation of distributional RL that is different from that of expected RL. However, the other aspects including stability, optimization, generalization analysis, and Sinkhorn loss are quite straightforward applications from existing works. Details are as follows: \n- The analysis in Section 4 is not significant nor relevant much to distributional RL. In particular, the paper considers a generalized linear model with softmax as the link function and analyzes the stability, optimization, and generalization in the supervised learning setting of such a generalized linear model with cross-entropy loss Eq. (8) which are quite a straightforward result from the literature. In addition, due to such simplification, their analysis's connection to distributional RL is weak. \n- In RL, we often consider bounded reward, so if we consider the expected RL with squared loss function, it still enjoys all the properties of KL loss arising from distributional RL that is discussed in Section 4. Thus, the analysis in Section 4 does not show a significant advantage of distributional RL over expected RL. \n\n**Clarity**: The paper is well structured though some paragraphs need to be improved for clarity. The paper studies several aspects of distributional RL (interpretation, optimization, generalization, representation, new algorithm) which are disconnected from each other. It might be clearer if the paper can focus on one aspect and go deeper analysis for that aspect instead of spreading out many disconnected ones. For example, with the proposed interpretation of distributional RL as entropy regularization, though being a nice interpretation, I am not sure what should we do with this fact to improve distributional RL? The paper can investigate deeply along this direction.\n\n**Empirical significant**: The proposed Sinkhorn distributional RL performs quite on par with MMD without a clear-cut difference on the selected set of games. Moreover, the paper only presents on a small set of games. These make it hard to conclude whether there is empirical advantages in the proposed algorithm. \n\n**Minor comments**: \n- In Eq. (5), $\\alpha$ should be explicitly written in terms of $\\epsilon$\n- The paragraph in Section 4 could be written better \n- Eq (8) is not correct; it is an approximation with an extra entropy term \n- Generalization guarantee is a consequence of uniform stability; thus Theorem 2 could be made less complicated in its statement. \n- \"our algorithm is not sensitive to the magnitude of $\\epsilon$ as long as ...\" in Section 8.2: not sure if I should interpret it as an advantage of the algorithm or a signal that something might have been wrong in your algorithm or implementation. It should be sensitive to $\\epsilon$", "summary_of_the_paper": "This paper interprets distributional RL as entropy-regularized MLE, explains distributional RL from the perspective of uniform stability, gives an insight into the representation of distributional RL as compared to that of expected RL, and proposes a Sinkhorn distributional RL algorithm that interpolates MMD and Wasserstein loss. ", "summary_of_the_review": "Though the paper proposes a nice interpretation of distributional RL as entropy regularization and the practical property of the representation of distributional RL, the other contributions seem do not significant or meaningful to me. However, the paper could be improved with the suggestions above as the idea presented in the paper is interesting and promising if it could be investigated more deeply. I feel the paper is not yet ready for publication at the current form, and I vote for rejection. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": "fGFufnL5pxG", "reviewer": "Reviewer_GhjJ", "summary": "", "questions": "", "limitations": "", "rating": 5, "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": "Generally speaking, the paper is poorly written; many parts were hard to follow, but more importantly, the different sections feel very disconnected. For example:\n\n- Eq. 4 is poorly explained; the term F_\\mu is not formally defined\n- Section 3.3 was hard to follow; it wasn’t clear exactly what connection is being drawn\n- Section 6 is very disconnected; there is no theory here but a small experiment, and the connection to the other results is unclear.\n- Section 7 claims that the Sinkhorn approach is motivated by the theoretical analysis in the previous sections, but I don’t see why this is the case.\n\nThe paper also does not give sufficient background for the Sinkhorn loss\n\nThe paper currently feels like a disjoint collection of results, none of which are properly motivated or connected to other results. I believe the authors should focus on one or two key results and describe them and their importance in more detail, rather than try to present so many results without sufficient explanation.\n\nAlong the same lines, the significance of the results is unclear. For example, consider Theorem 2 in Section 4. How does this result differ from a corresponding result in the non-distributional setting (or to a different algorithm)? It would be significantly more helpful if such a contrast is drawn, since it would help clarify what is special about distributional reinforcement learning compared to traditional reinforcement learning. A similar comment goes for Theorem 3 in Section 5; this theorem had a bit more discussion, but it was hard to follow due to the conciseness.\n\nFinally, the use of KL divergence in the objective is non-standard; a more typical approach (as the authors mention) is to use the Wasserstein or Cramer distances. Do the authors expect their insights to generalize to other objectives? A lot of their results appear to hinge pretty strongly on this formulation (e.g., the connection to maximum entropy RL).\n", "summary_of_the_paper": "This paper performs a theoretical analysis of distributional RL. Furthermore, they propose a new algorithm based on the Sinkhorn loss, and empirically demonstrate that their algorithm outperforms existing approaches.\n", "summary_of_the_review": "Pros\n- Important problem\n\nCons\n- Poor writing\n- Disjointed and poorly motivated results\n- Non-standard formulation\n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and 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": {}}]

Author rebuttal:

Meta-review: It appears that the reviewers have reached a consensus that the paper is not ready for publication at ICLR.
Decision: Reject