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

Original: Policy Optimization (PO) algorithms have been proven particularly suited to handle the high-dimensionality of real-world continuous control tasks.
Revised: Policy Optimization (PO) algorithms have been proven particularly suited to handle the 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 stabilize the policy updates.
Revised: Trust Region Policy Optimization methods represent a popular approach to stabilize the policy updates.

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

Original: The Wasserstein distance represents a natural alternative, in place of the KL divergence, to define trust regions or to regularize the objective function.
Revised: The Wasserstein distance represents a natural alternative, in place of the KL divergence, to deﬁne trust regions or to regularize the 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, reducing the 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, reducing the applicability of the method.

Original: In this paper, we explore optimal transport discrepancies (which include the Wasserstein distance) to define trust regions, and we propose 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 we propose a novel algorithm – Optimal Transport Trust Region Policy Optimization (OT-TRPO) – for continuous state-action spaces.

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

Original: We then analytically derive the optimal policy update given the solution of the dual problem.
Revised: We then analytically 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 optimal transport maps, which we implicitly characterize by solving the dual formulation.
Revised: This way, we bypass the computation of optimal transport costs and of optimal transport maps, which we implicitly characterize by solving the dual formulation.

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

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

Original:  Introduction
Revised:  Introduction

Original: Reinforcement Learning (RL) has achieved outstanding results in numerous fields, from resource management [16], recommendation systems [42], and optimization of chemical reactions [44], to video-games [18, 43, 10] and board games [39], without sparing the world’s champion of GO [32].
Revised: Reinforcement Learning (RL) has achieved outstanding results in numerous ﬁelds, from resource management [16], recommendation systems [42], and optimization of chemical reactions [44], to video-games [18, 43, 10] and board games [39], without sparing the world’s champion of GO [32].

Original: Many of these successful applications rely on Policy Optimization (PO) algorithms, a family of RL methods that are particularly suited to handle the high-dimensionality of real-world control tasks.
Revised: Many of these successful applications rely on Policy Optimization (PO) algorithms, a family of RL methods that are particularly suited to handle the high-dimensionality of real-world control tasks.

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: In this context, the main challenge is to provide policy improvement guarantees.
Revised: In this context, the main challenge is to provide policy improvement guarantees.

Original: One remarkable option in this direction  is represented by Trust Region Policy Optimization (TRPO) [29], which constrains the optimization problem to policies that are “close” to the current one, whereby the Kullback-Leibler (KL) divergence is used as a similarity measure.
Revised: One remarkable option in this direction ⇤ Equal contribution. 36th is represented by Trust Region Policy Optimization (TRPO) [29], which constrains the optimization problem to policies that are “close” to the current one, whereby the Kullback-Leibler (KL) divergence is used as a similarity measure.

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

Original: Recent work [22, 33, 20] proposed to replace the KL divergence with the Wasserstein distance, a particular instance of optimal transport discrepancy (or cost).
Revised: Recent work [22, 33, 20] proposed to replace the KL divergence with the Wasserstein distance, a particular instance of optimal transport discrepancy (or cost).

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

Original: In particular, (i) optimal transport discrepancies allow us to compare probability measures (and thus policies) not sharing the same support (for which the KL divergence is infinity); and (ii) they encapsulate the geometry encoded by the transport cost in the action space: the discrepancy between two actions coincides with the discrepancy between the corresponding deterministic policies (whereas the KL divergence is again infinity).
Revised: In particular, (i) optimal transport discrepancies allow us to compare probability measures (and thus policies) not sharing the same support (for which the KL divergence is inﬁnity); and (ii) they encapsulate the geometry encoded by the transport cost in the action space: the discrepancy between 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 discrepancies particularly attractive for RL.
Revised: These reasons make optimal transport discrepancies particularly attractive for RL.

Original: However, the mere evaluation of optimal transport discrepancies entails solving a transportation problem (e.g., see [23, 34]), which poses significant computational challenges for its deployment.
Revised: However, the mere evaluation of optimal transport discrepancies entails solving a transportation problem (e.g., see [23, 34]), which poses signiﬁcant computational challenges for its deployment.

Original: Most of the previous work on the topic [22, 20] overcomes the computational burden via approximation, effectively changing the original problem.
Revised: Most of the previous work on the topic [22, 20] overcomes the computational burden via approximation, effectively changing the original problem.

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

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

Original: In our work, we consider optimal transport discrepancies to construct the trust region 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 trust region in settings where actions and states take value in general compact Polish spaces.

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

Original: We derive and leverage a dual reformulation of the PO problem to ensure an optimal policy update within the trust region, without any additional need for line searches (conversely to [29]).
Revised: We derive and leverage a dual reformulation of the PO problem to ensure an optimal policy update within the trust region, without any additional need for line searches (conversely to [29]).

Original: We circumvent the computation of optimal transport discrepancies via an analytical expression of the transport maps, which are characterized thanks to the dual reformulation.
Revised: We circumvent the computation 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 practical and efficient algorithm that encompasses both discrete and continuous settings.
Revised: Notably, our analysis enables a practical and 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 that strong duality holds for general compact metric state-action spaces.
Revised: We derive the dual of the optimal transport trust region policy optimization problem and we show that strong duality holds for general compact metric state-action spaces.

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

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

Original: We propose a novel PO algorithm for continuous spaces, O ptimal T ransport T rust R egion P olicy O ptimization (OT-TRPO).
Revised: We propose a novel PO algorithm for continuous spaces, O ptimal T ransport T rust R egion P olicy O ptimization (OT-TRPO).

Original: Herein, we leverage the derived duality theory to provide policy updates that satisfy the optimal transport discrepancy constraint while circumventing its computation.
Revised: Herein, we leverage the derived duality theory to provide policy updates that satisfy the optimal transport discrepancy constraint while circumventing its computation.

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

Original: Our results show the effectiveness of our approach for PO and the benefits of using optimal transport discrepancies.
Revised: Our results show the effectiveness of 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 RL and in particular in PO algorithms.
Revised: Optimal transport, and in particular the Wasserstein distance, has found various applications in RL and in particular in PO algorithms.

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

Original: In [22], the authors propose Behavior Guided Policy Gradient (BGPG), whereby they replace the KL divergence trust region from TRPO [29] by a Wasserstein distance penalty in a behavioral space.
Revised: In [22], the authors propose Behavior Guided Policy Gradient (BGPG), whereby they replace the KL divergence trust region from TRPO [29] by a Wasserstein distance penalty in a behavioral space.

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

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

Original: In [20], the authors further suggested incorporating additional information about the local behavior of policies encapsulated in the so-called Wasserstein Information Matrix, in the attempt to speed up the PO using a Wasserstein Natural Policy Gradient (WNPG).
Revised: In [20], the authors further suggested incorporating additional information about the local behavior of policies encapsulated in the so-called Wasserstein 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 Optimization (PPO) and
Revised: However, these approaches are relatively slow, compared to the traditional Proximal Policy Optimization (PPO) and

Original: TRPO.
Revised: 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 defined 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, the closest related work to ours is the recent paper [33] which studied Wasserstein Policy Optimization (WPO) for discrete action spaces.
Revised: Accordingly, the closest related work to ours is the recent paper [33] which studied Wasserstein Policy Optimization (WPO) for discrete action spaces.

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

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

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

Original: Optimization (DRO) [8, 19, 4, 25, 13].
Revised: Optimization (DRO) [8, 19, 4, 25, 13].

Original: Albeit our duality results are inspired from this literature,
Revised: Albeit our duality results are inspired from this literature,

Original: DRO is concerned with quantifying the worst-case risk of a cost functional over an ambiguity set of probability measures, which is a fundamentally different setting; see Remark 1.
Revised: DRO is concerned with quantifying the worst-case risk of a cost functional over an 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 TRPO in continuous settings: we exploit optimal transport theory to circumvent the computational burden of evaluating optimal transport discrepancies while still performing exact policy updates within the trust regions.
Revised: Conversely to all the previous work, we show how to perform exact optimal transport-based TRPO in continuous settings: we exploit optimal transport theory to circumvent the computational burden of evaluating optimal transport discrepancies while still performing exact policy updates within the trust 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 briefly introduce useful background and notation for the remainder of the paper.
Revised: We brieﬂy introduce useful background and notation 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 point x ∈ X is denoted by δ x .
Revised: The Dirac measure at some point x 2 X is denoted by  x .

Original: Given two Polish spaces X , Y , a Borel probability measures µ ∈ P ( X ) , and a Borel map T : X → Y , the pushforward measure of µ , denoted by T # µ , is defined by ( T # µ )( A ) :
Revised: Given two Polish spaces X , Y , a Borel probability measures µ 2 P ( X ) , and a Borel map T : X ! Y , the pushforward measure of µ , denoted by T # µ , is deﬁned by ( T # µ )( A ) :

Original: = µ ( T − 1 ( A )) for all A ∈ 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 finite set X coincides with the probability simplex and will also be denoted by P ( X ) .
Revised: The set of probability measures on a ﬁnite set X coincides with the probability simplex and will also be denoted by P ( X ) .

Original: For a given function f : X → R , the notation ∥ f ∥ ∞ refers to sup x ∈X | f ( x ) | .
Revised: For a given function f : X ! 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 infinite-horizon discounted Markov Decision Process (MDP)
Revised: We consider an inﬁnite-horizon discounted Markov Decision Process (MDP)

Original: [24] M = ( S , A , P , r, ρ, γ ) , where S is the state space,
Revised: [24] M = ( S , A , P , r, ⇢ ,  ) , where S is the state space,

Original: A is the action space, P : S×A×S → R ≥ 0 is the state transition probability kernel, r : S×A → R is the reward function, ρ is the initial state probability distribution, and γ ∈ [0 , 1) is the discount factor.
Revised: A is the action space, P : S ⇥ A ⇥ S ! R  0 is the state transition probability kernel, r : S ⇥ A ! R is the reward function, ⇢ is the initial state probability distribution, and

Original:  A randomized stationary Markovian policy, which we will simply call a policy in the rest of the paper, is a mapping π : S → P ( A ) specifying for each s ∈ S a probability measure over the set of actions A by π ( ·| s ) ∈ P ( A ) .
Revised: [Equation] is the discount factor. A randomized stationary Markovian policy, which we will simply call a policy in the rest of the paper, is a mapping ⇡ : S ! P ( A ) 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 π ∈ Π induces a discrete-time Markov reward process { ( s t , r ( s t , a t )) } t ∈ N , where s t ∈ S represents the state of the system at time t and r ( s t , a t ) corresponds to the reward received when executing action a t ∈ A in state s t .
Revised: Each policy ⇡ 2 ⇧ induces a discrete-time Markov reward process { ( s t , r ( s t , a t )) } t 2 N , where s t 2 S represents the state of the system at time t and r ( s t , a t ) corresponds to the reward received when executing action a t 2 A in state s t .

Original: We denote by P ρ,π the probability distribution of the Markov chain ( s t , a t ) issued from the MDP controlled by the policy π with initial state distribution ρ .
Revised: We denote by P ⇢ , ⇡ the probability distribution of the Markov chain ( s t , a t ) issued from the MDP controlled by the policy ⇡ with initial state distribution ⇢ .

Original: The associated expectation is denoted by E ρ,π and the notation E π is used whenever there is no dependence on ρ .
Revised: The associated expectation is denoted by 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 → R and the action-value function Q π : S ×
Revised: V ⇡ : S ! R and the action-value function Q ⇡ : S ⇥ A !

Original: A → R are defined for all s ∈ S , a ∈ A by V π ( s ) :
Revised: R are deﬁned for all s 2 S , a 2 A by V ⇡ ( s ) :

Original: =
Revised: =

Revised: E ⇡ [

Original: E π [ (cid:80) ∞ t =0 γ{+=
 +} t r ( s t , a t ) | s [-
-] = s ] and Q π ( s, a ) [-:= E π [ (cid:80) ∞ t =0 γ t r ( s t , a t ) | s 
-]
Revised: P 1 t {+=
 +} t r ( s t , a t ) | s [-
-]0 = s ] and Q ⇡ ( s, a ) [-:= E π [ (cid:80) ∞ t =0 γ t r ( s t , a t ) | s 
-]:

Original: =
Revised: =

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

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

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

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

Original: Given an initial state probability distribution ρ , our goal is to find a policy π maximizing the expected long-term returnJ
Revised: Given an initial state probability distribution ⇢ , our goal is to ﬁnd a policy ⇡ maximizing the expected long-term return

Original: [Equation]
Revised: [Equation]

Original: , which is well-defined 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 only have access to the observed state, action, and reward s t , a t , r t at each time step t , whereas the state transition kernel P is unknown.
Revised: To solve this PO problem, we only have access to the observed state, action, and reward s t , a t , r t at each time step t , whereas the state transition kernel P is unknown.

Original: When the state and action spaces ( S and A ) are finite, 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 [24, Theorem 6.11.11, p. 262]) under appropriate assumptions on the state and action spaces, the reward function and the transition kernel; we will explicit these later on.
Revised: When S and A are continuous, a (measurable) optimal policy is also guaranteed to exist (see [24, Theorem 6.11.11, p. 262]) under appropriate assumptions on the state and action spaces, the reward function and the transition kernel; we will explicit these later on.

Original: In this paper, we focus on the continuous state-action space setting and comment on the discrete (non necessarily finite) setting as a special case.
Revised: In this paper, we focus on the continuous state-action space setting and comment on the discrete (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 → R ≥ 0 , referred to as transport cost .
Revised: Consider a Polish space X and a continuous non-negative function c : X ⇥ X ! R  0 , referred to as transport cost .

Original: Let µ, ν ∈ P ( X ) and define the set of joint probability measures on X × X with marginals µ and ν :
Revised: Let µ, ⌫ 2 P ( X ) and deﬁne the set of joint probability measures on X ⇥ X with marginals µ and ⌫ :

Original: [Equation]
Revised: [Equation]

Original: We define 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 definition 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 Wasserstein distance [37, 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 Wasserstein distance [37, 2].

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

Original:  Optimal Transport for Trust Region Policy Optimization
Revised:  Optimal Transport for Trust Region Policy Optimization

Original: In this section, we study the TRPO algorithm with a trust region defined using an optimal transport discrepancy 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 transport discrepancy as a measure of closeness between policies.

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

Original: Importantly, we show that, given the dual optimal solution, the primal solution has an analytical expression, which can lead to monotonic improvements of the performance index.
Revised: Importantly, we show that, given the dual optimal solution, the primal solution has an analytical expression, which can lead to monotonic improvements of 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: By the policy difference lemma [11, Lemma 6.1], the difference between the expected returns of two policies π, ˜ π ∈ Π reads
Revised: By the policy difference lemma [11, Lemma 6.1], the difference between the expected returns of two policies ⇡ , ˜ ⇡ 2 ⇧ reads

Original: [Equation]
Revised: [Equation]

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

Original: The complex dependency of the discounted visitation frequency ρ ˜ π on the policy ˜ π hampers the direct optimization of (2); see [29, Sec. 2].
Revised: The complex dependency of the discounted visitation frequency ⇢ ˜ ⇡ on the policy ˜ ⇡ hampers the direct optimization of (2); see [29, Sec. 2].

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

Original: [Equation]
Revised: [Equation]

Original: Observe that this approximation uses the discounted state-occupancy measure ρ π (which can be estimated) instead of ρ ˜ π (see (2)).
Revised: Observe that this approximation uses the discounted state-occupancy measure ⇢ ⇡ (which can be estimated) instead of ⇢ ˜ ⇡ (see (2)).

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

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

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

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

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

Original: Then, at each time step, our method solves
Revised: Then, at each time step, our method solves

Original: [Equation]
Revised: [Equation]

Original: where ε > 0 is a parameter defining 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 [29, Eq.
Revised: Similarly to [29, Eq.

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

Original: Accordingly, the OT-TRPO policy optimization results from iteratively solving Problem (P).
Revised: Accordingly, the OT-TRPO policy optimization results from iteratively solving 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 is continuous, it is an infinite-dimensional optimization problem.
Revised: First, as soon as the state or action space is continuous, it is an inﬁnite-dimensional optimization problem.

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

Original: However, inspired by prior works on Wasserstein DRO [8, 19, 41, 4], we show that problem (P) admits a tractable one-dimensional convex dual reformulation.
Revised: However, inspired by prior works on Wasserstein DRO [8, 19, 41, 4], we show that problem (P) admits a tractable one-dimensional convex dual reformulation.

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

Original: Before stating the result, we make the following assumptions.
Revised: Before stating the result, we make the following assumptions.

Original: Assumption 1.
Revised: Assumption 1.

Original: The state space S is a compact subset of an Euclidean space, the action space A is a compact subset of a Polish space, the reward function r is a continuous function and for every continuous function w on S , (cid:82) S w ( u )dP( u | s, a ) is continuous in both s and a .
Revised: The state space S is a compact subset of an Euclidean space, the action space A is a compact subset of a Polish space, the reward function r is a continuous function and for every continuous function w on S , [Equation] u | s, a ) is continuous in both s and a .

Original: Under this assumption, there exists an optimal measurable (stationary) policy to the PO problem formulated in Section 3.
Revised: Under this assumption, there exists an optimal measurable (stationary) policy to the PO problem formulated in Section 3.

Original: We refer the reader to [24, Theorem 6.11.11, p. 262] for a statement of this result and milder assumptions.
Revised: We refer the reader to [24, Theorem 6.11.11, p. 262] for a statement of this result and milder assumptions.

Original: In particular, our duality result continues to hold if S is a compact
Revised: In particular, our duality result continues to hold if S is a compact

Original: Polish space (i.e., not necessarily Euclidean).
Revised: Polish space (i.e., not necessarily Euclidean).

Original: Assumption 2.
Revised: Assumption 2.

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

Original: A → R is continuous.
Revised: R is continuous.

Original: Moreover, the transport cost c : A × A → R ≥ 0 is continuous and satisfies c ( a, a ) = 0{+
+} for all a ∈ A .
Revised: Moreover, the transport cost c : A ⇥ A ! R  0 is continuous and satisﬁes c ( a, a ) = {+
+} for all a 2 A .

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

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

Original: For every ε > 0 and for every policy π ∈ Π , under Assumptionsand 2 the following strong duality result holds:
Revised: For every " > 0 and for every policy ⇡ 2 ⇧ , under Assumptionsand 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: In particular, we derive a closed-form solution of problem (P) as a function of the optimal Lagrange multiplier λ ∗ solving problem (D) and the policy π defining the problem.
Revised: In particular, we derive a closed-form solution of problem (P) as a function of the optimal Lagrange multiplier  ⇤ solving problem (D) and the policy ⇡ deﬁning the problem.

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

Original: To do so, we introduce some additional notation, which is instrumental to derive a practical algorithm (similarly to [33]) [Equation]
Revised: To do so, we introduce some additional notation, which is instrumental to derive a practical algorithm (similarly to [33]).

Original: Under Assumptions 1 and 2, define for every [-λ ≥ 
-] the λ -regularized advantage Φ λ : S ×
Revised: Under Assumptions 1 and 2, deﬁne for every [-λ ≥ 
-]  0 the  -regularized advantage   : S ⇥ A !

Original: A → R and its associated set of maximizers for every s ∈ S [Equation] a ∈ A as follows:
Revised: R 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
Revised: Under the setting and assumptions of Theorem

Original: [Equation]
Revised: [Equation]

Original: for any policy π ∈ Π [Equation] let λ ∗ ≥ 0 be the minimizer of the dual problem (D) .
Revised: 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 → A and T λ : S×A
Revised: For every   0 , there exist two measurable selection maps T  : S ⇥ A ! A and T  : S ⇥ A !

Original: → A such that for every s∈ S , a ∈ A
Revised: A such that for every s

Original: [Equation]
Revised: [Equation]

Original: If λ ∗ > 0 , there exists t ∗ ∈ [0 , 1] such that 
Revised: If  ⇤ > 0 , there exists t ⇤ 2 [0 , 1] such that [Equation] S

Original: [Equation]
Revised: [Equation]

Original: There exists an optimal feasible policy ˜ π for problem (P) defined for every s ∈ 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: where t ∗ results from (6) if λ ∗ > 0 and t ∗ = 0{+
+} if λ ∗ = 0 .
Revised: where t ⇤ results from (6) if  ⇤ > 0 and t ⇤ = {+
+} if  ⇤ = 0 .

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 optimal dual 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 optimal dual solution.

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

Original: In the special case λ ∗ = 0 , the trust region constraint is either not active (i.e., the optimal policy lies within the trust region) or it does not affect the optimal trust region constraint (i.e., the optimal policy would lie at the boundary of the trust region even if the constraint is removed).
Revised: In the special case  ⇤ = 0 , the trust region constraint is either not active (i.e., the optimal policy lies within the trust region) or it does not affect the optimal trust region constraint (i.e., the optimal policy would lie at the boundary of the trust region even if the constraint is removed).

Original: In this case, it suffices to displace all probability mass to the closest maximizer of the advantage functions (i.e., T 0 ( s, a ) ∈ D 0 ( s, a ) ).
Revised: In this case, it sufﬁces to displace all probability mass to the closest maximizer of the advantage functions (i.e., T 0 ( s, a ) 2 D 0 ( s, a ) ).

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 [8, Theorem 1] and [4, Theorem 1]).

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

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

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

Original: s ).
Revised: s ).

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

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

Original: To transform (P) into a DRO, one might alternatively identify a policy as a probability measure over (cid:81) s ∈S A , and hope to deploy standard duality arguments in DRO.
Revised: To transform (P) into a DRO, one might alternatively identify a policy as a probability measure over Q s 2 S A , and hope to deploy standard 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 [37], 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 [37], inapplicable.

Original: Remark 2 .
Revised: Remark 2 .

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

Original: We highlight four major differences.
Revised: We highlight four major differences.

Original: First, in the discrete setting, problem (P) is a finite-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 derive the dual reformulation.
Revised: Thus, linear programming arguments can be used to derive the dual reformulation.

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

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

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

Original: Third, the optimal policy update [33] implicitly assumes that the set D λ ( s, a ) (see (4)) is a singleton, which is rarely satisfied in practice.
Revised: Third, the optimal policy update [33] implicitly assumes that the set D  ( s, a ) (see (4)) is a singleton, which is rarely satisﬁed in practice.

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

Original: This way, we do not need to resort to approximation techniques [33, Section 6.1] for the optimal dual multiplier.
Revised: This way, we do not need to resort to approximation techniques [33, Section 6.1] for the optimal dual multiplier.

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

Original: In the remainder of this section, we specialize our results to discrete (finite) state-action spaces (which trivially satisfy Assumptions 1 and 2).
Revised: In the remainder 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, we represent the state and action spaces by S
Revised: Without loss of generality, we represent 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 are two positive integers, and we describe any policy π ∈ Π and its corresponding state-occupancy measure ρ π as discrete measures:
Revised: . , a N } where M and N are two positive integers, and we describe any policy ⇡ 2 ⇧ and its corresponding state-occupancy measure ⇢ ⇡ as discrete measures:

Original: [Equation]
Revised: [Equation]

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

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

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

Original: , (cid:80) Mi 
Revised: P Mi =1 ⇢ i = 1 and P

Original: =1 ρ i = 1 and (cid:80) Nj = 1 π i,j = 1 for every i ∈ { 1 , . . .
Revised: Nj =1 ⇡ i,j = 1 for every 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 π ∈ Π , the following strong duality result holds:
Revised: For every policy ⇡ 2 ⇧ , the following strong duality result holds:

Original: [Equation]= [Equation]
Revised: [Equation]

Original: discrete [Equation] In particular, let λ ∗ ≥ 0 be a solution to (discrete-D) , and given for every i ∈ { 1 , . . .
Revised: In particular, let  ⇤  0 be a solution to (discrete-D) , and given for every i 2 { 1 , . . .

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

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

Original: [Equation][Equation]
Revised: [Equation]

Original: and let c :=M (cid:88) i = ρ i
Revised: and let c :=

Original: N (cid:88) j

Original: = π i,j c ( a j , b ∗ i,j ) , c :=

Original: M (cid:88) i =

Original: [Equation]ρ [Equation]
Revised: [Equation]

Original: (cid:88) j = [Equation] c (
Revised: Then, an optimal policy ˜ ⇡ is given by

Original: [Equation]∗ i [Equation]
Revised: [Equation]

Original: Then, an optimal policy ˜ π is given by

Original: [Equation]
Revised: [Equation]

Original: with t ∗ =
Revised: with t ⇤ =

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

Original: ( c − c ) ∈ [0 , 1] (and t ∗ ∈ [0 , 1] if c = c = ε ) if λ ∗ > 0 and t ∗ = 

Revised: ( c  c ) 2 [0 , 1] (and t ⇤ 2 [0 , 1] if c = c = " ) if  ⇤ > 0 and t ⇤ =

Original: if λ ∗ = 0 .
Revised: if  ⇤ = 0 .

Original: The proof of this result stems from substituting the discrete measures as defined in (8) in problems (P) and (D), and observing that the images of the mappings T λ ∗ and T λ ∗ have finite support in the current setting.
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 support in the current setting.

Original: Notably, Corollary 3 directly provides an implementable algorithm for the policy update, circumventing the difficulty of the mixed-integer optimization problem (discrete-P): solving the one-dimensional convex program (discrete-D) provides the optimal Lagrange multiplier associated to the trust region constraint of the primal problem which can be directly used to compute the actions b ∗ i,j , b ∗ i,j via (9), and thus the policy ˜ π via (10).
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 the one-dimensional convex program (discrete-D) provides the optimal Lagrange multiplier associated to the trust region constraint of the primal problem which can be directly used to compute the actions 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 [33] (see [33, Theorem 1, (5)] where f ∗ s ( i, j ) ∈ { 0 , 1 } with their notation).
Revised: The policy update suggested by (10) differs from the one in [33] (see [33, Theorem 1, (5)] where f ⇤ s ( i, j ) 2 { 0 , 1 } with their notation).

Original: Indeed, our policy update relies on “splitting the probability mass”: the probability mass π ( ·| s i ) is displaced to b ∗ i,j and b ∗ i,j with weights t ∗ and 1 − t ∗ , respectively.
Revised: Indeed, our policy update relies on “splitting the probability mass”: the probability mass ⇡ ( ·| s i ) is displaced to b ⇤ i,j and b ⇤ i,j with weights t ⇤ and 1  t ⇤ , 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 [33] corresponds to the particular case where b ∗ i,j = b ∗ i,j which amounts to supposing that the set D λ ( s i , a j ) defined in (4) is a singleton.
Revised: The result provided in [33] corresponds to the particular case where b ⇤ i, j = b ⇤ i,j which amounts to supposing that the set D  ( s i , a j ) deﬁned in (4) is a singleton.

Original: We provide further comments and examples in Appendix A.2 to illustrate the importance of this “mass splitting”.
Revised: We provide further comments and examples 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 π ∈ Π .
Revised: Let ⇡ 2 ⇧ .

Original: Consider solutions ˜ 
Revised: Consider solutions ˜ ⇡

Original: π ∗ ∈ Π and λ ∗ ≥of problems (P) and (D) , respectively.
Revised: ⇤ 2 ⇧ and  ⇤  0 of problems (P) and (D) , respectively.

Original: If the true advantage function A π is approximated by some estimated advantage function ˆ A π such that ∥ A π − ˆ A π ∥ ∞ < ∞ , then the following bound holds:
Revised: If the true advantage function A ⇡ is approximated by some estimated 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 constructed in 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 constructed in the proof of Theorem 1 (see Corollary 2).

Original: The analog of this result for a finite action space was proved in [33, Theorem 2, p. 5].
Revised: The analog of this result for a ﬁnite action space was proved in [33, Theorem 2, p. 5].

Original: To the best of our knowledge, this result is novel for the continuous state-action space setting.
Revised: To the best of our knowledge, this result is novel for the continuous state-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 a vector θ ∈ R d for some integer d > 0 .
Revised: Herein, we restrict the policy search set ⇧ to the set of policies ⇡ ✓ parametrized by a vector ✓ 2 R d for some integer d > 0 .

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

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

Original: Accordingly, the dual problem (D) can be reformulated as follows for every θ ∈ R d :
Revised: Accordingly, the dual 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 first 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 use the optimal dual multiplier λ ∗ to derive the optimal policy update within the trust region.
Revised: Then, we use the optimal dual multiplier  ⇤ to derive the optimal policy update within the trust region.

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

Original: Depending on the parametrization of the policy, the steps of Section 5 are as follows:
Revised: Depending on the parametrization of the policy, the steps of Section 5 are as follows:

Original: Algorithm 
Revised: Algorithm 1 OT

Original:  1: Initialize π θ2: for all t = 0 , 1 , . . .
Revised: TRPO. 1: Initialize ⇡ ✓2: for all t = 0 , 1 , . . .

Original: do 3:
Revised: do 3:

Original: Estimate A π θt and ρ π θt .
Revised: Estimate A ⇡ ✓ t and ⇢ ⇡ ✓ t .

Original: Compute λ ∗ ∈ argmin λ ≥ 0 G ( λ, θ t ) .
Revised: Compute  ⇤ 2 argmin   0 G (  , ✓ t ) .

Original: Update θ t to θ t +1 using λ ∗ .
Revised: Update ✓ t to ✓ t +1 using  ⇤ .

Original: 6: end for
Revised: 6: end for

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

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

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 to estimate the unknown advantage function via samples 3 :
Revised: We propose three ways to estimate the unknown advantage function via samples 3 :

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

Original: General Advantage Estimation (GAE) [30], using a neural network to approximate the value function like in standard actor-critic methods; and 3.
Revised: General Advantage Estimation (GAE) [30], using a neural network to approximate the value function like in standard actor-critic methods; and 3.

Original: Direct estimation via non-linear approximators (e.g., using directly a neural network for the advantage function).
Revised: Direct estimation via non-linear approximators (e.g., using directly a neural network for the 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 to evaluate G ( λ, θ ) .
Revised: Depending on the setting, we propose various ways to evaluate 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 maximization in (12) is over finitely many actions
Revised: Finite actions : Since the maximization in (12) is over ﬁnitely many actions, we can directly evaluate (12) for any   0 . 2.

Original: , we can directly evaluate (12) for any λ ≥ 0 . Gaussian policy parametrization : With m ( s ) being the mean of the Gaussian policy (with fixed variance), we approximate G ( λ, θ ) by
Revised: Gaussian policy parametrization : With m ( s ) being the mean of the Gaussian policy (with ﬁxed variance), we approximate G (  , ✓ ) by

Original: [Equation]
Revised: [Equation]

Original: if A π θ via NN , where ˆ S are the state visited in the trajectory and ˆ A ( s ) is a (possibly state-dependent) collection of samplescollected from A . 
Revised: where ˆ S are the states visited in the trajectory and ˆ A ( s ) is a (possibly state-dependent) collection of samples from A . 3.

Original: Arbitrary policy parametrization : For a neural network approximation of the advantage function, we approximate G ( λ, θ ) by
Revised: Arbitrary policy parametrization : For a neural network approximation of the advantage function, we approximate G (  , ✓ ) by

Original: [Equation]
Revised: [Equation]

Original: where ˆ S are the state visited in the trajectory and ˆ A i ( s ) are (possibly state-dependent) collections of samplescollected from A .
Revised: where ˆ S are the states visited in the trajectory and ˆ A i ( s ) are (possibly state-dependent) collections of samples from A .

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

Original: Since (D) is a one-dimensional convex optimization problem, λ ∗ can be found using any solver for convex optimization problems.
Revised: Since (D) is a one-dimensional convex optimization problem,  ⇤ can be found using any solver for convex optimization problems.

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

Original: Update the parameter vector θ .
Revised: Update the parameter vector ✓ .

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

Original: Direct parametrization 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: Direct parametrization 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 policy network by performing gradient descent on the loss L ( θ ) = (cid:80) s ∈ ˆ S ρ π θt ( s ) (cid:13)(cid:13)
Revised: Then, update the policy network by performing gradient descent on the loss L ( ✓ ) = [Equation]  [Equation]

Original: π θ ( ·| s ) − π ∗ θ t ( ·| s ) (cid:13)(cid:13)to steer π θ towards the optimal policy update π ∗ θ t within the trust region.
Revised: to steer ⇡ ✓ towards the optimal policy update ⇡ ⇤ ✓ t within the trust region.

Original: Arbitrary policy parametrization : Since there are infinitely many actions, the computation
Revised: Arbitrary policy parametrization : Since there are inﬁnitely many actions, the computation

Original: of the maximization is computationally demanding
Revised: of the maximization is computationally demanding

Original: , and 
Revised: and so

Original: so Corollary 2 cannot be directly utilized for the policy update.
Revised: Corollary 2 cannot be directly utilized for the policy update

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: according to Corollary 2, the optimal policy update attains the maximum max a ′ ∈A { A π θt ( s, a ′ )
Revised: The intuition is as follows: according to Corollary 2, the optimal policy update attains the maximum max a 0 2 A { A ⇡ ✓ t ( s, a 0 )

Original: − λ ∗ c ( a, a ′ ) } 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 fixed variance), combined with GAE estimate of the advantage function, one can maximize
Revised: In the particular case of a Gaussian policy with parametrized mean (and ﬁxed variance), combined with 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 for the optimal policy update.
Revised: Intuitively, this update implicitly estimates the transport maps T  ⇤ and T  ⇤ , which are needed for the optimal policy update.

Original: This way, we steer the policy network towards the optimal policy update within the trust region.
Revised: This way, we steer the policy network towards the optimal policy update within 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 to maximizing a regularized advantage function, where the value of the regularization λ ∗ ≥ 0 is based 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 to maximizing a regularized advantage function, where the value of the regularization  ⇤  0 is based 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, 36] 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, 36] of increasing complexity, ranging from discrete to continuous settings.

Original: We compare it to the classical TRPO [29, 9] and PPO [31, 9], with A2C [17, 9], with the recent approaches leveraging the Wasserstein distance, BGPG [22] and WNPG [20] (in continuous settings), and with WPO [33] (in discrete settings).
Revised: We compare it to the classical TRPO [29, 9] and PPO [31, 9], with Advantage Actor Critic (A2C) [17, 9], with the recent approaches leveraging the Wasserstein distance, BGPG [22] and WNPG [20] (in continuous settings), and with WPO [33] (in discrete settings).

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

Original: 1; see Appendix A.3 for implementation details, Appendix A.6 for further details on the experimental results, and Appendix A.4 for an ablation study on our algorithm.
Revised: 1; see Appendix A.3 for implementation details, Appendix A.6 for further details on the experimental results, and Appendix A.4 for an ablation study on our algorithm.

Original: [Figure]
Revised: [Figure]

Original: OT-TRPO

Original: [Equation]
Revised: [Equation]

Original: Our approach is shown to consistently improve over the other algorithms: OT-TRPO leads to larger final returns, with lower variance, and only in few cases at the expense of a slightly slower learning curve.
Revised: Our approach is shown to consistently improve over the other algorithms: OT-TRPO leads to larger ﬁnal returns, with lower variance, and only in few cases at the expense of a slightly slower learning curve.

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 confirms that trust regions help stabilize training, as already observed in [29].
Revised: WNPG conﬁrms that trust regions help stabilize training, as already observed in [29].

Original: Second, optimal transport discrepancies induce a more natural notion of “closeness” between policies compared to the
Revised: Second, optimal transport discrepancies induce a more natural notion of “closeness” between policies 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 optimal policy π ∗ and the candidate policy π depicted below, which differ at one state only (see figure).
Revised: For instance, in CliffWalking-v0, consider the optimal policy ⇡ ⇤ and the candidate policy ⇡ depicted below, which differ at one state only (see ﬁgure).

Original: [Figure]
Revised: [Figure]

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

Original: =
Revised: =

Original: ρ π ( s ) c ( Down , Right ) .
Revised: ⇢ ⇡ ( s ) c ( Down , Right ) .

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

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

Original: Third, OT-TRPO improves 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 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).

Original: This performance improvement results from the “mass splitting” (see Remark 3), which is the only difference between the two algorithms (in discrete settings).
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 best-performing alternative approach.
Revised: Fourth, in continuous settings, the performance of OT-TRPO is aligned with the best-performing alternative approach.

Original: In the environment
Revised: In the environment

Original: Swimmer-v3, it even yields an improvement of more than 50% in the performance of the trained agent.
Revised: Swimmer-v3, it even yields an improvement of more than 50% in the performance of the trained agent.

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

Original: We studied trust region policy optimization for continuous state-action spaces whereby the trust region is defined in terms of a general optimal transport discrepancy.
Revised: We studied trust region policy optimization for continuous state-action spaces whereby the trust region 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 the computational 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 the computational burden of evaluating optimal transport discrepancies.

Original: Moreover, we show that the policy update can yield a monotonic improvement of the performance index.
Revised: Moreover, we show that the policy update can yield a monotonic improvement of the performance index.

Original: Empowered by our theoretic results, we propose a novel algorithm, OT-TRPO, for trust region policy optimization with optimal transport discrepancies.
Revised: Empowered by our theoretic results, we propose a novel algorithm, OT-TRPO, for trust region policy optimization with optimal transport discrepancies.

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

Original: Our results reveal that trust regions defined by optimal transport discrepancies can offer advantages over the KL divergence or non-trust region methods.
Revised: Our results reveal that trust regions deﬁned by optimal transport discrepancies can offer advantages over the 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, transport costs provide us actionable knobs to shape the geometry of the trust region, and can be used to encode prior knowledge on the environment or preferred exploration strategies.
Revised: First, transport costs provide us actionable knobs to shape the geometry of the trust region, and can be used to encode prior knowledge on the environment or preferred exploration strategies.

Original: Second, we would like to study the convergence properties of the proposed algorithm.
Revised: Second, we would like to study the convergence properties of the proposed algorithm.

Original: Acknowledgements
Revised: Acknowledgements

Original: This project has received funding from Google Brain, Swiss National Science Foundation under theNCCR Automation (grant agreement 51NF40_180545), and it was partially supported by the ETH AI Center. Checklist
Revised: This project has received funding from Google Brain, Swiss National Science Foundation under the

Original: For all authors..

Original: (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?

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

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

Original: [Yes] We discuss both benefits and drawbacks of our methodology.

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

Original: [N/A] We do 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 them?

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

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

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

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

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

Original: If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?

Original: [Yes] Our code, with instructions, is part of the supplementary material.

Original: (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?

Original: [Yes] All training details are specified in the supplementary material.

Original: (c) Did you report error bars (e.g., with respect to the random seed after running experiments 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.g., type of GPUs, internal cluster, or cloud provider)?

Original: [Yes] We comment on the computational time in the supplementary material.

Original: If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?

Original: [Yes] We cited the creators of OpenAI, Mujoco (which we used for experiments), stable baselines (which provided 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 verifiable following its reference.

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

Original: [Yes] 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 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 identifiable information or offensive content?

Original: [N/A]

Original: The data we used does not contain personal information or offensive content.

Original: If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?

Original: [N/A]

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

Original: (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?

Original: [N/A]

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

Original: (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?

Original: [N/A]

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

Original: Appendix A Appendix A.1 More details on related work

Original: Optimal transport and in particular the Wasserstein distance has found various applications in RL, despite the computational challenge raised by its evaluation.

Original: Optimal transport for PO.

Original: In [26], the authors established a connection between policy gradient in Wasserstein trust regions and variational optimal transport, suggesting to solve the Fokker-Planck equations and to use diffusion processes.

Original: Concurrently, [40] formulated the PO problem as a gradient descent flow on the space of probability measures using Wasserstein Gradient Flows [2] which is then solved approximately via particles.

Original: BGPG [22].

Original: In [22], the authors proposed to replace the KL divergence trust region from TRPO [29] by a Wasserstein distance penalty in a behavioral space; the proposed approach embeds the policies in a latent behavioral space via a map acting on the trajectories induced by the policies, and leverages the Wasserstein distance to compare two embeddings.

Original: To approximate the Wasserstein distance, they exploit the dual formulation of the entropy-regularized Wasserstein distance.

Original: Although our alternating procedure in Section 5 may formally resemble the approach of [22] in spirit, our approach (i) builds on the idea of trust regions and (ii) does not consider policy embeddings or entropy regularization.

Original: Not surprisingly, our Algorithm 1 is fundamentally different from Algorithms 1 and 3 in [22].

Original: If

Original: BGPG can be seen as a variant of TRPO with a Wasserstein regularization, we propose a novel distinct variant of OT-based TRPO fully based on trust regions.

Original: WNPG [20].

Original: While [22] proposed to use the Wasserstein distance as a global penalty to the objective, [20] further suggested to incorporate additional information about the local behavior of policies encapsulated in the so-called Wasserstein Information Matrix.

Original: They proposed accordingly

Original: WNPG to speed up policy optimization using a Wasserstein natural gradient.

Original: The cornerstone of this algorithm is the estimation of a Wasserstein natural gradient stemming from a second-order expansion of the 2-Wasserstein distance between two (parametric) behavioral embedding distributions of two parameterized policies.

Original: WPO [33].

Original: The closest related work to ours is the recent paper [33] which studied PO with

Original: Wasserstein distance and Sinkhorn divergence-based trust regions for discrete (and finite) action spaces.

Original: In contrast to [33], the present work addresses the general setting of compact subsets of Polish spaces encompassing the cases of continuous and discrete state-action spaces as particular cases.

Original: While we also adopt a duality approach, our level of generality induces many challenges compared to the discrete action space setting (see Remark 2).

Original: Moreover, the methods proposed in [33] do not allow to restrict the set of policies to a particular family of distributions, which might lead to very large models.

Original: Conversely, the novel algorithm we propose can handle policy parametrization (including direct parametrization).

Original: We refer the reader to Section 4 for more detailed comparisons with [33].

Original: A.2 Comments on “mass splitting”

Original: In this section, we elaborate on the concept of “mass splitting”.

Original: Specifically, we propose two examples showing that splitting the probability mass is in general required to construct optimal policy updates.

Original: Example 1 .

Original: Let S =

Original: { s 1 } , A

Original: = { a 1 , a 2 } , and ε ∈ (0 , 1) .

Original: Suppose c ( a 1 , a 2 ) = 1 , c ( a 1 , a 1 ) =

Original: c ( a 2 , a 2 ) = 0 , π ( ·| s 1 )

Original: =

Original: [-δ a 1 , A π ( s 1 , a 1 ) = 
-]
Revised: [-δ a 1 , A π ( s 1 , a 1 ) = 
-]

Original: and A π ( s 1 , a 2 ) = 1 .

Original: Clearly, it is optimal to assigns as much probability mass as possible to action a 2 ; i.e., ˜ π ( ·| s 1 ) =

Original: (1 − ε )

Original: δ a 1 + εδ a 2 , which corresponds to t ∗ = 1 − ε , b 1 , 1 = a 1 and b 1 , 2 = a 2 using the notation of Corollary 3.

Original: However, if an optimal policy was to be described by a single transport map (as in [33, Theorem 1, Eq.

Original: (5), p. 4]), the only possible solutions would be π 1 ( ·| s 1 ) =

Original: δ a 1 and π 2 ( ·| s 1 )

Original: =

Original: δ a 2 , which are respectively sub-optimal and infeasible.

Original: Example 2 .

Original: Consider an agent in an initial state s 0 who can move left (L) or right (R).

Original: The rewards are r ( s 0 , L) = − 1 , and r ( s 0 , R)

Original: = +1 , and the task terminates whenever the agents reaches s 1 or at s 2 , as shown below.

Original: [Equation]

Original: Consider the initial policy π 0 (L | s 0 )

Original: = 1 and π 0 (R | s 0 )

Original: = 0 , and trust region defined by the binary distance, and let ε ∈ (0 , 1) .

Original: Then, Corollary 3 yields the new policy π 1 (L | s )

Original: [-= 1 − ε and π 
-]
Revised: [-= 1 − ε and π 
-]

Original: (R | s )

Original: = ε .

Original: Note that this update requires “mass splitting”.

Original: A second update yields π 2 (L | s )

Original: = 1 − 2 ε and π 2 (R | s )

Original: = 2 ε , and after 1 /ε the routine converges to the optimal policy π ∗ (L | s )

Original: [-= 
-]
Revised: [-= 
-]

Original: and π ∗ (R | s ) = 1 .

Original: Conversely, the update of [33] (i.e., without “mass splitting”) leads to either π 1 = π 0 , which is suboptimal, or to π 1 = π ∗ , which violates the trust region constraint.

Original: A.3 Implementation details

Original: In this section, we present the implementation details supporting our experimental results.

Original: We now discuss the details concerning the implementation and tuning of our algorithm. Our code, along with the instructions to set up the environment and run it, is available at https://gitlab. ethz.ch/lnicolas/ot-trpo .
Revised: NCCR Automation (grant agreement 51NF40_180545), and it was partially supported by the ETH

Original: Discrete settings.

Original: For the discrete settings, we developed a custom training and testing framework.

Original: We used TD-learning to estimate the advantage function, with learning rate α and discount factor γ .

Original: Namely, we first collect a set T of trajectory.

Original: [-Then, we initialize Q ( s, a ) = 
-]
Revised: [-Then, we initialize Q ( s, a ) = 
-]

Original: for all ( s, a ) ∈ S × A and for all τ ∈ T and every ( s t , a t , r t , s tt , a tt ) ∈

Original: τ we update Q as

Original: Q ( s t , a t ) =

Original: (1 − α ) Q ( s t , a t ) + α ( r t + γQ ( s tt , a tt )) .

Original: The transport cost used is the binary distance 4 (i.e., c ( a, a ′ )

Original: [-= 
-]
Revised: [-= 
-]

Original: if a = a ′ and c ( a, a ′ )

Original: = 1 otherwise), and the trust region radius is ε .

Original: Every n full environment simulations the policy is updated.

Original: The parameters for the different environments are reported in Table 1.

Original: [Table]

Original: With respect to the computational complexity of the training process, one can assess the following about the main steps of Algorithm 1:

Original: Step 3 .

Original: The complexity of TD-learning is linear in the number N of samples collected during the rollout of the current policy; i.e., O ( N ) .

Original: Step 4 .

Original: This step boils down to scipy.optimize.minimize_scalar 5 .

Original: The complexity of each function call can be shown to be O ( |S||A| 2 ) .

Original: Step 5 .

Original: The complexity of this step can be shown to be O ( |S||A| 2 ) .

Original: Continuous settings.

Original: For the continuous settings, we developed an agent which can be interfaced with stable baselines [9].

Original: In the benchmark provided, we estimate the advantage function via the GAE, with coefficient gae_lambda .

Original: We use the policy and value network provided in the

Original: ActorCriticPolicy class from [9], with network sizes [64 , 64] and the activation activation_fn .

Original: Accordingly, the policy is a Gaussian policy with fixed standard deviation whose mean is expressed by the policy network.

Original: That is, we do not use a neural network to approximate the advantage function.

Original: We perform stochastic gradient descent in batches of size batch_size , every n_steps timesteps, for n_epochs epochs, with learning rate learning_rate , and with maximum gradient norm max_grad_norm .

Original: The value function loss is multiplied by vf_coef .

Original: The transport cost is the square Euclidean distance 6 , and the trust region radius is ε .

Original: Along the lines of the implementation of PPO [31, 9], we introduce a Z-normalization of the advantage estimates if normalize is True , a state dependent exploration with sample frequency sde_sample_freq

Original: if this value is not − 1 , and an orthogonal initialization if ortho_init is True .

Original: The logarithm of the standard deviation of the Gaussian policy network is initialized to log_std_init .

Original: The parameters for the different environments are reported in Table 2.

Original: [Table]

Original: With respect to the computational complexity of the training process, one can assess the following about the main steps of Algorithm 1:

Original: Step 3 .

Original: The complexity of using a GAE is linear in the number N of samples collected during the rollout of the current policy; i.e., O ( N ) .

Original: Step 4 .

Original: This step boils down to scipy.optimize.minimize_scalar .

Original: The complexity of each function call can be shown to be O ( N ) .

Original: Step 5 .

Original: This corresponds to a gradient descent step.

Original: Therefore, the computational costs are in line with the algorithms to which we compare.

Original: A.3.2 Compared algorithms details

Original: We compare our algorithm to several baselines, which we implemented and tuned as follows:

Original: • The implementation, hyperparameters, and network sizes for TRPO [29], PPO [31], and A2C [17] can be

Original: found in stable baselines [9].
Revised: AI Center.

Original: • We thank the authors of [22] for privately providing us the code and tuning for BGPG.

Original: • For WNPG [20], we used the code and tuning publicly available at https://github.com/ tedmoskovitz/WNPG .

Original: • We thank the authors of [33] for pointing us to their implementation (with tuning parameters) of WPO, available at https://github.com/efficientwpo/EfficientWPO .

Original: This code has

Original: inspired in some parts our training and testing framework.

Original: Their primal update has been ported in our code to properly compare the only difference in the algorithms (in discrete settings).

Original: We use the same parameters for both WPO and OT-TRPO.

Original: To the best of our knowledge, we used the hyperparameters and network sizes with the best performances for every algorithm in every environment.

Original: The results obtained for the baselines considered are comparable or better to the ones found in the literature; e.g., see [33].

Original: We refrained from reporting the results of the algorithms whenever they were not reproducible or not comparable to the others.

Original: For BGPG [22] and WNPG [20], we used tensorflow.compat to ensure compatibility of their code (written with TensorFlow 1.x) with our installation of TensorFlow 2.x.

Original: The code reproduces the results reported in these works, but the comparisons provided in this paper refer to the updated environments; e.g., Hopper-v3 instead of Hopper-v2 .

Original: Finally, to the best of our knowledge, the timescale of the results obtained with the code in https: //github.com/efficientwpo/EfficientWPO and the results ported in our work might differ.

Original: In the former, the evaluation is performed after a certain number of episodes (full environment simulation), rather than timesteps, and the conversion is not clear from the code (as every episode might differ in length).

Original: We instead adhere to the standard provided in [9].

Original: A.4 Ablation study

Original: We investigate how different (i) trust region radii ε and (ii) transportation costs, affect the learning performances in the Taxi-v3 environment.

Original: These effects are shown in Fig.

Original: [Figure]

Original: Different trust region radii ε .

Original: As expected, larger radii lead to significantly steeper learning curves, at the expense of converging to suboptimal policies.

Original: Conversely, a smaller radius results in slower learning trajectories.

Original: Our trust region radius ε = 0 .

Original: 01 seems to balance the two effects.

Original: Different transportation costs.

Original: The Taxi-v3 environment is characterized by two radically different set of actions: Move = { Up , Left , Right , Down } and Passenger = { PickUp , DropOff } .

Original: We use this insight to provide a simple example of how transport costs can affect the geometry of the action space and in turn affect the learning process.

Original: Specifically, we consider different cost functions.

Original: Let m, m

Original: ′ ∈ Move , p, p ′ ∈ Passenger .

Original: [-Consider now the following transport costs (for all of them we set c ( x, x ) = 
-]
Revised: [-Consider now the following transport costs (for all of them we set c ( x, x ) = 
-]

Original: ∀ x ∈ Move ∪ Passenger ):

Original: Equally costly.

Original: In this cost term, c ( m, m ′ )

Original: =

Original: c ( m ′ , m ) =

Original: c ( p, p ′ )

Original: =

Original: c ( p ′ , p )

Original: =

Original: c ( m, p ) =

Original: c ( p, m ) .

Original: This is also a distance.

Original: Cheap [Expensive].

Original: In this cost term, c ( m, m ′ )

Original: =

Original: c ( m ′ , m ) , c

Original: ( p, p ′ )

Original: =

Original: c ( p ′ , p ) , but c ( m, p ) >

Original: c ( p, m ) [ c ( m, p ) <

Original: c ( p, m ) ].

Original: In particular, c ( p, m ) =

Original: c ( m, p ) − 0 .

Original: 2 [ c ( p, m ) =

Original: c ( m, p ) + 0 . 2 ].

Original: Since it is not symmetric, this is not a distance.

Original: Very cheap [expensive].

Original: In this cost term, c ( m, m ′ )

Original: =

Original: c ( m ′ , m ) , c

Original: ( p, p ′ )

Original: =

Original: c ( p ′ , p ) , but c ( m, p ) ≫ c ( p, m )

Original: [ c ( m, p ) ≪ c ( p, m ) ].

Original: In particular, c ( p, m ) =

Original: c ( m, p ) − 0 .

Original: 5 [ c ( p, m ) =

Original: c ( m, p ) + 0 . 5 ].

Original: Since it is not symmetric, this is not a distance.

Original: Notice that a “local minimum” in the Taxi-v3 environment is to never play actions in the Passenger set: such a strategy yields a score of − 200 .

Original: Indeed, performing a DropOff leads to large penalty when not at the right location.

Original: As a result, carrying a passenger is “risky”.

Original: However, these are also the actions that are required to cross the − 200 bar and learn the optimal policy.

Original: We can interpret the aforementioned costs as follows:

Original: Equally costly.

Original: This cost function ignores the distinction between the two classes of actions: moving from one to the other is equally costly.

Original: Cheap [Expensive].

Original: Moving mass towards “risky” actions is incentivized [penalized].

Original: This cost function shapes the landscape of the action space so that there is a downhill [uphill] from Move to Passenger , and withing these sets the geometry is analogous to the Equally costly one.

Original: Very cheap [expensive].

Original: Analogous to Cheap [ Expensive ].

Original: We observe that with Very cheap [ expensive ] the agent learns apparently worse than with the others.

Original: Intuitively, this is due to the fact that the Very cheap cost incentivizes too much the agent to drop off the passenger; as a result, it completes the task, but it also incurs many illegal DropOff s, which yield the penalty.

Original: Oppositely, the Very expensive cost is too conservative, and the agent learns to perform too few DropOff s, being unable to complete the task correctly consistently.

Original: On the other hand, the Cheap [ Expensive ] cost trains the agent slightly faster [slower] than the Equally costly option.

Original: Namely, these costs balance in a different way the trade-off between exploration and exploitation.

Original: This example highlights the impact of the geometry of the action space, which, in turn, is another tuning parameter available to practitioners.

Original: For instance, one could consider transport costs which embed an exploration penalties (e.g., moving mass from actions that are already rarely used or to others that are almost always used can be penalized), or physical considerations (e.g., an expert supervisor might know that certain actions are most likely the correct ones given certain configurations of the agent).

Original: Sensitivity to the advantage estimation Both the dual problem (D) and the “regularized” policy loss rely on the accuracy of the advantage function estimation, as well as on the solution of a maximization problem over a continuous space, which is approximated by finitely many evaluated points.

Original: That is, when training on continuous actions spaces, we rely on the approximation

Original: [Equation]

Original: where A ′ ⊂

Original: A is a (possibly state-dependent) finite set of actions.

Original: While smoothness of the objective guarantees convergence as |A ′ | → ∞ with uniformly chosen actions in A , we empirically investigate the sensitivity of the method to various practical advantage estimation schemes.

Original: In the experimental results so far presented, we use a GAE based estimate of the advantage function.

Original: In continuous action spaces, we explore a neural network approximation for the advantage function as well.

Original: Overall, this approach does not perform comparably to the single sample estimation of the objective with GAE.

Original: The training performances in the Hopper-v3 environment for different cardinalities of A ′ are juxtaposed in Fig.

Original: 3, and the trained agents scores are summarized in Table 3.

Original: In Fig.

Original: 4 we compare the performances of the neural network approximation of the advantage function with the GAE in the environments HalfCheetah-v3 , Hopper-v3 , and MountainCarContinuous-v0 .

Original: The number of sampled actions is |A ′

Original: | = 4 , |A

Original: | = 2 , and |A ′

Original: | = 16 , respectively.

Original: In general, we observe that (i) simultaneous training of an advantage network and the OT-TRPO policy update could be unstable, and (ii) in the case of convergence, OT-TRPO with neural advantage function approximation converges to a suboptimal policy.

Original: We hypothesize that the OT-TRPO objective is sensitive to biased value estimates in the neighborhood of the mean of the Gaussian policy network.

Original: It will be subject of future research if and to what extent these results can be improved, for instance via normalizing flow policy parametrizations [38].

Original: To support our observation regarding the failure mode of biased advantage estimates, we provide an ablation study for different λ parameters for the GAE.

Original: The training curves are shown in Fig.

Original: Learning steps ( 10 6 )

Original: J ( θ ∗ )

Original: Hopper-v

Original: GAE

Original: NN (2)

Original: NN (4)

Original: NN (8)

Original: NN (16)

Original: NN (32)

Original: [Figure]

Original: [Table]

Original: It is known that for λ = 1 , the GAE is equivalent to an unbiased Monte Carlo estimate of the advantage [30].

Original: The observation that state-of-the-art performances are achieved with a gae_lambda hyperparameter very close to 1 (i.e., an unbiased Monte Carlo estimate) supports the idea that approximating the maximization with a biased advantage estimator is prone to instability.

Original: A.5 More details on the policy update

Original: We shall now discuss how the implementation affects the exactness guarantees of Theorem 1.

Original: Indeed, albeit theoretically the algorithm presented guarantees an optimal policy update at every step, the practical implementation might resort to approximations.

Original: We consider as running example the simple setting in Example 2, where all the computations can be carried out analytically.

Original: We consider a generic current policy with π θ ( L | s 0 )

Original: = θ, π θ ( R | s 0 )

Original: = 1 − θ .

Original: For this example, we can pedagogically solve both (P) and (D).

Original: The Q -function is Q π θ ( s 0 , L ) = − 1 , Q π θ ( s 0 , R )

Original: = 1 .

Original: Then, the value function at s 0 can be expressed in terms of the current policy as

Original: [Equation]

Original: and the advantage reads

Original: [Equation]

Original: Trivially, the visitation frequency is fully described by ρ π θ ( s 0 )

Original: = 1 ; we can ignore the terminal states as they do not affect the problem.

Original: We consider the binary transportation cost c ( L, R )

Original: =

Original: c ( R, L ) = 1 , c

Original: ( L, L ) =

Original: c ( R, R ) = 0 , and we first solve (P).

Original: By direct inspection, it is easy to see that there are two cases.

Original: Let the new policy be π ˜ θ

Original: ( L | s 0 )

Original: =

Original: ˜ θ, π ˜ θ

Original: ( R | s 0 )

Original: = 1 − ˜ θ .

Original: If θ < ε , then the solution is to move all the mass from action L to action R : ˜ θ = 0 .

Original: Otherwise, we can move at most ε mass, and this is the optimal solution: ˜ θ

Original: = θ − ε .

Original: [Figure]

Original: HalfCheetah

Original: [Figure]

Original: [Figure]

Original: Second, we solve (D):

Original: [Equation]

Original: (cid:123)(cid:122) (cid:125) =

Original: 2 θ − min { 2 ,λ } +(1 − θ ) max { 2( θ − 1) − λ, 2 θ } (cid:124) (cid:123)(cid:122) (cid:125) =

Original: [-2 θ = arg min λ ≥λε − θ min { 2 , λ } , which yields λ ∗ = 
-]
Revised: [-2 θ = arg min λ ≥λε − θ min { 2 , λ } , which yields λ ∗ = 
-]

Original: if θ < ε , λ ∗ = 2 if θ ≥ ε .

Original: As expected, λ ∗ is not zero if the constraint is active, which is the case every time there is more mass in π ( L | s 0 ) than what we are allowed to move.

Original: In practice, we can solve this optimization problem to global optimality (with sufficient numerical precision).

Original: In discrete settings, the numerical approximations introduced by the solver are the only source of approximations.

Original: In the considered example, with finite state-action spaces, one should follow the first approach discussed for the parameter update; namely, one should update θ according to (10) and (9).

Original: For the sake of providing insights and shading light on both the inner working of the proposed algorithm and its limitations, we now study all the proposed options for step 5 in Algorithm 1.

Original: Direct parametrization (finite spaces)

Original: We distinguish two cases: θ > ε .

Original: The constraint is active: λ ∗ = 2 .

Original: We obtain D λ ∗ ( s, L ) =

Original: { L, R } and D λ ∗ ( s, R ) = { R } .

Original: The minimizers are not unique: we need “mass splitting”.

Original: In particular, T λ ∗ ( s 0 , L ) =

Original: L (the “closest” minimizer is the one that not require any mass displacement), T λ ∗ ( s 0 , L ) =

Original: R , and T λ ∗ ( s 0 , R )

Original: =

Original: T λ ∗ ( s 0 , R )

Original: = R .

Original: Finally, t ∗ is selected to “activate” the constraint.

Original: In this case, t ∗ = ε , and we obtain ˜ π

Original: ( L | s 0 )

Original: = θ − ε =:

Original: ˜ θ , as expected.

Original: θ ≤ ε .

Original: In this case, λ ∗ = 0 .

Original: Then we have D λ ∗ ( s, L ) =

Original: { R } and D λ ∗ ( s, R ) = { R } .

Original: That is, the minimizers are unique, the unique transport map T λ ∗ ( s 0 , · )

Original: =

Original: T λ ∗ ( s 0 , · ) =

Original: T is the constant map T ( L ) =

Original: T ( R ) =

Original: R and we do not need “mass splitting”.

Original: The new policy is then trivially ˜ π

Original: ( L | s 0 )

Original: =

Original: T

Original: # π ( L | s 0 )

Original: = 0 ; that is, ˜ θ = 0 , again as expected.

Original: Direct parametrization via policy network (continuous states, discrete actions)

Original: In this approach, we build on top of the result of the previous one, and we perform gradient descent on the loss

Original: L ( θ ) = (cid:80) s ∈ ˆ S ρ π θt ( s ) (cid:13)(cid:13)

Original: π θ ( ·| s ) − π ∗ θ t ( ·| s ) (cid:13)(cid:13) 2 to steer π θ .

Original: In particular, the minimization of L steers the policy network towards the optimal policy.

Original: Hence, if we were to perform gradient descent until convergence, the policy update would be exact up to numerical precision.

Original: Arbitrary policy parametrization (continuous states, continuous actions)

Original: Armed with λ ∗ , we construct the loss function

Original: [Equation]

Original: In particular, ˜ θ affects only π ˜ θ .

Original: The intuition behind this loss function is to increase the probability mass where the regularized advantage function has its maxima.

Original: However, in general, the solutions are multiple, while the gradient descent will possibly converge to one.

Original: As such, one limitation of this cost function is related with the “mass splitting” issue.

Original: That is, if we were to perform gradient descent until convergence, the policy update might violate the trust-region constraint.

Original: Empirically, the proposed loss function performs well, and allows the deployment of optimal transport trust-region methods in continuous spaces.

Original: Moreover, the gradient descent is not iterated until convergence in practice.

Original: Nonetheless, the arising optimization procedure deserves further study in terms of convergence and optimality guarantees, but our duality results provide an intuition on the method.

Original: Future work will focus on understanding how to implicitly describe transport plans , rather than transport maps.

Original: In the remainder of the discussion, we study the gradient for the running example, and we validate the intuition above:

Original: [Equation]

Original: We consider the cases θ > ε and θ ≤ ε separately: θ > ε .

Original: The gradient reads ∇ ˜ θ

Original: L (˜ θ )

Original: = − 2 : we are steering θ to 0 .

Original: Namely, the loss is embedding the transport map T λ ∗ ( s 0 , L )

Original: =

Original: R , T λ ∗ ( s 0 , R ) =

Original: R .

Original: Following the gradient corresponds to interpolating between π θ ( ·| s 0 ) and T λ ∗ ( s 0 , · ) # π θ ( ·| s 0 ) .

Original: However, as previously discussed, the optimal solution requires “mass splitting”: if we perform gradient descent until convergence, the policy update yields ˜ θ = 0 , which violates the trust-region constraint.

Original: θ ≥ ε .

Original: The gradient vanishes: ∇ ˜ θ

Original: L (˜ θ )

Original: = − λ ∗ = 0 .

Original: This issue is related with the policy parametrization.

Original: The issue can be addressed with a decreasing trust-region radius ε → 0 .

Original: This highlights a potential drawback of the proposed practical implementation: albeit our theoretical results are valid and parametrization-independent, the choice of the network architecture affects the learning process.

Original: Aside from the possibility of a vanishing gradient, it is possible that the optimal policy is not contained in the family described by the parametrization.

Original: A.6 Further details on the experimental results A.6.1 Additional numerical details

Original: In addition to the experimental results of Section 6, we provide the numerical results of the performance comparison among our algorithm and the baseline methods in Table 4.

Original: In particular, we average the expected scores obtained in the last 10% of the training episodes.

Original: [Table]

Original: Moreover, in Table 5 we report the training time in the Mujoco [36] environments, which are the most computationally demanding.

Original: The values in Table 5 have been collected under comparable conditions and similar hardware.

Original: While the absolute values here provide little insight, the relative comparisons highlight that our algorithm has computational performances only slightly slower than TRPO [29],

Original: PPO [31]

Original: [Table]

Original: and A2C [17].

Original: Instead, other methods levereging optiamal transport and, in particular, the

Original: Wasserstein distance (i

Original: [Table]

Original: e., BGPG [22] and WNPG [20]) require a considerably longer training time.

Original: [Table]

Original: A.6.2 Some considerations on convergence speed

Original: Finally, in some environments (e.g., Taxi-v3 , Hopper-v3 , and Swimmer-v3 ), TRPO seems to learn faster than the proposed algorithm (albeit not converging to a better solution).

Original: This is closely related to the choice of trust-region.

Original: With OT-TRPO, moving the “probability mass” is done at the cost c ( a 1 , a 2 ) .

Original: In TRPO, there is no notion of transport cost in the action space: all the actions are at the same “distance”.

Original: That is, they differ only based on the probability of using them when at a state s ∈ S ; namely log ( π ( a 1 | s ) /π

Original: ( a 2 | s )) .

Original: Whenever c ( a 1 , a 2 ) > log ( π ( a 1 | s ) /π

Original: ( a 2 | s )) , the convergence is faster with TRPO compared to OT-TRPO.

Original: Clearly, one could design a transport cost that speeds up the convergence (a “smaller” one), at the price of a potentially less stable and robust behavior of the algorithm.

Original: The study of the “optimal” choice of transport cost is indeed an interesting question for future research.

Original: As an example, in the Taxi-v3 environment there are 6 actions.

Original: The proposed method with a binary distance incurs a cost of 1 to bring a state from a uniform probability distribution over the actions to a deterministic one (which for most states is the case in the optimal policy).

Original: Instead,

Original: TRPO considers such policies only log (6) <

Original: 1 away.

Original: Similar (but more complicated) calculations can be carried out for the Mujoco environments.

Original: Another reason can be found in the lack of symmetry of the KL divergence.

Original: Such asymmetry can “enforce” a direction of exploration: going back to a previously explored point in the policy space might be more (or less) expensive.

Original: Oppositely, an optimal transport cost can be chosen to be symmetric (e.g., the Wasserstein distance and the costs used for our experiments).

Original: When combined with errors in the estimation of the advantage function, it can result in oscillations in the policy space, which might slow down the convergence.

Original: B Proofs B.1 Proof of Theorem

Original: We start with weak duality:

Original: Proposition 5 (Weak Duality) .

Original: Weak duality holds.

Original: Namely, (P) ≤ (D) .

Original: Proof.

Original: By minimax we have

Original: [Equation]

Original: ≤ inf

Original: [Equation]

Original:  

Original: [Equation]

Original: Second, for all s ∈ S , Kantorovich duality [2, Theorem 6.1.1] gives

Original: [Equation]

Original: Without loss of generality, we assume λ > 0 ; else, the result is straightforward.

Original: Then, we choose ψ ( · ) =

Original: A π ( s, · ) /λ and ϕ ( · ) = inf a ′ ∈A { c ( · , a ′ )

Original: − ψ ( a ′ ) } .

Original: By Assumptions 1 and 2 ϕ, ψ are continuous on a compact space.

Original: Hence, they are also bounded [21,

Original: Theorem 27.4].

Original: Moreover,

Original: [Equation]

Original: Thus, ( ϕ, ψ ) is a valid, possibly sub-optimal, choices for the supremum.

Original: Overall, we get the upper bound

Original: [Equation]

Original: Hence

Original: [Equation]

Original: we obtain

Original: [Equation]

Original: The proof of strong duality (i.e., equality) is more delicate, and requires some preliminary results.

Original: First, we recall the definition the regularization operator Φ λ , which represents intuitively a “regularized” advantage function:

Original: [Equation]

Original: The regularization operator Φ λ is well defined by Assumptions 1 and 2.

Original: Indeed, for all λ ∈ R , s ∈

Original: S , a ∈ A , a ′ (cid:55)→

Original: A π ( s, a ′ ) − λc ( a, a ′ ) is a continuous function on a compact space and thus it attains its maximum value (e.g., see [21, Theorem 27.4]).

Original: Accordingly, we recall the set of maximizers

Original: [Equation]

Original: Since minimizers are generally not unique, D λ ( s, a ) is indeed a set.

Original: In fact, it is also closed:

Original: Lemma 6.

Original: For all λ ∈ R , s ∈ S , a ∈ A , the set D λ ( s, a ) ⊂ A is non-empty and closed (and thus measurable).

Original: Proof.

Original: Since a ′ (cid:55)→ A π ( s, a ′ ) − λc ( a, a ′ ) is a continuous function on a compact space and thus it attains its maximum value (e.g., see [21, Theorem 27.4]), which directly implies that D λ ( s, a ) is non-empty.

Original: To prove that it is closed, it suffices to show that D λ ( s, a ) contains all its limit points (e.g., see [21, Corollary 17.7]).

Original: Let ( a n ) n ∈ N , a n ∈ D λ

Original: ( s, a ) be a sequence converging to ¯ a ∈ A .

Original: Let h λs,a : A → R be defined as h λs ( a )

Original: =

Original: A π ( s, a ′ ) − λc ( a, a ′ ) .

Original: Since a n ∈ D λ ( s, a ) , h λs,a ( a n ) =

Original: Φ λ ( s, a ) for all n ∈ N .

Original: Moreover, A π is continuous by Assumption 2, and thus so is h .

Original: By continuity, h λs,a (¯ a ) =

Original: lim n →∞ h λs,a ( a n ) =

Original: Φ λ ( s, a ) , and so ¯ a ∈ D λ ( s, a ) .

Original: Measurability of D λ ( s, a ) follows from closedness and π ( ·| s ) being a Borel probability measure.

Original: Among all possible minimizers, the closest and furthest apart will play an important role in the proof of strong duality.

Original: Thus, we define

Original: [Equation]

Original: where the max and min are well defined by Lemma 6, since c ( a, · ) is a continuous function for all a ∈ A and D

Original: λ ( s, a ) is compact (being the closed subset of a compact space, [21, Theorem 26.2]).

Original: These definitions allow us to study the properties of Φ more in detail:

Original: Lemma 7 (Properties of Φ ) .

Original: The regularization operator Φ has the following properties:

Original: It is uniformly (in s , a , and λ ) lower and upper bounded by some constant.

Original: It is lower semi-continuous, non-increasing, and convex in λ .

Original: For all λ > 0 the left and right derivatives are

Original: [Equation]

Original: There are measurable selections

Original: [Equation]

Original: Proof.

Original: The proof follows closely [8, Lemma 3].

Original: Specifically, we prove the steps separately:

Original: Since A π is continuous on a compact space, its absolute value is bounded by some C (e

Original: g., see [21,

Original: Theorem 27.4]).

Original: For

Original: the upper bound: Φ λ results from an infimum, λc ( a, a ′ ) ≥ 0 , and a ′ =

Original: a is always a valid choice, yielding Φ λ ( s, a ) ≥

Original: A π ( s, a ) − λc ( a, a ) ≥ − C

Original: For the lower bound;

Original: since

Original: transport costs are non

Original: -negative,

Original: we also

Original: have

Original: Φ

Original: λ

Original: ( s, a ) ≤ sup a ′ ∈A

Original: A π ( s, a ′ ) ≤ C .

Original: Thus, | Φ λ ( s, a ) | ≤ C for all λ ≥ 0 , s ∈ S , and a ∈ A . 2.

Original: For all s ∈ S , a ∈ A , λ (cid:55)→ Φ λ ( s, a ) is the supremum of non-increasing and affine functions.

Original: As a result, it is non-increasing, convex and lower semi-continuous in λ .

Original: The proof boils down to upper semi-continuity and lower semi-continuity of λ (cid:55)→ D λ ( s, a ) and λ (cid:55)→

Original: D λ ( s, a ) , respectively, whose proof is identical to the setting of DRO since both s and a are frozen; see [8, Corollary 1] and replace the distance d p with c .

Original: Then, the results follows from [8,

Original: Lemma 4.iii], since, by definition, the growth rate is 0 for compact spaces and Φ > −∞ .

Original: The set of minimizers D λ

Original: ( s, a ) is non-empty and closed by Lemma 6.

Original: Since it is a closed subset of a compact space, it is also compact [21, Theorem 26.2].

Original: Moreover, c ( a, · ) is continuous for all a ∈ A and c ( · , a ′ ) is continuous (and thus measurable) for all a ′ ∈ A , thus the result follows from [1, Theorem 18.19].

Original: We are now ready to prove strong duality:

Original: Proof of Theorem 1.

Original: In view of Proposition 5, we only need to prove (P) ≥ (D).

Original: We split the proof in two steps.

Original: First, we show existence of an optimal multiplier λ ∗ ∈ [0 , ∞ ) .

Original: Second, we study the case λ ∗ > 0 , whereas λ ∗ = 0 is straightforward from direct inspection of (P) and (D).

Original: [-Let h : R ≥ 
-]
Revised: [-Let h : R ≥ 
-]

Original: → R be defined as

Original: [Equation]

Original: By Lemma 7.1-2, h is convex and lower semi-continuous in λ .

Original: Moreover, h is coercive:

Original: [Equation]

Original: since Φ λ ( s, a ) ≥ A π ( s, a ) for all a ∈ A and s ∈ S .

Original: Thus, there exists λ ∗ ∈ arg min λ ≥ 0 h ( λ ) .

Original: Indeed, by definition of limit, there exists ¯ λ such that h ( λ ) > h (0) for all λ > ¯ λ .

Original: Hence, when looking for a minimum, we can restrict ourselves to λ ∈ [0 , ¯ λ ] .

Original: Since h is lower semi-continuous on a compact space [21, Corollary 27.2], it attains its minimum at some λ ∗ ∈ [0 , ¯ λ ] [21, Theorem 27.4].

Original: Assume now that λ ∗ > 0 .

Original: We study the first-order optimality conditions on h .

Original: Since h is generally non-smooth (due to the maximum in Φ λ ), we express them via the left and right derivatives, welldefined for convex functions [27, Theorem 23.1], as ∂h ( λ ) ∂λ − ≤ 0 and ∂h ( λ ) ∂λ + ≥ 0 .

Original: Hence, we have

Original: [Equation]

Original: 14) and ∂h

Original: [Equation]

Original: where in

Original: ♡ we use that, for all λ , D λ ( s, a )

Original: =

Original: c ( a, T λ ( s, a )) , D λ ( s, a ) =

Original: c ( a, T λ ( s, a )) result from the composition of a continuous (the cost) and a measurable function ( T λ ( s, a ) and T λ

Original: ( s, a ) ), and thus they are measurable [28, Theorem 1.8].

Original: Integrability follows directly from non-negativity [28,

Original: Theorem 1.17, Definition 1.23].

Original: Hence, we can use the differentiation lemma [12, Theorem 6.28] together with Lemma 7;

Original: ♢ we use Lemma

Original: ; and △ we use the definition of T λ ∗ , T λ ∗ and of

Original: D

Original: λ ∗ ( s, a ) , D λ ∗ ( s, a ) .

Original: Overall, the inequalities in (14) and (15) lead to

Original: [Equation]

Original: With this observation, we can construct a primal solution ˜ π ( ·| s ) and show its optimality.

Original: First, by (16), there exists t ∗ ∈ [0 , 1] such that

Original: [Equation]

Original: For all s ∈ S consider the policy

Original: [Equation]

Original: By construction, π ( ·| s ) is a feasible solution.

Original: Indeed, with Id being the identity mapping I

Original: d

Original: ( a ) =

Original: a , we can consider the (possibly sub-optimal) transport plan γ s ∈ Π( π ( ·| s ) ,

Original: ˜ π

Original: ( ·| s )) , defined as

Original: [Equation]

Original: Then, by definition of optimal transport discrepancy and monotonicity of the integral we have

Original: [Equation]

Original: To conclude

Original: [Equation]

Original: by definition

Original: [Equation]

Original: we have

Original: [Equation]

Original: and thus

Original: [Equation]

Original: P) ≥

Original: [Equation]

Original: = (D)

Original: [Equation]

Original: proving (P) ≥ (D).

Original: In particular

Original: [Equation]

Original: we have (P) =

Original: (D) and π is the primal solution.

Original: Thus, the supremum over Π is attained

Original: [Equation]

Original: The proofs of Corollaries 2 and 3 then follow directly.

Original: Remark 4 .

Original: Our theoretic results readily extends to transport costs c s,π parametrized by the state s or the policy π , as long as the mapping ( s, π, a 1 , a 2 ) (cid:55)→ c s,π ( a 1 , a 2 ) is continuous.

Original: A natural question is whether the results of Theorem 1 can be extended to Sinkhorn divergence (or entropy-regularized optimal transport discrepancies in general), in particular given the success of these in the context of RL [6, 23].

Original: Our duality results do not directly apply to Sinkhorn divergence (or entropy-regularized optimal transport discrepancies in general).

Original: For instance, our proof of weak duality (Proposition 5) leverages Kantorovich duality, which is optimal transport specific.

Original: However, duality results for entropy-regularized optimal transport [35], and recent results in DRO [7] suggest that our proof can be adapted to such constraints.

Original: Having said this, we remark that

Original: Sinkhorn divergence was successful in mitigating the burden of optimal transport computations.

Original: The proposed algorithm does not rely on any of such, and it is thus unclear if regularized optimal transport discrepancies yield any benefit in this context.

Original: B.2 Proof of Proposition

Original: Proof.

Original: First, recall the performance difference lemma [11, Lemma 6.1]:

Original: [Equation]

Original: Second, recall that from Corollary 2 that an optimal policy update is

Original: [Equation]

Original: Then, the proof follows from the properties of ˜ π ∗ :

Original: [Equation]

Original: where in ♡ we use that, by assumption, A π ( s, a ) ≥ ˆ A π ( s, a ) − ∥

Original: A π − ˆ A π ∥ ∞ and ˆ A π ( s, a ) ≥

Original: A π ( s, a ) − ∥

Original: A π − ˆ A π ∥ ∞ for all s ∈ S and a ∈ A ;

Original: ♢ by definition of the regularization operator Φ and of T λ ∗ ( s, a ) and T λ ∗ ( s, a ) (Lemma 7.4), for all s ∈ S and a ∈ A it holds

Original: [Equation]

Original: △ as in (17), it holds

Original: [Equation]

Original: and A π ( s, a ) ≤ Φ

Original: [Equation]

Original: ( s, a )

Original: [Equation]

Original: since for all λ ∈ R , a ∈ A , and s ∈ S

Original: [Equation]

Original: □ the expected value of the advantage function vanishes, since, by definition, for all s ∈ S

Original: [Equation]

Original: This concludes the proof.