Title: Distributional Reinforcement Learning with Regularized Wasserstein Loss
Venue: NeurIPS 2024
Accepted: yes
Keywords: distributional reinforcement learning

Sinkhorn divergence
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}.}.

Distributional Reinforcement Learning with Regularized Wasserstein Loss

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 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. Code is available in https://github.com/datake/SinkhornDistRL ..

1 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

38th Conference on Neural Information Processing Systems (NeurIPS 2024).

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).

Contributions. In this work, we propose a new family of distributional RL algorithms based on 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.

2 Preliminary Knowledge

2.1 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].

2.2 Divergences between Measures

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].

Maximum Mean Discrepancy [26]. Define two random variables X and Y . The squared Maximum Mean Discrepancy (MMD) MMD 2 k with the kernel k is formulated as MMD 2 k = E [ k ( X,X ′ )] + E [ k ( Y, Y ′ )] -2 E [ 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.

3 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.

4 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.

4.1 Sinkhorn Divergence and New Convergence Properties in Distributional RL

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 as

where 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:

where K is the Gibbs distribution and its density function satisfies d K ( 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 as

which 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 d p ( 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 (Π ∗ | µ ⊗ ν ) W c,ε ∈ (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 as W ∞ 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 (Π ∗ | µ ⊗ ν ) W c,ϵ ( µ,ν ) ∈ (0 , 1) with sup µ,ν λ ε ( µ, ν ) < 1 . Then, we have:

(1) ( ε → 0 ) W c,ε ( µ, ν ) → 2 W α α ( µ, ν ) . When ε = 0 , T π is γ α -contractive under W ∞ c,ε .

(2) ( ε → + ∞ ) W c,ε ( µ, ν ) → MMD 2 k α ( µ, ν ) . When ε = + ∞ , T π is γ α -contractive under W ∞ c,ε .

(3) ( ε ∈ (0 , + ∞ ) ) , T π is at least ∆ ε ( γ, α ) -contractive under W ∞ c,ε , where ∆ ε ( γ, α ) is an MDPdependent constant defined as ∆ ε ( γ, α ) = γ α (1 -sup µ,ν λ ε ( µ, ν )) + sup µ,ν λ ε ( µ, ν )) ∈ ( γ α , 1) .

Proof 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).

Consistency with Existing Contraction Conclusions. As Sinkhorn divergence interpolates between 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,

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 .

4.2 Extension to Multi-dimensional Return Distributions

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 -dimensional return 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 )

Corollary 1. For two joint distributions Z 1 and Z 2 , T π d is ∆ ε ( γ, α ) -contractive under W ∞ c,ε , i.e.,

Please 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.

4.3 SinkhornDRL Algorithm and Approximation

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 as Z θ ( 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 1 N ∑ N i =1 δ Z θ ( s,a ) i to approximate the

Algorithm 1 Generic Sinkhorn distributional RL Update

Require : Number of generated samples N , the cost function c , hyperparameter ε and the target network Z θ ∗ .

Input

: Sample transition ( s, a, r ′ , s ′ )

1: Policy evaluation : a ∗ ∼ π ( ·| s ′ ) or Control : a ∗ ← arg max a ′ ∈A 1 N ∑ N i =1 Z θ ( s ′ , a ′ ) i

2: T Z i ← r + γZ θ ∗ ( s ′ , a ∗ ) i , ∀ 1 ≤ i ≤ N

true 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.

5 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 implementation 2 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 QRDQN, 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.

5.1 Performance of SinkhornDRL

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

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

2

4

6

8

10

12

Mean of HNS

DQN

C51

MMD

Sinkhorn

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0.0

0.2

0.4

0.6

0.8

1.0

Median of HNS

DQN

C51

MMD

Sinkhorn

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

1

2

3

4

5

IQM(5.0%) of HNS

DQN

C51

MMD

Sinkhorn

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 QRDQN 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 of Appendix K.

In summary, compared with QR-DQN, the empirical success of SinkhornDRL can be attributed to several key factors: 1. Enhanced return distribution representation: SinkhornDRL captures return

Venture

Solaris

Seaquest

Tennis

MsPacman

Phoenix

Enduro

Zaxxon

VideoPinball

KungFuMaster

Alien

IceHockey

UpNDown

DoubleDunk

BattleZone

Skiing

Gopher

RoadRunner

Riverraid

StarGunner

Pong

Atlantis

SpaceInvaders

TimePilot

Breakout

PrivateEye

Bowling

Freeway

Boxing

MontezumaRevenge

CrazyClimber

ChopperCommand

Krull

Robotank

Amidar

Asteroids

NameThisGame

BeamRider

Jamesbond

Frostbite

Centipede

BankHeist

Gravitar

Hero

Tutankham

Asterix

YarsRevenge

DemonAttack

WizardOfWor

Kangaroo

Assault

69900

-50

0

50

100

150

200

250

Ratio Improvment (%)

Sinkhorn / QRDQN (%)

(a) SinkhornDRL vs QR-DQN

Venture

Atlantis

VideoPinball

Zaxxon

Tennis

Enduro

Solaris

Gopher

KungFuMaster

MsPacman

SpaceInvaders

DoubleDunk

Gravitar

CrazyClimber

Robotank

IceHockey

RoadRunner

Asteroids

Krull

YarsRevenge

NameThisGame

Amidar

Seaquest

Jamesbond

StarGunner

Alien

Riverraid

Freeway

Boxing

MontezumaRevenge

Pong

PrivateEye

Bowling

Tutankham

Breakout

ChopperCommand

BattleZone

Hero

Phoenix

Assault

Frostbite

BankHeist

UpNDown

TimePilot

Asterix

Centipede

WizardOfWor

DemonAttack

Kangaroo

BeamRider

Skiing

3600

-100

0

100

200

300

400

Ratio Improvment (%)

Sinkhorn / MMD (%)

(b) SinkhornDRL vs MMD-DQN

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.

5.2 Sensitivity Analysis and Computational Cost

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 MMDDQN 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.

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

5000

10000

15000

20000

Average Return

Seaquest

=1

=10

=100

=500

MMD

(a) Hyper-parameter

ε

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

5000

10000

15000

20000

Average Return

Seaquest

Samples=2

Samples=5

Samples=200

Samples=500

(b) Number of Samples

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

100

200

300

400

500

600

Average Return

Breakout

L=2

L=5

L=10

L=50

(c) Sinkhorn Iterations

L

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.

5.3 Modeling Joint Return Distribution for Multi-Dimensional Reward Functions

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.

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

1000

2000

3000

4000

5000

6000

Average Return

AirRaid

MMD

Sinkhorn

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

200

400

600

800

1000

1200

Average Return

Asteroids

MMD

Sinkhorn

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

1000

2000

3000

4000

Average Return

Gopher

MMD

Sinkhorn

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

500

1000

1500

2000

2500

Average Return

MsPacman

MMD

Sinkhorn

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

20

10

0

10

20

Average Return

Pong

MMD

Sinkhorn

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

5000

10000

15000

20000

25000

Average Return

UpNDown

MMD

Sinkhorn

6 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.

References

Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Marc Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems , 34:29304-29320, 2021.

Rohit Agrawal and Thibaut Horel. Optimal bounds between f-divergences and integral probability metrics. Journal of Machine Learning Research , 22(128):1-59, 2021.

Mokhtar Z Alaya, Maxime Berar, Gilles Gasso, and Alain Rakotomamonjy. Screening sinkhorn algorithm for regularized optimal transport. Advances in Neural Information Processing Systems , 32, 2019.

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning , pages 214-223. PMLR, 2017.

Marc G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective on reinforcement learning. International Conference on Machine Learning (ICML) , 2017.

Marc G Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan, Stephan Hoyer, and Rémi Munos. The cramer distance as a solution to biased wasserstein gradients. arXiv preprint arXiv:1705.10743 , 2017.

Nils Bjorck, Carla P. Gomes, and Kilian Q. Weinberger. Towards deeper deep reinforcement learning with spectral normalization. Advances in neural information processing systems , 34, 2021.

Tianshi Cao, Alex Bie, Arash Vahdat, Sanja Fidler, and Karsten Kreis. Don't generate me: Training differentially private generative models with sinkhorn divergence. Advances in Neural Information Processing Systems , 34:12480-12492, 2021.

Yu Chen, Xiangcheng Zhang, Siwei Wang, and Longbo Huang. Provable risk-sensitive distributional reinforcement learning with general function approximation. International Conference on Machine Learning , 2024.

Yutian Chen, Max Welling, and Alex Smola. Super-samples from kernel herding. UAI, 109-116. AUAI Press , 2012.

Taehyun Cho, Seungyub Han, Heesoo Lee, Kyungjae Lee, and Jungwoo Lee. Pitfall of optimism: Distributional reinforcement learning by randomizing risk criterion. Advances in Neural Information Processing Systems , 36, 2024.

Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems , 26, 2013.

Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. International Conference on Machine Learning (ICML) , 2018.

Will Dabney, Mark Rowland, Marc G Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. Association for the Advancement of Artificial Intelligence (AAAI) , 2018.

Christoph Dann, Yishay Mansour, and Mehryar Mohri. Reinforcement learning can be more efficient with multiple rewards. In International Conference on Machine Learning , pages 6948-6967. PMLR, 2023.

John N Darroch and Douglas Ratcliff. Generalized iterative scaling for log-linear models. The annals of mathematical statistics , pages 1470-1480, 1972.

Monroe D Donsker and SR Srinivasa Varadhan. Asymptotic evaluation of certain markov process expectations for large time-iii. Communications on pure and applied Mathematics , 29(4):389-461, 1976.

Stephan Eckstein and Marcel Nutz. Quantitative stability of regularized optimal transport and convergence of sinkhorn's algorithm. SIAM Journal on Mathematical Analysis , 54(6):59225948, 2022.

Odin Elie and Charpentier Arthur. Dynamic Programming in Distributional Reinforcement Learning . PhD thesis, Université du Québec à Montréal, 2020.

Kilian Fatras, Thibault Séjourné, Rémi Flamary, and Nicolas Courty. Unbalanced minibatch optimal transport; applications to domain adaptation. In International Conference on Machine Learning , pages 3186-3197. PMLR, 2021.

Jean Feydy, Thibault Séjourné, François-Xavier Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 2681-2690. PMLR, 2019.

Joel Franklin and Jens Lorenz. On the scaling of multidimensional matrices. Linear Algebra and its applications , 114:717-735, 1989.

Kenji Fukumizu, Arthur Gretton, Gert Lanckriet, Bernhard Schölkopf, and Bharath K Sriperumbudur. Kernel choice and classifiability for rkhs embeddings of probability distributions. Advances in neural information processing systems , 22, 2009.

Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complexity of sinkhorn divergences. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 1574-1583. PMLR, 2019.

Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics , pages 16081617. PMLR, 2018.

Arthur Gretton, Karsten Borgwardt, Malte Rasch, Bernhard Schölkopf, and Alex Smola. A kernel method for the two-sample-problem. Advances in neural information processing systems , 19, 2006.

Philippe Rigollet Jason Altschuler, Jonathan Weed. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration, 2017.

Edwin T Jaynes. Information theory and statistical mechanics. Physical review , 106(4):620, 1957.

Christian Léonard. A survey of the schr \ " odinger problem and some of its connections with optimal transport. arXiv preprint arXiv:1308.0215 , 2013.

Zichuan Lin, Derek Yang, Li Zhao, Tao Qin, Guangwen Yang, and Tie-Yan Liu. Rd2: Reward decomposition with representation decomposition. Advances in Neural Information Processing Systems , 33:11298-11308, 2020.

Zichuan Lin, Li Zhao, Derek Yang, Tao Qin, Tie-Yan Liu, and Guangwen Yang. Distributional reward decomposition for reinforcement learning. Advances in neural information processing systems , 32, 2019.

Daniel J Lizotte, Michael H Bowling, and Susan A Murphy. Efficient reinforcement learning with multiple reward functions for randomized controlled trial analysis. In International Conference on Machine Learning , volume 10, pages 695-702, 2010.

Yudong Luo, Guiliang Liu, Haonan Duan, Oliver Schulte, and Pascal Poupart. Distributional reinforcement learning with monotonic splines. In International Conference on Learning Representations , 2021.

Yecheng Ma, Dinesh Jayaraman, and Osbert Bastani. Conservative offline distributional reinforcement learning. Advances in neural information processing systems , 34:19235-19247, 2021.

Borislav Mavrin, Shangtong Zhang, Hengshuai Yao, Linglong Kong, Kaiwen Wu, and Yaoliang Yu. Distributional reinforcement learning for efficient exploration. International Conference on Machine Learning (ICML) , 2019.

Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature , 518(7540):529-533, 2015.

Thanh Tang Nguyen, Sunil Gupta, and Svetha Venkatesh. Distributional reinforcement learning with maximum mean discrepancy. Association for the Advancement of Artificial Intelligence (AAAI) , 2020.

Giorgio Patrini, Rianne Van den Berg, Patrick Forre, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. In Uncertainty in Artificial Intelligence , pages 733-743. PMLR, 2020.

Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learning , 11(5-6):355-607, 2019.

Aaditya Ramdas, Nicolás García Trillos, and Marco Cuturi. On wasserstein two-sample testing and related families of nonparametric tests. Entropy , 19(2):47, 2017.

Mark Rowland, Robert Dadashi, Saurabh Kumar, Rémi Munos, Marc G Bellemare, and Will Dabney. Statistics and samples in distributional reinforcement learning. International Conference on Machine Learning (ICML) , 2019.

Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Harutyunyan, Karl Tuyls, Marc G Bellemare, and Will Dabney. An analysis of quantile temporal-difference learning. Journal of Machine Learning Research (JMLR) , 2024.

Mark Rowland, Yunhao Tang, Clare Lyle, Rémi Munos, Marc G Bellemare, and Will Dabney. The statistical benefits of quantile temporal-difference learning for value estimation. International Conference on Machine Learning , 2023.

Ludger Rüschendorf and Wolfgang Thomsen. Closedness of sum spaces andthe generalized schrödinger problem. Theory of Probability & Its Applications , 42(3):483-494, 1998.

Richard Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. The American Mathematical Monthly , 74(4):402-405, 1967.

Ke Sun, Bei Jiang, and Linglong Kong. How does return distribution in distributional reinforcement learning help optimization? In ICML 2024 Workshop: Aligning Reinforcement Learning Experimentalists and Theorists .

Ke Sun, Yingnan Zhao, Shangling Jui, and Linglong Kong. Exploring the training robustness of distributional reinforcement learning against noisy state observations. European Conference on Machine Learning and Principles and Practice of Knowledge Discovery in Databases (ECML-PKDD) , 2023.

Ke Sun, Yingnan Zhao, Yi Liu, Enze Shi, Yafei Wang, Xiaodong Yan, Bei Jiang, and Linglong Kong. The benefits of being categorical distributional: Uncertainty-aware regularized exploration in reinforcement learning. arXiv preprint arXiv:2110.03155 , 2021.

Gábor J Székely. E-statistics: The energy of statistical samples. Bowling Green State University, Department of Mathematics and Statistics Technical Report , 3(05):1-18, 2003.

Harm Van Seijen, Mehdi Fatemi, Joshua Romoff, Romain Laroche, Tavian Barnes, and Jeffrey Tsang. Hybrid reward architecture for reinforcement learning. Advances in Neural Information Processing Systems , 30, 2017.

Kaiwen Wang, Owen Oertell, Alekh Agarwal, Nathan Kallus, and Wen Sun. More benefits of being distributional: Second-order bounds for reinforcement learning. International Conference on Machine Learning , 2024.

Kaiwen Wang, Kevin Zhou, Runzhe Wu, Nathan Kallus, and Wen Sun. The benefits of being distributional: Small-loss bounds for reinforcement learning. Advances in neural information processing systems , 2023.

Li Kevin Wenliang, Grégoire Déletang, Matthew Aitchison, Marcus Hutter, Anian Ruoss, Arthur Gretton, and Mark Rowland. Distributional bellman operators over mean embeddings. International Conference on Machine Learning , 2024.

Eric Wong, Frank Schmidt, and Zico Kolter. Wasserstein adversarial examples via projected sinkhorn iterations. In International Conference on Machine Learning , pages 6808-6817. PMLR, 2019.

Runzhe Wu, Masatoshi Uehara, and Wen Sun. Distributional offline policy evaluation with predictive error guarantees. International Conference on Machine Learning , 2023.

Derek Yang, Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tie-Yan Liu. Fully parameterized quantile function for distributional reinforcement learning. Advances in neural information processing systems , 32:6193-6202, 2019.

Pushi Zhang, Xiaoyu Chen, Li Zhao, Wei Xiong, Tao Qin, and Tie-Yan Liu. Distributional reinforcement learning for multi-dimensional reward functions. Advances in Neural Information Processing Systems , 34:1519-1529, 2021.

Fan Zhou, Jianing Wang, and Xingdong Feng. Non-crossing quantile regression for distributional reinforcement learning. Advances in Neural Information Processing Systems , 33, 2020.

Florian Ziel. The energy distance for ensemble and scenario reduction. arXiv preprint arXiv:2005.14670 , 2020.

Appendix

Table of Contents

A Smoother Transport Plan via Sinkhorn Divergence by Increasing ε

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).

-2

-1

0

1

2

-2

-1

0

1

2

: 1.0, Cost: 0.95

-2

-1

0

1

2

-2

-1

0

1

2

: 0.2, Cost: 0.40

-2

-1

0

1

2

-2

-1

0

1

2

: 0.02, Cost: 0.26

0

20

40

60

0

10

20

30

40

50

60

Probabilistic Transport Plan

0

20

40

60

0

10

20

30

40

50

60

Probabilistic Transport Plan

0

20

40

60

0

10

20

30

40

50

60

Probabilistic Transport Plan

0.0005

0.0010

0.0015

0.0020

0.0025

0.000

0.002

0.004

0.006

0.0000

0.0025

0.0050

0.0075

0.0100

0.0125

0.0150

Probability mass

(a) Enduro

-2

-1

0

1

2

-2

-1

0

1

2

: 1.0, Cost: 0.96

-2

-1

0

1

2

-2

-1

0

1

2

: 0.2, Cost: 0.41

-2

-1

0

1

2

-2

-1

0

1

2

: 0.02, Cost: 0.26

0

20

40

60

0

10

20

30

40

50

60

Probabilistic Transport Plan

0

20

40

60

0

10

20

30

40

50

60

Probabilistic Transport Plan

0

20

40

60

0

10

20

30

40

50

60

Probabilistic Transport Plan

0.0005

0.0010

0.0015

0.0020

0.000

0.001

0.002

0.003

0.004

0.005

0.006

0.0000

0.0025

0.0050

0.0075

0.0100

0.0125

0.0150

Probability mass

-2

-1

0

1

2

-2

-1

0

1

2

: 1.0, Cost: 0.89

-2

-1

0

1

2

-2

-1

0

1

2

: 0.2, Cost: 0.30

-2

-1

0

1

2

-2

-1

0

1

2

: 0.02, Cost: 0.17

0

20

40

60

0

10

20

30

40

50

60

Probabilistic Transport Plan

0

20

40

60

0

10

20

30

40

50

60

Probabilistic Transport Plan

0

20

40

60

0

10

20

30

40

50

60

Probabilistic Transport Plan

0.0005

0.0010

0.0015

0.0020

0.0025

0.000

0.002

0.004

0.006

0.008

0.010

0.0000

0.0025

0.0050

0.0075

0.0100

0.0125

0.0150

Probability mass

(c) Seaquest

B Definition of Distribution Divergences and Contraction Properties

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:

which 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 as

Further, ℓ p distance [19] is defined as

The ℓ 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 by

where 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 as

where X and Y are multivariate. Moreover, the energy distance is a special case of the maximum mean discrepancy (MMD), which is formulated as

where 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 ( X ) S×A → R formulated as

We 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].

C Proof of Proposition 1

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]:

where 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 ) whose pdf p Π ′ satisfies p Π ′ ( x, y ) = p Π ( x -a, y -a ) . Based on the two-dimensional variational representation of KL divergence MI Π ( U, V ) = sup f ∈L b { E Π [ f ( x, y )] -log ( E U,V [ e f ( x,y ) ]) } , we have:

where (a) is by the independence of A between X,Y , and the joint cdf Π . For instance, in the one-dimensional setting, we have E Z = A + X [ f ( z )] = ∫ a ∫ x f ( x + a ) p A ( a ) p X ( x ) dxda = E A [ E X [ f ( x + a )]] . (b) and (c) are by Jensen's inequality in terms of the convex function -log( x ) and sup f , and (d) is because the translated cdf is still within L b .

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:

Putting 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.

where (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 Π .

D Proof of Proposition 2

D.1 Sum Invariant Property

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 proving

According to [39], we have the dual form of W c,ε :

Therefore, we have:

where ( 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 φ, ψ .

D.2 A Variant of Scale Sensitive Property when c = -k α

General Conclusion. Let Π ∗ be the optimal coupling for W c,ε , we define a ratio λ ε ( U, V ) = ε KL (Π ∗ | µ ⊗ ν ) W c,ε ∈ (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:

where 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 (Π ∗ | µ ⊗ ν ) W c,ε ∈ (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 as d K 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:

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, Π + ∞ k =1 q k = ( 2 3 4 3 )( 3 4 5 4 ) · · · > 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 of ∆ U,V ε ( a, α ) :

where 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.

E Proof of Theorem 1

E.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 that

and 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 Course 4 and we invite interested readers for reference. It follows that

where 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:

where 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 π .

Our complete proof is inspired by [40, 25]. Recap the definition of squared MMD is

When the kernel function k degenerates to an unrectified k α ( x, y ) := -∥ x -y ∥ α for α ∈ (0 , 2) , the squared MMD would degenerate to

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 as

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 have

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,

Besides, by the compactness of Π( µ, ν ) , we can extract a converging subsequence Π n k → Π ∞ . Since KL is weakly lower-semicontinuous, it holds that

Hence Π ∞ = µ ⊗ ν . That being said that the optimal coupling is simply the product of the marginals, indicating that Π ε → µ ⊗ ν as ε →∞ . As a special case, when α = 1 , W -k 1 , ∞ ( u, v ) is equivalent to the energy distance

In 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:

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 satisfying

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 (Π ∗ | µ ⊗ ν ) W c,ε ∈ (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 as

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,ε :

Putting all together, we now derive the convergence of distributional Bellman operator T π under the supreme form of W c,ε , i.e., W ∞ c,ε :

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.

F Proof of Corollary 1

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:

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 have

where 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 RL

Algorithm 2 Sinkhorn Iterations to Approximate W c,ε ( { Z i } N i =1 , { T Z j } N j =1 )

Input : Two samples sequences { Z i } N i =1 , { T Z j } N j =1 , number of iterations L and hyperparameter

.

2: K i,j = exp( -ˆ c i,j /ε )

3: b 0 ← 1 N

4: for l = 1 , 2 , ..., L do

6: end for

Given two sample sequences { Z i } N i =1 , { T Z j } N j =1 in the distributional RL algorithm, the optimal transport distance is equivalent to the form:

where the empirical cost function is ˆ c i,j = c ( Z i , T Z 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:

where K i,j = e -ˆ c 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 as

ε

where · · indicates an entry-wise division. Combining Sinkhorn Iteration in Algorithm 2 and the generic update of Sinkhorn Distributional RL in Algorithm 1, we provide a full version of Sinkhorn Distributional RL algorithm in Algorithm 3.

Algorithm 3 Sinkhorn Distributional RL

Require: Number of generated samples N , the kernel k (e.g., unrectified kernel), discount factor γ ∈ [0 , 1] , learning rate α , replay buffer M , main network Z θ , target network Z θ ∗ , number of iterations L , hyperparameter ε , and a behavior policy π based on Z θ following an ϵ -greedy rule ∗

1: Initialize θ and θ ← θ

2: for t = 1 , 2 , . . . do

3: Take action a t ∼ π ( ·| s t ; θ ) , receive reward r t ∼ R ( ·| s t , a t ) , and observe s t +1 ∼ P ( ·| s t , a t )

4: Store ( s t , a t , r t , s t +1 ) to the replay buffer M

5: 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: T Z i ← r + γZ θ ∗ ( s ′ , a ∗ ) i , ∀ 1 ≤ i ≤ N

8: Evaluate Sinkhorn divergence via Sinkhorn Iterations in Algorithm 2:

9: Update the main network Z θ : θ ← θ -α ∇ θ W c,ε ( { Z θ ( s, a ) i } N i =1 , { T Z j } N j =1 )

10: Periodically update the target network θ ∗ ← θ

11: end for

H Summary Table for Human Normalized Scores (HNS)

I Learning Curves on 55 Atari Games

500

4000

3000

200

100

Sinkho

MMD

8000

6000

40000

20000

9M

200000

100000

MMD

Sinkhc

22

Milions of Frames (M)

Milions of Frames (M)

Milions of Frames (M)

Millions of Frames (M)

202530

35

Millions of Frames (M)

SC0C0CS1015

40

Milions of Frames (M)

Enduro

Freeway

Frostbite

Gopher

Gravitar

300

1800

3500

DQN

C51

MMD

Average Retum

08

QRDON

DQN

C51

MM

35

QRDON

DQN

C51

MMD

300

1500

1000

3500

500

C51

DQN

MMD

Sinkhern

12000

200

MMD

DQN

C51

Sinkhom

600

100

DQN

C51

MMD

Milions of Frames(M)

100

Millions o Frames (M)

40

Millions of Frames (M)

15

20

25

3035

Millions of Frames (M)

1015202530

35

40

Milions of Frames (M)

10 15 2025

35

Millions of Frames (M)

10  15 20

2530

35

Hero

IceHockey

Jamesbond

Kangaroo

Krull

KungFuMaster

0000

800

6000

MMD

Sinkhog

DON

C51

Average Return

12.5

15.0

5.0

10.0

17.5

2.5

7.5

DQN

C51

Retum

Average

1000

800

600

DQN

MMD

Retum

Average

1000

800

4000

6000

2000

12000

MMD

Sinkhom

DQN

C51

Average Return

8000

000

DQN

Return

40000

30000

0000

DQN

C51

ORRAN

Milions of Frames (M)

101520253035

Millions o Frames (M)

01520253035

Millions of Frames (M)

253035

Millions of Frames (M)

25

40

Millions of Frames (M)

10

20

253035

40

Millions of Frames (M)

10

15

O

25

30

MontezumaRevenge

MsPacman

NameThisGame

Phoenix

Pitfall

Pong

Average Return

0

DQN

C51

MMD

Sinkr

1500

1000

400

500

DQN

Average Return

1000

12000

800

6000

4000

0000

DQN

MMD

C51

Retum

ebe.

Aver

30000

2500

10000

1500

35000

2000

5000

C51.

MMD

DQN

Sinkho

MMN

Average Return

100

1200

600

Sinkhom

DON

Average Retum

10

10

DQN

C51

Milions of Frames (M)

Millions of Frames (M)

4

Milions of Frames (M)

40

Millions o Frames (M)

15

20

25

3035

Millions of Frames (M)

10

15

20

25

30

35

Millions of Frames (M)

10

25

30

PrivateEye

Riverraid

RoadRunner

Robotank

Seaquest

Average Retum

7000

1000

NOO

C51

Average Retum

1

16000

MMD

DQN

C51

Sinkhom.

Average Return

1000

8000

600

12000

4000

200

C51

DQN

MMD

Return

4000

20000

DQN

C51

DQN

C51

Retum

16000

4000

2000

DQN

C51

MMD

Sinkhon

Milions o Frames(M)

20

25

30

35

Millions of Frames (M)

Millions of Frames (M)

10

15

20

25

3035

Milons of Frames (M)

1015

20253035

40

Milions of Frames (M)

35

0

Millions of Frames (M)

10

15

25

30

SC

Solaris

Spacelnvaders

StarGunner

Tennis

TimePilot

Average Return

DON

Average Return

DQN

C51

E

180

200

DQN

C51

Retum

5000

4000

C51

DQN

Average Retum

C51

DQN

Retum

DQN

C51

25000

MMD

Sinkhorn

25

MMD

46067

Average

2000

10000

30000

MMD

Sinkhon

5

10

15

MMD

2000

MMD

Milions of Frames (M)

10

25

OC

Millions of Frames (M)

1015

30

00510

Milions of Frames (M)

Millions o Frames (M)

51015

20

253035

40

Millions of Frames (M)

10

25

Milions o Frames (M)

25

Tutankham

UpNDown

Venture

VideoPinball

WizardOfWor

YarsRevenge

Average Return

200

150

100

DQN.

Retum

Average

200

25000

1500

5000

ORDONAW

C51

DQN

一

MMD

Sinkhom

DQN

C51

AverageReturn

50000

2000

100000

000

00000

4000

70000

DQN

C51

MMD

Average

10000

2000

400

一

MMD

DQN

Sinkhom

Average Retum

200

0000I

15000

5000

DQN

C51

MMB

Millions of Frames (M)

Milions of Frames (M)

10

20

25

30

35

0

Millions of Frames (M)

40

0

Milions of Frames (M)

0

Milions of Frames (M)

10

15

20

25

40

5101520253035

Milions of Frames (M)

40

Zaxxon

Average Retum

16000

DQN

C51

MMD

Sink

0

Millions of Frames (M)

J Raw Score Table Across 55 Atari Games

K Features of Atari Games

L Sensitivity Analysis and Computational Cost

L.1 More results in Sensitivity Analysis

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

100

200

300

400

500

600

Average Return

Breakout

=0.075

=0.1

=0.3

=0.5

=10.0

(a) Small

ε

vs QR-DQN

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

100

200

300

400

500

600

Average Return

Breakout

=10

=1000

=5000

MMD

(b) Large

ε

vs MMD-DQN

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

100

200

300

400

500

Average Return

Breakout

=1

=10

=100

=500

(c)

ε

on Breakout

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

5000

10000

15000

20000

Average Return

Seaquest

=1

=10

=100

=500

MMD

(d)

ε

on Seaquest

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.

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

100

200

300

400

500

600

Average Return

Breakout

Samples=2

Samples=5

Samples=50

Samples=200

Samples=500

(a) Samples on Breakout

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

5000

10000

15000

20000

Average Return

Seaquest

Samples=2

Samples=5

Samples=200

Samples=500

(b) Samples on Seaquest

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.

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

10000

20000

30000

40000

50000

Average Return

StarGunner

=1

=10

=500

(a) StarGunner:

ε

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

10000

20000

30000

40000

50000

60000

Average Return

StarGunner

Samples=200

Samples=600

(b) StarGunner: Samples

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

2000

4000

6000

8000

10000

12000

14000

16000

Average Return

Zaxxon

=1

=10

=500

(c) Zaxxon:

ε

0

5

10

15

20

25

30

35

40

Millions of Frames (M)

0

2000

4000

6000

8000

10000

12000

14000

16000

Average Return

Zaxxon

Samples=2

Samples=200

Samples=500

(d) Zaxxon: Samples

L.2 Comparison with the Computational Cost

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.

 

























 

























































































M Experimental Setting in Multi-dimensional Return Distributions

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 MMDDQN [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

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: As the title states, we propose a new family of distributional reinforcement learning algorithm by leveraging Sinkhorn divergence, a regularized Wasserstein loss. The main content of our study accurately states our contributions from the perspectives of algorithm, theory and experiments.

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.

2. Limitations

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.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We have a few propositions and theorems in our study as our main theoretical contributions, and the complete proof of them are provided in Appendix, respectively.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We have provided the detailed experimental settings in both the experimental part of the main content and the appendix, including the reference code, which our experiments adapt from.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in

some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

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.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, 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.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: 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.

8. 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).

9. 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 ?

Answer: [Yes]

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).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [NA]

Justification: As 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).

11. 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 authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification:We believe our paper does not release new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: The paper does not involve crowdsourcing nor research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- The algorithm is clearly proposed and can be easily understood.\n- The authors study both theoretical (mostly contraction property) and empirical aspects of the proposed algorithm.\n- The experimental results seem comprehensive (on the entire atari games). They show that the algorithm performs better on complicated games.\n\nWEAKNESSES:\nAfter 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:\n\n- line 48-51: why does using samples instead of pre-specified statistics have better accuracy? Is there any reference?\n- 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.\n- line 53: \"smoother\" is very vague in this context. I would suggest using some different words.\n\nAs 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.\n\nFor 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.\n\nQUESTIONS:\nLine 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?\n\nLIMITATIONS:\nI didn't find any potential negative societal impact.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-08"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\n•\tThis 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.\n\nSTRENGTHS:\n•\tThe paper reviews the relevant literatures in distributional RL and analyzes the new Sinkhorn divergence metric as an alternative to Wasserstein loss. The author provides sufficient background and advantages of Sinkhorn divergence, demonstrating a well-motivated approach.\n\n•\tAlthough the evaluation scores are presented for 40 million frames rather than 200 million, the authors ensure transparency and reproducibility by reporting raw scores for performance comparisons. The authors also present a significant amount of experiment and ablation studies to provide sufficient experimental results for the effectiveness of their algorithm.\n\nWEAKNESSES:\n•\tTypos: In line 52, Sinkrhorn -> Sinkhorn \n\n•\tThe text and figures in the paper are quite dense and difficult to read. Some parts seem unnecessary to be included in the main paper.\n \no\tAlgorithm 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. \n\no\tThe 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].\n\n•\tThere 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. \n\n[1] Dabney, W., Rowland, M., Bellemare, M. G. & Munos, R. Distributional Reinforcement Learning with Quantile Regression. Arxiv (2017).\n\n[2] Hessel, M. et al. Rainbow: Combining Improvements in Deep Reinforcement Learning. arXiv (2017) doi:10.48550/arxiv.1710.02298.\n\nQUESTIONS:\n•\tCould 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?\n\n•\tThe 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?\n\nLIMITATIONS:\n•\tIn line 22, the reference [11] is cited as a risk-sensitive control, but this is close to exploration method for distributional RL.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n- Introduces a new family of distributional RL algorithms based on Sinkhorn divergence, expanding the toolkit for researchers and practitioners.\n\n- Provides theoretical analysis of Sinkhorn divergence in the context of distributional RL, including convergence guarantees.\n\n- Demonstrates the effectiveness of SinkhornDRL through extensive experiments on a standard benchmark, showing superior performance in many cases.\n\n- Specifically tackles the issues of inaccurate distribution capture and difficulty with multi-dimensional rewards that plague quantile regression methods.\n\nWEAKNESSES:\n- SinkhornDRL introduces some additional computational cost compared to simpler methods like C51 and QR-DQN.\n\n- The algorithm requires tuning of additional hyperparameters (e.g., the regularization strength), which might require extra effort in practice.\n\n- The paper acknowledges that a deeper connection between theoretical properties of divergences and practical performance in specific environments remains an open question.\n\n- Some typos:\n  - line 151 and line 647 Sinkrhon\n  - line 163 supremal from\n  - line 787 numeral\n\nQUESTIONS:\nCould the authors address the following points:\n\n- Are there existing works on distributional reinforcement learning that use the entropic regularized Wasserstein distance?\n- What specific advantages does the Sinkhorn divergence offer over the entropic regularized Wasserstein distance in the context of this study?\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}