D ISTRIBUTIONAL R EINFORCEMENT L EARNING VIA S INKHORN I TERATIONS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 Sinkhorn distributional RL (SinkhornDRL) algorithm that learns a ﬁnite 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 ﬁnds 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 suite of 55 Atari games.
I NTRODUCTION
Classical reinforcement learning (RL) algorithms are normally based on the expectation of discounted 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-art performance in a wide range of environments (Bellemare et al., 2017a; Dabney et al., 2018b;a; 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 connection 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 compare probability distributions with non-overlapping supports, has recently emerged as an appealing contender in various machine learning applications.
It is known that Wasserstein distance was long disregarded because of its computational burden in its original form to solve an expensive network
ﬂow problem.
However, recent works (Sinkhorn, 1967; Genevay et al., 2018) have shown that this cost can be largely mitigated by settling for cheaper approximations through strongly convex regularizers.
The beneﬁt of this regularization has opened the path to wider applications of the Wasserstein distance in relevant learning problems, including the design of distributional RL algorithms.
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 successfully applied in numerous crucial machine learning developments, including the SinkhornGAN (Genevay et al., 2018) and Sinkhorn-based adversarial training
(Wong et al., 2019).
More importantly, it 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 et 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 (Dabney et al., 2018b), aimed at approximating Wasserstein distance, suffers from 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, we propose a 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 a 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 SinkhornDRL algorithm with typical baselines on 55 Atari games, verifying the competitive performance of our proposal.
Our method inspires researchers to ﬁnd a trade-off that simultaneously leverages the geometry of the Wasserstein distance and the favorable unbiased gradient estimate property of MMD while designing new distributional RL algorithms in the future.
P RELIMINARY K NOWLEDGE 2. D ISTRIBUTIONAL R EINFORCEMENT L EARNING
In classical RL, an agent interacts with an environment via a Markov decision process (MDP), a 5-tuple ( S , A , R, P, γ ), where S and A are the state and action spaces, respectively.
P is the environment transition dynamics, R is the reward function and γ ∈ (0 , 1) is the discount factor.
From Value function to Value distribution.
Given a policy π , the discounted sum of future rewards is a random variable Z π ( s, a ) = (cid:80) ∞ t =
γ t R ( s t , a t ) , where s
=
s , a 0 = a , s t +1 ∼ P ( ·| s t , a t ) , and a t ∼ π ( ·| s t ) .
In the control setting, expectation-based RL is based on the action-value function Q π ( s, a ) , which is the expectation of Z π ( s, a ) , i.e., Q π ( s, a ) =
E [ Z π ( s, a )] .
By contrast, distributional RL focuses on the action-value distribution, the full distribution of Z π ( s, a ) .
The incorporation of additional distributional knowledge intuitively interprets its empirical success.
Distributional Bellman operators.
For the policy evaluation in expectation-based RL, the actionvalue function is updated via Bellman operator T π Q ( s, a ) =
E [ R ( s, a )] + γ E s (cid:48) ∼ p,π [ Q ( s (cid:48) , a (cid:48) )] .
In distributional RL, the distribution of Z π ( s, a ) is updated via the distributional Bellman operator T π
[Equation]
where s (cid:48) ∼ P ( ·| s, a ) and a (cid:48) ∼ π ( ·| s (cid:48) ) .
The equality implies random variables of both sides are equal in distribution.
The distributional Bellman operator T π is contractive under certain distribution divergence metrics.
We provide a detailed discussion about more related works in Appendix A.
2. D IVERGENCES BETWEEN M EASURES
Optimal Transport (OT) and Wasserstein Distance.
The optimal transport (OT) metric between two probability measures ( µ, ν ) is deﬁned as the solution of the linear program:
[Equation]
where c is the cost function and Π is the joint distribution with marginals ( µ, ν ) .
Wasserstein distance (a.k.a.
earth mover distance) is a special case of optimal transport with the Euclidean norm as the cost function.
In particular, given two scalar random variables X and Y , p -Wasserstein metric W p between the distributions of X and Y can be simpliﬁed as
[Equation]
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 of measures, but it suffers from the curse of dimension (Genevay et al., 2019; Arjovsky et al., 2017).
Maximum Mean Discrepancy.
The squared Maximum Mean Discrepancy (MMD) MMD 2 k with the kernel k is formulated as
[Equation]
[Equation]
where k ( · , · ) is a continuous kernel on X .
X (cid:48) (resp.
Y (cid:48) ) is a random variable independent of X (resp.
Y ).
If k is a trivial kernel, MMD degenerates to the energy distance.
Mathematically, the “ﬂat” geometry that MMD induces on 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 al., 2019).
We provide the detailed introduction of more distribution divergences in Appendix B.
R OLES OF D ISTRIBUTION D IVERGENCE AND R EPRESENTATION 3. D ISTRIBUTIONAL RL: F ROM N EURAL F ITTED Q TO Z I TERATION
Neural Fitted Q-Iteration.
It is known that Deep Q Networks (Mnih et al., 2015) can be simpliﬁed into Neural Fitted Q-Iteration (Fan et al., 2020) under tricks of experience replay and the target network, where we update Q θ ( s, a ) parameterized by θ in each iteration k :
[Equation] [Equation]
where the target y i
=
r ( s i , a i ) + γ max a ∈A Q kθ ∗ ( s (cid:48)
i , a ) is ﬁxed within every T target steps to update target network Q θ ∗ with parameters θ ∗ by letting θ ∗ =
θ and the experience buffer induces independent samples { ( s i , a i , r i , s (cid:48) i ) } i ∈ [ n ] .
In an ideal case that neglects the non-convexity and TD approximation errors, we have Q k +1 θ =
T Q kθ , 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 valuebased distributional RL methods based on a parameterized Z θ into a Neural Fitted Z-Iteration as
[Equation]
where the target Y i
=
R ( s i , a i ) + γZ kθ ∗ ( s (cid:48)
i , π Z ( s (cid:48) )) with π Z ( s (cid:48) )
= argmax a (cid:48) E (cid:2)
Z kθ ∗ ( s (cid:48) , a (cid:48) )
(cid:3) is ﬁxed within every T target steps to update target network Z θ ∗ , and d p is the distribution divergence.
3. K EY R OLES OF d p AND Z θ
Within the Neural Fitted Z-Iteration framework proposed in Eq.
6, we observe that the choice of representation manner on Z θ and the metric d p are pivotal for the distributional RL algorithms.
For instance, QR-DQN (Dabney et al., 2018b) approximates Wasserstein distance W p , which leverages quantiles to represent the distribution of Z θ .
C51 (Bellemare et al., 2017a) represents Z θ via a categorical distribution under the convergence of Cram ´ er distance (Bellemare et al., 2017b;
Rowland et al., 2018), while MMD distributional RL (MMDDRL) (Nguyen et al., 2020) learns samples to represent the distribution of Z θ based on MMD.
We compare characteristics of these distribution divergence, including the convergence rate and sample complexity, in Table 1.
Theoretical results regarding Sinkhorn divergence is based on (Genevay et al., 2019) and the detailed convergence proof of other distances is also provided in Appendix B.
In summary, we argue that d p and Z θ are two crucial factors in distributional RL design, based on which we introduce Sinkhorn distributional RL.
[Table]
S INKHORN D ISTRIBUTIONAL RL (S INKHORN DRL)
In this section, we ﬁrstly introduce Sinkhorn divergence and apply it in distributional RL.
Next, we conduct a 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. S INKHORN D IVERGENCE AND G ENETIC A LGORITHM
We design Sinkhorn distributional RL algorithm via Sinkhorn divergence.
Sinkhorn divergence (Sinkhorn, 1967) is a tractable loss to approximate the optimal transport problem by leveraging an entropic regularization to turn the original Wasserstein distance into a differentiable and more robust quantity.
The resulting loss can be computed using Sinkhorn ﬁxed point iterations, which is naturally suitable for modern deep learning frameworks.
In particular, the entropic smoothing generates a family of losses interpolating between MMD.
As such, it allows us to ﬁnd a sweet trade-off that simultaneously leverages the geometry of Wasserstein distance on the one hand, and the favorable high-dimensional sample complexity and unbiased gradient estimates of MMD.
We introduce the entropic regularized Wassertein distance W c,ε ( µ, ν ) as
[Equation]
where
KL (Π | µ ⊗ ν )
= (cid:82) log (cid:16) Π( x,y )
d µ ( x )d ν ( y ) (cid:17)
dΠ( x, y ) is a strongly convex regularization.
The impact of this entropy regularization is similar to (cid:96) 2 ridge regularization in linear regression.
Next, the sinkhorn loss (Feydy et al., 2019; Genevay et al., 2018) between two measures µ and ν is deﬁned as
[Equation]
As demonstrated by (Feydy et al., 2019), the Sinkhorn divergence W c,ε ( µ, ν ) is convex, smooth and positive deﬁnite that metrizes the convergence in law.
In statistical physics, W c,ε ( µ, ν ) can be re-factored as a projection problem:
[Equation]
where K is the Gibbs distribution and its density function satisﬁes d K ( x, y )
=
e −
[Equation]
ε dµ ( x ) dν ( y ) .
This problem is often referred to as the “static Schr ¨ odinger problem” (L ´ eonard, 2013; R ¨ uschendorf & 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 W c,ε to measure the distance between the current action-value distribution
Z θ ( s, a ) and the target distribution
T π Z θ ( s, a ) , yielding W c,ε ( Z θ ( s, a ) , T π
Z
θ
( s, a )) for each s, a pairs.
In terms of the representation for Z θ ( s, a ) , we employ the unrestricted statistics, i.e., deterministic samples, due to its superiority in MMDDRL (Nguyen et al., 2020), instead of using predeﬁned statistic functionals, e.g., quantiles in QR-DQN (Dabney et al., 2018b) or categorical distribution in C51 (Bellemare et al., 2017a).
Require :
generated samples
cost function
hyperparameter
Input
then 2: a ∗ ∼ π ( ·| s (cid:48) ) .
3: else 4: a ∗ ← arg max a (cid:48) ∈A 1 N (cid:80) Ni =
1 Z θ ( s (cid:48) , a (cid:48) )
i 5: end if 6: T Z i ← r + γZ θ ∗ ( s (cid:48) , a ∗ )
i , ∀ 1 ≤ i ≤ N
Output : W c,ε (cid:16) { Z θ ( s, a ) i } Ni =1 , { T Z θ ( s, a ) j } Nj =1 (cid:17)
More concretely, we use neural networks to generate samples that approximate the value distribution.
This can be expressed as Z θ ( s, a ) :
= { Z θ ( s, a ) i } Ni =1 , where N is the number of generated samples.
We refer to the samples { Z θ ( s, a ) i } Ni =1 as particles .
Then we leverage the Dirac mixture 1 N (cid:80) Ni
=1 δ Z θ
( s,a ) i to approximate the true density function of Z π ( s, 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 representation is provided in Algorithm 1.
Remark.
From the general algorithm framework in Algorithm 1, our SinkhornDRL generally modiﬁes the distribution divergence comparing the state-of-the-art 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 SinkhornDRL, and the follow-up works IQN (Dabney et al., 2018a) and FQF (Yang et al., 2019) can naturally extend both MMDDRL and SinkhornDRL as discussed in (Nguyen et al., 2020).
4. T HEORETICAL A NALYSIS UNDER S INKHORN D IVERGENCE
Convergence.
Firstly, we denote the supreme form of Sinkhorn divergence as W ∞ c,ε ( µ, ν ) :
[Equation]
We will use W ∞ c,ε ( µ, ν ) to establish the convergence of T π in Theorem 1.
Theorem 1.
If we leverage Sinkhorn loss W c,ε ( µ, ν ) in Eq.
8 as the distribution divergence in distributional RL, and choose the unrectiﬁed kernel k
α := −(cid:107) x − y (cid:107) α as − c ( α > 0 ) , it holds that
(1) ( ε → 0 ) W c,ε ( µ
ν ) → 2 W
α
( µ
ν
When ε = 0 , T π is a γ -contraction under W ∞ c,ε .
(2) ( ε → + ∞ ) W c,ε ( µ, ν ) → MMD 2 k α ( µ, ν ) .
When ε = + ∞ , T π is γ α/ 2 -contractive under W ∞ c,ε .
(3) ( ε ∈ (0 , + ∞ ) ) , T π is a contractive operator under W ∞ c,ε .
The related non-constant contraction factor ∆( γ, α ) < 1 also depends on the distribution sequence in distributional Bellman iterations.
We provide the long yet rigorous proof of Theorem 1 in Appendix C. Theorem 1 (1) and (2) are follow-up conclusions in terms of the convergence behavior of T π based on the interpolation relationship between Sinkhorn divergence with Wasserstein distance and MMD (Genevay et al., 2018).
Our key theoretical contribution is for the general ε ∈ (0 , ∞ ) , in which we conclude that T π is 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 c = − κ α , where the resulting non-constant scaling factor is also determined by the speciﬁed two probability measures.
Next, we propose a new distribution Contraction mapping theorem in Theorem 2 of Appendix C, based on which we eventually arrive at the convergence of distributional Bellman operator under W ∞ c,ε .
Intriguingly yet reasonably, the contraction factor ∆( γ, α ) is non-constant but a 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 to the optimal transport literature.
Consistency with Related Conclusions.
As Sinkhorn divergence interpolates between Wasserstein distance and MMD, its contraction property when the cost function holds c = − k α for the general ε ∈ [0 , ∞ ] is intuitive.
Note that if we choose Gaussian kernels as the cost function, there will be no concise and consistent contraction results as Theorem 1 (3).
This conclusion is also consistent with
MMDDRL (Nguyen et al., 2020), where T π is generally not a contraction operator under MMD equipped with Gaussian kernels as a counterexample has been pointed out in MMDDRL (when ε → + ∞ ).
To be consistent with the contraction property analyzed in our theory (Theorem 1 (3)), we employ the rectiﬁed kernel k α as the cost function in our experiments and set α = 2 , under which SinkhornDRL suggests a favorable performance in Section 5.
Regularized Moment Matching under Sinkhorn Divergence Associated with Gaussian Kernels.
We further examine the potential connection between SinkhornDRL with existing distributional RL families.
Inspired by the similar manner in MMDDRL (Nguyen et al., 2020), we ﬁnd that the Sinkhorn divergence with the Gaussian kernel can also promote to match all moments be-
tween two distributions.
More speciﬁcally, the Sinkhorn divergence can be rewritten as a regularized moment matching form in Proposition 1.
Proposition 1.
For ε ∈ (0 , + ∞ ) , Sinkhorn divergence W c,ε ( µ, ν ) associated with Gaussian kernels k
( x, y )
= exp
[Equation]
−
[Equation]
x − y ) 2 / (2 σ 2 )) as − c , is equivalent to
[Equation]
where Π ∗ ε denotes the optimal Π determined by ε by evaluating the Sinkhorn divergence via min Π
[Equation]
, and similarly for ˜ M n ( ν ) .
We provide the proof of Proposition 1 in Appendix D. Similar to MMDDRL associated with a
Gaussian kernel (Nguyen et al., 2020), Sinkhorn divergence approximately performs a regularized moment matching scaled by e − x 2 / (2 σ 2 ) .
Equivalence to Regularized MMD Distributional RL.
Based on Proposition 1, we can immediately establish the connection between Sinkhorn divergence and MMD in Corollary 1, indicating that minimizing Sinkhorn divergence between two distributions is equivalent to minimizing a regularized squared MMD.
Corollary 1.
For ε ∈ (0 , + ∞ ) and denote Π ∗ ε as the optimal Π by evaluating the Sinkhorn divergence, it holds that
[Equation]
where we use W c,ε to replace W c,ε ( µ, ν ) for short.
Proof of Corollary 1 is provided in Appendix D.
It is worthy of noting that this equivalence is established for the general case when ε ∈ (0 , + ∞ ) , and it does not hold in the limit cases when ε → 0 or + ∞ .
For example, when ε → + ∞ , the second part including ε in Eq. 12 is not expected to dominate.
This is owing to the fact that the regularization term would be 0 as Π ∗ ε → µ ⊗ ν when ε → + ∞ .
In summary, even though the Sinkhorn divergence was initially proposed to serve as an entropy regularized Wasserterin distance when the cost function c = κ α , it turns out that it is equivalent to a regularized MMD if associated with Gaussian kernels, as revealed in Corollary 1.
4. D ISTRIBUTIONAL RL VIA S INKHORN I TERATIONS
The theoretical analysis in Section 4.2 sheds light on the behavior of distributional RL with Sinkhorn divergence, but another crucial issue we need to address is how to evaluate the Sinkhorn loss effectively.
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 algorithm, i.e.,
Sinkhorn Iterations (Sinkhorn, 1967; Genevay et al., 2018), to evaluate the Sinkhorn loss.
Notably,
Sinkhorn iteration with L steps yields a differentiable and solvable efﬁciently 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 a DQN architecture.
Speciﬁcally, given two sample sequences { Z i } Ni =1 , { T Z j } Nj =1 in the distributional RL algorithm, the optimal transport distance is equivalent to the form:
[Equation]
where the empirical cost function ˆ c
i,j = c ( Z i , T Z j ) .
By adding entropic regularization on optimal transport distance, Sinkhorn divergence can be viewed to restrict the search space of P in the following scaling form:
[Equation] [Equation]
where K i,j = e − ˆ c
i,j /ε is the Gibbs kernel deﬁned in Eq.
This allows us to leverage iterations regarding the vectors a and b .
More speciﬁcally, we initialize b 0 = 1 N , and then the Sinkhorn iterations are expressed as
[Equation]
and
[Equation]
(15) where ·· indicates an entry-wise division.
It has been proven that Sinkhorn iteration asymptotically converges to the true loss in a linear rate (Genevay et al., 2018; Franklin & Lorenz, 1989; Cuturi, 2013; Jason Altschuler, 2017).
We provide a detailed algorithm description of Sinkhorn iterations in Algorithm 2.
With the efﬁcient and differential Sinkhorn iterations, we can easily evaluate the
Sinkhorn divergence and thus let our algorithm enjoy its theoretical advantages.
In practice, we need to choose L and ε , and we conduct a rigorous sensitivity analysis in Section 5.
E XPERIMENTS
We demonstrate the effectiveness of SinkhornDRL as described in Algorithm 1 on the full 55 Atari 2600 games.
Speciﬁcally, we leverage the same architecture as QR-DQN (Dabney et al., 2018b), and replace the quantiles output with N particles, i.e., samples.
In contrast to MMDDRL, SinkhornDRL only changes the distribution divergence from MMD to Sinkhorn divergence, and therefore the potential superiority in 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 RL algorithms as classic baselines, including
QR-DQN (Dabney et al., 2018b) that approximates the Wasserstein distance, C51 (Bellemare et 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 kernels k h ( x, y ) =
exp( − ( x − y ) 2 /h ) with the kernel mixture trick covering a range of bandwidths h , which is same as the basic setting in the original MMDDQN paper (Nguyen et 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.
Hyperparameter settings.
For a fair comparison with QR-DQN, C51 and MMDDRL, we used the same hyperparamters: the number of generated samples N = 200 , Adam optimizer with lr = 0 .
00005 , (cid:15) Adam = 0 .
We used a target network to compute the distributional Bellman target, which ﬁts well in the Neural Fitted Z-Iteration framework.
In addition, we choose number of
Sinkhorn iterations L = 10 and smoothing hyperparameter ε = 10 .
in Section 5.1 as they are not sensitive within a proper interval as demonstrated in Section 5.2.
We choose the unrectiﬁed kernel as the cost function, i.e., − c
=
k α , and select α = 2 in k α in our SinkhornDRL algorithm.
5. P ERFORMANCE OF S INKHORN DRL
Figure 1 illustrates that SinkhornDRL can achieve the competitive performance across 55 Atari games compared with various baseline algorithms with different metrics d p and representation manners on Z θ .
On a large number of games, e.g., Tennis, Seaquest and Atlantis, SinkhornDRL can
    
 7LPH6WHSVH          $ Y H U DJH  5 H W X U Q 7HQQLV '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH          $ Y H U DJH  5 H
W X U Q 6HDTXHVW '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH     
   $ Y H U DJH  5 H W X U Q
H $WODQWLV
[Table]
     7LPH6WHSVH 
     $ Y H U DJH  5 H
W
[Table]
U Q %UHDNRXW '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH      $ Y H U DJH  5 H W X U Q 3RQJ '41&45'4100'6LQNKRUQ 
   
 7LPH6WHSVH  
      $ Y H U DJH  5 H W X U Q 6SDFH,QYDGHUV '41&45'4100'6LQNKRUQ     
 7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 1DPH7KLV*DPH '41&45'4100'6LQNKRUQ 
   
 7LPH6WHSVH         $ Y H U DJH  5 H W X U Q 3KRHQL
[ '41&45'4100'6LQNKRUQ      7LPH6WHSVH  
    $ Y H U DJH  5 H W X U Q :L]DUG2I:RU '41&45'4100'6LQNKRUQ signiﬁcantly outperform other baselines, especially on Tennis where other algorithms even fail to converge.
The improvement of SinkhornDRL over MMDDRL empirically veriﬁes the regularization advantage of the Sinkhorn as analyzed in Corollary 1.
On some games, e.g., Breakout, Pong and
SpaceInvaders, SinkhornDRL is on par with MMDDRL and other baselines, while on the last row in Figure 1, SinkhornDRL is slightly inferior to the state-of-the-art algorithm.
We provide learning curves of all typical distributional RL algorithms on all 55 Atari games in Appendix F, where
SinkhornDRL still achieves the competitive performance in general.
9 H Q W X U H 6 HD TX H V W ' H P RQ $ WW DF N 7 H QQ L V 3 KR H Q L [ = D [[RQ $ WO D Q WL V * RSK H U H QGX U R % H U ]H U N 6 S DFH , QY D G H U V 5 L Y H UU D L G 6 R O D U L V 7 L P H 3 LO R W * U D Y LW D U 6 W D U * XQQ H U
& KRSS H U & R PP D QG $ OL H Q
. XQJ ) X 0 D V W H U SRQJ E U HD NRX W 8 S 1' R Z Q 5 R D G 5 XQQ H U 6 N LL QJ ) U HH Z D \ % R
[ L QJ $ V W H U L [ . U X OO 5 RER W D QN 0 V 3 DF P D Q % HD P 5 L G H U
$ P L G D U < D U V 5 H Y H QJ H % D WWO H = RQ H % D QN + H L V W 1 D P H 7 K L V * D P H ) L V K L QJ ' H U
E\ & H Q WL S H G H 7 X W D QNK D P : L ]D U G 2 I : R U + H U
R
[Table]
U L
[Table]
D W H ( \ H . D QJ
[Table]
U
[Table]
H
[Table]
W , FH
[Table]
\
[Table]
)
[Table]
H
[Table]
D
[Table]
0 RQ W H] X P D
[Table]
H Y H
[Table]
H ' RXE
[Table]
H
[Table]
        5 D WL R  , P
[Table]
RY
[Table]
6LQNKRUQ45'41
9 H Q W X U H $ WO D Q WL V 3 U L Y D W H ( \ H 7 H QQ L V = D [[RQ * RSK H U
[Table]
HD
[Table]
H V W 6 S DFH , QY D G H U V
* U D Y LW D
[Table]
RER W D QN 5 R D G
[Table]
XQQ H U
< D U V
[Table]
H Y H QJ H . XQJ ) X
[Table]
D V W H U ) L V K L QJ ' H U
E\ % D WWO H = RQ H 7 L P H 3 LO R W 5 L Y H UU D L G H QGX U R
. U X OO 0 V 3 DF P D Q - D P H V ERQGSRQJ ) U
HH Z D \ & KRSS H U & R PP D QG % H U ]H U N E U HD NRX W % R[ L QJ 6 W D U * XQQ H U , FH + R F N H \ $ P L G D U $ OL H Q + H U R % D QN + H L V W $ VV D X OW : L ]D U G 2 I : R U 1 D P H 7 K L V * D P H 6 N LL QJ . D QJ D U RR & H Q WL S H G H 3 KR H Q L
[ % HD P 5 L G H U ) U R V
W E LW H 9 L G H R 3 L QE D
OO $ V W H U R L G V $ V W H U L [ 7 X W D QNK D P 6 R O D U L V 8 S 1' R Z Q 4 E H U
W ' H P RQ $ WW DF N
& U D] \ & OL P E H U % R Z OL QJ 0 RQ W H] X P D 5 H Y H QJ H ' RXE O H ' XQN       5 D WL R  , P S U RY P H Q W     6LQNKRUQ00'
 (
[Figure]
We conduct a ratio improvement comparison across 55 Atari games between SinkhornDRL 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 superiority of SinkhornDRL is signiﬁcant across a large amount of games, including Venture, Seaquest, Tennis and Phoenix.
This empirical outperformance veriﬁes 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 a remarkable performance improvement for SinkhornDRL on a large amount of games can be easily observed.
We also report mean and median of best human-normalized scores in Table 2 of Appendix E, where SinkhornDRL achieves almost state-of-the-art performance as MMDDRL on average.
Therefore, we conclude that SinkhornDRL is competitive with the state-of-the-art distributional RL 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 simultaneously makes full use of the data geometry from Wasserstein distance and the unbiased gradient estimate property from MMD, which coincides with results in Theorem 1.
     7LPH6WHSVH      $ Y H U DJH  5 H
W X U Q 6HDTXHVW      45'4100'
[Figure]
ε      7LPH6WHSVH      $ Y H U DJH  5 H W X U Q
[Figure]
     7LPH6WHSVH 
     $ Y H U DJH  5 H W X U Q %UHDNRXW / / / / /  (c)
Sinkhorn Iterations L
5. S ENSITIVITY A NALYSIS AND C OMPUTATIONAL C OST
The limit behavior connection in Theorem1 (1) and (2) between SinkhornDRL with QR-DQN and
MMDDRL may not be rigorously veriﬁed in numerical experiments as an overly large or small ε will lead to numerical instability of Sinkhorn iterations in Algorithm 2, worsening its performance, as shown in Figure 3 (a).
In practice, we choose a proper ε = 10 across all games.
SinkhornDRL also requires a proper number of iterations L and samples N .
For example, a small N , e.g., N =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 G).
We conjecture that using larger networks to represent more samples is more likely to suffer from the overﬁtting issue, yielding the instability in the RL training (Bjorck et al., 2021).
Therefore, we choose N = 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.
D ISCUSSIONS AND C ONCLUSION
To extend our algorithm for better performance, implicit generative models, including parameterizing the cost function in Sinkhorn loss, can be further incorporated.
We leave it as the future work.
Moreover, other divergences, e.g., those that can also smooth Wassertein distance, can also be applied into the design of distributional RL algorithms in the future.
In this paper, a novel family of distributional RL algorithms based on Sinkhorn divergence is proposed that accomplishes a competitive performance compared with the-state-of-the-art distributional
RL algorithms on 55 Atari games.
Theoretical analysis about the convergence and moment matching behavior is provided along with a rigorous empirical veriﬁcation.
Sinkhorn distributional RL will lead to an important contribution among the research community.
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 F.
We also provide the detailed proof from Appendix C to Appendix D.
A R ELATED W ORK
Based on the choice of distribution divergence and the distribution representation manner of Z θ , 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 SinkhornDRL.
Categorical Distributional RL.
As the ﬁrst successful distributional RL family, categorical distributional RL (Bellemare et al., 2017a) represents the value distribution η by the categorical distribution
ˆ η = (cid:80) Ni =1 p i
δ z i , where { z i } Ni =1 ( z 1 < ... <
z N ) is the ﬁxed supports within the pre-speciﬁed interval [ l, u ] and p i 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 i and apply a projected KL divergence between the target and current categorical value distributions.
C51 has also been shown the theoretical contraction under the Cram´er 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) ﬁrstly proposed to use quantile regression to approximate
Wasserstein distance W p .
QR-DQN leverages quantiles to represent the distribution η of Z θ , i.e., ˆ η = (cid:80) Ni
=
δ z i , where { z i } Ni =1 is the learnable support atoms as the quantile values of a ﬁxed quantile { 2 i − 1 2 N } Ni =1 .
Implicit Quantile Networks (IQN) utilizes an implicit model to output quantile values { z i } Ni =1 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., 2019) improves IQN by proposing a more expressive quantile network, achieving better performance on Atari games.
Non-crossing issue in quantile-regression has been raised and addressed properly in (Zhou et al., 2020) that further improves QR-DQN.
The monotonic rational-quadratic splines are also used to learn smooth continuous quantile functions (Luo et al., 2021).
MMD Distributional RL.
MMD distributional RL (MMDDRL) (Nguyen et al., 2020) learns samples to represent the value distribution of Z θ based on maximum mean discrepancy (MMD) in Eq.
20, achieving the state-of-the-art performance on Atari games.
Discussion about SinkhornDRL.
As a complementary Wasserstein distance-based distributional
RL, our SinkhornDRL allows to solve Wasserstein Distance by incorporating the entropic regularization 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 as future works.
Meanwhile, SinkhornDRL is also closely linked with MMDDRL as Sinkhorn divergence interpolates between Wasserstein distance and MMD, and also learns the unrestricted statistics, i.e., samples, akin to MMDDRL.
B D EFINITION OF DISTANCES AND C ONTRACTION
Deﬁnition of distances.
Given two random variables X and Y , p -Wasserstein metric W p between the distributions of X and Y is deﬁned as
[Equation]
which F − 1 is the inverse cumulative distribution function of a random variable with the cumulative distribution function as F .
Further, (cid:96) p distance (Elie & Arthur, 2020) is deﬁned as
[Equation]
The (cid:96) p distance and Wassertein metric are identical at p = 1 , but are otherwise distinct.
Note that when p = 2 , (cid:96) p distance is also called Cram´er distance (Bellemare et al., 2017b)
d C ( X, Y ) .
Also, the Cram´er distance has a different representation given by
[Equation]
where X (cid:48) and Y (cid:48) are the i.i.d.
copies of X and Y .
Energy distance (Sz´ekely, 2003; Ziel, 2020) is a natural extension of Cram´er distance to the multivariate case, which is deﬁned as
[Equation]
[Equation]
where X and Y are multivariate.
Moreover, the energy distance is a special case of the maximum mean discrepancy (MMD), which is formulated as
[Equation]
(20) where k ( · , · ) is a continuous kernel on X .
In particular, if k is a trivial kernel, MMD degenerates to energy distance.
Additionally, we further deﬁne the supreme MMD, which is a functional
P ( X ) S×A × P
[Equation]
X
[Equation]
S×A → R deﬁned as
[Equation]
We further present the convergence rate under different distribution divergences
• T π is γ -contractive under the supreme form of Wassertein distance W p .
• T π is γ 1 /p -contractive under the supreme form of (cid:96) p distance.
• T π is γ α/ 2 -contractive under MMD ∞ with the kernel k α ( x, y )
=
−(cid:107) x − y (cid:107) α , ∀ α > 0 .
Proof of Contraction.
• Contraction under supreme form of Wasserstein diatance is provided in Lemma 3 (Bellemare et al., 2017a).
• Contraction under supreme form of (cid:96) p distance can refer to Theorem 3.4 (Elie & Arthur, 2020).
• Contraction under MMD ∞ is provided in Lemma
(Nguyen et al., 2020).
C P ROOF OF T HEOREM
Proof.
ε → 0 and c = − k α
It is obvious to observe that Sinkhorn loss degenerates to the wasserstein distance.
We also have the conclusion that the distributional Bellman operator T π is γ -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 ε = 0 , for an unspeciﬁed ε , we need a more rigorous proof.
We show that their distance difference is at most an inﬁnitesimal δ .
Firstly, as W c,ε → W α and the regularization term is non-negative, using the language of ( ε, δ ) deﬁnition, we have: for ∀ δ , there exists a small positive constant a , such that W c,ε − W α < δ when (cid:15) ≤ a .
Based on that, we have the contraction conclusion:
[Equation]
where the second term W ∞ α ( T π Z 1 , T π Z 2 ) is contractive, and thus for the unspeciﬁed ε , the only difference from the limting ε = 0 is an inﬁnitesimal δ , which will vanish as ε → 0 or a → 0 . 2.
ε → ∞ .
Our complete proof is inspired by (Ramdas et al., 2017; Genevay et al., 2018).
Recap the deﬁnition of squared MMD is
[Equation]
When the kernel function k degenerates to a unrectiﬁed k α ( x, y ) := −(cid:107)
x − y (cid:107) α for α ∈ (0 , 2) , the squared MMD would degenerate to
[Equation]
On the other hand, we have the Sinkhorn loss as
[Equation]
Denoting Π ε be the unique minimizer for W c,ε , it holds that Π ε → µ ⊗ ν as ε → ∞ .
That being said, W c, ∞ ( µ, ν ) → (cid:82) c ( x, y )d µ ( x )d ν ( y ) + 0 = (cid:82) c ( x, y )d µ ( x )d ν ( y ) .
If c = − k α = (cid:107) x − y (cid:107)
α , we eventually have W − k α , ∞ ( µ, ν ) →
(cid:82) (cid:107) x − y (cid:107) α d µ ( x )d ν ( y ) =
E (cid:107)
X − Y (cid:107) α .
Finally, we can have
[Equation]
α ,
[Equation]
which is exactly the form of squared MMD.
Now the key is prove that Π ε → µ ⊗ ν as ε → ∞ .
Firstly, it is apparent that W c,ε ( µ, ν ) ≤ (cid:82) c
( x, y )d µ ( x )d ν ( y ) as µ ⊗ ν ∈ Π( µ, ν ) .
Let { ε k } be a positive sequence that diverges to ∞ , and Π k be the corresponding sequence of unique minimizers for
W c,ε .
According to the optimality condition, it must be the case that (cid:82) c ( x, y )
dΠ k + ε k KL (Π k , µ ⊗ ν ) ≤ (cid:82) c ( x, y )
d µ ⊗ ν +
(when Π( µ, ν ) = µ ⊗ ν ).
Thus,
[Equation]
Besides, by the compactness of Π( µ, ν ) , we can extract a converging subsequence Π n k → Π ∞ .
Since KL is weakly lower-semicontinuous, it holds that
[Equation]
Hence Π ∞ = µ ⊗ ν .
That being said that the optimal coupling is simply the product of the marginals, indicating that Π ε → µ ⊗ ν as ε → ∞ .
As a special case, when α = 1 , W − k 1 , ∞ ( u, v ) is equivalent to the energy distance
[Equation]
In summary, if the cost function is the rectiﬁed kernel k α , it is the case that W − k α ,ε converges to the squared MMD as ε → ∞ .
According to (Nguyen et al., 2020), T π is γ α/ 2 -contractive in the supreme form of MMD with the rectiﬁed kernel k α .
For the unspeciﬁed ε , we can get the similar result to the case of ε → 0 .
For ∀ δ , there exists a large positive constant M , such that MMD 2 k
α − W c,ε < δ when (cid:15) ≥ M .
Based on that, we have the contraction conclusion:
[Equation]
where the ﬁrst term MMD 2 ∞ ( T π Z 1 , T π Z 2 ) is γ α 2 -contractive, and thus for the unspeciﬁed ε , the only difference from the limiting ε = ∞ is an inﬁnitesimal δ , which will vanish as ε → + ∞ or ( M → + ∞ ).
For ε ∈ (0 , + ∞ ) , the contraction property needs a long proof.
The proof pipeline is ﬁrstly 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 a new non-constant factor.
3.1 Properties of Sinkhorn Divergence.
We recap three crucial properties of a divergence metric.
The ﬁrst is scale sensitive (S) (of order β , β > 0 ), i.e., d p ( cX, cY ) ≤ | c | β d p ( X, Y ) .
The second property is shift invariant (I) , i.e., d p ( A + X, A + Y ) ≤
d p ( X, Y ) .
The last one is unbiased gradient (U) .
A key observation for the analysis is that the Sinkhorn divergence would degenerate to a two-dimensional KL divergence, and therefore embraces a similar convergence behavior to KL divergence.
Concretely, according to the equivalent form of W c,ε ( µ, ν ) in Eq.
9, it can be expressed as the KL divergence between an optimal joint distribution and a Gibbs distribution associated with the cost function:
[Equation]
where Π ∗ is the optimal joint distribution.
Thus, the total Sinkhorn divergence is expressed as
[Equation]
Due to the form of W c,ε ( µ, ν ) , the convergence behavior is determined by W c,ε ( µ, ν ) , which is similar to the behavior of KL divergence.
Thus, we will focus on the convergence analysis of
W c,ε ( µ, ν ) .
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 I can be extended to Sinkhorn divergence.
However, we ﬁnd the non scale sensitive ( S ) property can not directly apply to Sinkhorn divergence due to the minimum nature of
W c,ε ( µ, ν ) as the optimal joint distribution Π ∗ ( µ, ν ) could be different from Π 0 ( aµ, aν ) where a is the scale factor.
We need a new rigorous proof of scale sensitive property as follows.
3.2 Scale Sensitive Property of Sinkhorn Divergence.
We show Sinkhorn divergence satisﬁes a variant of scale sensitive property when c = − k α that corresponds to a non-constant scale factor ∆( a, α ) that is not only a function of the vanilla scale factor a
and α in k α , but also the two speciﬁed probability measures ( U, V ) .
By deﬁnition, the pdf of K ( U, V )
∝ e − c ( x,y ) ε µ ( x ) ν ( y ) .
After a scaling transformation, the pdf of aU and aV with respect to x and y would be 1 a µ ( xa ) and 1 a ν ( ya ) .
Thus K ( aU, aV ) ∝ e − c ( x,y ) ε 1 a µ ( xa ) 1 a ν ( ya ) .
We denote Π ∗ and Π 0 as the optimal joint distribution of W c,ε ( µ, ν ) and W c,ε ( aµ, aν ) .
[Equation]
where ∆ U,
[Equation]
( a,
[Equation]
[Equation]
∈ (0 , 1) for ε ∈ (0 , + ∞ ) and a < 1 due to the fact that 0 < (1 − | a | α ) (cid:82) ( x − y )
α dΠ ∗ ( x, y ) <
(cid:82) ( x − y )
α dΠ ∗ ( x, y ) <
W c,ε ( U, V ) .
∆ U,V ( a, α ) is function less than 1 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 a non-constant factor ∆ U,V ( a, α ) < 1 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 d in order to prove the convergence in 3.4.
Theorem 2.
(Distribution Contraction Mapping Theorem with a Non-constant Factor) Consider a distribution distance d and a function g : P → P .
The mapping d is a contraction: There exists a function q ( X, Y ) <
1 such that for ∀ distributions X and Y :
[Equation]
Then there exists a unique distribution X ∗ with g ( X ∗ ) =
X ∗ .
Proof.
We consider the convergence of the distribution sequence X k .
We have the updating rule as
[Equation]
where we use q k,k − 1 = q ( X k , X k − 1 ) for short.
Hence, we have
[Equation]
[Equation]
Let d 0 = d ( X 1 , X 0 ) .
From the triangle inequality, we have
[Equation]
For the inﬁnite series 1 + q k +1 ,k + ...
+ Π k + l − 1 i
= k +1 q
i,i − 1 + ... , which we denote as u i for i -the term, according to the ratio convergence judgment method of inﬁnite series, lim k →∞ u i +1 u
i < 1 .
Thus, the inﬁnite series is convergent.
Due to the fact Π ki =1 q i,i − 1 → 0 as k → ∞ , we have d ( X k +1 , X k )
→as k
→ ∞ .
Therefore, it must converge to a limit distribution X ∗ that satisﬁes g ( X ∗ ) =
X ∗ .
Contraction of Distributional Bellman Operator under Sinkhorn Divergence.
According to the equation of W c
[Equation]
ε
[Equation]
it holds the same properties as W c,ε , i
[Equation]
e., shift invariant and scale sensitive.
Thus, we derive the convergence of distributional Bellman operator T π under the supreme form of W c,ε , i.e., W ∞ c,ε :
[Equation]
(32) where the ﬁrst inequality comes from the scale sensitive property proof of Sinkhorn divergence and we let ∆ Z 1 ,Z 2 ( γ, α )
=
sup s (cid:48) ,a (cid:48) ∆ Z 1 ( s (cid:48) ,a (cid:48) ) ,Z 2 ( s (cid:48) ,a (cid:48) )
( γ ) .
If ∆ Z 1 ,Z 2 ( γ, α ) is only a constant function in terms of γ and α , we can directly arrive the conclusion that distributional Bellman operator is ∆ Z 1 ,Z 2 ( γ, α ) -contractive based on existing Banach ﬁxed point theorem.
However, the fact is that ∆ Z 1 ,Z 2 ( γ, α ) also depends on Z 1 and Z 2 , and thus we need a new contraction mapping theorem to guarantee the convergence of ﬁxed distribution iteration.
According to Theorem 2 in 3.3 that we speciﬁcally ﬁgure out for the our contraction proof, we have W ∞ c,ε can guarantee the convergence via distributional Bellman iterations.
In summary, we conclude that T π is a contractive operator when we use the − k α as the cost function and γ ≤ 1 , while the contraction factor, which is short for ∆( γ, α ) < 1 , is not only a function of α and γ , but also depends on distribution sequence in the while iterations.
D P ROOF OF P ROPOSITION 1 AND C OROLLARY
Proof.
As we leverage Π ∗ to denote the optimal Π by evaluating the Sinkhorn divergence via min Π ∈ Π ( µ,ν ) W c,ε ( µ, ν ; k ) , the Sinkhorn divergence can be composed in the following form:
[Equation]
where the cost function c in the Gibbs distribution K is minus Gaussian kernel, i.e., c ( x, y )
=
− k
( x, y )
=
e − ( x − y ) / (2 σ 2 ) .
Till now, we have shown the result in Corollary
[Equation]
Next
[Equation]
we use Taylor expansion to prove the moment matching of MMD.
Firstly
[Equation]
we have the following equation:
[Equation]
We expand the Gaussian kernel via Taylor expansion,
i
[Equation]
e.,
[Equation]
Therefore
[Equation]
we have
[Equation]
˜ M n ( µ ) =
E x ∼ µ (cid:104)
e − x 2 / ( 2 σ 2 ) x n (cid:105) , and similarly for ˜ M n ( ν ) .
The conclusion is the same as the moment matching in (Nguyen et al., 2020).
Finally, due to the equivalence of W c,ε ( µ, ν ) after multiplying ε
[Equation]
we have
[Equation]
This result is also equivalent to Theorem 1, where Π ∗ would degenerate to µ ⊗ ν as ε → + ∞ .
In that case, the ﬁrst regularization term would vanish, and thus the Sinkhorn divergence degrades to a
MMD loss, i.e., MMD 2 − c ( µ, ν ) .
E H UMAN - NORMALIZED S CORES
[Table]
Human normalized score equation is (algorithm - randomplay) /
(human - randomplay).
Our implementation 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 implemented by (Zhang, 2018).
The results about mean and median human-normalized scores of all considered distributional RL algorithms are reported in Table 2.
We also compare the performance of QRDQN(tf) and MMDDRL(tf) after the same 10M time steps based on tensorﬂow implementation on Dopamine framework (Castro et al., 2018).
These results are averaged over training data released in (Nguyen-Tang, 2021).
We argue that Human-normalized scores may be limited to evaluate 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 is superior to QR-DQN as Sinkhorn outperforms MMDDRL on a smaller portion of games compared with QRDQN, but the mean score for MMDDRL in Table 2 is lower than QRDQN.
By contrast, SinkhornDRL is superior in terms of mean score and competitive in terms of median.
We also provide all average results in Table 3 of
Appendix H and all learning curves of our implemented algorithms in Appendix F.
F M ORE EXPERIMENTAL R ESULTS
We provide learning curves of DQN, QRDQN, C51, MMD and SinkhornDRL algorithms on all 55 Atari games in Figures 4 5 6 7 8 9.
It illustrates that SinkhornDRL dramatically surpasses the other distributional RL algorithms on a large amount of environments, e.g., Venture, Atlantis, Tennis and SpaceInvader, and presents competitive performance or is only slightly inferior as opposed to the state-of-the-art baselines on other games.
Note that the average improvement of SinkhornDRL on Venture game is signiﬁcant 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 SinkhornDRL algorithm also presents huge potential on some complicated environments.
We leave the further exploration on the advantage and potential of SinkhornDRL algorithm as the future work.
G S ENSITIVITY A NALYSIS AND C OMPUTATIONAL C OST G. M ORE RESULTS IN S ENSITIVITY A NALYSIS
Decreasing ε .
We argue that the limit behavior connection as stated in Theorem 1 may not be able to be veriﬁed rigorously via numeral experiments due to the numerical instability of Sinkhorn
[Table]

[Table]

 $ Y H U DJH  5 H W X U Q 6SDFH,QYDGHUV '41&45'4100'6LQNKRUQ
     7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 8S1'RZQ '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q
[Table]
   
[Table]
   
[Table]
     $ Y H U DJH 
[Table]
H W X U Q
[Table]
   
[Table]
 
[Table]
   $ Y H U DJH 
[Table]
H W X U Q
[Table]
   
[Table]
 $ Y H U DJH
[Table]
H W
[Table]
U
[Table]
     7LPH6WHSVH    
    $ Y H U DJH  5 H W X U Q )LVKLQJ'HUE\
[Table]
     7LPH6WHSVH         $ Y H U DJH  5 H W X U Q )URVWELWH '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 5LYHUUDLG '41&45'4100'6LQNKRUQ
     7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 7LPH3LORW '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 6WDU*XQQHU '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH          $ Y H U DJH  5 H W
[Table]
U Q 6HDTXHVW '41&45'4100'6LQNKRUQ 
   
 7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 1DPH7KLV*DPH '41&45'4100'6LQNKRUQ 
   
 7LPH6WHSVH         $ Y H U DJH  5 H W X U Q 3KRHQL
[ '41&45'4100'6LQNKRUQ     

7LPH6WHSVH          $ Y H U DJH  5 H W X U Q 7HQQLV '41&45'4100'6LQNKRUQ     

7LPH6WHSVH  
    $ Y H U DJH  5 H W X U Q 7XWDQNKDP '41&45'4100'6LQNKRUQ 
   
 7LPH6WHSVH       $ Y H U DJH  5 H
W X U Q 9HQWXUH '41&45'4100'6LQNKRUQ       7LPH6WHSVH         $ Y H U DJH  5 H W X U Q 9LGHR3LQEDOO '41&45'4100'6LQNKRUQ
     7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 5RDG5XQQHU '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q -DPHVERQG '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH          $ Y H U DJH  5 H W X U Q ,FH+RFNH\ '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q +HUR '41&45'4100'6LQNKRUQ

    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q %DQN+HLVW '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH     
   $ Y H U DJH  5 H W X U Q
H $WODQWLV '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q :L]DUG2I:RU '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH        $ Y H U DJH  5 H W X U Q $PLGDU '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH        $ Y H U DJH  5 H W X U Q $VVDXOW '41&45'4100'6LQNKRUQ
     7LPH6WHSVH         $ Y H U DJH  5 H W X U Q %RZOLQJ '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH         $ Y H U DJH  5 H W X U Q %R[LQJ '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 'RXEOH'XQN '41&45'4100' 6LQNKRUQ 
    7LPH6WHSVH         $ Y H U DJH  5 H W X U Q )UHHZD\ '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH     $ Y H U DJH  5 H W X U Q .UXOO '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH      $ Y H U DJH  5 H W X U Q .XQJ)X0DVWHU '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH 

    $ Y H U DJH  5 H W X U Q 0RQWH]XPD5HYHQJH '41&45'4100'6LQNKRUQ
     7LPH6WHSVH        $ Y H U DJH  5 H W X U Q
3LWIDOO '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH          $ Y H U DJH  5 H W X U Q 3ULYDWH(\H '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q
5RERWDQN '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 6NLLQJ '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH       $ Y H U DJH  5 H W X U Q
6RODULV '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH          $ Y H U DJH  5 H W X U Q =D[[RQ '41&45'4100'6LQNKRUQ 
    7LPH6WHSVH        $ Y H U D JH  5 H W X U Q *RSKHU '41&45'4100'6LQNKRUQ     

7LPH6WHSVH      $ Y H U D JH  5 H W X U Q 'HPRQ$WWDFN '41&45'4100'6LQNKRUQ
     7LPH6WHSVH 
      $ Y H U DJH  5 H W X U Q %UHDNRXW           45'4 (a) Small ε in SinkrhornDRL vs QRDQN      7LPH6WHSVH        $ Y H U DJH   H W X U Q %UHDNRXW       00' (b) Large ε in SinkrhornDRL vs MMDDRL 
    7LPH6WHSVH       $ Y H U
DJH   H W
[Figure]
U Q %UHDNRXW         (c) ε on Breakout 
    7LPH6WHSVH      $ Y H U DJH  5 H W X U Q 6HDTXHVW         45'41 00'
[Figure]
Seaquest
Iteration in Algorithm 2.
From Figure 10 (a), we can observe that if we gradually decline ε to 0,
SinkhornDRL’s performance tends to degrade and approach to QR-DQN.
Note that an overly small ε will lead to a trivial almost 0 K i,j in Sinkhorn iteration in Algorithm 2, and will cause 10 numerical instability issue for a l and b l in Line 5 of Algorithm 2.
In addition, we also conducted experiments on Seaquest, the similar result is also observed in Figure 10 (d).
As shown in Figure 10 (d), the performance of SinkhornDRL is robust when ε = 10 , 100 , 500 , but a small (cid:15)
= 1 tends to worsen the performance.
Increasing ε .
Moreover, for breakout, if we increase ε , the performance of SinkhornDRL tends to degrade and be close to MMDDRL as suggested in Figure 10 (b).
It is also noted that an overly large ε will let the K i,j explode to ∞ .
This also leads to numerical instability issue in Sinkhorn iteration in Algorithm 2.
Samples N .
We ﬁnd that SinkhornDRL requires a proper number of samples N to perform favorably, and the sensitivity w.r.t N depends on the environment.
As suggested in Figure 11 (a), a smaller N , e.g., N = 2 on breakout has already achieved favorably performance and even accelerates the convergence in the early phase, while N = 2 on Seaquest will lead to the divergence issue.
Meanwhile, an overly large N worsens the performance across two games.
We conjecture that using larger network networks to generate more samples may suffer from the overﬁtting issue, yielding the training instability (Bjorck et al., 2021).
In practice, we choose a proper number of sample, i.e.,
N = 200 across all games.
     7LPH6WHSVH 
      $ Y H U DJH  5 H W X U Q %UHDNRXW 6DPSOHV  6DPSOHV  6DPSOHV  6DPSOHV  6DPSOHV  6DPSOHV  (a) Samples on Breakout 
    7LPH6WHSVH      $ Y H U
DJH   H W X U Q
6HDTXHVW 6DPSOHV  6DPSOHV  6DPSOHV  6DPSOHV  6DPSOHV  (b) Samples on Seaquest
G. C OMPARISON WITH THE C OMPUTATIONAL C OST
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 12.
It suggests that
SinkhornDRL 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 be tolerant in view of the signiﬁcant outperformance of SinkhornDRL in a large amount of environments.
In addition, we also ﬁnd that the number of Sinkhorn iterations L is negligible to the computation cost, while an overly large samples N , e.g., 500, will lead to a large computational burden as illustrated in Figure 13.
This can be intuitively explained as the computation complexity of the cost function
c
i,j is O ( N 2 ) in SinkhornDRL, which is particularly heavy in computation if N is large enough.
'41 45'4
& 00'6LQNKRUQ          $ Y H U D J H

& R P SX W D WL RQ  7 L P H %UHDNRXW '41 45'41 & 00'6LQNKRUQ          $ Y H U D J H

& R P SX W D WL RQ  7 L P H 4EHUW
[Equation]
For distributional Rl 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 implementation on Pytorch, which is adapted from the reliable open-source implementation (Zhang, 2018).
QRDQN(tf) and MMD(tf) are from training results after 40M frames from the implementation of
MMDDRL (Nguyen et al., 2020) in (Nguyen-Tang, 2021).
[Figure]


    $ Y H U D J H  & R P SX W D WL RQ  7 L P H   V  %UHDNRXW
[Figure]


   
[Figure]
U
[Figure]
[Table]
[Table]
KungFuMaster
[Table]
MontezumaRevenge
[Table]
[Table]
MsPacman
[Table]
Solaris 2346.
792.0 SpaceInvaders 136. 2302.5 StarGunner 10,250.
43820.0 Tennis -23.
13.3 TimePilot 5,229.
7060.0 Tutankham 15.
202.8 UpNDown 604.
20063.0 Venture 0.
1370.0 VideoPinball 15720.
164597.3 WizardOfWor 4,756.
3250.0 YarsRevenge 3271.
[Table]