Distributional Reinforcement Learning with Regularized Wasserstein Loss
Authors: Ke Sun
Yingnan Zhao
Wulong Liu
Bei Jiang
Linglong Kong
Venue: NeurIPS 2024 poster (2024)
Keywords: distributional reinforcement learning
Sinkhorn divergence
Primary area: reinforcement_learning

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
Ke Sun', Yingnan Zhao", Wulong Liu Bei Jiang', Linglong Kong'
University of Alberta, Edmonton, Canada Harbin Engineering University, China
Huawei Noah's Ark Lab
(ksun6, , Qualberta. ca
zhaoyingnan@hrbeu. edu. en Liuvulong@hhaweeiccoo com
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 RI. (Sinkhorn- DRL), which leverages Sinkhorn divergence regularized Wasserstein loss-10 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 Diserepancy (MMD). The introduced SinkhomDRL en- riches the family of distributional RL algorithms, contributing to interpreting the algorithm behaviors compared with existing approaches by our investigation into their relationships. Empirically, show that SinkhornDRL. consistently outperforms or matches existing algorithms on the Atari games suite and partic- ularly stands out in the multi-dimensional reward setting. Code is available in com datake/ com./attalee/Snnnooonitttt.. SiinkhornDistRR.
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,
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). Mhantil-RRereesioon DQN (QR-DQN) [4]. EDRL [1], Implicit Quantile Networks (IQN) [3] Fully Parameterized Quantile Function (FQF) [56], Crossinn QR-DQN (58). Maximum Mean
Discrepancy (MMD-DQN) (37), Spline (SPL-DON) [33] and Bketch-DQN [53]. Beyond the
performance advantage, distributional RL has also possessed benefits in risk- sensitvve control 13
exploration 335 [1 48 offline setting 34 [55 statistical value estimation 4
robustness
47 and
optimization 46 52 51
Limitations of Typical Distributional RL Algorithms. Despite the gradual introduction of
numerous algorithms, quantile regression-based algorithms [14, [3] [6)] 41 33 42 43) dominate attention and research in the realm of distributional RL. These algorithms utilize quantile regres- sion to approximate the dieeeinnnsionl 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).
"Corresponingg author

---

generating quantiles of return distributions via neural networks often suffers from the non non-crossing issue [38], 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-specifiet statistics [41], while unrestricted statis- ties based on deterministic samples can be potentially more effective in complex environments [37]. 2) Difficulties in Extension to Multt-dimensional Rewards. Many RL tasks involve multiple sources of rewards (32 15). hybrid reward architecture [50, (3), sub-reward structures after reward de- composition B1 171, which require learning multi-dimensional return distributions to reduce the intrinsic uncertainty of the environments. However, it remains elusive how quantile regressions to upproximate a multi-dimensional Wasserstein distance, while circumventing the computational intractability issue in the related multi-dimensionn output space.
Motivation of Sinkhorn Divergence: a Regularized Wasserstein loss. Sinkhorn divergence 45
has emerged as theoretically principled and computationally efficient alternative for approximating
Wasserstein distance. It has gained increasing attention in the field of optimal transport 44
24 21
39
and has been successfully applied in various areas of machine learning [38 25
54
20
8 By intro-
ducing entropic regularization, Sinkhorn divergence can efficiently approximate multi-dimensional
Wasserstein distance using computationally efficient matrix scaling algorithms A5 39) This makes
it feasible to apply optimal transport distances to RL tasks with multi-dimensse dimensional rewards (see experiments in Section [..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 RI. Beyond addressing the two main limitations mentioned above, the well-controlled regularization introduced in Sinkrhorn divergence helps 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-entrope principle [28 [6], which aims to maximize entropy while keeping the transportation cost constrained. Furthermore, the resulting strongly convex loss function [3] and the induced smoothness by regu- larization facilitate faster and more stable convergence in the deep RL setting (see more details in Sections 4 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 regression- based algorithms while promoting more stable training. As Sinkhorn divergence interpolates between Wasserstein distance and MMD (26 21 39] we probe this relationship the RL context, character- izing 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 SinkhorDDRI, which overcomes the primary shortcomings of predominantly utilized quantile regression- based algorithms. SinkhomDDRL 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 SinkhoraDRL in the multi- multi-dimensional reward setting.
2 Preliminary Knowledge
2.1 Distributional Reinforcement Learning
In classical RI, an agent interacts with an environment via Markov decision process (MDP), 5-tuple (S, A, R, P,7), 7), where S and are the state and action spaces, P is the environment transition dynamics, R is the reward function, and y (0, 1) is the discount factor.
2

---

Given a policy т, the discounted sum of future rewards Z is a random variable with (²(s.a) = Et= 2 R(s,, at), where so = 8, ap = a, $141 P( st, a,), and a T(-|s;). ) In expectation- based RI,, the action-value function Q" is defined as (s.a] = E 2 (s, a)), which is iteratively updated via Bellman operator T through T (*)(s a) E R(s a) + 7.u..p.. [Q (s',a'),, where s' P(-|s, a) and a' ~ (=Is)). In contrast, distributional RL focuses on the action-return distribution, the full distribution of 2 ("(s,)) a). The return distribution is iteratively updated by applying the
distributional Bellman operator I through Z a) :2 R(s, a) + 7Z (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 metries [9
2.2 Divergences between Measures
Optimal Transport (OT) and Wasserstein/ Earth Mover's Distance. The optimal transport (OT) metric We defines powerful geometry compare two probability measures (J,, p), i.e. W. infneи()))) f c(1 (x, y). where is the cost function, IT is the joint distribution with marginals (H, and the minimizer II is called the optimal transport plan or optimal coupling. The p- Wasserstein distance Wp = (inf f | - dll'd((( y)) is a special case of optimal transport with the Euclidean norm as the cost function. Relative conventional divergences, including Hellinger, total variation or Kullback-Leii distance inherently integrates the spatial 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 (d3 log(()) when comparing two histograms of dimension d [12]
divergences, the formulation of OT and Wasserstein
Maximum Mean Diserepancy [26). Define two random variables X and Y. The squared Maximum
Mean Discrepancy (MMD) MMD with the kernel k is formulated MMD = E [k (X, X')] E [k (Y, 2E [xx(X, Y)), where .(. 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 ka 1| y|ا in algorithm analysis. With a slight abuse of notation, we also use Zo to denote 0 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 [ e.g.. C51, represents the return distribution using a categorical distribution with discrete fixed supports within a predefined interval.
(2) Quantile Regression (Wasserstein Distanee) Distributional RL. QR-DQN [4 employs quan- tile 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 more expressive quantile network. However, mentioned in Section [ quantile regression distributional RL struggles with accurately capturing return distribution characteristics and handling dinliiddimesionall reward settings. SinkhomDRL, 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 statisties, i.e., samples, to represent the return distribution and optimizes under MMD, which can manage multi-dimeesional rewards. However, the data geometry captured by MMD with a specific
3

---

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]. contrast, SinkhomDRL is fundamentally based on OT, inherently capturing the spatial and geometric layout of return distributions. This enables SinkhornDRI. 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
Sinkhorn Distributional RL (SinkhornDRL)
The algorithmic evolution of distributional RL can be primarily viewed along two dimensions 37
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 particular family by increasing their model capacity, e.g., IQN and FQF. Concretely, SinkhorDDRL 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 sweet trade-off that simultaneously leverages the geometry property 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 We,s (14, v) formulated as
/ ((x,y)dl((a v) + =KL(TT] & v), TETT))))))
(o, min
(1)
where the entropic regularization KL(IT /4 @ v) f log ( m(x.g)
dII v), also known as
(d(())d(())
mutual information, makes the optimization strongly convex and differential [ 21), allowing for
efficient matrix scaling algorithms for approximation, such as Sinkhorn Iterations 45 In statistical
physics, (e v) can be re-factored as a projection problem:
(e,(H,,) v) == min
KL (II)/C,
ne (p..)
where is the Gibbs distribution and its density function satisfies dK(x,, y) clc..v// du
This problem is often referred to as the "static Schrödinger problem" [29] [4] as it was initially considered in statistical physics. Formally, the Sinkhorn divergence is defined
We,(u,v (u, We.,(p, 14) We,e(,
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 I" under distribution divergence dp depends on its seale sensitive (S) and sum invariant (I) properties (6] 5]. We say dp is scale sensitive (of order т) if there exists T > O, such that for all random S Ja dX, Y). dp has the sum invariant property if whenever random variable A is independent from X, Y. we have (A + X, A + Y) < dp (X, Y). Based on these properties, (5] shows that Ihh -ycontractivv under the supremal form of Wasserstein distance W,, which is regarding the first term of We,e or directly letting in Eq. [ When examining the regularized loss form of Wo,s a natural question arises: What is the influence of the incorporated regularization term on the contraction of IT? We begin to address this question in Proposition mutual information (In( (s), a)) = KL[II] a) & v(s. a)) and its supremal form MI (p, v) = (a.)Es.) (I(II) a) @ v(s,a) given a joint distribution II.
variables X, Y and real value > 0, (aX,
0
[ focusing on the separate regularization term. Here, we define
Proposition 1. I is non-expansive under MI for any non-trivial joint distribution II.
Please refer to Appendix for the proof, where we investigate both (S) and (I) properties. The non-trivial II rules out the independence case of μ and v, where KL & v) would degenerate

---

to zero. Although the non-expansive nature of the introduced regularization term, as shown in Proposition may potentially slow the convergence Sinkhorn divergence compared with W without the regularization, we will demonstrate that I is still contractive under the full Sinkhorn divergence in Theorem ! Before introducing Theorem [ we first present the sum-invariant and a new variant of scale-sensitive properties in Proposition P which acts as the foundation for Theorem [
Proposition 2. Considering We.s with the unrectified kernel ki- := == yll" as -c > 0) and scaling factor a E (0, 1). We,s is sum-invariant (1) and satisfies We.saau av) I (a We,, (p, v) (S) with a scaling constant (a a) E [|a]", 1) for any and v in a finite set of probability measures.
Proof Sketch. The detailed proof is provided in Appendix D Let IT be the optimal coupling of
We,s. we define a ratio de ( u) that satisfies >e(ue, = ekLill" mBN) E (0, 1) for a generally non- Wer
zero We.s- The ratio X (,, 2/) measures the proportion of the entropie regularization term over the whole loss term We,e- Therefore, the contraction factor (a, a) is defined as Ae a) = [a]a(1 sup,.. (e, 2)) + supg, y (e(,,)) E ([a]", l 1) with sup,.. As (p, v/ < 1, which is determined by the scale factor a, the order a, the hyperparameter E, and the set of interested probability measures.
Contraction Guarantee and Interpolation Relationship. Proposition 2 reveals that We,e with an unrectified kernel satisfies (I) and a variant of (S) properties. While the scaling constant (a a) in (S) has complicated form, it remains strictly less than one, even considering a non-expansive nature of the entropic regularization as shown in Proposition divergence as Wa. c.c (1,, 2/) : W our (H, 2) = sup(s.a)ESx. Wa.(H(s, a), (/(s, a)). In Theorem 1 we will integrate all these properties to demonstrate the contraction property of distributional dynamic programming under We,. specifically highlighting the interpolation property of Sinkhorn divergence between MMD and Wasserstein distance in the context of distributional RL.
| We denote the supremal form of Sinkhorn
Theorem 1. Considering We.e (1,, v) with an unrectified kernel hi := ==-1 yll as cc > 0), where p. v E the distribution set of (=(ss, (s,a)} for E S, A in a finite MDP. We define the ratio Xe(s, v) as Te(1,, v) = =KL/II' (NSv) E 1) with supp.y (e v) < 1. Then, we have: (H.))
(1) + 0) We,,(,, v) - (w. When e 0, T is 7thh contracctive under W e.e. (2 - +0o) Wo,((u,) v) - MMD (1. When +00, I" is 78 contrrcctive under Wo o,
(3) (E E (0, +o0)), I" is at least Ae ccontraccvvve under W where Te a) is an MDP- dependent constant defined as Ac(r, = (1 supp, de v)) + supp., de u) E (77,
Proof Sketch. The detailed proof of Theorem I can be found in Appendix E Theorem (1) and (2) follow-up conclusions in terms of the convergence behavior of I based on the interpolation relationship between Sinkhorn divergence with Wasserstein distance and MMD [25). We also provide rigorous analysis within the context of distributional RL for completeness. Our critical theoretical contribution is the part (3) for the general 5 E (0, oo), where
show that I is at least A (7, a)-
contractive operator. The contraction factor Ar(7, a)) E (7 1) depends on the return distribution set (Z* (s, }}} of the considered MDP, and is also a function of 7, and a. Due to the influence of the regularization term in Sinkhorn loss, A.(7, a) larger than |7|%, the contraction factor for Wasserstein distance without the regularization. Thus, (2(7,a) can be seen as an interpolation between 7 and 1, with the coefficient sup,,. (e v) E (0, defined in Theorem The ratio (e(.,,) v) measures the proportion of the KL regularization term relative to We,e- As E + or +oc, sup,, Xe(u, u) - 0. leading to "0 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 E [0, oo) also aligns well with the existing distributional RL algorithms when 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 I (3). This conclusion is consistent with MMD-DQN [37 (s - +00), where TT 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 kig as the cost function and set a = 2 in our experiments, ensuring that Ix retains the contraction property guaranteed by we also summarize the main properties of distribution divergences in typical distributional RL algorithms, including the convergence rate of I" and sample complexity.
Theorem (3). In Table [
5

---

Algorithm C51
de Distribution Divergence Cramér distance
Representation Zo Categorical Distribution
Convergence Rate of
Sample Complexity of dz
VT
QR-DQN-1
Wasserstein distance
Quantiles Samples
in Oin
MMD-DQN
MMD
2//2 (ka)
SinkhomDRL
Sinkhorn divergence
Samples
2.
Ope 0) O(n +)-+t)))0) oc)
(ours)
k.)
2/2 (+
Table Properties of different distribution divergences in typical distributional RL algorithms. is the sample dimension and k = 28d + ||e|| where the cost function с is 3-Lipschitz Sample complexity is improved to (1/nn) using the kernel herding technique 10 in MMD.
ie., the convergence rate of given metric between measure and its empirical counterpart function of the number of samples n.
4.2 Extension to Multi-dimennional Return Distributions
As the ability to extend to the multi-diman cilttiddimesiooall reward setting is one of the major advantages of SinkhomDDRL over quantile regression based algorithms, we next demonstrate that the joint distribu- tional Bellman operator in the multi-dimensional reward case contractive under Sinkhorn divergence Was First, define d-dimensional reward function R. 8 A P(TR)), where d represents the number of reward sources, Consequently, we have joint return distributions of the d-dimensional return vector 2 (s, a) Ltmo 7 (s, (t), where 2 (s, a) (ZT (s, a..... 2g (s, a)) The joint distributional Bellman operator "y applied on the joint distribution of the random vector Z(s, a)
is defined as [ a) _ R(s, a) + 7 (s', d'), where s' ~ P(-|s, a), a' ~ x-115).
Corollary 1. For two joint distributions Z₁ and Ze. I7 De (%, a) (ereaaa)ccoooa contracctive under W i.e.
Wh ( Z1, Z2) S (2(7, a) (Z), Zz).
(4)
Please refer to Appendix for the proof. The contraction guarantee of Sinkhorn divergence enables us effectively deploy our SinkhoraDRL 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 [3] 17 We compare SinkhorDDRL with MMD-DQN in multiple reward sources
setting in Section 5.3 where SinkhornDRL significantly outperforms MMD-DQP by leveraging its ability 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 RI. is to leverage the Sinkhorn loss Wx.e to measure the distance between the current action-return distribution (s(s, a)) and the target distribution I (o (s, a). This yields (Zo a), Ie (o a)) for each a pair. For the representation of Zg(s, a), employ the unrestricted statistics, i.e., deterministic samples, akin to MMD-DQN, instead of predefined statistic functionals like quantiles in QR-DQN categorical distributions in C51. More concretely, we use neural networks to generate samples to approximate the return distributions, expressed as (0 a) = (Z N where N is the number of generated samples. We refer to these samples {Zo(s, a), N as particles. We then use the Dirac mixture
1 EN 87u(s,a), to approximate the N
Algorithm Generic Sinkhorn distributional RL Update
Require: Number of generated samples N, the cost function c, hyperparameter and the target network Zo.
Input: Sample transition as, r','))
1 1: Policy evaluation: a" ~ (--ss' or Control: a" +- arg max N EN Zo (s', d');
2: I7 f- + 7 (s',a)) LA S i N
Output: We. (1Zo(s,a))) ,

---

true density function of (*,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
Efficiint 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 44 25 Notably, burden is the matrix-vecti differentiable layers on the typical deep neural network, such as DQN architecture. It has been proven that Sinkhorn iterations asymptotically converge to the true loss in a linear rate [25] 12] 27.. We provide detailed description of Sinkhorn iterations in Algorithm 2 and full version in Algorithm 3 of Appendix G practice, selecting proper values of and is crucial. To this end, we conduct rigorous sensitivity analysis, detailed in Section 5
Sinkhorn iteration with L steps yields a differentiable and solvable efficient loss function as the main
multiplication, which streams well on the GPU by simply adding extra
Remark: Relationship with IQN and FQE, In the realm of distributional RL algorithms, it is important to highlight that QR-DQN and MMD-DQN are direct counterparts to SinkhorDRL 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 SinkhorDRRL. For instance, we can implicitly generate (Zo } by applying neural network to N samples of base sampling distribution, as in IQN. We can also use a proposal network 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 SinkhomDRRI as described in Algorithm 1 on the entire 55 Atari 2600 games. Without increasing the model capacity for fair comparison, we leverage the same architecture 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 (3), QR-DQN [4 and MMD-DQN (37] For fair comparison. we build SinkhornDRI. and all baselines based on well-accepted PyTorch implementatior of distributional RL. algorithms. We re-implement MMD-DQN based on its original TensorFlow implementation and keep the same setting. For example, our MMD-DQN still employs Gaussian kernels hn (z, y) = exp(-(e (- (x - -y)²/h with the same kernel mixture trick covering a range of bandwidths h as adopted in MMD-DQN (37]
SinkhoraDRL Implementation and Hyperparameter Settings. For a fair comparison with QR- DQN, C51, and MMD-DQN, we use the same hyperparameters: the number of generated samples N 200, Adam optimizer with 0.00005, CAdam = 0.01 0.01/32. In SinkhormDRL, we choose the number of Sinkhorn iterations L = 10 and smoothing hyperparameter E = 10.0 in Section [5. after conducting sensitivity analysis Section 5.2, Guided by the contraction guarantee analyzed in Theorem we use the unrectified kernel, specifically setting - ka and choosing = 2. This choice ensures our implementation is consistent with the theoretical results regarding the contraction guarantee in Theorem (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) [ across all considered distributional RL algorithms in Figure I 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

---

Mcan of HNS
12
Le
CS1 GRDON
DGN
Mlf
0.5
CS1 QRDON
0%% di S0
CS1
GRDON
MMD
Modian of
015
MMD
MMD
Sinkmerr
114
10
15
10
30 35
25
50
Millions of Frames (M)
Millions of Frames (M)
Millions of Frames (M)
Figure 1: Mean (left), Median (middle), and IQM (5%) (right) of Human Normalier Scores (HNS)
summarized over 55 Atari games. We run 3 seeds for each algorithm.
extremely high scores 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 SinkhomDRL 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 D. This slower initial convergence can be explained by the slower contraction factor ( a) > ? in Theorem 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 [ Furthermore, 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 SinkhomDRL generall 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 [ a pertinent question arises: In which environments does SinkhorDDRI, potentially perform better? We empirically address this question by conducting a ratio improvement comparison between SinkhornDRL and both QR- DQN and MMD-DQN across all games. Figure 2 showcases that SinkhornDRL surpasses both QR-DQN and MMD-DQN in more than half of the games and significantly excels at them in large propertion of games. Notably, the games where SinkhomDRL achieves considerable improvemens tend to have larger action spaces and more complex dynamics. In particular, illustrated in Figure 2 these games include Venture, Seaquest, Solaris, Temnis, Phoenix, Atlantis, and Zaxxon. Most of these games have an 18 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 Sin For detailed comparison, we provide the features of all games, including the number of action spaces, and complexity of environment dynamics in Table Appendix K
In summary, compared with QR-DQN, the empirical success of SinkhorDDRL can be attributed to several key factors: 1, Enhanced return distribution representation: SinkhorDRLL captures return
Sinkhorn QRDQN (%)
Sinkhom MMD (%)
Ratio Improvment (%)
Ratio Improvment (%)
(a) SiL vs QR-DQN
(b) SinkhomDRL MMD-DQN
Figure 2: Ratio improvement of return for SinkhorDDRL over QR-DQN (left) and MMD-DQN (right) averaged over seeds. The ratio improvement is calculated by (SinkhomDRLL- QRR-QN)))) QR-DQN in (a) and (SinkhomDRL. MMD-DQN) MMD-DQN in (b), respectively.
8

---

distribution characteristics more accurately by directly using samples, avoiding the non-crosse con-crossing issue of learned quantile curves 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 SinkhornDRI. arise from its richer data representation capability when comparing return distributions, rooted in the OT nature. This is in comparison to the potentially restricted kerne-speccfic distances, such MMD,
5.2 Sensitivity Analysis and Computational Cost
Sensitivity Analysis. In practice, a proper is preferable an overly large or small e will lead to numerical instability of Sinkhorn iterations in Algorithm 2 (see the discussion in Section 4.4 of (39]) therefore worsening performance, as shown Figure (a). This implies that the potential interpolation nature of limiting behaviors between SinkhomDRL with QR-DQN and MMD- DQN revealed in Theorem may not be able to rigorously verified in numerical experiments. SinkhomDRI. also requires proper number of iterations I. and samples N. For example, a small N, N N 2 in Seaquest in Figure 3 (b) leads the divergence of algorithms, while an overly large N can degrade the performance and meanwhile increases the computational burden (Appendix (..)) 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 more detailed sensitivity analysis and more results StarGunner and Zaxxon in Appendix L.
Seaquest
Seaquest
Breakout
Average Relum
30000
20000
Sumplee=0
600
#+10
Samples75
510
35000
1.000
Sanstea-22
are
30000
e-500
GROON
Average
80000
Sanphea-500
Samples-230
Average
210
:10
Millions (M)
Millions Fremes (M)
Millions Frames
(a) Hyper-parameter
(b) Number of Samples
(c) Sinkhorn Iterations L
Figure 3: Sensitivity analysis of SinkhornDRL on Breakout and Seaquest in terms of E, number of samples, and number of iteration L. Learning curves reported over three seeds.
Computation Cost. In terms of the computation cost, SinkhornDRL slightly increases the com- putational overhead compared with C51, QR-DQN, and MMD-DQN. For instance, Sin 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
5.3 Modeling Joint Return Distribution for malti-Dimensional Reward Functions
Many RL tasks involve modeling multivariate return distributions. Following the nulti-dimensional reward setting in [57]. we compare SinkhoraDRL 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 SinkhorDDLL outperforms MMD-DQN in most cases for ditltiidmmeniinall reward functions. Of particular note, it remains open question to directly approximate multi-dimensional Wasserstein distances via quantile regression or other efficient algorithms, particularly in RL tasks.
Neerol Redum
AirRaid
Asteroids
Gopher
McPacman
Pong
LipNDown
Mub
VVD Senthmm
MNO
Si eihur
Sie Sention
Sireborr
Santhars
1000
romes (M)
mamme (M)
(M)
(M)
(K))
Figure 4: Performance of SinkhornDRL six Atari games with multi-dimennional reward functions.

---

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 SinkhomDRL achieves competitive performance, it relatively increases the computational cost and requires tuning additional hyperparameters. This hints that the enhanced performance offered by SinkhornDRI. may come with slightly greater efforts practical deployment. Additionally, it remains elusive for a deeper connection between the theoretical properties of diver- gences and the practical performance of distributional RL. algorithms given 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 parame- terizing the cost function and increasing model capacity. Moreover, Sinkhorn distributional RL also opens 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.
Acknowledgements
Yingnan Zhao and Ke Sun were supported by the State Scholarship Fund from China Scholarship Council (No:20200612)) 120405 and No:202006010088), Bei Jiang and Linglong Kong were partially supported by grants from the Canada CIFAR AI Chairs program, the Alberta Machine Intelligence Institute (AMII), and Natural Sciences and Engineering Council of Canada (NSERC), and Linglong Kong was also partially supported by grants from the Canada Research Chair program from NSERC. The authors express their gratitude for the insightful feedback provided by all reviewers and the area chairs, which significantly enhanced the initial version of this paper.
References
[1] Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron C Courville, and Mare Belle- mare. Deep reinforcement learning at the edge of the statistical precipice. Advances in neural information processing systems, 14:29304-29322 2021.
[2] Rohit Agrawal and Thibaut Horel. Optimal bounds between f-divergences and integral proba- bility metries. Journal of Machine Learning Research, 22(128):1-53) 2021.
[3] Mokhtar Z Alaya, Maxime Berar, Gilles Gasso, and Alain Rakotomamonjy. Screening sinkhorn algorithm for regularized optimal transport. Advances Neural Information Processing Systems, 32, 2019.
[4] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pages 214-223. PMLR, 2017.
[5] Mare G Bellemare, Will Dabney, and Rémi Munos. distributional perspective on reinforce- ment learning. International Conference on Machine Learning (ICML), 2017.
[6] 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
770. 2017.
[7] Nils Bjorek, Carla P. Gomes, and Kilian Q. Weinberger. Towards deeper deep reinforcement learning with spectral normalization. Advances in neural information processing systems, 34, 2021.
[8] Tianshi Cao, Alex Bie, Arash Vahdat, Sanja Fidler, and Karsten Kreis. Don't generate me: Training differentially private generative models with sinkhorn divergence. Advances Neural Information Processing Systems, 34: 12400-12222, 2021.
10

---

[9] Yu Chen, Xiangcheng Zhang, Siwei Wang, and Longbo Huang. Provable risk-sensitive distribu- tional reinforcement learning with general function approximation. International Conference on Machine Learning, 2024.
[10] Yutian Chen, Max Welling, and Alex Smola. Super-samples from kernel herding. UAI, 109-116. AUAI Press, 2012.
[11] 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.
[12] Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
[13] Will Dabney, Georg Ostrovski, David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. International Conference on Machine Learning (ICML), 2018.
[14] Will Dabney, Mark Rowland, Mare G Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. Association for the Advancement of Artificial Intelligence (AAAT), 2018.
[15] 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.
[16] John N Darroch and Douglas Ratcliff. Generalized iterative scaling for log-linear models. The annals of mathematical statistics, pages 1470-1480, 1972.
[17] 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.
[18] Stephan Eckstein and Marcel Nutz. Quantitative stability of regularized optimal transport and convergence of sinkhorn's 's algorithm. SIAM Journal on Mathematical Analysis, 54(6):5922- 5948, 2022.
[19] Odin Elie and Charpentier Arthur. Dynamic Programming in Distributional Reinforcement Learning. PhD thesis, Université du Québec à Montréal, 2020.
[20] 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.
[21] 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.
[22] Joel Franklin and Jens Lorenz, On the scaling of multidimennsinal matrices. Linear Algebra and its applications, 114:717-735, 1989.
[23] Kenji Fukumizu, Arthur Gretton, Gert Lanckriet, Bernhard Schölkopf, and Bharath K Sripe- rumbudur. Kernel choice and classifiability for rkhs embeddings of probability distributions. Advances in neural information processing systems, 22, 2009.
[24] Aude Genevay, Lénaic Chizat, Francis Bach, Marco Cuturi, and Gabriel Peyré. Sample complexity of sinkhorn divergences. In The 22nd International Conference Artificial Intelligence and Statistics, pages 1574-1583. PMLR, 2019.
[25] Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pages 1608- 1617. PMLR, 2018.
11

---

[26] Arthur Gretton, Karsten Borgwardt, Malte Rasch, Bernhard Schölkopf. and Alex Smola. A kernel method for the wwo-sample-probles Advances neural information processing systems, 19, 2006.
[27] Philippe Rigollet Jason Altschuler, Jonathan Weed. Near-linear time approximation algorithms for optimal transport via sinkhorn iteration, 2017.
[28] Edwin T Jaynes. Information theory and statistical mechanics. Physical review, 106(4):620, 1957.
[29] Christian Léonard. A survey of the schr)) odinger problem and some of its connections with optimal transport. arXiv preprint arXiv: 1330.02255 2013.
[30] 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:1 1291-11308, 2020.
[31] 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.
[32] Daniel Lizotte, Michael Bowling, and Susan 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.
[33] Yudong Luo, Guiliang Liu, Haonan Duan, Oliver Schulte, and Pascal Poupart. Distributional reinforcement learning with monotonic splines. In International Conference on Learning Representations, 2021.
[34] Yecheng Ma, Dinesh Jayaraman, and Osbert Bastani. Conservative offline distributional reinforcement learning, Advances in neural information processing systems, 11219235-19244 2021.
[35] 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.
[36] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, al. Human-level control through deep reinforcement learning, nature, 518(7540):599-33,, 2015.
[37] Thanh Tang Nguyen, Sunil Gupta, and Svetha Venkatesh. Distributional reinforcement learning with maximum mean discrepancy. Association for the Advancement of Artificial Intelligence (AAAI), 2020.
[38] 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.
[39] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning. (55-6))355-607, 2019.
[40] Aaditya Ramdas, Nicolás Garefa Trillos, and Marco Cuturi. On wasserstein two-sample testing and related families of nonparametric tests. Entropy, 19(2):47, 2017.
[41] Mark Rowland, Robert Dadashi, Saurabh Kumar, Rémi Munos, Mare G Bellemare, and Will Dabney, Statistics and samples in distributional reinforcement learning. International Confer- ence on Machine Learning (ICMI), 2019.
[42] Mark Rowland, Rémi Munos, Mohammad Gheshlaghi Azar, Yunhao Tang, Georg Ostrovski, Anna Harutyunyan, Karl Tuyls, Mare G Bellemare, and Will Dabney. analysis of quantile cemporal-differeet learning. Journal of Machine Learning Research (JMLR), 2024.
12

---

[43] Mark Rowland, Yunhao Tang, Clare Lyle. Rémi Munos, Mare G Bellemare, and Will Dab- ney. The statistical benefits of quantile temporal-differee difference learning for value estimation. International Conference on Machine Learning, 2023.
[44] 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.
[45] Richard Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. The American Mathematical Monthly, 74(4):402 405, 1967.
[46] Ke Sun, Bei Jiang, and Linglong Kong. How does return distribution in distributional reinforce- ment learning help optimization? In ICML 2024 Workshop: Aligning Reinforcement Learning Experimentalists and Theorists
[47] 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.
[48] Ke Sun, Yingnan Zhao, Yi Liu, Enze Shi, Yafei Wang, Xiaodong Yan, Bei Jiang, and Ling- long Kong. The benefits of being categorical distributional: Uncertainty-awar regularized exploration in reinforcement learning. arXiv preprint arXiv:21 10.03155, 2021.
[49] Gábor J Székely. E-statistics: The energy of statistical samples. Bowling Green State University, Department of Mathematics and Statistics Technical Report, 3005):1-18, 2003.
[50] Harm Van Seijen, Mehdi Fatemi, Joshua Romoff, Romain Laroche, Tavian Bames, and Jeffrey Tsang. Hybrid reward architecture for reinforcement learning. Advances in Neural Information Processing Systems, 30, 2017.
[51] 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.
[52] 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.
[53] Li Kevin Wenliang, Grégoire Délctang, Matthew Aitchison, Marcus Hutter, Anian Ruoss, Arthur Gretton, and Mark Rowland. Distributional bellman operators over mean embeddings. International Conference on Machine Learning, 2024.
[54] 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.
[55] Runzhe Wu, Masatoshi Uchara, and Wen Sun. Distributional offline policy evaluation with predictive eтror guarantees. International Conference Machine Learning, 2023.
[56) Derek Yang. Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Yan Liu. Fully parameterized quantile function for distributional reinforcement learning. Advances in neural information processing systems, 32:6193-6200 2019.
[57] Pushi Zhang, Xiaoyu Chen, Li Zhao, Wei Xiong, Tao Qin, and Tie-Yan Yan Liu. Distributional reinforcement learning for dinlti-dimensionall reward functions. Advances in Neural Information Processing Systems, 34:1 $4:1519-11 1199- 2021.
[58] Fan Zhou, Jianing Wang. and Xingdong Feng. Non-crossing quantile regression for distri- butional reinforcement learning. Advances in Neural Information Processing Systems, 33, 2020.
[59] Florian Ziel. The energy distance for ensemble and scenario reduction. arXiv preprint arXiv:2005. /4670, 2020.
13

---

Appendix
Table of Contents
A Smoother Transport Plan via Sinkhorn Divergence by Increasing
15
Definition of Distribution Divergences and Contraction Properties
16
Proof Proposition
D) Proof of Proposition 2
18
D.I Sum Invariant Property
18
D.2 A Variant of Scale Sensitive Property when
Ko
18
E Proof of Theorem 1
20
& + 0 and c
20
E - and
................
21
(0, +000 and
................................
21
F Proof of Corollary
22
G Algorithm: Sinkhorn Iterations and Sinkhorn Distributional RL.
23
Summary Table for Human Normalized Scores (HNS)
24
Learning Curves on Atari Games
25
Raw Score Table Across 55 Atari Games
26
K Features of Atari Games
27
L Sensitivity Analysis and Computational Cost
28
L.I More results Sensitivity Analysis
28
Comparison with the Computational Cost
29
M Experimental Setting in Multi-dimensional Return Distributions
30
14

---

A Smoother Transport Plan via Sinkhorn Divergence by Increasing E
We visualize the optimal transport plans by solving Sinkhorn divergence with different e in well- trained SinkhornDRL. models across three games F 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 S 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).
E: 1.0, Cost: 0.95
0.2, Cost: 0.40
0.02, Cost: 0,26
Profubiliniic Traneuer Plan
Probabilistic Teanapuet Plan
Probubilatic Traraport Plan
0006
(a) Enduro
E: L.0, Cost: 0.96
0.2, Cost: 0.41
0.02, Cost: 0,26
Probubilittc Transport Plan
Protakllstic Trassport Plan
Probobillttk Trangort Plan
30
0000
0008
(b) Qbert
E: 1.0, Cost: 0.89
E: 0.2, Cost: 0.30
E: 0.02, Cost: 0.17
Transport
Probabilists Transport
Transport
000
20
50
(c) Seaquest
Figure 5: Optimal transport plans for via Sinkhorn Iterations in SinkhornDRL on three Atari games. The first row denotes the (two-dimensional) spatial transport plans across different data points, while the second row represents the heat map of the obtained transport plan (optimal coupling).
15

---

B Definition of Distribution Divergences and Contraction Properties
Definition of distances. Given two random variables X and Y, one-dimensional p- Wassersttin metric W between the distributions of X and Y has simplified form via the quantile functions:
1/8 = =11xx Fy'l
W,(X x)
S -
which F-1 the quantile function, also known as inverse cumulative distribution function, of random variable with the cumulative distribution function as F. The supremal form of W denoted by wee, is defined as
(u,
sup (s.a)E5xA
W (m((,, ((),s
(6)
Further, Ep distance [19] is defined as
4, (X, Y) = [
1/v
[x(ww) - FY du = =I - Fy
The ép distance and Wasserstein metric identical p but are otherwise distinct. Note that
when p = 2, fp distance is also called Cramér distance 6 dc(X,Y) Y). Also, Cramér distance has different representation given by
= =
E X
E Y - Y
where X' and Y' are the i..d. copics of X and Y. Energy distance Cramér distance the multivariate case, which is defined
49 59 is natural extension of
dx X. Y) = E YI
-
E -
E Y'
where X and Y are multivariate. Moreover, the energy distance is special case of the maximum mean discrepancy (MMD), which is formulated as
MMD(X, Y; k) = (E [k X, X +EKK Y') - 2E Y)
(10)
where k(::) -) is a continuous kernel on X' particular, if I is a trivial kernel, also called the unrectified kernel, MMD degenerates energy distance. Additionally, further define the supreme MMD, which is functional SxA x P(X)Sxx
formulated as
(M..
sup (s,a)ESxA
(11)
We further summarize the convergence rates of the distributional Bellman operator I" under different distribution divergences.
I is 7-contractive under the supreme form of Wassertein distance W.,
Txx is //1 ////ccoorracei under the supreme form of Ep distance.
T" is 1 under MMD with the kernel (x, v) = -| - 3/1a, Va 0.
Proof of Contraction in Distributional Dynamic Programming.
Contraction under the supreme form of Wasserstein distance is provided in Lemma
5
Contraction under supreme form of é, distance can refer to Theorem 3.4
Contraction under MMD is provided in Lemma 6 [37].
16

---

C Proof of Proposition I
Proof. We denote two marginal random variables U and V with the pdf u(²) and ply). We next denote the (r, y) as the pdf for II in MI V) KL(MIU @ V). We first prove that the Ml V) is sum-invariant, which is based on the dual form of KL divergence via the variational representation [7
(x)(
sup fEC
- log [By [eftu))
(12)
where _ is the space of bounded measurable functions. The mutual information involves two- dimensional random variables in the KI. divergence. Let U' = + U and V = a+ -+ with pdf y' and v. we denote the joint distribution with margins (x( yía - a) and e'(y) = w(y- - a) as [x, y) whose pdf Pir satisfies Pir (z, y) = pn(æ a, a). Based on the two-dimensional variational representation of KL divergence MIn(U,V) supject (En u)] have:
log (Ee.v [ef(ww)))) }. we
(In + +U,
sup {Eir [f v)] log (EA+ULAtv (ellaw))): SeC
(a)
sup A (B((--ay-a)) [f w)]l log (RA [Bo-uatv (clloan]))
fEC
sup (EA (T(())) [f(x ( a)]] log (EA [Bu.y (oltttaattt))) fEc"
(8)
< sup [E
- EA log (Bu.v (e(lotaatt))))
(13)
SEL*
sup (EA [EA[En]f + a)] log (Evv (croatootta)) ) rech
[2) EA sup tec
[En]s (2 + + a)] log EuE (elletastal))
(
E sup [Euig(z, v)1 - log Eu.y (o0tew0))) NEcb
V),,
where (a) is by the independence of A between X, Y, and the joint edf II. For instance, in the one-dimensional setting, we have Ex=a+x [[[=]] Sa Sa J(x + a)pa(a)ps (x)d EA [Ex [f + a)]]- (b) and (c) are by Jensen's 's inequality in terms of the convex function log(x) and supy. and (d) is because the translated edf is still within L..
Next, we show that MIn is NOT scale-sensitive or with the zero-order T. This result is directly based on the similar property of KL divergence. With slight abuse of notations, we denote U" = at' and V = aV, whose pdfs are pz) 1u( ( and (y( Lv(()) respectively. The scaled joint distribution IT' with the pdf Prr satisfying Prr (z. y) P y/a). Therefore, its marginal distributions are J, 1 pu(x// (all///y//)dg = a n( () and fa pu y/ a/a)dy = vv()), ). We thus have the following result:
MIn(aU, KL( (), & V¹)
f pr (z, v) log (1 v) ('(x)x'())
drdy
= / / p w/a) log p w/a) drdy
(14)
(u(z/a))(//a)
/ pn(z, v) log Pn v) dxdy
m(x)x(y)
MIn V).
17

---

Putting the two properties together and given two return distributions (1 a) and Zz(s, we have the non-expansive contraction property of the supremal form of MIn as follows.
M ( ,,, , 72) ssup MIn (T" (2
9,0
sup (R 8,0 + .'), R(s,a) + 7Z2(6',
(15)
(2)
sup (I d') (2 d'))
5.4
= MIf (2,, 22),
where (a) relies on the sum invariant property of MIn and (b) utilizes the non-scale sensitive property Mln. By applying the well-known Banach fixed point theorem, we have unique return distribution when convergence of distributional dynamic programming under MIn for any non-trivial joint distribution II.
D Proof of Proposition
D.1 Sum Invariant Property
Given two random variables U and with the marginal distributions as and v, and random variable A that is independent of them, we aim at proving
We, (A A V) S (U,
(16)
According to
39 we have the dual form of We,si
{ p(o)) + J Hu)e,dy - exp
(0(x) + v1(1)) - c(x,3)
(U, = sup
M dady
9.00
(17)
En (p((()))] + [+(w)] eEu,p exp p(x) + v1y) c(x,
sup
8
Therefore, we have:
(a.(A U,
(0 + ((4) c(x. 11
sup 2..
EAtU [p(x))] + EA+V [v(()] CEA-U.A4V
exp
60 + a) + (b + a) - c(x, 8)
sup
EA
E2 [ip( a)] + E, [v(2 a)] - =BE,,
exp
v.v
(18)
(0)
60 + a) + 1/ + a) c(2,
RA
sup
E [1p a)) + F, (v + a)] ETR
exp
10.0
E
(c)
f(x)
g(v)
c( m)
sup
E2 [f())) R. [o())] EBN, exp
f.g
=
(U.V..
where (a) relies 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 (p, v.
D.2 A Variant of Scale Sensitive Property when == -kia
General Conclusion. Let II be the optimal coupling for We.s. define ratio A-U V) #KL(II =KL(II)HHe)) [4@M] (0, 1) for any considered U, V with measures to compare, where the denominator
We.e
18

---

Wa,e is generally non-zero. We thus have the following result:
We,(aU,,V) S de W),
(19)
where the scaling factor Ae(a, a) = lal (1-spury \U, +supuu Ae(U.V) suud,y \U, > 0. The ratio >= V) measures the proportion of the entropic regularization
E (lal", 1) with
term over the whole divergence term We,e, i.e., (U, = cKI (1(W)) E (0,1). 1). Under the mild assumption of a finite set of probability measures, we have supu, (U, > 0. To elaborate the reason behind it, we first know that (U, < 1 for any U and V with their measures on the probability measure set. If this set finite, the ratio set that contains all {>= V)} is also finite. Based on the fact that the real set is dense, we can directly find a positive lower bound X* for the ratio set, such that {Ae(U, V)} I X < 1. This implies that supu,v A= V) = maxy,v Àe (U, V) < 1. Notably, this finite set property of the ratio avoids the extreme case that may lead to conservative conclusion about a non-expansive distribution Bellman operator, which we will give more details
Was
later.
Scale-sensitive Property. By definition of Sinkhorn divergence [ 39), the pdf of Gibbs kernel in c(a.))
the equivalent form of Sinkhorn divergence is KC V), which satisfies KC(U, V) x dIC
p(z)v(u)
In particular, the pdf of Gibbs kernel is defined as π[pाс (z, y) Texp(- -e/e -/eedduusu) where the denominator is the normalization factor. After a scaling transformation, the pdf of al and aV with respect to x and y would be +, = and (²)()) Thus K(aU, aV) -e(x.u)
exp(-c/a)
= (v())) In the following proof, we use the change variable formula (multivariate version) constantly, while changing the joint pdf x (Ir, y) and keep the cost function term c(2r,y). y). In particular, we denote II" and 17 as the optimal joint distribution of We,e(1, v) and Wape(ap, av). Then we have:
/ c(z, w)dlI (z, + sKL(II aa @ av)
Wa,(aU, aV) =
/ cxx, y)dII v) + =IL lop av)
c= ha
(I y) 1 _ 인 dxddy + 8 / ( a= a
az --= T y l
E dody
( 2)
a
= |a|0 / [(o y) "0, yxxdy +e / n" (",,,)) log T" v) dxdy
u(z)u(y)
= |a|% f(o - y)a² (x. dxddy + (la) +1
xt (z, y) dady x (x, log m(z)v(y)
(U, V) + (1 (1 (al)) ()
= - A.VV (o.,))
(20)
where AV. a)) =
X. ++ E ([a] 1) for
E (0, +o0) and non-trivial Wo, (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, A.. (a, a) is a function less than I, which depends the two margins, including their independence and distribution similarity, the scale factor a and the order a.
due to the fact that Xe V) 10, for any non-trivial We.e (U, V). The
Ruling Out Extreme Cases in the Convergence via Finite Set. However, the fact that A. (a, a) < 1 can only guarantee a "conservative" non-expansive contraction rather than a de- sirable 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ı for the k-th distributional Bellman update, where qk < 1. We can construct a counterexample as qk = 1 - 1/(k- + 2)2. In this case, 11 qx = ( 3)) > instead of the convergence to 0 and the non-zero limit can not guaran- tee the contraction. It also intuitively implies that iteratively applying distribution Bellman operator under We,, may not lead to convergence in general by considering all possible return distributions given the non-constant factor A (a, at). Although we know these extreme cases are very unlikely to happen, we have to rule out these extreme cases for rigorous proof. As we have the assumption
19

---

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 s (U, V) the uniform upper bound across the whole set of interested probability measures. Under this condition, we can immediately find universal upper bound of A U.V (a, a):
sup A U.V (a, a) | + (1 U.V
(al") sup Ag (U, V) U.V
=aa|0(1 sup (U, V))
sup (U,V) U.V
(21)
U.V
= a.)
where the upper bound su AU. (a, ax) has an interpolation form, which can be viewed as the convex combination between |a|a and I with the coefficient supu, (U, V) determined by the probability measure set. More importantly, suPu,v A.. (a, is strictly less than 1, which is guaranteed by the finite set of {A=(U, V)} thanks to a finite set of interested probability measures. Finally, have the variant of scale-sensitive property follows, where the factor De Aa, a) depends on a, a and the probability measure set.
Wa,CaU,,Va S a)
(22)
E Proof of Theorem I
E.1 E - 0 and -ka-
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, Wo. 2 (x "II'III" (x, y)dzdy= 2 W as KL 2 0. We want to provide the inequality on the other side. Denote TT' the minimizer in the Wasserstein distance Wo. For any >>> O, there always exists joint distribution 11 such that
1f (c -) v) / y) 4) (x, v)1 <
(23)
and KL([T'|u @ v) < +o0, i.e., f(x y) (x, v) fff=)) u) dl'' (x. y) S ö. One possible way to find 11s is provided in notes of Lecture 6 in Optimal Transport Course and we invite interested readers for reference. It follows that
Wa S We,5 S f(0- 1)" (ddd (z, v) + EKL(I
( ( =KL( v),
(24)
where the RHS fx (x- dII' (II' + ( + (x=))))' (a. y) + W + 6
as e > 0. As > 0 is arbitrary. combing the two sides, it shows that Wc.e + W as E - 0. Thus, Sinkhorn divergence maintains the properties of Wasserstein distance when E - 0.
When & = 0, has been shown that Wa can guarantee a y-contraction property for distributional Bellman operator 5 The crux of proof is that W., is 7-scale sensitive:
(rentie inf (TEMaaaaaa) f a0(x== (0(x=y)ddlla) 8)
1/a
Wa(aU,aV)
S <a (TTTTT / ( 1) ((-
1/a
(25)
inf
(U,VV,
where the inequality comes from the change of optimal joint distribution. Therefore, W. (aU, aV) S aW. (U, V) guarantees --contraction property for the distributional Bellman operator. As such, for Wa when 0, it suggest that We Wo corresponds to 7 ccotrrcction for the distributional Bellman operator T
https://lchizat github. 1o/ott020orsay.bbtt html html
20

---

E.2 & - do and -ka-
Our complete proof is inspired by 40 25 Recap the definition of squared MMD is
E
lk
(X,x' | (Y, Y') - 2E[ Y)
(26)
When the kernel function & degenerates to unrectified ka(2,)) == -||x yll for (0, the squared MMD would degenerate to
X X Y -EJX X E Y Y
(27)
where X, X' i.d. ~ Y, Y Lid. ~ V and X, Y, Y' are mutually independent. On the other hand,
by definition, have the Sinkhorn loss as
We,0 (1, v) = (1, - W1 p)) - (V,
(28)
Denoting 11g be the unique minimizer for We.. it holds that 1I; 1 2 u ase a 00, which is the product of two marginal distributions. That being said, We (u, - c(a, ((z,))dp(zddd0) +
c(x, y v)dp((z)dv(y)) One important proof insight here is although & - +00, the KL term tendy to zero, which is faster than e. Therefore, the whole regularization term still tends to 0 as & - +00. If -i ||x-yll" yll", we eventually have =o v) ||a--y] dll"dd((=))dlll where fa and v can be inherently correlated, although the minimizer degenerates to the product of the two marginal distributions. Finally, can have
= X Y|
W W + 2B X -
X
(29)
which is exactly the form of squared MMD with the unrectified kernel k., Now the key is to prove that II, + μ @ v as + co. We give the detailed proof as follows.
Firstly, it is apparent that (e,(p, < c(x,, 3)d as @ E I[(,,)). Let {ER} be positive sequence that diverges to co, and I1, be the corresponding sequence of unique minimizers for Wo.s According to the optimality condition, it must be the case that f c(x, y)dll, + ELL,, / &
v) S c(æ, y)dp & v - (when 11( v) & v). Thus,
(fcduove V- ( /dllx - 0. Ek
1
KL (4 8 &
Besides, by the compactness of I[ v), we can extract converging subsequence II, + II Since KI. weakly lover-semicontimuos, holds that
KL (1 @ v)
< lim inf KL (IIm 2 v) = 0 k-000
Hence 11x = 4 @ v. That being said that the optimal coupling is simply the product of the marginals,
indicating that II, + l @ v as E + 0o. As a special case, when a 1, W oo (u, v) is equivalent to the energy distance
d (X, Y) 2E X YI E X B Y'I.
(30)
In summary, if the cost function is the rectified kernel kn, it is the case that - converges the in the supremal form of MMD
squared MMD co. According [37] T"d is a with the unrectified kernel ka. As Wa.s (H, v) + MMD (u, v), which is squared MMD instead of MMD. implies that I is -ontttractevv under the squared MMD W We.+oo.
E.3 (0, +00) and -Ka
In the proof of Proposition] 2 we have shown that the Sinkhorn loss Wa,s satisfies the sum- aum-invariant (I) and new variant of scale-sensitive properties as follows:
Wa.e (A U, A V) S (U,)(
(31)
Wa, aV) < (c.,,.))...(TVV)
The Sinkhorn divergence War is defined by additionally subtracting two self-distance terms (W,.(P, and ( v)) based on We,c(u, v) order to guarantee the non-negativity, tri- angularity and metric properties. These two self-distance terms do not change the (I) and (S)
21

---

properties when extending We,e to We.e, and some proof techniques can refer to Section 2 in 21
The only difference is that the scaling factor will be A U.V (a, a), which is the counterpart of Eq 20
satisfying
U.V (ae S A (a,a)
(32)
where A (a, = |a| (0, + (U, E (la|",1) for E (0, +00) and a < 1
due to the fact that (V, E 10, for any non-trivial (V, The new ratio Xe(U, = (48V]
(0, for any considered U, V with measures jau in the interested probability measure set. In particular, in the context of distributional RL, the set of interested probability
Wo.
measures would be the return distribution set of {Z(s, a)} for s E S and a E A in a given finite MDP. We now want to find the universal upper bound Ag (a, a), which is defined as
(a ar) = |a (1 sup (U, + sup Xe (U, V)) E (la)", 1). U.V
(33)
U.V
Following the proof in Appendix D the finite MDP guarantees a finite ratio set of (Ae V)), and
thus can find universal upper bound X of the ratio set such that (X, v)) I X < 1. This also implies that supu,v Ae( V) E (0, and thus the scaling factor (a a) (la|", 1), which is strictly less than Therefore, we have the (I) and (S) properties of We.::
We, (A + U, U (U,(U, As (l,,,a))...(UVV)
(34)
Wa,
Putting all together, we now derive the convergence of distributional Bellman operator I" under the supreme form of We,e, i..., W
Wat (I" Z1, x- Z2) = sup We. (F..
3,
+ (2
8,a
(35)
A 7i(s' . 7.7aa (7, a) W. (Z) a'), (5))
Zi(s' .a)) Z')' aa) (7, a) sup W. (1 d'), Zz(5', ,'))
I sup A s',d'
aa
2 = -
c.e (Z 2z)
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 = = Notably, (7 a) E (hil", 1) is dependenndnnt constant (determined by the return distribution set), which is also determined by 7.8 and a. As such, we conclude that distributional Bellman operator is at least Ae (7, a --conttra,,, where the contraction factor A ar) is strictly less than I in given finite MDP. Based on the existing Banach fixed point theorem, we have unique optimal return distribution by applying the distributional Bellman operator ITH in the distributional dynamic programming when convergence.
F Proof of Corollary I
Proof. The contraction conclusion that extends to the multi-dimmnsional 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(xr, v) (x y) to c(x, ||x-y|l and re-define two d-dimensional random vector and V with the d-dimensional probability measure μ and Therefore, the (I) and (S) properties in the multi-dimensional reward settings are:
We.c(A + U. A + V) < W..(U, (aU,
(36)
Wo. aV) - Az Wo. V),
22

---

where A is d-dimensional random vector independent of U and V.
By leveraging these properties, now derive the convergence of distributional Bellman operator
2 under W in the joint return distribution setting. Given two d-dimensional return distributions Z and Zg. we have
We L1, 12 = sup W (c.cc s,a
sup We e, [ ( R (3((3)))
(37)
A 2l)s'aa)),zz(t (1, a W c.e (Z (z'' Zz('',)
Zi(d's aa' 22 (7. a) sup W c.e (Z , d²), (s's,))))
I sup 8'
A
A7 a: W (Z, Zz)
where the inequality (a) is based 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 -k.. where the contraction factor (e a)) will depend on the set of d-dimensional probability meassurs/disstriuutions Notably, the analysis of (7, a)) in the one-dimensional return setting established in Appendix D and Appendix also applicable in the multi-dimensional setting.
G Algorithm: Sinkhorn Iterations and Sinkhorn Distributional RL
Algorithm 2 Sinkhorn Iterations Approximate Wa,e ,
Input: Two samples sequences {2,) {32z) y number of iterations L and hyperparameter
1: G.j clZi, 52,) for Vi 1, N,j 1,.... N
2: K1,; =exp(-Cr,j 22
3: bo f- 11N
4: for = 1, 2, ..., do
5: 01 Rb
IN bi t
1v
Kai
6: end for
7: Was
(2)) ( (32,),1)) ((K 0)6,,
Return: Wo.s [12.) (521)
Given two sample sequences (2)),- (32,) in the distributional RL algorithm, the optimal transport distance is equivalent to the form:
min PER
{P, PIN 1N. PT 1N 1x).
(38)
where the empirical cost function is G.j c(Z,, IZ,). By adding entropic regularization optimal transport distance, Sinkhorn divergence can be viewed to restrict the search space of P in the following scaling form:
P.y=axK..bb
(39)
where Ki,y - A is the Gibbs kernel defined in Eq. This allows to leverage iterations regarding the vectors and b. More specifically, we initialize by = 1x, and then the Sinkhor iterations are expressed as
01+1 f- In and bit1 f KTa441"
1N
(40)
Kbb
23

---

where : indicates an entry-wise division. Combining Sinkhorn Iteration in Algorithm 2 and the generic update Sinkhorn Distributional RL in Algorithm ! Distributional RI. algorithm in Algorithm
we provide full version of Sinkhorn
Algorithm 3 Sinkhorn Distributional RL
Require: Number of generated samples N, the kernel & (e.g., unrectified kernel,, discount factor y [0, 1], learning rate a, replay buffer M, main network Zo. target network Zø number of iterations L, hyperparameter E, and behavior policy x based on Zg following an e-greedy rule
1: Initialize 0 and 0 +- 0
2: for = 2, do
3: Take action at T[-1(1; f), receive reward r, R(-18t, az), and observe 8t+1 P(-(st, 04)
4: Store (st, dg, rt, 8t+1) to the replay buffer
5:
Randomly draw batch of transition samples (s, a, 8') from the replay buffer M
6:
arg max N Di=1 Zg (8', (6',')): 7: Compute the target Bellman return distribution: IZ; r 7 Zo- (s' ".,a);; 1A SiSN
Compute greedy action:
8:
Evaluate Sinkhorn divergence via Sinkhorn Iterations in Algorithm
War
(52,)
9: Update the main network Ze: 0 t - aVoWe. (Zooss ( (
Periodically update the target network 0" + 0
11: end for
H
Summary Table for Human Normalized Scores (HNS)
Mean
IQM (5%
Median
>QQ
DQN
452.6 %
181.2 %
32.8 %
0
C51
640.2 %
368.5 %
68.5 %
35
QR-DQN-1
780.9 %
401.9 %
85.8 %
38
MMD-DQN
781.7
428.4
96.6 %
37
SinkhorDDRL
1306.1 %
477.0 %
91.1 %
41
Table Evaluation best Human Normalized Scores (HNS) across 55 Atari games. Results are averaged over 3 seeds. Our proposed SinkhornDRL achieves the best performance in terms of Mean and IQM(5%) HNS as well as the DQN" metric, and is on par with MMD-DQN in terms of Median of HNS.
24

---

I Learning Curves on 55 Atari Games
Arrdtrr
Apes
Actertt
Autorrice
Alaerti
Dock
Blowling
Baxiing
Preman
Fremees
Broakout
CropparCammmac
CruzyCliber
Domonkttcck
Dout
Frehuro
Freetway
Frrethiw
Goophn
Cravitar
IxaHookky
Kangaanoo
KungF Kuheeihhheder
Movengu
NameTiaGame
Piral
Framas PHO
PrivloEye
Oteert
Rivoraid
RoadRunter
Roberons
Seaquest
Sodaere
Space modders
StanGerneer
Tennnte
Timeliiott
Tutoridian
UsNDown
VideoFinball
WizaroOfNioe
YarsRovenge
Mnltorn of
Zaxxon
Figure 6: Learning curves of SinkhornDRL 55 Atari games after training 40M frames averaged over seeds.
25

---

J Raw Score Table Across 55 Atari Games
GAMES
RANDOM
HUMAN
DQN
C51
QR-DQN
MMD-DQN
SinkhornDRL
Alien
211.9
7.127.7
T030
1510
1030
1480
1560
Amidar
2.34
1,719.5
341
424
677
510
588
Assault
283.5
742.0
3232
3647
12943
3295
2960
Asterix
268.5
8,503.3
3000
34900
11500
14900
6500
Asteroids
1008.6
47,388.7
1180
780
1650
1080
1370
Atlantis
22188
29,028.1
15500
84900
3316700
93600
3447100
BankHeist
14
753.1
570
960
980
880
700
BattleZone
3000
37,187.5
15000
19000
26000
35000
32000
BeamRider
414.3
16,926.5
8200
7476
7642
25602
6022
Berzerk
165.6
2,630.4
970
650
640
860
910
Bowling
23.48
160.7
54
Boxing
12.1
100
100
Breakout
1.5
30.5
343
452
414
432
418
Centipede
2064.77
12,017.0
7551
4133
5388
9342
4070
ChopperCommand
794
7.387.8
1500
3600
3500
3600
3400
CrazyClimbe
8043
35,829.4
94300
153100
139500
98500
137400
DemonAttack
162.25
1,971.0
31420
5U240
240660
407030
105185
DoubleDunk
-18.14
-16.4
-16
-20
-18
-22
-12
Enduro
0.01
860.5
1387
1086
1972
1953
4608
FishingDerby
93.06
-38.7
23
-1
31
31
61
Freeway
0.01
29.6
31
32
34
33
34
Frostbite
73.2
4.334.7
3330
3690
3470
3250
2640
Gopher
364
2,412.5
11400
14780
5440
3740
6620
Gravitar
226.5
3,351.4
350
350
750
350
500
Hero
551
30,826.4
3440
8535
10155
7195
6540
TceHockey
-0033
0.9
-13
-10
-3
-2
Jamesbond
27
302.8
350
600
650
450
500
Kangaroo
54
3,035.0
1300
14600
14800
3600
Krull
1,566.59
2,665.5
8892
9336
10053
7762
9630
KungFuMasstr
451
22,736.3
46500
38000
27900
26900
43600
MontezumaRevenge
0.0
4,753.3
400
1
MsPacman
242.6
6,951.6
3230
2440
1860
3130
5120
TameThiGamme
2404.9
8.049.0
6160
5750
13580
9350
11250
Phoenix
757.2
7,242.6
18780
9390
25690
23300
Pitfall
-265
6,463.7
Pong
-20334
14.6
21
20
20
21
21
PrivateRye
34.49
69,571.3
100
100
100
100
100
Obert
188.75
13,455.0
7425
16375
7800
16225
7750
RaverRaid
1575.4
17.118.0
8470
13310
8710
9190
9530
RoadRunner
7,845.0
45500
60900
$2500
45600
59500
Robotank
2.24
11.9
11
58
39
54
Seaquest
88.2
42,054.7
1740
5940
2640
7370
8350
Skiing
-16267.9
4,336.9
-13681
-20495
299770
89966
-23455
Solaris
2346.6
12.326.7
1640
660
2200
3380
7720
SpaceInvaders
136.15
1,668.7
1170
770
1200
StarGunner
631
10.250.0
1200
17200
$2900
52500
57500
Tennis
-3.92
-8.3
-23
-1
-7
-8
5
TimePilot
3682
5,229.2
800
4100
4400
8000
4500
Tutankham
15.56
167.6
201
213
220
141
137
UpNDown
604.7
11,693,2
14560
18440
13710
27370
18910
Venture
1,187.5
700
VideoPinball
15720.98
17,667.9
155165
576843
189460
69175
347700
WizardOlWor
534
4.756.5
1400
2400
14300
11500
4300
YarsRevenge
3271.42
54,576.9
28048
7882
17729
7520
9120
Zaxxon
9,173.3
3900
9100
4300
19500
Number Best
Number Second Best
Table Best score of all algorithms over seeds across 55 Atari games after training 40M Frames. Bold denotes the best performance, while the underline represents the second best performance. The number of games with the best and second best performance substantiate the superiority of our SinkhorDDRL across all considered baseline algorithms.
26

---

K Features of Atari Games
GAMES
Action Space
Dynamics
Alien
Complex
Amidar
Simple
Assault
Complex
Asterix
Complex
Asteroids
Simple
Atlantis
Simple
BankHeist
Simple
BattleZone
Simple
BeamRider
Complex
Berzerk
Complex
Bowling
Continuous
Simple
Boxing
Simple
Breakout
Simple
Centipede
Complex
ChopperCommant
Continuous
Complex
CrazyClimber Climber
Complex
DemonAttack
Complex
DoubleDunk
Simple
Enduro
Simple
FishingDerby
Simple
Freeway
Simple
Frostbite
Complex
Gopher
Simple
Gravitar
Continuous
Complex
Hero
Simple
IceHockey
Continuous
Simple
Jamesbond
Complex
Kangaroo
Complex
Krull
Complex
KungFuMaster
Complex
MontezunaRevenge
Complex
MsPacman
Simple
TameThiGGame
Complex
Phoenix
Complex
Pitfall
Complex
Pong
Simple
PrivateEye
Complex
Qbert
Complex
Riverraid
Complex
RoadRunner
Simple
Robotank
Simple
Seaquest
18
Complex
Skiing
Simple
Solaris
Complex
Spacelnvaders
Simple
StarGunner
18
Complex
Tennis
18
Simple
TimePilot
Complex
Tutankham
Complex
UpNDown
18
Complex
Venture
18
Complex
VideoPinball
Simple
WizardOrWor
Complex
VarsRevenge
Complex
Zaxxon
18
Complex
Table Number of Action space and difficulty of environmental dynamics of 55 Atari games.
27

---

L Sensitivity Analysis and Computational Cost
L.1 More results in Sensitivity Analysis
Breekout
Breakout
Breakout
Seaquest
Averags Ratum
2008
CTT0
20100
ons Frames (M)
Miloon Framnes (M)
Millore Frames (M)
Millions Frames (M)
(a) Small E vs QR-DQN (b) Large& vs MMD-DQN
(c) Breakout
(d) on Seaquest
Figure 7: (a) Sensitivity analysis w.r.t. a small level of SinkhomDRL to compare with QR-DQN that approximates Wasserstein distance Breakout. (b) Sensitivity analysis w.r.t. large level of SinkhomDRL algorithm to compare with MMD-DQN on Breakout. All learning curves are reported over 2 seeds. (c) and (d) are results for general Breakout and Seaquest, respectively.
Decreasing 5. We argue that the limit behavior connection as stated in Theorem 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 E to SinkhornDRL's performance tends to degrade and approach QR-DQN. Note that overly small will lead to trivial almost KC., in Sinkhorn iteration Algorithm 2 and will cause numerical instability issue for and by 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 e tends to worsen the performance.
Increasing e. Moreover, for breakout, if we increase E, the performance of SinkhornDRL tends degrade and be close to MMD-DQN as suggested in Figure (b). It is also noted that an overly large will let the Ks.; explode to oo. This also leads to the numerical instability issue in Sinkhorn iteration in Algorithm?
Samples N. We find that SinkhomDRL 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 g (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.
Breakout
Seaquest
Bumples
Samplie-3
Assrage Ratant
Samplee-
Samg Samplea-500
Millons of (MI
Millons of Famms (N)
(a) Samples Breakout
(b) Samples Seaquest
Figure 8: Sensitivity analysis of Sinkhorn in terms of the number of samples N on Breakout (a) and Seaquest (b).
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., and overall large samples tend to degrade the performance, especially on Zaxxon. Although the two games are robust to and we find a small large & hurts the performance in Seaquest. Thus, considering all games, we set samples 200, and & 10.0 moderate range across all games, although more careful tuning in each game will improve the performance further.
28

---

StarGunner
StarGunner
Zaxxron
Zaxxon
Sa Tarpheiit
Sam Damplennt
Sa rpieerctt
Aver 1g0 Rooum
Aveeage Reture
Samphers
+0004
Samplea-500
Samphee:201
Aoarage
Samr pinentoo
Milions Frames
Milions Frames
Millons ofFiomes (M)
Millons Frames (M)
(a) StarGunner: E
(b) StarGunner: Samples
(c) Zaxxom: E
(d) Zaxxon: Samples
Figure 9: Sensitivity analysis of SinkhornDRL on StarGunner and Zaxxon in terms of E, and number of samples. Learning curves are reported over 3 seeds.
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 comparison Figure T0 suggests that SinkhomDRL indeed increases around 50% computation cost compared with QR-DQN and C51, but only slightly increases the cost 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 a large computational burden illustrated in Figure IT This can be intuitively explained as the computation complexity of the cost function G.g (()) in SinkhornDRL, which particularly heavy the computation if N is large enough.
Breakout
Qbert
Average Computation Time
Average Computation Time.
DQN QR-DQN C51
MIMD-Sinkhor
DQN QR-DQN C51
MMD Sinkhom
Figure 10: Average computational cost per 10,000 iterations of all considered distributional RL algorithm, where we select 10, and the number of samples N 200 in SinkhornDRL. algorithm.
Breakout
Qbert
Avargg: Computaaion Time
Avamg: Computaaioo
ineg
Figure 11: Average computational cost per 10,000 iterations of SinkhornDRL. algorithm over different samples.
29

---

M Experimental Setting in Multi-dimensional Return Distributions
Reward Structure and Decomposition. In practice, the reward function can be multi- 15 57 3II, where distributional RL is aimed at modeling multivariate return distribution with multiple reward sources. We follow the multi-dimeensonal 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.
dimensional 550, 30 32
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-dimennionall rewards [100, O, O, 0], [0. 75, O, 0]. [0, O, 50, 0],[0, 0, 0, 25], [0, O, O, O] respectively for the primitive reward 100, 75, 50, 25 and O.
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 [ri, T2, ra]. If (r - 20) mod 50- = 0, we let r = 20, otherwise = 0. If (r - 50) mod 100- = 0, we let r2 50, otherwise r2 = 0. We let 73 = r T1 T2-
Gopher. For primitive rewards, the agent gets +80 reward for killing monster and +00
reward after removing holes on the ground. We denote the primitive reward as r, and decompose it into the two-dimensions as [ri, r2.]. If (r - 20) mod 100 = 0, we re = 20,
otherwise 11 O. We let 12 = T1-
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 decom- position, we decompose primitive reward denoted as into four dimensions [ri, r2, F3, ra]. If (r 10) mod 50 = O, we let rı = 10, otherwise r1 O. If (r 1ı 50) mod 100 0, we let 72 = 50, otherwise r2 = 0. If (r
r2 100) mod 200 O, we let r3- = 100, otherwise r3- = O. We let = r ---r-- - r2 T3.
Pong. For primitive rewards, the agent gets +1 if wins a round, and -1 for losing the round, We decompose the reward into -iw-dimennion: the agent will get [-1, 0] for a -1 reward, [0. 1] for a +1 reward; otherwise. [0.
For primitive rewards, the agent gets +400 reward for killing an energy car, +100 for reaching a flag, and + reward for being alive. We denote the primitive reward as r, and decompose it into the dinee-iimesiioal reward as [r1, 72, r3]. (r 10) mod 100- 0, we let r1 = 10. otherwise T1 = 0. If (r T1- 100) mod 200 = O, we let re = 100, otherwise r2 O. We let r3- - T1 - T2-
UpNDown.
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 clipping reward normalizer to clip the miltiidimmesioonll 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, AAl 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 W-saseerstein 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 miltii-dimensional Wasserstein distance. This say, is still unclear how extend the quantile regression distributional RL (QR-DQN) into dinlii-dmeesionall return distribution setting, resulting in no proper baseline in our experiment. Despite that, we directly compare SinkrhornDRL with MMD- DQN [77 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 Din
has already surpassed MMD-DQN across almost all the considered games, and thus excels over
30

---

MD3QN and HRA, correspondingly. 2) The primary focus of our study is the comprehensive advantages of SinkhorDRL over other distributional RL classes, especially in the more common setting within one-dimensional return distributions. The extension capability of SinkhornDRI. into the multi-dimensional reward setting is one of its merits, which is not the primary focus of our study.
31

---

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 distributional reinforcement learning algorithm by leveraging Sinkhorn divergence, 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 the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made 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.
is fine to include aspirational goals as motivation as long as is clear that these goals 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 separate paragraph of the last section.
Guidelines:
The answer NA means that the paper has 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 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 few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails 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 fairesss.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged
the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant 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
32

---

Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We have few propositions and theorems in our study our main theoretical contributions, and the complete proof 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 cross- referenced.
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 provide 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
Does the paper fully disclose all the information needed to reproduce the main perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We have provided the detailed experimental settings in both the experimen- tal 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 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 dataset and/or model, the authors should describe the steps taken make their results reproducible 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 be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to model. In general. releasing code and data often one good way to accomplish this, but reproducibilityy can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, other means that are appropriate to the research performed.
specific model and empirical evaluation, may
While NeurlPS does not require releasing code, the conference does require all submis- sions 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 reproduce that algorithm. (b) If the contribution is primarily the architecture clearly and fully.
model architecture, the paper should describe
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open- 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 describe the particular way they provide for reproducibiliity In the case of closed-source models, it may be that access to the model is limited in
33

---

some way (e.g., to registered users), but it should be possible for other researchers have some path 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 instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
We have provided the code the supplementary files for faithful reproducibil-
ity.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips. ec/
public/guides httaa//guuiis//oooeuubiiissssss oolliccyyyy for more details.
While 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 open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips. //nips.c/pput httpol/i/ggiidee//oooooeo
ssmeiissinniiii ssionPolicy for more details.
The authors should provide instructions 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. only 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 much information possible in supplemental material (appended the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/eetails
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary 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 include experiments.
The experimental setting should be presented in the core 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, appendix, 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.
34

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall 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).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 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 eπror 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 tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We run experiments 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 compute required for each of the individual experimental rums well 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 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 http:::/hhorrrpp..cpppl/////tttt
Answer: [Yes]
Justification:
We have preserved anonymity and conformed with the NeurlPS 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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [NA]
35

---

Justification: As our study is to propose a new family of reinforcement learning algorithm and 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 partieular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out, For example, it is legitimate to point that 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 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 attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe 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)?
Question:
Answer: [NA]
Justification:
We believe our paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 dataset.
36

---

The authors should state which version of the asset is used and, if possible, include 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 souree should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperssi thoode com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 encouraged to reach out to the asset's ' creators.
13. New Assets
Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Question:
Answer: [NA]
Justification: We believe our paper does not release new assets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the datasct/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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create 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 sereenshots, 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 contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According to the NeurIPS Code of Ethics, workers involved 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 research with human subjects.
Guidelines:
37

---

The answer NA means that the paper does not involve crowdsoureing 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 expect authors adhere 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.
38


Reviews:
[{"review_id": "KJKSCSVoco", "reviewer": "Reviewer_nXJK", "summary": "This paper introduces Sinkhorn Distributional Reinforcement Learning (SinkhornDRL), a new algorithm designed to address the limitations of current distributional RL methods, particularly those relying on quantile regression. Existing methods often struggle with accurately capturing the characteristics of return distributions and extending to scenarios with multi-dimensional rewards. SinkhornDRL leverages Sinkhorn divergence, a regularized Wasserstein loss, to minimize the difference between current and target Bellman return distributions. This approach combines the geometric advantages of Wasserstein distance with the computational efficiency of Maximum Mean Discrepancy (MMD). The paper provides theoretical proof of SinkhornDRL's contraction properties, demonstrating its convergence behavior and relationship to other distance metrics. Empirical evaluations on the Atari games suite show that SinkhornDRL consistently outperforms or matches existing algorithms, especially in settings with multi-dimensional rewards.", "questions": "Could 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?", "limitations": "NA", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- 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.", "weaknesses": "- 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper. We appreciate your positive assessment and insightful feedback and would like to address the concerns you raised in your review.\n\n\n>Question 1: Are there existing works on distributional reinforcement learning that use the entropic regularized Wasserstein distance?\n\nTo our knowledge, our study is the first to investigate the entropic regularized Wasserstein distance in the context of distributional RL.\n\n>Question 2: What specific advantages does the Sinkhorn divergence offer over the entropic regularized Wasserstein distance in the context of this study?\n\nAs highlighted in our introduction, the specific advantages of Sinkhorn divergence, which is an entropic regularized Wasserstein distance, over the vanilla Wasserstein distance in distributional RL can be summarized as follows.\n\n* **Enhanced accuracy in representing return distributions** Sinkhorn divergence uses the samples to depict return distributions, offering more flexibility than the pre-specified statistics, e.g., quantiles, used in vanilla Wasserstein distance. Consequently, our approach naturally circumvents the non-crossing issues of the learned quantile curves.\n\n* **Applicability to multi-dimensional rewards** Many RL tasks involve multi-dimensional reward structures, but it remains elusive how to use quantile regressions to approximate a multi-dimensional Wasserstein distance. In contrast, Sinkhorn divergence can effectively approximate a multi-dimensional Wasserstein distance, thus providing an efficient solution to address these RL tasks,\n\n* **Increased robustness to noises via smoother transport plans** The entropic regularization incorporated in Sinkhorn divergence fosters smoother transport plans compared with those derived from unregularized Wasserstein distance. This makes it less sensitive to noises in the learning process.\n\n* **Stable optimization in RL learning** The entropic regularization turns the objective into a strongly convex problem, and the induced smoothness facilitates faster and more stable convergence in RL learning.\n\nWe acknowledge the typos you pointed out and will correct them. Thank you once again for the time and effort you dedicated to reviewing our work.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VVWHgqqnIK", "reviewer": "Reviewer_Fdqh", "summary": "•\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.", "questions": "•\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?", "limitations": "•\tIn line 22, the reference [11] is cited as a risk-sensitive control, but this is close to exploration method for distributional RL.", "rating": 5, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "•\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.", "weaknesses": "•\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper. We appreciate your positive assessment and insightful feedback and would like to address the concerns you raised in the Weakness and Question parts of your review.\n\n>Weakness 1: The text and figures in the paper are quite dense and difficult to read. Some parts seem unnecessary to be included in the main paper ....\n\n* We acknowledge that Algorithm 1 does not involve too many technical details as it is aimed at outlining a generic update procedure of SinkhornDRL. As mentioned in Line 264, we have also provided a full version of SinkhornDRL in Algorithm 3 of Appendix G, in which we add more necessary technical details. In terms of Algorithm 2, given that Sinkhorn iteration is a well-established algorithm with guaranteed convergence, we thus leave it in the Appendix. We appreciate your suggestion and will add more explanations about Algorithm 2 in the appendix.\n\n* Thanks for this great advice. Following your suggestion, we have also **included a table to summarize the results in the global response**. It shows that our proposed SinkhornDRL achieves the best performance in terms of Mean and IQM(5%) HNS as well as the \"> DQN\" metric, and is on par with MMD-DQN in terms of Median of HNS.\n\n>Weakness 2: There may be a technical error regarding Theorem 1 (3). ...uniform upper bound and finite set assumption.\n\nThe contraction conclusion indeed requires a uniform upper bound (also after taking the supremum) over all possible return distributions of $\\Delta_{\\epsilon}^{U, V}$ to remain strictly below 1, which can be satisfied by straightforwardly assuming a finite (**yet arbitrarily large** in the state and action spaces) MDP. The finite MDP ensures a finite return distribution set, further implying a finite set of $\\lambda_{\\epsilon}(U, V) \\in (0, 1)$. Consequently, it is straightforward to establish that $\\sup_{U, V} \\Delta_{\\epsilon}^{U, V} < 1$ as shown in Eq.21 under this finite condition. We argue that the finite MDP is not a strong condition but a common setting as most RL theory is established on it. In addition, we can also relax this finite MDP assumption by ruling out the extreme case, where there exists a series of elements in this set that can be arbitrarily close to 1. We have provided an example of this extreme case in Line 666. Since expressing this extreme case mathematically is complicated, we, therefore, retain the more common (and mild) finite MDP condition.\n\n\n>Question 1: Could the author elaborate on the statement in Lines 129-134 that 'SinkhornDRL inherently captures the spatial and geometric layout of return distributions'? Does this imply that the KL divergence term can leverage a richer representation of data geometry?\n\nBy the nature of optimal transport distances, Sinkhorn divergence inherently captures the spatial and geometric layout, which can be demonstrated by its definition. In particular, Sinkhorn divergence and other optimal transports are defined by measuring the cost of moving mass from one point to another in a space, where the cost directly depends on the spatial distances between points. This is because the optimal transport plan is not just about matching quantiles but also **considers where mass is located and where it needs to go**, which inherently reflects the spatial relationships and geometric layout of distributions. This is also the primary advantage of optimal transport distances over others, e.g., MMD. In addition, the KL divergence also involves the joint distribution (transport plan), which retains the properties of optimal transport distances.\n\n\n>Question 2: The definition of Sinkhorn divergence for multi-dimensional reward settings is unclear. Specifically, in Line 756, the cost function applies equal weights among reward sources. Can this be generalized to a weighted sum?\n\nYes, it is feasible to generalize the cost function to a weighted form, akin to an extension from Euclidean distance to Mahalanobis distance. By using the weighted cost function, the optimal transport distances can be **more closely aligned with specific applications.** The weighted cost function can potentially lead to benefits, provided that we have access to accurate or relevant weights, which typically require prior knowledge.\n\n>Limitation:  In line 22, the reference [11] is cited as a risk-sensitive control, but this is close to the exploration method for distributional RL.\n\nWe acknowledge this citation issue and will revise it accordingly. Thanks for bringing this to our attention.\n\nWe appreciate your insights and suggestions, which will guide our revisions to improve the clarity and accuracy of our manuscript. Please feel free to let us know if you have any further questions.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Net9w0wCOy", "reviewer": "Reviewer_uu9w", "summary": "This paper proposes a new RL algorithm that leverages Sinkhorn divergence, which they claimed as a regularized Wasserstein loss. Theoretically, they showed the contraction properties that align with the interpolation nature of Sinkhorn divergence between Wasserstein distance and MMD. Empirically, it outperforms or matches existing algorithms on many atari games.", "questions": "Line 319: the algorithm is much better on complicated games, which is very interesting. In other words, the algorithm is probably inferior on easy games. Do you have more explanation on this?", "limitations": "I didn't find any potential negative societal impact.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- 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.", "weaknesses": "After reading the paper I am still not sure about the benefits of using sinkhorn divergence, although I agree that the authors established some theoretical results (which I also have some concerns about below) and showed experimental results. My question is more intuitive---why should we pick Sinkhorn over others at a high level of intuition? I appreciate that the authors have spent many words explaining it from various perspectives. However, I didn't totally understand them and am not convinced at this moment. Some of my specific confusion are below:\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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for taking the time to review our paper. We appreciate your positive assessment and insightful feedback, and we would like to address the concerns you raised in your review.\n\n\n>My question is more intuitive---why should we pick Sinkhorn over others at a high level of intuition? \n\nWe summarize the explanations in the introduction to provide a high-level intuition. Exiting distributional RL algorithms that rely on Wasserstein distance often struggle with 1) inaccuracies in representing distribution using pre-specified statistics, i.e., quantiles, and 2) difficulties in extending to the multi-dimensional reward setting. **Sinkhorn divergence stands out for efficiently approximating multi-dimensional Wasserstein distance** (by introducing an entropic regularization), simultaneously overcoming the two limitations. Additionally, the smoother transport plan and induced smoothness by the entropic regularization can also benefit the optimization in the RL learning process, making it a preferred choice.\n\n\n>line 48-51: why does using samples instead of pre-specified statistics have better accuracy? Is there any reference?\n\nYes, using samples provides a more direct and accurate representation of return distributions. The MMD-DQN paper [1] discussed this extensively, and a more recent study [2] also adopts sample-based representation. These approaches, rooted in kernel methods, contrast with our method's foundation in optimal transport, offering a different approach to compare the distribution difference.\n\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\nAccording to the Sinkhorn divergence literature, the KL regularization in essence **amounts to vanilla entropy regularization** in terms of the joint distribution / optimal coupling. This is because:\n$$\\begin{aligned}\n\\text{KL}(\\Pi|\\mu \\otimes \\nu) = \\mathcal{H}(\\mu) + \\mathcal{H}(\\nu) - \\mathcal{H}(\\Pi).\n\\end{aligned} \n$$\nwhere $\\mathcal{H}$ is the entropy. Thus, the objective function is equivalent to\n$$\\begin{aligned}\n \\min _{\\Pi \\in \\mathbf{\\Pi}(\\mu, \\nu)} \\int c(x, y) \\mathrm{d} \\Pi(x, y) + \\varepsilon \\text{KL}(\\Pi|\\mu \\otimes \\nu) \\iff  \\min _{\\Pi \\in \\mathbf{\\Pi}(\\mu, \\nu)} \\int c(x, y) \\mathrm{d} \\Pi(x, y) - \\varepsilon \\mathcal{H}(\\Pi)\n\\end{aligned} \n$$\nwhere the $\\mathcal{H}(\\mu)$ and $\\mathcal{H}(\\nu)$ are constant for fixed marginal distributions, which are independent of the minimizer. In this equivalent objective function form, minimizing the Sinkhorn divergence also encourages to maximize the entropy $\\mathcal{H}(\\Pi)$, leading to a more uniformly distributed optimal transport plan / joint distribution. This aligns well with the maximum entropy principle.\n\n\n>line 53: \"smoother\" is very vague in this context. I would suggest using some different words.\n\nThanks for bringing this to our attention. The term \"smoother\" refers to the effect of the regularization in Sinkhorn divergence, which **encourages a more uniformly distributed transport plan**, as analyzed above. This more uniformly distributed transport plan is in contrast to the potentially sparse plan resulting from optimizing the unregularized Wasserstein distance. We will refine this terminology and provide a clearer explanation in the revised version.\n\n>As for the theoretical results, I am not sure how novel they are. ... 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) ...\n\nWe want to clarify that **Theorem 1(3) is precisely the contraction conclusion under the general $\\epsilon \\in (0, +\\infty)$**, which is the main theoretical contribution of this paper. We encourage a further review of Theorem 1(3) and the subsequent proof sketch on Page 5 of our paper.\n\n>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\nWe acknowledge the **asymmetry** of the ratio improvement/percentage increase metric, i.e., $(A-B)/B$, in the visual representation, but it is intuitively interpretable and thus more commonly used in practice.\n\nAnother metric is the log difference metric, i.e., $\\log(A) - \\log(B)$, which is symmetric as $\\log(A/B)=\\log(k)$ vs $\\log(A/B)=\\log(1/k)=-\\log(k)$. However, the log differences are undefined for zero or negative values, which is its main limitation. As **the raw score in each Atari game can be either positive or negative**, the ratio improvement/percentage increase metric is more broadly applicable and particularly preferable in evaluating the RL algorithm across a wide range of Atari games than other metrics like the log difference.\n\n>Line 319: the algorithm is much better on complicated games, which is very interesting. In other words, the algorithm is probably inferior on easy games. Do you have more explanation on this?\n\nThe finding that the algorithm is much better on complicated games does not imply that it is inferior on easy games. After checking the raw score table across 55 Atari games, we find our algorithm is also on par with or even performs better than other baselines on easy games, such as Atlantis, Enduro, FishingDerby, and MsPacman.\n\nWe appreciate your insights and suggestions, which will guide our revisions to improve the clarity and accuracy of our manuscript. Please feel free to let us know if you have any further questions.\n\n## Reference \n\n[1] Distributional Reinforcement Learning via Moment Matching (AAAI 2021)\n\n[2] Distributional Bellman Operators over Mean Embeddings (ICML 2024)", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
Dear Reviewers,

We would like to thank all the reviewers for their thoughtful and constructive feedback on our paper. We deeply appreciate your positive assessment and have thoroughly provided our response to address each of your concerns.  We are committed to enhancing the quality of our work and remain at your disposal for any further clarifications or questions you might have.

(Attached is the Summary Table result in response to Reviewer Fdqh)

Sincerely,  
Authors

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

The reviewers are uniformly positive about the paper. I have read the paper and agree that the idea is interesting and will be appreciated by the community.
Decision: Accept (poster)