Distributional Reinforcement Learning via Sinkhorn Iterations
Authors: Ke Sun
Yingnan Zhao
Yi Liu
Wulong Liu
Bei Jiang
Linglong Kong
Venue: Submitted to ICLR 2023 (2023)
Keywords: distributional reinforcement learning
sinkhorn divergence

Abstract: Distributional reinforcement learning~(RL) is a class of state-of-the-art algorithms that estimate the entire distribution of the total return rather than only its expectation.  The empirical success of distributional RL is determined by the representation of return distributions and the choice of distribution divergence. In this paper, we propose a new class of \textit{Sinkhorn distributional RL~(SinkhornDRL)} algorithm that learns a finite set of statistics, i.e., deterministic samples, from each return distribution and then uses Sinkhorn iterations to evaluate the Sinkhorn distance between the current and target Bellmen distributions. Sinkhorn divergence features as the interpolation between the Wasserstein distance and Maximum Mean Discrepancy~(MMD). SinkhornDRL finds a sweet spot by taking advantage of the geometry of optimal transport-based distance and the unbiased gradient estimate property of MMD. Finally, compared to state-of-the-art algorithms, SinkhornDRL's competitive performance is demonstrated on the suit of 55 Atari games.

Under review as a conference paper at ICLR 2023
DISTRIBUTIONAL REINFORCEMENT LEARNING VIA SINKHORN ITERATIONS
Anonymous authors Paper under double-blind review
ABSTRACT
Distributional reinforcement learning (RL) is a class of state-of-the-ar the- algorithms that estimate the entire distribution of the total return rather than only its expecta- tion. The empirical success of distributional RL is determined by the representa- tion of return distributions and the choice of distribution divergence. In this paper, we propose a new class of Sinkhorn distributional RL (SinkhornDRL) algorithm that learns finite set of statistics, i.e., deterministic samples, from each return distribution and then uses Sinkhorn iterations to evaluate the Sinkhorn distance between the current and target Bellmen distributions. Sinkhorn divergence fea- tures as the interpolation between the Wasserstein distance and Maximum Mean Discrepancy (MMD). SinkhorDDRI. finds sweet spot by taking advantage of the geometry of optimal transport-based distance and the unbiased gradient estimate property of MMD. Finally, compared to state-of-the-art algorithms, Sinkhorn- DRL's competitive performance is demonstrated on the suite of 55 Atari games.
INTRODUCTION
Classical reinforcement learning (RL) algorithms are normally based on the expectation of dis- counted cumulative rewards that an agent observes while interacting with the environment. Recently, a new class of RL algorithms called distributional RL estimates the full distribution of total returns and has exhibited the state-of-the of-the-art performance in a wide range of environments (Bellemare al., 2017a Dabney et al. 201 8b Yang et al.) 2019 Zhou et al., 2020 Nguyen et al. 2020).
In distributional RL literature, it is easily recognized that algorithms based on either Wasserstein distance or MMD have gained great attention due to their superior performance. Their mutual con- nection from the perspective of mathematical properties intrigues us to explore further in order to design new algorithms. Particularly, Wasserstein distance, long known to be a powerful tool to com- pare probability distributions with non-overlapping supports, has recently emerged as an appealing contender in various machine learning applications. It is known that Wasserstein distance disregarded because of its computational burden in its original form to solve an expensive network flow problem. However, recent works (Sinkhorn) 1967 Genevay al., 2018) have shown that this cost can be largely mitigated by settling for cheaper approximations through strongly convex regular- izers. The benefit of this regularization has opened the path wider applications of the Wasserstein distance in relevant learning problems, including the design of distributional RL algorithms.
long
The Sinkhorn divergence (Sinkhorn) 1967) introduces the entropic regularization on the Wassertein distance, allowing it tractable for the evaluation especially in high-dimensions. It has been suc- cessfully applied in numerous crucial machine learning developments, including the Sinkhorn- GAN (Genevay et al. 2018) and Sinkhorn- based adversarial training (Wong et al.) (2019). More importantly, has been shown that Sinkhorn divergence interpolates Wasserstein ditance and MMD, and their equivalence form can be well established in the limit cases (Feydy et al. [2019] Ramdas cet al., 2017 Nguyen et al. 2020). However, a Sinkhorn-based distributional RL algorithm has not yet been formally proposed and its connection with algorithms based on Wasserstein distance and MMD is also less studied. Therefore, a natural question is can we design a new class of distributional RL algorithms via Sinkhorn divergence, thus bridging the gap between existing two main branches of distributional RL algorithms? Moreover, the dominant quantile regression- based algorithms, e.g. QR-DQN (Danney et al. 2018b), aimed at approximating Wasserstein distance, suffers from

---

Under review as a conference paper at ICLR 2023
the non-crossing issue in the quantile estimation (Zhou et al. 2020), while sample-based Sinkhorn distributional algorithm can naturally circumvent this problem.
In this paper, propose novel distributional RL family based on Sinkhorn divergence. Firstly, we show key roles of distribution divergence and value distribution representation in the design of distributional RL algorithms. After detailed introduction of our proposed SinkhornDRL. algorithm, with a non-trivial proof, we theoretically analyze the convergence property of distributional Bellman operators under Sinkhorn divergence. A regularized MMD equivalence with Sinkhorn divergence is also established, interpreting its empirical success in real applications. Finally, we compare the performance of our SinkhorDDRL algorithm with typical baselines on 55 Atari games, verifying the competitive performance of our proposal. Our method inspires researchers to find trade-off that simultaneously leverages the geometry of the Wasserstein distance and the favorable unbiased gradient estimate property of MMD while designing new distributional RI algorithms in the future.
2
PRELIMINARY KNOWLEDGE
2.1
DISTRIBUTIONAL REINFORCEMENT LEARNING
In classical RL, an agent interacts with an environment via a Markov decision process (MDP), a 5-tuple (S, R, P))), where S and A are the state and action spaces, respectively. P is the environment transition dynamics, R the reward function and 7 E (0, 1) the discount factor.
From Value function to Value distribution. Given a policy the discounted sum of future rewards random variable 2" (s,
Σ 7 R(s,, at), where So 8. au a, 81+1 P(-|s,, at), and x (-18)). In the control setting, expectation-based RL is based on the action-value func- tion ("(ssa), a), which is the expectation of (*(s,, a), i.e., (* a) = EE (2* a)}. By contrast, distributional RI. focuses on the action-value distribution, the full distribution of 7* (s, The incorporation of additional distributional knowledge intuitively interprets its empirical success.
Distributional Bellman operators. For the policy evaluation in expectation- RL, the action- value function is updated via Bellman operator 7 Q=Q a) = E[ a)] + +E,,'p.x (s', .a)]. In distributional RL, the distribution of (s, a) is updated via the distributional Beliman operator T
TT Z(, a) _ R(s, a) + + (s', (s',a'),
(1)
where ~ P(-|s, a) and d' π (-1s'). The equality implies random variables of both sides are equal in distribution. The distributional Bellman operator IT is contractive under certain distribution divergence metrics. We provide a detailed discussion about more related works in Appendix
2.2 DIVERGENCES BETWEEN MEASURES
Optimal Transport (OT) and Wasserstein Distance. The optimal transport (OT) metric between two probability measures (1, v) is defined as the solution of the linear program:
/ c(xx))dd((z,)))
min IEI().))
where is the cost function and II is the joint distribution with marginals (H. v). Wasserstein distance (a.k.a. earth mover distance) is special case of optimal transport with the Euclidean norm as the cost function. In particular, given two scalar random variables and Y, p- Wasserstein metric W,, between the distributions of X and Y can be simplified as
W W,(X = (()))
1/P
(3)
where F=1 is the inverse cumulative distribution function of a random variable. The desirable geometric property of Wasserstein distance allows it to recover full support measures, but suffers from the curse of dimension (Genevay al. 2019 Arjovsky et al. 2017).
Maximum Mean Diserepancy. The squared Maximum Mean Discrepancy (MMD) MMD? with the kernel & is formulated as
MMD
E [k (X, X')) E(k E(K -2 (X),
2

---

Under review as a conference paper at ICLR 2023
where k(:,.) .) is continuous kernel X. X' (resp. Y' is random variable independent of X (resp. Y). If & is a trivial kernel, MMD degenerates to the energy distance. Mathematically, the "flat" geometry that MMD induces the space of probability measures does not faithfully lift the ground distance (Feydy et al. 2019), but MMD is cheaper to compute than OT and has a smaller sample complexity, i.e., approximating the distance with samples of measures (Genevay et 2019). We provide the detailed introduction of more distribution divergences in Appendix
3 ROLES OF DISTRIBUTION DIVERGENCE AND REPRESENTATION
3.1 DISTRIBUTIONAL RL: FROM NEURAL FITTED Q TO Z ITERATION
Neural Fitted Q-Iteration. is known that Deep Networks (Mnih al. 2015 can be simplified into Neural Fitted Q-Iteration (Fan et al. 2020) under tricks of experience replay and the target network, where we update Qo a) parameterized by 0 in each iteration
Q = argmin Qa
Σ [v: - Q (st, a;)]² a=1
(5)
72
where the target yi = r(8;; 0z) + m Q/- (s/, a) is fixed within every Turget steps to up- date target network Qa- with parameters 0r by letring 6" 0 and the experience buffer induces independent samples {(s,, a,, Tı, s()) (E(m) In an ideal case that neglects the non-convexity and TD
approximation errors, we have Q = TQ2. which is exactly equivalent to the updating rule under Bellman optimality operator.
Neural Fitted Z-Iteration. Analogous to Neural Fitted Q-iteration, we can also simplify value- based distributional RL methods based on a parameterized Ze into Neural Fitted Z- Z-Iteration
1 n
Z zh+1 argmin
Σ Za n
dp(Y,, 2z (s),
(6)
where the target Y; R(si, a.) + 7 7%z. (0/s (+()) with xz(8') argmax," EE [ (8', d')]
fixed within every Tragea steps to update target network Zgo, and d, is the distribution divergence.
3.2 KEY ROLES OF dp AND Zo
Within the Neural Fitted Z-Iteration framework proposed in Eq. 6 we observe that the choice of representation manner Zo and the metric d, are pivotal for the distributional RL algorithms. For instance, QR-DQN (Dabney et al. 20185) approximates Wasserstein distance W which leverages quantiles represent the distribution of Zg. CS1 (Bellemare et al. 2017a) represents Ze via categorical distribution under the convergence of Cramér distance (Bellemare et al. 2017b] Rowland al. 2018), while MMD distributional RL (MMDDRL) (Nguyen et al.) (020) learns samples to represent the distribution of Za based on MMD. We compare characteristics of these distribution divergence, including the convergence rate and sample complexity, in Table [ Theoretical results regarding Sinkhorn divergence is based on (Genevay et 2019) and the detailed convergence proof of other distances is also provided in Appendix B In summary, we argue that d, and Zo are two crucial factors in distributional RL design, based on which we introduce Sinkhorn distributional RL.
Algorithm
d, Distribution Divergence Cramér distance Wassenstein distance
Representation Zo Convergence Rate "
Sample Complexity of d,
CS1 QR-DQN MMDDRI.
Histogram
Quantiles
O(n O(1/n)
MMD
Samples
20/2 with kernel kg
SinkhornDRI.
(E 1+ 0)
O(natiivv (s O(n (e o0)
Sinkhorn divergence
Samples
(ours)
y (c o0) oo)
Table 1: Comparison between typical distributional RI. algorithms under different distribution di- || y/ in MMDDRL. d is the sample dimension and k 23d + ||c]]oc where the cost function e is 3-Lipschitz (Genevay et al. 2019). Sample complex- ity of MMD can be improved to O(1/n) using kernel herding technique (Chen et al. 2012).
vergences and represention of Za. Ka -
3

---

Under review as a conference paper at ICLR 2023
4 SINKHORN DISTRIBUTIONAL RL (SINKHORNDRL)
In this section, we firstly introduce Sinkhorn divergence and apply it in distributional RL. Next, conduct theoretical analysis about the convergence and a new moment matching manner of our algorithm under the Sinkhorn divergence. Finally, a practical Sinkhorn iteration algorithm is introduced to evaluate the Sinkhorn divergence.
4.1 SINKHORRN DIVERGENCE AND GENETIC ALGORITIMM
We design Sinkhorn distributional RL algorithm via Sinkhorn divergence. Sinkhorn diver- gence (Sinkhorn 1967) is a tractable loss to approximate the optimal transport problem by leverag- ing an entropic regularization to turn the original Wasserstein distance into a differentiable and more robust quantity. The resulting loss can be computed using Sinkhorn fixed point iterations, which is naturally suitable for modern deep learning frameworks. In particular, the entropic smoothing gen- erates family of losses interpolating between MMD. As such, it allows us to find sweet trade-off' that simultancously leverages the geometry of Wasserstein distance on the one hand, and the favor- able nigh-dimensional sample complexity and unbiased gradient estimates of MMD. We introduce the entropic regularized Wassertein distance We,((, (1,, 2) as
/ c(x, ((x,))dl(zz) v) KLL(II
min nEN(
diI(,, y) is a strongly convex regularization. The impact of this entropy regularization is similar l ridge regularization in linear regression. Next, the sinkhorn loss (Feydy et al.) 2019; Genevay et al. 2018) between two measures μ and v is defined as
where (L( = log
I(a.y)
ap(r)de(())
W (e,((uvv) (,,,) We,(m, We,evv
(8)
As demonstrated by (Feydy et al. 2019). the Sinkhorn divergence We.eu,, is convex, smooth and positive definite that metrizes the convergence in law. In statistical physics, Wa,e (p,, v) can be re-factored as projection problem:
We,c(1, min KL (T)KK), NEN(P))
where IC is the Gibbs distribution and its density function satisfies dK. c(x.p) dp(x) (u(x)d((( dv
This problem is often referred to as the "static Schrödinger problem" (Léonard, 2013 Rüschendori
Thomsen 1998) as it was initially considered in statistical physics.
Distributional RL with Sinkhorn Divergence and Particle Representation. The key of applying Sinkhorn divergence in distributional RL is to simply leverage the Sinkhorn loss Wo,s to mea- sure the distance between the current action-value distribution Ze(s, a) and the target distribution I Za(s, a), yielding We.r(Zo(6 a), ITh Zo a)) for each 8, pairs. In terms of the representation for (s a). we employ the unrestricted statistics, i.c., deterministic samples, due to its superiority in MMDDRL (Nguyen et al. 2020). instead of using predefined statistic functionals, e.g., quantiles in QR-DQN (Dabney et al."
2018b) or categorical distribution in C51 (Bellemare et al. 2017a
Algorithm 1 Generic Sinkhorn distributional RI. Update
Require: Number of generated samples N, the cost function and hyperparameter e. Input: Sample transition (s, r', r')
1: if Policy evaluation then
2: a" *--ss'")
3: else
4: a" <- arg max.EA N ΣN Zo (s','')
5: end if
6: SZ; &- r 770 ",,a");, siS N
Output: Wo. (IZo(s,a))) - (Z0(s,))))))
4

---

Under review as a conference paper at ICLR 2023
More concretely,
neural networks to generate samples that approximate the value distribu- s
tion. This can be expressed as Zø(s, a) == [Zo a),) where is the number generated samples. We refer to the samples { (7)(s,))))2
as particles. Then we leverage the Dirac mix- ture N Ei=1 (z.(s,a), to approximate the true density function of (= a). thus minimizing the Sinkhorn divergence between the approximate distribution and its distributional Bellman target. A detailed and generic distributional RL algorithm with Sinkhorn divergence and particle representa- tion is provided Algorithm |
Remark. From the general algorithm framework in Algorithm our SinkhomDRRL, generally mod- ifies the distribution divergence comparing the state-of-the-ar MMDDRL (Nguyen et al.] 2020). but SinkhornDRL fundamentally falls into Wasserstein distance-based distributional RL family as discussed in Appendix [A] As such, QR-DQN and MMDDRL are direct counterparts for Sinkhorn- DRL, and the follow-up up works IQN (Dabney al.) 2018a) and FQF (Yang et al. 2019) can naturally extend both MMDDRL and SinkhomDRL as discussed in (Nguyen et al. 2020
4.2 THEORETICAL ANALYSIS UNDER SINKHORN DIVERGENCE
Convergence. Firstly, we denote the supreme form of Sinkhorn divergence as W (,,
V), sup (e) (x,a)ESXA
(10)
We will use Wae (p, v) to establish the convergence of T" in Theorem
Theorem 1. If we leverage Sinkhorn loss We,(H, v) in Eq. 8 as the distribution divergence in distributional RL, and choose the unrectified kernel Ko == - a vll as -c > 0), it holds that
(1) + 0) We,(H, v) + 2Wa(u, v). When E = 0, I a y-contraction under W c,e"
(z + +00) We, v) + MMD (H,V). When +oo, Tnth is 72/2 -contraccte under W e.s-
(E) E (0, )oo)), I" a contractive operator under W The related non-constant contraction factor a(7,a) < also depends on the distribution sequence in distributional Bellman iterations.
We provide the long yet rigorous proof of Theorem in Appendix C Theorem [ (1) and (2) are follow-up conclusions in terms of the convergence behavior of I" based on the interpolation rela- tionship between Sinkhorn divergence with Wasserstein distance and MMD (Genevay et al.) 2018). Our key theoretical contribution is for the general E (0, oo), in which we conclude that X a contractive operator. The crux of the proof is two-fold. Firstly, we show the a variant of scale sensitive property of Sinkhorn divergence when factor is also determined by the specified two probability measures. Next, we propose a new distri- bution Contraction mapping theorem in Theorem 2 of Appendix C based on which eventually arrive at the convergence of distributional Bellman operator under W Intriguingly yet reason- ably, the contraction factor (17, ox) is non-constant but function less than 1 that also depends on the distribution sequence while iteratively applying distribution Bellman updates. Our non-trivial proof about Sinkhorn divergence can even contribute the optimal transport literature.
Ka. where the resulting non-constant scaling
Consistency with Related Conclusions. As Sinkhorn divergence interpolates between Wasserstein distance and MMD, its contraction property when the cost function holds -ka for the general E E [0, oc] is intuitive. Note that if we choose Gaussian kernels as the cost function, there will be no concise and consistent contraction results as Theorem
(3). This conclusion is also consistent with MMDDRL. (Nguyen et al. 2020), where Tx is generally not contraction operator under MMD equipped with Gaussian kernels as counterexample has been pointed out in MMDDRI. (when +oo). To be consistent with the contraction property analyzed in theory (Theorem I (3)), we employ the rectified kernel li as the cost function in our experiments and set a = 2, under which SinkhomDRL suggests favorable performance in Section
Regularized Moment Matching under Sinkhorn Divergence Associated with Gaussian Ker- nels. We further examine the potential connection between SinkhomDRL with existing distribu- tional RL families. Inspired by the similar manner in MMDDRL (Nguyen et al., 2020). we find that the Sinkhorn divergence with the Gaussian kernel can also promote match all moments be-
5

---

Under review as a conference paper at ICLR 2023
tween two distributions. More specifically, the Sinkhorn divergence can be rewritten regularized moment matching form in Proposition !
Proposition 1. For E (0, +00). Sinkhorn divergence We,e (1, v) associated with Gaussian kernels k(z,)) exp( (pp--xx- y)² (0/0(222)) as -c, equivalent to
(,, v) = Σ
2mm!! (Mi())) (Mau)) ) + cE log XX)XXXXII
(TT
(11)
n=0
where 11 denotes the optimal II determined by by evaluating the Sinkhorn divergence via minnen(pw) Wae(H, Ma(u) = Eapp x===//1 1200) and similarly for Mn(v).
We provide the proof of Proposition 1 Appendix D Similar to MMDDRL associated with
Gaussian kernel (Nguyen et al. 2020). Sinkhorn divergence approximately performs regularized
moment matching scaled by 7(200)
Equivalence to Regularized MMD Distributional RL. Based on Proposition I we can immedi- ately establish the connection between Sinkhorn divergence and MMD in Corollary [ indicating that minimizing Sinkhorn divergence between two distributions is equivalent to minimizing a regu- larized squared MMD.
Corollary 1. For E (0, +o0) and denote 11; the optimal 1I by evaluating the Sinkhorn diver- gence, holds that
(C(X,Y) We = MMD2 ( v) + =E log XX X')
(12)
where We,s replace Wa,s (H, v) for short.
Proof of Corollary is provided in Appendix D It is worthy of noting that this equivalence is established for the general case when E + 0 or +o0. For example, when E + +o0, the second part including & in Eq. 12 is not expected dominate. This is owing to the fact that the regularization term would be 0 as IT. μ & when + +00. In summary, even though the Sinkhorn divergence was initially proposed to serve as an entropy regularized Wasserterin distance when the cost function = Ka. it turns out that is equivalent to regularized MMD if associated with Gaussian kernels, as revealed in Corollary
(00), and it does not hold in the limit cases when
4.3 DISTRIBUTIONAL RL VIA SINKHORN ITERATIONS
The theoretical analysis in Section 4.2 sheds light on the behavior of distributional RL with Sinkhorn divergence, but another crucial issue need to address is how evaluate the Sinkhorn loss effec- tively. Due to the advantages of Sinkhorn divergence that both enjoys geometry property of optimal transport and the computational effectiveness of MMD, we can utilize Sinkhorn's 's algorithm, i.e., Sinkhorn Iterations (Sinkhorn 1967: Genevay et al. 2018), 3), to evaluate the Sinkhorn loss. Notably, Sinkhorn iteration with L steps yields a differentiable and solvable efficiently loss function as the main burden involved in it is the matrix-vector multiplication, which streams well on the GPU with simply adding extra differentiable layers on the typical deep neural network, such as DQN archi- tecture.
Specifically, given two sample sequences ( , in the distributional RL algorithm, the optimal transport distance equivalent to the form:
min NER
{ (); PIN = IN. 1 = 1n},
(13)
where the empirical cost function c.,5 = c(Zi, IZ,). By adding entropic regularization on opti- mal transport distance, Sinkhom divergence can be viewed restrict the search space of P in the following scaling form:
P., bbj.
(14)
where K.,j = e As/l is the Gibbs kernel defined in Eq. This allows us to leverage iterations regarding the vectors and b. More specifically, we initialize bn 1v, and then the Sinkhorn

---

Under review as a conference paper at ICLR 2023
Algorithm Sinkhorn Iterations to Approximate Was
(Z,);=1
(
Input: Two samples sequences {Z, N
i=1 (
1 1 j=1- number of Sinkhorn iterations L and hyperpa-
rameter E.
I: Cijj cZi, x2j) for Vi 1, ,,.., N, = 1,,,,,
2: Ki,; =exp( (exp(-cij/EE)
3: bn 1-1N
4: Iorl=1,2 1,,, ..., L do
5: 01 f Kbi-1 bi t-
1v
Kai
6: end for
7: Was
(Z:))=1: (57)) ((K O)b,
Return: Wo. (121) (52
iterations are expressed as
1x KCb
1N
01+1
and
bi+
(15)
IC a1411
where indicates an entry-wise division. It has been proven that Sinkhorn iteration asymptotically converges the true loss in linear rate (Genevay al.) 20188 Franklin & Lorenz 1989 Cuturi 2013 Jason Altschuler] (017). We provide detailed a algorithm description of Sinkhorn iterations Algorithm 2] With the efficient and differential Sinkhorn iterations, we can easily evaluate the Sinkhorn divergence and thus let our algorithm enjoy its theoretical advantages. In practice, we need choose and E, and we conduct a rigorous sensitivity analysis sis in Section
EXPERIMENTS
We demonstrate the effectiveness of SinkhomDRL described in Algorithm T the full 55 Atari 2600 games. Specifically, we leverage the same architecture as QR-DQN (Dabney et al. (018b) 8b), and replace the quantiles output with N particles, i.e., samples. In contrast MMDDRL, Sinkhorn- DRL only changes the distribution divergence from MMD to Sinkhorn divergence, and therefore the potential superiority the performance can be attributed to the advantages of Sinkhorn divergence.
Baselines. Due to the interpolation feature of Sinkhorn divergence between Wassertein distance and MMDDRL, we choose three typical distributional RI. algorithms as classic baselines, including QR-DQN (Dabney et al. 2018b) that approximates the Wasserstein distance, C51 (Bellemare al., 2017a) and MMDDRL (Nguyen et al. 2020), as well as DQN (Mnih et al.] 2015). MMDDRL algorithm is implemented with the same architecture as QRDQN, and leverages Gaussian kernlss y)² /h) with the kernel mixture trick covering a range of bandwidths h, which same the basic setting in the original MMDDQN paper (Nguyen al. 2020). We deploy all algorithms on 55 Atari 2600 games, and reported results are averaged over 3 seeds with the shade indicating the standard deviation. We runs 10M time steps (40M frames) for the computation cost reason, but we report learning curves across all games to make results convincing enough.
ky y) = cxp(-
Hyperparameter settings. For fair comparison with QR-DQN, C51 and MMDDRL, used the same hyperparamters: the number of generated samples N = 200, Adam optimizer with 0.00005, Adam 0.01/32. We used target network to compute the distributional Bellman target, which fits well in the Neural Fitted Z-Iteration framework. In addition, we choose number of Sinkhorn iterations L 10 and smoothing hyperparameter sensitive within proper interval as demonstrated in Section 5.2 We choose the unrectified kernel
10.0 in Section [5.] as they are not
as the cost function, i.e, -c kin and select a = 2 k in our SinkhornDRI. algorithm.
PERFORMANCE OF SINKHORNDRL
Figure I illustrates that SinkhorDDRL achieve the competitive performance across 55 Atari games compared with various baseline algorithms with different metrics d, and representation man- ners Zg. On large number of games, e.g., Tennis, Seaquest and Atlantis, SinkhornDRI. can
7

---

Under review as a conference paper at ICLR 2023
Tennis
Se Seaquest
Atlantis
DDN
OH1 VMA
14000
CC1
Average Relum
HMF
Awerage Retum
8000
0.9
022
LA4
8.6
9.
1.3
60
Time Sleps (1e7)
Time Steps (1e7)
Time Staps (1w7)
Spacelnvaders
Breakout
Pong
2000
600
DON Mill
Average Return
CSt
Average Retum
1758
1508
NMD
Seansm
100
Average
1000
200
750
100
10
500
-20
Time Steps (1=7)
Steps (1c7)
Time Sleps (1e7)
Phoenix
This
WizardOWorr
2000
Average Retum
CS1
ORDON
MMD
Sindhom
4000
2000
200
0.6 Time Steps (1e7)
04
Time Steps (1e7)
Time Staps (1677
Figure 1:
Learning curves of SinkhomDRL algorithm compared with DQN, C51, QR-DQN and MMD, on nine typical Atari games over 3 seeds.
significantly outperform other baselines, especially on Tennis where other algorithms even fail to converge. The improvement of SinkhornDRL over MMDDRL empirically verifies the regulariza- tion advantage of the Sinkhorn as analyzed in Corollary On some games, e.g., Breakout, Pong and SpaceInvaders, SinkhorDRLL is par with MMDDRL and other baselines, while on the last row in Figure [] SinkhoraDRL is slightly inferior to the state-of-the-ar algorithm. We provide learn- ing curves of all typical distributional RL algorithms on all 55 Atari games in Appendix F where SinkhornDRL still achieves the competitive performance in general.
Sinkhorn / QRDQN (%)
Sinkhom MMD (%)
Ratio Improvment (%
Ratio Improvment (%)
608
408
100
204
-108
-20-
200
(a) Sinkhorn vs QRDQN
(b) Sinkhorn vs MMDDRL
Figure 2: Ratio improvement of return for Sinkhom distributional RL algorithm over QRDQN (left) and MMDDRL (right) over seeds. For example, the ratio improvement is calculated by (Sinkhorn -QRDON) ORDON in the left.
8

---

Under review as a conference paper at ICLR 2023
We conduct ratio improvement comparison across 55 Atari games between SinkhomDRR with QRDQN and MMDDRL, respectively. Figure 2 showcases that by comparing with QRDQN (left), SinkhornDRL achieves better performance across almost half of considered games and the superior- ity of SinkhornDRI. is significant across a large amount of games, including Venture, Seaquest, Ten- nis and Phoenix. This empirical outperformance verifies the effectiveness of smoothing Wassertein distance in distributional RL. In contrast with MMDDRL, the advantage of SinkhornDRL. is reduced with the performance improvement on a smaller proportion of games, but remarkable performance improvement for SinkhornDRL. on a large amount of games can be easily observed. We also report mean and median of best soman-normalized scores Table Of Appendix where SinkhomDRL achieves almost state-of-the thrt art performance as MMDDRL on average.
Therefore, we conclude that SinkhornDRI competitive with the state-of-the-ar distributional RI. algorithms, e.g., MMDDRL, and can be extremely superior over existing algorithms on a large proportion of games. This empirical success can be owing to theoretical advantage of Sinkhorn divergence that simultancously makes full use of the data geometry from Wasserstein distance and the unbiased gradient estimate property from MMD, which coincides with results in Theorem
Seaquest
Seaquest
Breekout
2006
2000
Samplec-2
L-10
Avve 100e Relumn
Samples45
Averigh Retum
ORDON
$amples-2200
Samples-20
Ave. Mworspo Rocun
MND
Samples-500
2006
0.
62
0.8
0.5
Time Staps (1e7)
Time Stops (1c7)
Titss Stoos (to7s
(a) Hyper-parameter
(b) Number of Samples
(c) Sinkhorn Iterations L
Figure 3: Sensitivity analysis of SinkhornDRL on Breakout regarding E, number of samples, and number of iteration L. Learning curves are reported over 3 seeds.
5.2 SENSITIVITY ANALYSIS AND COMPUTATIONAL COST
The limit behavior connection in [ (1) and (2) between SinkhornDRL with QR-DQN and MMDDRL may be rigorously verified in numerical experiments as an overly large small will lead to numerical instability of Sinkhorn iterations in Algorithm 2 worsening its performance, as shown Figure (a). In practice, we choose a proper ε 10 across all games. SinkhomDRL also requires proper number of iterations L and samples N. For example, a small N, c.g., N = 2 in Seaquest in Figure 3 (b) leads to the divergence of algorithms, while an overly large N can degrade the performance and meanwhile increases the computational burden (Appendix GG). We conjecture that using larger networks to represent more samples is more likely to suffer from the overfitting issue, yielding the instability in the RL training (Bjorck et al. 2021). Therefore, choose 200 to attain an appealing performance with the computational effectiveness. For the computation cost (Appendix G). SinkhornDRL increases around 50% computation cost compared with QR-DQN and C51, but only slightly increases the overhead (by around 20%) in contrast to MMDDRL. Please refer to Appendix G for more detailed results and discussion.
DISCUSSIONS AND CONCLUSION
To extend our algorithm for better performance, implicit generative models, including parameteriz- ing the cost function in Sinkhorn loss, can be further incorporated. We leave it as the future work. Moreover, other divergences, e.g., those that also smooth Wassertein distance, can also be ap- plied into the design of distributional RL algorithms the future.
In this paper, novel family of distributional RL algorithms based on Sinkhorn divergence is pro- posed that accomplishes a competitive performance compared with the-state-of-ther distributional RL algorithms on 55 Atari games. Theoretical analysis about the convergence and moment match- ing behavior is provided along with a rigorous empirical verification. Sinkhorn distributional RL will lead to an important contribution among the research community.
9

---

Under review as a conference paper at ICLR 2023
Ethics Statement. Our study is about the design of distributional RL algorithms, which is not involved with any ethics issue.
Reproducibility Statement. Our results is based on the public implementation released in (Zhang 2018) with necessary implementation details given in Appendix We also provide the detailed proof from Appendix to Appendix
REFERENCES
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214 223. PMLR, 2017.
Mare G Bellemare, Will Dabney, and Rémi Munos. A distributional perspective reinforcement learning. International Conference on Machine Learning (ICML), 2017a.
Mare G Bellemare, Ivo Danihelka, Will Dabney, Shakir Mohamed, Balaji Lakshminarayanan, Stephan Hoyer, and Rémi Munos. The cramer distance as solution to biased wasserstein gradi- ents. arXiv preprint arXiv: 10743, 2017b.
Nils Bjorck, Carla P. Gomes, and Kilian Q. Weinberger. Towards deeper deep reinforcement learning with spectral normalization. Advances in neural information processing systems, 34, 2021.
Pablo Samuel Castro, Subhodeep Moitra, Carles Gelada, Saurabh Kumar, and Marc G Bellemare. Dopamine: A research framework for deep reinforcement learning. CoRR abs/1 222066 10, 2018.
Yutian Chen, Max Welling, and Alex Smola. Super-samples from kernel herding. UAI, 109-116.
AUAI Press, 2012.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. Advances in neural information processing systems, 26, 2013.
Will Dabney, Georg Ostrovski. David Silver, and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. International Conference on Machine Learning (ICML), 2018a.
Will Dabney, Mark Rowland, Marc G Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. Association for the Advancement of Artificial Intelligence (AAAI), 2018b.
Odin Elie and Charpentier Arthur. Dynamic Programming in Distributional Reinforcement Learn- ing. PhD thesis, Université du Québec à Montréal, 2020.
Jianqing Fan, Zhaoran Wang, Yuchen Xie, and Zhuoran Yang. A theoretical analysis of deep q- learning. In Learning for Dynamics and Control, pp. 486-489, PMLR, 2020.
Jean Feydy, Thibault Séjourné, François-Xavies Vialard, Shun-ichi Amari, Alain Trouvé, and Gabriel Peyré. Interpolating between optimal transport and mmd using sinkhorn divergences. In The 22nd International Conference Artificial Intelligence and Statistics, pp. 2681-2690. PMLR, 2019.
Joel Franklin and Jens Lorenz. On the scaling of multidimensional matrices. Linear Algebra and its applications, 11::717-735, 1989.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with sinkhorn di- vergences. In International Conference on Artificial Intelligence and Statistics, pp. 1608-1617.
PMLR, 2018.
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, pp. 1574-1583, PMLR, 2019.
Philippe Rigollet Jason Altschuler, Jonathan Weed. Near-lincar time approximation algorithms for optimal transport via sinkhom iteration, 2017.
10

---

Under review as a conference paper at ICLR 2023
Christian Léonard. A survey of the schr\" odinger problem and some of its connections with optimal transport. arXiv preprint arXiv: 2013.
Yudong Luo, Guiliang Liu, Haonan Duan, Oliver Schulte, and Pascal Poupart. Distributional rein- forcement learning with monotonie splines. In International Conference on Learning Represen- tations, 2021.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Mare G Belle- mare, Alex Graves, Martin Riedmiller, Andreas Fidjeland, Georg Ostrovski, et al. Human-level control through deep reinforcement learning. nature, 51 3(77540)552--333 2015.
Thanh Tang Nguyen, Sunil Gupta, and Svetha Venkatesh. Distributional reinforcement learning with maximum mean discrepancy. Association for the Advancement of Artificial Intelligence (AAAI), 2020.
Bauyen-Tang Distributional reinforcement learning via moment matching. https://githun
com/ Ehanhnguyent ang/ angg/mdrrr mmdrl 2021.
Aaditya Ramdas, Nicolás García Trillos, and Marco Cuturi. On wasserstein two-sample testing and related families of nonparametric tests. Entroy, 19(2):47, 2017.
Mark Rowland, Mare Bellemare, Will Dabney, Rémi Munos, and Yee Whye Teh. An analysis of categorical distributional reinforcement learning. In International Conference on Artificial Intelligence and Statisticx, pp. 29-37. PMLR, 2018.
Ludger Rüschendorf and Wolfgang Thomsen. Closedness of sum spaces andthe generalized schrödinger problem. Theory of Probability & Its Applications, 42(3))483 494, 1998.
Richard Sinkhorn. Diagonal equivalence to matrices with prescribed row and column sums. The American Mathematical Monthly, 7404):402-405, 1967.
Gábor J Székely. E-statistics: The energy of statistical samples. Bowling Green State University,
Department of Mathematics and Statistics Technical Report, 3(05):1-18, 2003.
Eric Wong, Frank Schmidt, and Zico Kolter. Wasserstein adversarial examples via projected sinkhorn iterations. In International Conference on Machine Learning, pp. 6808-6817. PMLR, 2019.
Derek Yang, Li Zhao, Zichuan Lin, Tao Qin, Jiang Bian, and Tie- Yan Liu. Fully parameterized quan- tile function for distributional reinforcement learning. Advances neural information processing systems,
2019.
Shangtong Zhang. Modularized implementation of deep rl algorithms in pytorch. https:// 2018.
github.ccom com hwwam////hanggtoggggggggg/eeeeee
Fan Zhou, Jianing Wang, and Xingdong Feng. Non-crossing quantile regression for distributional reinforcement learning. Advances Neural Information Processing Systems, 33, 2020.
Florian Ziel. The energy distance for ensemble and scenario reduction.
arXiv preprint
147700 2020.

---

Under review as a conference paper at ICLR 2023
A RELATED WORK
Based on the choice of distribution divergence and the distribution representation manner of Ze. distributional RL algorithms can be mainly categorized into three classes, including categorical, Wasserstein Distance and MMD distributional RL. Finally, we discuss their relationships with our proposed SinkhomDRL.
Categorical Distributional RL. As the first successful distributional RL family, categorical distri- butional RL (Bellemare et al. 2017a) represents the value distribution rj by the categorical distribu- tion f Σ p;o=,, where } (2) < ... < 2N) is the fixed supports within the pre-specified interval [7, md) and p; is the approximated categorical probability in each bin, respectively. Within this algorithm family, C51 (Bellemare et al. 2017a) leverages a neural network to approximate the categorical probabilities p; and apply projected KL divergence between the target and current cat- egorical value distributions. C51 has also been shown the theoretical contraction under the Cramér distance (Bellemare et al. 2017b Rowland et al. 2018), and empirically performs favorably in the suite of Atari games.
Wasserstein Distance Distributional RL. As directly solving wasserstein distance in Eq. 16 is tricky, QR-DQN (Dabney et al. 2018b) firstly proposed to use quantile regression to approximate Wasserstein distance Wp. QR-DQN leverages quantiles to represent the distribution of Ze. i.c.,
Ei= ò2,, where { the learnable support atoms as the quantile values of fixed quantile 2N fi=1. Implicit Quantile Networks (IQN) utilizes an implicit model output quantile values f more expressively, instead of the given ones in QRDQN. IQN also incorporates the risk measure in the framework of distributional RL. A follow-up work Fully parameterized Quantile Function (FQF) (Yang et al. (019)) improves IQN by proposing a more expressive quantile network, achieving better performance on Atari games. Non-crossing issue in quantileregrression has been raised and addressed properly in (Zhou et al., al.] (2020) that further improves QR-DQN. The monotonic rational-quadran splines also used to learn smooth continuous quantile functions (Luo al. 2021
MMD Distributional RL. MMD distributional RL (MMDDRL) (Nguyen et al. 2020 learns samples represent the value distribution of Ze based on maximum mean discrepancy (MMD) in Eq. 20 achieving the state-of- -oe--ofthherrr performance on Atari games.
Discussion about SinkhornDRL. As a complementary Wasserstein distance-based distributional RL, our SinkhoraDRL allows to solve Wasserstein Distance by incorporating the entropic regular- ization and circumstance the non-crossing issue in quantile regression intrinsically. Moreover, the cost function in SinkhornDRL can be further parameterized similar to IQN and FQF, which can intuitively achieve better performance. We leave the investigation towards this direction future works. Meanwhile, SinkhomDRL is also closely linked with MMDDRL as Sinkhorn divergence interpolates between Wasserstein distance and MMD, and also learns the unrestricted statisties, i.e., samples, akin to MMDDRL.
B DEFINITION OF DISTANCES AND CONTRACTION
Definition of distances. Given two random variables X and Y, p- i--asssesttein metric W, between the distributions of X and Y is defined as
W,(X, Y) = (( (1) Fy
1/p
dw)
=11 Fy 1
(16)
which F-1 is the inverse cumulative distribution function of a random variable with the cumulative distribution function as F. Further, fp distance (Elie & Arthur 2020) is defined as
(,( = ( [Fx(w) - Fx dw
1/r
= Fx - Fy
(17)
12

---

Under review as a conference paper at ICLR 2023
The Ep distance and Wassertein metric identical 1, but otherwise distinct. Note that when p = 2, Ep distance is also called Cramér distance (Bellemare et al. 2017b) dc Y). Also, the Cramér distance has different representation given by
do (c(X,y) x - Y 1 F X
2 E Y Y
(18)
where X'' and Y' are the i..d. copies of X and Y. Energy distance (Szčkely) 2003 Ziel 2020 is natural extension of Cramér distance to the multivariate case, which is defined as
dæ Y) = E Y| X x
E Y'Il,
(19)
where X and Y are multivariate. Moreover, the energy distance is a special case of the maximum mean discrepancy (MMD), which is formulated as
MMD(X, Y; (E (X) Y') - Y)1)
(20)
where k(:-)) is a continuous kernel on X. In particular, if l: is a trivial kernel, MMD degener- ates to energy distance, Additionally, we further define the supreme MMD, which is functional P(X))SX P(X)SXA + R defined as
(M,))
sup (MD,) ()) (x,a)eSxA
(21)
We further present the convergence rate under different distribution divergences.
T is y-contractive under the supreme form of Wassertein distance Wn.
T is /////cootracctive contractive under the supreme form of lp distance.
T is / contractive under MMD with the kernel (x, --1< y11 wa > 0.
Proof of Contraction.
Contraction under supreme form of Wasserstein diatance is provided in Lemma 3 (Belle- mare et al. (017a).
Contraction under supreme form of tp distance refer Theorem 3.4 (Elie & Arthur 2020
Contraction under MMD is provided in Lemma (Nguyen et al. 2020
C PROOF OF THEOREM
Proof. 1. and - It is obvious observe that Sinkhorn loss degenerates to the wasserstein distance. We also have the conclusion that the distributional Bellman operator I" is y-contractive under the supreme form of Wasserstein diatance, the proof of which is provided in Lemma 3 (Bellemare et al. 2017a). Since the above conclusion is made directly based on the limiting case when O, for an unspecified e, we need more rigorous proof. We show that their distance difference is at most an infinitesimal 8.
Firstly, as Wa. > W. and the regularization term is non-negative, using the language of (E, definition, we have: for V6, there exists small positive constant a, such that We,e War < 8 when e < a. Based on that, we have the contraction conclusion:
W
W
T 71, 1272))
W 71 72)
W
2 ",,T22 72)
-R...
< W
T 71, T 72),
(22)
where the second term W Z1, Iz contractive, and thus for the unspecified s, the only difference from the limting 8 = 0 is an infinitesimal 6, which will vanish as s - 0 or + 0.
2. Oo. Our complete proof is inspired by (Ramdas et al. 2017 the definition of squared MMD is
Genevay al. 2018).
Recap
E k (X,X')
2 k X, Y) Y')
13

---

Under review as a conference paper at ICLR 2023
When the kernel function i degenerates to unrectified ka (2. -| y/1 for (0, the squared MMD would degenerate to
X
E Y-11 2 Y
On the other hand, we have the Sinkhorn loss as
We (1, v) 2We.co (1, (1, v)
We,
We (1,
Denoting 1Lg be the unique minimizer for W c.s. it holds that 11; + p2 as - 30. That being said, (1, v) - c(x, v)dp(xx)dx(x)) +
f cl2, (d If c -aa = || yll", || y / (d(pdd)))d(((
we eventually have W W-
E X Y1 Finally, we can have
W 00
2E X Y - E|X X E Y Y
which is exactly the form of squared MMD. Now the key is prove that ITg + a co.
Firstly, it is apparent that (1, v) I c(x, (addu())dvv()) as pov E [[,, v). Let {ex) be pos- itive sequence that diverges to oc, and II, be the corresponding sequence of unique minimizers for We,s. According to the optimality condition, must be the case that f c(æ, y)dlly + +eLKL v) S c(x y)dp & v + (when [T v) = _a 8 v). Thus,
f - O. 1 - Ek
KL (T 8 8D) <
Besides, by the compactness of [1,, v), we can extract a converging subsequence Im, + I Since KL weakly lower-semicontinoous holds that
KL (1 < lim inf KL (Ilm, / 2 v) =0
Hence II vv That being said that the optimal coupling is simply the product of the marginals,
indicating that II, + @ v as E + co. As special case, when a = W oo, v) is equivalent the energy distance
d X = 2 X Y
E X Y Y
(23)
In summary, if the cost function the rectified kernel ka, it is the case that W W-R.. converges to the squared MMD as E + oc. According to (Nguyen et al. 2020), I is /// coottcccea in the supreme form of MMD with the rectified kernel k..
For the unspecified e, we can get the similar result to the case of g + 0. For Vo, there exists large positive constant M, such that MMD contraction conclusion:
Wa. 8 when 2 M. Based on that, have the
Z1, T Z2) W W
I 21, T 22) MMD 2. 22) MMD2 (I"Z, Zz)
Nast
< MMD (I 21, 55
(24)
where the first term MMD only difference from the limiting = 00 an infinitesimal 6, which will vanish as & >> +00 or (M + +00).
(I* 00 (I 21, 21,3"22 (22 is c and thus for the unspecified E, the
For E (0, +o0), the contraction property needs long proof. The proof pipeline firstly we prove the three properties of Sinkhorn divergence, and then we show the contraction of distributional Bellman operator under Sinkhorn divergence based on its properties. Most importantly, we analyzed the contraction under
non-constant factor.
3.1 Properties of Sinkhorn Divergence. We recap three crucial properties of a divergence metric. The first is scale sensitive (S) (of order 3, > 0), i.e., dp(cX, < dcl The second property is shift invariant (I), i.e., X, x & dx The last one unbiased gradient (U). A key observation for the analysis is that the Sinkhorn divergence would degenerate to a wwo-dimensional KL divergence, and therefore embraces a similar convergence behavior to KL divergence. Concretely, according the equivalent form of Wa,(p, v) in Fq.. 9 it can be expressed
14

---

Under review as a conference paper at ICLR 2023
as the KL divergence between optimal joint distribution and a Gibbs distribution associated with the cost function:
(H, == KL (II" (H, ()(,,,)
(25)
where II" is the optimal joint distribution. Thus, the total Sinkhorn divergence is expressed as W (e,(m, 2KL (II" (m. (p.,)))((,,,) KL (IT" (4, ()) K((, pi) KL (II" (v, (K(,,,)
(26)
vi,, the convergence behavior is determined by We.e (1, v), which is similar to the behavior KL divergence. Thus, we will focus on the convergence analysis of (e,,(H, z/). According to the fact that KL divergence has unbiased gradient estimates (U) and shift invariant (I), and Sinkkhorn divergence can be viewed as a two-dimensional KL divergence, both properties of U and can be extended to Sinkhorn divergence. However, we find the non scale sensitive (S) property can not directly apply to Sinkhorn divergence due to the minimum nature of We, (H, 2/) as the optimal joint distribution II" (2, v/) could be different from II (ap, av) where is the scale factor. We need a new rigorous proof of scale sensitive property as follows.
Due to the form of We,e (,,
3.2 Scale Sensitive Property of Sinkhorn Divergence.
We show Sinkhorn divergence satisfies a variant of scale sensitive property when с = -ka that corresponds to non-constant scale factor A(a, a) that is not only a function of the vanilla scale factor and a in ka. but also the two specified probability measures (U, V). By definition, the pdf (V)) After a scaling transformation, the pdf of aU and aV with respect
of ICU V) x e -clr.u)
to and y would be (=(()) and (,)) Thus K(aU,a x e - 2 1u(²) ] ((²) We denote II" and II as the optimal joint distribution of Wa, (p, v) and We,(apu,
/ cx,, y) (z. v) + cKL @ av)
Wa Wa,(aU,aV))
/ dx, y)Il (z,y) + =KL lap @ av)
c= ha (x- w) Te पt (
1 _ 2 ddxdy + 8
x "
log 1 (2. 2)
के (
{ dady
a
|a| [(0 y)) (z, ydddyy +E / n" (x, log T' v) dody
u(z)v(y)
1) -1- (I, dryy + (KL - - (1-- / (x 1)(² a-" (2, v)dzdy =/(0 (1 -((o-1)
=
-1
We.clU,V (1 lal%) f(x a) (a
(27)
where A U.V (a, = 1 - (1-1a () (a.v) (0,1) for E (0, 600) and < 1 due to the (VV)
fact that 0 < (1- (d())((a y) (I, < f(x dil'' (z, y) < VV. A U.V (a, a) function less than that depends on the two margin distributions and the scale factor a. The result implies that we have a new variant of scale sensitive property of Sinkhorn divergence with non-constant factor A U.V (a, a) < when we choose c = -k,, and [a] < 1.
3.3 A New Contraction Mapping Theorem.
We derive a new contraction mapping theorem based on the distribution distance in order to prove the convergence in 3.4.
Theorem 2. (Distribution Contraction Mapping Theorem with a Non-constant Factor) Consider distribution distance d and a function g : P P. The mapping d is contraction: There exists function g(,,)) < such that for V distributions and Y:
dig(X), S ((XYY)d(XX Y,d(Yd(XX)
(28)
Then there exists unique distribution X with (x*) X*.
15

---

Under review as a conference paper at ICLR 2023
Proof. We consider the convergence of the distribution sequence X We have the updating rule as
d(x xx) d(g((*)), 9 S Id xx*-1)
(29)
where we use 0k,k-1 = g(x*, x x for short. Hence, we have
d(x xx) - ,
(30)
Let do - d(xt, ))). From the triangle inequality, we have
,(*))
- d(xx+1 xx*) + d(x R+1 , xxtii-).
11 - + II - ii,--
(31)
π 1= + +1, + -
t = -- + 11+ + II b-1 1= 11. ...)da
For the infinite series + Q++1. + +... according to the ratio convergence judgment method of infinite series, lime--oo infinite series convergent. Due the fact II -1 - as k oc, we have d(X Xx)
i+ - + which we denote as u; for i-the term,
Bi+1 < 1. Thus, the
as k + co. Therefore, it must converge to a limit distribution X that satisfies g(X*) X".
3.4 Contraction of Distributional Bellman Operator under Sinkhorn Divergence.
According to the equation of We.s. it holds the same properties as We,ss i.e., shift invariant and scale sensitive. Thus, we derive the convergence of distributional Bellman operator I under the supreme form of We,e. i... We.:
Wee I 2I'
=sup HaeWee (31,
8.a
(Ro,(B(s,a) + ( ((s,,
5" ==
(32)
sup A A1('' (7, sup We. ( Z2(')
8'
5'.0"
21,22 (7, a) sup W k.. (7 0.) Z2(5'
A Zu (7, a: Wh ka (7, 72)
where the first inequality comes from the scale sensitive property proof of Sinkhorn divergence and we let A Zi a7 = sup.,,a A Z1 a'), a'),Zat'',a' (7). A Z1. a) is only constant function in terms of y and a, we can directly arrive the conclusion that distributional Bellman operator is A 72 a contractive based on existing Banach fixed point theorem. However, the fact that Zi and Zz, and thus we need a new contraction mapping theorem
A 2 2..Za a)) also depends to guarantee the convergence of fixed distribution iteration. According to Theorem 2 in 3.3 that we specifically figure out for the our contraction proof, we have W via distributional Bellman iterations. In summary, we conclude that I" is contractive operator when we use the kk, as the cost function and & 1, while the contraction factor, which is short for ((%, a) 1, is not only function of and 7, but also depends distribution sequence in the while iterations.
c.s can guarantee the convergence
16

---

Under review as a conference paper at ICLR 2023
D PROOF OF PROPOSITION 1 AND COROLLARY 1
Proof. As we leverage I1 to denote the optimal II by evaluating the Sinkhorn divergence via min (H, vi k), the Sinkhorn divergence can be composed in the following form:
We.(Hu, k;
2KL (IT (1, (I, v)) KL (IT' (1, M)]K p)) KL (IT' (v, (=)((,)))
' 1 E EX. ())X, (xx,x log IT (4. vI] + Ex, (c(X, Y)J)
y Dlog IT' (1, +
1
-(Ey, log IT (v, I)
+ By, (o(V, )'I)
&
(X)²
-XX,XYY''
log IT XX)XX) Y!)
+ (Ex,x (6XxX') + Eyy [(YYY")) - 2xx,x' (x(X, Y)))
(X))
=EX,XYY''
log
2MD22.(..)
(X),,
(33)
where the function e in the Gibbs distribution K. is minus Gaussian kernel, i.e., c(x,
-k(x, v) (=--v)/((2"") Till now, we have shown the result in Corollary |
Next, we use Taylor expansion to prove the moment matching of MMD. Firstly, we have the follow- ing equation:
MMD (H, = Ex.x [k(X, + Ey,y [k(Y,"') 2Ex.x [k(X, Y)]
= Ex [6(X) [x(X)T((X')] + Ey,,, [ - 2Ex,x [o(X) o(Y)] (34)
EXlØ(X) - 2(
We expand the Gaussian kernel via Taylor expansion, i.c.,
k(x,, v) e ((0-)))((20))
2
22
20
e
2x
)r
n!.
Vn!
(35)
1=0
Σ
e 20
26
y Vnl
1
ym
Vn!
n=0
<(x) Tby()
Therefore, we have
0
MMD2 v)
BBany [e (200(2))))))))))))()))))))))(1
Braw Te- (2=0)(20)))
0 n!
n=0
(36)
1
(Ma(H)- -Ma(v))
nn!
n=0
Ma(1) = Egany [e-
(20
x²
and similarly for Mr The conclusion is the same as the
moment matching in (guuyen et al. 2020 Finally, due to the equivalence of We,c v) after multiplying s, we have
(X))
Wee = MMD2 +EE
11, X')II (Y,
(37)
(Ma(()
MM(())) 2 + +eE
(X))
2mn
X'XX,X")IIXXXX))
n=0
17

---

Under review as a conference paper at ICLR 2023
This result is also equivalent to Theorem [ where II would degenerate too. In that case, the first regularization term would vanish, and thus the Sinkhorn divergence degrades to a MMD loss, i.., MMD*<(u. v))
E HUMAN-NORMALIZED SCORES
Mean
Median
Human
DDQ
DQN
438.7 %
43.6 %
17
0
C51
1043.4 %
103.7
26
42
QR-DQN-1
1286.4 %
108.6 %
31
47
MMDDRL
924.6 %
117.5
43
(RDQN(tf)
535.1 %
108.2 %
28
40
MMDDRL(1f)
665.0 %
99.8
27
39
SinkhorDDRL
1435.8
113.0
42
Table 2: Mean and median of best noman-normaaliedd scores across 55 Atari 2600 games. The
results for all considered algorithms are averaged over seeds. DQN, C51, QR-DQN-1 MMDDRL based on Pytorch implementation after 10M time steps adapted from (Zhang 2018), while QRDQN((f)) and MMDDRL are training results also after 10M time steps from the tensorflow Dopamine framework of MMDDRL (Nguyen al., 2020 released in (Nguyen-Tang 2021
Human normalized score equation (algorithm randomplay) (human randomplay). Our im- plementation of DQN, QRDQN-1, C51, MMDDRL, Sinkhorn is based on (Zhang 2018) and all the experimental settings, including parameters are identical to the distributional RL baselines im- plemented by (Zhang] 2018). The results about mean and median human- normalized scores of all considered distributional RL algorithms are reported Table 2 We also compare the performance of QRDQN(1f) and MMDDRL after the same 10M time steps based tensorflow implemen- tation on Dopamine framework (Castro et al. 2018). These results are averaged over training data released in (Nguyen-Tange (021). We argue that coamannnorrmalzeen scores may be limited to eval- uate the superiority of algorithms as mean can be highly affected by the performance on games with high-level returns. For example, in Figure 2 MMDDRL superior to QR-DQN as Sinkhorn out- performs MMDDRL on a smaller portion of games compared with QRDQN, but the mean score for MMDDRL Table[ii lower than QRDQN. By contrast, SinkhornDRL is superior terms of mean score and competitive in terms of median. We also provide all average results in Table Jo Appendix and learning curves our implemented algorithms in Appendix
F MORE EXPERIMENTAL RESULTS
We provide learning curves of DON, QRDQN, C51, MMD and SinkhornDRL algorithms on all It illustrates that SinkhornDRL dramatically surpasses the
55 Atari games in Figures 45 6 8 9 other distributional RI. algorithms on large amount of environments, e.g., Venture, Atlantis, Tennis and Spacelnvader, and presents competitive performance or is only slightly inferior as opposed to the state-of-the-arr baselines on other games. Note that the average improvement of SinkhornDRL on Venture game is significant owing to one to two times convergence of SinkhornDRL algorithm over 3 seeds, while the other baselines do not converge over the considered seeds. Although this improvement may also suffer from the instability issue. its occasional success for our DiL algorithm also presents huge potential on some complicated environments. We leave the further exploration the advantage and potential of SinkhornDRI. algorithm as the future work.
G SENSITIVITY ANALYSIS AND COMPUTATIONAL COST
G.1 MORE RESULTS IN SENSITIVITY ANALYSIS
Decreasing E. We argue that the limit behavior connection as stated Theorem I may not be able to be verified rigorously via numeral experiments due the numerical instability of Sinkhorn
18

---

Under review as a conference paper at ICLR 2023
Breakout
Enduro
Pong
2000
GSI
Aversge Return
-00
MM4
3003
OODON
Averaga Return
300
2300
200-
0.
0.2
0.2
0.2
8.5 Time Steps (1c7)
Time Sleps (1c7)
Time Sleps (1e7)
RersReeenge
Alien
BattleZone
DCN
2259
Averaaee Return
20000
Average Relum
3750
Jueriig Rolum
22000
NAG
WIVI
1250
18400
730
10000
3000
250
0.
0.4
6.5
0.4
0.0
04 Time Stepe (1e7)
Time Steps (1c7)
Time Steps (1e7)
Spacelnvaders
Berzerk
Average Return
1893
Average Retum
16000
12000
DGB
CS1
GADONN
Qbert
Averago Return
1500
1750
DON
MVD
1240
8000
Bekham
1900
400
750
4000
210
0.0
0.2
0.4
0.5
0.5
0.0
0.2
0.4
055
0.5
3.0
Steps (1o7)
Time Steps (1e7)
Time Steps (1e7)
Figure 4: Performance of SinkhomDRL compared with DQN, C51, QRDQN and MMD on Break- out, Enduro, Pong, YarRevenge, Alien, BattleZone, Berzerk, Qbert and SpaceInvader.
19

---

Under review as a conference paper at ICLR 2023
UpNDown
DON
C51
W
Asterix
Asteroids
Average Relur
20000
Auosuge Ratun
E230
10000
5000
Time Stops (1o7)
Time Steps Steps((167)
Thmm Steps (1=T)
CrazyClimber
BeamRider
Centipede
20.0
Avarage Relum
0200
200000
2000
400
8:4
Thme (1o7
Time Stops (167)
Time Steps (1e7)
Derby
Frostbite
PGN CM
Riverraid
UVD
MAD
Cintemm
Timl Stope (107)
Time Stops (167)
Time Steps S1e7)(167)
Figure 5: Performance of SinkhornDRL compared with DQN, C51, QRDQN and MMD on UpN- Down, Asterix, Asteriods, BeamRider, Centipede, FishingDerby, Frostbite and Riverraid.
20

---

Under review as a conference paper at ICLR 2023
TimePilot
StarGunner
Seaquest
1600
Average Rotuum
5900
12000
CRooss
Avarage Rehum
1980
HAO
12:00
R004
8000
00
02
Steps (1e7)
Steps (1m7)
(107)
Phoenix
Tennis
750
OREON
TOO
3.00
Averaage Retur
GEDON
18000
Borkhorn
2004
1594
Stepps
Steps
Time Steps (167)
Venture
VideoPinball
Tutankham
700000
+50.00
Average Return
CEITION
500
Saiduar
8.0
Average
206
1808]
60
Time Steps (1c7)
Time Steps (1e7)
Figure 6: Performance of SinkhomDRR compared with DQN, C51, QRDQN and MMD on TimePi- lot, StarGuner, Seaquest, LameThisGmme Phoenix, Tennix, Tutankham, Venture and VideoPinball.
21

---

Under review as a conference paper at ICLR 2023
RoadRunner
James bond
IceHockey
2010
Cat
DGN
DON
VMC
ORDON
Auvenge Patten
OREINN VINU
2000
Time Stapa (1=7)
Time Steps (107)
Time Steps [Tw7]
BankHeist
Hero
DON
Atlantis
1889
MND
CD
GROON KMM Sindhore
CS1
Avoorago Retum
CHESONN
CC1
Mi
(1w7)
Time Steps (1a7)
Time Steps (167)
Amidar
Assault
1204
DGN
WizardOlWor
C91
DGN
GRDON MVD
Averaga Ruruugaauuuuu Richum
700
GRDON
MT4D
CREOON
CST
Sinthor
6808
Sealhom
MMD
Avsrraga
Benitom
400
0.2
Time Steps (167)
Time Stops (1o75 (107)
Timo Stops (167)
Figure 7: Performance of SinkhorDDRL compared with DQN, C51, QRDQN and MMD on Road- Runner, Jamesbond, IceHockey, Hero, HankHeist, Atlantis, WizardOfWor, Amidar and Assault.
22

---

Under review as a conference paper at ICLR 2023
Boxing
Bowling
DON CRNON
DoubleDunk
CSI GRIOON
Averrge Retumn
WNC
Avenga Rahum.
16
Averopp
64
25 (1a7)
114
94
(1w7)
(1o7)
Freeway
Kangaroo
Gravitar
r00
DON
CS1
Average Roturn
GHDON MMD
Return
Average Retum
CRDCN
MMD
MND
8000
Seikeom
Sinthom
200
100
2000
Time Stepei (1e7)
Stepe (1e7)
8.
0.2
Tiim Stops (1o7)
E-4
0.5
KungFuMasser
Krull
MontezumalReennge
DDN
40000
Avorag Rotunn
Seaivare
Average
2bee
1008
Time Steps (107)
Time Steps. (1w7)
Time Stops (107)
Figure 8: Performance of SinkhoraDRI. compared with DQN, C51, QRDQN and MMD on Bowl- ing, Boxing, DoubleDunk, Freeway, Gravitar, Kangaroo, Krull, KunFuMaster and MontezumaRe- venge.
23

---

Under review as a conference paper at ICLR 2023
MsPacman
Pitfall
PrivateEye
DGN
NMD
GIRON
200
Average Relunn
5900
Average
2000
2390
Shikhom
Average
-000
1590
2000
0.5 Steps (1c7)
Time Steps (1e7)
Time Steps (1e7)
Solaris
Robolank
Skiing
-1909
GREINN
C51
Ausraga Rotum
GRIGN HM4
300
Avtrage
18
<10000
-1080
Time Stops (107)
Time Stass (107)
Time Steps (167)
Zaxxon
ChopperCommand
Gopher
Evaraga Return
1964
19009
12000
CS1
DGN
GROON
5000
GRDON
DGN
13808
18804
CS1
GIEDINN
D0N
MMD
MMD
8068
Average
3000
Strkhom
Average
0009
1000
2K0
02
6.4
0.2
9.4
(1o7)
Time Steps (1e7)
Steps (107)
DemonAttack
09000
CR1
DALDON
500000
Sensuun
Average
XX090
63
0.
E-6 0.8
Time Staps (107)
Figure 9: Performance of SinkhornDRL compardd with DQN, C51, QRDQN and MMD on MsPac- man, Pitfall, PrivateEye, Robotank, Skiing, Solaris, Zaxxon, ChopperCommand, Gopher and De- monAttack.
24

---

Under review as a conference paper at ICLR 2023
Breakout
Breakout
600
E=10
Average Return
500
990
r-0.3
E=1000
Retur
499
400
E=5000
ORDON
300
MMD
200
100
180
0.0
0.2
0.4
0.6
0.8
Time Steps (1e7)
Steps (1e7)
(a) Small e in SinkrhomDRL vs QRDQN
(b) Large in SinkrhomDRLL vs MMDDRL
Breakout
Seaquest
500
1m1
20000
E=10
Average Return
400
1=100
Average Retur
c=100
5=00
15000
5=500
300
QRDON
200
Time Steps (1e7)
Time Steps (1e7)
(c) E on Breakout
(d) on Seaquest
Figure 10: (a) Sensitivity analysis w.r.t. a small level of SinkhomDRL compare with QR-DQN that approximates Wasserstein distance on Breakout. (b) Sensitivity analysis w.r.t. a large level of SinkhoraDRL algorithm to compare with MMDDRL on Breakout. All learning curves are reported over 2 seeds. (c) and (d) are results for general E on Breakout and Seaquest, respectively.
observe that if gradually decline O, SinkhornDRL's performance tends to degrade and approach to QR-DQN Note that an overly small will lead trivial almost 0 KC. in Sinkhorn iteration in Algorithm 2 and will cause 1 numerical instability issue for and by in Line 5 of Algorithm 2 In addition, we also conducted experiments on Seaquest, the similar result also observed in Figure IO (d). As shown Figure I0 (d), the performance of SinkhornDLL is robust when E = 10, 100, 500, but small = 1 tends to worsen the performance.
Iteration in Algorithm 2 From Figure 10 (a),
Increasing Moreover, for breakout, if we increase e, the performance of SinkhomDRL tends to degrade and be close to MMDDRL as suggested in Figure is also noted that an overly large will let the Kij explode to co. This also leads numerical instability issue in Sinkhorn iteration in Algorithm
Samples N. We find that SinkhorDRRL requires proper number of samples to perform fa- vorably, and the sensitivity w.r. N depends on the environment. As suggested in Figure [1 (a), smaller N, e.g., N 2 on breakout has already achieved favorably performance and even acceler- ates the convergence in the early phase, while N 2 Seaquest will lead to the divergence issue. Meanwhile, an overly large worsens the performance across two games. We conjecture that using larger network networks generate more samples may suffer from the overfitting issue, yielding the training instability (Bjorek al. 202I). In practice, we choose a proper number of sample, i.c.,
N 200 across all games.
25

---

Under review as a conference paper at ICLR 2023
Seaquest
Breakout
Samples22
20000
Samples=2
Samplz:5
Average Retur
Samples=5
Average Retur
400
Samples-200
Sampler20
S6
15000
Samples=20
Samples=200
300
Samp
10000
Samples=500
0.
Time Steps (1e7)
Time Steps (1e7)
(a) Samples Breakout
(b) Samples Seaquest
Figure 11: Sensitivity analysis of Sinkhorn in terms of the number of samples N on Breakout (a)
and Seaquest (b).
G.2 COMPARISON WITH THE COMPUTATIONAL COST
We evaluate the computational time every 10,000 iterations across the whole training process of all considered distributional RL algorithms and make a comparison in Figure [2] It suggests that SinkhoraDRL indeed increases around 50% computation cost compared with QR-DQN and C51, but only slightly increases the the cost in contrast to MMDDRL on both Breakout and Qbert games. We argue that this additional computational burden can tolerant in view of the significant outper- formance of SinkhoraDRL. in a large amount of environments.
In addition, we also find that the number of Sinkhorn iterations I, is negligible to the computation cost, while overly large samples N, c.g., 500, will lead to a large computational burden as il- lustrated Figure 13 This can be intuitively explained as the computation complexity of the cost function Czj is ON in SinkhomDDLL, which is particularly heavy in computation if N is large enough.
Brcakout
Qbert
Average Computation Time
Average Computation Time
121
100
DQN QR-DON C51
Sinhorn
DQN QR-DQN C51 MMD Sinkhorn
Figure 12: Average computational cost per 10,000 iterations of all considered distributional RI. algorithm, where we select = 10, L = 10 and number of samples N 200 in SinkhomDRL. algorithm.
H RAw SCORE TABLES ACROSS ALL ATARI GAMES AFTER 10 TIMESTEPS (40M FRAMES)
For distributional RI algorithms, results are averaged over 3 seeds after 10 timeSteps. i.e., 40M Frames. Results of DQN, C51, QRDQN, MMD and Sinkhorn are based on our implementa- tion Pytorch, which is adapted from the reliable open-source implementation (Zhang 2018). QRDQN((1) and MMD(tf) are from training results after 40M frames from the implementation of MMDDRL (Nguyen et al. 2020) in (Nguyen-Tang 2021
26

---

Under review as a conference paper at ICLR 2023
Breakout
Qbert
Average Computation Time (
Average Computation Time (s)
290
Sinkhom (samples-20)
Sinkhorn (sample=-200)
Sinkhom (samples-350)
Sinkhorn (samples-500)
Sinkhorn (samples-20)
Sinkhom (smppss-200)
Sinkbom (samples-350)
Sinkhorn (sample-500))
Figure 13: Average computational cost per 10,000 iterations of SinkhomDRL algorithm over differ-
samples.
27

---

Under review as a conference paper at ICLR 2023
GAMES
RANDOM
HUMAN
DQN
CS1
QRDQN
QRDQN(1)
MMD
MMD(th)
Sinkhorn
Allen
211.9
7,127.7
1334.0
1946.0
1625.0
[425.8
2218.0
1683.1
1873.0
Amidar
2.34
1,719.5
400.2
354,5
554.6
870,2
706.4
694.7
506.7
Assault
283.5
742.0
5651.8
3368.1
7593.6
79318
6001.5
8947.5
3771.0
Asterix
268.5
8,503.3
$490.0
31860.0
7660.0
17126.8
15890.0
3641 1.5
7610.0
Asteroids
1008.6
47.388.7
826.0
1660.0
1610.7
1095.0
1460.9
624.0
Atlantis
22188
29.028.
18990.0
1490040.0
2520080.0
859419.7
80920.0
935273.3
3417430.0
BankHeist
753.1
657.0
948.0
1000.0
936.7
1034.0
982.8
849.0
BattleZome
3000
18.5
22100.0
28400.0
37800.0
25223.2
28400.0
20089
27000.0
BeamRider
414.3
16.926.5
9519.0
13069.2 2
8043.8
9728.9
14072.6
16889.
9865.6
Berzerk
165.6
2,630.4
746.0
824.0
928.0
766.3
959.9
802.3
1029.0
Bowling
23.48
160.7
29.6
30.3
35.5
33.8
60.0
37.3
12.6
Boxing
-0.69
12.1
96.0
91.8
98.3
97.9
96.9
92.3
96.7
Breakout
1.5
30,5
313.4
373.0
361
391.0
405.9
486.9
402.5
Centipede
2064.77
12,017.0
6090.9
5508.0
5866.0
5152.0
4952.2
Chopper
794
7,387,8
2780.0
4360.0
5490.0
3575.7
6760.0
2465 8
6520.00
Crazy Climber
35,829.4
15960.0
58070.0
69430.0
85543.6
111330.0
115997.2
16000.0
Demon Attack
162.25
1,971.0
58324.5
41656.5
63889.0
94433.0
437760.5
98849.8
195827.0
DoubleDunk
-18.14
-164
0.6
-15.1
-14.9
-2.2
Enduro
0.01
860.5
1961.3
1507
2812.5
1912.4
3248.2
1839.5
4272.0
DerbyngDerby
-38.7
21.6
Freeway
0.01
29.6
30.9
32.6
34.0
31.1
33.66
33.2
34.0
Frustbite
Gopher
364
2,412.5
7058.0
9314.0
6466.U
3854.7
6412.0
4966.6
15168.0
Gruvilar
470.0
Hero
551
30,826.4
4657.5
8098
11673.5
9779.7
7215.0
8382.4
7476.0
IeeHockey
-10.3
0.9
-13.0
-11.4
-3.7
-4.5
-2.4
-4.6
Jameshond
302.8 8
320.0
625.0
1995.0
694.0
480.0
657.4
450.0
Kangaroo
3,035.0
660.0
9870.0
13440.00
14398.2
14720.0
8548.
10680.0
Krall
1,566.59
2,665.5
9191.
9366
9918.7
9293.
8732.7
6069 8
9549.0
KungFuMaster
451
22,736.3
62800.0
55060.0
36020.0
27786.9
36940.0
28394.6
42600.0
Montezuma Revenge
4,753,3
1.0
1.0
0.0
1.0
0.0
0.0
M&Pacman
242.6
6,951.6
3230.0
2168.0
2673.0
2557.7
2568.0
3244.8
2568.0
NameThisGame
2404.9
8,049.0
4702.0
6278.0
11739.0
11161.7
12394,0
10859.2
9200.0
Phoenix
757.2
7,242.6
$398.0
12043.0
12324.0
21813.7
32086.0
30561.7
18558.0
Piffall
-265
6,463.7
1.0
1.0
-37.4
1.0
-82.8
0.0
Pong
-20.34
14.6
20.0
20.7
20.8
20.5
20.9
20.3
21.0
PrivateKyye
34.49
69.571.3
100.0
100.0
100.0
62.1
100.0
284.2
100.0
Obert
188.75
13,455.0
8150.0
16575.0
13830.0
12307
15782.5
13325.7 7
6530.0
RiverRaid
17.118.0
8350.0
10232.0
10975.4
ReadRunner
7,845.0
44950.0
$4490.0
54620.0
43574.9
42530.0
37909.9
$2600.0
Robotank
11.9
22.5
48.1
52.4
48.4
48.1
Seaquest
88.2
42.054.7
1444.0
10666.0
2640.0
4424.3
11685.0
1745.2
14795.0
Skiing
4,336.9
-23701.6
29970.3
Solaris
2346.6
12,326.7
582.0
192.0
956.0
1220.5
3336.0
996.8
792.0
Spucelnvaders
1.668.7
StarGunner
631
10,250.0
1270.0
22600.0
38380.0
$2937.6
$2050.0
46658.3
43820.0
Tennis
-8.3
13.3
TimePilot
3682
5,229,2
1420.0
3260.0
6030.0
6338.7
7900.0
4516.2
7060.0
Tutankham
15.56
167.6
206.6
186.0
178.3
153.8
205.2
190.5
202.6
UpNDown
604.7
11,693.2
19145.0
16046.0
17074.0
14402.2
44746.0
19350.3
20063.0 O
Venture
1,187.5
1,0
39.3
1370,0
VidenPinhall
15720.98
17.667.9
270050.9 9
477206.8
388106.7
44599.9
288137.2
152686.0
164597.3
WizardorWor
4,756.5
1440.0
1620.0
4890.0
8670.
4480.0
5095.3
3250.0
YarsRevenge
3271.42
54,576.9
12507.9
15954.4
17593.8
12336.
8516.8
13420.7
13507.3
Zasxon
9,173.3
1.0
5910.0
7410.0
9077.0
4640.0
8557.9
10320.0
Table 3: Scores of all algorithms averaged over 3 seeds across 55 Atari games.
28


Reviews:
[{"review_id": "JS7BY2WxNB", "reviewer": "Reviewer_tLLY", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper introduces a new method for solving distributional RL problems. The main idea behind the paper is to use the Sinkhorn distance (wasserstein with entropic regularization) as the basis for the distributional updates. Several theoretical results are proven concerning the contractive property of the Sinkhorn update as well as equivalences with other updates in certain settings. Extensive experimental results are presented with comparison to existing distributional RL algorithms", "summary_of_the_review": "In summary, this paper presents an interesting idea and does an interesting theoretical analysis and extensive experimental evaluation. However, it seems to fall short in a few categories as described in the above sections, with several unanswered questions. The lack of clarity is also a drawback.\n\nOther notes:\n- More exact convergence results exist for the general sinkhorn algorithm: “Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration”\n- I encourage the authors to also consider more in depth analysis of experimental results as prescribed by the following paper considering only 3 seeds are used: \"Deep Reinforcement Learning at the Edge of the Statistical Precipice\"\n", "strength_and_weaknesses": "Strengths:\n- The update rule appears to be novel for distributional RL; however, it is already known that you can use any distance. This is just one particular example that has also seen applications and success in other areas.\n- The experiments are fairly extensive, evaluating the algorithm against existing DRL methods across 50+ Atari tasks\n- The theory seems reasonable, but I didn’t check carefully.\n\nWeaknesses:\n- Although probably new in application, the algorithmic idea is not particularly novel, so the results should probably be judged more on significance of improvement and analytical insights.\n\n- The clarity could be greatly improved. There are numerous inconsistencies, typos, undefined notation, etc. See specifics below.\n\n- The novel part of Theorem 1 (c) is not actually a contraction in the usual sense as it depends on the iterates. It’s not clear that this is truly significant and actually yields the desired convergence.\n\n- The rates in Table 1 are not proven (there is no theorem) or discussed or interpreted. This is especially important since the first one is convoluted.\n\n- The empirical results do not appear to yield that significant of an improvement over existing algorithms. But there is improvement in some domains.\n\n- There are claims made in the experiments section that don’t appear to be supported by the results. It says “SinkhornDRL achieves better performance across more than half” but it looks like it is the opposite. It’s really only half the games for (a) and less than half for (b). Am I just reading the plot backwards? Blue (negative) is worse for Sinkhorn, no? \n\n“the performance of SinkhornDRL tends to QR-DQN or MMDDRL on Seaquest when we decrease or increase epsilon.” This doesn’t seem to be supported by the plots. The performance decreases, but it’s not at all clear that it’s leveling out at the performance of the other algorithms.\n\n\n\nQuestions:\n- Why was DQN not included in the ratio plots?\n\n- In Fig 3 (b), since the performance doesn’t seem to be effected much by N, why not choose N to be very small such as N = 1, 2? Theoretically it doesn’t make much sense, but it seems it would be useful to see when performance finally drops off.\n- “Sample complexity” in Table 1 is never defined. Sample complexity to achieve what? What is the objective? Sub-optimality in distribution, in expectation, etc? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "7hwNf0ytIAZ", "reviewer": "Reviewer_vqCE", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new DRL method leveraging Sinkhorn divergence. Analogous to DQN, they consider learning the action value distribution by minimizing the divergence between the current distribution and the target distribution. The main difference with previous DRL methods is they replace the divergence to be Sinkhorn divergence. Theoretically, they study the convergence of the distributional\nBellman operator, the moment matching meaning of Sinkhorn divergence, and the connection with MMD DRL. Although the paper is generally well-written, the experimental results does not show significant improvements over the current state-of-the art methods. ", "summary_of_the_review": "This paper proposes a new DRL method leveraging Sinkhorn divergence. Analogous to DQN, they consider learning the action value distribution by minimizing the divergence between the current distribution and the target distribution. The main difference with previous DRL methods is they replace the divergence to be Sinkhorn divergence. Theoretically, they study the convergence of the distributional\nBellman operator, the moment matching meaning of Sinkhorn divergence, and the connection with MMD DRL. Although the paper is generally well-written, the experimental results does not show significant improvements to the current state-of-the art methods. ", "strength_and_weaknesses": "Strengths: 1. The new contributions of this paper is to replace the divergence in DRL to be Sinkhorn divergence, which can be viewed\nas an interpolation between Wasserstein distance and MMD.\n\nWeakness: 1. \"The main limitation of our proposal is that the superiority over existing state-of-the-art algorithms\nmay not be sufficiently significant.\" The authors themselves point this out. So what is the advantage of using Sinkhorn iterations?\n2. Figure 3 (b) Why does the performance degrade with increase in sample size? Increasing the samples should better approximate the return distribution.\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "010OJ8erPm8", "reviewer": "Reviewer_VKzt", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes a new distributional RL algorithm named SinkhornDRL, which proposes to solve the distribution matching problem in distributional RL with Sinkhorn divergence. \n\nSinkhorn divergence in formulation is equivalent to an entropy-regularized version of MMD, and it could be featured as an interpolation between Wassersteon distance  and MMD, so that it simultaneously leverages the geometry of Wasserstein distance and unbiased gradient estimate property of MMD in moment matching.  SinkhornDRL is highly related to MMDDRL, as they are both moment matching approaches that use samples to approximate the distribution associated with kernel tricks. \n\nThe main new theoretical insight introduced in this paper is to prove the contraction property of SinkhornDRL in a short section from Appendix (Equation 32 in Appendix 3.4).  The other proof on the statistical properties of Sinkhorn, such as its equivalence form to regularized MMD, is not novel. \n\nFor empirical comparison, the authors compare SinkhornDRL with DQN, C51, QRDRN and MMDDQN on their modified 55 Atari games benchmark, where they evaluate the agents under 10M frame regime. The results show that SinkhornDQN could outperform MMDDQN occasionally and overall results in inferior median and median HNS than MMDDQN. \n\n", "summary_of_the_review": "This paper introduces a novel distributional RL method where the distributional divergence is estimated from deterministic samples with Sinkhorn divergence. The authors evaluate the method on a Atari 2600 domain under a customized frame number 10M, where I feel the results reveal MMDDRL outperforms SInkhornDRL and the settings are falsely claimed. \n\nI'll re-evaluate the recommendation score to this paper if  my concerns on (1) technical novelty, (2) presentation on empirical settings and results,  (3)  sufficient details, atari wrappers, code or pseudocode  to implement the experiment could be provided to convince me the result is reproducible. It's even better if the authors could replace the results under 10M with those under 200M. ", "strength_and_weaknesses": "**Strength**\n* Sinkhorn divergence has not been formally considered for distributional learning in the existing distributional RL literature.\n* The paper is clearly written and easy to understand.\n\n**Weaknesses**\n\n[Motivation] \n* Overall I feel this work is not well motivated from distributional RL perspective. The motivations  introduced for the work are mainly related to: (1) introducing a new distributional RL algorithm  (2) SinkhornDRL bridges between Wasserstein distance and MMD. It remains unclear what challenge in distributional RL could be  tackled by Sinkhorn, and why bridging the two major types of distributional attempts could be more beneficial than employing each single one. (I think motivating the RL perspective from the entropy-regularization might also be valid but there is little statement of this type.)\n* The related works discussed/compared in this paper are not inclusive enough. In the sections before methodology, I only see two major distributional RL methods, QR-DRL and MDDRL,  introduced in details, while it would be unfair to describe the literature through only two works. In some important part in methodology part. the authors also compare SinkhornDRL with a limited subset of distributional RL baselines, e.g., the complexity and convergence rate shown in Table 1. I refer the authors to the related works presented in the MDDRL paper, where there is a more inclusive and informative categorization/description on related works.\n\n\n[Formulation] \n\n* The overall SinkhornDRL solution comes with limited novelty because the algorithmic approach of solving moment matching with unrestricted statistics of deterministic samples highly depends on the existing work MMDDRL. The main change of SinkhornDRL from  MMDDRL is very small, i.e., to minimize a regrarized squared MMD, while the major idea like applying unrestricted statistics and approximating the distance with samples of measures, are the same.  \n*  The statistical relationship between Sinkhorn and MMD is not newly developed, since it's well adopted from existing statistical literature. \n* The authors make rather strong claim on the theoretical contribution of this paper, but I feel it is not very convincing as the only new proof is the proof of contraction theorem only, which itself is not very complicated. Furthermore, considering the limited algorithmic change has not been well justified by the empirical evaluation, I feel the significance of the proposed method is not enough for being published in ICLR conference.\n\n[Empirical Evaluation] \n\n* The authors only adopt a single experimental domain Atari 2600, but they could not align the well adopted Atari 2600 setting with the other baselines which all adopt the same domain, e.g., C51, QRDQN and even MMD. The authors only adopt a much smaller frame number, which they claim is **10M**, 1/20 from the most widely adopted standard **200M**.  I understand the computational resource required for Atari 2600 is a lot, but if not accessible to such resources, should turn to alternative domains which have lots of candidates (e.g., Mujoco and ViZDoon), as well as proof-of-the-concept toy domains. To me, lack of the computational budget is not a valid excuse for using such small frame budget and make the setting biased from its most related baselines.\n\n* There is one important issue the experiment, which is the **10M** frame budget the authors claim they use. $\\textcolor{red}{\\mbox{I suspect the authors actually use a much larger number than 10M}}$. The authors never mention *frame repeat* or  *frame stack* in their paper, so based on the well-adopted standard, **10M** should correspond to the total number of frames, not that number before multiplying frame repeat. So it's 1/20 to C51, QRDQN, MMDDQN and many other classic DRL methods like IMPALA, RAINBOW and MUZERO. It's important to make that much progress within 10M only. For example, with Breakout, SinkhornDRL uses less than 5M frames to score > 350, but that progress in C51 paper's Fig 3, QRDQN paper's Fig 6 and even Rainbow's Fig 5, it takes way more than 10M frames for Breakout agent to progress to that standard. [Highly recommend reviewers and AC to take a look into aforementioned papers]. \n\n* I think the reported median/mean HNS scores for Sinkhorn as well as its baselines (in Table 2) are not responsible. I have computed the median HNS for C51 and QR-DQN from a very reliable  open-source experiment logs released in a public repository dqn-zoo. Within 10M frame, the median HNS for both methods are bellow 50%, and within 10M * 4 repeat (assume they forget to multiply frame repeat),  the scores for both methods > 100%. For neither case, the standards of median HNS aligns with the reported numbers. Therefore, I'm not convinced if SinkhornDRL could be identified as an effective distributional RL method. $\\textcolor{red}{\\mbox{Even their reported numbers for Sinkhorn underperform MMDQN in both median and mean HNS metrics. }}$\nIt seems that the empirical results for their sole benchmark domain reveals the performance of MMDQN is better than Sinkhorn. \n\n* In Fig 2, the ratio improvement (%) curve seems to show that both QRDQN and MMDDRL could dominate Sinkhorn in the number of games they outperform Sinkhorn, i.e., width for orange bars are shorter than blue bars.\n\n* The baselines are not inclusive. Note that both C51 and QRDQN are not methods with the unrestricted distributions like SinkhornDRL. Instead, their follow-up yet well-adopted works IQN and FQF are methods with unrestricted distributions. I feel the authors should include those unrestricted distribution baselines (apart from MMDQN) for compression. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "-3fjfF-I-D", "reviewer": "Reviewer_LQMW", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes the sinkhorn algorithm for optimal transport to learn the distributional Q-estimate via a bootstrapped wasserstein distance between distributions. It builds on top of existing distributional off-policy model-free RL methods which estimate a distribution over the discounted sum of rewards ahead and leverage it for policy learning. Additionally, the paper proposes a theoretical framework wherein their method interpolates between the wasserstein distance and MMD between distributions. ", "summary_of_the_review": "Overall, while this is a nice applied paper, the Sinkhorn algorithm used in this context requires further study, in particular why it would work well for certain tasks and not others. I would recommend rejecting the paper in its current state but would increase my score if more justification could be provided for how beneficial this approach is vs existing ones for other tasks.", "strength_and_weaknesses": "Strengths:\n- The paper provides a nice theoretical framework for distributional RL via sinkhorn iterations and how the optimal transport problem it solves in distributional RL interpolates between Wasserstein and MMD distances, both of which are appealing for their contractive divergences in Bellman equations. Furthermore, the free choice of cost function is an interesting one which allows more flexibility than the contractive requirement for divergences needed for conventional distributional RL. It is shown that sinkhorn algorithm using Gaussian kernels is equivalent to regularized MMD. \n- There is an implementation to utilize the sinkhorn algorithm for distributional value-distribution learning in the RL setting, which requires Sinkhorn to be differentiable and efficiently computed.\n- For several tasks, the proposed algorithm significantly outperforms a state-of-the-art distributional RL algorithm in MMDRL, which is impressive.\n\nWeaknesses:\n- The advantages of OT vs MMD is known to be mostly taking advantage of the data geometry and the geometry of the flat distance measures when following the gradient of the divergences. However, it is unclear to me whether this advantage is applicable to the setting of estimating a 1-dimensional value-distribution. Furthermore, OT suffers from high computational cost, with the same number of design choices required to tune (cost function replacing kernel choice) as MMD. \n- Along the lines of the above, it is unclear whether the benefits of this algorithm outweigh its implementation difficulty along with overall complexity. The results for which MMDRL is superior seems to indicate that for a majority of Atari games, this does not do much better\n- It would be have interesting to see if these results generalize to harder tasks either in discrete control or even continuous control. ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject