Original: Trust Region Policy Optimization with Optimal Transport Discrepancies: Duality and Algorithm for Continuous Actions
Revised: Trust Region Policy Optimization with Optimal Transport Discrepancies: Duality and Algorithm for Continuous Actions

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

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

Original: Abstract
Revised: Abstract

Original: Policy Optimization (PO) algorithms have been proven particularly suited to handlethe high-dimensionality of real-world continuous control tasks.
Revised: Policy Optimization (PO) algorithms have been proven particularly suited to handlethe high-dimensionality of real-world continuous control tasks.

Original: In this context,
Revised: In this context,

Original: Trust Region Policy Optimization methods represent a popular approach to stabilizethe policy updates.
Revised: Trust Region Policy Optimization methods represent a popular approach to stabilizethe policy updates.

Original: These usually rely on the Kullback-Leibler (KL) divergenceto limit the change in the policy.
Revised: These usually rely on the Kullback-Leibler (KL) divergenceto limit the change in the policy.

Original: The Wasserstein distance represents a naturalalternative, in place of the KL divergence, to deﬁne trust regions or to regularizethe objective function.
Revised: The Wasserstein distance represents a naturalalternative, in place of the KL divergence, to deﬁne trust regions or to regularizethe objective function.

Original: However, state-of-the-art works either resort to its approximations or do not provide an algorithm for continuous state-action spaces, reducingthe applicability of the method.
Revised: However, state-of-the-art works either resort to its approximations or do not provide an algorithm for continuous state-action spaces, reducingthe applicability of the method.

Original: In this paper, we explore optimal transport discrepancies (which include the Wasserstein distance) to deﬁne trust regions, and wepropose a novel algorithm – Optimal Transport Trust Region Policy Optimization(OT-TRPO) – for continuous state-action spaces.
Revised: In this paper, we explore optimal transport discrepancies (which include the Wasserstein distance) to deﬁne trust regions, and wepropose a novel algorithm – Optimal Transport Trust Region Policy Optimization(OT-TRPO) – for continuous state-action spaces.

Original: We circumvent the computationalcomplexity of the inﬁnite-dimensional optimization problem for PO by providinga one-dimensional dual reformulation for which strong duality holds.
Revised: We circumvent the computationalcomplexity of the inﬁnite-dimensional optimization problem for PO by providinga one-dimensional dual reformulation for which strong duality holds.

Original: We thenanalytically derive the optimal policy update given the solution of the dual problem.
Revised: We thenanalytically derive the optimal policy update given the solution of the dual problem.

Original: This way, we bypass the computation of optimal transport costs and of optimaltransport maps, which we implicitly characterize by solving the dual formulation.
Revised: This way, we bypass the computation of optimal transport costs and of optimaltransport maps, which we implicitly characterize by solving the dual formulation.

Original: Finally, we provide an experimental evaluation of our approach across variouscontrol tasks.
Revised: Finally, we provide an experimental evaluation of our approach across variouscontrol tasks.

Original: Our results show that optimal transport discrepancies can offer anadvantage over state-of-the-art approaches.
Revised: Our results show that optimal transport discrepancies can offer anadvantage over state-of-the-art approaches.

Original:  Introduction
Revised:  Introduction

Original: Reinforcement Learning (RL) has achieved outstanding results in numerous ﬁelds, from resourcemanagement [15], recommendation systems [40] and optimization of chemical reactions [42], tovideo-games [17, 41, 10] and board games [37], without sparing the world’s champion of GO [31].
Revised: Reinforcement Learning (RL) has achieved outstanding results in numerous ﬁelds, from resourcemanagement [13], recommendation systems [37] and optimization of chemical reactions [39], tovideo-games [15, 38, 8] and board games [34], without sparing the world’s champion of GO [29].

Original: Many of these successful applications rely on Policy Optimization (PO) algorithms, a family of
Revised: Many of these successful applications rely on Policy Optimization (PO) algorithms, a family of

Original: RL methods that are particularly suited to handle the high-dimensionality of real-world controltasks.
Revised: RL methods that are particularly suited to handle the high-dimensionality of real-world controltasks.

Original: PO algorithms approach the RL setting as an optimization problem in the policy space.
Revised: PO algorithms approach the RL setting as an optimization problem in the policy space.

Original: Inthis context, the main challenge is to provide policy improvement guarantees.
Revised: Inthis context, the main challenge is to provide policy improvement guarantees.

Original: One remarkableoption in this direction is represented by Trust Region Policy Optimization (TRPO) [28], whichconstrains the optimization problem to policies that are “close” to the current policy, whereby the
Revised: One remarkableoption in this direction is represented by Trust Region Policy Optimization (TRPO) [26], whichconstrains the optimization problem to policies that are “close” to the current policy, whereby the

Original: Kullback-Leibler (KL) divergence is used as a similarity measure.
Revised: Kullback-Leibler (KL) divergence is used as a similarity measure.

Original: Nevertheless, “closeness” inthe policy space can also be quantiﬁed via other functions.
Revised: Nevertheless, “closeness” inthe policy space can also be quantiﬁed via other functions.

Original: Recent work [21, 32, 19] proposedto replace the KL divergence with the Wasserstein distance, a particular instance of a so-calledoptimal transport discrepancy (or cost).
Revised: Recent work [19, 30, 17] proposedto replace the KL divergence with the Wasserstein distance, a particular instance of a so-calledoptimal transport discrepancy (or cost).

Original: Besides being very natural and expressive, optimal transport discrepancies enjoy powerful topological, geometrical, computational, and statistical features and guarantees [3, 36].
Revised: Besides being very natural and expressive, optimal transport discrepancies enjoy powerful topological, geometrical, computational, and statistical features and guarantees [3, 33].

Original: In particular, (i) optimal transport discrepancies allow us to compare probabilitymeasures (and so policies) not sharing the same support (for which the KL divergence is inﬁnity); and(ii)
Revised: In particular, (i) optimal transport discrepancies allow us to compare probabilitymeasures (and so policies) not sharing the same support (for which the KL divergence is inﬁnity); and(ii)

Original: they encapsulate the geometry encoded by the transport cost in the action space: the discrepancybetween two actions coincides with the discrepancy between the corresponding deterministic policies(whereas the KL divergence is again inﬁnity).
Revised: they encapsulate the geometry encoded by the transport cost in the action space: the discrepancybetween two actions coincides with the discrepancy between the corresponding deterministic policies(whereas the KL divergence is again inﬁnity).

Original: These reasons make optimal transport discrepanciesparticularly attractive for RL.
Revised: These reasons make optimal transport discrepanciesparticularly attractive for RL.

Original: However, the mere evaluation of optimal transport discrepanciesentails solving a transportation problem (e.g., see [22, 33]), which poses signiﬁcant computationalchallenges for its deployment in practical applications.
Revised: However, the mere evaluation of optimal transport discrepanciesentails solving a transportation problem (e.g., see [20, 31]), which poses signiﬁcant computationalchallenges for its deployment in practical applications.

Original: Most previous work [21, 19] overcomes thecomputational burden via approximation, effectively changing the original problem.
Revised: Most previous work [19, 17] overcomes thecomputational burden via approximation, effectively changing the original problem.

Original: Conversely, [32]proposes two algorithms to solve the PO problem exactly, studying the trust regions described viathe Wasserstein distance and the Sinkhorn divergence.
Revised: Conversely, [30]proposes two algorithms to solve the PO problem exactly, studying the trust regions described viathe Wasserstein distance and the Sinkhorn divergence.

Original: However, their analysis is limited to discrete(and ﬁnite) settings.
Revised: However, their analysis is limited to discrete(and ﬁnite) settings.

Original: In our work, we consider optimal transport discrepancies to construct the trustregion in settings where actions and states take value in general compact Polish spaces.
Revised: In our work, we consider optimal transport discrepancies to construct the trustregion in settings where actions and states take value in general compact Polish spaces.

Original: This allowstackling many applications of interest involving continuous domains such as physical control tasks.
Revised: This allowstackling many applications of interest involving continuous domains such as physical control tasks.

Original: We exploit and build on optimal transport duality theory to ensure an optimal policy update withinthe trust region, without any additional need for line searches, conversely to [28].
Revised: We exploit and build on optimal transport duality theory to ensure an optimal policy update withinthe trust region, without any additional need for line searches, conversely to [26].

Original: We circumvent thecomputation of optimal transport discrepancies via an analytical expression of the transport maps,which are characterized thanks to the dual reformulation.
Revised: We circumvent thecomputation of optimal transport discrepancies via an analytical expression of the transport maps,which are characterized thanks to the dual reformulation.

Original: Notably, our analysis enables a practicaland efﬁcient algorithm that encompasses both discrete and continuous settings.
Revised: Notably, our analysis enables a practicaland efﬁcient algorithm that encompasses both discrete and continuous settings.

Original: Contributions.
Revised: Contributions.

Original: Our contributions are summarized as follows:
Revised: Our contributions are summarized as follows:

Original: We derive the dual of the optimal transport trust region policy optimization problem and we show
Revised: We derive the dual of the optimal transport trust region policy optimization problem and we show

Original: that strong duality holds for general compact metric state-action spaces.
Revised: that strong duality holds for general compact metric state-action spaces.

Original: We further characterizethe optimal policy update given the solution to the dual problem.
Revised: We further characterizethe optimal policy update given the solution to the dual problem.

Original: We show that policy updates canresult in monotonic improvement of the performance function.2.
Revised: We show that policy updates canresult in monotonic improvement of the performance function.2.

Original: We propose a novel PO algorithm for continuous
Revised: We propose a novel PO algorithm for continuous

Original: spaces, O ptimal T ransport T rust R egion P olicy
Revised: spaces, O ptimal T ransport T rust R egion P olicy

Original: O ptimization (OT-TRPO).
Revised: O ptimization (OT-TRPO).

Original: Herein, we leverage the derived duality theory to provide policy updatesthat satisfy the optimal transport discrepancy constraintwhile circumventing its computation.3.
Revised: Herein, we leverage the derived duality theory to provide policy updatesthat satisfy the optimal transport discrepancy constraint circumventing its computation.3.

Original: We conduct experiments in several RL benchmarks in both discrete and continuous state-actionspaces, comparing our method to state-of-the-art approaches.
Revised: We conduct experiments in several RL benchmarks in both discrete and continuous state-actionspaces, comparing our method to state-of-the-art approaches.

Original: Our results show the effectivenessof our approach for PO and the beneﬁts of using optimal transport discrepancies.
Revised: Our results show the effectivenessof our approach for PO and the beneﬁts of using optimal transport discrepancies.

Original:  Related works
Revised:  Related works

Original: Optimal transport, and in particular the Wasserstein distance, has found various applications in RLand in particular in PO algorithms.
Revised: Optimal transport, and in particular the Wasserstein distance, has found various applications in RLand in particular in PO algorithms.

Original: In this section, we discuss the most relevant for our work; abroader overview is postponed to Appendix A.1.
Revised: In this section, we discuss the most relevant for our work; abroader overview is postponed to Appendix A.1.

Original: In [21], the authors propose Behavior Guided
Revised: In [19], the authors propose Behavior Guided

Original: Policy Gradient (BGPG), whereby they replace the KL divergence trust region from TRPO [28]by a Wasserstein distance penalty in a behavioral space.
Revised: Policy Gradient (BGPG), whereby they replace the KL divergence trust region from TRPO [26]by a Wasserstein distance penalty in a behavioral space.

Original: Although our alternating procedurein Section 5 may resemble the approach of [21] in spirit, our approach is fundamentally different;cfr.
Revised: Although our alternating procedurein Section 5 may resemble the approach of [19] in spirit, our approach is fundamentally different;cfr.

Original: [21, Algorithms 1 and 3] with Algorithm 1.
Revised: [19, Algorithms 1 and 3] with Algorithm 1.

Original: In [19], the authors further suggested incorporatingadditional information about the local behavior of policies encapsulated in the so-called Wasserstein
Revised: In [17], the authors further suggested incorporatingadditional information about the local behavior of policies encapsulated in the so-called Wasserstein

Original: Information Matrix, in the attempt to speed up the PO using a Wasserstein Natural Policy Gradient(WNPG).
Revised: Information Matrix, in the attempt to speed up the PO using a Wasserstein Natural Policy Gradient(WNPG).

Original: However, these approaches are relatively slow, compared to the traditional Proximal Policy
Revised: However, these approaches are relatively slow, compared to the traditional Proximal Policy

Original: Optimization (PPO) and TRPO.
Revised: Optimization (PPO) and TRPO.

Original: Conversely to our work, they do not build on the idea of trust regions:we instead guarantee that the policy update is “close” to the previous one, where the “closeness”is deﬁned via an optimal transport discrepancy (e.g., the Wasserstein distance).
Revised: Conversely to our work, they do not build on the idea of trust regions:we instead guarantee that the policy update is “close” to the previous one, where the “closeness”is deﬁned via an optimal transport discrepancy (e.g., the Wasserstein distance).

Original: Accordingly, theclosest related work to ours is the recent paper [32] which studied Wasserstein Policy Optimization(WPO) for discrete action spaces.
Revised: Accordingly, theclosest related work to ours is the recent paper [30] which studied Wasserstein Policy Optimization(WPO) for discrete action spaces.

Original: In contrast to [32], the present work addresses the general settingof compact Polish spaces encompassing the cases of continuous and discrete state-action spacesas particular cases.
Revised: In contrast to [30], the present work addresses the general settingof compact Polish spaces encompassing the cases of continuous and discrete state-action spacesas particular cases.

Original: While we also adopt a duality approach, our level of generality induces manychallenges compared to the discrete action space setting (see Remark 2), and it is compatible evenwith non-direct policy parametrizations.
Revised: While we also adopt a duality approach, our level of generality induces manychallenges compared to the discrete action space setting (see Remark 2), and it is compatible evenwith non-direct policy parametrizations.

Original: Finally, our work is closely connected with Wasserstein
Revised: Finally, our work is closely connected with Wasserstein

Original: Distributionally Robust Optimization (DRO) [8, 18, 4, 24, 13].
Revised: Distributionally Robust Optimization (DRO) [6, 16, 4, 22, 11].

Original: Albeit our duality results are inspiredfrom this literature, DRO is concerned with quantifying the worst-case risk of a cost functional overan ambiguity set of probability measures, which is a fundamentally different setting; see Remark 1.
Revised: Albeit our duality results are inspiredfrom this literature, DRO is concerned with quantifying the worst-case risk of a cost functional overan ambiguity set of probability measures, which is a fundamentally different setting; see Remark 1.

Original: Conversely to all the previous work, we show how to perform exact optimal transport-based TRPOin continuous settings: we exploit optimal transport theory to circumvent the computational burden of evaluating optimal transport discrepancies while still performing exact policy updates within thetrust regions.
Revised: Conversely to all the previous work, we show how to perform exact optimal transport-based TRPOin continuous settings: we exploit optimal transport theory to circumvent the computational burden of evaluating optimal transport discrepancies while still performing exact policy updates within thetrust regions.

Original: Accordingly, to the best of our knowledge, our practical algorithm is completely novel.
Revised: Accordingly, to the best of our knowledge, our practical algorithm is completely novel.

Original:  Preliminaries
Revised:  Preliminaries

Original: We brieﬂy introduce useful background and notations for the remainder of the paper.
Revised: We brieﬂy introduce useful background and notations for the remainder of the paper.

Original: Notation.
Revised: Notation.

Original: For every Polish space X (i.e., completely metrizable separable topological space),the set of Borel probability measures on X is denoted by P ( X ) .
Revised: For every Polish space X (i.e., completely metrizable separable topological space),the set of Borel probability measures on X is denoted by P ( X ) .

Original: The Dirac measure at some 96 point x 2 X is denoted by  x .
Revised: The Dirac measure at some 96 point x 2 X is denoted by  x .

Original: Given two Polish spaces X , Y , a Borel probability measures 97 µ 2 P ( X ) , and a Borel map T : X !
Revised: Given two Polish spaces X , Y , a Borel probability measures 97 µ 2 P ( X ) , and a Borel map T : X !

Original: Y , the pushforward measure of µ , denoted by T # µ , is deﬁnedby ( T # µ )( A ) :
Revised: Y , the pushforward measure of µ , denoted by T # µ , is deﬁnedby ( T # µ )( A ) :

Original: = µ ( T  1 ( A )) for all A 2 B ( Y ) , where B ( Y ) is the collection of Borel subsets of Y .
Revised: = µ ( T  1 ( A )) for all A 2 B ( Y ) , where B ( Y ) is the collection of Borel subsets of Y .

Original: The set of probability measures on a ﬁnite set X coincides with the probability simplex and will alsobe denoted by P ( X ) .
Revised: The set of probability measures on a ﬁnite set X coincides with the probability simplex and will alsobe denoted by P ( X ) .

Original: For a given function f : X !
Revised: For a given function f : X !

Original: R , the notation k f k 1 refers to sup x 2 X | f ( x ) | .
Revised: R , the notation k f k 1 refers to sup x 2 X | f ( x ) | .

Original: Markov Decision Process.
Revised: Markov Decision Process.

Original: We consider an inﬁnite-horizon discounted Markov Decision Process(MDP)
Revised: We consider an inﬁnite-horizon discounted Markov Decision Process(MDP)

Original: [23] M = ( S , A , P , r, ⇢ ,  ) , where S is the state space, A is the action space, P : S ⇥
Revised: [21] M = ( S , A , P , r, ⇢ ,  ) , where S is the state space, A is the action space, P : S ⇥

Original: A ⇥ S !
Revised: A ⇥ S !

Original: R  0 is the state transition probability kernel, r : S ⇥ A !
Revised: R  0 is the state transition probability kernel, r : S ⇥ A !

Original: R is the reward function, ⇢ is the initial state 104 probability distribution, and  2 [0 , 1) is the discount factor.
Revised: R is the reward function, ⇢ is the initial state 104 probability distribution, and  2 [0 , 1) is the discount factor.

Original: A randomized stationary Markovianpolicy, which we will simply call a policy in the rest of the paper, is a mapping ⇡ : S !
Revised: A randomized stationary Markovianpolicy, which we will simply call a policy in the rest of the paper, is a mapping ⇡ : S !

Original: P ( A ) 106 specifying for each s 2 S a probability measure over the set of actions A by ⇡ ( ·| s ) 2 P ( A ) .
Revised: P ( A ) 106 specifying for each s 2 S a probability measure over the set of actions A by ⇡ ( ·| s ) 2 P ( A ) .

Original: The set of all policies is denoted by ⇧ .
Revised: The set of all policies is denoted by ⇧ .

Original: Each policy ⇡ 2 ⇧ induces a discrete-time Markovreward process { ( s t , r ( s t , a t )) }
Revised: Each policy ⇡ 2 ⇧ induces a discrete-time Markovreward process { ( s t , r ( s t , a t )) }

Original: t 2 N , where s t 2 S represents the state of the system at time t 109 and r ( s t , a t ) corresponds to the reward received when executing action a t 2 A in state s t .
Revised: t 2 N , where s t 2 S represents the state of the system at time t 109 and r ( s t , a t ) corresponds to the reward received when executing action a t 2 A in state s t .

Original: Wedenote by P ⇢ , ⇡ the probability distribution of the Markov chain ( s t , a t ) issued from the MDPcontrolled by the policy ⇡ with initial state distribution ⇢ .
Revised: Wedenote by P ⇢ , ⇡ the probability distribution of the Markov chain ( s t , a t ) issued from the MDPcontrolled by the policy ⇡ with initial state distribution ⇢ .

Original: The associated expectation is denotedby E ⇢ , ⇡ and the notation E ⇡ is used whenever there is no dependence on ⇢ .
Revised: The associated expectation is denotedby E ⇢ , ⇡ and the notation E ⇡ is used whenever there is no dependence on ⇢ .

Original: The state-value function
Revised: The state-value function

Original: V ⇡ : S !
Revised: V ⇡ : S !

Original: R and the action-value function Q ⇡ : S ⇥ A !
Revised: R and the action-value function Q ⇡ : S ⇥ A !

Original: R are deﬁned for all s 2 S , a 2 A 114 by V ⇡ ( s ) :=
Revised: R are deﬁned for all s 2 S , a 2 A 114 by V ⇡ ( s ) :=

Original: E ⇡ [
Revised: E ⇡ [

Original: P 1 t =

Revised: P 1 t =

Original:  t r ( s t , a t ) | s 0 = s ] and Q ⇡ ( s, a ) :
Revised:  t r ( s t , a t ) | s 0 = s ] and Q ⇡ ( s, a ) :

Original: =
Revised: =

Original: E ⇡ [
Revised: E ⇡ [

Original: P [-1 t =
-]
Revised: P [-1 t =
-][Equation]

Original:  t r ( s t , a t ) | s
Revised: t r ( s t , a [Equation] | s

Original: [Equation]
Revised: [Equation]

Original: s, a 0 =a ] .
Revised: s, [Equation] =

Revised: [Equation]

Original: We also deﬁne the advantage function A ⇡ : S ⇥ A !
Revised: We also deﬁne the advantage function A ⇡ : S ⇥ A !

Original: R by A ⇡ ( s, a ) :=
Revised: R by A ⇡ ( s, a ) :

Original:  Q ⇡ ( s, a )  V ⇡ ( s ) .
Revised: = Q ⇡ ( s, a )  V ⇡ ( s ) .

Original: Givenan initial state probability distribution ⇢ , our goal is to ﬁnd a policy ⇡ maximizing the expectedlong-term return
Revised: Givenan initial state probability distribution ⇢ , our goal is to ﬁnd a policy ⇡ maximizing the expectedlong-term return

Original: [Equation]
Revised: [Equation]

Original: which is well-deﬁned when, e.g., the reward function is bounded.
Revised: which is well-deﬁned when, e.g., the reward function is bounded.

Original: To solve this PO problem, we onlyhave access to the observed state and action s t , a t at each time step t and the reward function r 1whereas the state transition kernel p is unknown.
Revised: To solve this PO problem, we onlyhave access to the observed state and action s t , a t at each time step t and the reward function r 1whereas the state transition kernel p is unknown.

Original: When the state and action spaces ( S and A ) areﬁnite, an optimal policy ⇡ ⇤ is guaranteed to exist.
Revised: When the state and action spaces ( S and A ) areﬁnite, an optimal policy ⇡ ⇤ is guaranteed to exist.

Original: When S and A are continuous, a (measurable)optimal policy is also guaranteed to exist (see [23, Theorem 6.11.11, p. 262]) under some assumptionson the state and action spaces, the reward function and the transition kernel; we will explicit theselater on.
Revised: When S and A are continuous, a (measurable)optimal policy is also guaranteed to exist (see [21, Theorem 6.11.11, p. 262]) under some assumptionson the state and action spaces, the reward function and the transition kernel; we will explicit theselater on.

Original: In this paper, we will focus on the continuous state-action space setting and comment on thediscrete (non necessarily ﬁnite) setting as a special case.
Revised: In this paper, we will focus on the continuous state-action space setting and comment on thediscrete (non necessarily ﬁnite) setting as a special case.

Original: Optimal transport.
Revised: Optimal transport.

Original: Consider a Polish space X and a continuous non-negative function c : X ⇥ X !
Revised: Consider a Polish space X and a continuous non-negative function c : X ⇥ X !

Original: R  0 , referred to as transport cost .
Revised: R  0 , referred to as transport cost .

Original: Let µ, ⌫ 2 P ( X ) and deﬁne the set of joint probability measures 128 on
Revised: Let µ, ⌫ 2 P ( X ) and deﬁne the set of joint probability measures 128 on

Original: X ⇥ X with marginals µ and ⌫ :
Revised: X ⇥ X with marginals µ and ⌫ :

Original: [Equation]
Revised: [Equation]

Original: We deﬁne the optimal transport discrepancy on P ( X ) for every probability measures µ and ⌫ by
Revised: We deﬁne the optimal transport discrepancy on P ( X ) for every probability measures µ and ⌫ by

Original: [Equation]
Revised: [Equation]

Original: Notice that this deﬁnition is valid for both discrete and continuous measures.
Revised: Notice that this deﬁnition is valid for both discrete and continuous measures.

Original: When c = d p ,where d is a distance on X and p  1 , then C ( µ, ⌫ ) 1 /p reduces to the celebrated type- p Wassersteindistance [36, 2].
Revised: When c = d p ,where d is a distance on X and p  1 , then C ( µ, ⌫ ) 1 /p reduces to the celebrated type- p Wassersteindistance [33, 2].

Original: In our PO context, we will use this discrepancy to compare two probabilitymeasures ⇡ ( ·| s ) 2 P ( A ) and ˜ ⇡ ( ·| s ) 2 P ( A ) for every s 2 S , where ⇡ , ˜ ⇡ 2 ⇧ are two policies.
Revised: In our PO context, we will use this discrepancy to compare two probabilitymeasures ⇡ ( ·| s ) 2 P ( A ) and ˜ ⇡ ( ·| s ) 2 P ( A ) for every s 2 S , where ⇡ , ˜ ⇡ 2 ⇧ are two policies.

Original: for
Revised: for

Original: Optimization
Revised: Optimization

Original: In this section, we study the TRPO algorithm with a trust region deﬁned using an optimal transportdiscrepancy as a measure of closeness between policies.
Revised: In this section, we study the TRPO algorithm with a trust region deﬁned using an optimal transportdiscrepancy as a measure of closeness between policies.

Original: We prove that the arising optimizationproblem admits an amenable dual reformulation.
Revised: We prove that the arising optimizationproblem admits an amenable dual reformulation.

Original: Importantly, we show that, given the dual optimalsolution, the primal solution has an analytical expression, which can lead to monotonic improvementsof the performance index.
Revised: Importantly, we show that, given the dual optimalsolution, the primal solution has an analytical expression, which can lead to monotonic improvementsof the performance index.

Original: 4.1 Policy iteration algorithm with optimal transport-based trust regions
Revised: 4.1 Policy iteration algorithm with optimal transport-based trust regions

Original: We start reporting the performance difference lemma [11, Lemma 6.1] expressing the differencebetween the expected returns of two policies ⇡ , ˜ ⇡ 2 ⇧ using the advantage function of the policy ⇡ :
Revised: We start reporting the performance difference lemma [9, Lemma 6.1] expressing the differencebetween the expected returns of two policies ⇡ , ˜ ⇡ 2 ⇧ using the advantage function of the policy ⇡ :

Original: [Equation]
Revised: [Equation]

Original: where ⇢ ˜ ⇡ is the discounted state-occupancy measure [11].
Revised: where ⇢ ˜ ⇡ is the discounted state-occupancy measure [9].

Original: The complex dependency of the discountedvisitation frequency ⇢ ˜ ⇡ on the policy ˜ ⇡ hampers the direct optimization of (2); see [28, Sec. 2].
Revised: The complex dependency of the discountedvisitation frequency ⇢ ˜ ⇡ on the policy ˜ ⇡ hampers the direct optimization of (2); see [26, Sec. 2].

Original: Following previous work, we consider instead a local approximation of the expected return J , deﬁnedby
Revised: Following previous work, we consider instead a local approximation of the expected return J , deﬁnedby

Original: [Equation]
Revised: [Equation]

Original: Observe that this approximation uses the discounted state-occupancy measure ⇢ ⇡ (which can beestimated) instead of ⇢ ˜ ⇡
Revised: Observe that this approximation uses the discounted state-occupancy measure ⇢ ⇡ (which can beestimated) instead of ⇢ ˜ ⇡

Original: (see (2)).
Revised: (see (2)).

Original: In other words, the inﬂuence of a policy change on the discountedstate-occupancy measure is neglected.
Revised: In other words, the inﬂuence of a policy change on the discountedstate-occupancy measure is neglected.

Original: Moreover, this surrogate function coincides with the expectedreturn J when ˜ ⇡ = ⇡ .
Revised: Moreover, this surrogate function coincides with the expectedreturn J when ˜ ⇡ = ⇡ .

Original: Then, (3) motivates a policy update rule maximizing at each time step theapproximation L ⇡ (˜ ⇡ ) over ˜ ⇡ , where ⇡ is the current policy that we want to improve upon (seealso [32, Section 2, Eq.
Revised: Then, (3) motivates a policy update rule maximizing at each time step theapproximation L ⇡ (˜ ⇡ ) over ˜ ⇡ , where ⇡ is the current policy that we want to improve upon (seealso [30, Section 2, Eq.

Original: To ensure stability of the update, we conservatively update the policyusing a discrepancy constraint between the current and the new one.
Revised: To ensure stability of the update, we conservatively update the policyusing a discrepancy constraint between the current and the new one.

Original: Unlike TRPO, we do not usethe KL divergence to deﬁne the trust region, but instead an optimal transport discrepancy.
Revised: Unlike TRPO, we do not usethe KL divergence to deﬁne the trust region, but instead an optimal transport discrepancy.

Original: Then, ateach time step, our method solves
Revised: Then, ateach time step, our method solves

Original: [Equation]
Revised: [Equation]

Original: where " > 0 is a parameter deﬁning the radius of the trust region T " ( ⇡ ) .
Revised: where " > 0 is a parameter deﬁning the radius of the trust region T " ( ⇡ ) .

Original: Similarly to [28, Eq.
Revised: Similarly to [26, Eq.

Original: (12)]and [32, Problem (4)], we consider the average optimal transport discrepancy over the state spaceas optimization constraint.
Revised: (12)]and [30, Problem (4)], we consider the average optimal transport discrepancy over the state spaceas optimization constraint.

Original: Accordingly, the OT-TRPO policy optimization results from iterativelysolving Problem (P).
Revised: Accordingly, the OT-TRPO policy optimization results from iterativelysolving Problem (P).

Original: 4.2 Dual of the trust-region constrained problem (P)
Revised: 4.2 Dual of the trust-region constrained problem (P)

Original: Problem (P) is intractable for two main reasons.
Revised: Problem (P) is intractable for two main reasons.

Original: First, as soon as the state or action space iscontinuous, it is an inﬁnite-dimensional problem.
Revised: First, as soon as the state or action space iscontinuous, it is an inﬁnite-dimensional problem.

Original: Second, the mere evaluation of the trust-regionconstraint (e.g., for line search as in TRPO [28]) needs (possibly) inﬁnitely many computations ofthe optimal transport discrepancy, which is itself already challenging to estimate.
Revised: Second, the mere evaluation of the trust-regionconstraint (e.g., for line search as in TRPO [26]) needs (possibly) inﬁnitely many computations ofthe optimal transport discrepancy, which is itself already challenging to estimate.

Original: However, inspiredby prior works on Wasserstein DRO [8, 18, 39, 4], we show that problem (P) admits a tractableone-dimensional convex dual reformulation.
Revised: However, inspiredby prior works on Wasserstein DRO [6, 16, 36, 4], we show that problem (P) admits a tractableone-dimensional convex dual reformulation.

Original: This duality theorem is the cornerstone of the design ofour algorithm.
Revised: This duality theorem is the cornerstone of the design ofour algorithm.

Original: Before stating the result, we introduce some useful assumptions.
Revised: Before stating the result, we introduce some useful assumptions.

Original: Assumption 1.
Revised: Assumption 1.

Original: The state space S is a compact Borel subset of an Euclidean space, the action space Ais a compact subset of a Polish space, the reward function r is a continuous function and for everybounded continuous function w on S ,
Revised: The state space S is a compact Borel subset of an Euclidean space, the action space Ais a compact subset of a Polish space, the reward function r is a continuous function and for everybounded continuous function w on S ,

Original: [Equation]
Revised: [Equation]

Original: is continuous in both s and a .
Revised: continuous in both s and a .

Original: Notice that under this assumption, there exists an optimal measurable (stationary) policy to the POproblem formulated in Section 3.
Revised: Notice that under this assumption, there exists an optimal measurable (stationary) policy to the POproblem formulated in Section 3.

Original: We refer the reader to [23, Theorem 6.11.11, p. 262] for a statementof this result and milder assumptions.
Revised: We refer the reader to [21, Theorem 6.11.11, p. 262] for a statementof this result and milder assumptions.

Original: Assumption 2.
Revised: Assumption 2.

Original: For every policy ⇡ 2 ⇧ , the advantage function A ⇡ : S ⇥ A !
Revised: For every policy ⇡ 2 ⇧ , the advantage function A ⇡ : S ⇥ A !

Original: R is continuous.
Revised: R is continuous.

Original: Moreover, the transport cost c: A ⇥ A ! R  0 is continuous and satisﬁes [-c ( a, a ) = 
-]
Revised: Moreover, the transport cost c and satisﬁes [-c ( a, a ) = 
-]all

Original: for all a 2 A . In the next theorem, we show that under these assumptions Problem (P) admits a dual reformulationfor which strong duality holds.
Revised: In the next theorem, we show that under these assumptions Problem (P) admits a dual reformulationfor which strong duality holds.

Original: Theorem 1 (Dual formulation) .
Revised: Theorem 1 (Dual formulation) .

Original: Let " > 0 , and let Assumptions 1 and 2 hold.
Revised: Let " > 0 , and let Assumptions 1 and 2 hold.

Original: For every policy ⇡ 2 ⇧ ,the following strong duality result holds:
Revised: For every policy ⇡ 2 ⇧ ,the following strong duality result holds:

Original: [Equation]
Revised: [Equation]

Original: Moreover, the primal and dual problems (P) and (D) admit a maximizer and a minimizer, respectively.
Revised: Moreover, the primal and dual problems (P) and (D) admit a maximizer and a minimizer, respectively.

Original: Remarkably, Problem (D) is one-dimensional and convex, and it only involves the current policy⇡ , advantage function A ⇡ , and visitation frequency ⇢ ⇡ .
Revised: Remarkably, Problem (D) is one-dimensional and convex, and it only involves the current policy⇡ , advantage function A ⇡ , and visitation frequency ⇢ ⇡ .

Original: The proof of Theorem 1 is constructive.
Revised: The proof of Theorem 1 is constructive.

Original: Inparticular, we derive a closed-form solution of problem (P) as a function of the optimal Lagrangemultiplier  ⇤ solving problem (D) and the policy ⇡ deﬁning the problem.
Revised: Inparticular, we derive a closed-form solution of problem (P) as a function of the optimal Lagrangemultiplier  ⇤ solving problem (D) and the policy ⇡ deﬁning the problem.

Original: Even if the closed-formpolicy is part of Theorem 1 and its proof, we present it separately for clarity and for later reference.
Revised: Even if the closed-formpolicy is part of Theorem 1 and its proof, we present it separately for clarity and for later reference.

Original: We also introduce a few additional useful notations, which are instrumental to to derive a practicalalgorithm (similarly to [32]).
Revised: We also introduce a few additional useful notations, which are instrumental to to derive a practicalalgorithm (similarly to [30]).

Original: Under Assumptions 1 and 2, deﬁne for every   0 the  -regularized advantage   : S ⇥ A !
Revised: Under Assumptions 1 and 2, deﬁne for every   0 the  -regularized advantage   : S ⇥ A !

Original: R 189 and its associated set of maximizers for every s 2 S , a 2 A as follows:
Revised: R 189 and its associated set of maximizers for every s 2 S , a 2 A as follows:

Original: [Equation]
Revised: [Equation]

Original: Corollary 2 (Optimal policy) .
Revised: Corollary 2 (Optimal policy) .

Original: Under the setting and assumptions of Theorem 1, for any policy ⇡ 2 ⇧ ,let  ⇤  0 be the minimizer of the dual problem (D) .
Revised: Under the setting and assumptions of Theorem 1, for any policy ⇡ 2 ⇧ ,let  ⇤  0 be the minimizer of the dual problem (D) .

Original: Then, the following statements hold:
Revised: Then, the following statements hold:

Original: For every   0 , there exist two measurable selection maps T  : S ⇥ A !
Revised: For every   0 , there exist two measurable selection maps T  : S ⇥ A !

Original: A and T  : S ⇥ A !
Revised: A and T  : S ⇥ A !

Original: Asuch that for every s 2 S , a 2 A :
Revised: Asuch that for every s 2 S , a 2 A :

Original: [Equation]
Revised: [Equation]

Original: There exists t ⇤ 2 [0 , 1] such that
Revised: There exists t ⇤ 2 [0 , 1] such that

Original: [Equation]
Revised: [Equation]

Original: There exists an optimal feasible policy ˜ ⇡ for problem (P) deﬁned for every s 2 S by
Revised: There exists an optimal feasible policy ˜ ⇡ for problem (P) deﬁned for every s 2 S by

Original: [Equation]
Revised: [Equation]

Original: Intuitively, Corollary 2 suggests that the optimal policy results from displacing the probability mass⇡ ( a | s ) to the maximizers of the  ⇤ -regularized advantage   ⇤ ( s, a ) , where  ⇤  0 is the optimaldual solution.
Revised: Intuitively, Corollary 2 suggests that the optimal policy results from displacing the probability mass⇡ ( a | s ) to the maximizers of the  ⇤ -regularized advantage   ⇤ ( s, a ) , where  ⇤  0 is the optimaldual solution.

Original: Since maximizers are generally not unique (i.e., D  ⇤ ( s, a ) is not a singleton), one needsto balance between the closest (i.e., T  ⇤ ( s, a ) ) and the furthest apart (i.e., T  ⇤ ( s, a ) ) to satisfy thetrust region constraint.
Revised: Since maximizers are generally not unique (i.e., D  ⇤ ( s, a ) is not a singleton), one needsto balance between the closest (i.e., T  ⇤ ( s, a ) ) and the furthest apart (i.e., T  ⇤ ( s, a ) ) to satisfy thetrust region constraint.

Original: The complete proof of the results of this section is deferred to Appendix B.1.
Revised: The complete proof of the results of this section is deferred to Appendix B.1.

Original: Remark 1 .
Revised: Remark 1 .

Original: Similar results were previously established in the literature in the context of DRO (e.g.,see [8, Theorem 1] and [4, Theorem 1]).
Revised: Similar results were previously established in the literature in the context of DRO (e.g.,see [6, Theorem 1] and [4, Theorem 1]).

Original: While these results closely inspire our proof, there isa major difference: in DRO, one seeks to evaluate the worst-case cost over an ambiguous set ofprobability distributions, expressed in terms of the optimal transport discrepancy.
Revised: While these results closely inspire our proof, there isa major difference: in DRO, one seeks to evaluate the worst-case cost over an ambiguous set ofprobability distributions, expressed in terms of the optimal transport discrepancy.

Original: As such, theaverage optimal transport discrepancy in the trust region constraint is replaced by a single optimaltransport discrepancy.
Revised: As such, theaverage optimal transport discrepancy in the trust region constraint is replaced by a single optimaltransport discrepancy.

Original: Thus, one does not need to ensure the regularity of the problem with respect tothe state (e.g., measurability of T  w.r.t.
Revised: Thus, one does not need to ensure the regularity of the problem with respect tothe state (e.g., measurability of T  w.r.t.

Original: s ).
Revised: s ).

Original: This is reﬂected in our assumption of joint continuity ofthe advantage in state and action and the state space being compact.
Revised: This is reﬂected in our assumption of joint continuity ofthe advantage in state and action and the state space being compact.

Original: To readily deploy existing resultsin DRO, one needs (i) to consider a single state only (i.e., S = { s } ) or (ii) to deﬁne a trust regionfor each state (at the price of inﬁnitely many constraints).
Revised: To readily deploy existing resultsin DRO, one needs (i) to consider a single state only (i.e., S = { s } ) or (ii) to deﬁne a trust regionfor each state (at the price of inﬁnitely many constraints).

Original: To transform (P) into a DRO, one mightalternatively identify a policy as a probability measure over
Revised: To transform (P) into a DRO, one mightalternatively identify a policy as a probability measure over

Original: Q s 2 S A , and hope to deploy standard 212 duality arguments in DRO.
Revised: Q s 2 S A , and hope to deploy standard 212 duality arguments in DRO.

Original: However, the uncountable product of Polish spaces is not a Polish space,which makes all results in DRO, and more generally in optimal transport [36], inapplicable.
Revised: However, the uncountable product of Polish spaces is not a Polish space,which makes all results in DRO, and more generally in optimal transport [33], inapplicable.

Original: Remark 2 .
Revised: Remark 2 .

Original: A similar result in the discrete case was presented in [32].
Revised: A similar result in the discrete case was presented in [30].

Original: We highlight four majordifferences.
Revised: We highlight four majordifferences.

Original: First, in the discrete setting, problem (P) is a ﬁnite-dimensional linear optimization problem, for which strong duality holds.
Revised: First, in the discrete setting, problem (P) is a ﬁnite-dimensional linear optimization problem, for which strong duality holds.

Original: Thus, linear programming arguments can be used to derivethe dual reformulation.
Revised: Thus, linear programming arguments can be used to derivethe dual reformulation.

Original: In the continuous setting, the same proof strategy would require to mobilizethe abstract machinery of inﬁnite-dimensional linear programming [14].
Revised: In the continuous setting, the same proof strategy would require to mobilizethe abstract machinery of inﬁnite-dimensional linear programming [12].

Original: Second, in the discretesetting, continuity and measurability of all functions are “for free”.
Revised: Second, in the discretesetting, continuity and measurability of all functions are “for free”.

Original: On the contrary, the continuouscase imposes a careful analysis of these issues.
Revised: On the contrary, the continuouscase imposes a careful analysis of these issues.

Original: Third, the optimal policy update [32] implicitlyassumes that the set D  ( s, a ) (see (4)) is a singleton, which is rarely satisﬁed in practice.
Revised: Third, the optimal policy update [30] implicitlyassumes that the set D  ( s, a ) (see (4)) is a singleton, which is rarely satisﬁed in practice.

Original: Fourth, as abyproduct of our proof, we show that (D) is a (one-dimensional) convex optimization problem, whichcan be solved efﬁciently via off-the-shelf solvers.
Revised: Fourth, as abyproduct of our proof, we show that (D) is a (one-dimensional) convex optimization problem, whichcan be solved efﬁciently via off-the-shelf solvers.

Original: This way, we do not need to resort to approximationtechniques [32, Section 6.1] for the optimal dual multiplier.
Revised: This way, we do not need to resort to approximationtechniques [30, Section 6.1] for the optimal dual multiplier.

Original: Discrete state-action spaces.
Revised: Discrete state-action spaces.

Original: In the rest of this section, we specialize our results to discrete (ﬁnite)state-action spaces (which trivially satisfy Assumptions 1 and 2).
Revised: In the rest of this section, we specialize our results to discrete (ﬁnite)state-action spaces (which trivially satisfy Assumptions 1 and 2).

Original: Without loss of generality werepresent the state and action spaces by S
Revised: Without loss of generality werepresent the state and action spaces by S

Original: = { s 1 , . . .
Revised: = { s 1 , . . .

Original: , s M } and A = { a 1 , . .
Revised: , s M } and A = { a 1 , . .

Original: . , a N } where M and N 228 are two positive integers, and represent any policy ⇡ 2 ⇧ and its corresponding state-occupancymeasure ⇢ ⇡ as discrete measures:
Revised: . , a N } where M and N 228 are two positive integers, and represent any policy ⇡ 2 ⇧ and its corresponding state-occupancymeasure ⇢ ⇡ as discrete measures:

Original: [Equation]
Revised: [Equation]

Original: where ⇢ i , ⇡ i,j  0 for every i 2 { 1 , . . .
Revised: where ⇢ i , ⇡ i,j  0 for every i 2 { 1 , . . .

Original: , M } , j 2 { 1 , . . .
Revised: , M } , j 2 { 1 , . . .

Original: , N } ,
Revised: , N } ,

Original: P
Revised: P

Original: M i [Equation] 1 ⇢ i = 1 and
Revised: M i =1 ⇢ i = 1 and

Original: P
Revised: P

Original: N j =1 ⇡ i,j = 1 forevery i 2 { 1 , . . .
Revised: N j =1 ⇡ i,j = 1 forevery i 2 { 1 , . . .

Original: , M } .
Revised: , M } .

Original: The analogous results to Theorem 1 and Corollary 2 are as follows.
Revised: The analogous results to Theorem 1 and Corollary 2 are as follows.

Original: Corollary 3 (Dual formulation - discrete setting) .
Revised: Corollary 3 (Dual formulation - discrete setting) .

Original: Let " > 0 .
Revised: Let " > 0 .

Original: For every policy ⇡ 2 ⇧ , the followingstrong duality result holds:
Revised: For every policy ⇡ 2 ⇧ , the followingstrong duality result holds:

Original: [Equation]
Revised: [Equation]

Original: In particular, let  ⇤  0 be a solution to (discrete-D) , and given for every i 2 { 1 , . . .
Revised: In particular, let  ⇤  0 be a solution to (discrete-D) , and given for every i 2 { 1 , . . .

Original: , M } , j 2{ 1 , . . .
Revised: , M } , j 2{ 1 , . . .

Original: , N } , select any
Revised: , N } , select any

Original: [Equation]
Revised: [Equation]

Original: ) and let
Revised: ) and let

Original: [Equation]
Revised: [Equation]

Original: Then, c  "  c and with t ⇤ 2 [0 , 1] if c = c and t ⇤ =
Revised: Then, c  "  c and with t ⇤ 2 [0 , 1] if c = c and t ⇤ =

Original: ( c  " ) /
Revised: ( c  " ) /

Original: ( c  c ) otherwise, an optimal policy ˜ ⇡is given by
Revised: ( c  c ) otherwise, an optimal policy ˜ ⇡is given by

Original: [Equation]
Revised: [Equation]

Original: The proof of this result stems from substituting the discrete measures as deﬁned in (8) in problems (P)and (D), and observing that the images of the mappings T  ⇤ and T  ⇤ have ﬁnite supports in the currentsetting.
Revised: The proof of this result stems from substituting the discrete measures as deﬁned in (8) in problems (P)and (D), and observing that the images of the mappings T  ⇤ and T  ⇤ have ﬁnite supports in the currentsetting.

Original: Notably, Corollary 3 directly provides an implementable algorithm for the policy update,circumventing the difﬁculty of the mixed-integer optimization problem (discrete-P): solving theone-dimensional convex program (discrete-D) provides the optimal Lagrange multiplier associatedto the trust region constraint of the primal problem which can be directly used to compute theactions
Revised: Notably, Corollary 3 directly provides an implementable algorithm for the policy update,circumventing the difﬁculty of the mixed-integer optimization problem (discrete-P): solving theone-dimensional convex program (discrete-D) provides the optimal Lagrange multiplier associatedto the trust region constraint of the primal problem which can be directly used to compute theactions

Original: b ⇤ i,j , b ⇤ i,j via (9), and thus the policy ˜ ⇡ via (10).
Revised: b ⇤ i,j , b ⇤ i,j via (9), and thus the policy ˜ ⇡ via (10).

Original: Remark 3 .
Revised: Remark 3 .

Original: The policy update suggested by (10) differs from the one in [32] (see [32, Theorem 1, (5)]where f ⇤ s ( i, j ) 2 { 0 , 1 } with their notations).
Revised: The policy update suggested by (10) differs from the one in [30] (see [30, Theorem 1, (5)]where f ⇤ s ( i, j ) 2 { 0 , 1 } with their notations).

Original: Indeed, our policy update relies on “splitting theprobability mass”: the probability mass ⇡ ( ·| s i ) is displaced to b
Revised: Indeed, our policy update relies on “splitting theprobability mass”: the probability mass ⇡ ( ·| s i ) is displaced to b

Original: ⇤ i,j and b
Revised: ⇤ i,j and b

Original: ⇤ i,j with weights t ⇤ and 1  t ⇤ , 249 respectively.
Revised: ⇤ i,j with weights t ⇤ and 1  t ⇤ , 249 respectively.

Original: This result is consistent with the Wasserstein DRO literature (e.g., see [4, Remark 2]).
Revised: This result is consistent with the Wasserstein DRO literature (e.g., see [4, Remark 2]).

Original: The result provided in [32] corresponds to the particular case where b ⇤ i,
Revised: The result provided in [30] corresponds to the particular case where b ⇤ i,

Original: j = b
Revised: j = b

Original: ⇤ i,j which amounts tosupposing that the set D  ( s i , a j ) deﬁned in (4) is a singleton.
Revised: ⇤ i,j which amounts tosupposing that the set D  ( s i , a j ) deﬁned in (4) is a singleton.

Original: We provide further comments andexamples in Appendix A.2 to illustrate the importance of this “mass splitting”.
Revised: We provide further comments andexamples in Appendix A.2 to illustrate the importance of this “mass splitting”.

Original: 4.3 Policy improvement
Revised: 4.3 Policy improvement

Original: In the next result, we show that our policy update leads to a monotonic improvement of the performance function J up to the advantage function estimation error.
Revised: In the next result, we show that our policy update leads to a monotonic improvement of the performance function J up to the advantage function estimation error.

Original: Proposition 4 (Performance improvement) .
Revised: Proposition 4 (Performance improvement) .

Original: Let ⇡ 2 ⇧ .
Revised: Let ⇡ 2 ⇧ .

Original: Consider solutions ˜ ⇡
Revised: Consider solutions ˜ ⇡

Original: ⇤ 2 ⇧ and  ⇤  0 257 of problems (P) and (D) , respectively.
Revised: ⇤ 2 ⇧ and  ⇤  0 257 of problems (P) and (D) , respectively.

Original: If the true advantage function A ⇡ is approximated by someestimated advantage function ˆ A ⇡ such that k A ⇡  ˆ A ⇡ k 1 < 1 , then the following bound holds:
Revised: If the true advantage function A ⇡ is approximated by someestimated advantage function ˆ A ⇡ such that k A ⇡  ˆ A ⇡ k 1 < 1 , then the following bound holds:

Original: [Equation]
Revised: [Equation]

Original: Proposition 4 indicates that optimal transport-based trust region policy optimization leads to monotonic improvement of the performance function when we have access to the true advantage function.
Revised: Proposition 4 indicates that optimal transport-based trust region policy optimization leads to monotonic improvement of the performance function when we have access to the true advantage function.

Original: The proof, postponed to Appendix B.2, of this result builds on the performance difference lemma(see (2)) and uses the closed-form expression of the optimal policy solving problem (P) as constructedin the proof of Theorem 1 (see Corollary 2).
Revised: The proof, postponed to Appendix B.2, of this result builds on the performance difference lemma(see (2)) and uses the closed-form expression of the optimal policy solving problem (P) as constructedin the proof of Theorem 1 (see Corollary 2).

Original: The analog of this result for a ﬁnite action space wasproved in [32, Theorem 2, p. 5].
Revised: The analog of this result for a ﬁnite action space wasproved in [30, Theorem 2, p. 5].

Original: To the best of our knowledge, this result is novel for the continuousstate-action space setting.
Revised: To the best of our knowledge, this result is novel for the continuousstate-action space setting.

Original:  Practical Optimal Transport Trust Region Policy Optimization Algorithm
Revised:  Practical Optimal Transport Trust Region Policy Optimization Algorithm

Original: In this section, we use the duality results on Section 4 to derive a practical algorithm for OTTRPO.
Revised: In this section, we use the duality results on Section 4 to derive a practical algorithm for OTTRPO.

Original: Herein, we restrict the policy search set ⇧ to the set of policies ⇡ ✓ parametrized by avector ✓ 2 R d for some integer d > 0 .
Revised: Herein, we restrict the policy search set ⇧ to the set of policies ⇡ ✓ parametrized by avector ✓ 2 R d for some integer d > 0 .

Original: We require the policy parametrization to be continuouslydifferentiable with respect to ✓ (for every state and action).
Revised: We require the policy parametrization to be continuouslydifferentiable with respect to ✓ (for every state and action).

Original: This way, we simultaneously coverthe direct parametrization (for which Corollary 3 directly provides a policy update) as well ascommonly used policy parametrizations (e.g., softmax and the Gaussian policies).
Revised: This way, we simultaneously coverthe direct parametrization (for which Corollary 3 directly provides a policy update) as well ascommonly used policy parametrizations (e.g., softmax and the Gaussian policies).

Original: Accordingly, thedual problem (D) can be reformulated as follows for every ✓ 2 R d :
Revised: Accordingly, thedual problem (D) can be reformulated as follows for every ✓ 2 R d :

Original: [Equation]
Revised: [Equation]

Original: Given a current policy represented by the vector ✓ ,we ﬁrst solve the one-dimensional convex problem (12) to obtain its solution  ⇤ .
Revised: Given a current policy represented by the vector ✓ ,we ﬁrst solve the one-dimensional convex problem (12) to obtain its solution  ⇤ .

Original: Then, we usethe optimal dual multiplier to derive the optimalpolicy update within the trust region.
Revised: Then, we usethe optimal dual multiplier to derive the optimalpolicy update within the trust region.

Original: The procedure is summarized in Algorithm 1.
Revised: The procedure is summarized in Algorithm 1.

Original: Dependingon the parametrization of the policy, the stepsof Section 5 are as follows:
Revised: Dependingon the parametrization of the policy, the stepsof Section 5 are as follows:

Original: Algorithm
Revised: Algorithm

Original: Algorithm 1 - step 3.
Revised: Algorithm 1 - step 3.

Original: In the discrete setting, the visitation frequency is estimated via Monte Carlomethods.
Revised: In the discrete setting, the visitation frequency is estimated via Monte Carlomethods.

Original: In the continuous case, we weight every visited state equally.
Revised: In the continuous case, we weight every visited state equally.

Original: We propose three ways toestimate the unknown advantage function via samples 3 :
Revised: We propose three ways toestimate the unknown advantage function via samples 3 :

Original: Monte Carlo methods or TD-learning (for discrete settings only);
Revised: Monte Carlo estimation (for discrete settings only);

Original: General Advantage Estimation (GAE) [29], using a neural network to approximate the value
Revised: General Advantage Estimation (GAE) [27], using a neural network to approximate the value

Original: function like in standard actor-critic methods; and
Revised: function like in standard actor-critic methods; and

Original: Direct estimation via non-linear approximators (e.g., using directly a neural network for the
Revised: Direct estimation via non-linear approximators (e.g., using directly a neural network for the

Original: advantage function).
Revised: advantage function).

Original: Algorithm 1 - step 4 (evaluation of G ).
Revised: Algorithm 1 - step 4 (evaluation of G ).

Original: Depending on the setting, we propose various ways toevaluate G (  , ✓ ) .
Revised: Depending on the setting, we propose various ways toevaluate G (  , ✓ ) .

Original: They all apply to both continuous and discrete states.
Revised: They all apply to both continuous and discrete states.

Original: Finite actions : Since the “max” in (12) is over ﬁnitely many actions, we can directly evaluate (12)
Revised: Finite actions : Since the “max” in (12) is over ﬁnitely many actions, we can directly evaluate (12)

Original: for any   0 .
Revised: for any   0 .

Original: Gaussian policy parametrization : With m ( s ) being the mean of the Gaussian policy (with ﬁxedvariance), we approximate G (  , ✓ ) by
Revised: Gaussian policy parametrization : With m ( s ) being the mean of the Gaussian policy (with ﬁxedvariance), we approximate G (  , ✓ ) by

Original: [Equation]
Revised: [Equation]

Original: where ˆ S are the state visited in the trajectory and ˆ A is a collection of samples collected from A .
Revised: where ˆ S are the state visited in the trajectory and ˆ A is a collection of samples collected from A .

Original: Arbitrary policy parametrization : For a neural network approximation of the advantage function:
Revised: Arbitrary policy parametrization : For a neural network approximation of the advantage function:

Original: [Equation]
Revised: [Equation]

Original: where ˆ S are the state visited in the trajectory and ˆ A i are collections of samples collected from A .
Revised: where ˆ S are the state visited in the trajectory and ˆ A i are collections of samples collected from A .

Original: Algorithm 1 - step 4 (solving for  ⇤ )
Revised: Algorithm 1 - step 4 (solving for  ⇤ )

Original: Since
Revised: Since

Original: (D) is a one-dimensional convex optimization problem,
Revised: (D) is a one-dimensional convex optimization problem,

Original: ⇤
Revised: ⇤

Original: can be found using
Revised: can be found using

Original: any solver
Revised: any solver

Original: for
Revised: for

Original: convex optimization problems.
Revised: convex optimization problems.

Original: Algorithm
Revised: Algorithm

Original: Direct parametrization (ﬁnite spaces) : Update ✓ according to (10) and (9).2.
Revised: Direct parametrization (ﬁnite spaces) : Update ✓ according to (10) and (9).2.

Original: Direct parametrization
Revised: Direct parametrization

Original: via policy network (continuous states, discrete actions) : Use (10) and (9)to compute the optimal policy update at the visited states, denoted by ⇡ ⇤ ✓ t .
Revised: via policy network (continuous states, discrete actions) : Use (10) and (9)to compute the optimal policy update at the visited states, denoted by ⇡ ⇤ ✓ t .

Original: Then, update the policynetwork by performing gradient descent on the loss L ( ✓ ) =
Revised: Then, update the policynetwork by performing gradient descent on the loss L ( ✓ ) =

Original: P s 2 ˆ S ⇢ ⇡ ✓ t ( s )  ⇡ ✓ ( ·| s )  ⇡ ⇤ ✓ t ( ·| s )  2to steer ⇡ ✓ to the optimal policy update ⇡ ⇤ ✓ t within the trust region.3.
Revised: P s 2 ˆ S ⇢ ⇡ ✓ t ( s )  ⇡ ✓ ( ·| s )  ⇡ ⇤ ✓ t ( ·| s )  2to steer ⇡ ✓ to the optimal policy update ⇡ ⇤ ✓ t within the trust region.3.

Original: Arbitrary policy parametrization : Since there are inﬁnitely many actions, the evaluation of the“max” is computationally demanding, and so Corollary 2 cannot be directly utilized for the policyupdate.
Revised: Arbitrary policy parametrization : Since there are inﬁnitely many actions, the evaluation of the“max” is computationally demanding, and so Corollary 2 cannot be directly utilized for the policyupdate.

Original: Yet, we can update the policy via gradient ascent.
Revised: Yet, we can update the policy via gradient ascent.

Original: The intuition is as follows: accordingto Corollary 2, the optimal policy update attains the maximum max a 0 2 A A ⇡ ✓ t ( s, a 0 )
Revised: The intuition is as follows: accordingto Corollary 2, the optimal policy update attains the maximum max a 0 2 A A ⇡ ✓ t ( s, a 0 )

Original:   ⇤ c ( a, a 0 )at each state.
Revised:   ⇤ c ( a, a 0 )at each state.

Original: Thus, we can steer the policy ⇡ ✓ to maximize
Revised: Thus, we can steer the policy ⇡ ✓ to maximize

Original: [Equation]
Revised: [Equation]

Original: In the particular case of a Gaussian policy with parametrized mean (and ﬁxed variance), combinedwith GAE estimate of the advantage function, one can maximize
Revised: In the particular case of a Gaussian policy with parametrized mean (and ﬁxed variance), combinedwith GAE estimate of the advantage function, one can maximize

Original: [Equation]
Revised: [Equation]

Original: Intuitively, this update implicitly estimates the transport maps T  ⇤ and T  ⇤ , which are needed forthe optimal policy update.
Revised: Intuitively, this update implicitly estimates the transport maps T  ⇤ and T  ⇤ , which are needed forthe optimal policy update.

Original: This way, we steer the policy network towards the optimal policy updatewithin the trust region.
Revised: This way, we steer the policy network towards the optimal policy updatewithin the trust region.

Original: Among others, this policy update allows for the following interpretation:imposing an optimal transport-based trust region constraints is, at least formally, equivalent tomaximizing a regularized advantage function, where the value of the regularization  ⇤  0 isbased on the transport cost c and the radius of the trust region " .
Revised: Among others, this policy update allows for the following interpretation:imposing an optimal transport-based trust region constraints is, at least formally, equivalent tomaximizing a regularized advantage function, where the value of the regularization  ⇤  0 isbased on the transport cost c and the radius of the trust region " .

Original:  Experiments and Insights
Revised:  Experiments and Insights

Original: In this section, we evaluate the performance of OT-TRPO across a variety of environments [5, 35]of increasing complexity, ranging from discrete to continuous settings.
Revised: In this section, we evaluate the performance of OT-TRPO across a variety of environments [5, 32]of increasing complexity, ranging from discrete to continuous settings.

Original: We compare it to theclassical TRPO [28, 9] and PPO [30, 9], with A2C [16, 9], with the recent approaches leveraging the
Revised: We compare it to theclassical TRPO [26, 7] and PPO [28, 7], with A2C [14, 7], with the recent approaches leveraging the

Original: Wasserstein distance, BGPG [21] and WNPG [19] (in continuous settings), and with WPO [32] (indiscrete settings).
Revised: Wasserstein distance, BGPG [19] and WNPG [17] (in continuous settings), and with WPO [30] (indiscrete settings).

Original: The training curves are shown in Fig.
Revised: The training curves are shown in Fig.

Original: 1; see Appendix A.3 for implementationdetails and Appendix A.6 for the scores of the trained agents.
Revised: 1; see Appendix A.3 for implementationdetails and Appendix A.5 for the scores of the trained agents.

Original: Our approach is shown to consistently improve over the other algorithms: OT-TRPO leads to betterﬁnal performance, with lower variance, in few cases at the expense of a slightly less steep learningcurve.
Revised: Our approach is shown to consistently improve over the other algorithms: OT-TRPO leads to betterﬁnal performance, with lower variance, in few cases at the expense of a slightly less steep learningcurve.

Original: Four remarks are in order.
Revised: Four remarks are in order.

Original: First, the performance gain of OT-TRPO compared to BGPG and
Revised: First, the performance gain of OT-TRPO compared to BGPG and

Original: WNPG conﬁrms that trust regions help stabilize training, as already observed in [28].
Revised: WNPG conﬁrms that trust regions help stabilize training, as already observed in [26].

Original: Second, optimaltransport discrepancies induce a more natural notion of “closeness” between polices compared to the
Revised: Second, optimaltransport discrepancies induce a more natural notion of “closeness” between polices compared to the

Original: KL divergence (e.g., see [3, Example 2.1]).
Revised: KL divergence (e.g., see [3, Example 2.1]).

Original: For instance, in CliffWalking-v0, consider the optimalpolicy ⇡ ⇤ and the candidate policy ⇡ depicted below, which differ at one state only (see ﬁgure).
Revised: For instance, in CliffWalking-v0, consider the optimalpolicy ⇡ ⇤ and the candidate policy ⇡ depicted below, which differ at one state only (see ﬁgure).

Original: [Figure]
Revised: [Figure]

Revised: WNPG

Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: The optimal transport discrepancy between ⇡ and ⇡ ⇤ is 335 ⇢ ⇡ ( s ) C ( ⇡ ( ·| s ) , ⇡ ⇤ ( ·| s ))
Revised: The optimal transport discrepancy between ⇡ and ⇡ ⇤ is 335 ⇢ ⇡ ( s ) C ( ⇡ ( ·| s ) , ⇡ ⇤ ( ·| s ))

Original: =
Revised: =

Original: ⇢ ⇡ ( s ) c ( Down ,
Revised: ⇢ ⇡ ( s ) c ( Down ,

Original: When using the KL divergence, instead, the discrepancy is inﬁnite, sincethe two policiesdo not  share the same support.
Revised: When using the KL divergence, instead, the discrepancy is inﬁnite, sincethe two policies not do share the same support.

Original: In particular, ifinitialized with ⇡ , TRPO cannot converge to the optimal policy,regardless of the radius of the trust region.
Revised: In particular, ifinitialized with ⇡ , TRPO cannot converge to the optimal policy,regardless of the radius of the trust region.

Original: Third, OT-TRPOimproves on WPO, in two ways: (i) it leads to superior performances of the trained policies and (ii) it does not violate the trust region constraint (which, e.g., in Taxi-v3 is the case for 72% of the updates of WPO).
Revised: Third, OT-TRPO improves on WPO, in two ways: (i) it leads to superior asymptoticperformance and (ii) it does not violate the trust region constraint (which, e.g., in Taxi-v3 is the casefor 72% of the updates of WPO).

Original: This performance improvement results from the“mass splitting” (see Remark 3), which is the only difference between the two algorithms (in discretesettings).
Revised: This performance improvement results from the “mass splitting”(see Remark 3), which is the only difference between the two algorithms (in discrete settings).

Original: Fourth, in continuous settings, the performance of OT-TRPO is aligned with the bestperforming alternative approach.
Revised: Fourth,in continuous settings, the performance of OT-TRPO is aligned with the best-performing alternativeapproach.

Original: In the environment Swimmer-v3, it even leads to an improvementof more than 50% in the ﬁnal performance.
Revised: In the environment Swimmer-v3, it even leads to an improvement of more than 50% in theﬁnal performance.

Original:  Conclusion and Future Work
Revised:  Conclusion and Future Work

Original: We studied trust region policy optimization for continuous state-action spaces whereby the trustregion is deﬁned in terms of a general optimal transport discrepancy.
Revised: We studied trust region policy optimization for continuous state-action spaces whereby the trustregion is deﬁned in terms of a general optimal transport discrepancy.

Original: Our analysis bases on a onedimensional convex dual reformulation of the optimization problem for the policy update which(i) enjoys strong duality and (ii) directly characterizes the optimal policy update, bypassing thecomputational burden of evaluating optimal transport discrepancies.
Revised: Our analysis bases on a onedimensional convex dual reformulation of the optimization problem for the policy update which(i) enjoys strong duality and (ii) directly characterizes the optimal policy update, bypassing thecomputational burden of evaluating optimal transport discrepancies.

Original: Moreover, we show that thepolicy update leads to monotonic improvements of the performance index.
Revised: Moreover, we show that thepolicy update leads to monotonic improvements of the performance index.

Original: Empowered by ourtheoretic results, we propose a novel algorithm, OT-TRPO, for trust region policy optimization withoptimal transport discrepancies.
Revised: Empowered by ourtheoretic results, we propose a novel algorithm, OT-TRPO, for trust region policy optimization withoptimal transport discrepancies.

Original: We evaluate its performance across several environments.
Revised: We evaluate its performance across several environments.

Original: Ourresults reveal that trust regions deﬁned by optimal transport discrepancies can offer advantages overthe KL divergence or non-trust region methods.
Revised: Ourresults reveal that trust regions deﬁned by optimal transport discrepancies can offer advantages overthe KL divergence or non-trust region methods.

Original: There are several research directions that merit further investigation.
Revised: There are several research directions that merit further investigation.

Original: We highlight two.
Revised: We highlight two.

Original: First, transportcosts provide us actionable knobs to shape the geometry of the trust region, and can be used to encodeprior knowledge on the environment or preferred exploration strategies.
Revised: First, transportcosts provide us actionable knobs to shape the geometry of the trust region, and can be used to encodeprior knowledge on the environment or preferred exploration strategies.

Original: Second, we would like tostudy the convergence properties of our algorithm.Checklist
Revised: Second, we would like tostudy the convergence properties of our algorithm.

Original: For all authors...

Original: (a) Do the main claims made in the abstract and introduction accurately reﬂect the pa-

Original: per’s contributions and scope?

Original: [Yes] All claims are in line with the paper and

Original: its

Original: contributions.

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

Original: [Yes] We discuss both beneﬁts and

Original: drawbacks of our methodology.

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

Original: [N/A] We do

Original: not believe our work entails potential negative societal impacts.

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

Original: them?

Original: [Yes] We carefully read ethics review guidelines.

Original: If you are including theoretical results...

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

Original: [Yes] All assumptionsare stated before the corresponding theoretic results.

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

Original: [Yes] All proofs are includedin the supplementary material.

Original: If you ran experiments...

Original: (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)?

Original: [Yes] Our code, with

Original: instructions, is part of the supplementary material.

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

Original: how they

Original: were chosen)?

Original: ] All training details are

Original: supplementary

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experi-

Original: ments multiple times)?

Original: [Yes] Our plots include error bars (mean ± standard deviation).

Original: (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)?

Original: [Yes] We comment on the computational

Original: time in the supplementary material.

Original: If you are using existing assets (e

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

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

Original: [Yes] We cited the creators

Original: of OpenAI, Mujoco (which we used for experiments), stable baselines (which provided

Original: algorithms), and the creators of state-of-the-art algorithms.

Original: (b) Did you mention the license of the assets?

Original: [N/A]

Original: The license of the assets is easily

Original: veriﬁable following its reference.

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

Original: [Yes]

Original: Our code is included in the supplementary material.

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

Original: using/curating?

Original: [N/A] All assets are publicly available or asked to the authors.

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

Original: information or offensive content?

Original: [N/A]

Original: The data we used does not contain personal

Original: information or offensive content.

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

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

Original: applicable?

Original: [N/A]

Original: We did not use crowdsourcing or conducted research with human

Original: subjects.

Original: (b) Did you describe any potential participant risks, with links to Institutional Review Board

Original: (IRB) approvals, if applicable?

Original: [N/A]

Original: We did not use crowdsourcing or conducted

Original: research with human subjects.

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount

Original: spent on participant compensation?

Original: [N/A]

Original: We did not use crowdsourcing or conducted

Original: research with human subjects.