Original: Distributional Reinforcement Learning via Sinkhorn Iterations
Revised: Distributional Reinforcement Learning via Sinkhorn Iterations

Original: Anonymous Author(s)
Revised: Anonymous Author(s)

Original: AfﬁliationAddress email
Revised: AfﬁliationAddress email

Original: Abstract
Revised: Abstract

Original: Distributional reinforcement learning (RL) is a class of state-of-the-art algorithmsthat estimate the whole distribution of the total return rather than only its expectation.
Revised: Distributional reinforcement learning (RL) is a class of state-of-the-art algorithmsthat estimate the whole distribution of the total return rather than only its expectation.

Original: The representation manner of each return distribution and the choice ofdistribution divergence are pivotal for the empirical success of distributional RL.
Revised: The representation manner of each return distribution and the choice ofdistribution divergence are pivotal for the empirical success of distributional RL.

Original: 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 leverages Sinkhorn iterations to evaluatethe Sinkhorn distance between the current and target Bellman distributions.
Revised: 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 leverages Sinkhorn iterations to evaluatethe Sinkhorn distance between the current and target Bellman distributions.

Original: Remarkably, Sinkhorn divergence interpolates between the Wasserstein distance and
Revised: Remarkably, Sinkhorn divergence interpolates between the Wasserstein distance and

Original: Maximum Mean Discrepancy (MMD).
Revised: Maximum Mean Discrepancy (MMD).

Original: This allows our proposed SinkhornDRLalgorithm to ﬁnd a sweet spot leveraging the geometry of optimal transport baseddistance and the unbiased gradient estimates of MMD.
Revised: This allows our proposed SinkhornDRLalgorithm to ﬁnd a sweet spot leveraging the geometry of optimal transport baseddistance and the unbiased gradient estimates of MMD.

Original: Finally, experiments onthe suit of 55 Atari games reveal the competitive performance of SinkhornDRLalgorithm as opposed to existing state-of-the-art algorithms.
Revised: Finally, experiments onthe suit of 55 Atari games reveal the competitive performance of SinkhornDRLalgorithm as opposed to existing state-of-the-art algorithms.

Original:  Introduction
Revised:  Introduction

Original: Classical reinforcement learning (RL) algorithms are normally based on the expectation of discountedcumulative rewards that an agent observes while interacting with the environment.
Revised: Classical reinforcement learning (RL) algorithms are normally based on the expectation of discountedcumulative rewards that an agent observes while interacting with the environment.

Original: Recently, a newclass of RL algorithms called distributional RL estimates the full distribution of total returns and hasexhibited the state-of-the-art performance in a wide range of environments [2, 8, 7, 24, 26, 17].
Revised: Recently, a newclass of RL algorithms called distributional RL estimates the full distribution of total returns and hasexhibited the state-of-the-art performance in a wide range of environments [2, 8, 7, 24, 26, 17].

Original: From the literature of distributional RL it is easily recognized that algorithms based on either
Revised: From the literature of distributional RL it is easily recognized that algorithms based on either

Original: Wasserstein distance or MMD have gained great attention due to their superior performance.
Revised: Wasserstein distance or MMD have gained great attention due to their superior performance.

Original: As such,their mutual connection from the perspective of mathematical properties intrigues us to explore furtherin order to design new algorithms.
Revised: As such,their mutual connection from the perspective of mathematical properties intrigues us to explore furtherin order to design new algorithms.

Original: Particularly, Wasserstein distance, long known to be a powerfultool to compare probability distributions with non-overlapping supports, has recently emerged as anappealing contender in various machine learning applications.
Revised: Particularly, Wasserstein distance, long known to be a powerfultool to compare probability distributions with non-overlapping supports, has recently emerged as anappealing contender in various machine learning applications.

Original: It is known that Wasserstein distancewas long disregarded because of its computational burden in its original form to solve an expensivenetwork ﬂow problem.
Revised: It is known that Wasserstein distancewas long disregarded because of its computational burden in its original form to solve an expensivenetwork ﬂow problem.

Original: However, recent works [21, 14] have shown that this cost can be largelymitigated by settling for cheaper approximations through strongly convex regularizers.
Revised: However, recent works [21, 14] have shown that this cost can be largelymitigated by settling for cheaper approximations through strongly convex regularizers.

Original: The beneﬁt ofthis regularization has opened the path to wider applications of the Wasserstein distance in relevantlearning problems, including the design of distributional RL algorithms.
Revised: The beneﬁt ofthis regularization has opened the path to wider applications of the Wasserstein distance in relevantlearning problems, including the design of distributional RL algorithms.

Original: The Sinkhorn divergence [21] introduces the entropic regularization on the Wassertein distance,allowing it tractable for the evaluation especially in high-dimensions.
Revised: The Sinkhorn divergence [21] introduces the entropic regularization on the Wassertein distance,allowing it tractable for the evaluation especially in high-dimensions.

Original: It has been successfully appliedin numerous crucial machine learning developments, including the Sinkhorn-GAN [14] and Sinkhornbased adversarial training [23].
Revised: It has been successfully appliedin numerous crucial machine learning developments, including the Sinkhorn-GAN [14] and Sinkhornbased adversarial training [23].

Original: More importantly, it has been shown that Sinkhorn divergenceinterpolates Wasserstein ditance and MMD, and their equivalence form can be well established in thelimit cases [11, 18, 17].
Revised: More importantly, it has been shown that Sinkhorn divergenceinterpolates Wasserstein ditance and MMD, and their equivalence form can be well established in thelimit cases [11, 18, 17].

Original: However, a Sinkhorn-based distributional RL algorithm has not yet been formally proposed and its connection with algorithms based on Wasserstein distance and MMD isalso less studied.
Revised: However, a Sinkhorn-based distributional RL algorithm has not yet been formally proposed and its connection with algorithms based on Wasserstein distance and MMD isalso less studied.

Original: Therefore, a natural question is can we design a new class of distributional RLalgorithms via Sinkhorn divergence, thus bridging the gap between existing two main branches ofdistributional RL algorithms?
Revised: Therefore, a natural question is can we design a new class of distributional RLalgorithms via Sinkhorn divergence, thus bridging the gap between existing two main branches ofdistributional RL algorithms?

Original: Moreover, the dominant quantile-based algorithms, e.g., QR-DQN [8],aimed at approximating Wasserstein distance, suffers from the non-crossing issue in the quantileestimation [26], while sample-based Sinkhorn algorithm can naturally circumvent this problem.
Revised: Moreover, the dominant quantile-based algorithms, e.g., QR-DQN [8],aimed at approximating Wasserstein distance, suffers from the non-crossing issue in the quantileestimation [26], while sample-based Sinkhorn algorithm can naturally circumvent this problem.

Original: In this paper, we propose a novel distributional RL algorithm based on Sinkhorn divergence .
Revised: In this paper, we propose a novel distributional RL algorithm based on Sinkhorn divergence .

Original: Firstly,we point out the key roles of distribution divergence and representation of value distribution in thedesign of distributional RL.
Revised: Firstly,we point out the key roles of distribution divergence and representation of value distribution in thedesign of distributional RL.

Original: After a detailed introduction of our proposed SinkhornDRL algorithm,we theoretically analyze its convergence guarantee and moment matching behavior of distributional
Revised: After a detailed introduction of our proposed SinkhornDRL algorithm,we theoretically analyze its convergence guarantee and moment matching behavior of distributional

Original: Bellman operators under Sinkhorn divergence.
Revised: Bellman operators under Sinkhorn divergence.

Original: Thus, a regularized MMD equivalence form of
Revised: Thus, a regularized MMD equivalence form of

Original: Sinkhorn divergence is derived, interpreting the emprical success of our algorithms in real applications.
Revised: Sinkhorn divergence is derived, interpreting the emprical success of our algorithms in real applications.

Original: Finally, we compare the performance of our SinkhornRL algorithm with typical baselines on 55 Atarigames, verifying the competitive performance of our proposal.
Revised: Finally, we compare the performance of our SinkhornRL algorithm with typical baselines on 55 Atarigames, verifying the competitive performance of our proposal.

Original: Our approach inspires researchersto ﬁnd a trade-off that simultaneously leverages the geometry of the Wasserstein distance and thefavorable unbiased gradient estimate property of MMD while designing new distributional RLalgorithms in the future.
Revised: Our approach inspires researchersto ﬁnd a trade-off that simultaneously leverages the geometry of the Wasserstein distance and thefavorable unbiased gradient estimate property of MMD while designing new distributional RLalgorithms in the future.

Original:  Preliminary Knowledge
Revised:  Preliminary Knowledge

Original: 2.1 Distributional Reinforcement Learning
Revised: 2.1 Distributional Reinforcement Learning

Original: In the classical RL setting, 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.
Revised: In the classical RL setting, 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.

Original: Pis the environment transition dynamics, R is the reward function and γ ∈ (0 , 1) is the discount factor.
Revised: Pis the environment transition dynamics, R is the reward function and γ ∈ (0 , 1) is the discount factor.

Original: From Value function to Value distribution.
Revised: From Value function to Value distribution.

Original: Given a policy π , the discounted sum of futurerewards is a random variable Z π ( s, a ) = (cid:80) ∞ t =

Revised: Given a policy π , the discounted sum of futurerewards is a random variable Z π ( s, a ) = (cid:80) ∞ t =

Original: γ t R ( s t , a t ) , where s 

Revised: γ t R ( s t , a t ) , where s

Original: =
Revised: =

Original: s , a 0 = a , s t +1 ∼
Revised: s , a 0 = a , s t +1 ∼

Original: P ( ·| s t , a t ) , and a t ∼ π ( ·| s t ) .
Revised: P ( ·| s t , a t ) , and a t ∼ π ( ·| s t ) .

Original: In the control setting, expectation-based RL is based on the actionvalue function Q π ( s, a ) , which is the expectation of Z π ( s, a ) , i.e., Q π ( s, a ) =
Revised: In the control setting, expectation-based RL is based on the actionvalue function Q π ( s, a ) , which is the expectation of Z π ( s, a ) , i.e., Q π ( s, a ) =

Original: E [ Z π ( s, a )] .
Revised: E [ Z π ( s, a )] .

Original: Bycontrast, distributional RL focuses on the action-value distribution, the full distribution of Z π ( s, a ) ,and the incorporation of additional distributional knowledge intuitively interprets its empirical success.
Revised: Bycontrast, distributional RL focuses on the action-value distribution, the full distribution of Z π ( s, a ) ,and the incorporation of additional distributional knowledge intuitively interprets its empirical success.

Original: Distributional Bellman operators.
Revised: Distributional Bellman operators.

Original: For the policy evaluation in expectation-based RL, the actionvalue function is updated via the Bellman operator T π Q ( s, a ) =
Revised: For the policy evaluation in expectation-based RL, the actionvalue function is updated via the Bellman operator T π Q ( s, a ) =

Original: E [ R ( s, a )] + γ E s (cid:48) ∼ p,π [ Q ( s (cid:48) , a (cid:48) )] .
Revised: E [ R ( s, a )] + γ E s (cid:48) ∼ p,π [ Q ( s (cid:48) , a (cid:48) )] .

Original: In distributional RL, the action-value distribution of Z π ( s, a ) is updated via the distributional Bellmanoperator T π
Revised: In distributional RL, the action-value distribution of Z π ( s, a ) is updated via the distributional Bellmanoperator T π

Original: [Equation]
Revised: [Equation]

Original: where s (cid:48) ∼ P ( ·| s, a ) and a (cid:48) ∼ π ( ·| s (cid:48) ) .
Revised: where s (cid:48) ∼ P ( ·| s, a ) and a (cid:48) ∼ π ( ·| s (cid:48) ) .

Original: The equality in Eq.
Revised: The equality in Eq.

Original: 1 implies that random variables ofboth sides are equal in distribution.
Revised: 1 implies that random variables ofboth sides are equal in distribution.

Original: The distributional Bellman operator T π is contractive undercertain distribution divergence metrics, but the distributional Bellman optimality operator T can onlyconverge to a set of optimal non-stationary value distributions in a weak sense [9].
Revised: The distributional Bellman operator T π is contractive undercertain distribution divergence metrics, but the distributional Bellman optimality operator T can onlyconverge to a set of optimal non-stationary value distributions in a weak sense [9].

Original: 2.2 Divergences between Measures
Revised: 2.2 Divergences between Measures

Original: Optimal Transport (OT) and Wasserstein Distance
Revised: Optimal Transport (OT) and Wasserstein Distance

Original: The optimal transport (OT) metric betweentwo probability measures ( µ, ν ) supported on two metric spaces is deﬁned as the solution of the linearprogram:
Revised: The optimal transport (OT) metric betweentwo probability measures ( µ, ν ) supported on two metric spaces is deﬁned as the solution of the linearprogram:

Original: [Equation]
Revised: [Equation]

Original: where c is the cost function and Π is the joint distribution with marginals ( µ, ν ) .
Revised: where c is the cost function and Π is the joint distribution with marginals ( µ, ν ) .

Original: Wasserstein distance(a.k.a.
Revised: Wasserstein distance(a.k.a.

Original: earth mover distance) is a special case of optimal transport with the Euclidean norm as thecost function.
Revised: earth mover distance) is a special case of optimal transport with the Euclidean norm as thecost function.

Original: In particular, given two scalar random variables X and Y , p -Wasserstein metric W pbetween the distributions of X and Y can be simpliﬁed as
Revised: In particular, given two scalar random variables X and Y , p -Wasserstein metric W pbetween the distributions of X and Y can be simpliﬁed as

Original: [Equation]
Revised: [Equation]

Original: where F − 1 is the inverse cumulative distribution function of a random variable.
Revised: where F − 1 is the inverse cumulative distribution function of a random variable.

Original: The desirablegeometric property of Wasserstein distance allows it to recover full support of measures, but it suffersfrom the curse of dimension [13, 1].
Revised: The desirablegeometric property of Wasserstein distance allows it to recover full support of measures, but it suffersfrom the curse of dimension [13, 1].

Original: Maximum Mean Discrepancy
Revised: Maximum Mean Discrepancy

Original: The squared Maximum Mean Discrepancy (MMD) MMD 2 k withthe kernel k is formulated as
Revised: The squared Maximum Mean Discrepancy (MMD) MMD 2 k withthe kernel k is formulated as

Original: [Equation]
Revised: [Equation]

Original: where k ( · , · ) is a continuous kernel on X .
Revised: where k ( · , · ) is a continuous kernel on X .

Original: X (cid:48) (resp.
Revised: X (cid:48) (resp.

Original: Y (cid:48) ) is a random variable independent of X(resp.
Revised: Y (cid:48) ) is a random variable independent of X(resp.

Original: Y ).
Revised: Y ).

Original: If k is a trivial kernel, MMD degenerates to the energy distance.
Revised: If k is a trivial kernel, MMD degenerates to the energy distance.

Original: Mathematically, the “ﬂat”geometry that MMD induces on the space of probability measures does not faithfully lift the grounddistance [11], but MMD is cheaper to compute than OT and has a smaller sample complexity, i.e.,approximating the distance with samples of measures [13].
Revised: Mathematically, the “ﬂat”geometry that MMD induces on the space of probability measures does not faithfully lift the grounddistance [11], but MMD is cheaper to compute than OT and has a smaller sample complexity, i.e.,approximating the distance with samples of measures [13].

Original: We provide the detailed introduction ofmore distribution divergences in Appendix A.
Revised: We provide the detailed introduction ofmore distribution divergences in Appendix A.

Original:  Roles of Distribution Divergence and Representation in distributional RL3.1 Distributional
Revised:  Roles of Distribution Divergence and Representation in distributional RL3.1 Distributional

Original: RL: From Neural Q-Fitted Iteration to Neural Z-Fitted Iteration
Revised: RL: From Neural Q-Fitted Iteration to Neural Z-Fitted Iteration

Original: Neural Q-Fitted Iteration.
Revised: Neural Q-Fitted Iteration.

Original: It is known that Deep Q Learning [16] can be simpliﬁed into Neural
Revised: It is known that Deep Q Learning [16] can be simpliﬁed into Neural

Original: Q-Fitted Iteration [10] under tricks of experience replay and the target network Q θ ∗ , where we updateparameterized Q θ ( s, a ) in each iteration k :
Revised: Q-Fitted Iteration [10] under tricks of experience replay and the target network Q θ ∗ , where we updateparameterized Q θ ( s, a ) in each iteration k :

Original: [Equation]
Revised: [Equation]

Original: where the target y i
Revised: where the target y i

Original: =
Revised: =

Original: r ( s i , a i ) + γ max a ∈A Q kθ ∗ ( s (cid:48)
Revised: r ( s i , a i ) + γ max a ∈A Q kθ ∗ ( s (cid:48)

Original: i , a ) is ﬁxed within every T target steps to updatetarget network Q θ ∗ by letting θ ∗ =
Revised: i , a ) is ﬁxed within every T target steps to updatetarget network Q θ ∗ by letting θ ∗ =

Original: θ and the experience buffer induces independent samples{ ( s i , a i , r i , s (cid:48) i ) } i ∈ [ n ] .
Revised: θ and the experience buffer induces independent samples{ ( s i , a i , r i , s (cid:48) i ) } i ∈ [ n ] .

Original: In an ideal case that neglects the non-convexity and TD approximation errors,we have Q k +1 θ =
Revised: In an ideal case that neglects the non-convexity and TD approximation errors,we have Q k +1 θ =

Original: T Q kθ , which is exactly equivalent to updating under Bellman optimality operator.
Revised: T Q kθ , which is exactly equivalent to updating under Bellman optimality operator.

Original: Neural Z-Fitted Iteration.
Revised: Neural Z-Fitted Iteration.

Original: Analogous to neural Q-ﬁtted iteration, we can also simplify value-baseddistributional RL methods based on a parameterized Z θ into a Neural Z-ﬁtted Iteration as
Revised: Analogous to neural Q-ﬁtted iteration, we can also simplify value-baseddistributional RL methods based on a parameterized Z θ into a Neural Z-ﬁtted Iteration as

Original: [Equation]
Revised: [Equation]

Original: where the target Y i
Revised: where the target Y i

Original: =
Revised: =

Original: R ( s i , a i ) + γZ kθ ∗ ( s (cid:48)
Revised: R ( s i , a i ) + γZ kθ ∗ ( s (cid:48)

Original: i , π Z ( s (cid:48) )) with π Z ( s (cid:48) )
Revised: i , π Z ( s (cid:48) )) with π Z ( s (cid:48) )

Original: = argmax a (cid:48) E (cid:2)
Revised: = argmax a (cid:48) E (cid:2)

Original: 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 a divergence metric betweentwo distributions.
Revised: 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 a divergence metric betweentwo distributions.

Original: 3.2 Key Roles of d p and Z θ
Revised: 3.2 Key Roles of d p and Z θ

Original: Within the Neural Z-ﬁtted Iteration framework proposed in Eq.
Revised: Within the Neural Z-ﬁtted Iteration framework proposed in Eq.

Original: 6, we observe that the choice ofrepresentation manner on Z θ and the metric d p are pivotal for the distributional RL algorithms.
Revised: 6, we observe that the choice ofrepresentation manner on Z θ and the metric d p are pivotal for the distributional RL algorithms.

Original: Forinstance, QR-DQN [8] approximates Wasserstein distance W p , which leverages quantiles to represent
Revised: Forinstance, QR-DQN [8] approximates Wasserstein distance W p , which leverages quantiles to represent

Original: [Table]
Revised: [Table]

Original: the distribution of Z θ .
Revised: the distribution of Z θ .

Original: C51 [2] represents Z θ via a categorical distribution under the convergence of
Revised: C51 [2] represents Z θ via a categorical distribution under the convergence of

Original: Cramér distance [3, 19], while MMD distributional RL (MMDDRL) [17] learns samples to representthe distribution of Z θ based on MMD.
Revised: Cramér distance [3, 19], while MMD distributional RL (MMDDRL) [17] learns samples to representthe distribution of Z θ based on MMD.

Original: We compare characteristics of these distribution divergence,including the convergence rate and sample complexity, in Table 1.
Revised: We compare characteristics of these distribution divergence,including the convergence rate and sample complexity, in Table 1.

Original: Theoretical results regarding
Revised: Theoretical results regarding

Original: Sinkhorn divergence is based on [13] and the detailed convergence proof of other distances is alsoprovided in Appendix A.
Revised: Sinkhorn divergence is based on [13] and the detailed convergence proof of other distances is alsoprovided in Appendix A.

Original: In summary, we argue that d p and Z θ are two crucial factors in distributional
Revised: In summary, we argue that d p and Z θ are two crucial factors in distributional

Original: RL design, based on which we introduce our Sinkhorn distributional RL.
Revised: RL design, based on which we introduce our Sinkhorn distributional RL.

Original:  Sinkhorn Distributional RL (SinkhornDRL)
Revised:  Sinkhorn Distributional RL (SinkhornDRL)

Original: In this section, we ﬁrstly introduce Sinkhorn divergence and apply it in distributional RL.
Revised: In this section, we ﬁrstly introduce Sinkhorn divergence and apply it in distributional RL.

Original: Next, weconduct a theoretical analysis about the convergence speed and a new moment matching manner ofour algorithm under the Sinkhorn divergence.
Revised: Next, weconduct a theoretical analysis about the convergence speed and a new moment matching manner ofour algorithm under the Sinkhorn divergence.

Original: Finally, a practical Sinkhorn iteration algorithm isintroduced to evaluate the Sinkhorn divergence.
Revised: Finally, a practical Sinkhorn iteration algorithm isintroduced to evaluate the Sinkhorn divergence.

Original: 4.1 Sinkhorn Divergence and Genetic Algorithm
Revised: 4.1 Sinkhorn Divergence and Genetic Algorithm

Original: We design Sinkhorn distributional RL algorithm via Sinkhorn divergence.
Revised: We design Sinkhorn distributional RL algorithm via Sinkhorn divergence.

Original: Sinkhorn divergence [21] isa tractable loss to approximate the optimal transport problem by leveraging an entropic regularizationto turn the original Wasserstein distance into a differentiable and more robust quantity.
Revised: Sinkhorn divergence [21] isa tractable loss to approximate the optimal transport problem by leveraging an entropic regularizationto turn the original Wasserstein distance into a differentiable and more robust quantity.

Original: The resultingloss can be computed using Sinkhorn ﬁxed point iterations, which is naturally suitable for modern deeplearning frameworks.
Revised: The resultingloss can be computed using Sinkhorn ﬁxed point iterations, which is naturally suitable for modern deeplearning frameworks.

Original: In particular, the entropic smoothing generates a family of losses interpolatingbetween Wasserstein distance and Maximum Mean Discrepancy (MMD).
Revised: In particular, the entropic smoothing generates a family of losses interpolatingbetween Wasserstein distance and Maximum Mean Discrepancy (MMD).

Original: As such, it allows us to ﬁnda 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.
Revised: As such, it allows us to ﬁnda 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.

Original: Weintroduce the entropic regularized Wassertein distance W c,ε ( µ, ν ) as
Revised: Weintroduce the entropic regularized Wassertein distance W c,ε ( µ, ν ) as

Original: [Equation]
Revised: [Equation]

Original: ) where KL (Π | µ ⊗ ν ) = (cid:82) log (cid:16) Π( x,y ) d µ ( x )d ν ( y ) (cid:17) dΠ( x, y ) is a strongly convex regularization.
Revised: ) where KL (Π | µ ⊗ ν ) = (cid:82) log (cid:16) Π( x,y ) d µ ( x )d ν ( y ) (cid:17) dΠ( x, y ) is a strongly convex regularization.

Original: The impactof this entropy regularization is similar to (cid:96) 2 ridge regularization in linear regression.
Revised: The impactof this entropy regularization is similar to (cid:96) 2 ridge regularization in linear regression.

Original: Next, thesinkhorn loss [11, 14] between two measures µ and ν is deﬁned as
Revised: Next, thesinkhorn loss [11, 14] between two measures µ and ν is deﬁned as

Original: [Equation]
Revised: [Equation]

Original: As demonstrated by [11], the Sinkhorn divergence W c,ε ( µ, ν ) is convex, smooth and positive deﬁnitethat metrizes the convergence in law.
Revised: As demonstrated by [11], the Sinkhorn divergence W c,ε ( µ, ν ) is convex, smooth and positive deﬁnitethat metrizes the convergence in law.

Original: In statistical physics, W c,ε ( µ, ν ) can be re-factored as aprojection problem:
Revised: In statistical physics, W c,ε ( µ, ν ) can be re-factored as aprojection problem:

Original: [Equation]
Revised: [Equation]

Original: where K is the Gibbs distribution with the density function satisﬁes d
Revised: where K is the Gibbs distribution with the density function satisﬁes d

Original: [Equation]
Revised: [Equation]

Original: This problem is often referred to as the “static Schrödinger problem” [15, 20] as it was initiallyconsidered in statistical physics.
Revised: This problem is often referred to as the “static Schrödinger problem” [15, 20] as it was initiallyconsidered in statistical physics.

Original: Distributional RL with Sinkhorn Divergence and Particle Representation.
Revised: Distributional RL with Sinkhorn Divergence and Particle Representation.

Original: 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
Revised: 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

Original: Z θ ( s, a ) and the target distribution
Revised: Z θ ( s, a ) and the target distribution

Original: T π Z θ ( s, a ) , yielding W c,ε ( Z θ ( s, a ) , T π
Revised: T π Z θ ( s, a ) , yielding W c,ε ( Z θ ( s, a ) , T π

Original: Z
Revised: Z

Original: θ
Revised: θ

Original: ( s, a )) for each s, a pairs.
Revised: ( s, a )) for each s, a pairs.

Original: In terms of the representationfor Z θ ( s, a ) , we employ the unrestricted statistics, i.e., deterministic samples, due to its superiority in
Revised: In terms of the representationfor Z θ ( s, a ) , we employ the unrestricted statistics, i.e., deterministic samples, due to its superiority in

Original: MMDDRL [17], instead of using predeﬁned statistic functionals, e.g., quantiles in QR-DQN [8] orhistogram partitions in C51 [2].
Revised: MMDDRL [17], instead of using predeﬁned statistic functionals, e.g., quantiles in QR-DQN [8] orhistogram partitions in C51 [2].

Original: More concretely, we use neural networks to generate samples thatapproximate the value distribution.
Revised: More concretely, we use neural networks to generate samples thatapproximate the value distribution.

Original: This can be expressed as Z θ ( s, a ) :
Revised: This can be expressed as Z θ ( s, a ) :

Original: = { Z θ ( s, a ) i } N i =1 , where Nis the number of generated samples.
Revised: = { Z θ ( s, a ) i } N i =1 , where Nis the number of generated samples.

Original: We refer to the samples { Z θ ( s, a ) i } Ni =1 as particles .
Revised: We refer to the samples { Z θ ( s, a ) i } Ni =1 as particles .

Original: Then weleverage the Dirac mixture 1 N (cid:80) Ni
Revised: Then weleverage the Dirac mixture 1 N (cid:80) Ni

Original: =1 δ Z θ
Revised: =1 δ Z θ

Original: ( 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
Revised: ( 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

Original: Bellman target.
Revised: Bellman target.

Original: A detailed and generic distributional RL algorithm with Sinkhorn divergence andparticle representation is provided in Algorithm 1.
Revised: A detailed and generic distributional RL algorithm with Sinkhorn divergence andparticle representation is provided in Algorithm 1.

Original: Require :
Revised: Require :

Original: Number of generated samples N , the cost function c and hyperparameter ε .
Revised: Number of generated samples N , the cost function c and hyperparameter ε .

Original: Input : Sample transition ( s, a, r (cid:48) , s (cid:48) ) 1: if Policy evaluation then 2: a ∗ 
Revised: Input : Sample transition ( s, a, r (cid:48) , s (cid:48) ) 1: if Policy evaluation then 2: a ∗ s

Original: 3: else 4: a ∗ ← arg max a (cid:48) ∈A 1 N (cid:80) Ni =
Revised: 3: else 4: a ∗ ← arg max a (cid:48) ∈A 1 N (cid:80) Ni =

Original: 1 Z θ ( s (cid:48) , a (cid:48) )
Revised: 1 Z θ ( s (cid:48) , a (cid:48) )

Original: i 5: end if 6: T Z i ← r + γZ θ ∗ ( s (cid:48) , a ∗ )
Revised: i 5: end if 6: T Z i ← r + γZ θ ∗ ( s (cid:48) , a ∗ )

Original: i , ∀ 1 ≤ i ≤ N
Revised: i , ∀ 1 ≤ i ≤ N

Original: Output : W c,ε (cid:16) { Z θ ( s, a ) i } Ni =1 , { T Z θ ( s, a ) j } Nj =1 (cid:17)
Revised: Output : W c,ε (cid:16) { Z θ ( s, a ) i } Ni =1 , { T Z θ ( s, a ) j } Nj =1 (cid:17)

Original: Remark.
Revised: Remark.

Original: By comparing the state-of-the-art MMDDRL algorithm [17], our Sinkhorn distributional
Revised: By comparing the state-of-the-art MMDDRL algorithm [17], our Sinkhorn distributional

Original: RL simply modiﬁes the distribution divergence.
Revised: RL simply modiﬁes the distribution divergence.

Original: Hence, we can also easily extend our generic
Revised: Hence, we can also easily extend our generic

Original: Sinkhorn algorithm to DQN-like architecture as well as IQN [7] and FQF [24].
Revised: Sinkhorn algorithm to DQN-like architecture as well as IQN [7] and FQF [24].

Original: A following questionis whether there is any theoretical connection between Sinkhorn distributional RL and algorithmsbased on MMD and Wasserstein distance.
Revised: A following questionis whether there is any theoretical connection between Sinkhorn distributional RL and algorithmsbased on MMD and Wasserstein distance.

Original: We provide this crucial analysis in Section 4.

Revised: We provide this crucial analysis in Section 4.

Original: 4.2 Theoretical Analysis under Sinkhorn Divergence
Revised: 4.2 Theoretical Analysis under Sinkhorn Divergence

Original: Convergence Analysis.
Revised: Convergence Analysis.

Original: Firstly, we denote the supreme form of Sinkhorn divergence as W ∞ c,ε ( µ, ν ) :
Revised: Firstly, we denote the supreme form of Sinkhorn divergence as W ∞ c,ε ( µ, ν ) :

Original: [Equation]
Revised: [Equation]

Original: We will use W ∞ c,ε ( µ, ν ) to establish the convergence ofdistributional Bellman operator T π in Theorem 1.
Revised: We will use W ∞ c,ε ( µ, ν ) to establish the convergence of T π in Theorem 1.

Original: Theorem 1.
Revised: Theorem 1.

Original: If we leverage Sinkhorn loss W c,ε ( µ, ν ) in Eq.
Revised: If we leverage Sinkhorn loss W c,ε ( µ, ν ) in Eq.

Original: 8 as the distribution divergence indistributional RL, it holds that
Revised: 8 as the distribution divergence indistributional RL, it holds that

Original: (1) As ε → 0 , W c,ε ( µ, ν ) → 2{+→
+} W α ( µ, ν ) , and thus T π is a γ -contraction under W ∞ c,ε .
Revised: (1) As ε → 0 as we choose the unrectiﬁed kernel k α := −(cid:107) x − y (cid:107) α as − c ( α > 0 ), W c,ε ( µ, ν ) {+→
+} W α ( µ, ν ) , and thus T π is a γ -contraction under W ∞ c,ε .(2)

Original: (2) As ε → + ∞ , W c,ε ( µ, ν ) 
Revised: As ε → + ∞ , W c,ε ( µ, ν ) →

Original: → MMD 2 − c ( µ, ν ) .
Revised: MMD 2 − c ( µ, ν ) .

Original: If we choose the unrectiﬁed kernel k α := −(cid:107) x − y (cid:107) αas − c
Revised: If we choose the unrectiﬁed kernel k α := −(cid:107) x − y (cid:107) αas − c

Original: ( α > 0 ), T π is a γ α/ 2 -contraction under W ∞ c,ε .
Revised: ( α > 0 ), T π is a γ α/ 2 -contraction under W ∞ c,ε .(3)

Original: (3) For ε ∈ (0 , + ∞ ) , T π is a non-expansive operator under W ∞ c,ε , and the expectation of Z π is still γ -contraction, i.e., (cid:107) E T π Z 1 − E T
Revised: For ε ∈ (0 , + ∞ ) , T π is a non-expansive operator under W ∞ c,ε , and the expectation of Z π is stillγ -contraction, i.e., (cid:107) E T π Z 1 − E T

Original: π Z 2 (cid:107) ∞ ≤ γ (cid:107) E Z 1 − E Z 2 (cid:107) ∞ 
Revised: π Z 2 (cid:107) ∞ ≤ γ (cid:107) E Z 1 − E Z 2 (cid:107) ∞ .

Original: Proof is provided in Appendix B. Theorem 1 (1) and (2) are follow-up conclusions in terms of theconvergence behavior of distributional Bellman operator based on the interpolation relationshipbetween Sinkhorn divergence with Wasserstein distance and MMD [14].
Revised: Proof is provided in Appendix B. Theorem 1 (1) and (2) are follow-up conclusions in terms of theconvergence behavior of distributional Bellman operator based on the interpolation relationshipbetween Sinkhorn divergence with Wasserstein distance and MMD [14].

Original: Note that T π under MMDwith Gaussian kernel is not a contraction [17], but it is γ α/ 2 -contraction when MMD is equippedwith the unrectiﬁed kernel k α .
Revised: Note that T π under MMDwith Gaussian kernel is not a contraction [17], but it is γ α/ 2 -contraction when MMD is equippedwith the unrectiﬁed kernel k α .

Original: Our key theoretical contribution is for the general ε ∈ (0 , ∞ ) , theconvergence behavior is determined by the “joint” KL divergence in Eq.
Revised: Our key theoretical contribution is for the general ε ∈ (0 , ∞ ) , theconvergence behavior is determined by the “joint” KL divergence in Eq.

Original: 9 between the optimal jointdistribution Π ∗ and the Gibbs distribution associated with the cost function c , and therefore T π is atmost a non-expansive operator.
Revised: 9 between the optimal jointdistribution Π ∗ and the Gibbs distribution associated with the cost function c , and therefore T π is atmost a non-expansive operator.

Original: Please refer to Appendix B for details.
Revised: Please refer to Appendix B for details.

Original: Theorem 1 indicates that the convergence behavior of distributional Bellman operator T π under
Revised: Theorem 1 indicates that the convergence behavior of distributional Bellman operator T π under

Original: Sinkhorn divergence interpolates between Wasserstein distance and MMD.
Revised: Sinkhorn divergence interpolates between Wasserstein distance and MMD.

Original: In the general case whenε ∈ (0 , + ∞ ) , we can only obtain that T π is non-expansive in spite of the cost function form in theory,but the expection of Z π remains a γ -contraction and in practice we also empirically demonstrate theappealing empirical performance of our SinkhornDRL algorithm in Section 5.
Revised: In the general case whenε ∈ (0 , + ∞ ) , we can only obtain that T π is non-expansive in spite of the cost function form in theory,but the expection of Z π remains a γ -contraction and in practice we also empirically demonstrate theappealing empirical performance of our SinkhornDRL algorithm in Section 5.

Original: Regularized Moment Matching under Sinkhorn Divergence.
Revised: Regularized Moment Matching under Sinkhorn Divergence.

Original: We further examine the potentialreason behind the empirical success for SinkhornDRL, although only a non-expansive contractioncan be guaranteed for the general case when ε ∈ (0 , + ∞ ) as shown in Theorem 1.
Revised: We further examine the potentialreason behind the empirical success for SinkhornDRL, although only a non-expansive contractioncan be guaranteed for the general case when ε ∈ (0 , + ∞ ) as shown in Theorem 1.

Original: Inspired by thesimilar manner in MMDDRL [17], we ﬁnd that the Sinkhorn divergence with the Gaussian kernelcan also promote to match all moments between two distributions.
Revised: Inspired by thesimilar manner in MMDDRL [17], we ﬁnd that the Sinkhorn divergence with the Gaussian kernelcan also promote to match all moments between two distributions.

Original: More speciﬁcally, the Sinkhorndivergence can be rewritten as a regularized moment matching form in Proposition 1.
Revised: More speciﬁcally, the Sinkhorndivergence can be rewritten as a regularized moment matching form in Proposition 1.

Original: Proposition 1.
Revised: Proposition 1.

Original: For ε ∈ (0 , + ∞ ) , Sinkhorn divergence W c,ε ( µ, ν ) associated with Gaussian kernelsk ( x, y ) = exp
Revised: For ε ∈ (0 , + ∞ ) , Sinkhorn divergence W c,ε ( µ, ν ) associated with Gaussian kernelsk ( x, y ) = exp

Original: [Equation]
Revised: [Equation]

Original: −
Revised: −

Original: [Equation]
Revised: [Equation]

Original: x − y ) 2 /
Revised: x − y ) 2 /

Original: [Equation]
Revised: [Equation]

Original: 2 σ 2 )) as − c , can be equivalent to
Revised: 2 σ 2 )) as − c , can be equivalent to

Original: [Equation]
Revised: [Equation]

Original: where Π ∗ ε denotes the optimal Π determined by ε by evaluating the Sinkhorn divergence via
Revised: where Π ∗ ε denotes the optimal Π determined by ε by evaluating the Sinkhorn divergence via

Original: [Equation]
Revised: [Equation]

Original: , and similarly for ˜ M n ( ν ) .
Revised: , and similarly for ˜ M n ( ν ) .

Original: We provide the proof of Proposition 1 in Appendix C. Similar to MMDDRL associated with a
Revised: We provide the proof of Proposition 1 in Appendix C. Similar to MMDDRL associated with a

Original: Gaussian kernel [17], Sinkhorn divergence approximately performs a regularized moment matchingscaled by e − x 2 / (2 σ 2 ) .
Revised: Gaussian kernel [17], Sinkhorn divergence approximately performs a regularized moment matchingscaled by e − x 2 / (2 σ 2 ) .

Original: This similar moment matching impact intuitively explains the empirical successof SinkhornDRL as MMDDRL, although the contraction of both MMD with Gaussian kernel [17]and Sinkhorn divergence for general (cid:15) ∈ (0 , + ∞ ) may not be guaranteed.
Revised: This similar moment matching impact intuitively explains the empirical successof SinkhornDRL as MMDDRL, although the contraction of both MMD with Gaussian kernel [17]and Sinkhorn divergence for general (cid:15) ∈ (0 , + ∞ ) may not be guaranteed.

Original: Equivalence to Regularized MMD distributional RL.
Revised: Equivalence to Regularized MMD distributional RL.

Original: Based on Proposition 1, we can immediately establish the connection between Sinkhorn divergence and MMD in Corollary 1, indicating thatminimizing Sinkhorn divergence between two distributions is equivalent to minimizing a regularizedsquared MMD.
Revised: Based on Proposition 1, we can immediately establish the connection between Sinkhorn divergence and MMD in Corollary 1, indicating thatminimizing Sinkhorn divergence between two distributions is equivalent to minimizing a regularizedsquared MMD.

Original: Corollary 1.
Revised: Corollary 1.

Original: For ε ∈ (0 , + ∞ ) and denote Π ∗ ε as the optimal Π by evaluating the Sinkhorn divergence,it holds that
Revised: For ε ∈ (0 , + ∞ ) and denote Π ∗ ε as the optimal Π by evaluating the Sinkhorn divergence,it holds that

Original: [Equation]
Revised: [Equation]

Original: where we use W c,ε to replace W c,ε ( µ, ν ) for short.
Revised: where we use W c,ε to replace W c,ε ( µ, ν ) for short.

Original: Proof of Corollary 1 is provided in Appendix C. It is worthy of noting that this equivalence isestablished for the general case when ε ∈ (0 , + ∞ ) , and it does not hold in the limit cases whenε → 0 or + ∞ .
Revised: Proof of Corollary 1 is provided in Appendix C. It is worthy of noting that this equivalence isestablished for the general case when ε ∈ (0 , + ∞ ) , and it does not hold in the limit cases whenε → 0 or + ∞ .

Original: For example, when ε → + ∞ , the second part including ε in Eq. 12 is not expected todominate.
Revised: For example, when ε → + ∞ , the second part including ε in Eq. 12 is not expected todominate.

Original: This is owing to the fact that the regularization term would be 0 as Π ∗ ε → µ ⊗ ν whenε → + ∞ .
Revised: This is owing to the fact that the regularization term would be 0 as Π ∗ ε → µ ⊗ ν whenε → + ∞ .

Original: In summary, even though the Sinkhorn divergence was initially proposed to serve as anentropy regularized Wasserterin distance, it turns out that it is equivalent to a regularized MMD, asrevealed in Corollary 1.
Revised: In summary, even though the Sinkhorn divergence was initially proposed to serve as anentropy regularized Wasserterin distance, it turns out that it is equivalent to a regularized MMD, asrevealed in Corollary 1.

Original: This connection provides strong evidence for our empirical results, in which
Revised: This connection provides strong evidence for our empirical results, in which

Original: SinkhornDRL achieves competitive performance as opposed to MMDDRL.
Revised: SinkhornDRL achieves competitive performance as opposed to MMDDRL.

Original: 4.3 Distributional RL via Sinkhorn Iterations
Revised: 4.3 Distributional RL via Sinkhorn Iterations

Original: The theoretical analysis in Section 4.2 sheds light on the behavior of distributional RL with Sinkhorndivergence, but another crucial issue we need to address is how to evaluate the Sinkhorn losseffectively.
Revised: The theoretical analysis in Section 4.2 sheds light on the behavior of distributional RL with Sinkhorndivergence, but another crucial issue we need to address is how to evaluate the Sinkhorn losseffectively.

Original: Due to the advantages of Sinkhorn divergence that both enjoys geometry property ofoptimal transport and the computational effectiveness of MMD, we can utilize Sinkhorn’s algorithm,i.e., Sinkhorn Iterations [21, 14], to evaluate the Sinkhorn loss.
Revised: Due to the advantages of Sinkhorn divergence that both enjoys geometry property ofoptimal transport and the computational effectiveness of MMD, we can utilize Sinkhorn’s algorithm,i.e., Sinkhorn Iterations [21, 14], to evaluate the Sinkhorn loss.

Original: Notably, Sinkhorn iteration with
Revised: Notably, Sinkhorn iteration with

Original: L steps yields a differentiable and solvable efﬁciently loss function as the main burden involvedin it is the matrix-vector multiplication, which streams well on the GPU with simply adding extradifferentiable layers on the typical deep neural network, such as a DQN architecture.
Revised: L steps yields a differentiable and solvable efﬁciently loss function as the main burden involvedin it is the matrix-vector multiplication, which streams well on the GPU with simply adding extradifferentiable layers on the typical deep neural network, such as a DQN architecture.

Original: Speciﬁcally, given two sample sequences { Z i } Ni =1 , { T Z j } Nj =1 in the distributional RL algorithm, theoptimal transport distance is equivalent to the form:
Revised: Speciﬁcally, given two sample sequences { Z i } Ni =1 , { T Z j } Nj =1 in the distributional RL algorithm, theoptimal transport distance is equivalent to the form:

Original: [Equation]
Revised: [Equation]

Original: where the empirical cost function ˆ c
Revised: where the empirical cost function ˆ c

Original: i,j = c ( Z i , T Z j ) .
Revised: i,j = c ( Z i , T Z j ) .

Original: By adding entropic regularization on optimaltransport distance, Sinkhorn divergence can be viewed to restrict the search space of P in thefollowing scaling form:
Revised: By adding entropic regularization on optimaltransport distance, Sinkhorn divergence can be viewed to restrict the search space of P in thefollowing scaling form:

Original: [Equation]
Revised: [Equation]

Original: where K i,j = e − ˆ c
Revised: where K i,j = e − ˆ c

Original: i,j /ε is the Gibbs kernel deﬁned in Eq.
Revised: i,j /ε is the Gibbs kernel deﬁned in Eq.

Original: This allows us to leverage iterationsregarding the vectors a and b .
Revised: This allows us to leverage iterationsregarding the vectors a and b .

Original: More speciﬁcally, we initialize b 0 = 1 N , and then the Sinkhorniterations are expressed as
Revised: More speciﬁcally, we initialize b 0 = 1 N , and then the Sinkhorniterations are expressed as

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: where ·· indicates an entry-wise division.
Revised: where ·· indicates an entry-wise division.

Original: It has been proven that Sinkhorn iteration asymptoticallyconverges to the true loss in a linear rate [14, 12, 6].
Revised: It has been proven that Sinkhorn iteration asymptoticallyconverges to the true loss in a linear rate [14, 12, 6].

Original: We provide a detailed algorithm description of
Revised: We provide a detailed algorithm description of

Original: Sinkhorn iterations in Algorithm 2.
Revised: Sinkhorn iterations in Algorithm 2.

Original: With the efﬁcient and differential Sinkhorn iterations, we caneasily evaluate the Sinkhorn divergence and thus let our algorithm enjoy its theoretical advantages.
Revised: With the efﬁcient and differential Sinkhorn iterations, we caneasily evaluate the Sinkhorn divergence and thus let our algorithm enjoy its theoretical advantages.

Original: Inpractice, we need to choose L and ε , and we conduct a rigorous sensitivity analysis in Section 5.
Revised: Inpractice, we need to choose L and ε , and we conduct a rigorous sensitivity analysis in Section 5.

Original:  Experiments
Revised:  Experiments

Original: We demonstrate the effectiveness of SinkhornDRL as described in Algorithm 1 on the full 55 Atari2600 games.
Revised: We demonstrate the effectiveness of SinkhornDRL as described in Algorithm 1 on the full 55 Atari2600 games.

Original: Speciﬁcally, we leverage the same architecture as QR-DQN [8], and replace the quantilesoutput with N particles, i.e., samples.
Revised: Speciﬁcally, we leverage the same architecture as QR-DQN [8], and replace the quantilesoutput with N particles, i.e., samples.

Original: In contrast to MMDDRL, SinkhornDRL only changes thedistribution divergence from MMD to Sinkhorn divergence, and therefore the potential superiority inthe performance can be attributed to the advantages of Sinkhorn divergence.
Revised: In contrast to MMDDRL, SinkhornDRL only changes thedistribution divergence from MMD to Sinkhorn divergence, and therefore the potential superiority inthe performance can be attributed to the advantages of Sinkhorn divergence.

Original: In Section 5.1, we makea rigorous comparison between SinkhornDRL with other typical distributional RL algorithms fromthe perspectives of learning curves and ﬁnal ratio improvement of returns.
Revised: In Section 5.1, we makea rigorous comparison between SinkhornDRL with other typical distributional RL algorithms fromthe perspectives of learning curves and ﬁnal ratio improvement of returns.

Original: An extensive sensitivityanalysis in terms of multiple hyperparameters in SinkhornDRL is provided in Section 5.2.
Revised: An extensive sensitivityanalysis in terms of multiple hyperparameters in SinkhornDRL is provided in Section 5.2.

Original: Baselines.
Revised: Baselines.

Original: Due to the interpolation characteristic of Sinkhorn divergence between Wasserteindistance and MMDDRL, we choose three typical distributional RL algorithms as classic baselines,including QR-DQN [8] that approximates the Wasserstein distance, C51 [2] and MMDDRL [17], aswell as DQN [16].
Revised: Due to the interpolation characteristic of Sinkhorn divergence between Wasserteindistance and MMDDRL, we choose three typical distributional RL algorithms as classic baselines,including QR-DQN [8] that approximates the Wasserstein distance, C51 [2] and MMDDRL [17], aswell as DQN [16].

Original: MMDDRL algorithm is implemented with the same architecture as QRDQN, andleverages Gaussian kernels k h ( x, y )
Revised: MMDDRL algorithm is implemented with the same architecture as QRDQN, andleverages Gaussian kernels k h ( x, y )

Original: =
Revised: =

Original: exp( − ( x − y ) 2 /h ) with the kernel mixture trick covering arange of bandwidths h , which is same as the basic setting in the original MMDDQN paper [17].
Revised: exp( − ( x − y ) 2 /h ) with the kernel mixture trick covering arange of bandwidths h , which is same as the basic setting in the original MMDDQN paper [17].

Original: Wedeploy all algorithms on 55 Atari 2600 games, and reported results are averaged over 3 seeds withthe shade indicating the standard deviation.
Revised: Wedeploy all algorithms on 55 Atari 2600 games, and reported results are averaged over 3 seeds withthe shade indicating the standard deviation.

Original: Hyperparameter settings.
Revised: Hyperparameter settings.

Original: For a fair comparison with QR-DQN, C51 and MMDDRL, we usedthe same hyperparamters: the number of generated samples N = 200 , Adam optimizer withlr = 0 .
Revised: For a fair comparison with QR-DQN, C51 and MMDDRL, we usedthe same hyperparamters: the number of generated samples N = 200 , Adam optimizer withlr = 0 .

Original: 00005 , (cid:15)
Revised: 00005 , (cid:15)

Original: Adam = 0 .
Revised: Adam = 0 .

Original: We used a target network to compute the distributional Bellmantarget, which ﬁts well in the neural Z-ﬁtted iteration framework.
Revised: We used a target network to compute the distributional Bellmantarget, which ﬁts well in the neural Z-ﬁtted iteration framework.

Original: In addition, we choose number of
Revised: In addition, we choose number of

Original: Sinkhorn iterations L = 10 and smoothing hyperparameter ε = 10 . 

Revised: Sinkhorn iterations L = 10 and smoothing hyperparameter ε = 10 .

Original: in Section 5.1 as they are notsensitive within a proper interval as demonstrated in Section 5.2.
Revised: in Section 5.1 as they are notsensitive within a proper interval as demonstrated in Section 5.2.

Original: We choose the unrectiﬁed kernel asthe cost function, i.e., − c = k α , and select α = 2 in k α in our SinkhornDRL algorithm.
Revised: We choose the unrectiﬁed kernel asthe cost function, i.e., − c = k α , and select α = 2 in k α in our SinkhornDRL algorithm.

Original: 5.1 Performance of SinkhornDRL
Revised: 5.1 Performance of SinkhornDRL

Original: Figure 1 illustrates that SinkhornDRL can achieve the competitive performance across 55 Atari gamescompared with various baseline algorithms with different metrics d p and representation manners on
Revised: Figure 1 illustrates that SinkhornDRL can achieve the competitive performance across 55 Atari gamescompared with various baseline algorithms with different metrics d p and representation manners on

Original: Z θ .
Revised: Z θ .

Original: On a large number of games, e.g., Tennis, Seaquest and Atlantis, SinkhornDRL can signiﬁcantlyoutperform other baselines, especially on Tennis where other algorithms even fail to converge.
Revised: On a large number of games, e.g., Tennis, Seaquest and Atlantis, SinkhornDRL can signiﬁcantlyoutperform other baselines, especially on Tennis where other algorithms even fail to converge.

Original: Theimprovement of SinkhornDRL over MMDDRL empirically veriﬁes the regularization advantage ofthe Sinkhorn as analyzed in Corollary 1.
Revised: Theimprovement of SinkhornDRL over MMDDRL empirically veriﬁes the regularization advantage ofthe Sinkhorn as analyzed in Corollary 1.

Original: On some games, e.g., Breakout, Pong and SpaceInvaders,
Revised: On some games, e.g., Breakout, Pong and SpaceInvaders,

Original:     
Revised:     

Original:  7LPH6WHSVH          $ Y H U DJH  5 H W X U Q 7HQQLV '41&45'4100'6LQNKRUQ 
Revised:  7LPH6WHSVH          $ Y H U DJH  5 H W X U Q 7HQQLV '41&45'4100'6LQNKRUQ 

Original:     7LPH6WHSVH          $ Y H U DJH  5 H
Revised:     7LPH6WHSVH          $ Y H U DJH  5 H

Original: W X U Q 6HDTXHVW '41&45'4100'6LQNKRUQ 
Revised: W X U Q 6HDTXHVW '41&45'4100'6LQNKRUQ 

Original:     7LPH6WHSVH     
Revised:     7LPH6WHSVH     

Original:    $ Y H U DJH  5 H W X U Q
Revised:    $ Y H U DJH  5 H W X U Q

Original: H $WODQWLV '41&45'4100'6LQNKRUQ 
Revised: H $WODQWLV '41&45'4100'6LQNKRUQ 

Original:     7LPH6WHSVH       $ Y H U DJH  5 H W X U Q %UHDNRXW '41&45'4100'6LQNKRUQ 
Revised:     7LPH6WHSVH       $ Y H U DJH  5 H W X U Q %UHDNRXW '41&45'4100'6LQNKRUQ 

Original:     7LPH6WHSVH      $ Y H U DJH  5 H W X U Q 3RQJ '41&45'4100'6LQNKRUQ 
Revised:     7LPH6WHSVH      $ Y H U DJH  5 H W X U Q 3RQJ '41&45'4100'6LQNKRUQ 

Original:    
Revised:    

Original:  7LPH6WHSVH  
Revised:  7LPH6WHSVH  

Original:       $ Y H U DJH  5 H W X U Q 6SDFH,QYDGHUV '41&45'4100'6LQNKRUQ     
Revised:       $ Y H U DJH  5 H W X U Q 6SDFH,QYDGHUV '41&45'4100'6LQNKRUQ     

Original:  7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 1DPH7KLV*DPH '41&45'4100'6LQNKRUQ 
Revised:  7LPH6WHSVH       $ Y H U DJH  5 H W X U Q 1DPH7KLV*DPH '41&45'4100'6LQNKRUQ 

Original:    
Revised:    

Original:  7LPH6WHSVH         $ Y H U DJH  5 H W X U Q 3KRHQL
Revised:  7LPH6WHSVH         $ Y H U DJH  5 H W X U Q 3KRHQL

Original: [ '41&45'4100'6LQNKRUQ      7LPH6WHSVH  
Revised: [ '41&45'4100'6LQNKRUQ      7LPH6WHSVH  

Original:     $ Y H U DJH  5 H W X U Q :L]DUG2I:RU '41&45'4100'6LQNKRUQ
Revised:     $ Y H U DJH  5 H W X U Q :L]DUG2I:RU '41&45'4100'6LQNKRUQ

Original: SinkhornDRL is on par with MMDDRL and other baselines, while on the last row in Figure 1,
Revised: SinkhornDRL is on par with MMDDRL and other baselines, while on the last row in Figure 1,

Original: SinkhornDRL is slightly inferior to the state-of-the-art algorithm.
Revised: SinkhornDRL is slightly inferior to the state-of-the-art algorithm.

Original: We provide learning curves of alltypical distributional RL algorithms on all 55 Atari games in Appendix E, where SinkhornDRL stillachieves the competitive performance in general.
Revised: We provide learning curves of alltypical distributional RL algorithms on all 55 Atari games in Appendix E, where SinkhornDRL stillachieves the competitive performance in general.

Original: To further demonstrate theoretical properties of SinkhornDRL in Theorem 1, we conduct a ratio improvement comparison across 55 Atari games between SinkhornDRL with QRDQN and MMDDRL,respectively.
Revised: To further demonstrate theoretical properties of SinkhornDRL in Theorem 1, we conduct a ratio improvement comparison across 55 Atari games between SinkhornDRL with QRDQN and MMDDRL,respectively.

Original: Figure 2 showcases that by comparing with QRDQN (left), SinkhornDRL achievesbetter performance across more than half of considered games.
Revised: Figure 2 showcases that by comparing with QRDQN (left), SinkhornDRL achievesbetter performance across more than half of considered games.

Original: More importantly, the superiorityof SinkhornDRL is signiﬁcant across a large amount of games, including Venture, Seaquest, Tennisand Phoenix.
Revised: More importantly, the superiorityof SinkhornDRL is signiﬁcant across a large amount of games, including Venture, Seaquest, Tennisand Phoenix.

Original: This empirical outperformance veriﬁes the effectiveness and potential of smoothing
Revised: This empirical outperformance veriﬁes the effectiveness and potential of smoothing

Original: Wassertein distance in distributional RL, e.g., Sinkhorn divergence.
Revised: Wassertein distance in distributional RL, e.g., Sinkhorn divergence.

Original: In contrast with MMDDRL, thesuperiority of SinkhornDRL is reduced with the performance improvement only on a small proportionof games, while a remarkable boost of performance for SinkhornDRL on a large amount of gamescan be easily observed.
Revised: In contrast with MMDDRL, thesuperiority of SinkhornDRL is reduced with the performance improvement only on a small proportionof games, while a remarkable boost of performance for SinkhornDRL on a large amount of gamescan be easily observed.

Original: We also report mean and median of best human-normalized scores in Table 2of Appendix D, where SinkhornDRL achieves almost state-of-the-art performance as MMDDRL onaverage.
Revised: We also report mean and median of best human-normalized scores in Table 2of Appendix D, where SinkhornDRL achieves almost state-of-the-art performance as MMDDRL onaverage.

Original: Therefore, we conclude that SinkhornDRL is competitive with the state-of-the-art distributional
Revised: Therefore, we conclude that SinkhornDRL is competitive with the state-of-the-art distributional

Original: RL algorithms, e.g., MMDDRL, and can be extremely superior over existing algorithms on a largeproportion of games.
Revised: RL algorithms, e.g., MMDDRL, and can be extremely superior over existing algorithms on a largeproportion of games.

Original: This empirical success can be owing to theoretical advantage of Sinkhorndivergence that simultaneously makes full use of the data geometry from Wasserstein distance andthe unbiased gradient estimate property from MMD, which coincides with results in Theorem 1.
Revised: This empirical success can be owing to theoretical advantage of Sinkhorndivergence that simultaneously makes full use of the data geometry from Wasserstein distance andthe unbiased gradient estimate property from MMD, which coincides with results in Theorem 1.

Original: 5.2 Sensitivity Analysis and Computational Cost
Revised: 5.2 Sensitivity Analysis and Computational Cost

Original: Figure 3 (a) suggests the performance of our algorithm is robust to ε in a certain range, e.g., [1 , 500] ,facilitating its deployment in practice.
Revised: Figure 3 (a) suggests the performance of our algorithm is robust to ε in a certain range, e.g., [1 , 500] ,facilitating its deployment in practice.

Original: If we increase ε , SinkhornDRL’s performance tends to MMD,
Revised: If we increase ε , SinkhornDRL’s performance tends to MMD,

Original: [Table]
Revised: [Table]

Original:         5 D WL R  , P S U RY P H Q W     6LQNKRUQ45'41
Revised:         5 D WL R  , P S U RY P H Q W     6LQNKRUQ45'41

Original: 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 6 HD TX H V W 6 S DFH , QY D G H U V
Revised: 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 6 HD TX H V W 6 S DFH , QY D G H U V

Original: * U D Y LW D U 5 RER W D QN 5 R D G 5 XQQ H U
Revised: * U D Y LW D U 5 RER W D QN 5 R D G 5 XQQ H U

Original: < D U V 5 H Y H QJ H . XQJ ) X 0 D V W H U ) L V K L QJ ' H U
Revised: < D U V 5 H Y H QJ H . XQJ ) X 0 D V W H U ) L V K L QJ ' H U

Original: 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
Revised: 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

Original: . U X OO 0 V 3 DF P D Q - D P H V ERQGSRQJ ) U
Revised: . U X OO 0 V 3 DF P D Q - D P H V ERQGSRQJ ) U

Original: 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
Revised: 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

Original: [ % HD P 5 L G H U ) U R V
Revised: [ % HD P 5 L G H U ) U R V

Original: W E LW H 9 L G H R 3 L QE D
Revised: W E LW H 9 L G H R 3 L QE D

Original: 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
Revised: 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

Original: W ' H P RQ $ WW DF N
Revised: W ' H P RQ $ WW DF N

Original: & U D] \ & OL P E H U % R Z OL QJ 0 RQ W H] X P D  H Y H QJ H ' RXE O H ' XQN       5 D WL R  , P
Revised: & U D] \ & OL P E H U % R Z OL QJ 0 RQ W H] X P D  H Y H QJ H ' RXE O H ' XQN       5 D WL R  , P

Original: [Figure]
Revised: [Figure]

Original: while if we gradually decline ε , SinkhornDRL’s performance tends to QR-DQN.
Revised: while if we gradually decline ε , SinkhornDRL’s performance tends to QR-DQN.

Original: It is also noted that
Revised: It is also noted that

Original: Sinkhorn iterations in Algorithm 2 will suffer from the numerical instability issue under an overlysmall or large ε .
Revised: Sinkhorn iterations in Algorithm 2 will suffer from the numerical instability issue under an overlysmall or large ε .

Original: More results with the discussion are provided in Appendix F.
Revised: More results with the discussion are provided in Appendix F.

Original: It is also illustratedthat our algorithm is insensitive to the number of iterations L and samples N as well, but an overlylarge N can slightly worsen the performance of SinkhornDRL, and at the same time increases thecomputational burden.
Revised: It is also illustratedthat our algorithm is insensitive to the number of iterations L and samples N as well, but an overlylarge N can slightly worsen the performance of SinkhornDRL, and at the same time increases thecomputational burden.

Original: Therefore, a proper number of samples, e.g., 200, is sufﬁcient to attain anappealing performance with the computational effectiveness.
Revised: Therefore, a proper number of samples, e.g., 200, is sufﬁcient to attain anappealing performance with the computational effectiveness.

Original: For the computation cost, SinkhronDRL indeed increases around 50 % computation cost comparedwith QR-DQN and C51, but only slightly increases the cost (by around 20 % ) in contrast to MMDDRL.
Revised: For the computation cost, SinkhronDRL indeed increases around 50 % computation cost comparedwith QR-DQN and C51, but only slightly increases the cost (by around 20 % ) in contrast to MMDDRL.

Original: Detailed comparison is given in Appendix F.   
Revised: Detailed comparison is given in Appendix F.   

Original:  7LPH6WHSVH       $ Y H U DJH  5 H W X U Q %UHDNRXW      (a) Hyper-parameter ε      7LPH6WHSVH        $ Y H U DJH  5 H W X U Q %UHDNRXW 6DPSOHV 6DPSOHV 6DPSOHV 6DPSOHV  (
Revised:  7LPH6WHSVH       $ Y H U DJH  5 H W X U Q %UHDNRXW      (a) Hyper-parameter ε      7LPH6WHSVH        $ Y H U DJH  5 H W X U Q %UHDNRXW 6DPSOHV 6DPSOHV 6DPSOHV 6DPSOHV  (

Original: [Figure]
Revised: [Figure]

Original: Samples 
Revised: Samples 

Original:     7LPH6WHSVH       $ Y H U DJH  5 H W X U Q %UHDNRXW / / / / / 
Revised:     7LPH6WHSVH       $ Y H U DJH  5 H W X U Q %UHDNRXW / / / / / 

Original: [Figure]
Revised: [Figure]

Original:  Discussions and Conclusion
Revised:  Discussions and Conclusion

Original: [Figure]
Revised: [Figure]

Original: The main limitation of our proposal is that the superiority over existing state-of-the-art algorithms maynot be sufﬁciently signiﬁcant.
Revised: The main limitation of our proposal is that the superiority over existing state-of-the-art algorithms maynot be sufﬁciently signiﬁcant.

Original: To extend our algorithm for better performance, implicit generativemodels, including parameterizing the cost function in Sinkhorn loss, can be further incorporated.
Revised: To extend our algorithm for better performance, implicit generativemodels, including parameterizing the cost function in Sinkhorn loss, can be further incorporated.

Original: Weleave it as the future work.
Revised: Weleave it as the future work.

Original: Moreover, other divergences, e.g., thoses that can also smooth Wasserteindistance, can also be applied into the design of distributional RL algorithms in the future.
Revised: Moreover, other divergences, e.g., thoses that can also smooth Wasserteindistance, can also be applied into the design of distributional RL algorithms in the future.

Original: In this paper, a novel family of distributional RL algorithms based on Sinkhorn Divergence is proposedthat accomplishes a competitive performance compared with the-state-of-the-art distributional RLalgorithms on 55 Atari games.
Revised: In this paper, a novel family of distributional RL algorithms based on Sinkhorn Divergence is proposedthat accomplishes a competitive performance compared with the-state-of-the-art distributional RLalgorithms on 55 Atari games.

Original: Theoretical analysis about the convergence and moment matchingbehavior is provided along with a rigorous empirical veriﬁcation.
Revised: Theoretical analysis about the convergence and moment matchingbehavior is provided along with a rigorous empirical veriﬁcation.

Original: Albeit being associated with MMDalgorithm, distributional RL with Sinkhorn divergence is complementary to previous algorithms,leading to an important contribution among the research community.
Revised: Albeit being associated with MMDalgorithm, distributional RL with Sinkhorn divergence is complementary to previous algorithms,leading to an important contribution among the research community.

Original: Checklist
Revised: Checklist

Original: For all authors...
Revised: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s
Revised: (a) Do the main claims made in the abstract and introduction accurately reﬂect the paper’s

Original: contributions and scope?
Revised: contributions and scope?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you describe the limitations of your work?
Revised: (b) Did you describe the limitations of your work?

Original: [Yes] We provide the discussion about
Revised: [Yes] We provide the discussion about

Original: the limitation of our proposal in Section 6.
Revised: the limitation of our proposal in Section 6.

Original: (c) Did you discuss any potential negative societal impacts of your work?
Revised: (c) Did you discuss any potential negative societal impacts of your work?

Original: [N/A]
Revised: [N/A]

Original: (d) Have you read the ethics review guidelines and ensured that your paper conforms to
Revised: (d) Have you read the ethics review guidelines and ensured that your paper conforms to

Original: them?
Revised: them?

Original: [Yes]
Revised: [Yes]

Original: If you are including theoretical results...
Revised: If you are including theoretical results...

Original: (a) Did you state the full set of assumptions of all theoretical results?
Revised: (a) Did you state the full set of assumptions of all theoretical results?

Original: [Yes] Please refer to
Revised: [Yes] Please refer to

Original: Appendix B and C.
Revised: Appendix B and C.

Original: (b) Did you include complete proofs of all theoretical results?
Revised: (b) Did you include complete proofs of all theoretical results?

Original: [Yes] Please refer to
Revised: [Yes] Please refer to

Original: Appendix B and C.
Revised: Appendix B and C.

Original: If you ran experiments...
Revised: If you ran experiments...

Original: (a) Did you include the code, data, and instructions needed to reproduce the main experi-
Revised: (a) Did you include the code, data, and instructions needed to reproduce the main experi-

Original: mental results (either in the supplemental material or as a URL)?
Revised: mental results (either in the supplemental material or as a URL)?

Original: [Yes]
Revised: [Yes]

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters,
Revised: (b) Did you specify all the training details (e.g., data splits, hyperparameters,

Original: how they
Revised: how they

Original: were chosen)?
Revised: were chosen)?

Original: [Yes] Our implementation is adapted from Pytorch distributional RL
Revised: [Yes] Our implementation is adapted from Pytorch distributional RL

Original: modules [25].
Revised: modules [25].

Original: (c) Did you report error bars (e
Revised: (c) Did you report error bars (e

Original: g., with respect to the random seed after running experiments multiple times)?
Revised: g., with respect to the random seed after running experiments multiple times)?

Original: [Yes]
Revised: [Yes]

Original: (d) Did you include the total amount of compute and the type of resources used (e
Revised: (d) Did you include the total amount of compute and the type of resources used (e

Original: g., typeof GPUs, internal cluster, or cloud provider)?
Revised: g., typeof GPUs, internal cluster, or cloud provider)?

Original: [Yes]
Revised: [Yes]

Original: We provide the comparison
Revised: We provide the comparison

Original: of
Revised: of

Original: computational cost
Revised: computational cost

Original: in
Revised: in

Original: Figure 12 of
Revised: Figure 12 of

Original: Appendix
Revised: Appendix

Original: F
Revised: F

Original: If you are using existing assets (e
Revised: If you are using existing assets (e

Original: g., code, data, models) or curating/releasing new assets...
Revised: g., code, data, models) or curating/releasing new assets...

Original: (a) If your work uses existing assets, did you cite the creators?
Revised: (a) If your work uses existing assets, did you cite the creators?

Original: [N/A]
Revised: [N/A]

Original: ) Did you mention the license of the assets?
Revised: ) Did you mention the license of the assets?

Original: [N/A]
Revised: [N/A]

Original: c) Did you include any new assets either in the supplemental material or as a URL?
Revised: c) Did you include any new assets either in the supplemental material or as a URL?

Original: [N/A]
Revised: [N/A]

Original: ) Did you discuss whether and how consent was obtained from people whose data you’re
Revised: ) Did you discuss whether and how consent was obtained from people whose data you’re

Original: using/curating?
Revised: using/curating?

Original: [N/A]
Revised: [N/A]

Original: ) Did you discuss whether the data you are using/curating contains personally identiﬁable
Revised: ) Did you discuss whether the data you are using/curating contains personally identiﬁable

Original: information or offensive content?
Revised: information or offensive content?

Original: [N/A]
Revised: [N/A]

Original: If you used crowdsourcing or conducted research with human subjects...
Revised: If you used crowdsourcing or conducted research with human subjects...

Original: Did you include the full text of instructions given to participants and screenshots, if
Revised: Did you include the full text of instructions given to participants and screenshots, if

Original: applicable?
Revised: applicable?

Original: [N/A]
Revised: [N/A]

Original: Did you describe any potential participant risks, with links to Institutional Review
Revised: Did you describe any potential participant risks, with links to Institutional Review

Original: Board
Revised: Board

Original: ) approvals, if applicable?
Revised: ) approvals, if applicable?

Original: [N/A]
Revised: [N/A]

Original: Did you include the estimated hourly wage paid to participants and the total amount
Revised: Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?
Revised: spent on participant compensation?

Original: [N/A]
Revised: [N/A]