Distributional Reinforcement Learning with Regularized Wasserstein Loss

Authors: ke  sun
yingnan  zhao
wulong  liu
bei  jiang
linglong  kong
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-13 17:01:02

Abstract: The empirical success of distributional reinforcement learning (RL) highly relies on the choice of distribution divergence equipped with an appropriate distribution representation. In this paper, we propose \textit{Sinkhorn distributional RL (SinkhornDRL)}, which leverages Sinkhorn divergence—a regularized Wasserstein loss—to minimize the difference between current and target Bellman return distributions. Theoretically, we prove the contraction properties of SinkhornDRL, aligning with the interpolation nature of Sinkhorn divergence between Wasserstein distance and Maximum Mean Discrepancy (MMD). The introduced SinkhornDRL enriches the family of distributional RL algorithms, contributing to interpreting the algorithm behaviors compared with existing approaches by our investigation into their relationships. Empirically, we show that SinkhornDRL consistently outperforms or matches existing algorithms on the Atari games suite and particularly stands out in the multi-dimensional reward setting. \thanks{Code is available in \url{https://github.com/datake/SinkhornDistRL}.}.

Main paper: 
 introduction: 
 The design of classical reinforcement learning (RL) algorithms primarily focuses on the expectation of cumulative rewards that an agent observes while interacting with the environment. Recently, a new class of RL algorithms called distributional RL estimates the full distribution of total returns and has exhibited state-of-the-art performance in a wide range of environments, such as C51 [5] , Quantile-Regression DQN (QR-DQN) [14] , EDRL [41] , Implicit Quantile Networks (IQN) [13] , Fully Parameterized Quantile Function (FQF) [56] , Non-Crossing QR-DQN [58] , Maximum Mean Discrepancy (MMD-DQN) [37] , Spline (SPL-DQN) [33] , and Sketch-DQN [53] . Beyond the performance advantage, distributional RL has also possessed benefits in risk-sensitive control [13, 9] , exploration [35, 11, 48] , offline setting [34, 55] , statistical value estimation [43] , robustness [47] and optimization [46, 52, 42, 51] .Limitations of Typical Distributional RL Algorithms. Despite the gradual introduction of numerous algorithms, quantile regression-based algorithms [14, 13, 56, 41, 33, 42, 43] dominate attention and research in the realm of distributional RL. These algorithms utilize quantile regression to approximate the one-dimensional Wasserstein distance to compare two return distributions. Nevertheless, two major limitations hinder their performance improvement and wider practical deployment. 1) Inaccuracy in Capturing Return Distribution Characteristics. The way of directly generating quantiles of return distributions via neural networks often suffers from the non-crossing issue [58] , where the learned quantile curves fail to guarantee a non-decreasing property. This leads to abnormal distribution estimates and reduced model interpretability. The inaccurate distribution estimate is fundamentally attributed to the use of pre-specified statistics [41] , while unrestricted statistics based on deterministic samples can be potentially more effective in complex environments [37] .2) Difficulties in Extension to Multi-dimensional Rewards. Many RL tasks involve multiple sources of rewards [32, 15] , hybrid reward architecture [50, 30] , or sub-reward structures after reward decomposition [31, 57] , which require learning multi-dimensional return distributions to reduce the intrinsic uncertainty of the environments. However, it remains elusive how to use quantile regressions to approximate a multi-dimensional Wasserstein distance, while circumventing the computational intractability issue in the related multi-dimensional output space.Motivation of Sinkhorn Divergence: a Regularized Wasserstein loss. Sinkhorn divergence [45] has emerged as a theoretically principled and computationally efficient alternative for approximating Wasserstein distance. It has gained increasing attention in the field of optimal transport [4, 24, 21, 39] and has been successfully applied in various areas of machine learning [38, 25, 54, 20, 8] . By introducing entropic regularization, Sinkhorn divergence can efficiently approximate a multi-dimensional Wasserstein distance using computationally efficient matrix scaling algorithms [45, 39] . This makes it feasible to apply optimal transport distances to RL tasks with multi-dimensional rewards (see experiments in Section 5.3). Moreover, Sinkhorn divergence enables the leverage of samples to approximate return distributions instead of relying on pre-specified statistics, e.g., quantiles, thereby increasing the accuracy in capturing the full data complexity behind return distributions and naturally avoiding the non-crossing issues in distributional RL. Beyond addressing the two main limitations mentioned above, the well-controlled regularization introduced in Sinkrhorn divergence helps to find a "smoother" transport plan relative to Wasserstein distance, making it less sensitive to noises or small perturbations when comparing two return distributions (see Appendix A for the visualization). The term "smoother" refers to the effect of regularization in Sinkhorn divergence to encourage a more uniformly distributed transport plan. This regularization also aligns with the maximum-entropy principle [28, 16] , which aims to maximize entropy while keeping the transportation cost constrained. Furthermore, the resulting strongly convex loss function [3] and the induced smoothness by regularization facilitate faster and more stable convergence in the deep RL setting (see more details in Sections 4 and 5).Sinkhorn divergence, a regularized Wasserstein loss, to address the limitations of quantile regressionbased algorithms while promoting more stable training. As Sinkhorn divergence interpolates between Wasserstein distance and MMD [26, 21, 39] , we probe this relationship in the RL context, characterizing the convergence properties of dynamic programming under Sinkhorn divergence and revealing the connections of different distances. Our study enriches the class of distributional RL algorithms, making them more effective for a broader range of scenarios and potentially inspiring advancement in other related areas of distribution learning. Our key contributions are summarized as follows:(1) Algorithm. We introduce a Sinkhorn distributional RL algorithm, called SinkhornDRL, which overcomes the primary shortcomings of predominantly utilized quantile regression-based algorithms. SinkhornDRL can be seamlessly integrated into existing model architectures and easily implemented.(2) Theory. We establish the properties of Sinkhorn divergence within distributional RL and derive the relevant convergence results for (multi-dimensional) distributional dynamic programming.(3) Experiments. We conduct an extensive comparison of SinkhornDRL with typical distributional RL algorithms across 55 Atari games, performing rigorous sensitivity analyses and computation cost assessments. We also verify the efficacy of SinkhornDRL in the multi-dimensional reward setting. 
distributional reinforcement learning: 
 In classical RL, 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, P is the environment transition dynamics, R is the reward function, and γ ∈ (0, 1) is the discount factor. Given a policy π, the discounted sum of future rewards Z π is a random variable with Z π (s, a) = ∞ t=0 γ t R(s t , a t ), where s 0 = s, a 0 = a, s t+1 ∼ P (•|s t , a t ), and a t ∼ π(•|s t ). In expectation-based RL, the action-value function Q π is defined as Q π (s, a) = E [Z π (s, a)], which is iteratively updated via Bellman operator T π through T π Q(s, a) = E[R(s, a)] + γE s ′ ∼p,π [Q (s ′ , a ′ )], where s ′ ∼ P (•|s, a) and a ′ ∼ π (•|s ′ ). In contrast, distributional RL focuses on the action-return distribution, the full distribution of Z π (s, a). The return distribution is iteratively updated by applying the distributional Bellman operator T π through T π Z(s, a) : D = R(s, a) + γZ (s ′ , a ′ ), where D denotes the distribution and the equality implies random variables of both sides are equal in distribution. The distributional Bellman operator T π is contractive under certain distribution divergence metrics [19] .Optimal Transport (OT) and Wasserstein / Earth Mover's Distance. The optimal transport (OT) metric W c defines a powerful geometry to compare two probability measures (µ, ν), i.e., W c = inf Π∈Π(µ,ν) c(x, y)dΠ(x, y), where c is the cost function, Π is the joint distribution with marginals (µ, ν), and the minimizer Π * is called the optimal transport plan or optimal coupling. The p-Wasserstein distance W p = (inf Π∈Π(µ,ν) ∥x -y∥ p dΠ(x, y)) 1/p is a special case of optimal transport with the Euclidean norm as the cost function. Relative to conventional divergences, including Hellinger, total variation or Kullback-Leibler divergences, the formulation of OT and Wasserstein distance inherently integrates the spatial or geometric relationships between data points and allows them to recover the full support of measures. This theoretical advantage comes, however, with a heavy computational price tag, especially in the high-dimensional space. Specifically, finding the optimal transport plan amounts to solving a linear program and the cost scales at least in O(d 3 log(d)) when comparing two histograms of dimension d [12] . [26] . Define two random variables X and Y . The squared Maximum Mean Discrepancy (MMD) MMD 2 k with the kernel k is formulated asMMD 2 k = E [k (X, X ′ )] + E [k (Y, Y ′ )] -2E [k(X, Y )],where k(•, •) is a continuous kernel and X ′ (resp. Y ′ ) is a random variable independent of X (resp. Y ). Mathematically, the "flat" geometry that MMD induces on the space of probability measures does not faithfully lift the ground distance [21] , potentially inferior to OT when comparing two complicated distributions. However, MMD is cheaper to compute than OT with a smaller sample complexity, i.e., the number of samples for measures to approximate the true distance [24] . We provide more details of various distribution divergences as well as their existing contraction properties in distributional RL in Appendix B.Notations. We constantly use the unrectified kernel k α = -∥x -y∥ α in our algorithm analysis. With a slight abuse of notation, we also use Z θ to denote θ parameterized return distribution.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: We have explicitly discussed the limitations of our study in a separate paragraph of the last section.Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
related work: 
 Based on the choice of distribution divergences and the distribution representation, distributional RL algorithms can be classified into three categories.(1) Categorical Distributional RL. As the first successful class, categorical distributional RL [5] , e.g., C51, represents the return distribution using a categorical distribution with discrete fixed supports within a predefined interval.(2) Quantile Regression (Wasserstein Distance) Distributional RL. QR-DQN [14] employs quantile regression to approximate the one-dimensional Wasserstein distance. It learns the quantile values for a series of fixed quantiles, offering greater flexibility in the support compared with categorical distributional RL. IQN [13] enhances this approach by utilizing an implicit model to produce more expressive quantile values, instead of fixed ones in QR-DQN, while FQF [56] further advances IQN by introducing a more expressive quantile network. However, as mentioned in Section 1, quantile regression distributional RL struggles with accurately capturing return distribution characteristics and handling multi-dimensional reward settings. SinkhornDRL, with the assistance of an entropy regularization, offers an alternative approach that addresses the two limitations simultaneously.(3) MMD Distributional RL. Rooted in kernel methods [26, 53] , MMD-DQN [37] learns unrestricted statistics, i.e., samples, to represent the return distribution and optimizes under MMD, which can manage multi-dimensional rewards. However, the data geometry captured by MMD with a specific kernel may be limited, as it is highly sensitive to the characteristics of kernels and the induced Reproducing Kernel Hilbert space (RKHS) [25, 26, 23] . In contrast, SinkhornDRL is fundamentally based on OT, inherently capturing the spatial and geometric layout of return distributions. This enables SinkhornDRL to potentially surpass MMD-DQN by leveraging a richer representation of data geometry. In Section 5, we present extensive experiments to demonstrate the advantage of SinkhornDRL over MMD-DQN, particularly in the multi-dimensional reward scenario in Section 5.3.some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results. 
sinkhorn distributional rl (sinkhorndrl): 
 The algorithmic evolution of distributional RL can be primarily viewed along two dimensions [37] . (1) Introducing new distributional RL families beyond the three established ones, leveraging alternative distribution divergences combined with suitable density estimation techniques. (2) Enhancing existing algorithms within a particular family by increasing their model capacity, e.g., IQN and FQF. Concretely, SinkhornDRL falls into the first dimension, aiming to expand the range of distributional RL algorithm families.Sinkhorn divergence [45] efficiently approximates the optimal transport problem by introducing an entropic regularization. It aims at finding a sweet trade-off that simultaneously leverages the geometry property of Wasserstein distance (optimal transport distances) and the favorable sample complexity advantage and unbiased gradient estimates of MMD [25, 21] . For two probability measures µ and ν, the entropic regularized Wasserstein distance W c,ε (µ, ν) is formulated asEQUATIONwhere the entropic regularization KL(Π|µ ⊗ ν) = log Π(x,y) dµ(x)dν(y) dΠ(x, y), also known as mutual information, makes the optimization strongly convex and differential [3, 21] , allowing for efficient matrix scaling algorithms for approximation, such as Sinkhorn Iterations [45] . In statistical physics, W c,ε (µ, ν) can be re-factored as a projection problem:EQUATION)where K is the Gibbs distribution and its density function satisfies dK(x, y) = e -c(x,y)/ε dµ(x)dν(y). This problem is often referred to as the "static Schrödinger problem" [29, 44] as it was initially considered in statistical physics. Formally, the Sinkhorn divergence is defined asEQUATIONwhich is smooth, positive definite, and metricizes the convergence in law [21] . This definition subtracts two self-distance terms to ensure non-negativity and metric properties.Properties for Convergence. The contraction analysis of distributional Bellman operator T π under a distribution divergence d p depends on its scale sensitive (S) and sum invariant (I) properties [6, 5] . We say d p is scale sensitive (of order τ ) if there exists a τ > 0, such that for all random variables X, Y and a real value a > 0, d p (aX, aY ) ≤ |a| τ d p (X, Y ). d p has the sum invariant property if whenever a random variable A is independent from X, Y , we have dp (A + X, A + Y ) ≤ d p (X, Y ).Based on these properties, [5] shows that T π is γ-contractive under the supremal form of Wasserstein distance W p , which is regarding the first term of W c,ε or directly letting ε = 0 in Eq. 1. When examining the regularized loss form of W c,ε , a natural question arises: What is the influence of the incorporated regularization term on the contraction of T π ? We begin to address this question in Proposition 1, focusing on the separate regularization term. Here, we define mutual information as MI Π (µ(s, a), ν(s, a)) = KL(Π|µ(s, a) ⊗ ν(s, a)) and its supremal form MI ∞ Π (µ, ν) = sup (s,a)∈S×A KL(Π|µ(s, a) ⊗ ν(s, a)) given a joint distribution Π. Proposition 1. T π is non-expansive under MI ∞ Π for any non-trivial joint distribution Π.Please refer to Appendix C for the proof, where we investigate both (S) and (I) properties. The non-trivial Π rules out the independence case of µ and ν, where KL(Π|µ ⊗ ν) would degenerate to zero. Although the non-expansive nature of the introduced regularization term, as shown in Proposition 1, may potentially slow the convergence in Sinkhorn divergence compared with W p without the regularization, we will demonstrate that T π is still contractive under the full Sinkhorn divergence in Theorem 1. Before introducing Theorem 1, we first present the sum-invariant and a new variant of scale-sensitive properties in Proposition 2, which acts as the foundation for Theorem 1. Proposition 2. Considering W c,ε with the unrectified kernel k α := -∥x -y∥ α as -c (α > 0) and a scaling factor a ∈ (0, 1), W c,ε is sum-invariant (I) and satisfies W c,ε (aµ, aν) ≤ ∆ ε (a, α)W c,ε (µ, ν) (S) with a scaling constant ∆ ε (a, α) ∈ (|a| α , 1) for any µ and ν in a finite set of probability measures.Proof Sketch. The detailed proof is provided in Appendix D. Let Π * be the optimal coupling of W c,ε , we define a ratio λ ε (µ, ν) that satisfies λ ε (µ, ν) = εKL(Π * |µ⊗ν) Wc,ε ∈ (0, 1) for a generally nonzero W c,ε . The ratio λ ε (µ, ν) measures the proportion of the entropic regularization term over the whole loss term W c,ε . Therefore, the contraction factor ∆ ε (a, α) is defined as∆ ε (a, α) = |a| α (1 - sup µ,ν λ ε (µ, ν)) + sup U,V λ ε (µ, ν)) ∈ (|a| α , 1) with sup µ,ν λ ε (µ, ν) < 1,which is determined by the scale factor a, the order α, the hyperparameter ε, and the set of interested probability measures.Contraction Guarantee and Interpolation Relationship. Proposition 2 reveals that W c,ε with an unrectified kernel satisfies (I) and a variant of (S) properties. While the scaling constant ∆ ε (a, α) in (S) has a complicated form, it remains strictly less than one, even considering a non-expansive nature of the entropic regularization as shown in Proposition 1. We denote the supremal form of Sinkhorn divergence asW ∞ c,ε (µ, ν) : W ∞ c,ε (µ, ν) = sup (s,a)∈S×A W c,ε (µ(s, a), ν(s, a)).In Theorem 1, we will integrate all these properties to demonstrate the contraction property of distributional dynamic programming under W c,ε , specifically highlighting the interpolation property of Sinkhorn divergence between MMD and Wasserstein distance in the context of distributional RL. Theorem 1. Considering W c,ε (µ, ν) with an unrectified kernel k α := -∥x -y∥ α as -c (α > 0), where µ, ν ∈ the distribution set of {Z π (s, a)} for s ∈ S, a ∈ A in a finite MDP. We define the ratioλ ε (µ, ν) as λ ε (µ, ν) = εKL(Π * |µ⊗ν)Wc,ϵ(µ,ν) ∈ (0, 1) with sup µ,ν λ ε (µ, ν) < 1. Then, we have:(1) (ε → 0) W c,ε (µ, ν) → 2W α α (µ, ν). When ε = 0, T π is γ α -contractive under W ∞ c,ε .(EQUATIONEQUATIONProof Sketch. The detailed proof of Theorem 1 can be found in Appendix E. Theorem 1 (1) and ( 2) are follow-up conclusions in terms of the convergence behavior of T π based on the interpolation relationship between Sinkhorn divergence with Wasserstein distance and MMD [25] . We also provide a rigorous analysis within the context of distributional RL for completeness. Our critical theoretical contribution is the part (3) for the general ε ∈ (0, ∞), where we show that T π is at least a ∆ ε (γ, α)contractive operator. The contraction factor ∆ ε (γ, α) ∈ (γ α , 1) depends on the return distribution set {Z π (s, a)} of the considered MDP, and it is also a function of γ, ε and α. Due to the influence of the regularization term in Sinkhorn loss, ∆ ε (γ, α) is larger than |γ| α , the contraction factor for Wasserstein distance without the regularization. Thus, ∆ ε (γ, α) can be seen as an interpolation between γ α and 1, with the coefficient sup µ,ν λ ε (µ, ν) ∈ (0, 1) defined in Theorem 1. The ratio λ ε (µ, ν) measures the proportion of the KL regularization term relative to W c,ε . As ε → 0 or +∞, sup µ,ν λ ε (µ, ν) → 0, leading to γ α -contraction. This aligns with parts (1) and (2).Wasserstein distance and MMD, its contraction property for ε ∈ [0, ∞] also aligns well with the existing distributional RL algorithms when c = -k α . It is worth noting that using Gaussian kernels in the cost function does not yield concise or consistent contraction results like those in Theorem 1 (3). This conclusion is consistent with MMD-DQN [37] (ε → +∞), where T π is generally not a contraction operator under MMD with Gaussian kernels, as counterexamples exist (Theorem 2) in [37] . Guided by our theoretical results, we employ the rectified kernel k α as the cost function and set α = 2 in our experiments, ensuring that T π retains the contraction property guaranteed by Theorem 1 (3) . In Table 1 , we also summarize the main properties of distribution divergences in typical distributional RL algorithms, including the convergence rate of T π and sample complexity, γ O(n -1 d ) MMD-DQN MMD Samples γ α/2 (k α ) O(n -1 ) SinkhornDRL (ours) Sinkhorn divergence (c = -k α ) Samples γ (ε → 0) γ α/2 (ε → ∞) O(n e κ ε ε ⌊d/2⌋ √ n ) (ε → 0) O(n -1 2 ) (ε → ∞)Table 1 : Properties of different distribution divergences in typical distributional RL algorithms. d is the sample dimension and κ = 2βd + ∥c∥ ∞ , where the cost function c is β-Lipschitz [24] . Sample complexity is improved to O(1/n) using the kernel herding technique [10] in MMD.i.e., the convergence rate of a given metric between a measure and its empirical counterpart as a function of the number of samples n.As the ability to extend to the multi-dimensional reward setting is one of the major advantages of SinkhornDRL over quantile regression-based algorithms, we next demonstrate that the joint distributional Bellman operator in the multi-dimensional reward case is contractive under Sinkhorn divergence W ∞ c,ε . First, we define a d-dimensional reward function as R : S × A → P (R d ), where d represents the number of reward sources. Consequently, we have joint return distributions of the d-dimensionalreturn vector Z π (s, a) = ∞ t=0 γ t R(s t , a t ), where Z π (s, a) = (Z π 1 (s, a), • • • , Z π d (s, a)) ⊤ .The joint distributional Bellman operator T π d applied on the joint distribution of the random vector Z(s, a) is defined as T π d Z(s, a) :EQUATIONPlease refer to Appendix F for the proof. The contraction guarantee of Sinkhorn divergence enables us to effectively deploy our SinkhornDRL algorithm in various RL tasks that involve multiple sources of rewards [32, 15] , hybrid reward architecture [50, 30] , or sub-reward structures after reward decomposition [31, 57] . We compare SinkhornDRL with MMD-DQN in multiple reward sources setting in Section 5.3, where SinkhornDRL significantly outperforms MMD-DQN by leveraging its ability to capture richer data geometry, a key advantage of optimal transport distances.Equipping Sinkhorn Divergence and Particle Representation. The key to applying Sinkhorn divergence in distributional RL is to leverage the Sinkhorn loss W c,ε to measure the distance between the current action-return distribution Z θ (s, a) and the target distribution T π Z θ (s, a). This yields W c,ε (Z θ (s, a), T π Z θ (s, a)) for each s, a pair. For the representation of Z θ (s, a), we employ the unrestricted statistics, i.e., deterministic samples, akin to MMD-DQN, instead of predefined statistic functionals like quantiles in QR-DQN or categorical distributions in C51. More concretely, we use neural networks to generate samples to approximate the return distributions, expressed asZ θ (s, a) := {Z θ (s, a) i } N i=1, where N is the number of generated samples. We refer to these samples {Z θ (s, a) i } N i=1 as particles. We then use the Dirac mixture a * ← arg max a ′ ∈A 1 N N i=1 Z θ (s ′ , a ′ ) i 2: TZ i ← r + γZ θ * (s ′ , a * ) i , ∀1 ≤ i ≤ N Output: W c,ε {Z θ (s, a) i } N i=1 , {TZ j } N j=1true density function of Z π (s, a), thus minimizing the Sinkhorn divergence between the approximate distribution and its distributional Bellman target. A generic Sinkhorn distributional RL algorithm with particle representation is provided in Algorithm 1.Efficient Approximation via Sinkhorn Iterations with Guarantee. By introducing an entropy regularization, Sinkhorn divergence renders optimal transport computationally feasible, especially in the high-dimensional space, via efficient algorithms, e.g., Sinkhorn Iterations [45, 25] . Notably, Sinkhorn iteration with L steps yields a differentiable and solvable efficient loss function as the main burden is the matrix-vector multiplication, which streams well on the GPU by simply adding extra differentiable layers on the typical deep neural network, such as a DQN architecture. It has been proven that Sinkhorn iterations asymptotically converge to the true loss in a linear rate [25, 22, 12, 27] . We provide a detailed description of Sinkhorn iterations in Algorithm 2 and a full version in Algorithm 3 of Appendix G. In practice, selecting proper values of L and ε is crucial. To this end, we conduct a rigorous sensitivity analysis, detailed in Section 5.Remark: Relationship with IQN and FQF. In the realm of distributional RL algorithms, it is important to highlight that QR-DQN and MMD-DQN are direct counterparts to SinkhornDRL within the first dimension of algorithmic evolution. In contrast, IQN and FQF enhance QR-DQN and position them in the second modeling dimension, which are orthogonal to our work. As discussed in [37] , the techniques from IQN and FQF can naturally extend both MMD-DQN and SinkhornDRL. For instance, we can implicitly generate {Z θ (s, a) i } N i=1 by applying a neural network to N samples of a base sampling distribution, as in IQN. We can also use a proposal network to learn the weights of each generated sample as in FQF. We leave these modeling extensions as future works and our current study focuses on rigorously investigating the simplest modeling choice via Sinkhorn divergence. 
experiments: 
 We substantiate the effectiveness of SinkhornDRL as described in Algorithm 1 on the entire 55 Atari 2600 games. Without increasing the model capacity for a fair comparison, we leverage the same architecture as QR-DQN and MMD-DQN, and replace the quantiles output in QR-DQN with N particles (samples). In contrast to MMD-DQN, SinkhornDRL only changes the distribution divergence from MMD to Sinkhorn divergence. As such, the potential performance improvement of our algorithm is directly attributed to the theoretical advantages of Sinkhorn divergence over MMD.Baseline Implementation. We choose DQN [36] and three typical distributional RL algorithms as classic baselines, including C51 [5] , QR-DQN [14] and MMD-DQN [37] . For a fair comparison, we build SinkhornDRL and all baselines based on a well-accepted PyTorch implementation2 of distributional RL algorithms. We re-implement MMD-DQN based on its original TensorFlow implementation 3 , and keep the same setting. For example, our MMD-DQN still employs Gaussian kernels k h (x, y) = exp(-(x -y) 2 /h) with the same kernel mixture trick covering a range of bandwidths h as adopted in MMD-DQN [37] .SinkhornDRL Implementation and Hyperparameter Settings. For a fair comparison with QR-DQN, C51, and MMD-DQN, we use the same hyperparameters: the number of generated samples N = 200, Adam optimizer with lr = 0.00005, ϵ Adam = 0.01/32. In SinkhornDRL, we choose the number of Sinkhorn iterations L = 10 and smoothing hyperparameter ε = 10.0 in Section 5.1 after conducting sensitivity analysis in Section 5.2. Guided by the contraction guarantee analyzed in Theorem 1, we use the unrectified kernel, specifically setting -c = k α and choosing α = 2. This choice ensures our implementation is consistent with the theoretical results regarding the contraction guarantee in Theorem 1 (3) . We evaluate all algorithms on 55 Atari games, averaging results over three seeds. The shade in the learning curves of each game represents the standard deviation.Learning Curves of Human Normalized Scores (HNS). We compare the learning curves of the Mean, Median, and Interquartile Mean (IQM) [1] across all considered distributional RL algorithms in Figure 1 summarized over 55 Atari games. The IQM (x%) computes the mean from the x% to (1 -x)% range of HNS, providing a robust alternative to the Mean that mitigates the impact of extremely high scores on specific games and is more statistically efficient than the Median. For computational feasibility, we evaluate the algorithms over 40M training frames. Our findings reveal that SinkhornDRL achieves state-of-the-art performance in terms of mean, median, and IQM (5%) of HNS across most training phases. Notably, SinkhornDRL exhibits slower convergence during the early training phase, as indicated by the Mean of HNS (left panel of Figure 1 ). This slower initial convergence can be explained by the slower contraction factor ∆ ε (γ, α) > γ α in Theorem 1, as opposed to MMD-DQN. To ensure the reliability of our results, we also provide the learning curves for each Atari game in Figure 6 in Appendix I. Furthermore, a table summarizing all raw scores is available in Table 3 in Appendix J. This table highlights that SinkhornDRL achieves the highest numbers of best and second-best performance of all games among all baseline algorithms. A summary table of Mean, IQM, and Median HNS is also given in Table 2 of Appendix H. Overall, we conclude that SinkhornDRL generally outperforms existing distributional RL algorithms.Ratio Improvement Analysis across All Games. Given the interpolation nature of Sinkhorn divergence between Wasserstein distance and MMD, as analyzed in Theorem 1, a pertinent question arises: In which environments does SinkhornDRL potentially perform better? We empirically address this question by conducting a ratio improvement comparison between SinkhornDRL and both QR-DQN and MMD-DQN across all games. Figure 2 showcases that SinkhornDRL surpasses both QR-DQN and MMD-DQN in more than half of the games and significantly excels at them in a large proportion of games. Notably, the games where SinkhornDRL achieves considerable improvement tend to have larger action spaces and more complex dynamics. In particular, as illustrated in Figure 2 , these games include Venture, Seaquest, Solaris, Tennis, Phoenix, Atlantis, and Zaxxon. Most of these games have an 18-dimensional action space and intricate dynamics, except for Atlantis, which has a 4-dimensional action space and simpler dynamics where MMD-DQN is substantially inferior to SinkhornDRL. For a detailed comparison, we provide the features of all games, including the number of action spaces, and complexity of environment dynamics in Table 4 distribution characteristics more accurately by directly using samples, avoiding the non-crossing issue of learned quantile curves or the potential limitations of quantile representation. 2. Smooth transport plan and stable convergence. The induced smoother transport plan (see Appendix A for visualization) and the inherent smoothness of Sinkhhorn divergence contribute to more stable convergence, leading to performance improvement. In contrast to MMD-DQN, the benefits of SinkhornDRL arise from its richer data representation capability when comparing return distributions, rooted in the OT nature. This is in comparison to the potentially restricted kernel-specific distances, such as MMD.Sensitivity Analysis. In practice, a proper ε is preferable as an overly large or small ε will lead to numerical instability of Sinkhorn iterations in Algorithm 2 (see the discussion in Section 4.4 of [39] ), therefore worsening its performance, as shown in Figure 3 (a) . This implies that the potential interpolation nature of limiting behaviors between SinkhornDRL with QR-DQN and MMD-DQN revealed in Theorem 1 may not be able to be rigorously verified in numerical experiments.SinkhornDRL also requires a proper number of iterations L and samples N . For example, a small N , e.g., N = 2 in Seaquest in Figure 3 (b) leads to the divergence of algorithms, while an overly large N can degrade the performance and meanwhile increases the computational burden (Appendix L.2). We conjecture that using larger networks to represent more samples is more likely to suffer from overfitting, yielding the instability in the RL training [7] . Therefore, we choose N = 200 to attain favorable performance and guarantee computational effectiveness simultaneously. We provide a more detailed sensitivity analysis and more results on StarGunner and Zaxxon in Appendix L.1. Computation Cost. In terms of the computation cost, SinkhornDRL slightly increases the computational overhead compared with C51, QR-DQN, and MMD-DQN. For instance, SinkhornDRL increases the average computational cost compared with MMD-DQN by around 20%. Due to the space limit, we provide more computation cost comparison in terms of L and N in Appendix L.2.Many RL tasks involve modeling multivariate return distributions. Following the multi-dimensional reward setting in [57] , we compare SinkhornDRL with MMD-DQN on six Atari games with multiple sources of rewards. In these tasks, the primitive scalar-based rewards are decomposed into reward vectors based on the respective reward structures (see Appendix M for more details). Figure 4 showcases that SinkhornDRL outperforms MMD-DQN in most cases for multi-dimensional reward functions. Of particular note, it remains an open question to directly approximate multi-dimensional Wasserstein distances via quantile regression or other efficient algorithms, particularly in RL tasks. Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?Answer: [Yes]Justification: We have provided the code in the supplementary files for faithful reproducibility.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
conclusion, limitations and future work: 
 In this work, we propose a novel family of distributional RL algorithms based on Sinkhorn divergence that accomplishes competitive performance compared with the typical distributional RL algorithms on the Atari games suite. Theoretical results about the properties of this regularized Wasserstein loss and its convergence guarantee in the context of RL are provided with rigorous empirical verification.Limitations. While SinkhornDRL achieves competitive performance, it relatively increases the computational cost and requires tuning additional hyperparameters. This hints that the enhanced performance offered by SinkhornDRL may come with slightly greater efforts in practical deployment. Additionally, it remains elusive for a deeper connection between the theoretical properties of divergences and the practical performance of distributional RL algorithms given a specific environment.Future work. Along the two dimensions of distributional RL algorithm evolution, we can further improve Sinkhorn distributional RL by incorporating implicit generative models, including parameterizing the cost function and increasing model capacity. Moreover, Sinkhorn distributional RL also opens a door for new applications of Sinkhorn divergence and more optimal transport approaches in RL. It also becomes increasingly crucial to design a quantitative criterion for a given environment to recommend the choice of a specific distribution divergence before conducting costly experiments. We visualize the optimal transport plans by solving Sinkhorn divergence with different ε in welltrained SinkhornDRL models across three games in Figure 5 We evaluate (randomly selected 64) current and target state features to be compared and then apply t-SNE to reduce their dimension from 512 to 2 associated with a normalization for visualization. In each game of Figure 5 , as we increase the regularization strength ε (from right to left), the resulting transport plans tend to be smoother, less concentrated, and more uniformly distributed by transporting the point mass between two distributions (in red and blue). Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes] Justification: We provide the links of the reference code, detailed experimental setting in the appendix, and the code in the supplementary files.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: We report learning curves on Atari games with the shading indicating the standard deviation.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors). • It should be clear whether the error bar is the standard deviation or the standard error of the mean. • It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified. • For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates). • If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text. 
experiments compute resources: 
 Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?Answer: [Yes] Justification: We run our experiments on multiple NVIDIA 3090 Ti GPUs and compare the computational costs of different algorithms we considered in Section 5.2 and Appendix L.Guidelines:• The answer NA means that the paper does not include experiments.• The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage. • The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute. • The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper). 
code of ethics: 
 Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines? 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [NA] Justification: As our study is to propose a new family of reinforcement learning algorithm and we test the performance on widely used Atari environments, we do not think there is any societal impact of our work.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations. • The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We believe our paper poses no such risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: We believe our paper does not use existing assets.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The answer NA means that the paper does not involve crowdsourcing nor research with human subjects. • Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Definition of distances. Given two random variables X and Y , one-dimensional p-Wasserstein metric W p between the distributions of X and Y has a simplified form via the quantile functions:EQUATIONwhich F -1 is the quantile function, also known as inverse cumulative distribution function, of a random variable with the cumulative distribution function as F . The supremal form of W p , denoted by W ∞ p , is defined asEQUATIONFurther, ℓ p distance [19] is defined asEQUATIONThe ℓ p distance and Wasserstein metric are identical at p = 1, but are otherwise distinct. Note that when p = 2, ℓ p distance is also called Cramér distance [6] d C (X, Y ). Also, Cramér distance has a different representation given byEQUATIONwhere X ′ and Y ′ are the i.i.d. copies of X and Y . Energy distance [49, 59] is a natural extension of Cramér distance to the multivariate case, which is defined asEQUATIONwhere X and Y are multivariate. Moreover, the energy distance is a special case of the maximum mean discrepancy (MMD), which is formulated asEQUATIONwhere k(•, •) is a continuous kernel on X . In particular, if k is a trivial kernel, also called the unrectified kernel, MMD degenerates to energy distance. Additionally, we further define the supreme MMD, which is a functional P(X ) S×A × P(XEQUATIONWe further summarize the convergence rates of the distributional Bellman operator T π under different distribution divergences.• T π is γ-contractive under the supreme form of Wassertein distance W p .• 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∥ α , ∀α > 0.Proof of Contraction in Distributional Dynamic Programming.• Contraction under the supreme form of Wasserstein distance is provided in Lemma 3 [5] .• Contraction under supreme form of ℓ p distance can refer to Theorem 3.4 [19] .• Contraction under MMD ∞ is provided in Lemma 6 [37] .Proof. We denote two marginal random variables U and V with the pdf µ(x) and ν(y). We next denote the p Π (x, y) as the pdf for Π in MI Π (U, V ) = KL(Π|U ⊗ V ). We first prove that the MI Π (U, V ) is sum-invariant, which is based on the dual form of KL divergence via the variational representation [17, 2] :EQUATIONwhere L b is the space of bounded measurable functions. The mutual information involves twodimensional random variables in the KL divergence. Let U ′ = a + U and V = a + V with pdf µ ′ and ν ′ , we denote the joint distribution with margins µ ′ (x) = µ(x -a) and ν ′ (y) = ν(y -a) as Π ′ (x, y)EQUATIONwhere (a) is by the independence of A between X, Y , and the joint cdf Π. For instance, in the one-dimensional setting, we have Next, we show that MI Π is NOT scale-sensitive or with the zero-order τ . This result is directly based on the similar property of KL divergence. With a slight abuse of notations, we denote U ′ = aU and V ′ = aV , whose pdfs are µ ′ (x) = 1 a µ( x a ) and ν ′ (y) = 1 a ν( y a ), respectively. The scaled joint distribution Π ′ with the pdf p Π ′ satisfying p Π ′ (x, y) = 1 a 2 p Π (x/a, y/a). Therefore, its marginal distributions are y 1 a 2 p Π (x/a, y/a)dy = 1 a µ( x a ) and x 1 a 2 p Π (x/a, y/a)dy = 1 a ν( y a ). We thus have the following result: E Z=A+X [f (z)] = a x f (x + a)p A (a)p X (x)dxda = E A [E X [f (x +MI Π (aU, aV ) = KL(Π ′ (x, y)|U ′ ⊗ V ′ ) = p Π ′ (x, y) log p Π ′ (x, y) µ ′ (x)ν ′ (y) dxdy = 1 a 2 p Π (x/a, y/a) logEQUATIONPutting the two properties together and given two return distributions Z 1 (s, a) and Z 2 (s, a), we have the non-expansive contraction property of the supremal form of MI Π as follows.EQUATIONwhere (a) relies on the sum invariant property of MI Π and (b) utilizes the non-scale sensitive property of MI Π . By applying the well-known Banach fixed point theorem, we have a unique return distribution when convergence of distributional dynamic programming under MI Π for any non-trivial joint distribution Π.Given two random variables U and V with the marginal distributions as µ and ν, and a random variable A that is independent of them, we aim at provingEQUATIONAccording to [39] , we have the dual form of W c,ε :EQUATIONTherefore, we have:EQUATIONwhere (a) relies on the same techniques used in the proof of Eq. 13 in Appendix C, (b) utilizes the Jensen inequality of sup, and (c) is based on the fact that the translation operator is still within the same functional space of φ, ψ.General Conclusion. Let Π * be the optimal coupling for W c,ε , we define a ratio λ ε (U, V ) = εKL(Π * |µ⊗ν) Wc,ε ∈ (0, 1) for any considered U, V with measures µ, ν to compare, where the denominator W c,ε is generally non-zero. We thus have the following result:EQUATIONwhere the scaling factor ∆ ε (a, α)= |a| α (1-sup U,V λ ε (U, V ))+sup U,V λ ε (U, V )) ∈ (|a| α , 1) with sup U,V λ ε (U, V ) > 0.The ratio λ ε (U, V ) measures the proportion of the entropic regularization term over the whole divergence term W c,ε , i.e., λ ε (U, V ) = εKL(Π * |µ⊗ν) Wc,ε ∈ (0, 1). Under the mild assumption of a finite set of probability measures, we have sup U,V λ ε (U, V ) > 0. To elaborate the reason behind it, we first know that λ ε (U, V ) < 1 for any U and V with their measures on the probability measure set. If this set is finite, the ratio set that contains all {λ ε (U, V )} is also finite. Based on the fact that the real set is dense, we can directly find a positive lower bound λ * for the ratio set, such that {λε (U, V )} ≤ λ * < 1. This implies that sup U,V λ ε (U, V ) = max U,V λ ε (U, V ) < 1.Notably, this finite set property of the ratio avoids the extreme case that may lead to a conservative conclusion about a non-expansive distribution Bellman operator, which we will give more details later.Scale-sensitive Property. By definition of Sinkhorn divergence [18, 39] , the pdf of Gibbs kernel in the equivalent form of Sinkhorn divergence is K(U, V ), which satisfies K(U, V ) ∝ e -c(x,y) ε µ(x)ν(y). In particular, the pdf of Gibbs kernel is defined asdK d(µ⊗ν) (x, y) = exp(-c/ε) exp(-c/ε)d(µ⊗ν), where the denominator is the normalization factor. After a scaling transformation, the pdf of aU and aV with respect to x and y would be 1 a µ( x a ) and 1 a ν( y a ). Thus K(aU, aV ) ∝ e -c(x,y) ε 1 a µ( x a ) 1 a ν( y a ). In the following proof, we use the change variable formula (multivariate version) constantly, while changing the joint pdf π(x, y) and keep the cost function term c(x, y). In particular, we denote Π * and Π 0 as the optimal joint distribution of W c,ε (µ, ν) and W c,ε (aµ, aν). Then we have:W c,ε (aU, aV ) = c(x, y)dΠ 0 (x, y) + εKL(Π 0 |aµ ⊗ aν) ≤ c(x, y)dΠ * (x, y) + εKL(Π * |aµ ⊗ aν) c=-kα = (x -y) α 1 a 2 π * ( x a , y a )dxdy + ε 1 a 2 π * ( x a , y a ) log 1 a 2 π * ( x a , y a ) 1 a 2 µ( x a )ν( y a ) dxdy = |a| α (x -y) α π * (x, y)dxdy + ε π * (x, y) log π * (x, y) µ(x)ν(y) dxdy = |a| α (x -y) α π * (x, y)dxdy + (|a| α + 1 -|a| α )ε π * (x, y) log π * (x, y) µ(x)ν(y) dxdy = |a| α W c,ε (U, V ) + (1 -|a| α )εKL(Π * |µ ⊗ ν) = ∆ U,V ε (a, α)W c,ε (U, V ) (20) where ∆ U,V ε (a, α) = |a| α + (1 -|a| α )λ ε (U, V ) = |a| α (1 -λ ε (U, V )) + λ ε (U, V ) ∈ (|a| α , 1)for ε ∈ (0, +∞) and a < 1 due to the fact that λ ε (U, V ) ∈ (0, 1) for any non-trivial W c,ε (U, V ). The non-trivial W c,ε (U, V ) rules out the case when the regularization term is zero, e.g., ϵ = 0 or the optimal coupling is the product of two margins. In other words, ∆ U,V ε (a, α) is a function less than 1, which depends on the two margins, including their independence and distribution similarity, the scale factor a and the order α.Ruling Out Extreme Cases in the Convergence via a Finite Set. However, the fact that ∆ U,V ε (a, α) < 1 can only guarantee a "conservative" non-expansive contraction rather than a desirable contraction of the distributional Bellman operator. This is because there will be extreme cases in the power of series in general, although it is very unlikely to occur given a certain MDP in practice. For example, denote the non-constant factor as q k for the k-th distributional Bellman update, where q k < 1. We can construct a counterexample as q k = 1 -1/(k + 2) 2 . In this case,EQUATION4 ) • • • > 0 instead of the convergence to 0 and the non-zero limit can not guarantee the contraction. It also intuitively implies that iteratively applying distribution Bellman operator under W c,ε may not lead to convergence in general by considering all possible return distributions given the non-constant factor ∆ U,V ε (a, α). Although we know these extreme cases are very unlikely to happen, we have to rule out these extreme cases for a rigorous proof. As we have the assumption of a finite set of probability measures, the set of {λ ε (U, V )} is also finite. As the real set is dense, we can always find a positive constant that can be used as the contraction factor. Alternatively, we can directly use the sup U,V λ ε (U, V ) as the uniform upper bound across the whole set of interested probability measures. Under this condition, we can immediately find a universal upper bound ofEQUATIONwhere the upper bound sup U,V ∆ U,V ε (a, α) has an interpolation form, which can be viewed as the convex combination between |a| α and 1 with the coefficient sup U,V λ ε (U, V ) determined by the probability measure set. More importantly, sup U,V ∆ U,V ε (a, α) is strictly less than 1, which is guaranteed by the finite set of {λ ε (U, V )} thanks to a finite set of interested probability measures. Finally, we have the variant of scale-sensitive property as follows, where the factor ∆ ε (a, α) depends on α, a and the probability measure set.EQUATIONE Proof of Theorem 1E.1 ε → 0 and c = -k α .We study the uniform convergence when ε → 0. The proof is summarized from the optimal transport literature [25, 21] and we here provide the detailed proof for completeness. On the one hand, W c,ε ≥ (x -y) α dΠ * (x, y)dxdy ≥ W α α as KL ≥ 0. We want to provide the inequality on the other side. Denote Π ′ as the minimizer in the Wasserstein distance W α α . For any δ > 0, there always exists a joint distribution Π δ such thatEQUATIONand KL(Π δ |µ ⊗ ν) < +∞, i.e., (x -y) α dΠ δ (x, y) -(x -y) α dΠ ′ (x, y) ≤ δ. One possible way to find Π δ is provided in notes of Lecture 6 in Optimal Transport Course4 and we invite interested readers for reference. It follows thatEQUATIONwhere the RHS (x -y) α dΠ ′ (x, y) + δ + εKL(Π δ |µ ⊗ ν) → (x -y) α dΠ ′ (x, y) + δ = W α α + δ as ε → 0. As δ > 0 is arbitrary, combing the two sides, it shows that W c,ϵ → W α α as ε → 0. Thus, Sinkhorn divergence maintains the properties of Wasserstein distance when ε → 0.When ε = 0, it has been shown that W α can guarantee a γ-contraction property for distributional Bellman operator [5] . The crux of proof is that W α is γ-scale sensitive:EQUATIONwhere the inequality comes from the change of optimal joint distribution. Therefore, W α (aU, aV ) ≤ aW α (U, V ) guarantees a γ-contraction property for the distributional Bellman operator. As such, for W α α , when ε = 0, it suggest that W c,0 = W α α corresponds to a γ α -contraction for the distributional Bellman operator T π .E.2 ε → ∞ and c = -k α .Our complete proof is inspired by [40, 25] . Recap the definition of squared MMD isEQUATIONWhen the kernel function k degenerates to an unrectified k α (x, y) := -∥x -y∥ α for α ∈ (0, 2), the squared MMD would degenerate toEQUATION)where X, X ′ i.i.d. ∼ µ, Y, Y ′ i.i.d. ∼ ν and X, X ′ , Y, Y ′ are mutually independent. On the other hand, by definition, we have the Sinkhorn loss asEQUATION)Denoting Π ε be the unique minimizer for W c,ε , it holds that Π ε → µ ⊗ ν as ε → ∞, which is the product of two marginal distributions. That being said, W c,∞ (µ, ν) → c(x, y)dµ(x)dν(y) + 0 = c(x, y)dµ(x)dν(y). One important proof insight here is although ε → +∞, the KL term tends to zero, which is faster than ε. Therefore, the whole regularization term still tends to 0 as ε → +∞. If c = -k α = ∥x-y∥ α , we eventually have W -kα,∞ (µ, ν) → ∥x-y∥ α dµ(x)dν(y) = E∥X-Y∥ α , where µ and ν can be inherently correlated, although the minimizer degenerates to the product of the two marginal distributions. Finally, we can haveEQUATION)which is exactly the form of squared MMD with the unrectified kernel k α . Now the key is to prove that Π ε → µ ⊗ ν as ε → ∞. We give the detailed proof as follows.Firstly, it is apparent that W c,ε (µ, ν) ≤ c(x, y)dµ(x)dν(y) as µ ⊗ ν ∈ Π(µ, ν). Let {ε k } be a positive sequence that diverges to ∞, and Π k be the corresponding sequence of unique minimizers for W c,ε . According to the optimality condition, it must be the case that c(x, y)dΠ k + ε k KL(Π k , µ ⊗ ν) ≤ c(x, y)dµ ⊗ ν + 0 (when Π(µ, ν) = µ ⊗ ν). Thus, KL (Π k , µ ⊗ ν) ⩽ 1 ε k c dµ ⊗ ν -c dΠ k → 0.Besides, by the compactness of Π(µ, ν), we can extract a converging subsequenceΠ n k → Π ∞ .Since KL is weakly lower-semicontinuous, it holds thatKL (Π ∞ , µ ⊗ ν) ⩽ lim k→∞ inf KL (Π n k , µ ⊗ ν) = 0Hence Π ∞ = µ ⊗ ν. That being said that the optimal coupling is simply the product of the marginals, indicating that Π ε → µ ⊗ ν as ε → ∞. As a special case, when α = 1, W -k1,∞ (u, v) is equivalent to the energy distanceEQUATIONIn summary, if the cost function is the rectified kernel k α , it is the case that W -kα,ε converges to the squared MMD as ε → ∞. According to [37] , T π is γ α/2 -contractive in the supremal form of MMD with the unrectified kernel k α . As W c,ε (µ, ν) → MMD 2 kα (µ, ν), which is a squared MMD instead of MMD, it implies that T π is γ α -contractive under the squared MMD / W c,+∞ .In the proof of Proposition 2, we have shown that the Sinkhorn loss W c,ε satisfies the sum-invariant (I) and a new variant of scale-sensitive properties as follows:EQUATION)The Sinkhorn divergence W c,ε is defined by additionally subtracting two self-distance terms (W c,ε (µ, µ) and W c,ε (ν, ν)) based on W c,ε (µ, ν) in order to guarantee the non-negativity, triangularity and metric properties. These two self-distance terms do not change the (I) and (S)properties when extending W c,ε to W c,ε , and some proof techniques can refer to Section 2 in [21] .The only difference is that the scaling factor will be ∆ U,V ε (a, α), which is the counterpart of Eq. 20 satisfyingEQUATION)where∆ U,V ε (a, α) = |a| α (1 -λ ε (U, V )) + λ ε (U, V ) ∈ (|a| α , 1)for ε ∈ (0, +∞) and a < 1 due to the fact that λ ε (U, V ) ∈ (0, 1) for any non-trivial W c,ε (U, V ). The new ratio λ ε (U, V ) = εKL(Π * |µ⊗ν) Wc,ε ∈ (0, 1) for any considered U, V with measures µ, ν in the interested probability measure set. In particular, in the context of distributional RL, the set of interested probability measures would be the return distribution set of {Z(s, a)} for s ∈ S and a ∈ A in a given finite MDP. We now want to find the universal upper bound ∆ ε (a, α), which is defined asEQUATION)Following the proof in Appendix D, the finite MDP guarantees a finite ratio set of {λ ε (U, V )}, and thus we can find a universal upper bound λ * of the ratio set such that {λ ε (U, V )} ≤ λ * < 1. This also implies that sup U,V λ ε (U, V ) ∈ (0, 1) and thus the scaling factor ∆ ε (a, α) ∈ (|a| α , 1), which is strictly less than 1. Therefore, we have the (I) and (S) properties of W c,ε :EQUATION)Putting all together, we now derive the convergence of distributional Bellman operator T π under the supreme form ofW c,ε , i.e., W ∞ c,ε : W ∞ c,ε (T π Z 1 , T π Z 2 ) = sup s,a W c,ε (T π Z 1 (s, a), T π Z 2 (s, a)) = sup s,a W c,ε (R(s, a) + γZ 1 (s ′ , a ′ ), R(s, a) + γZ 2 (s ′ , a ′ )) (a) ≤ W c,ε (γZ 1 (s ′ , a ′ ), γZ 2 (s ′ , a ′ )) (b) ≤ ∆ Z1(s ′ ,a ′ ),Z2(s ′ ,a ′ ) ε (γ, α)W c,ε (Z 1 (s ′ , a ′ ), Z 2 (s ′ , a ′ )) ≤ sup s ′ ,a ′ ∆ Z1(s ′ ,a ′ ),Z2(s ′ ,a ′ ) ε (γ, α) sup s ′ ,a ′ W c,ε (Z 1 (s ′ , a ′ ), Z 2 (s ′ , a ′ )) = ∆ ε (γ, α)W ∞ c,ε (Z 1 , Z 2 ) (35)where the inequality (a) is based on the sum invariant property (I) of Sinkhorn divergence. (b) is based on the new variant of scale-sensitive property (S) of Sinkhorn divergence and the leverage of c = -k α . Notably, ∆ ε (γ, α) ∈ (|γ| α , 1) is an MDP-dependent constant (determined by the return distribution set), which is also determined by γ, ε and α. As such, we conclude that distributional Bellman operator is at least ∆ ε (γ, α)-contractive, where the contraction factor ∆ ε (γ, α) is strictly less than 1 in a given finite MDP. Based on the existing Banach fixed point theorem, we have a unique optimal return distribution by applying the distributional Bellman operator T π in the distributional dynamic programming when convergence.Proof. The contraction conclusion that extends to the multi-dimensional return distributions is straightforward. As the definition of Sinkhorn divergence inherently allows the multi-dimensional measures, the sum-invariant and the variant of scale-sensitive properties hold naturally. Specifically, after recapping to proof of these properties, we only need to change c(x, y) = (x -y) α to c(x, y) = ∥x-y∥ α and re-define two d-dimensional random vector U and V with the d-dimensional probability measure µ and ν. Therefore, the (I) and (S) properties in the multi-dimensional reward settings are:EQUATION)where A is a d-dimensional random vector independent of U and V.By leveraging these two properties, we now derive the convergence of distributional Bellman operator T π d under W ∞ c,ε in the joint return distribution setting. Given two d-dimensional return distributions Z 1 and Z 2 , we haveEQUATIONwhere the inequality (a) is based on the sum invariant property (I) of Sinkhorn divergence that cancels the additive d-dimensional random vector R(s, a). (b) is based on the new variant of scale-sensitive property (S) of Sinkhorn divergence and the leverage of c = -k α , where the contraction factor ∆ ε (γ, α) will depend on the set of d-dimensional probability measures/distributions. Notably, the analysis of ∆ ε (γ, α) in the one-dimensional return setting established in Appendix D and Appendix E is also applicable in the multi-dimensional setting.G Algorithm: Sinkhorn Iterations and Sinkhorn Distributional RLAlgorithm 2 Sinkhorn Iterations to Approximate W c,ε {Z i } N i=1 , {TZ j } N j=1Input: Two samples sequences{Z i } N i=1 , {TZ j } N j=1, number of iterations L and hyperparameter ε.1: ĉi,j = c(Z i , TZ j ) for ∀i = 1, ..., N, j = 1, ..., N 2: K i,j = exp(-ĉ i,j /ε) 3: b 0 ← 1 N 4: for l = 1, 2, ..., L do 5:a l ← 1 N Kb l-1 , b l ← 1 N Ka l 6: end for 7: W c,ε {Z i } N i=1 , {TZ j } N j=1 = ⟨(K ⊙ ĉ)b, a⟩ Return: W c,ε {Z i } N i=1 , {TZ j } N j=1Given two sample sequences{Z i } N i=1 , {TZ j } N j=1in the distributional RL algorithm, the optimal transport distance is equivalent to the form:EQUATIONwhere the empirical cost function is ĉi,j = c(Z i , TZ j ). By adding entropic regularization on optimal transport distance, Sinkhorn divergence can be viewed to restrict the search space of P in the following scaling form:EQUATIONwhere K i,j = e -ĉi,j /ε is the Gibbs kernel defined in Eq. 2. This allows us to leverage iterations regarding the vectors a and b. More specifically, we initialize b 0 = 1 N , and then the Sinkhorn iterations are expressed asEQUATIONwhere Randomly draw a batch of transition samples (s, a, r, s ′ ) from the replay buffer M 6:Compute a greedy action:a * = arg max a ′ ∈A 1 N N i=1 Z θ * (s ′ , a ′ ) i 7:Compute the target Bellman return distribution:TZ i ← r + γZ θ * (s ′ , a * ) i , ∀1 ≤ i ≤ N 8:Evaluate Sinkhorn divergence via Sinkhorn Iterations in Algorithm 2:W c,ε {Z θ (s, a) i } N i=1 , {TZ j } N j=1 9:Update the main network Bold denotes the best performance, while the underline represents the second best performance.Z θ : θ ← θ -α∇ θ W c,ε {Z θ (s, a) i } N i=1 , {TZ j } N j=1The number of games with the best and second best performance substantiate the superiority of our SinkhornDRL across all considered baseline algorithms.L Sensitivity Analysis and Computational Cost Decreasing ε. We argue that the limit behavior connection as stated in Theorem 1 may not be able to be verified rigorously via numeral experiments due to the numerical instability of Sinkhorn Iteration in Algorithm 2. From Figure 7 (a), we can observe that if we gradually decline ε to 0, SinkhornDRL's performance tends to degrade and approach QR-DQN. Note that an overly small ε will lead to a trivial almost 0 K i,j in Sinkhorn iteration in Algorithm 2, and will cause 1 0 numerical instability issue for a l and b l in Line 5 of Algorithm 2. In addition, we also conducted experiments on Seaquest, a similar result is also observed in Figure 7 (d). As shown in Figure 7 (d), the performance of SinkhornDRL is robust when ε = 10, 100, 500, but a small ϵ = 1 tends to worsen the performance.Increasing ε. Moreover, for breakout, if we increase ε, the performance of SinkhornDRL tends to degrade and be close to MMD-DQN as suggested in Figure 7 (b). It is also noted that an overly large ε will let the K i,j explode to ∞. This also leads to the numerical instability issue in Sinkhorn iteration in Algorithm 2.Samples N . We find that SinkhornDRL requires a proper number of samples N to perform favorably, and the sensitivity w.r.t N depends on the environment. As suggested in Figure 8 (a), a smaller N , e.g., N = 2 on breakout has already achieved favorable performance and even accelerates the convergence in the early phase, while N = 2 on Seaquest will lead to the divergence issue. Meanwhile, an overly large N worsens the performance across two games. We conjecture that using larger network networks to generate more samples may suffer from the overfitting issue, yielding the training instability [7] . In practice, we choose a proper number of samples, i.e., N = 200 across all games. More Games on StarGunner and Zaxxon. Beyond Breakout and Seaquest, we also provide sensitivity analysis on StarGunner and Zaxxon games in Figure 9 . It suggests overly small samples, e.g., 1 and overall large samples tend to degrade the performance, especially on Zaxxon. Although the two games are robust to ε, and we find a small or large ε hurts the performance in Seaquest. Thus, considering all games, we set samples 200, and ε = 10.0 in a moderate range across all games, although a more careful tuning in each game will improve the performance further. We evaluate the computational time every 10,000 iterations across the whole training process of all considered distributional RL algorithms and make a comparison in Figure 10 . It suggests that SinkhornDRL indeed increases around 50% computation cost compared with QR-DQN and C51, but only slightly increases the cost in contrast to MMD-DQN on both Breakout and Qbert games. We argue that this additional computational burden can be tolerant given the significant outperformance of SinkhornDRL in a large number of environments.In addition, we also find that the number of Sinkhorn iterations L is negligible to the computation cost, while an overly large sample N , e.g., 500, will lead to a large computational burden as illustrated in Figure 11 . This can be intuitively explained as the computation complexity of the cost function c i,j is O(N 2 ) in SinkhornDRL, which is particularly heavy in the computation if N is large enough. Reward Structure and Decomposition. In practice, the reward function can be multidimensional [50, 30, 32, 15, 57, 31] , where distributional RL is aimed at modeling multivariate return distribution with multiple reward sources. We follow the multi-dimensional return distribution setting in [57] , which construct six Atari games with multiple sources of rewards by decomposing the scalar-valued primitive rewards into multi-dimension. For completeness, we introduce the respective reward structure and the decomposing method of the six considered Atari games, including AirRaid, Asteroids, Gopher, MsPacman, UpNDown, and Pong. The reward is decomposed while keeping the total reward unchanged.• AirRaid. For primitive rewards, the agent kills different kinds of monsters and then receive discrete values of the rewards. The scalar-based primitive rewards are decomposed into four dimensions. The agent will get multi-dimensional rewards [100, 0, 0, 0], [0, 75, 0, 0], [0, 0, 50, 0],[0, 0, 0, 25], [0, 0, 0, 0] respectively for the primitive reward 100, 75, 50, 25 and 0.• Asteroids. For primitive rewards, the agent kills different kinds of monsters and then receive values of the rewards. We denote the primitive reward as r, and decompose it into the three-dimensional reward as [r 1 , r 2 , r 3 ]. If (r -20) mod 50 = 0, we let r 1 = 20, otherwise r 1 = 0. If (r -r 1 -50) mod 100 = 0, we let r 2 = 50, otherwise r 2 = 0. We let r 3 = r -r 1 -r 2 .• Gopher. For primitive rewards, the agent gets +80 reward for killing a monster and +20 reward after removing holes on the ground. We denote the primitive reward as r, and decompose it into the two-dimensions as [r 1 , r 2 , ]. If (r -20) mod 100 = 0, we let r 1 = 20, otherwise r 1 = 0. We let r 2 = r -r 1 .• MsPacman. For primitive rewards, the agent gets {+200, +400, +800, +1, 600} rewards after killing different monsters and +10 rewards after eating beans. In the reward decomposition, we decompose primitive reward denoted as r into four dimensions [r 1 , r 2 , r 3 , r 4 ].If (r -10) mod 50 = 0, we let r 1 = 10, otherwise r 1 = 0. If (r -r 1 -50) mod 100 = 0, we let r 2 = 50, otherwise r 2 = 0. If (r -r 1 -r 2 -100) mod 200 = 0, we let r 3 = 100, otherwise r 3 = 0. We let r 4 = r -r 1 -r 2 -r 3 .• Pong. For primitive rewards, the agent gets +1 if it wins a round, and -1 for losing the round. We decompose the reward into two-dimension: the agent will get [-1, 0] for a -1 reward, [0, 1] for a +1 reward; otherwise, [0, 0].• UpNDown. For primitive rewards, the agent gets +400 reward for killing an energy car, +100 for reaching a flag, and +10 reward for being alive. We denote the primitive reward as r, and decompose it into the three-dimensional reward as [r 1 , r 2 , r 3 ]. If (r -10) mod 100 = 0, we let r 1 = 10, otherwise r 1 = 0. If (r -r 1 -100) mod 200 = 0, we let r 2 = 100, otherwise r 2 = 0. We let r 3 = r -r 1 -r 2 .Detailed Experimental Setup. Our implementation extends our code in one-dimensional return setting to multi-dimensional return scenario and adopts the key aspects in [57] . For instance, similar to [57] , we leverage a clipping reward normalizer to clip the multi-dimensional rewards into [-1, 1] after applying the reward decomposition procedure mentioned above to the primitive rewards. We keep the same model architecture except only modifying the output of the last layer from (B, |A|, N ) to (B, |A|, D, N ), where B is the batch size within each batch training, and D is the dimension of the decomposed mutivariate reward function in each game.Baseline Algorithms. Quantile regression can be used to approximate 1-Wasserstein distance in one-dimensional setting [14] as the one-dimensional Wassertein distance has a closed-form expression via the quantile function. However, it remains elusive how to use quantile regression to approximate multi-dimensional Wasserstein distance. This is to say, it is still unclear how to extend the quantile regression distributional RL (QR-DQN) into multi-dimensional return distribution setting, resulting in no proper baseline in our experiment. Despite that, we directly compare SinkrhornDRL with MMD-DQN [37] as MMD is applicable and computationally tractable in the multi-dimensional setting. Notably, we did not introduce other baselines, such as Hybrid Reward Architecture (HRA) [50] , or MD3QN [57] . This is because 1) [57] shows that their proposed MD3QN and HRA do not outperform MMD-DQN in most of the six Atari games. By contrast, as suggested in Figure 4 , our SinkhornDRL has already surpassed MMD-DQN across almost all the considered games, and thus excels over MD3QN and HRA, correspondingly.2) The primary focus of our study is the comprehensive advantages of SinkhornDRL over other distributional RL classes, especially in the more common setting within one-dimensional return distributions. The extension capability of SinkhornDRL into the multi-dimensional reward setting is one of its merits, which is not the primary focus of our study.NeurIPS Paper Checklist Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.Justification: We have preserved anonymity and conformed with the NeurIPS Code of Ethics.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).https://github.com/ShangtongZhang/DeepRLhttps://github.com/thanhnguyentang/mmdrlhttps://lchizat.github.io/ot2021orsay.html 


Decision: Distributional RL approximates the full distribution of returns, which contrasts with classic expected-value RL. Quantiles are the method of choice for quantifying the distribution of returns, but are limited in their expressiveness. The paper proposes Sinkhorn distributional RL to address limitations in expressiveness. Theoretically, the paper presents a contraction result and empirically the results show improvements over previous Distributional RL results over 40M training frames. 

The reviewers are uniformly positive about the paper. I have read the paper and agree that the idea is interesting and will be appreciated by the community.

Reviews: review: paper_summary: paper_summary: This paper introduces Sinkhorn Distributional Reinforcement Learning (SinkhornDRL), a new algorithm designed to address the limitations of current distributional RL methods, particularly those relying on quantile regression. Existing methods often struggle with accurately capturing the characteristics of return distributions and extending to scenarios with multi-dimensional rewards. SinkhornDRL leverages Sinkhorn divergence, a regularized Wasserstein loss, to minimize the difference between current and target Bellman return distributions. This approach combines the geometric advantages of Wasserstein distance with the computational efficiency of Maximum Mean Discrepancy (MMD). The paper provides theoretical proof of SinkhornDRL's contraction properties, demonstrating its convergence behavior and relationship to other distance metrics. Empirical evaluations on the Atari games suite show that SinkhornDRL consistently outperforms or matches existing algorithms, especially in settings with multi-dimensional rewards.
strength_weakness: strength_weakness: - SinkhornDRL introduces some additional computational cost compared to simpler methods like C51 and QR-DQN.

- The algorithm requires tuning of additional hyperparameters (e.g., the regularization strength), which might require extra effort in practice.

- The paper acknowledges that a deeper connection between theoretical properties of divergences and practical performance in specific environments remains an open question.

- Some typos:
  - line 151 and line 647 Sinkrhon
  - line 163 supremal from
  - line 787 numeral
questions: questions: Could the authors address the following points:

- Are there existing works on distributional reinforcement learning that use the entropic regularized Wasserstein distance?
- What specific advantages does the Sinkhorn divergence offer over the entropic regularized Wasserstein distance in the context of this study?
limitations: limitations: NA
paper_summary: paper_summary: •	This paper proposes a novel distributional RL algorithm, called SinkhornDRL, which interpolates between Wasserstein distance and MMD. The authors aim to estimate the distribution using unrestricted statistics, enhancing stability and facilitating extension to multi-dimensional reward settings. The authors also provide some theoretical guarantees on convergence of SinkhornDRL. The proposed method shows decent experimental performance in 55 Atari games and several multi-dimensional reward settings.
strength_weakness: strength_weakness: •	Typos: In line 52, Sinkrhorn -> Sinkhorn 

•	The text and figures in the paper are quite dense and difficult to read. Some parts seem unnecessary to be included in the main paper.
 
o	Algorithm 1 lacks certain technical details, and may rather be placed in Appendix. In addition, explaining the details of Algorithm 2 on Line 769 would help a lot in implementing a practical algorithm. 

o	The text in Figure 2 is not legible. It would be better to include a table showing the mean and median of the best HNS scores, as commonly done in other distributional RL papers [1,2].

•	There may be a technical error regarding Theorem 1 (3). The authors state in Eq (33) that a universal upper bound $\bar{\Delta}_{\epsilon}(a,\alpha)$ is strictly less than 1. But if $\mu$ and $\nu$ are close enough, even if non-trivial, doesn’t the scaling factor become 1 by taking the supremum? In Line 660, the authors recognize the non-expansion case, but it seems to rely on the strong assumption that the set $\{\lambda_{\epsilon}(U,V)\}$ is finite. This assumption essentially implies that the set $\{\lambda_{\epsilon}(\(mathcal{T}^{\pi}_D)^n U, (\mathcal{T}^{\pi}_D)^n V )\}$ where $n \in \mathbb{N}$ is finite. 

[1] Dabney, W., Rowland, M., Bellemare, M. G. & Munos, R. Distributional Reinforcement Learning with Quantile Regression. Arxiv (2017).

[2] Hessel, M. et al. Rainbow: Combining Improvements in Deep Reinforcement Learning. arXiv (2017) doi:10.48550/arxiv.1710.02298.
questions: questions: •	Could the author elaborate on the statement in Lines 129-134 that 'SinkhornDRL inherently captures the spatial and geometric layout of return distributions'? Does this imply that the KL divergence term can leverage a richer representation of data geometry?

•	The definition of Sinkhorn divergence for multi-dimensional reward settings is unclear. Specifically, in Line 756, the cost function applies equal weights among reward sources. Can this be generalized to a weighted sum?
limitations: limitations: •	In line 22, the reference [11] is cited as a risk-sensitive control, but this is close to exploration method for distributional RL.
paper_summary: paper_summary: This paper proposes a new RL algorithm that leverages Sinkhorn divergence, which they claimed as a regularized Wasserstein loss. Theoretically, they showed the contraction properties that align with the interpolation nature of Sinkhorn divergence between Wasserstein distance and MMD. Empirically, it outperforms or matches existing algorithms on many atari games.
strength_weakness: strength_weakness: After reading the paper I am still not sure about the benefits of using sinkhorn divergence, although I agree that the authors established some theoretical results (which I also have some concerns about below) and showed experimental results. My question is more intuitive---why should we pick Sinkhorn over others at a high level of intuition? I appreciate that the authors have spent many words explaining it from various perspectives. However, I didn't totally understand them and am not convinced at this moment. Some of my specific confusion are below:

- line 48-51: why does using samples instead of pre-specified statistics have better accuracy? Is there any reference?
- line 55: how is the regularization aligning with the max entropy principle? This point is mentioned multiple times in the paper but not explained in detail.
- line 53: "smoother" is very vague in this context. I would suggest using some different words.

As for the theoretical results, I am not sure how novel they are. For instance, the authors show convergence rate and sample complexity only under limit (when epsilon is approaching zero or infinity). Considering that the algorithm reduces to either QRDQN or MMD under the limit, the results are not so surprising and are actually quite natural. In addition, I didn't find general convergence results (I mean under any value of epsilon). Hence I am doubting if there is anything significant enough in the theory.

For the experiments, it generally looks good to me. I only have a minor comment for now: figure 2 seems a bit misleading--specifically, I feel plotting the ratio "Sinkhorn/QRDQN" is misleading. Think about two cases: if Sinkhorn is twice QRDQN, then the value will be 200%; on the other hand if QRDQN is twice Sinkhorn, then the ratio will be 50%. The two cases are symmetric but 200% will be visually four times higher than 50% in the plot. This may be why the Sinkhorn visually has a huge advantage over others in the figure.
questions: questions: Line 319: the algorithm is much better on complicated games, which is very interesting. In other words, the algorithm is probably inferior on easy games. Do you have more explanation on this?
limitations: limitations: I didn't find any potential negative societal impact.

Comments: title: Please respond to authors
Reviewer Fdqh please respond
comment: Thank you for the detailed response! It looks good to me, so I will maintain my positive score.
Hello Reviewer uu9w: The authors have responded to your comments. I would expect that you would respond in kind.
Hello Reviewer Fdqh: The authors have responded to your comments. I would expect that you would respond in kind.
Thanks! That makes a lot of sense.
We apologize for our misunderstanding. In practice, including in various applications where Sinkhorn divergence is used, the corrected entropic regularization Wasserstein distance $\overline{W}_ {c, \varepsilon}$ is typically preferred over the uncorrected one $W_{c, \varepsilon}$. 

As mentioned in Line 157 in our paper, $\overline{W}_ {c, \varepsilon}$ subtracts two self-distance terms to **correct the bias** introduced by the entropy regularization. This correction is necessary as $W_{c, \varepsilon}$ introduces a bias such that $W_{c, \varepsilon}(\mu, \mu) \neq 0$ for any $\mu$, which is also discussed in  [1] and [2]. Additionally, subtracting the two self-distances also ensures **non-negativity and metric properties**.  

On the contrary, directly using the uncorrected $W_{c, \varepsilon}$ will introduce bias when evaluating the distance between the current and target return distributions during the optimization of distributional RL, thus undermining the algorithm's performance. Therefore, it is essential to leverage the corrected version $\overline{W}_ {c, \varepsilon}$  in real applications to **ensure accurate and unbiased distance measurements**.

### Reference

[1] Feydy, Jean, et al. "Interpolating between optimal transport and mmd using sinkhorn divergences." (AISTATS 2019). 

[2] Genevay, Aude, et al. "Sample complexity of Sinkhorn divergences." (AISTATS 2019)
Thank you for your responses! It seems there was a misunderstanding. For both questions 1 and 2, I was referring to the entropic regularized Wasserstein distance W_{c, \varepsilon}. In the paper, you adopted the Sinkhorn divergence \bar{W}_{c, \varepsilon}. I was curious why the entropic regularized Wasserstein distance W_{c, \varepsilon} was not used instead.