Year: 2021

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 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. Finally, we propose a 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.

heading: 1 INTRODUCTION
text: The intrinsic characteristics of classical reinforcement learning (RL) algorithms, such as temporaldifference (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-based RL, a new branch of algorithms called distributional RL estimates the full distribution of total returns and has demonstrated stateof-the-art performance in a wide range of environments (Bellemare et al., 2017a; Dabney et al., 2018b;b; 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 et al., 2019; Rowland et al., 2019) and robsutness (Sun et al., 2021).
Despite the existence of numerous algorithmic variants of distributional RL with remarkable empirical success, theoretical studies of advantages of distributional RL over expectation-based 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-based RL. Lyle et al. (2021) investigated the impact of distributional RL from the perspective of representation dynamics. Martin et al. (2020) recently mapped distributional RL problems to a Wasserstein gradient flow problem, treating the distributional Bellman residual as a potential energy functional. Offline distributional RL (Ma et al., 2021) has also been proposed to investigate the efficacy of distributional RL in both risk-neutral and risk-averse domains. Recent works have tended towards closing the gap between theory and practice in distributional RL.
From an algorithmic perspective, the Sinkhorn loss (Sinkhorn, 1967) can be used to tractably approximate the Wasserstein distance and has been successfully applied in numerous crucial machine learning developments, including the Sinkhorn-GAN (Genevay et al., 2018) and Sinkhorn-based adversarial training (Wong et al., 2019). Inspired by the distributional RL literature, Martin et al.
(2020) argues for the use of a second-order stochastic dominance relation to select among a multiplicity of competing solutions via Sinkhorn iteration (Sinkhorn, 1967), which can be useful to manage stochastic uncertainty in RL paradigms. However, a Sinkhorn-based distributional RL algorithm has not yet to be formally proposed and investigated.
In this paper, we theoretically illuminate the superiority of distributional RL over expectation-based RL from the perspectives of regularization, optimization, acceleration, and representation. Specifically, we simplify distributional RL as a neural Z-fitted iteration, within which we establish an equivalence between distributional RL and a form of entropy-regularized maximum likelihood estimation (MLE). We also demonstrate that the resulting novel cross entropy regularization correlates strongly with the behavior of maximum entropy RL. By incorporating a histogram distributional loss, we further achieve stable optimization and guaranteed generalization of distributional RL, attributable to desirable smoothness properties of the distribution loss. We further characterize the effect of acceleration on distributional RL and discuss when distributional RL algorithms are effective 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 summary, our analysis opens the door to a deeper understanding of theoretical advantage of distributional RL.

heading: 2 PRELIMINARY KNOWLEDGE
text: In the standard RL setting, an agent interacts with an environment via a Markov decision process (MDP), a 5-tuple (S,A, R, P, γ), where S and A are the state and action spaces, respectively. P is the environment transition dynamics, R is the reward function and γ ∈ (0, 1) is the discount factor.
State-value function vs. state-value distribution. Given a policy π, the discounted sum of future rewards is a random variable Zπ(s, a) = ∑∞ t=0 γ
tR(st, at), where s0 = s, a0 = a, st+1 ∼ P (·|st, at), and at ∼ π(·|st). In the control setting, expectation-based RL focuses on the state-value function Qπ(s, a), which is the expectation of Zπ(s, a), i.e., Qπ(s, a) = E [Zπ(s, a)]. Distributional RL, on the other hand, focuses on the state-value distribution, the full distribution of Zπ(s, a). Leveraging knowledge on the entire distribution can better capture the uncertainty of returns in the MDP beyond the expectation of return (Dabney et al., 2018a; Mavrin et al., 2019).
Bellman operators vs. distributional Bellman operators. For the policy evaluation in expectation-based RL, the value function is updated via the Bellman operator T πQ(s, a) = E[R(s, a)] + γEs′∼p,π [Q (s′, a′)]. In distributional RL, the state-value distribution of Zπ(s, a) is updated via the distributional Bellman operator Tπ
TπZ(s, a) = R(s, a) + γZ (s′, a′) , (1)
where s′ ∼ P (·|s, a) and a′ ∼ π (·|s′). 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 Tπ is contractive under certain distribution divergence metrics, but the distributional Bellman optimality operator T can only converge to a set of optimal non-stationary value distributions in a weak sense (Elie & Arthur, 2020).

heading: 3 EFFECT OF REGULARIZATION ON DISTRIBUTIONAL RL
text: 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 expectationbased RL has been less studied. In this section, we attribute the superiority of distributional RL into its regularization effect.

heading: 3.1 DISTRIBUTIONAL RL: NEURAL Z-FITTED ITERATION
text: Neural Q-Fitted Iteration. In the function approximation setting, Deep Q Learning (Mnih et al., 2015) can be simplified into Neural Q-Fitted Iteration (Fan et al., 2020) under tricks of experience replay and the target network Qθ∗ , where we update parameterized Qθ(s, a) in each iteration k:
Qk+1θ = argmin Qθ
1
n n∑ i=1 [ yi −Qkθ (si, ai) ]2 , (2)
where the target yi = r(si, ai) + γmaxa∈AQkθ∗ (s ′ i, a) is fixed within every Ttarget steps to update target network Qθ∗ by letting θ∗ = θ and the experience buffer induces independent samples {(si, ai, ri, s′i)}i∈[n]. In an ideal case that neglects the non-convexity and TD approximation errors, we have Qk+1θ = T Qkθ , which is exactly the updating under Bellman optimality operator. Under the two conditions, the optimization problem in Eq. 2 can be viewed as Least Square Estimation (LSE) in a 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 on a parameterized Zθ into a Neural Z-fitted Iteration as
Zk+1θ = argmin Zθ
1
n n∑ i=1 dp(Yi, Z k θ (si, ai)), (3)
where the target Yi = R(si, ai) + γZkθ∗ (s ′ i, πZ(s ′)) with πZ(s′) = argmaxa′ E [ Zkθ∗(s ′, a′) ]
is fixed within every Ttarget steps to update target network Zθ∗ , and dp is a divergence metric between two distributions. Notably, the options of representation manner on Zθ and the metric dp are pivotal for the empirical success of distributional RL algorithms. For instance, QR-DQN (Dabney et al., 2018b) approximates Wasserstein distance Wp, which leverages quantiles to represent the distribution of Zθ. C51 (Bellemare et al., 2017a) represents Zθ via a categorical distribution under the convergence of Cramér distance (Bellemare et al., 2017b; Rowland et al., 2018), a special case with p = 2 of the `p distance (Elie & Arthur, 2020), while Moment Matching (Nguyen et al., 2020) learns deterministic samples to represent the distribution of Zθ based on Maximum Mean Discrepancy (MMD). Contractive properties under typical metrics dp can be summarized as follows
• T π is γ-contractive under the supreme form of Wassertein distance Wp. • T π is γ1/p-contractive under the supreme form of `p distance. • T π is γα/2-contractive under MMD∞ with the kernel kα(x, y) = −‖x− y‖α,∀α ∈ R.
For the completeness, the definition of mentioned distances and the proof of contraction are provided in Appendix A. Although the widely used Kullback–Leibler (KL) divergence is not a contraction (Morimura et al., 2012), we show in Proposition 1 that the KL divergence still enjoys desirable properties in distributional RL context, which can be reasonable for the theoretical analysis. We assume Zθ is absolutely continuous and has joint supports, under which the KL divergence is welldefined. Proof of Proposition 1 and the definition of supreme DKL are provided in Appendix B. Proposition 1. Denote the supreme of DKL as D∞KL, we have: (1) Tπ is a non-expansive operator underD∞KL, i.e., D ∞ KL(T πZ1,T πZ2) ≤ D∞KL(Z1, Z2), (2)D∞KL(Zn, Z)→ 0 impliesWp(Zn, Z)→ 0, (3) the expectation of Zπ is still γ-contractive, i.e., ‖ETπZ1 − ETπZ2‖∞ ≤ γ ‖EZ1 − EZ2‖∞.

heading: 3.2 DISTRIBUTIONAL RL: A NOVEL ENTROPY-REGULARIZED MLE
text: The reasonable properties of KL divergence in Proposition 1 allows us to leverage it to conduct the theoretical analysis. To separate the impact of additional distribution information from the expectation of Zπ , we leverage the variant technique of gross error model from robust statistics (Huber, 2004), similar to the technique to analyze Label Smoothing (Müller et al., 2019) and Knowledge Distillation (Hinton et al., 2015). Specifically, we denote the one-dimensional full distribution of Zπ as F , and the distribution on the remaining support getting rid of E [Zπ] as Fµ. Hence, we can obtain the distribution decomposition for Zπ(s, a) as
F s,a(x) = (1− )1{x≥E[Zπ(s,a)]}(x) + F s,aµ (x), (4)
where controls the proportion of F s,aµ (x) and the indicator function 1{x≥E[Zπ(s,a)]} = 1 if x ≥ E [Zπ(s, a)], otherwise 0. After taking derivatives on both sides, we attain the relationship of their density functions as ps,a(x) = (1 − )δ{x=E[Zπ(s,a)]}(x) + µs,a(x), where µs,a(x) is the density function related to Zπ(s, a) on remaining supports removing E [Zπ(s, a)]. It is worth noting that the existence of µs,a(x) can be simply guaranteed by directly computing µs,a(x) = ps,a(x)/ − (1 − )δ{x=E[Zπ(s,a)]}/ as long as ps,a(x) and the expectation of Zπ(s, a) exist. Next, we use ps,a(x) and qs,aθ (x) to denote the density distributions behind {Yi}i∈[n] and Zkθ (s, a) in neural Zfitted iteration via Eq. 3, respectively. Therefore, we can derive the following result in Proposition 2. Proposition 2. Let H(P,Q) as the cross entropy, i.e., H(P,Q) = − ∫ x∈X P (x) logQ(x) dx. Let α be a positive constant, and based on the decomposition in Eq. 4 and DKL as dp, Neural Z-fitted iteration in Eq. 3 can be reformulated as
Zk+1θ = argmin Zθ
1
n n∑ i=1 H(δ{x=E[Zπ(si,ai)]}, q si,ai θ ) + αH(µ si,ai , qsi,aiθ ). (5)
We provide the proof in Appendix C. For the uniformity of notation, we still use s, a in the following analysis instead of si, ai in Eq. 5. Importantly, the first term in Eq. 5 can be further simplified as − ∫ x∈X log q s,a θ (E [Z(s, a)]). Minimizing this first term can be viewed as a variant of Maximum Likelihood Estimation (MLE) on the expectation E [Z(s, a)] rather than the traditional MLE directly on observed samples. The cross entropy regularization in the second term pushes qs,aθ to approximate the distribution µs,a 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 µs,a and qs,aθ is different from the classical entropy regularization used in RL, which we further analyze their connection and discrepancy in Section 3.3. In summary, distributional RL can be simplified as a novel entropy-regularized MLE within neural Z-fitted iteration framework in stark contrast to the Least-Square estimation of expectation-based RL in the neural Q-fitted iteration.

heading: 3.3 CONNECTION WITH MAXIMUM ENTROPY RL
text: We establish the connection between the derived novel cross entropy regularization in Eq. 5 in distributional RL with the classical maximum entropy RL (Williams & Peng, 1991). Maximum entropy RL, including Soft Q-Learning (Haarnoja et al., 2017), greedily maximizes the entropy of the policy π(·|s) in each state:
J(π) = T∑ t=0 E(st,at)∼ρπ [r (st, at) + βH(π(·|st))] , (6)
where H (πθ (·|st)) = − ∑ a πθ (a|st) log πθ (a|st) and ρπ is the generated distribution following π. The temperature parameter β determines the relative importance of the entropy term against the cumulative rewards, and thus controls the stochasticity of the optimal policy. This maximum entropy regularization has various conceptual and practical advantages. Firstly, the policy encourages exploration, avoiding situations in which the agent might fall into a local optimum. Secondly, it considerably improves learning speed over classical methods and therefore are widely used in stateof-the-art algorithms, e.g., Soft Actor-Critic (SAC) (Haarnoja et al., 2018) and PPO (Schulman et al., 2017). Lastly, maximum entropy can yield more robustness to abnormal or rare events while developing a task. Similar robustness superiority of distributional RL against noisy state observations has been investigated in (Sun et al., 2021).
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 in SAC 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 πQ (st, at) , r (st, at) + γEst+1∼ρπ [V (st+1)], where V (st+1) = Eat+1∼π [Q (st+1, at+1)− log π (at+1|st+1)] is the soft value function (Haarnoja et al., 2018). Interestingly, we still have the similar convergence result related to our novel cross entropy regularization in Eq. 5, which is shown in Theorem 1: Theorem 1. Consider the soft Bellman operator T πd in Eq. 7 and a mapping Q0 : S × A → R with |A| ≤ ∞, and define Qk+1 = T πd Qk. Given the true distribution µst,at for each t, then the
sequence Qk will converge to a soft Q-value of π as k →∞ with the new entropy objective function as J ′(θ) = ∑T t=0 E(st,at)∼ρπ [r (st, at)− γH (µst,at , q st,at θ )]:
T πd Q (st, at) , r (st, at) + γEst+1∼ρπ [V (st+1)] , (7) where V (st+1) = Eat+1∼π,x∼µst+1,at+1 [ Q (st+1, at+1) + log q st+1,at+1 θ (x) ] and the policy π follows the classical greedy policy rule, i.e., π(·|s) = arg maxa E [Zθ(s, a)], determined by θ.
Please refer to Appendix D for the proof. The new entropy augmented objective function J ′(θ) in Theorem 1 reveals that different from adding the entropy regarding the policy π(·|st) in maximum entropy RL, distributional RL can be viewed as subtracting a cross entropy term between µ and qθ. Moreover, the entropy in maximum entropy RL is state-wise, while our cross entropy regularization is state-action-wise, which is a more fine-grained characterization on the variability of the full distribution of return Z.
Notably, we highlight that our derived cross entropy regularization is risk-aware. If the true state-action distribution Fµ has higher degree of dispersion, e.g., a larger variance, than the current qθ, minimizing our entropy automatically encourages the policy to explore the uncertainty of the environment through changing Zθ. A sketch map in Figure 1 is 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 πZ(s) = argmaxa E [Zθ(s, a)] even though qθ changes towards µ as exhibited in Figure 1. A risk-aware action rule on the current greedy version might be preferred in the future.

heading: 4 STABLE OPTIMIZATION AND GUARANTEED GENERALIZATION
text: In this section, we further characterize qs,aθ in order to attain stable optimization and guaranteed generalization properties of distributional RL. Concretely, we leverage the histogram function fs,a to parameterize the density function of Zθ(s, a), i.e., q s,a θ , still based on the KL 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 on each state s, and we let the support of x(s) be uniformly partitioned into k bins. We let the function fs,a : X → [0, 1]k provide a k-dimensional vector fs,a(x(s)) of the coefficients indicating the probability the target is in that bin given the state s and action a pair, and x(s). We use softmax based on the linear approximation x(s)>θi to express fs,a, i.e., fs,a,θi (x(s)) = exp ( x(s)>θi ) / ∑k j=1 exp ( x(s)>θj ) . For simplicity, we use fθi (x(s)) to replace fs,a,θi (x(s)). Therefore, the resulting histogram distributional loss Lθ is formulated as:
Lθ(s, a) = DKL(ps,a, qs,aθ ) = ∫ x∈X ps,a(x) log (ps,a(x)/qs,aθ (x)) dx = − k∑ i=1 ps,ai log f θ i (x(s))
(8) where θ = {θ1, ..., θk} and ps,ai is the histogram probability of ps,a(x) defined in Eq. 4. The derivation of the histogram distributional loss is given in Appendix E. To attain the stable optimization property of distributional RL, we firstly derive Lemma 1 as follows Lemma 1. (Properties of Histogram Distributional Loss) Assume each state feature is bounded, i.e., ‖x(s)‖ ≤ l, then Lθ is kl2-smooth, convex and 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 classical Least-Squared estimation in neural-Q-fitted iteration for expectation-based RL does not enjoy these smoothness properties. In particular, for histogram distributional loss we have ‖∇θLθ‖ ≤ kl, while the counterpart in expectation-based is |yi−Qkθ(s, a)|‖x(s)‖, which might be arbitrarily large
due to the unbounded |yi − Qkθ(s, a)|. To further derive the uniform stability of distributional RL while running stochastic gradient descend (SGD), we introduce its definition in Definition 1.
Definition 1. (Uniform Stability) (Hardt et al., 2016) Consider a function g, a randomized algorithm M is uniformly stable if for all data sets S, S′ such that S, S′ differ in at most one example, we have
sup x
EM [g(M(S);x)− g (M (S′) ;x)] ≤ stab . (9)
Given the definition of uniform stability, we show that distributional RL under histogram distributional loss while running SGD is stab-uniformly stable and generalization guaranteed in Theorem 2.
Theorem 2. (Stable Optimization and Guaranteed Generalization) Suppose that we run SGD under Lθ in Eq. 8 with step sizes λt ≤ 2/kl2 for T steps, and ‖x(s)‖ ≤ l for each state s, we have:
(1) SGD under Lθ satisfies the uniform stability in Definition 1 with stab ≤ 4kTn , (2) Suppose that the data distribution is stable within each neural Z-fitted iteration, the generalization gap for solving the entropy regularized MLE in Eq. 5 is stab -bounded.
Please refer to the proof of Theorem 2 in Appendix G. In summary, the optimization to solve the entropy regularized MLE of distributional RL within the update of neural Z-fitted iteration is stable with the stability errors shrinking at the rate of O(n−1). This stability optimization properties is ascribed to the smooth and convex properties from histogram distributional loss (non-convex case still holds in (Hardt et al., 2016)) and can further control the generalization gap. By contrast, the least-estimation in neural Q-fitted iteration for expectation-based RL, which is without these smooth properties, can not yield the stable optimization and guaranteed generalization directly.

heading: 5 ACCELERATION EFFECT
text: Based on the setting in Section 4, we further incorporate the decomposition of F s,a(x) proposed in Eq. 4 into the analysis, where ps,a(x) = (1 − )δ{x=E[Zπ(s,a)]}(x) + µs,a(x), in order to derive the acceleration effect of distributional RL. Within each update in the neural Z-fitted iteration, the target is to minimize 1n ∑n i=1 Lθ(si, ai). We denote Gk(θ) as the expectation of Lθ, i.e., Gk(θ) = E(s,a)∼ρπ [Lθ(s, a)]. As such, the convex and smooth properties in Lemma 1 still hold for Gk(θ). We use G(θ) for Gk(θ) 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)∼ρπ [ ‖∇Lθ(δ{x=E[Zπ(s,a)]}, qs,aθ ))−∇G(θ)‖2 ] = σ2. Next, we characterize the approximation degree of qs,aθ to the ground-truth distribution µ s,a defined in Eq. 5 by measuring its variance as κσ2:
E(s,a)∼ρπ [ ‖∇Lθ(µs,a, qs,aθ ))−∇G(θ)‖ 2 ] = κσ2. (10)
Notably, a favorable approximation of qs,aθ to µ s,a θ would lead to a small κ, in which case an acceleration effect of distributional RL can be derive as shown in Theorem 3. Based on Eq. 10, we immediately have the following lemma.
Lemma 2. As ps,a(x) = (1− )δ{x=E[Zπ(s,a)]}(x) + µs,a(x), we have: E(s,a)∼ρπ [ ‖∇Lθ(ps,a, qs,aθ ))−∇G(θ)‖ 2 ] ≤ (1− )σ2 + κσ2. (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 in expectation-based RL, we need the following definition of the first-order τ -stationary point.
Definition 2. (First-order τ -Stationary Point) While solving minθ G(θ), the updated parameters θT after T steps is a first-order τ -stationary point if ‖∇G(θT )‖ ≤ τ , where the small τ is in (0, 1).
We formally show our Theorem 3 to characterize the acceleration effect of distributional RL determined by the variance magnitude τ .
Theorem 3. (Acceleration Effect) While running SGD to solve the entropy-regularized MLE within neural Z-fitted iteration in Eq. 5 with the step size λ = 1/kl2 and = 1/(1 + τ), we have:
(1) The sample complexity is O( 1τ4 ) if we only consider the expectation in p s,a, i.e., δ{x=E[Zπ(s,a)]}.
(2) When κ ≤ τ 2 4σ2 and let T = 4G(θ0) λτ2 , the regularized MLE of distributional RL within each neural Z-fitted iteration converges to a τ -stationary point in expectation with sample complexity O( 1τ2 ).
(3) When κ > τ 2 4σ2 and let T = G(θ0) λκτ2 , the regularized MLE of distributional RL does not converge to a τ -stationary point, but we have E [ ‖∇G(θ)‖2 ] ≤ O(κ).
The proof is provided in Appendix I. Theorem 3 is inspired by (Xu et 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 3 in fact reveals the reason why distributional RL algorithms can achieve inconsistent superiority across different Atari games. This phenomenon is mainly owing to the different degree of TD target approximation to ps,a, especially µs,a for various environments. In particular, a small approximation error of qθ to p (or µ) corresponds to a small κ, This can be normally ≤ τ2/4σ2, yielding better sample efficiency and accelerating the algorithm relative to expectation-based RL. Conversely, an unsatisfactory approximation 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.
Extension of Representation Effect. Apart from the acceleration effect, we also conduct some empirical analysis of distributional RL from the perspective of representation. In particular, distributional RL encourages state representation from the same action class classified by the policy in tighter clusters. Please refer to Appendix K for more details.

heading: 6 ALGORITHM: SINKHORN DISTRIBUTIONAL RL
text: Through the theoretical analysis, we can gain insights into the advantages of distribution RL over expectation-based RL, attributing to entropy-based regularization, stable optimization and acceleration effect determined by the distributional approximation error (measured by κ). From the algorithmic perspective, the choice of representation manner on Zθ and the distributional divergence metric dp mentioned in Section 3.1 are pivotal for the final performance of distributional RL algorithms.
we further design Sinkhorn distributional RL algorithm. Sinkhorn loss (Sinkhorn, 1967) is a tractable loss to approximate optimal transport problem by leveraging an entropic regularization to turn the original Wasserstein distance into a differentiable and more robust quantity. The resulting loss can be computed using Sinkhorn fixed point iterations, which is naturally suitable for modern deep learning frameworks. In particular, the entropic smoothing generates a family of losses interpolating between Wasserstein distance and Maximum Mean Discrepancy (MMD). Thus it allows us to find a sweet trade-off that simultaneously leverages the geometry of Wasserstein distance on the one hand, and the favorable high-dimensional sample complexity and unbiased gradient estimates of MMD. We introduce the entropic regularized Wassertein distance as
min Π∈Π(u,v)
∫ c(x, y)dΠ(x, y) + ε ∫ log ( Π(x, y)
du(x)dv(y)
) dΠ(x, y), (12)
where c is the cost function in optimal transport. This objective function associated with cost function c can be rewritten as Wc,ε(µ, ν) = ∫ c(x, y)dΠε(x, y). Therefore, the sinkhorn loss between two measures u and v is defined as
Wc,ε(u, v) = 2Wc,ε(u, v)−Wc,ε(u, u)−Wc,ε(v, v) (13)
Theorem 4. If we leverage Sinkhorn loss as the metric in distributional RL and choose c as unrectified kernel kα, i.e., kα(x, y) := −‖x− y‖α, for ∀α ∈ R, we have:
(1) As → 0,Wc,ε(u, v)→ 2Wα(u, v), and thus Tπ is a γ-contraction.
(2) As →∞,Wc,ε(u, v)→ MMD−kα(u, v), and thus Tπ is a γα/2-contraction.
Proof is provided in Appendix J. Theorem 4 indicates that if we choose c as the unrectified kernel, the limiting behaviors of distributional Bellman operator Tπ are both contractive under Sinkhorn loss. Note that without the limitation, the entropic regularization restricts the search space on the optimal
transport that eventually corresponds to a Gibbs kernel, based on which Tπ may not be a contraction. However, similar to MMD method, we show that our approach can still achieve empirical success and is very competitive across a wide range of Atari games in Section 7.
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∗ ∼ π(·|s′). 3: else 4: a∗ ← arg maxa′∈A 1N ∑N i=1 Zθ (s
′, a′)i 5: end if 6: TZi ← r + γZθ− (s′, a∗)i ,∀1 ≤ i ≤ N
Output: Wk2,ε ( {Zθ(s, a)i}Ni=1 , {TZi} N i=1 )
In the algorithm design, similar to in MMD distributional RL (Nguyen et al., 2020), we apply particle representation to represent Zθ(s, a) by directly generating multiple deterministic samples. Finally, we minimize the Sinkhorn loss between the approximate distribution via multiple samples and its distributional Bellman target. A detailed description of algorithm is provided in Algorithm 1.
Remark. MMD enjoys a convergence rate of O(n−1/2) 1 regardless of the underlying dimension while 1-Wasserstein distance has a convergence rate of O(n−1/d) if d > 2, which is slower for large d. Thus, Sinkhorn loss has the potential to enjoy the faster convergence of MMD.

heading: 7 EXPERIMENTS
text: We demonstrate the effectiveness of Sinkhorn distributinal RL (SinkhornDRL) as described in Algorithm 1 on a wide range of Atari 2600 games. Specifically, we leverage the same architecture as QR-DQN (Dabney et al., 2018b) for a fair comparison. More advanced techniques that can expand 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 SinkhornDRL, we choose 3 typical distributional RL algorithms as classic baselines, including QR-DQN (Dabney et al., 2018b), C51 (Bellemare et al., 2017a) and MMD (Nguyen et al., 2020), as well as DQN (Mnih et al., 2015). MMD algorithm is implemented with the same architecture as QRDQN, and leverages Gaussian kernels kh(x, y) =
exp(−(x − y)2/h) with the kernel mixture trick covering a range of bandwidths h, which is same as the basic setting in the original MMD paper (Nguyen et al., 2020).
Hyperparameter settings. For a fair comparison with QR-DQN, C51 and MMD, we used the same hyperparamters: the number of generated samples N = 200, Adam optimizer with lr = 0.00005, Adam = 0.01/32. 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 L = 10 and smoothing hyperparameter ε = 10.0 in Section 7.1 as they are not sensitive within a proper interval demonstrated in Section 7.2. We choose α = 2 in kα.

heading: 7.1 PERFORMANCE OF SINKHORNDRL
text: Figure 2 illustrates that SinkhornDR can achieve the state-of-the-art or competitive performance on typical Atari games compared with various baseline algorithms (DQN, QR-DQN, C51, MMD) in different metrics dp and representation manners on Zθ. Even though C51 outperforms others on breakout game, it is significantly inferior on other games, e.g., Enduro, Pong and Berzerk. Moreover, SinkhornDRL significantly outperforms MMD on breakout and SpaceInvader, and this superiority can be owing to the theoretical advantage of Sinkhorn loss over MMD. This coincides with the theoretical results as demonstrated in Theorem 4 that 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 L.

heading: 7.2 SENSITIVITY ANALYSIS
text: We further examine the impact of different hyperparameters in SinkhornDRL on the final performance, including the smoothing hyperparameter ε in Eq. 12 the number of deterministic samples N in Algorithm 1. From the left diagram in Figure 3, we can observe that our algorithm is not sensitive to the magnitude of ε as long as ε is within an appropriate interval, e.g., [1, 100]. Mean-
while, it turns out that an overly large number of samples N can even slightly worsen the performance of SinkhornDRL.

heading: 8 DISCUSSIONS AND CONCLUSION
text: Implicit generative models can be further incorporated into SinkhornDRL, including parameterizing the cost function in Sinkhorn loss, which can be naturally expected to achieve more promising performance in the future. Moreover, a direct analysis on Wasserstein or `p 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, a 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 RL algorithms in real scenarios. As distributional RL enjoys the robustness against
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 fairness issues.
Reproducibility Statement. For the theoretical part, we clearly state the related assumption and detailed proof process in the appendix. In terms of the algorithm implementation, our Sinkhorn algorithm is directly adapted from the public distributional RL algorithms, such as MMD.

heading: A DEFINITION OF DISTANCES AND CONTRACTION
text: Definition of distances. Given two random variables X and Y , p-Wasserstein metric Wp between the distributions of X and Y is defined as
Wp(X,Y ) = (∫ 1 0 ∣∣F−1X (ω)− F−1Y (ω)∣∣p dω)1/p = ‖F−1X − F−1Y ‖p, (14) which F−1 is the inverse cumulative distribution function of a random variable with the cumulative distribution function as F . Further, `p distance (Elie & Arthur, 2020) is defined as
`p(X,Y ) := (∫ ∞ −∞ |FX(ω)− FY (ω)|p dω )1/p = ‖FX − FY ‖p (15)
The `p distance and Wassertein metric are identical at p = 1, but are otherwise distinct. Note that when p = 2, `p distance is also called Cramér distance (Bellemare et al., 2017b) dC(X,Y ). Also, the Cramér distance has a different representation given by
dC(X,Y ) = E|X − Y | − 1 2 E |X −X ′| − 1 2 E |Y − Y ′| , (16)
where X ′ and Y ′ are the i.i.d. copies of X and Y . Energy distance (Székely, 2003; Ziel, 2020) is a natural extension of Cramér distance to the multivariate case, which is defined as
dE(X,Y) = E‖X−Y‖2 − 1 2 E‖X−X′‖ − 1 2 E‖Y −Y′‖, (17)
where X and Y are multivariate. Moreover, the energy distance is a special case of the maximum mean discrepancy (MMD), which is formulated as
MMD(X,Y; k) = (E [k (X,X′)] + E [k (Y,Y′)]− 2E[k(X,Y)])1/2 (18)
where k(·, ·) is a continuous kernel on X . In particular, if k is a trivial kernel, MMD degenerates to energy distance. Additionally, we further define the supreme MMD, which is a functional P(X )S×A × P(X )S×A → R defined as
MMD∞(µ, ν) = sup (x,a)∈S×A MMD∞(µ(x, a), ν(x, a)) (19)
Proof of Contraction.
• Contraction under supreme form of Wasserstein diatance is provided in Lemma 3 (Bellemare et al., 2017a).
• Contraction under supreme form of `p distance can refer to Theorem 3.4 (Elie & Arthur, 2020).
• Contraction under MMD∞ is provided in Lemma 6 (Nguyen et al., 2020).

heading: B PROOF OF PROPOSITION 1
text: Proof. (1) We recap three crucial properties of a divergence metric. The first is scale sensitive (S) (of order β, beta > 0), i.e., dp(cX, cY ) ≤ |c|βdp(X,Y ). The second property is shift invariant (I), i.e., dp(A + X,A + Y ) ≤ dp(X,Y ). The last one is unbiased gradient (U). We use p and q to denote the density function of two random variables X and Y , and thus DKL(X,Y ) is defined as DKL(X,Y ) = ∫∞ −∞ p(x) p(x) q(x) dx. Firstly, we show that DKL(X,Y ) is NOT scale sensitive:
DKL(cX, cY ) = ∫ ∞ −∞ 1 a p( x a ) 1 ap( x a ) 1 aq( x a ) dx
= ∫ ∞ −∞ p(y) p(y) q(y) dy
= DKL(X,Y ), with β = 0
(20)
We further show that DKL(X,Y ) is shift invariant:
DKL(A+X,A+ Y ) = ∫ ∞ −∞ p(x−A)p(x−A) q(x−A) dx
= ∫ ∞ −∞ p(y) p(y) q(y) dy
= DKL(X,Y )
(21)
Moreover, it is well-known that KL divergence has unbiased sample gradients (Bellemare et al., 2017b). The supreme DKL is a functional P(X )S×A × P(X )S×A → R defined as
D∞KL(µ, ν) = sup (x,a)∈S×A DKL(µ(x, a), ν(x, a)) (22)
Therefore, we prove Tπ is at best a non-expansive operator under the supreme form of DKL:
D∞KL(T πZ1,T πZ2)
= sup s,a
DKL(T πZ1(s, a),T πZ2(s, a))
= DKL(R(s, a) + γZ1(s ′, a′), R(s, a) + γZ2(s ′, a′)) = DKL(R(s, a) + Z1(s ′, a′), Z2(s ′, a′))
≤ sup s′,a′ DKL(R(s, a) + Z1(s ′, a′), Z2(s ′, a′)) = D∞KL(Z1, Z2)
(23)
There we have D∞KL(T πZ1,T πZ2) ≤ D∞KL(Z1, Z2), implying that Tπ is a non-expansive operator under D∞KL.
(2) By the definition of D∞KL, we have sups,aDKL(Zn(s, a), Z(s, a)) → 0 implies DKL(Zn, Z) → 0. DKL(Zn, Z) → 0 implies the total variation distance δ(Zn, Z) → 0 according to a straightforward application of Pinsker’s inequality
δ (Zn, Z) ≤ √ 1
2 DKL (Zn, Z)→ 0 δ (Z,Zn) ≤ √ 1
2 DKL (Z,Zn)→ 0
(24)
Based on Theorem 2 in WGAN (Arjovsky et al., 2017), δ(Zn, Z) → 0 implies Wp(Zn, Z) → 0. This is trivial by recalling the fact that δ and W give the strong an weak topologies on the dual of (C(X ), ‖ · ‖∞) when restricted to Prob(X ). (3) The conclusion holds because the Tπ degenerates to T π regardless of the metric dp. Specifically, due to the linearity of expectation, we obtain that
‖ETπZ1 − ETπZ2‖∞ = ‖T πEZ1 − T πEZ2‖∞ ≤ γ‖EZ1 − EZ2‖∞. (25)
This implies that the expectation of Z under DKL exponentially converges to the expectation of Z∗, i.e., γ-contraction.

heading: C PROOF OF PROPOSITION 2
text: Proof. Firstly, given a fixed p(x) we know that minimizing DKL(p, qθ) is equivalent to minimizing H(p, q) by following
DKL(p, qθ) = ∫ +∞ −∞ p(x) log p(x) qθ(x) dx
= − ∫ +∞ −∞ p(x) log qθ(x) dx− (− ∫ +∞ −∞ p(x) log p(x) dx)
= H(p, qθ)−H(p) ∝ H(p, qθ)
(26)
Based onH(p, qθ), we can derive the objective function within each neural Z-fitted iteration as
1
n n∑ i=1 H(psi,ai(x), qsi,aiθ (x))
= 1
n n∑ i=1 ( − ∫ +∞ −∞ psi,ai(x) log qsi,aiθ (x) dx )
= 1
n n∑ i=1 ( − ∫ +∞ −∞ ( (1− )δ{x=E[Zπ(s,a)]}(x) + µsi,ai(x) ) log qsi,aiθ (x) dx )
= 1
n n∑ i=1 [ (1− ) ( − ∫ +∞ −∞ δ{x=E[Zπ(s,a)]}(x) log q si,ai θ (x) dx ) + H(µsi,ai , qsi,aiθ ) ]
= 1
n n∑ i=1 [ (1− )H(δ{x=E[Zπ(si,ai)]}, q si,ai θ ) + H(µ si,ai , qsi,aiθ ) ]
∝ 1 n n∑ i=1 H(δ{x=E[Zπ(si,ai)]}, q si,ai θ ) + αH(µ si,ai , qsi,aiθ ), where α = 1− > 0
= 1
n n∑ i=1 ( − ∫ +∞ −∞ log qsi,aiθ (E [Z π(s, a)]) dx+ αH(µsi,ai , qsi,aiθ ) )
(27)

heading: D PROOF OF THEOREM 1
text: Proof. Firstly, we have: V (st+1) = Eat+1∼π,x∼µst+1,at+1 [ Q (st+1, at+1) + log q st+1,at+1 θ (x) ] = Eat+1∼π [Q (st+1, at+1)] + Eat+1∼π,x∼µst+1,at+1 [ log q st+1,at+1 θ (x)
] = Eat+1∼π [Q (st+1, at+1)]− Eat+1∼π [ − ∫ ∞ −∞ µst+1,at+1(x) log q st+1,at+1 θ (x) dx
] = Eat+1∼π [Q (st+1, at+1)]− Eat+1∼π [ H ( µst+1,at+1 , q st+1,at+1 θ
)] (28)
Further, we plug in V (st+1) into RHS of the iteration in Eq. 7, then we obtain
T πd Q (st, at) = r (st, at) + γEst+1∼ρπ [V (st+1)] = r (st, at)− γE(st+1,at+1)∼ρπ [ H ( µst+1,at+1 , q st+1,at+1 θ )] + γE(st+1,at+1)∼ρπ [Q (st+1, at+1)]
, rπ (st, at) + γE(st+1,at+1)∼ρπ [Q (st+1, at+1)] , (29) where rπ (st, at) , r (st, at) − γE(st+1,at+1)∼ρπ [ H ( µst+1,at+1 , q st+1,at+1 θ )] is the entropy augmented reward we redefine. Thus, repeatedly updating via the above equations is equivalent to maximizing the new objective function J ′(θ) = ∑T t=0 E(st,at)∼ρπ [r (st, at)− γH (µst,at , q st,at θ )]. Applying the standard convergence results for policy evaluation (Sutton & Barto, 2018), we can attain that this Bellman updating under T πd is convergent under the assumption of |A| < ∞ and bounded entropy augmented rewards rπ .

heading: E DERIVATION OF HISTOGRAM DISTRIBUTIONAL LOSS
text: We show the derivation details of the Histogram distribution loss starting from KL divergence between p and qθ. pi is the cumulative probability increment of target distribution {Yi}i∈[n] within the i-th bin, and qθ corresponds to a (normalized) histogram, and has density values fθi (x(s)) wi per bin.
Thus, we have:
DKL (p s,a, qs,aθ ) = − ∫ b a ps,a(y) log qs,aθ (y)dy
= − k∑ i=1 ∫ li+wi li ps,a(y) log fθi (x(s)) wi dy
= − k∑ i=1 log fθi (x(s)) wi (F s,a (li + wi)− F s,a (li))︸ ︷︷ ︸
pi ∝ − k∑ i=1 ps,ai log f θ i (x(s))
(30)
where the last equality holds because the width parameter wi can be ignored for this minimization problem.

heading: F PROOF OF LEMMA 1
text: Proof. For the histogram distributional loss below,
Lθ(s, a) = − k∑ i=1 ps,ai log f θ i (x(s)), where f θ i (x(s)) =
exp ( x(s)>θi )∑k j=1 exp (x(s) >θj)
we firstly prove its convexity. Note that − log exp(x(s) >θi)∑k
j=1 exp(x(s) >θj)
= log ∑k j=1 exp ( x(s)>θj ) −
x(s)>θi, 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, Lθ(s, a) is convex. Secondly, we show that Lθ(s, a) is kl-Lipschitz. We compute the gradient of the Histogram distributional loss regarding θi:
∂
∂θi k∑ j=1 ps,aj log f θ j (x(s))
= k∑ j=1 ps,aj 1 fθj (x(s)) ∇θifθj (x(s))
= k∑ j=1 ps,aj 1 fθj (x(s)) fθi (x(s))(δij − fθj (x(s)))x(s)
= ps,ai (1− fθi (x(s)))− k∑ j 6=i ps,aj f θ i (x(s)) x(s) = ( ps,ai − p s,a i f θ i (x(s))− (1− p s,a i )f θ i (x(s)) ) x(s)
= ( ps,ai − f θ i (x(s)) ) x(s)
(31)
where δij = 1 if i = j, otherwise 0. Then, as we have ‖x(s)‖ ≤ l, we bound the norm of its gradient
‖ ∂ ∂θ k∑ j=1 pj log f θ j (x(s))‖
≤ k∑ i=1 ‖ ∂ ∂θi k∑ j=1 pj log f θ j (x(s))‖
= k∑ i=1 ‖ ( ps,ai − f θ i (x(s)) ) x(s)‖
≤ k∑ i=1 |ps,ai − f θ i (x(s))|‖x(s)‖
≤ kl
(32)
The last equality satisfies because |pi − fθi (x(s))| is less than 1 and even smaller. Therefore, we obtain that Lθ is kl-Lipschitz. Lastly, we show that Lθ is kl2-Lipschitz smooth. A lemma is that log(1 + exp(x)) is 14 -smooth as its second-order gradient is bounded by 14 , and if g(w) is β-smooth w.r.t. w, then g(〈x,w〉) is β‖x‖2-smooth. Based on this knowledge, we firstly focus on the 1-dimensional case of function log fθj (z). As we have derived, we know that ∂ ∂θi log fθj (z) = δij − fθi (z). Then the secondorder gradient is ∂2 log fθj (z)
∂θi∂θk = −(δik − fθk (z)) = fθk (z) − 1 if i = k, otherwise fθk (z). Clearly,
|∂ 2 log fθj (z)
∂θi∂θk | ≤ 1, which implies that log fθj (z) is 1-smooth. Thus, log fθj (〈x, θi〉) is ‖x‖2-smooth, or l2-smooth. Further, ∑k j=1 p s,a j log f θ j (x(s)) is also l 2-smooth as we have
‖∇θi k∑ j=1 ps,aj log f θ j (µ)−∇θi k∑ j=1 ps,aj log f θ j (ν)‖
≤ k∑ j=1 ps,aj ‖∇θi log f θ j (µ)−∇θi log fθj (ν)‖ ≤ k∑ j=1 ps,aj · l 2‖µ− ν‖
= l2‖µ− ν‖
(33)
for each µ and ν. Therefore, we further have
‖∇θ k∑ j=1 ps,aj log f θ j (µ)−∇θ k∑ j=1 ps,aj log f θ j (ν)‖
≤ k∑ i=1 ‖∇θi k∑ j=1 ps,aj log f θ j (µ)−∇θi k∑ j=1 ps,aj log f θ j (ν)‖ ≤ k∑ i=1 l2‖µ− ν‖
= kl2‖µ− ν‖
(34)
Finally, we conclude that Lθ(s, a) is kl2-smooth.

heading: G PROOF OF THEOREM 2
text: Proof. Consider the stochastic gradient descent rule asGλ,L(θ) = θ−λ∇θL(θ). Firstly, we provide two definitions about Lθ for the following proof.
Definition 3. (σ-bounded) An update rule is σ-bounded if supθ ‖θ − λ∇θL(θ)‖ ≤ σ.
Definition 4. (η-expansive) An update rule is η-expansive if supv,w ‖Gλ,L(v)−Gλ,L(w)‖ ‖u−w‖ ≤ η.
Lemma 3. (Grow Recursion, Lemma 2.5 (Hardt et al., 2016)) Fix an arbitrary sequence of updates G1, ..., GT and another sequence G′1, ..., G ′ T . Let θ0 = θ ′ 0 be the starting point and define δt = ‖θ′i − θt‖, where θt and θ′t are defined recursively through
θt+1 = Gλ,L(θt), θ ′ t+1 = G ′ λ,L(θ ′ t)
Then we have the recurrence relation: δt+1 ≤ { ηδt Gt = G ′ t is η-expansive
min(η, 1)δt + 2σt Gt and G′t are σ-bounded , Gt is η expansive
Lemma 4. (Lipschitz Continuity) Assume Lθ is L-Lipschitz, the gradient update Gλ,L is (λL)bounded.
Proof. ‖θ −Gλ,L(θ)‖ = ‖λ∇θL(θ)‖ ≤ λL
Lemma 5. (Lipschitz Smoothness) Assume Lθ is β-smooth, then for any λ ≤ 2β , the gradient update Gλ,L is 1-expansive.
Proof. Please refer to Lemma 3.7 in (Hardt et al., 2016) for the proof.
Based on all the results above, we start to prove Theorem 2. Our proof is largely based on (Hardt et al., 2016), but it is applicable in distributional RL setting as well as considering desirable properties of histogram distributional loss. According Lemma 1, we attain that Lθ is kl-Lipschitz as well as kl2-smooth, and thus based on Lemma 4, Gλ,L is (λkl)-bounded, and 1-expansive if λ ≤ 2kl2 . In the step t, SGD selects samples that are both in S and S′, with probability 1 − 1n . In this case, Gt = G ′ t, and thus δt+1 ≤ δt as Gt is 1-expansive based on Lemma 3. The other case is that samples selected are different with probability 1n , where δt+1 ≤ δt + 2λtkl based on Lemma 3. Thus, if λt ≤ 2kl2 we have:
E |L(θT ;x)− L(θ′T ;x)| ≤ klE [δT ] , where δT = ‖θT − θ′T ‖ ≤ kl (
(1− 1 n )E [δT−1] + 1 n E [δT−1] + 2λT−1kl n ) = kl ( E [δT−1] + 2λT−1kl
n ) = kl ( E [δ0] +
T−1∑ t=0 2λtkl n
)
≤ 2k 2l2
n T−1∑ t=0 2 kl2
= 4kT
n
(35)
Since this bounds hold for all S, S′ and x, we attain the uniform stability in Definition 1 for our histogram distributional loss applied in distributional RL.
Define the population risk as: R [θ] = ExL(θ;x)
and the empirical risk as:
RS [θ] = 1
n n∑ i=1 L(θ;xi)
According to Theorem 2.2 in (Hardt et al., 2016), if an algorithmM is stab-uniformly stable, then the generalization gap is stab-bounded, i.e.,
|ES,A [RS [M(S)]−R[M(S)]]| ≤ stab

heading: H PROOF OF LEMMA 2
text: E(s,a)∼ρπ [ ‖∇Lθ(ps,a, qs,aθ ))−∇G(θ)‖ 2 ] ≤ (1− )σ2 + κσ2. (36)
Proof. As we know that ps,a(x) = (1− )δ{x=E[Zπ(s,a)]}(x) + µs,a(x), then we have: ∇Lθ(ps,a, qs,aθ ) = (1− )∇Lθ(δ{x=E[Zπ(s,a)]}, q s,a θ ) + ∇Lθ(µ s,a, qs,aθ )
Therefore, E(s,a)∼ρπ [ ‖∇Lθ(ps,a, qs,aθ ))−∇G(θ)‖ 2 ]
≤ E(s,a)∼ρπ [ (1− )‖∇Lθ(δ{x=E[Zπ(s,a)]}, qs,aθ ))−∇G(θ)‖ 2 + ‖∇Lθ(µs,a, qs,aθ ))−∇G(θ)‖ 2 ]
= (1− )σ2 + κσ2, (37)
where the first inequality uses the triangle inequality of norm, i.e., ‖(1 − )a + b‖2 ≤ (1 − )‖a‖2 + ‖b‖2, and the last equality uses the definition of the variance of Lθ(δ{x=E[Zπ(s,a)]}, qs,aθ ) and Lθ(µs,a, qs,aθ ).

heading: I PROOF OF THEOREM 3
text: Proof. (1) If we only consider the expectation of Zπ(s, a), the entropy-regularized MLE would degenerate to the pure MLE regarding δ{x=E[Zπ(s,a)]}. As Lθ(δ{x=E[Zπ(s,a)]}, qs,aθ ) is kl2-smooth, we have
G(θt+1)−G(θt)
≤ 〈∇G(θt), θt+1 − θt〉+ kl2
2 ‖θt+1 − θt‖2
= −λ 〈 ∇G(θt),∇Lθ(δ{x=E[Zπ(s,a)]}, qs,aθ ) 〉 + kl2λ2
2 ‖∇Lθ(δ{x=E[Zπ(s,a)]}, qs,aθ )‖ 2
(38)
where the last first equation is according to the definition of Lipschitz-smoothness, and the last second one is based on the updating rule of θ. Next, we take the expectation on both sides,
E [G(θt+1)−G(θt)]
≤ −λE [ ‖∇G(θt)‖2 ] + kl2λ2 2 E [ ‖∇Lθ(δ{x=E[Zπ(s,a)]}, qs,aθ )−∇G(θt) +∇G(θt)‖ 2 ]
≤ −λE [ ‖∇G(θt)‖2 ] + kl2λ2 2 E [ ‖∇Lθ(δ{x=E[Zπ(s,a)]}, qs,aθ )−∇G(θt)‖ 2 ] + kl2λ2 2 E [ ‖∇G(θt)‖2 ] = λ(kl2λ− 2)
2 E [ ‖∇G(θt)‖2 ] + kl2λ2 2 σ2
≤ −λ 2 E [ ‖∇G(θt)‖2 ] + kl2λ2 2 σ2
(39) where the first two equation hold because ∇G(θ) = E [∇Lθ] and the last inequality comes from λ ≤ 1kl2 . Through the summation, we obtain that
E [G(θT )−G(θ0)] ≤ − λ
2 T−1∑ t=0 E [ ‖∇G(θt)‖2 ] + kl2λ2T 2 σ2
We let E [G(θT )] = 0, we have
1
T T−1∑ t=0 E [ ‖∇G(θt)‖2 ] ≤ 2G(θ0) λT + kl2λσ2
By setting λ ≤ τ 2 2kl2σ and T = 4G(θ0) λτ2 , we can have 1 T ∑T−1 t=0 E [ ‖∇G(θt)‖2 ] ≤ τ2, implying that the degenerated MLE can achieve τ -station point if the sample complexity T = O( 1τ4 ).
(2) and (3) We are still based on the kl2-smoothness of L(ps,a, qs,aθ ).
G(θt+1)−G(θt)
≤ 〈∇G(θt), θt+1 − θt〉+ kl2
2 ‖θt+1 − θt‖2
= −λ 〈∇G(θt),∇Lθ(ps,a, qs,aθ )〉+ kl2λ2
2 ‖∇Lθ(ps,a, qs,aθ )‖ 2
= −λ 2 ‖∇G(θt)‖2 + λ 2 ‖∇G(θt)−∇Lθ(ps,a, qs,aθ )‖ 2 + λ(kl2λ− 1) 2 ‖∇Lθ(ps,a, qs,aθ )‖ 2
≤ −λ 2 ‖∇G(θt)‖2 + λ 2 ‖∇G(θt)−∇Lθ(ps,a, qs,aθ )‖ 2
(40)
where the second equation is based on 〈a,−b〉 = 12 ( ‖a− b‖2 − ‖a‖2 − ‖b‖2 ) , and the last inequality is according to λ ≤ 1kl2 . After taking the expectation, we have
E [G(θt+1)−G(θt)] ≤ −λ 2 E [ ‖∇G(θt)‖2 ] + λ 2 E [ ‖∇G(θt)−∇Lθ(ps,a, qs,aθ )‖ 2 ]
≤ −λ 2 E [ ‖∇G(θt)‖2 ] + λ 2 ( (1− )σ2 + κσ2 ) (41) where the last inequality is based on Lemma 2. We take the summation, and therefore,
E [G(θT )−G(θ0)] ≤ − λ
2 T−1∑ t=0 E [ ‖∇G(θt)‖2 ] + Tλ 2 ( (1− )σ2 + κσ2 ) We let E [G(θT )] = 0 and = 11+κ , then,
1
T T−1∑ t=0 E [ ‖∇G(θt)‖2 ] ≤ 2G(θ0)
λT + (1− )σ2 + κσ2
= 2G(θ0)
λT +
2κ
1 + κ σ2
≤ 2G(θ0) λT + 2κσ2
(42)
If κ ≤ τ 2 4σ2 and let T = 4G(θ0) λτ2 , this leads to 1 T ∑T−1 t=0 E [ ‖∇G(θt)‖2 ] ≤ τ2, i.e., τ -stationary point, with the sample complexity asO( 1τ2 ). Thus, (2) has been proved. On the other hand, if κ > τ2
4σ2 , we set T = G(θ0)λκσ2 . This implies that 1 T ∑T−1 t=0 E [ ‖∇G(θt)‖2 ] ≤ 4κσ2 = O(κ). Therefore, the degree of stationary point is determined the degree of distribution approximation measured by κ. Thus, we obtain (3).

heading: J PROOF OF THEOREM 4
text: Proof. 1. As ε → 0, it is obvious to observe that Sinkhorn loss degenerates to the wasserstein distance.
2. As ε → ∞, Sinkhorn loss turns to be ∫ log (
Π(x,y) du(x)dv(y)
) dΠ(x, y). We can consider its dual
problem by introducing Lagrange multipliers u and v. By additionally considering the primal-dual relation, we can solve the dual problem, which gives u = v = 0, and thus the optimal coupling is simply the product of the marginals , ie., Π = u ⊗ v. Please refer to more detailed proof in Sinkhorn-GAN (Genevay et al., 2018). Since the cost function is the rectified kernel kα, under with MMD (Nguyen et al., 2020) can be γα/2-contractive.

heading: K REPRESENTATION EFFECT OF DISTRIBUTION RL
text: 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.
The intrinsic characteristic of distributional RL is that it enables to learn a richer and more faithful representation of the environment, leading to more stable and efficient learning. To investigate the more informative representation resulting from distributional RL, we visualize how distributional RL changes the representation learned in the penultimate layer of the value network. We collect state fea-
tures in the penultimate layer classified by the policy π in different action classes, and perform t-SNE to reduce them to the two-dimensional space. In Figure 4, points in different colors represent state features classified into different action classes. It illustrates that both QR-DQN and C51 encourage the representation of state observations from the same action class to group in tighter clusters relative to expectation-based DQN. Interestingly, this phenomenon is similar to label smoothing (Müller et al., 2019), which leverages additional distributional knowledge in soft labels relative to hard labels. Therefore, we conclude that similar to the benefit of label smoothing, distributional RL also encourages the activations of the penultimate layer to be close to the template of the correct action class and distant to the templates of the incorrect action classes.
L MORE EXPERIMENTAL RESULTS

Review summaries:
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.