Provably Convergent Policy Optimization via Metric-aware Trust Region Methods
Authors: Jun Song
Niao He
Lijun Ding
Chaoyue Zhao
Venue: Accepted by TMLR ()

Abstract: Trust-region methods based on Kullback-Leibler divergence are pervasively used to stabilize policy optimization in reinforcement learning. In this paper, we exploit more flexible metrics and examine two natural extensions of policy optimization with Wasserstein and Sinkhorn trust regions, namely Wasserstein policy optimization (WPO) and Sinkhorn policy optimization (SPO). Instead of restricting the policy to a parametric distribution class, we directly optimize the policy distribution and derive their close-form policy updates based on the Lagrangian duality. Theoretically, we show that WPO guarantees a monotonic performance improvement, and SPO provably converges to WPO as the entropic regularizer diminishes. Moreover, we prove that with a decaying Lagrangian multiplier to the trust region constraint, both methods converge to global optimality. Experiments across tabular domains, robotic locomotion, and continuous control tasks further demonstrate the performance improvement of both approaches, more robustness of WPO to sample insufficiency, and faster convergence of SPO, over state-of-art policy gradient methods.

Published in Transactions on Machine Learning Research (06/2023)
Provably Convergent Policy Optimization via Metric-aware Trust Region Methods
Jun Song
juns 10.011. edu
Deportment of Industrial and Systems Engineering University of Washington
Niao He
niao.he etw.f..ff..ccc........
Department of Computer Science ETH Zürich
Lijun Ding
Iding47@wisch edu
Wisconsin Institute for Discovery University of Wisconsin Madison
Chaoyue Zhao
cyshaoour. edu
Deportment of Industrial and Systems Engineering University of Washinglon
Reviewed on OpenReview: https: // openreview. net/ form? id= JJGMGS
Abstract
Trust-region methods based on Kullback-Leibler divergence are pervasively used to stabilize policy optimization in reinforcement learning. In this paper, we exploit more flexible metrics and examine two natural extensions of policy optimization with Wasserstein and Sinkhorn trust regions, namely Wasserstein policy optimization (WPO) and Sinkhorn policy optimization (SPO). Instead of restricting the policy to a parametric distribution class, directly optimize the policy distribution and derive their closed-form policy updates based on the Lagrangian duality. Theoretically, we show that WPO guarantees monotonic performance improvement, and SPO provably converges to WPO as the entropic regularizer diminishes. Moreover, we prove that with a decaying Lagrangian multiplier to the trust. region constraint, both methods converge to global optimality. Experiments across tabular domains, robotic locomotion, and continuous control tasks further demonstrate the performance improvement of both approaches, more robustness WPO to sample insufficiency, and faster convergence of SPO, over state-of-art, policy gradient methods.
1 Introduction
Policy-based reinforcement learning (RL) approaches have received remarkable success in many domains,
including video games (Wang al.) 2017 Wu et al. 2017 Mnih et al.
2016), robotics (Grudic et al. 2003
Levine et al. 2016), and continuous control tasks (Duan et al. 2016
Schulman al. 2016 Hoess al,
2003 Peters Schaal 2006 [illicrap et al. 2016 Sutton et al.) 1999: Williams 1992 Mnih et al. 2016 Silver et al. 2014). The core idea is to represent the policy with probability distribution TT)as)) = P(a s; 0), such that the action in state s is chosen stochastically following the policy my controlled by parameter 0. Determining the right step size to update the policy crucial for maintaining the stability of policy gradient methods: too conservative choice of stepsizes result in slow convergence, while too large stepsizes may lead to catastrophically bad updates.
2015). One prominent example policy gradient method (Grudic al.
To control the size of policy updates, Kullback-Leible (KL) divergence is commonly adopted to measure the difference ween two policies. For example, the seminal work on trust region policy optimization (TRPO)

---

Published in Transactions on Machine Learning Research (06/2023)
by Schulman et al. (2055) introduced KL divergence based constraints (trust region constraints) to restrict the size of the policy update; see also Peng et al. (2019) Abdolmaleki et al. (2018). Kakade (2001) and
Schulman et al. 2017 introduced a KL-based penalty term to the objective to prevent excessive policy shift.
Though KL-based policy optimization has achieved promising results, it remains interesting whether using other metrics to gange the similarity between policies could bring additional benefits. Recently, a few work (Richemond & Maginnis 2017] Zhang al. 2018; Moskovitz al,, 2021 Pacchiano al., 2020) has explored the Wasserstein metric restrict the deviation between consccutive policies. Compared with KL divergence, the Wasserstein metric has several desirable properties. Firstly, it is true symmetric distance measure. Secondly, it allows flexible user-defined costs between actions and is less sensitive to ill-posed likelihood ratios. Thirdly but most importantly, the Wasserstein metric takes into account the geomelry of the metric space Paanaretos Zemel 2019) and allows distributions to have different or even non-overlapping supports.
Motivating Example: Below we provide an example of a grid world (see Figure D that illustrates the advantages of using the Wasserstein metric over the KL divergence to construct trust regions and policy updates. The grid world consists of 5 regular grids and 2 goal grids, and there are three possible actions: left, right, and pickup. The player always starts from the middle grid, and making a left or right move results in a reward of 11 Picking up yields a reward of -3 at regular grids, +5 at the blue goal grid, and +10 at the red goal grid. An episode terminates either at the maximum length of 10 or immediately after picking up. We define the geometric distance between left and right actions
and 4 between other actions.
Figure 1: Motivating grid world example
Figure 2 shows the Wasserstein distance and KL divergence for different policy shifts of this grid world that Wasserstein metric utilizes the geometric distance between actions to distinguish
example. We the shift of policy distribution to close action (policy distribution + in Figure 2a) from the shift to a far action (policy distribution 1 + 3 in Figure 2b while KL divergence does not. Figure
3 demonstrates the
constrained policy updates based on Wasserstein distance and KT. divergence respectively with a fixed trust region size We can see that Wasserstein- based policy update finds the optimal policy faster than KL-based policy update. This because KL distance is larger than Waserstein when considering policy shifts of close actions (see Figure 2al). Therefore, Wasserstein policy update is able shift action (from left right) in multiple states, while KL update is only allowed to shift action in single state. Besides, KL policy update keeps using suboptimal short-sighted solution between the 2nd and 4th iteration, which further slows down the convergence.
Policy distribution
Policy distribution 2
Policy distribution
Policy distribution 3
0.8
0.8
0.8
0.6
dg : 0.85
0.6
0.6
dw : 3.4
0.6
0.4
dRz.
0.4
0.4
dgt 2.5
0.4
0.2
0.2
0.2
0.2
0.0
0.0
right pickup Action
left
right Action
pickup
right pickup Action
right pickup Action
(a) Policy shift of close action
(b) Policy shift of far action
Figure 2: Wasserstein utilizes geometric feature of action space
2

---

Published in Transactions on Machine Learning Research (06/2023)
KL constrained policy update
1st iteration:
Wasserstein constrained policy update
2nd iteration:
P
1st iteration:
lteration:
P
2nd iteration:
4th iteration:
P
5th iteration:
Figure 3: Demonstration of policy updates under different trust regions
However, the challenge of applying the Wasserstein metric for policy optimization is also evident: evaluating the Wasserstein distance requires solving an optimal transport problem, which could be computationally expensive. To avoid this computation hurdle, existing work resorts to different techniques to approzimate the policy updule under Wasserstein regularization. For example, Richemond & Maginnis (2117) solved the resulting RL problem using Pokker-Plancee equations; Zhang et al. (2018) introduced particle approximation method to estimate the Wasserstein gradient flow. Recently, Moskovitz et al. (2021) instead considered the sccond-order Taylor expansion of Wasserstein distance based on Wasserstein information matrix to characterize the local behavioral struct ure of policies. Pacchiano et al. (2020) tackled behavior-guided policy optimization with smooth Wasserstein regularization by solving an approximate dual reformulation defined reproducing kernel Hilbert spaces. Aside from such approximation, some of these work also limits the policy representa ation lo particular parametric distribution class, As indicated in Tessler et. al. 2011)), since parametric distributions are not convex in the distribution space, optimizing over such distributions results in local movements in the action space and thus leads to convergence to sub-optimal solution. Until now, the theoretical performance of policy optimization under the Wasserstein metric remains elusive light of these approximation errors.
In this paper, we study policy optimization with trust regions based on Wasserstein distance and Sinkhorn divergence. The latter is smooth variant of Waserstein distance by imposing entropic regularization to the optimal transport problem (Cuturi 2013). We call them, Wasserstein Policy Optimization (WPO) and Sinkhorn Policy Optimization (SPO), respectively. Instead of confining the distribution of policy to a particular distribution class, we work on the space of policy distribution directly, and consider all admissible policies that are within the trust regions with the goal of avoiding approximation errors. Unlike existing work, we focus on eract characterization of the policy updates. We would like emphasize that our methodology and theoretical analysis in Section and primarily concentrate on discrete action space. However, we also present extension of method to accommodate continuous action space, detailed in Section 7.5 We highlight our contributions follows:
1. Algorithms:
We develop closed-form expressions of the policy updates for both WPO and SPO based on the corresponding optimal Lagrangian multipliers of the trust region constraints. To the best of our knowledge, this is the first explicit closed-form updates for policy optimization based on Wasserstein and Sinkhorn trust regions. In particular, the optimal Lagrangian multiplier of SPO admits a simple form and can be computed efficiently. A practical on-policy actor-critic algorithm proposed based on the derived expressions of policy updates and ad advantage value function estimation.
We theoretically show that WPO guarantecs a monotonsc performance improvement through the iterations, even with non-optimal Lagrangian multipliers. We also prove that SPO converges WPO the entropic regularizer diminishes. Moreover, we prove that with decaying schedule of the multiplier,
Theory:
3

---

Published in Transactions on Machine Learning Research (06/2023)
SPO and WPO converge to global optimality, and with constant multiplier, both methods converge linearly up to neighborhood of the optimal value. To our best knowledge, this appears to be the first. convergence rate analysis of policy optimization based on Wasserstein-tyye metrics.
3. Experiments: We provide comprchensive evaluation on the efficiency of WPO and SPO under several types of lesting environments including tabular domains, robotic locomotion lasks, and further exlend it to continuous control tasks. Compared to state-of-art policy gradients approaches that use KL divergence such as TRPO and PPO and those use Wasserstein metric such as Wasserstein Natural Policy Gradient (WNPG) (Moskovitz et al. 2021) and Behavior Guided Policy Gradients (BGPG) (Pacchiano et al.) 2020 our methods achieve better sample efficiency, faster convergence, and improved final performance. Numerical study indicates that by properly choosing the weight of entropic regularizer, SPO achieves a better trade-off between convergence and final performance than WPO.
Related work: Wasserstein-like metrics have been explored in a number of works in the context reinforcement learning. Ferns et al. (2004) first introduced bisimulation metrics based on Wasserstein distance to quantify
behavioral similarity between states for the purpose of state aggregation. Such bisimulation metrics were
recently utilized for representation learning of RL; see e.g.. Castro (2000)
; Agarwal et al. (2021b). In addition,
a few recent work has also exploited Wasserstein distance for imitation learning (see e.g., Xiao el al. 2019
Dadashi et al. (2021p) and unsupervised RL (see e.g., He et al.
(2022)). ). Our work is closely related several
previous studies, including Richemond & Maginnis (2017 Zhang et al. (2018); Moskovitz et al. (2021)
Pacchiano et al. (2020), which also utilize Wasserstein distance to measure the proximity of policies. However, unlike the aforementioned studies that solely employ Wasserstein distance as an explicit penalty function, we additionally utilize it as a trust region constraint. Moreover, we consider nonparametrie policies and dcrive explicit policy update forms, whereas these studics update paramotric policies using policy gradicnts. Furthermore, we remonstrate monotonic performance improvement and global convergence with our policy update, which is not provided in these previous works. Regarding the use of Sinkhorn divergence in RL, Pacchiano et al. (2020 the only related work to our best knowledge, where the entropy regularization is used to mitigate the computational burden of computing Wasserstein metric. However, policy update is provided in this work, while we derive an explicit Sinkhorn policy update and demonstrate its advantage in convergence speed. Additionally, we use Wasserstein distance to directly measure the proximity of nonparametric policics in the distribution space, while Pacchiano et al. (2020;; Moskovitz et al. 2021 measure the similarity of parametric policies in the behavioral space.
explicit form of
Wasserstein-llik metrics are also pervasively studied in distributionally robust optimization (DRO); ; see e.g., Esfahani & Kuhn (2018); Gao & Kleywegl (2016); Zhao & Guan (2018); Blanchet Murthy (2019). We also point out that a recent concurrent work by Waan et al. (2021a) studied DRO using the Sinkhorn distance. Our duality formulations are largely inspired from existing work in DRO. However, we note that constrained policy optimization is conceptually different from DRO. Constrained policy optimization focuses on finding the optimistic policy that falls in a trust region, whereas DRO (e.g., the KL DRO) aims to optimize some worst-case loss given by the adversarial distribution of unknown parameters within some ambiguity set.
2 Background and Notations
Markov Decision Process (MDP): We consider an infinite-horizon discounted MDP, defined by the tuple (S, P,r,U, 7), where S is the state space, A is the action space, S: x A x S + R is the transition probability, r x A + R is the reward function, v: S + R is the distribution of the initial state so, and E (0, 1) is the discount factor. We define the return of timestep t as the accumulated discounted reward from f, R. = > **((*t*,, at+x), and the value function as V*(s) = E[[g = s; The performance of a stochastic policy π is defined ( a.....n....... I2 7'r(s,, at)] where a ~ ((ayls)), St41 ~ P(s a;). As shown in Kakade & Langford (2022). the expected return of a new policy n' can be expressed in terms of the advantage over the old policy π: J(T') = J(T)- +B Brpp'' .anniaaasooo [A (8, a)), where AT a) [R = 8, ar = a; E[R,| = s; m] represents the advantage function and P represents the unnormalized discounted visitation frequencies with initial state distribution U, i.e. (T((s) = Esowe EEe p(ps s|so)].
4

---

Published in Transactions on Machine Learning Research (06/2023)
Trust Region Policy Optimization (TRPO): In TRPO (Schulman et al. 2015). the policy is parameterized as xo with parameter vector 0. For notation brevity, we use 0 to represent the policy Then, the new policy 0' is found in each iteration to maximize the expected improvement J - J(T), or equivalently, the expected value of the advantage function:
Ee (A (s,aa) a)]
max
Ø'
s.t. Banp (6 0) S 6,
where dRL. represents the KL divergence and ō is the threshold of the distance between new and old policies.
Wasserstein Distance: Given two probability distributions of policies æ and a' on the discrete action space A = {ai, .2.............................. aag...,aa, (22.),,ay), avy, the Wasserstein distance between the policies is defined
dw( T)) inf (Q. (ETT(=(m)))
where (:-)) denotes the Frobenius inner product. The infirumm is taken over all joint distributions Q with marginals n' and n, and D is the cost matrix with Dij d(a;,a)); where d(a;; ay) is defined the distance between actions 0g and uj. Its largest entry in magnitude is denoted by || In implementation, our choice
of distance d is task-dependent and is reported in Table 3 in Appendix A
2013 provides smooth approximation of the Wasserstein distance by adding an entropic regularizer. The Sinkhorn divergence is defined as:
Sinkhorn Divergence: Sinkhorn divergence (Cuturi
QEll(T(.))) { - 1 m0))
ds( x|{))= = inf
(3)
where =(Q)= Σ5= Qig log represents the entropy term, and À >> 0 is a regularization parameter. The intuition of adding the entropic regularization is: since most elements of the optimal joint distribution Q will be 0 with high probability, by trading the sparsity with entropy, smoother and denser coupling between distributions can be achieved (Courty et all 2014 2016). Therefore, when the weight of the entropic regularization decreases (i.c., X increases), the sparsity of the divergence increases, and the Sinkhorn divergence converges to the Wasserstein metric, i.e., lim ds =TA) divergence is useful to mitigate the computational burden of computing Wasserstein distance. In fact, the efficiency improvement that Sinkhorn divergence and the related algorithms brought paves the way to utilize Wasserstein-like metrics in many machine learning domains, including online Icarning (Cesa-Biandd
(''','/))))))) = dw(=', More critically, Sinkhorn
Lugosi 2006 model selection (Juditsky et al. 2008 Rigollet. & Tsybakov (201]), generative modeling (Genevay
et al. 2018 Petzka et al. 2017 Patrini et al. 2019), dimensionality reduction (Huang et al. 2021 Lin et al.
2020
Wang et al. 2021b
3 Wasserstein Policy Optimization
Molivated by TRPO, here consider trust region based on the Wasserstein metric. Moreover, we lift the restrictive assumption that a policy has to follow a parametric distribution class by allowing all admissible policies. Then, the new policy x' is found in each iteration to maximize the estimated expected value of the advantage function. Therefore, the Wasserstein Policy Optimization (WPO) framework is shown as follows:
max E'ED where D (T' (dw ( (+1))) < <8),
(s,, (s.a)]
(4)
where the Wasserstein distance dw(..)) is defined
2
In most practical cases, the reward r is bounded and correspondingly, the accumulated discounted reward R, is bounded. So without loss of generality, we make the following assumption:
Assumption 1. Assume A () is bounded, Le., SuPaeA. | (s, & A mms for some ume 0.
5

---

Published in Transactions on Machine Learning Research (06/2023)
With Wasserstein metric based trust region constraint, update shown in Theorem The main idea is to form the Lagrangian dual of the constrained optimization problem presented above, which is inspired by the way to obtain the extremal distribution in Wasserstein DRO literature, see c.g.: Kuhn et al. (0199) Blanchet & Murthy (2019) Zhao & Guan (0118) The detailed proof can be found in Appendix
able to derive the closed-form of the policy
Theorem 1. (Closed-form policy update) Let n² (3,5) = kt N (. (s,()) BDzz), where D
denotes the cost matrir. If Assumption holds, then an optimal solution to ( is:
xt' (a(1/)) = Σ, (ayl)))1111))
(5)
where (i.j) 1 if= :) and fz (i.j) =0 0 otherwise, and 8 is an optimal Lagrangian multiplier corresponds to the following dual formulation:
N min {,5 + Eo-pr Ex=, =(azls) max (A" - 820
min F(8) 320
(6)
i=1...N
A A Morcover, have 8 3, where B mox.sskk.st...N (Dks)-'(A -
Remark 1. For ease of nolation and simplicily,, we ussume the uniqqueees of K. (B. order to form the simple expression of fž in Theorem 4 When not unique, a necessary condition for the optimality of x* in 5 is LiExy 1 (i, 1, and (i,5) = for & K, (B) where K,, (3, argmars=1...v (s, ax) Day. The weight f; (i.j) for E (B. could be determined through linear programming (see details in 170 in Appendir
The exact policy update for WPO in requires computing the optimal Lagrangian multiplier 8* by solving the one-dimensional subproblem 6 closed-form 3 is not easy to obtain in general, except for special we provide the closed-form of 3 for the case cases of the distance d(x,3) y) or cost matrix D. In Appendix C
when d(x, y) if x = y and 1 otherwise.
WPO Policy Update: Based Theorem 1 introduce the following WPO policy updating rule:
N
Th41(a;1 = WPO (TT) = Σ (1
(WPO)
where =f and 0 otherwise. Note that different from allow Be to be chosen arbitrarily and time dependently. We show that such policy update always leads to a monotonic improvement of the performance even when B₁ is not the optimal Lagrangian multiplier. In particular, we propose two strategies to update multiplier S4:
(i) Approximation of optimal Bg: To improve the convergence, we can approximately solve the optimal Lagrangian multiplier based Sinkhorn divergence. More details in Section
(ii) Timo-dependent 3r: To improve the computational efficiency, we can simply treat Bg as a time- dependent parameter, e.g., we can set 3& as a diminishing sequence. In this setting, (WPO) produces the solution the following penalty version of problem 4 (with d= dw):
max E Soupth (4) 36+1
(A*M Bx [d(Tx (-1s), TA (-1s))
(7)
4 Sinkhorn Policy Optimization
In this section, we introduce Sinkhorn policy optimization (SPO) by constructing trust region with Sinkhorn divergence. In the following theorem, we derive the optimal policy update each step when using Sinkhorn divergence based trust region. Detailed proofs are provided in Appendix [

---

Published in Transactions on Machine Learning Research (06/2023)
Theorem 2. If Assumption 1 holds, then the optimal solution to
with Sinkhorn divergence is: i
(ai))) = Σ n (aayls)55,)(01
(8)
and 5% is an optimal solution to the exp (re A (8.0,)
where D denotes the cost matriz, fa (1,3)
Ei 2 exp ( 1
following dual formulation:
N
min Fx(8) 020
min D20
p6
Bemup
Σ
T(a, 8) In(T((,|s)) - In Σ exp A
(s)
1=1
N
N
B exp (3A (s. ADb) x(a,1s))
Expp
X Σ exp (s) AD2))
Moreover, we have BX 2A
In contrast to the Wasserstein dual formulation 6 the objective in the Sinkhorn dual formulation 19 is
differentiable in and admits closed-form gradients (shown in Appendix F). With this gradient information.
can use gradient-based global optimization algorithms (Wales & Doye 1998 Zhan et al. 2006 Leary 2000 to find global optimal solution 3x to 19).
Next, we show that the entropic regularization parameter À is large enough, then the optimal solution 8% is a close approximation to the 3 of Wasscrstein dual formulation. Proof is provided in Appendix
Theorem 3. Define BuB = max m B}. We have:
(3) converyes to F(B) uniformly on [0, BuBl: lim sup x og85Ave
FA(3)- ((13)-F(B)) F(B)
S lim A
N
X00 X
2. lim argmingg85san (x S argmingsss@0m F(3).
>
Although difficult to obtain the exact value of the optimal solution 3* to the Wasserstein dual formulation 6 the above theorem suggests that we approximate 8 via 8 by setting up relative large X. In practice, we can also adopt smooth homotopy approach by setting an increasing sequence X for cach iteration and letting d + xc.
SPO Policy Update: Based on Theorem 2 we introduce the following SPO policy updating rule:
(a:ls) (TR)) Σ Ta(agls) (i.5).
(SPO)
Here S2 (i, En, exp (x: ds 0 and Br are some control parameters. The
parameter 3g be either computed via solving the one-dimensional subproblem 9 simply set diminishing sequence. The proper setup of d can effectively adjust the trade-off between convergence speed and final performance. More details are provided in the ablation study in Section
5 Theoretical Analysis
We first show that SPO policy update converges to WPO policy update the regularization parameter increases (i.e., X + oo). The detailed proof is provided in Appendix
Lemma 1. As d oo, SPO updale converges to WPO update: lim to0 to0 (Tx) E B WPWOO

---

Published in Transactions on Machine Learning Research (06/2023)
We then provide theoretical justification that WPO policy update (and SPO with À - o0) are always guaranteed to improve the true performance J monotonically if we have access to the true advantage function. If the advantage function can only be evaluated inexactly with limited samples, then an extra estimation error (measured by the largest absolute entry || lloo) will occur. Proof can be found in Appendix !
Theorem 4. (Performance improvement) For any initial state distribution vs and any B& 0, if || - A S for some e > O, let (3x.)) = argmas:=1..... (A (8, a:) - Dist WPO policy update (and SPO with À + 20) guarantee the following performance improvement bound when the inaccurate aduantage function A is used,
J( 2 J(Ts) + B&T.
Σ (s(alss)
(1,D)) (B..)
Σ
2e 1
(10)
smpp 111
The value of e, which quantifies the approximation error of the advantage function, is dependent various factors such as the advantage estimation algorithm used and the number of samples (Schulman et al. 2016 It is worth noting that the improvement bound of NPG/TRPO (Cen et al. 2021 includes the same additional term - T based updates. In the following, we show that with a decreasing schedule of the multiplier 3e, both WPO and SPO policy updates have their values J(T&)) converging to the optimal J = max, J(r) on the tabular domain. To start, for k-th iteration, we consider (WPO) and (SPO) (with arbitrary À >> 0) whose updates TT++1 are optimal solutions to 7 with d being dw and ds respectively.
2z_ which indicates that our methods offer comparable theoretical performance guarantees to KL
Assumption 2. The state space and the action space are both finit, the reward function is non-negative, and the initiol distribution covers all state.
Note that once state and action spaces both finite, the reward can be assumed non-negative without loss of generality, as we can always add maxs, [r(s, a)] to the reward function without changing the optimal
policy and the order of the policies. Defining the optimal value function (*(s) = max, B,R, st = s]. we have
the following theorem, whose proof is in Appendix J and is inspired by Bhandari & Russo 2021
Theorem 5. (Global convergence) Under Assumption
2
we have for any Br 2 0,
WPO satisfies that
I|V V=++i 1l00 S7IVV*
- 1lo + DI
(11)
and SPO
satisfies that
Ax
(ID) + +2 log N A
V V V 1l00 S ~/ - V 11 + +2
(12)
$21 1-
If limm 250 3k= 0, we further have lim J (TT) J Remark 2. Note the convergence is geometric. If we keep BR as a constant, then S J -- V + I- where B = =I111 for (WPO) and = Dl Non To achieve an e optimality gap, we only need to take 3 = 2B
J(xTT] TT)
||V* - - 1100
I 1-y
for SPO
(1-7)e and let T> log(s/2)
Remark 3. The study of global non-asymplotic convergence of noncoweer policy optimization algorithms has been an active research topic. Recent theoretical work has mostly centered on PG and natural policy gradient (NPG) Kakade] (2001)- a close relative of TRPO; see e.g., [Agarwal al. (2021a); Cen et al.) (2021); Lan (202)). To our best knowledge, a few work has discussed the global convergence of TRPO. Neu et al. 2017| 1 and Geist et al. (2019) established the connection of TRPO to Mirror Descent, but did not provide any non-asymptotic rate; Shani et al.) (2020) showed that adaptive TRPO with decaying stepsize achieved O(IVT)) convergence rate for unregularized MDPs in the tabular setting (finite state and finite action). Our resull seems to be the first non-asymplooic analysis of policy oplimization based on Wasserslein ond to function
Sinkhorn divergence. It remains interesting to eztend the convergence theory of ERPO/WPO/SS approzimation regime following recent advance Agarwal et al.) (2021a). However, this is beyond the scope of our current work, as we focus on explicit dosed-form update of WPO/SPO, which con be a viable alternative to TRPO in practice.
8

---

Published in Transactions on Machine Learning Research (06/2023)
6 A Practical Algorithm
In practice, the advantage value functions are often estimated from sampled trajectorics. In this section, we provide a practical on-policy actor-critic algorithm, described in Algorithm ] that combines WPO/SPO with advantage function estimation.
At cach itcration, the first step collect trajectorics, which can be either complete or partial. If the trajectory is complete, the total return can be directly expressed as the accumulated discounted rewards T If the trajectory is partial, it can be estimated by applying the multi-step temporal k= -
R₁ = E difference (TD) methods DD Asis et al. 2017): fuxt+ = Ek= " + +7"V(se+a). Then for the advantage estimation, we can use Monte Carlo advantage estimation, i.e., A, = Rs - Vv, (sx) or Generalized Advantage Estimation (GAE) (Schulman et al.) 2016). which provides a more explicit control over the bias-variance trado- off. In the value update step, we use a neural net to represent the value function, where & is the parameler that specifies the value net 8 - V(s). Then, we can update v by using gradient descent, which significantly reduces the computational burden of computing advantage directly. The computational complexity of the algorithm is discussed in Appendix K
Algorithm 1: On-policy WPO/SPO algorithm
Input: number of iterations K, learning rate Initialize policy To and value network Vv with random parameter v/o for k 1.,....A. do Collect trajectory set D2 on policy For each timestep in each trajectory, compute total returns G2 and estimate advantages Ar Update value:
4/++1 < 4/R VV E(G. Vxr)))))
Update policy:
Tk+1 + (TTR) via WPO/ SPO with Ap
end
7 Experiments
In this section, we evaluate the proposed WPO and SPO approaches presented in A Algorithm 1 We compare the performance of our methods with benchmarks including TRPO (Schulman al. 2015). PPO (Schulman ct al. 2017), A2C (Mnih al. 2016); and with BGPG (Pacchiano al.] 2020), WNPG (Moskovitz al. 2021 for continuous control. The code of our WPO/SPO can be found here We adopt the implementations of TRPO, PPO and A2C from OpenAT Baselines (Dhariwal et al.) 2017) for MuJuCo tasks and Stable
Baselines Hill et al. 2018 for other tasks. For BGPG, we adopt the same implementation as (Pacchiano
ct al. 2020
Our experiments include (1) ablation study that focuses on sensitivity analysis of WPO and SPO; (2) tabular domain tasks with discrete state and action including the Taxi, Chain, and Cliff Walking environments; (3)
locomotion tasks with continuous state and discrete action including the CartPole, Acrobot environments; (4) comparison of KL and Wasserstein trust regions under tabular domain and locomotion tasks; and (5) extension to continuous control tasks with continuous action including HalfCheetah, Hopper, Walker, and Ant environments from MuJuCo. See Table in Appenddi A for a summary of performance. The setting of hyperparameters and network sizes of our algorithms and additional results are provided in Appendix
7.1 Ablation Study
In this experiment, we first. examine the sensitivity of WPO in terms of different strategies of 8x- We test four settings of 3 value for WPO policy update: (1) Setting 1: Computing optimal 3 value for all policy update; (2) Setting 2: Computing optimal B value for first 20% of policy updates and decaying 3 for the
9

---

Published in Transactions on Machine Learning Research (06/2023)
remaining; (3) Setting 3: Computing optimal 3 value for first 20% of policy updates and fix B as its last updated value for the remaining; (4) Setting 4: Decaying & for all policy updates (e.g,, 3 = ( log k)). In particular, Setting 2 is rooted in the observation that 8 decays slowly in the later stage of the experiments carried out in the paper. Small perturbations are added to the approximate values to avoid any stagnation in updating. Taxi task (Dietterich 1998) from tabular domain selected for this experiment.
Table 1: Runtime for different 3 settings, average across 5 runs with random initialization
Runtime
Taxi (s)
CartPPlc (s)
Setting 1 (optinal 8)
1224.3 +: 105.7
129.7 1: 15.2
Setting 2 -heenoooyy))
648.4 1: 55.7
63.2 1 8.3
Sotting 3 (optimal-then-fix)
630.2 1 67.4
67.1 + 9.7
Setting 4 (decaying (9)
522.7 + 49.5
44.3 + 6.2
The performance comparisons and average run times are shown in Figure 4 and Table respectively. Figure 4a and Table [ clearly indicate a tradeoff between computation efficiency and accuracy in terms of different choices of 3 value. Setting 2 is the most effective way to balance the tradeoff between performance and run time. For the rest of experiments, we adopt this setting for both WPO and SPO (see Appendix [A.2 for how Sctting 2 is tuned for each task). Figure 4b shows that as increases, the convergence of SPO becomes slower but the final performance of SPO improves and becomes closer to that of WPO, which verifies the convergence property of Sinkhorn to Wasserstein distance shown in Theorem 3 Therefore, the choice of X can effectively adjust the tradc-off between convergence and final performance. Similar results are observed when using time-varying À on Taxi, Chain and CartPole tasks, presented in Figure in Appendix
Taxi
Taxi
0
-100
-100
Average Return
Average Return
-20
-200
-300
-300
-400
WPO: Optimal
-400
WPO: Optimal decay
-500
SPO(A 4)
-500
WPO: Optimal @
SPO (A 2)
SPO(A-5) 5)
-600
-600
WPO: Decay
SPO (A 3)
SPO 6)
-700
0.0
0.5
1.0
1.5
2.0
2.5
0.0
0.5
1.0
1.5
2.0 x10
Timesteps
x10
Timesteps
(a) Different settings of /
(b) Different choices
Figure 4: Episode rewards for Taxi with different 3 and settings, averaged across 5 runs with random initialization. The shaded area depicts the mcan + the standard deviation.
7.2 Tabular Domains
We evaluate WPO and SPO on tabular domain tasks and test the exploration ability of the algorithms on several environments including Taxi, Chain, and Cliff Walking. We use table of size [S] x |A] to represent. the policy x[a]s). For the value function, we use neural net smoothly update the values. The performance of WPO and SPO are compared to the performance of TRPO, PPO and A2C under the same neural net. structure. Results on Taxi, Cliff and Chain are reported in Figure
As shown in Figure the performances of WPO, SPO and TRPO are manifestly better than A2C and PPO. Among the trust region based methods, WPO and SPO outperform TRPO in Taxi and Cliff Walking, whereas in Chain, the performances of these three methods are comparable. In all of the test cases, SPO converges faster than WPO but to a lower optimum. As further shown in Table 2 for the Taxi environment, WPO has a higher successful drop-off rate and a lower task completion time while the original TRPO reaches the time limit with drop-off rate 0, suggesting that WPO finds better policy than the original TRPO. In Figure 7
10

---

Published in Transactions on Machine Learning Research (06/2023)
also compare the performance of WPO under Wasserstein and KL divergences given different number of samples NA used to estimate the advantage function, and the result suggests that using Wasserstein metric is more robust than KL divergence under inaccurate advantage values.
Taxi
Chain
Cliff
4000
-100
Averago Return
-200
Average Return
-500
Averaçe Return
3500
-300
-1000
3000
-500
WPD
-1500
2500
8
2200
-600
-700
12C
-2500
2000
PPO
3000
1500
0.0
0.5
1.0
1.5
2.0
2.5
0.05
0.15
0.25
0.30
0.0
0.5
1.0
2.0 25
3.0 3.5
4.0
Timesleps
x10
Timestaps
x10>
Timesleps
x10
Figure 5: Episode rewards during training for tabular domain tasks, averaged across 5 runs with random initialization. The shaded area depicts the mean the standard deviation.
Table 2:
Trained agents performance on Taxi
Success (+20) Fail (-10)
Steps (-1)
Return
WPO
0.753
0.232
70.891
-58.151
TRPO
200
-200
7.3 Robotic Locomotion Tasks
We now integrate deep neural network architecture into WPO and SPO and evaluate their performance on several locomotion tasks (with continuous state and discrete action), including CartPolc (Barto et al., 1983) and Acrobot (Geramilard et al., 2015). We use two separale neural nets to represent the policy and the value. The policy neural net receives state s as an input and outputs the categorical distribution of T(als). A random subset of states Sg E S is sampled each iteration perform policy updates.
Figure O shows that WPO and SPO outperform TRPO, PPO and A2C in most tasks in terms of final performance, except in Acrobot where PPO performs the best. Tn most cases, SPO converges faster but WPO has a better final performance. To train 105 timesteps in the discrete locomotion tasks, the training wall-clock time 63.2 +8.2s for WPO, 64.7 #7.8s for SPO, 59.4 110.s for TRPO and 69.9 10..5s for PPO. Therefore, WPO has a similar computational efficiency as TRPO and PPO.
500
CartPole-vr1
CartPole
Acrobot-v1
Acrobot-v1
8PO
400
-100
ADC
TRIPO
Average Relurn
300
Retum
-150
-200
150
2010
PRO
-250
250
230
200
330
-350
350
1001
100
4400
-400
450
0
550
25
50
75
125 Falenda
100
150
175
200
50
100
150
200
250
10
20
25
30
35
Runtima (seconds)
Episode
Runlima (secridd)
Figure 6: Episode rewards during the training process for the locomotion tasks, averaged across 5 runs with random initialization. The shaded area depicts the mean + the standard deviation.
7.4 Comparison of Wasserstein and KL Trust Regions
We show that compared with the KL divergence, the utilization of Wasserstein metric can cope with the inaccurate advantage estimations caused by the lack of samples. Let NA denote the number of samples used
11

---

Published in Transactions on Machine Learning Research (06/2023)
to estimate the advantage function. We evaluate the performance of WPO framework 44) with Wasserstein and KT, constraints (as derived in Peng et al. (2019)). ). We consider the Chain task and different NA. As shown in Figure 7 when NA is 1000, KL performs slightly better than WPO. However, when NA decreases to 100 or 250, WPO outperforms KL. These results indicate that WPO is more robust than KL under inaccurate advantage values. This finding is consistent with our observations on the policy update formulations of Wasserstein and KT. For the Wasserstein update in 5 policy will be updated only when the advantage difference between two actions significant, i.e., AT a.) - Bij 2 (s, However, for the KL update in Peng et al. (2019), policy will be updated as long as the current advantage function has single non-zero value. Therefore, KT. update is more sensitive; while Wasserstein update is more robust and more tolerant to advantage inaccuracies. Similar results are obtained for the locomotion tasks (Figure 10 in Appendix
The runtime of Wasserstein and KL. updates are reported in Table 5 in Appendix
Chain
Chain
Chain
4000
4000
4000
3500
3500
3500
Average Return
3000
Averaee
2500
Average
2000
2000
2000
WPO
WPO
WPO
1500
1900
1500
1000
0.0
3.5 06
1000
0.0
0.4 0.5
100g
0.4 0.5
Timestops
K10"
Timestops
K10"
Timestops
x10*
(a) N4 100
(b) NA 250
(c) Na 1000
Figure 7: Episode rewards during training for the Chain task, where advantage value function is estimated under different number of samples, averaged across 5 runs with random initialization. The shaded area depicts the mean + the standard deviation.
7.5 Extension to Continuous Control
To extend to environments with continuous action, we use Implicit Quantile Networks (IQN) WWil Dabney &
Munos 2018 actor that can represent arbitrary complex non-parametric policy. Let F-1 (p) represent the quantile function associated with policy (-1ss.. The TQN actor takes state and probability E [0, 1] as input, and outputs the corresponding quantile value F==FF-10 (p). IQN actor can be trained to approach pre-defined
target policy distributions through quantile regression (Will Dabney & Munos 2018
Tessler et al. 2019).
Define the action support for state & in k-th iteration as 1 (s) = {a' : (s, a') > min (x) A** ())
Then, the WPO/SPO target policy distribution to guide IQN update in the k-th iteration is:
(p((s)(a'(s)= -E (6) m ((dals)fa(d',a))
(13)
where for WPO update aa) = 1 if = arr (AT (s) - Bed(d',a) a)} and fa a) = 0
2x (saa (a))
otherwise; for SPO update, fa(d'aa)= E Exterbron exp( s1 -xedaaaa))) In implementation, we sample
batch of states Sy E S at each iteration to perform policy updates, and for each E Sk, we sample [Ax actions to approximate the support (s) and the target policy distribution (-)l(s)
We additionally compare WPO and SPO with BGPG (Pacchiano et al. 2020 and WNPG (Moskovitz et al. 2021) that are specially designed to address the continuous control with Wasserstein metric, for several MuJuCo tasks including HalfCheetah, Hopper, Walker, and Ant. Figure g shows that WPO and SPO have consistently better performances than other benchmarks. Similar results are obtained for the challenging Humanoid task, presented in Figure I in Appendix A. We also provide the runtime of each algorithm in
Table K in Appendix A
12

---

Published in Transactions on Machine Learning Research (06/2023)
HaliCheetah--2
Hopper-v2
Walker2d-v2
Ant-22
3600
2500
WPO
ABC
1500
Average Retur
3000
2600
3000
370
2000
800
T9PO
1900
2000
1900
Awersge
TRPO
1500
1500
500
MP0
42C
A01
500
50PO
500
0
TIm estast
x10
Timeztops
x100
Timestaps
x10
Timestapa
x10
Figure Episode rewards during training for MuJuCo continuous control tasks, averaged across 10 runs with random initialization. The shaded area depicts the mean the standard deviation.
8 Conclusion
In this paper, we present two policy optimization frameworks, WPO and SPO, which can exactly characterize the policy updates instead of confining their distributions to particular distribution class requiring any approximation. Our methods outperform TRPO and PPO with better sample efficiency, faster convergence, and improved final performance. Our numerical results show that the Wasserstein metric is more robust to the ambiguity of advantage functions, compared with the KL divergence. Our strategy for adjusting 8 value for WPO can reduce the computational time and boost the convergence without noticeable performance degradation. SPO improves the convergence speed of WPO by properly choosing the weight of the entropic regularizer. Performance improvement and global convergence for WPO are discussed. For future work, it remains interesting extend the idea PPO and natural policy gradients, which penalize the policy update instead of imposing trust region constraint, and extend off-policy frameworks.
13

---

Published in Transactions on Machine Learning Research (06/2023)
References
Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a posteriori policy optimisation. ArXiv Preprint, pp. arXiv:1806.0692) 2018.
Alekh Agarwal, Sham M Kakade, Jason D Lee, and Gaurav Mahajan. On the theory of policy gradient methods: Optimality, approximation, and distribution shift. J. Mach. Learn. Res., 22(98):1 76, 2021a.
Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Mare G Bellemare. Contrastive behavioral similarity embeddings for generalization in reinforcement learning. In Proceedings of the 9th International Conference on Learning Representations, 2021b.
Andrew G. Barto, Richard S. Sutton, and Charles W. Anderson. Neuronlike adaptive elements that can solve difficult learning control problems. IEFE Transactions on Systems, Man, and Cybernetics, SMC-13(5): 534-846, 1983.
Jalaj Bhandari and Daniel Russo. On the linear convergence of policy gradient methods for finite mdps. In International Conference on Artificial Intelligence and Statistics, pp. 2386 2394. PMLR, 2021.
Jose Blanchet and Karthyek Murthy. Quantifying distributional model risk via optimal transport. Mathematics of Operations Research, 44(2):565 600, 2019.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic Markov decision processes. In Proceedings of the Fhirty-Fourth AAAI Conference on Artificial Intelligence, volume 34, pp. 2020.
Shicong Cen, Chen Cheng, Yuxin Chen, Yuting Wei, and Yuejie Chi. Fast global convergence of natural policy gradient methods with entropy regularization. Operations Research, 2021.
Nicolo Cesa-Bianchi and Gábor Lugosi. Prediction, learning, and games. Cambridge University Press, 2006.
Nicolas Courty, Rémi Flamary, and Devis Tuia. Domain adaptation with regularized optimal transport. In Joinl. Faropean Conference on Machine Learning and Knowledge Discovery in Dalabases, pp. 274-289. Springer, 2014.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Anolysis and Machine Intelligence, 39(9):1853 1865, 2016.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transport. In Advances in Neural Information Processing Systems, volume 26, pp. 2292 2300, 2013.
Robert Dadashi, Léonard Hussenot, Matthieu Geist, and Olivier Pietquin. Primal Wasserstein imitation learning. In Proceedings of the 9th International Conference on Learning Representations, 2021.
Kristopher De Asis, J. Fernando Hernandez-Garcia. G. Zacharias Holland, and Richard S. Sutton. Multi-step reinforcement learning: A unifying algorithm. ArXiv Preprint, pp. arXiv:1703.0132. 2017.
Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plapport, Alec Radford, John Schulman, Szymon Sidor, Yuhuai Wu, and Peter Zhokhov. OpenAl baselines. httpp:://itthbb.com openai/baselinene 2017.
Thomas G. Dietterich. The MAXQ method for hierarchical reinforcement learning. In Proceedings of the 15th International Conference on Machine Learning, pp. 118- 126, 1998.
Yan Duan, Xi Chen, Rein Houthooft, John Schulman, and Pieter Abbeel. Benchmarking deep reinforcement. learning for continuous control. In Proceedings of the 33rd International Conference on Machine Learning, pp. 1339-1338, 2016.
Peyman Mohajerin Esfahani and Daniel Kuhn. Data-driven distributionally robust optimization using the Wasserstein metric: Performance guarantees and tractable reformulations. Mathematical Programming. 171 (1):115-166, 2018.
14

---

Published in Transactions on Machine Learning Research (06/2023)
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite Markov decision processes. In Uncertainly in Artificial Intelligence, volume 4, pp. 162-169, 2004.
Rui Gao and Anton J Kleywegt. Distributionally robust stochastic optimization with Wasserstein distance. arXiv preprinl arx: 0614000999 02199, 2016.
Matthieu Geist, Bruno Scherrer, and Olivier Pielquin. A theory of regularized markov decision processes. In International Conference on Machine Learning, pp. 2160 2169. PMLR, 2019.
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with Sinkhorn divergences. In International Conference on Artificial Intelligence and Stalistics, pp. 1608-1617, 2018.
Alborz Geramifard, Christoph Dann, Robert H. Klein, William Dabney, and Jonathan How. RLPy: A value-functiion-basee reinforcement learning framework for education and research. Journal of Machine Learning Research, 16(46):1573-155) 1578,, 2015.
Gregory Z. Grudic, Vijay Kumar, and Lyle H. Ungar. Using policy gradient reinforcement learning on autonomous robot controllers. In Proceedings of the 2003 IEEE/RS] International Conference Intelligent Robots and Systems, pp. 406-411, 2003.
Shuncheng He, Yuhang Jiang, Hongchang Zhang, Jianzhun Shao, and Xiangyang Ji. Wasserstein Unsupervised Reinforcement Learning. In Proceedings of the Chirty-Sizth AAAI Conference on Artificial Intelligence, 2022.
Nicolas Heess, Gregory Wayne, David Silver, Timothy Lillicrap, Tom Erez, and Yuval Tassa. Learning continuous control policies by stochastic value gradients. In Advances in Neural Information Processing Systems, volume 28, 2015.
Ashley Hill, Antonin Raffin, Maximilian Ernestus, Adam Gleave, Anssi Kanervisto, Rene Traore, Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, and Yuhuai Wu. Stable baselines.
baseliness 2018.
Minhui Huang, Shiqian Ma, and Lifeng Lai. A Riemannian block coordinate descent method for computing the projection robust Wasserstein distance. In Proceedings of the 38th International Conference on Machine Learning, pp. 4446 4465, 2021.
Anatoli Juditsky, Philippe Rigollet, and Alexandre B Tsybakov. Learning by mirror averaging. The Annals of Statistics, 36(5):2183-220) 2008.
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. In Proceedings of the 19th International Conference on Machine Learning, pp. 267-274, 2002.
Sham M Kakade. A natural policy gradient. In Advances in Neural Information Processing Systems, volume 14, 2001.
Daniel Kuhn, Peyman Mohajerin Esfahani, Viet Anh Nguyen, and Soroosh Betteezaaee-Abaadee Wasserstein distributionally robust optimization: Theory and applications in machine learning. In Operations Research & Management Science in the Age of Analytics, pp. 130 166. INFORMS, 2019.
Guanghui Lan. Policy mirror descent for reinforcement learning: Linear convergence, new sampling complexity, and generalized problem classes. Mathematical programming, pp. 1-48, 2022.
Robert Leary. Global optimization on funneling landscapes. Journal of Global Optimization, 18:367 383, 2000.
Sergey Levine, Chelsca Finn, Trevor Darrell, and Picter Abbeel. End-to-ond training of deep visuomotor policies. Journal of Machine Learning Research, 17(39):1 40, 2016.
Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Hoess, Tom Erez, Yuval Tassa, David Silver, and Daan Wicrstra. Continuous control with deep reinforcement learning. In Proceedings of the 4th International Conference on Learning Representations, 2016.
15

---

Published in Transactions on Machine Learning Research (06/2023)
Tianyi Lin, Chenyou Fan, Nhat Ho, Marco Cuturi, and Michael I Jordan. Projection robust Wasserstein distance and Riemannian optimization. ArXiv Preprint, pp. arXiv:2006.07458, 2020.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Proceedings of the 33rd International Conference on Machine Learning, pp. 1928-1937, 2016.
Ted Moskovitz, Michael Arbel, Ferenc Huszar, and Arthur Gretton. Efficient Wasserstein natural gradients for reinforcement learning, In Proceedings of the 9th International Conference on Learning Representations, 2021.
Gergely Neu, Anders Jonsson, and Viceng Gómez. A unified view of entropy-rreuularized markov decision processes. ar preprint ArXiv:1705. 07798, 2017.
Aldo Pacchiano, Jack Parker-Holes Holder, Yunhao Tang, Krzysztof Choromanski, Anna Choromanska, and Michael Jordan. Learning to score behaviors for guided policy optimization. In Proceedings of the 37th International Conference on Machine Learning, pp- 7445-74444 7754,, 2020.
Victor M. Panaretos and Yoav Zemel. Statistical aspects of wasserstein distances. Annual Review of Statistics and Its Application, 6(1):405-493) 2019.
Giorgio Patrini, Rianne van den Berg, Patrick Forre, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. In Uncerlainly in Artificial Intelligence, pp. 733 743, 2019.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantaage-wiighted regression: Simple and scalable off-policy reinforcement learning. ArXiv Preprint, pp. arXiv:1910.00177, 2019.
Jan Peters and Stefan Schaal. Policy gradient methods for robotics. In Proceedings of the 2006 IEEE/RS, International Conference on Intelligent Robots and Systems, pp. 2219-2225, 2006.
Henning Petzka, Asja Fischer, and Denis Lukovnicov. On the regularization of Wasserstein GANs, ArXiv Preprinl, pp. arXiv: 1rXivv1779.000 709.08894,...0 2017.
Pierre H. Richemond and Brendan Maginnis. On Wasserstein reinforcement learning and the Bokker-Plandee equation. ArXiv Preprint, pp. arXiv:1712.071855, 2017.
Philippe Rigollet and Alexandre Tsybakov. Exponential screening and optimal rates of sparse estimation. The Annals of Slatistics, 39(2):731 771, 2011.
R. Tyrrell Rockafellar and Roger J. B. Wets. Variational Analysis. Springer, 1998.
Johannes O. Royset. Approximations and solution estimates in optimization. Mothematical Programming, 170:479 506, 2018.
John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan, and Pieter Abbeel. Trust region policy optimization. In Proceedings of the 32nd International Conference on Mochine Learning, pp. 1889-188 1897, 2015.
John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. In Proceedings of the 4th International Conference on Learning Representations, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv Preprint, pp. arXiv:1707.06347, 2017.
Lior Shani, Yonathan Efroni, and Shic Mannor. Adaptive trust region policy optimization: Global convergence and faster rates for regularized mdps. In Proceedings of the Thirty-Fourree AAAI Conference on Artificial Intelligence, volume 34, pp. 5668-56666 2020.
16

---

Published in Transactions on Machine Learning Research (06/2023)
David Silver, Guy Lever, Nicolas Heess, Thomas Degris, Daan Wierstra, and Martin Riedmiller. Deterministic policy gradient algorithms. In Proceedings of the 31st International Conference on Machine Learning, pp- 387 395, 2014.
Maurice Sion. On general minimax theorems. Pacific Journal of Mathematics, 8(1):171-176, 1958.
Student. The probable error of a mean. Biometrika, 6(1):1-25, 1908.
Richard S. Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Advances in Newral Information Processing Systems, pp. 1057 1063, 1999.
Chen Tessler, Guy Tennenholtz, and Shie Mannor. Distributional policy optimization: An alternative approach for continuous control. In Advances in Neural Information Processing Systems, pp. 1350-1360, 2019.
David Wales and Jonathan Doye. Global optimization by basin-hopping and the lowest energy structures of Lennard-Jones clusters containing up 110 atoms. The Journal of Physical Chemistry A, 101(28): 5111-5116, 1998.
Jie Wang, Rui Gao, and Yao Xie. Sinkhorn distributionally robust optimization. ArXiv Preprint, pp- 1991 2021a.
Jie Wang, Rui Gao, and Yao Xic. Two-sample test using projected Wassorstein distance. In Proceedings of IEEE International Symposium on Information Theory, volume 21, 2021b.
Ziyu Wang, Victor Bapst, Nicolas Heess, Volodymyr Mnih, Remi Munos, Koray Kavukcuoglu, and Nando de Freitas. Sample efficient actor-critic with experience replay. In Proceedings of the 5th International Conference on Learning Representations, 2017.
David Silver Will Dabney, Georg Ostrovski and Rémi Munos. Implicit quantile networks for distributional reinforcement learning. In International Conference on Machine Learning, 2018.
Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist. reinforcement learning. Machine Learning, 8(3 4):229 256, 1992.
Yuhuai Wu, Elman Mansimov, Shun Liao, Roger Grosse, and Jimmy Ba. Scalable trust-region method for deep reinforcement learning using chonecker-factoree approximation. In Advances in Neural Information Processing Systems, pp. 2295, 2017.
Huang Xiao, Michael Herman, Joerg Wagner, Sebastian Ziesche, Jalal Etesami, and Thai Hong Linh. Wasserstein adversarial imitation learning. ArXio Preprint, pp. arXiv:1906.08115 2019.
Lixin Zhan, Jeff Chen, and Wing-Ki Ki Liu. Monte Carlo basin paving: An improved global optimization method. Physicol Review E, 73:015701, 2006.
Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin. Policy optimization as Wasserstein gradient flows. In Proceedings of the 35th International Conference on Machine Learning, pp. 5737 5746, 2018
Chaoyue Zhao and Yongpei Guan. Data-driven risk-averse stochastic optimization with Wasserstein metric. Operations Research Letters, 46(2):262-267, 2018.
17

---

Published in Transactions on Machine Learning Research (06/2023)
A Implementation Details and Additional Results
The implementation of our WPO/SPO can be found in https://github..co http://iinbooooccccctttttttoooop Efficient We
use the implementations of TRPO, PPO and A2C from OpenAl Baselines (Dhariwal et al. 2017 for MuJuCo tasks and Stable Baselines Hill et al. 2018) for other tasks. For BGPG, we adopt the same implementation as Pacchinao et al., (2020) based on the released code https://itthub.comm http:////dttbboooo////rrrrrrrrrd
A.1 Visitation Frequencies Estimation:
The unnormalized discounted visitation froquencics are needed to compute the global optimal 3 At the k-th iteration, the visitation frequencies P are estimated using samples of the trajectory set Dx. Specifically, we first. initialize (s) = O, Vs E S. Then for each timestep each trajectory from Dx, we update P as (st) + P (5c) /Dxl-
A.2 Optimal-then-deeee Beta Strategy:
During the training multiple tasks, including Taxi, Chain and CartPole, we observe a consistent trend in the behavior of the optimal value during the policy updates: It initially fluctuates, then stabilizes and decays slowly towards 0. In the Taxi task, the optimal stabilizes after approximately 18% of the total training iterations. If decay 3 before this stabilization point (e.g, using optimal beta for only first 5% 10% updates), we observe drop in performance. However, we do not observe any notable performance difference when we decay 3 after this stabilization point (e.g., using optimal 3 for first 20% or 30% updates). We also observe that the optimal 3 decays at a very slow rate, and 6(1 ((]/log(k)) matches this trend best. If we employ a faster decaying function, such as ((1/k) or @(1/k²), we observe drop in performance.
Based on these findings, when implementing the optimal-theee then-decay 8 strategy on other tasks, we compute the optimal B for each policy update until we observe that its value stabilizes across updates. At this point, stop calculating the optimal 3 and decay it using ((1/ ([[/lo(k))) for the remaining policy updates. The specific iteration at which the optimal value stabilizes varies across tasks, and we denote this point as kg, which is reported in Table 3
A.3 Hyperparameters and Performance Summary
Our main experimental results are reported in section 7 In addition, we provide the setting of hyperparametors and network sizes of our WPO/SPO algorithms in Table 3 and summary of performance in Table 4
Table 3: Hyperparameters and network sizes
Taxi-v3
NChain-vO
CartPole-vl
Acrobot-vl
MuJuCo tasks
CliffWalking-v0
7
0.9
0.9
0.95
0.95
0.99
bx
10-2
x 10-3
10-4
Irvatue
10-2
10-2
10-2
x 10-3
10-3
[Del
60 (Taxi)
(Chain)
2
3
partial
(ClifWWalking)
size
2D array
2D array
664, 64]
[64, 64]
[400,300]
Q/v size
[10, 7,5]
[10, 7,
(64, 64)
[64, 64]
[400,300]
[Sal
all states,
[S]
all states, [S]
128
128
64
[A
all actions, [4] all actions,
[A]
all actions, [A] all actions, [4]
32
18

---

Published in Transactions on Machine Learning Research (06/2023)
d(a,a')
0-1 distance 3
0-1 distance
0-1 distance
0-1 distance
L1 distance
kg
250
100 (Chain)
150
150
1000
50 (ChifWalking)
Table 4: Averaged rewards over last 10% episodes during the training process
Environment
WPO
SPO
TRPO
PPO
A2C
BGPG
WNPG
Taxi-v3
45 + 27
87 11
200 3
-381 + 31
-338 4: 30
NChain-vO
3549 + 197
3432 + 131
3522 258
3506 237
1606
ClifWalking-vo)
-35 + 15
-5 +
-559 + 94
-3290 + 2106
-5587 1942
CartPole-vl
388 + 54
370 + 30
297 + 65
193 + 45
267 + 61
Acrobot-vl
-162 =
185 + 15
248 = 33
1003 +
337 + 39
HallCheetal-v2
2050 + 108
1750 +: 172
1158 + 35
1628 = 136
-45 +: 31
1697 1: 195
1832 = 125
Hopper-v2
3208 + 259
2834 + 305
2035 248
2321 233
43
1982 + 218
2361
272
Walker2d-v2
3739 + 298
3489 + 257
2535 369
3290 354
28 1
2775 + 301
3059 = 209
Ant-v2
1863 + 271
1780 + 257
21 = 10
1487 = 206
39 + 8
1622 + 235
1587 = 221
Humanoid-v2
965 + 76
914 +: 93
725 + 112
632 +: 73
107 +: 15
797 1 85
820 1:: 91
A.4
Additional Results for Ablation Studies
500
CartPole-v1
Chain
0
Taxi
WPO Optial B
3500
Average Return
400
WPO. Optinal comdant
WPO Optnal + dacay F
Average Retum
3000
Average Return
-100
-200
WPO. Deayy
300
2500
-300
200
-400
WPO
WPO
SPO (1- 50)
SPO (oog0
100
1500
SPO
SPO
-500
SPO A--0
SPQ
logt)
SPO (-
0
25
50
75
150 175 200 Episode
100
125
1000
0.0
0.1
0.2
0.3
0.4
0.5
-600
0.0
0.5
1.0
1.5
2.0
Timesteps
x 100
Timesteps
x10
Figure 9: Episode rewards during the training process for different: and X settings, averaged across 5 runs with a random initialization. The shaded area depicts the mean + the standard deviation.
We note that specifying distance based control relevance leads higher performance this test i.c., = 1 to distinct actions from A move north, move south, move west, move east)) d distinct actions from B pickup, dropoff), and actions from different
19

---

Published in Transactions on Machine Learning Research (06/2023)
A.5 Additional Comparison of Wasserstein and KL Trust Regions
Ca
500
CartPole-v1
-103
Acrobot-v1
-100
Acrobot-v1
WPO
WPO
Nerrag Retur
KL
Retum
RL
-203
-200
-200
-303
-300
200
-350
100
100
400
-400
0
500
550
Corodde
Epsodd
Escrade
(a) NA = 100
(b) N 500
(c) NA = 100
(d) NA = 500
Figure 10: Episode rewards during the training process for the locomotion tasks, averaged across 5 runs with a random initialization. The shaded area depicts the mean = the standard deviation.
Table 5: Average runtime (seconds) of WPO, SPO and KL
WPO
SPO
K1.
Taxi-v3 (per 10a steps)
71.0 +: 7.3
69.5 1: 8.7
74.3 9.5
NChain-vO (per 10a steps)
58.4 9.1
63.1 7.4
59.9 8.7
CartPole-vl (per 10 steps)
11.4 1.8
10.2 2.3
9.7 1.9
Acrobot-vl (per 10 steps)
10.4 # 1.9
9.7 +
10.9 2.3
Humanoid-v2 (per 10 steps)
422.7 65.4
409.1 + 46.5
438.5 61.2
A.6 Additional Results for Large-scale Continuous Control
Humanoid-v v2
1000
900
Average Return
800
700
600
500
WPO
A2C
400
SPO
TRPO
WNPG
PPO
300
BGPG
200
5
10
15
20
25
30
35
Timesteps
x10
Figure 11: Episode rewards during training for MuJuCo Humanoid task, averaged across 10 runs with random initialization. The shaded area depicts the mcan the standard deviation.
Table 6: Average runtime (seconds per 105 timesteps) for the MuJuCo continuous control tasks
Eaviromment
WPO
SPO
TRPO
PPO
A2C
BGPG
WNPG
HallCheetah-v2
297
289 +: 25
290 28
292
293 +: 27
298 22
Hopper-v2
233 :
226 4: 42
212 56
167 -
254
197 = 31
Walker2d-v2
289 :: 55
312 1:
253 -: 39
307
259
322 4:
214 : 45
Ant-v2
307
290 +: 57
296 +: 63
251 47
291 41
286 4:
269 + 54
Humanoid
423 + 65
401 +:
446 +: 52
395 57
230 + 31
425 +:
398 + 49
20

---

Published in Transactions on Machine Learning Research (06/2023)
B Proof of Theorem i
Theorem 1. (Closed-form policy update) Let n2 (3,5) argmank=1.... I.... (sssss)) ax) BDkj}, where D
denoles the cost malrir. If Assumplion holds, then an optimal solution to (4) is:
tt' (a.ls) E (aa3))))
(5)
where (i,))) -f x3*,,, j) and (i,j) = 0 otherwise, and 8 is optimal Lagrangian multiplier corresponds to the following dual formulation:
min F(B) B20
min {35 + Esmp: Σ j=l T(azls) max 820
BDi))).
(6)
i=
Moreover, we have 8 < 8, where 3 == max.es b. (Day) (A" (s)))))))
(Dz,)
Proof of Theorem 1 First, we denote Qa as the joint distribution (-1s) and n' (-1s) with Σn= Qi ((,)s))
and 2j=1 0ig T' (a:ss) Also, let j)( represent the conditional distribution of n (aals) under ((a)ss).
Then 0i =(ay)s)ff(i,.)) ('(ass) j== Qi = =(j
In addition:
N
N
N N Dis Di and
dw (w(()(())),=((11)))
= min Qi
ΣΣ
Dig
min fa(i.)
i=1 j=1
E Eas (1s)AA
Σ
A ((s.)))((_,s))f...).)
j=1
Thus, the WPO problem in 4 can reformulated as:
max f.(J))20
E Egnip
Σ>
(14a)
i=1 j=1
s.t. E E
Σ
(14b)
Σ f
Vs S 1...N.
(14c)
Note here that (14b) equivalent to Eswpy ming,(.))) Ey (isT(alls)fA(i S 8 because we have a
feasible fi, to make (14b) hold, we must have Eg-pg min EN E (1 s 8.
Since both the objective function and the constraint are linear in fa(i.j). 14 is a convex optimization problem. Also, Slater's is condition holds for [14p the feasible region has an interior point, which is f. (i, i) Vi, and fi,j) 0 Vi / j. Meanwhile, since A A a) is bounded based Assumption [ the objective is bounded above. Therefore, strong duality holds for 14 At this point we can derive the dual problem of 14 as its equivalent reformulation:
min $20.33
38 +
Σ dss
BES
1=1
(15)
5 00,
s.t. A A3 - (ayss)) -DD))
Vs VsE,,i,) =1....... I...N.
(7(()
We observe that with a fixed 8, the optimal &j will be achioved at:
(3)
max p (s) (A) a.) -3D,)).
(16)
-=
21

---

Published in Transactions on Machine Learning Research (06/2023)
Denote 3 optimal solution to 11 and (i, as an optimal solution to
14 Due to the complimentary
slackness, the following equations hold:
59 (8") PJ(()
- 8r (1y=(()1)))
(1, = 0,
Vs,i,
In this case, (i, can have non-zero values only when ( (.) 3 (ayT((y)ss) 53 (8¹)
03(s) 0, which
means (3*) = p. (s) a.) - Dig). Given the expression of the optimal Cj in (16 (1
can have non-zero values only when i E K2 (0*, j). where kg argmaa A (8,
- 3Daj-
When there exists unique optimizer, i.e., (8* :j]] 1, let x. (B*, denote the optimizer. Since Ees fi(1,) = 1 as indicated in 14ch, the only optimal solution is:
if i i (3", otherwise.
f1(1.)) (
When there exists multiple optimizers, i.c., (B* 1, the optimal weights f1(1,) for i KC could be determined by solving the following linear programming:
max (
Ex-ps
T(alss)
Σ
(ss,aa)(((i))
j=1
(B
Egpp
T T(alss)
Σ
(17)
(B-.))))
Σ (E
f((ij) 1
Vs I.
And then the corresponding optimal solution is, n" (ails) = Σÿ ((agls){( (i.j)
Last, by substituting (B) = PT((s)(((a)l)) max (A ( a,) 3Dis) into the dual problem 15 we can
reformulate 15 into:
N
min B8
Σ (J)(ssss) = min Eanpy
Σ Taalss) max (A" - (Di))
(18)
#20
LES
j=1
#20
i=I...N
The optimal 3 can then be obtained by solving 18
We will further show that 3* < == max M.aX,ES.Aj=IL.... N.
In the general case, ie., B 2 0, 14a) non- nog-negative because:
AT (371))))) Eawpg Σ Σ
(19a)
N (s) Σ
x(azl>)
A" (s, (i.))
(19b)
f=1
i=1
N
Es-pp
x(azls)
Σ
(s,,a)
(19c)
jml
iEK (A+
N
> Ex-or
x(azls)
Σ
(x,,,)
+ + Dz)
(19d)
22

---

Published in Transactions on Machine Learning Research (06/2023)
N = B
N
7(ayls)4 + Eampt
Σ T(ylls) 1=1
Σ (2(1,))3 Dij
(19e)
j=l
N
= BB
Σ
x(ayls)8*
Σ
f1(,5)D
(19f)
(
2 0,
(19g)
where 19d holds since for i (5" - 2 - D3 A When 8 >
M.... A (=.0x)-A" (0 ) we have that for all S, K3 (3" (3' :) j. Thus, (i = Vi and
Dkj
(i, = 0, vi j. The objective value 14a will be 0 because Exnpy Σi=1 F
j= A (4
EgnpT Σi= (s, (()a)s)) = The left hand side of (14b) equals to Exupp Σ Ei=1 Ej=1 (ay
Eamp: Ei (a:ls) O. Thus, for any 8 >> 0, (14b) is always satisfied.
Since the objective of the primal Wasserstein trust-region constrained problem in 6 constantly evaluates to when 8r > max
A ("(s.at)-AA"aaa (((0)))))) and is non-negative when 8 S A ( as an upper bound
Daj
A (s.0r) (-.4)-A4'(0 (7(0))))))) }, we can use maX S, k.j=1... Net)
D2y for the optimal dual variable 3"
D2,
C Optimal Beta for Special Distance
Proposition 1. Let ki = arymar AT (8; 0z), we have:
(1). If the initial point 3o is in [ max,5 (AT (s. ax,))=)))) (s, a,)).
- (s,,, +oo), the local optimal B solution is
(2). If the initial point Bo in [0. min {AT ak,) as)s,a))) 8 px (s) x (((ax,ls)dde O, the local optimal B is mins.j4k, (A" (s, A a;);) otherwise, the local optimal B solution is O.
(3). If the initial point Bo is in (mins.jyR, {A (A5(s, (s,)),,)) AT (;)) max. (A" (* an.) - A (s construct sets 12 and 12 as:
for S E S.J. E (1,2...NN 30 2 (5 - (T aj) then Add j to I! clse Add j to P2. Then, if 8- E Enp Ejer (ayls) 0, the local optimal B is min (A (s, - A" A"(s, az)}; olherwise, the local optimal 8 is max, (s, - A (s))
1 (1). When 3 E [max.. (A ax,) - as(s, +oo), we have AT > A (s, a - 3 for all s s S, I. Since ( ax,) - B 2 AT ax) - B for all k == I...N, we have A (s. az) A" (1s,a))-88 3 for all s s 1,j=1.... N, == 1... N. Thus, j E argmax N( (A"(s, - -DD,,) for all s E S, ===.... I. Therefore,
Proof of Proposition
can be reformulated as:
min + EE EsNpT Σ 820
((ayls))5(ssa)))
j=1
Since d 2 O, we have the local optimal 3= = max,, , (A (3
(2). When B E [0, mik, {AT ag.) A A" 1)) we have A a,) < A (s, (s,)),) - for all E S,
1 N, k.. Thus h2 E argmaxx=1.... ( (A (8,ax) DDx} for all S, 1...N. N. The inner part of
23

---

Published in Transactions on Machine Learning Research (06/2023)
then is:
36 Boupet {
Σ (s),
-3) (so,, ax,)
j=1,fsk,,
B(6 Bgmps
Σ m(azls))
Egpup Σ
((aa)ss)A"(s.0.
j=L,i=h,
j=l
B(6 / pT(s)(1
(ak, (s))ds) + Bsmpp
Σ
{(a)ls)A" (s,
.EE
j=1
If 8 SAES Po (1) - x T(ax, 1s) < 0, we have the local optimal B mins, jth, (s7(s, - A (s, If Jaes (1) Taa, ss) 2 O, we have the local optimal B O.
(3). For an initial point 3o in (min {A" ak.) - A" az)} max { (A*( (8, az )- ) ( aj))), we construct partitions 1 and I of the set 21, in the way described in Proposition for all s S. Consider in the neighborhood of Su: i.c., 3 2 (s, ak.) (s, for s S. I and B S A5 as.) (s, ay) for S 13. Then the inner part of can be reformulated as:
B6 + E Eampy
m(ayls)A (s,
Σ (A -3} jE?
jE
= 3(6 - E Esnpy
Σ
π ((as (s) - E E
Σ jET!
saa (5) + Σ
T(agls)A"* ox.)}. SET3
JET
If 8 Benup Ejerz (ayls)) < O, we have the local optimal 3 = min
(A (A ax,) As,,,,))}, If 8 Ex-py Ejer? πI (agls) 2 0, we have the local optimal 8 max.es.jer) (s, ok.) (s)
0
Proof of Theorem
Theorem 2. If Assuupptio holds, then the optimal solution to
4 with Sinkhorn divergence is:
<(aasss) = Σ Taalls)f2
(8)
exp (s.0),
where D denotes the cost matria, (i,j)
and B% is an optimal solution the exp ( 3 1 (s,aan))
following dual formulation:
3 In Σ exp d ( (3), AD.))) B8 Expps Σ T(a, T(aris)( In(
min Fx
min #20
#20
N
N
I exp (3 a;) (Di)) x(azls)
Σ
Exset
Σn exp A (3( >Dks)
t=1
1=1
Moreover, have 3,, < m0
Proof of Theorem Invoking the definition of Sinkhorn divergence in 3 the trust region constrained problem with Sinkhorn divergence can be reformulated as:
Egspy Σ
N i=1 (s, E Qil
(20a)
DiD0 Di + E,
N
s.t. E Eappy
Qi log Q3,1 < <
(20b)
24

---

Published in Transactions on Machine Learning Research (06/2023)
Σn 1= Oi = =(asss), = ,j=1,...,N,sc.
(00c)
Let B and w represent the dual variables of constraints 20b and duality of (20) can be derived as:
20c
respectively, then the Lagrangian
N a4)
max min L(Q,B,w)
max min B
(
Qid
Q $20,0
Q 920,
j=1
N
N
1 0 log Qiy))
of π (ajls)))s + 3(0 - E Eampy
Di
(21a)
SES
3=1
i=1
j=1
N
N
N
N
= max min E Egwpy 320,0
A (8, az)
Σ
Qigl +
Q (s)p
(5())
j=1
SES
j=l
N
N
Σ
w (
+ 38
B.app
1 Qi + x
2i log Q1,))
(21b)
SES
max
min
38-
Σ
(afs(a
Q $20,
SES j=1
N
+ wj
8 Qi log Qi
+ E pu
- -Biis
(21c)
pT({)
A
i= i=
N
= min max 38-
Σ wfr(ayls)ds
$20,
Q
+ES
jml
w/ pT(s) x 2i log Qigh,
$
E
-A
Qii
(21d)
where 21d holds since the Lagrangian function L(Q, B,w) w) is concave in Q and linear in B and w, and we can exchange the max and the min following the Minimax theorem (Sion 1958).
Note that the inner max problem of (2d is an unconstrained concave problem, and we can obtain the optimal Q by taking the derivatives and setting them to O. That is,
aL as, - BDi + (T))) X (log Qi + = 0, = ,N,E, ES i0;
(22)
Therefore, we have the optimal Qij as:
A
Awf
Q =xp ( 1 (8, - >D)s) exp
1), = vi.j=1,...,N,a N,sES.
(23)
(p)(s)
In addition, since Σm Q7 = T(ayls), we have the following hold:
Awo
<(aiss)
exp
1)
(24)
Bp((s)
ΣN exp (2 (),T(s ADr;))
By substituting the left hand side of 24 into
23
we can further reformulate the optimal Q as:
exp A A (8, az) - ADiz)
2
T(azls), Vi,j = 1,,".,,,sE. NNsES.
(25)
ΣN exp A ( (8,
2Dz))
25

---

Published in Transactions on Machine Learning Research (06/2023)
To obtain the optimal dual variables, based on
24 we have the optimal w" as:
wi = P(() 3 B
B
In( s))
In
exp
a"ss ADz))ll, -j=1,-. N,SE (26)
By substituting 25 and
26
into
21d
we can obtain the optimal 8" via:
N
+ In(T(a)1s)) 8 In
B
min 820
35 Bs
p
(a
x
exp
B
A" a)) -DD;))l>
N
N
exp
(s, (s,a,) - >D1) <(a)ls)
+ E Eg
p
X Σ exp ( A5 ax) AD1,)
The proof for the upper bound of sinkhorn optimal 3 can be found in Appendix E
E Upper bound of Sinkhorn Optimal Beta
In this section, we will derive the upper bound of Sinkhorn optimal 3. First, for a given 3, the optimal
Q to the Lagrangian dual L(Q, 3,w) can be expressed in (25). With this, we will present the following two lemmas:
Lemma 1. The objective function 20a with respect to Qij (B) decreases as the dual variable 3 increases.
Lemma 2. If Assumption holds, then for every ā > 0, Qij (24-)) feasible to 200b for any X.
We provide proofs for Lemma 1 and Lemmna in Appendix bJ and Appendix respectively. Given the above two lemmas, we are able to prove the following proposition on the upper bound of Sinkhorn optimal 8: Proposition 2. If B8 is the optimal dual solution to the Sinkhorn dual formulation then B3 S 2A for any X.
Proof of Proposition We will prove it by contradiction. According to Lemma
Qi3 24mm feasible according to Lemma
20b Since 3 is the optimal dual solution, 2 ($5) is optimal to (20). If 53
(5 the objective value in (20g) with respect mux smaller than the objective value (20a) with respect
8y which contradicts the fact that Q (85) is the optimal solution to (20
E.1 Proof of Lemma
Lemma 1. The objective function 200 with respect Qif (3) decreases as the dual variable 3 increases.
Proof of Lemma 7 Let Gx(3) represent the objective function 20ah. By substituting the optimal Q in
25
into
20a
we have:
cxp ( AT (s) - ADzj)
GX(B) Rgnop
AT 04)
[[all8)]
(27a)
k= exp ( AF (s, ADxz)
N
N
exp A (8, - >Dis)
Egpp po
T(alls)
AT ( ai)
(27b)
Σ exp ( A3 ADvz)
For any 82 > 31 > 0, we have:
G2(B1) G2(32)
exp 1 A (s, a;) AD:;)
EBNpp Σ T(a (s) Σ
A ( a
ΣN exp 2/7 A (s ax) >Dki)
1=1
26

---

Published in Transactions on Machine Learning Research (06/2023)
exp A A a,) - -DDs))
B
(28a)
Σ
A= exp 2 Ar (8, AT ax)) - -2Dss)
N
N
exp A A A5 aqa) (D
= R
Σ x(agl))
(8)
Σ
A (T(s,
8,
2k=1 exp (p) ( an)) - ADm1)))
A
j=1
i=1
exp
( 32
13 (s, qn) - >D(33))
(28b)
> exp 1/t Ar and) \Dbs))
where [i] denotes sorted indices that satisfy (8, anj) 2 (p)) 2... (T (s))) Let
f.(1) = exp ( (s, - \D()))
exp A AT (s, (D((1))
(29a)
Σ exp 3 AT (14) - 71Hs5
Σ exp ( AT (14))) - 1p1)s
exp (1 8. b. ) (1) exp 2 (T)) - (D})))
Σ 1= exp (()) 2/ A (s, an)) exp (3) (T axt) pDp))
exp 3, A (s, 9z) - AD
(29b)
Σ exp
A A (T a(1)) 1D1)
B2
For notation brevity, we let (i) = exp (0A1 - 32
A ) (s, ap)) > 0, wx(1) = -exp A A (7( - AD(A)) > 0
and (1)
Then we have
ma, (R)
Σ.m m. (.), (1)
ma ma(i)ux(()
w.(() (i)
296
(30a)
Em (k)
Ex= (1))
1
(.)()a(()) (
(30b)
Ex m, (k)))s())
Σx my (i): (bc)
(1)
(30c)
Since A > > 0, (i) decreases as increases. Thus, (i) decreases as i increases. Since m, m, 2 m, (k)) and A
(N) S (k) for all k = I,...,,, N, we have (=(1)= = En m. (k)w,, (R) Ex (1)wss(( 2 Ex (k)es
Ex L my (k) = 0, and go (N) = Eim, (M)m)
Ee ma(N)oo Ex (A) Σe_ (M( O. Since (1)= 2 0, (N) SO and (i) decreases as i increases, there exists an index I S k; S N such
that (i) 2 0 for S k. and q)(i) 0 for i > ka. Since mi((), (i) > O, we have fa(1) 2 0 for < ke and fi) < for i> kg- In addition, we have ΣNV fa( directly follows from the definition. Thus, ΣN (i) = Σ [sa(1)) Ema,+t (sa(1) = 0. Therefore,
Σ N
N A (s. (p)
GX(31) G2 = E Esups
x(ayls) Σ
(31a)
5=1
N
k
N
= E Exppt
T(alls)
ATH (s, ( (g)((1((1) - Σ
A assaaa)))f
(31b)
1=1
i=
N
k.
N
2 E Egmp5 p5
Σ
T(arls)
Ath (8, (R. ())((()))) - Σ
A (5, (1
(31c)
j=1
N
4e Σ 15a(0) - (s
N
Egppp
Σ
(A
Σ
15001)
(31d)
j=1
1=1
i=k++1
-E-pp Σ
(A (p.))
Σ
[fa(1) A (1))))))))
Σ
(f(()))
(31e)
27

---

Published in Transactions on Machine Learning Research (06/2023)
N -
= Es
(s a1,+1))) lfa(i)
(31f)
2 O.
(31g)
where (3ch and 31g hold since A (s, azd) is non-increasing as increases. Furthermore, at least one inequality
of (31c) and 31g will not hold at equality since Σi= (r(als))A a,) O, Vs S, and for non-trivial
cases, Pr Pr(A A 5, a 0, E S, Ya A) 1, which means Pr(3s1, S,a,, ag A, s.t. (s))), + A (52, 02) } 0. Therefore, we have Gx(31) Gx(32) 0.
E.2 Proof of Lemma
Lemma 2. If Assumplion holds, then for every 8 > 0, Qij (2)) is feasible (00b) for any X.
Proof of Lemmal of 20b) follows:
2
By substituting the optimal Qij in 25 into 20bb, we can reformulate the left hand side
N
N + Qij log Qijl
E Egpo ΣΣ
Qi
(32a)
i=1 i=jj=1
Eggup:
N
Qij + x Qz3 X A
(T A
>Dz + log
x(a,ls)
11
(32b)
Σ Ek=l exp (g (s, ADkj)
j=1
N
<(azls)
Egppp
2ij AT (s)
Qij log
(32c)
Σ k= exp (g) A (8, ax)) ADki)
Now
we
prove
that when B = 2A
E Egppy
EE N Σn (()(
2
and
Exnpy
(1 (B) log
x(a,10)
s hold. For the first part, we have:
Am (x
E
Ex
A A a))
(33a)
B
N
0 (s, (s,
(33b)
B
j=1
1
B
T' A'"a
(33c)
3
<
EgNp5
T' (s,
(33d)
B
A mar
(33e)
8
For the second part, the followings hold:
N
N
x(ayls)
Banpy
E
Q33 log
(34a)
Σ exp (T, (s, - >Dxs)
N
((ayls)
EENpp
255 log
(34b)
Σk=1 exp {/4 A (s, (x) \Dkj)
1=1
1=1
1
π T(a 3)
Bawpz Σ
[(alls) log
(34c)
k= exp ( A- ax) \Dxs)
28

---

Published in Transactions on Machine Learning Research (06/2023)
N
log T(a,ls)
T(asle [s)
(34d)
X
exp ( A" (x,
==1
N
a(aa]s) log
(34c)
X
exp ( (s),))))
==1
1
T(asls){ A (T(s)
Bgpp
B
(34f)
X
N
E
r(a5ls)[4
(34g)
B
P
==1
A mec
8
(34h)
3
2
Therefore, Qij (2 (24mm) is feasible to (20b) for any X.
F Gradient of the Objective in the Sinkhorn Dual Formulation
The closed-form gradient of the objective in the Sinkhorn dual formulation as follows:
6 Ex
Σ x(aa)1))
(n(T((y)s))
Σ
exp
A A (s. - -)).))))D)))
3
jm1
B
exp /タ A" (s) - -WDs) x -AA" (s,
DN exp (3)7(s AD;))
N
T(ag/a) exp { A7 (s, ADis)
+ Ea
ΣN exp ( ( as) ADas)
+ (ajls))) exp (s, a:) >D1)) x -1A (s. a=) Ex exp (s, ADwz)
(E exp ( A0 (8, ax) ADk)))² ))
(T(a))s) exp ( - \D2)) x Σ exp ( (s, ax) ADks) -AA
4= exp A (5 aD2)))
G Proof of Theorem
Given the upper bound of Wasscrtein optimal / in Theorem and the upper bound of Sinkhorn optimal in Proposition 2 able to derive the following theorem:
Theorem 3. Define Bun = max ( 22 B). We have:
(8) converges to F(B) uniformly on [0, BuB]:
lim
sup
Fx
lim 200
Bua NnN 0.
700 oSR
lim argming (x) C argminggasse F(8). >
Proof of Theorem
To show that (3) converges to F(8) uniformly on [0, BuB]. it is equivalent to show that
lim suPoS6S8cs Fx F(3) ((3))-F(3) = Let (B, = (A 0x) BDx;) [A* a;) - 3D
and eg (B, 2 First, we have
[3) - F(3) F(3))
29

---

Published in Transactions on Machine Learning Research (06/2023)
= |30 =|36- E Esupp
Σ T(a 18){
N
- A + x In(x(a,1))) - _ Inl
x
I exp 3 (s, - -)D
N
A
1=1
i=1
N N
3 cxp (g A (s, a.) (5,
AD2)) x(azls) B8
+ Rapp
A ki=l exp / Ar (8, xD2))
;=
- Eg
Σ T(alss)
max (AT ai) BDw) i=L...
(35a)
8 E Egspp
+12/ Egmpy
3 +
N
(aals))
Σ
(((a))))l((((((ss)
1
j=1
N
N
B exp (g) ( a:) >Dij)) T(a)ls)
+
E Egupy
A
Σ exp (g) A" ax) - >Daj)
1=1
8 In Σ exp 8 1" (s. a.) AD:))]
+
(aglls)
x
j=1
Es Px p5
Σ
(ayls)
max (A (A" (s) ai) BDis)
(35b)
i=I...
N
Σ
3 [(alls) +
Σ T(cyls))[((((als))
S 2
p
j=1
5=1
N
+
p5
((azls)
x
In|
Σ
exp
A - -Dii)l
j=1
1=1
Ex-ox 1=1
T(alls) max (A, a:) BDa) FI....
(35c)
In addition,
Eanpz Σ
E
T
((yls))
3 In
Σ exp A A
(s, ai) (s)
ADDa)]
F=1
BeNDp
Σ (agls) max (s)) - BDib)l
(36a)
=
/=1
Erm-p p3
Σ
N
3 π (ayls)) Infexp ~ max (AT(s, - -8Dks))
Σ 1=1
N
exp (- (-
x
k=I..N
1=1
N
E Egppp
T(azls) max (s) - -(B.s))))BDB))))))) i=I...
(36b)
1=1
3 [Inlexp(
(- X
Esmpz pz
T (agls)) X
/ max (AT BDDx 11 In Σ exp (
B (B,
k=1..N
E Egnp
(aalss) max (1' (s, - -3Dss)
(36c)
i=...."
1=1
N
N
exp - >
Egmpx ph
π (ayl8)
In
B (B,
(36d)
1=1
30

---

Published in Transactions on Machine Learning Research (06/2023)
Therefore,
lim
sup 0SBSPuB
Fx(B) - F(3)
(37a)
23UB
N
BUB
N
< lim t
Esmpp Σ T(azlo)
+ lim
Bampy
Σ f=1
(n()())l((T((s1)))
X
x000
f=1
X
+ lim
sup
E B
x (a,l)))
Σ exp (1 (3,
(37b)
X OSASBUB
j=1
lim 2 ogg
sup
Sxpp
T(azls) In
2 exp (- ((,
(37c)
In addition, VB E [0, BUB] and AY, j)) is bounded since
(B,,
max (A" a&) 3Dks) - (s, 3Dbl k=
(38)
S max A" (s, a) + BUB max Dij i.j
8,0
< 2A max
BUB max Dij 1.j
(39)
Then, Esnps Σj= T(a,1s) In Di exp (- J (B, bounded. Therefore in (37 the optimal / can be
achieved. Let 3 argmax A Es-py Σe= T(u, (8) In In exp (- (B, and then we have:
lim t00
sup 05B$BUB
E
S T(ayls)
In Σ exp 3 (3,
(40a)
N
lim
E Esups
T(a
exp
(40b)
200
i, Then since c0 (B, > 0 for $ KT (3,j) based on its definition, we have o,( O. On one hand, we have
Let 10g (3.j) argmax k- I...N A A5,s, ax) 3Dz-- Define oz(i) minogss I... 5 (B,ú,,
N
lim In 700
exp
Bx
Bx " (1,)),
(41a)
N
N Σ
exp (- A
lim In -oc
Σ
exp
BI (B (i,1)) +
3x (BA (i,)))
,BA,,,))))
(41b)
im ligg (8xa)
im lkkee (9xJ)
X
< lim In Aoo
Σ
exp
(i) +
exp (0)
(41c)
BUB
k (8x.))
i (8x.))
X
lim In 200
Σ
exp
0z (i) (Cx,(B)))
(41d)
BUB
lim In[ (8x, J)]-
(41c)
On the other hand, we have
N
lim In exp *
Σ
BA C00 ((5))
(42a)
1=1
31

---

Published in Transactions on Machine Learning Research (06/2023)
N
N
= lim In toc
Σ
exp
BA S C. (,5)
+
Σ
exp
33 35
(3*,,,, (2,5))
(42b)
i= (8x.))
i= (8xJ)
N
2
lim In
Σ
exp (-
BA (A
(42c)
A +00
== (3x.))
lim In[
Σ
(x)
(42d)
(1
= lim X
In[1K] (3x,
(42e)
Therefore, limx
In ii=
In
exp
Ba
(BA, ((,5))
-iim ~ In|1K (Bx,
Based on that, we have
gr
N
N
lim t0c
Bgmps
(ayls))In
In Σ exp (-
B 533 (((3)
(43a)
X
j=1
1=1
B*
lim
exp
3a (
(43b)
A 700
j=l
B^
exp (- A
< lim
In
Ba C. (.,))
(43c)
B^
= lim
(43d)
x000
j=1
BUB Nmn 0.
s lim
(43c)
which means limx-+> suPosagōm Fx - F(3) $ 0. Furthermore, since limx supos [Fx F(B) holds naturally, we have lim suPoçaSAup [F) - F(B) = 0. Therefore, (x) converges to F(B) uniformly
[0, Bum], which also indicates Fx epi-converges to F(B) on
[0, BuB] (Royset 2018 Rockafellar & Wets 1998). By properties of epi-convergence, we have that lim argminosaSfup (3) S argmingsassus F F(B) (Rockafellar & Wets 1998).
H Proof of Lemma 1
Lemma 1. As Au + 0o, SPO updale converges to WPO updale: limA, t00 R800 (TT)) E FF @FWPO (Tx).
Proof of Lemma 7 Let (i,5) update with X equals to:
Dz (max - Bin - [A** D.yi The SPO
exp
(s, ADis))
(ati(a lim (TR) lim
A (s, - ADis)) Th (a)lss)
(44a)
200
1==1
exp
3r
N
exp a3t4 43Th (B05)) Bx D
R, (Phn)))) exp (i,
T (ayls)
(44b)
exp (A (s,
BrD (Bn Σn (1.j))
exp (- (i.
exp( (l.(5 (ays)s))
(44c)
700
exp (i.)) (a
lim
(44d)
700
ΣieRT (8r.)
expl- (1. + ΣişkeT (Bu.)) exp(-xt (1,)))
32

---

Published in Transactions on Machine Learning Research (06/2023)
N
LIERIH (Bx-) lima limx xx(ay(s)
(44e)
Zigke (AxJ) (BaJ)) lima-x exp( (1.)))
Tes* (Bn (i) IRT 5)l
Th (T(a)))),
(44f)
where / denotes the indicator function; 44h holds because X - oo, (i,)) oo for $ RT (3x, and 0 otherwise, thus limx->oc exp (- = 0 for i o Rg (6 and 1 otherwise.
1 if K (Bk,)), and (i. = 0 otherwise. Therefore, SPO update with À - 00 Let (i,)) = ( equals to the following WPO update, (WPO(TT) Σド= ( (i.))
Proof of Theorem
Theorem 4. (Performance improvement) For any initial state distribution v and any Ba 2 0, if || l c for some e> 0, let (3x. = argmaz;.1.... (A== (A.*(s,a) BxD WPO policy update (and SPO with À 00) guarantee the following performance improvement. bound when the inaccurate advantage function A is used,
J(T441) 2 (TTK) + Ba
(a(a))))
Σ
2c
(10)
snipa
1++1
-1-7
f=1
"ERT (8..5)
Proof of Theorem
4
J(Tx+1) - (TTK) = E
T++1 on
BeN AA (s,aa
(45a)
= E
(45b)
srupu
iml
E
scupe
k+1
T (1.
(45c)
= E
(a(a)s)) Σ (s.
(45d)
snupy
1=1
N
= E
erupo
1+1
Th Tk(ay)s)
Σ
(s..))
(45c)
j=l
DERT*(BE.d)
> E
TR(alls)
Σ
(h + BeD - 2c]
(45f)
sruph
iER (Bk.3)
N
= B&E
Σ
(asss)
Σ
(1,5) 1 1-7"
(45g)
snup
j=1
(
where 45a
holds due to the performance difference lemma in Kakade & Langford (002) 45f
follows from
the definition of iT (Bx-j) and the fact that -| A A Ilxe I c, therefore for E K (B..))
[ th (s, a.) + holds since
c) - Di 2 (s, BiDi, 2 A (s, a,) Bx A (s, 2 (s, - G 45g F E A (s, 0.
33

---

Published in Transactions on Machine Learning Research (06/2023)
J Proof of Theorem 5
Theorem 5. (Global convergence) Under Assumption 2 we have for any Bx 2 0,
WPO
satisfies that
I|V - V 110 S <>Ilv* v 1120 + B110
(11)
and SPO
satisfies that
Bx. -- (IDII1o + log N
I 1 30 S < -V 11:0 +
(12)
If limg- 3x = O, we further have lim (TTK) = J
Proof of Theorem Our proof is inspired by the work Bhandari & Russo (2021).
We use the shorthand т, for the probability distribution C(. 8) on the actions and denole the probability distribution on the action space A as A. To save notations, we rewrite Tk+1; Wk and f as xt, x and B
respectively. We use d for either dw ds in the following derivation. Note d D| D for both case and ds -2 log
Since a policy ı is indeed just member of I A, we find that the problem 7 optimization problems. For each S, we need to solve
can be split into |S many
p" ((()))nsr(()))) TEA
max
3p" (s) T.).
(46)
Denote the quality function of (=(ssa) E[Ri s,a, a; T]. and the value function of (*) EiR,|s, s; r) we find that A (s. (s) - (s).). Since the second term is only function
of the current policy ı and the state we find that Problem 46 is further equivalent to (in the sense of the same solution set):
max EA
Eaaws (s, - 3d(m',, m.).
(47)
Here we use Pu(s) > 0 for all s. Let स be a solution the policy iteration:
Fg E arg Eawx; (Q" a)].
(48)
Also define the bellman operator T: T
R S and the operator T RRS R R| for any V RIS!
(TV), 8 max ( (s. a) + PBs (V)))) (V(')),
(49)
(T V). = Eawn; (r(s + + =... V(s)).
(50)
Using the relation between the quality function and the value function, Q (s, a) =r(s +(s (V (s)) we can rewrite the above equations in terms of the quality function for V = V=:
(TV*), max r(s, a) BBPPP(( (V [s)) (V aEA
max Q(s,a) -T*V= aeA
(51)
(T V*), Enna; (8,(
(52)
Forr Sinkhorn divergence, note that the entropy function always nonnegative.
8This
lower
bound
is obtained via ds (s' x|2)
2 min min Σ., -z D) D) [M())}
(Q, D)
(6) 2 N Here in the step (a), the Lagrangian mul-
((())
O. APg) - log exp(-
tiplier method to derive the optinal Qı = En, exp(-ADis)
exp(-AD;
In the step (b), we use the fact that log Er exp(x.)) &
max(211...,, Fn logn any 2i.-.. Fn E and D. any
34

---

Published in Transactions on Machine Learning Research (06/2023)
Let us consider d = dw first. Using the optimality of x for the problem 46). we know that
B aen [Q" (s. - =d(= To) 2 EaER, E a)] o,
(53)
B aen (8"( 2 Baes, (x, 3D.
and
B a [Q (s. a)] - _d(= TT)) 2 E (Q" Bd Te)
(54)
R Ren [Q" (8, a)] 2 Baes, (8=(s, a)] (8).
Using the notation in 51 and
52
53 and
become
7
TV BD1
(55)
T V 2 V*.
(56)
Here [| is vector of all one entries and the incquality 2 mcans entrywisely larger than equal to. By iteratively applying 7n* to (66 and use the fact that 7 is a monotone and contraction map with V=* as the unique fixed point, we have
Vo 2 2... .....(T (T y2v
T V 2V
(57)
Hence we have
0 (0) V V yxt (b)
(c) S V TV + 1BD (8)
v -
(58)
Here the inequality (a) is due to the optimality of V*. The inequality (b) due to 57 and the inequality (c) is due to (55). Now using the fact V* is the unique fixed point of T, and T is a monotone and contraction map, we have from 588 that
IV
V" 10 S TV IV + BD S +/ V* loc + 3D.
(59)
Next consider d = dg. The optimality of m reveals that for के or π.:
| E B (Q" I a)] - 3d(=J To) 2 EaER, (Q" a)] - Bd ,, T
Ban (s. > Eaex, (x) 3DD + log N
(60)
X
Thus we have the following
rx ya 2 TV - -B(D+ 2 2og N
)1ss
(61)
X
T v > V - 3(D+ + 2
X 1[s)-
(62)
By iteratively applying * to (62) and use the fact that - is monotone and contraction map with V the unique fixed point, have
8 -D +
log N
V+
2 V
1ss
(63)
1-7
Hence we have
(b)
8
log N
0 < V*
Val
ZV- - TT V2 +
(D + -2-
[s]
1 1-7
(64)
(c) <V* TV +2
8
(D 2 log N
Here the incquality (a) due to the optimality of VA The incquality (b) due to 63 and the incquality
(c) due 61 A similar derivation as 59 shows the inequality the theorem. Hence the theorem is established.
35

---

Published in Transactions on Machine Learning Research (06/2023)
K
Computational Complexity of the Algorithm 1
Our overall algorithm applies general actor-critic framework: the actor follows the proposed WPO or SPO update while the critic follows TD methods. The computational complexity depends on (i) the per-iteration computation cost of the policy and critic update and (ii) the iteration complexity of the actor-critic method. Here we mainly discuss the per-iteration computation cost of the policy update, as studies on the iteration complexity of critic framework for constrained policy optimization are limited.
The computation cost of WPO and SPO updates at each iteration depends on the selection of Sx. If Be is chosen time dependently, the compulation cost of WPO/SPO policy update is O(nān,), where and n, are the number of actions and states to perform policy update. If we set By as the dual optimizer, there will be additional cost to run gradient descent to solve the one-dimensional dual formulation. As discussed in our experiments, we can set 9x to be the dual oplimizer only the first a few iterations and use a decaying afterward. Therefore, the average computational complexity of a policy update step can be Onn2n))
L Difference between
and Other Exponential Style Updates
Sinkhorn divergence smooths the original Wasserstein by adding an entropy term, which causes the SPO update to contain exponential components similar to standard exponential style updates such as NPG (Kakade) 2001: Peng et al. ) Thus, SPO can be viewed as a smoother version of WPO update. Nonetheless, it's important to note that SPO/ PPO/POO updates differ fundamentally from standard exponential style updates that are based solely on entropy or KL divergence, In both SPO and WPO, the probability mass at action is redistributed to neighboring actions with high value (i.c., those a' with high A (" a') Bd(aa,, (a(aa"a)) a)). In cantrast, in these standard exponential style updates, probability mass at action a is reweighted according to its exponential advantage or Q value.
M Exploration Properties of WPO/SPO
Compared to the Wasserstein metric, the KL divergence between policies is often larger, especially when considering the policy shifts of elosely related actions, shown in Figure 2 In practice, when employing the same trust region size 8, Wasserstein metric allows for more admissible policies within the trust region compared to KL, thereby leading to better exploration. This advantage is demonstrated in our motivating example in Figure 3
Furthermore, Sinkhorn divergence has even more exploration advantages than using Wasserstein. As Sinkhorn smooths the original Wasserstein with an entropy term, it includes additional smoother (more uniform) policies in the trust region, leading even faster exploration.
Our numerical results in Section also support that WPO/SPO explores better than KL; and SPO achieves faster exploration than WPO.
N Policy Parametrization, Prior Work on Nonparametric Policy
As noted in (Tssseer et al. 2019); the suboptimality of policy gradient not due to parametrization (c.g., neural work) but is a result of the parametrie distribution assumption imposed on policy, which constrains policies to a predefined set. In our work, we strive to avoid suboptimality by circunventing the parametric distribution assumption imposed on policy, while still allowing for parametrization of policy in our empirical studies.
Previous research, such as (Abdolmaleki et al. 2018 Peng et al. 2019), has investigated theoretical policy update rules basod on KL divergence without making explicit parametric assumptions about the policy being used. However, to our best knowledge, no prior work has explored theoretical policy update rules based on Wasserstein metric Sinkhorn divergence.
36

---

/ Published in Transactions on Machine Learning Research (06/2023)
O T-tests to Compare the Performance of WPO, SPO with BGPG and WNPG
We conduct independent two-sample one-tailed t-tests (Student 1908) to compare the mean performance of our proposed methods (WPO and SPO) with two other Wasserslein-based policy optimization approaches:
BGPG (Pacchiano et al. 2020) and WNPG (Moskovitz et al.
2021 Specifically, we formulate four alternative
hypotheses for each task: JwPO > JBGPG, JWPo > JWNPG: JSPO > JBBPGG, and Jspo > JWNPG.
MuJuCo continuous control tasks are considered for the t-tests, with sample size of for each algorithm. All t-tests conducted at confidence level of 90%. The results of the t-tests are presented in Table 7 where checkmark (v) indicates that the alternative hypothesis is supported with 90% confidence, and dash indicates failure to support the alternative hypothesis.
Based on the results presented in Table we can conchude the following:
The mean performance WPO is higher than BGPG with 90% confidence for all tasks.
The mean performance of WPO is higher than WNPG with 90% confidence for all tasks.
The mean performance of SPO is higher than BGPG with 90% confidence for all tasks except Ant-v2.
The mean performance of SPO is higher than WNPG with 90% confidence for all tasks except.
HalfCheetah-v2.
We note that though SPO's performance not statistically significantly higher than BGPG or WNPG in Ant-v2 and HalfCheetah-v2 tasks, SPO demonstrates faster convergence speed than WNPG and BGPG in these two tasks.
Table 7: T-tests results on the performance of WPO, SPO, BGPG and WNPG
Enviromment
JWPO > JBGPG
JWPO > JWNPG JsPO > JBCPG JsPO > JWNPG
HallCheetah-v2 V
Hopper-v2 v
Walker2d-v2 /
/
Ant-v2
/ V
Humanoid-v2
37


Reviews:
[{"review_id": "TGElfHuJsp", "reviewer": "Reviewer_xF5g", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Strengths*\n- Pursuing better methods for making use of the Wasserstein-based metrics is definitely an important direction.\n- The derivation of the policy update rules and the theoretical results on monotone improvement and convergence appear to be new, at least for this specific problem.\n- The experimental results are fairly convincing even though the improvement is marginal, but there are a few missing results that might make the case stronger.\n\n*Weaknesses*\nWhile I generally feel positive about the paper, there are some weaknesses that I hope can be fixed.\n- The related work indeed cites a lot of the relevant literature, but falls short in properly contrasting the contributions and significance of this paper in light of that literature.\n- The motivating example is a little shaky because it requires one to know in advance how the actions are related by “distance.”\n- A major claim of the paper is that the update rules do not explicitly require a certain parametric class which can be restrictive (see prior work). But in practice it seems that some either implicit or explicit parameterization of the policy is used in all but the tabular settings. Theoretical update rules for KL-based methods are also do not require explicit parameterization, if I recall correctly.\n- The theoretical results are not particularly concerned with exploration and generally just assume the problem away with sufficient coverage. While there is certainly more to be explored here, it may be alright considering this is on par with similar papers that study this sort of thing.\n- The experimental performance improvement (although definitely there) seems marginal in practice.\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "First, I apologize to the authors for the tardiness of this review. I hope the feedback can still be useful.\n\nThis paper studies the problem of policy optimization for reinforcement learning via constraining the policy using metrics beyond the KL divergence. The paper contributes to a line of work that explores Wasserstein-based metrics for the trust-region part of policy optimization. In contrast to other works that have investigated the Wasserstein-based metrics, this work is concerned with providing theoretical guarantees and avoiding potentially suboptimal approximations or use of penalty functions. The variants of the algorithm are given, one based on the Wasserstein metric directly and one using a Sinkhorn-like version. Both involve solving their respective dual problem. A monotone improvement guarantee is established and a convergence theorem is also proven in the tabular setting.\n\nThe method is evaluated on a number of classic reinforcement learning benchmarks, covering tabular settings and continuous control settings. The proposed methods appear to outperform the baselines, although marginally.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have some questions and suggestions.\n\n- I’m confused about how the cost matrix M is defined. In the motivating example, the distances are simply “assigned” by the authors. Fair enough, if you have knowledge of the problem setting in this way. However, later it goes on to say that $M_{ij} = d(a_i, a_j) = ||a_i - a_j||_1$. If these are discrete abstract actions, this would be undefined, such as in some of the experiments. How exactly do you define subtraction between two such abstract actions? In general, what would one do in the case where nothing is known about the discrete action space? How does this choice affect performance?\n- How does the SPO update rule compare to NPG/SAC or other sort of “exponential weights” style algorithms for policy optimization?\nIt would be helpful to have more extended discussion about the policy parameterization as mentioned in the weaknesses section.\n- As mentioned, the related work could be revised to include better distinction between this work and prior works rather than just discussion what prior works did.\n- Potentially several more difficult control settings (e.g. Humanoid) or maybe difficult exploration problems would make the experimental results more convincing.\n- Theorem 3: it would be helpful to write directly the inequality of the bound.\n- Page 3: “Henceforth” used incorrectly. Consider: “Until now” or “Heretofore”\n- Page 11: “coherent” used awkwardly. Consider: “consistent”", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Ap2jwYZvf4", "reviewer": "Reviewer_br9i", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n1. Using Wasserstein and Sinkhorn in trust region policy optimization seems reasonable (from the motivating example) and new.\n\n2. The proposed ideas are guaranteed by theory and practical algorithms.\n\n3. The experimental results verify the proposed methods.\n\n**Weaknesses**\n\n1. From the closed form update of WPO (Theorem 1), it seems it is more complicated than KL divergence. This might suggest the method is computationally inefficient comparing to using KL divergence, given sometimes the proposed methods have comparable or worse performance (e.g., in Figure 6).\n\n2. From the closed form update of SPO (Theorem 2), it looks very similar to standard exponentiate gradient update from using entropy or KL divergence. This might suggest the proposed SPO eventually is similar to existing methods, which could possibly reduce the novelty.\n\n3. Other closely related baselines such as SAC should be included, since SAC also use entropy / KL divergence.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work proposes to replace Kullback-Leibler divergence in trust-region methods with Wasserstein and Sinkhorn trust regions, namely Wasserstein policy optimization (WPO) and Sinkhorn policy optimization (SPO). Algorithm 1 presents the updates for WPO/SPO.\n\nThe authors firstly use an example of grid world to intuitively compare Wasserstein and KL divergences. Figure 2 shows that Wasserstein can characterize distances between difference probability distributions over action spaces, while KL cannot. Figure 3 shows that Wasserstein-based policy update finds the optimal policy faster than KL-based policy update.\n\nIn Sections 3 and 4, the authors develop closed-form expressions of the policy updates for both WPO and SPO, using optimal Lagrangian multipliers of the trust region constraints. Results are given in Theorems 1 and 2. \n\nSection 5 provides theoretical analysis. Theorem 4 shows that WPO has monotonic performance improvement, the same results apply for SPO with multipliers approaching infinity $\\lambda \\to \\infty$. Theorem 5 shows that WPO has global convergence with decaying regularizer coefficient $\\beta_k \\to 0$. \n\nThe authors then conduct experiments to compare the proposed WPO/SPO with TRPO, PPO, A2C, and BGPG, WNPG, using several tasks, including tabular domains with discrete state and action (Taxi, Chain, and Cliff Walking), locomotion tasks with continuous state and discrete action (CartPole, Acrobot), and continuous control tasks with continuous action (HalfCheetah, Hopper, Walker, and Ant). The authors also did ablation study to examine the hyperparameter sensitivity of the proposed methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1.  I am wondering how computationally efficient the proposed method is comparing to using KL divergence. Maybe this could be shown by comparing the performance of different methods using the number of samples or actual running time.\n\n2. In the motivating example, why is it the case that $1 \\to 2$ in Figure 2a is a close action and $1 \\to 3$ in Figure 2b is a far action?\n\n3. Please include some closely related baselines such as SAC in the experiments.\n\n4. It is mentioned that $\\beta_k = 1/k^2$ is used in Section 7.1. Why is this specific choice and how does it compare to other possibility (e.g., $1/\\sqrt{k}$, $1/k$, etc)?", "broader_impact_concerns": "The experiments are conducted using publicly available benchmarks such as Mujoco, and I do no see ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "8vcBxOInRz", "reviewer": "Reviewer_dwsm", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "### Strengths\n\nThe paper provides a relatively comprehensive account of the trust region optimization algorithm using Wasserstein distance or Sinkorn divergence as the local metrics. The paper details the properties of the policy updates (e.g., convergence of Sinkorn updates to Wasserstein distance updates), local improvement properties and global convergence properties of the algorithm. \n\n### Weaknesses\n\nMaybe a major weakness of the paper is its lack of novelty -- the idea of using Wasserstein distance or Sinkorn divergence as the trust region metric is not new in the RL literature (see the multiple references in the paper too). The experimental results in the paper are not convincing enough to show that such a new metric leads to significantly better results than baseline KL divergence. It is also not completely clear that the theory brings much novelty and insights to the established literature.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes to use Wasserstein distance or Sinkorn divergence as the trust region distance in policy optimization, compared to the KL divergence in the original TRPO algorithm. The paper studies the theoretical properties of such an algorithm, and provides some empirical evidence that this new trust region performs on par or better than previous baselines.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have multiple questions and comments below. The authors may consider making changes to the paper based on these questions and comments.\n\n### Novelty of the method\n\nIt is not quite clear how much novelty this current paper brings compared to a few papers already in the literature [1,2] (both are referenced in the paper). Using metrics alternative to KL divergence is a very natural idea, and introducing distance metrics based trust region is a reasonable next step, however, this seems to have already been accomplished by prior work. If the aim here is to establish better theoretical properties / empirical performance, the results feel to me not strong enough, as I will detail below.\n\n[1] Pacchiano et al, Learning to score behaviors for guided policy optimization, 2020\n[2] Moskovitz et al, Efficient Wasserstein natural gradientsfor reinforcement learning, 2021\n\n### Theoretical result: Thm 4\n\nThm 4 is a local improvement property popularized in trust region policy optimization literature. I think a major weakness here is that obtaining such a local improvement property is not very challenging since the improvement bound always contain something like $-2\\epsilon/(1-\\gamma)$ as shown in Eqn 10 of the paper. Such a penalty term stems from the fact that we cannot control higher order approximation error from the improvement step. It is not clear what is the benefit of such a local improvement bound compared to a KL based improvement bound. In its current form, the result does not bring much technical insight compared to existing results.\n\n### Theoretical result: Thm 5\n\nThe global convergence property in Thm 5 is not a very strong result unfortunately. Essentially what it says is as $\\lambda\\rightarrow\\infty$ we have $\\gamma$-contraction of the value function $V^{\\pi_k}$ to the optimal value function $V^\\ast$. This is quite obvious since if $\\lambda\\rightarrow \\infty$, we have $\\pi_{k+1}=\\arg\\max A^{\\pi_k}$. In other words, the algorithm reduces to policy iteration, which indeed achieves a $\\gamma$-contraction. Having a more general result in Eqn 11, which considers the case when $\\lambda$ and $\\beta_k$ are finite, does not bring much more compared to Thm 4.\n\n### Choice of distance metric\n\nI think a major challenge in adopting Wasserstein distance based trust region update is that it is not clear apriori what is a good choice of the action metric $d(a_i,a_j)$. In this paper this is not discussed extensively (only that the paper suggests to implement L-1 distance throughout). Obviously, the choice of such a metric would be critical to the practical performance of the RL algorithm.\n\n### Experiment results\n\nThe paper has studied empirical results in a few setups. Maybe the most practically relevant result is in Fig 8 where the testbeds are higher-dimensional control problems. Though it seems that there is some marginal gains of WPO/SPO compared to baseline algorithms, I do feel the performance of baseline PPO to be a bit worse than expected, compared to results reported in e.g., [3]. It is also not clear to me why the training step is generally <1M steps since such near on-policy optimization algorithms are usually benchmarked with 3-5M steps. This brings into question whether other baseline algorithms are not properly tuned with the 1M regime; or rather, it is just more convenient to run all algorithms longer for 3-5M steps and see what happens.\n\n[3] Openai spinning up, https://spinningup.openai.com/en/latest/spinningup/bench_ppo.html\n\n### Minor: notation\n\nInstead of using $M$ as the distance matrix, it might be more clear to use $D$, which is compatible with the definition $D_{ij}=d(a_i,a_j)$.", "broader_impact_concerns": "No ethical impact.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "2Kp2cQcdbR", "reviewer": "Reviewer_yvPF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n\n- Overall the paper is clearly written and presented. The toy problem in Figures 1-3 is interesting and illustrative. It’s interesting that the Wasserstein distance’s sensitivity to “close” vs. “far” actions permits a *bigger* policy change in this case, which skips over the local minimum of “picking up” in the blue state. How sensitive is this phenomenon to the specific values of the geometric distances among actions, and can we think of this insight as applying to high-dimensional problem settings? \n\n- Related work is appropriately cited and discussed. \n\n- Most of the empirical evaluation is fairly convincing wrt the performance of both WPO and SPO, and the pros and cons of each approach is discussed nicely.\n\n- The theoretical analysis seems thorough and rigorous, and reinforces intuition about how the proposed approaches work, though see below for more discussion. \n\n\nWeaknesses\n\n- Regarding the performance improvement guarantee, my impression is that the challenge of the exploration problem is packed away into the assumed bound on the advantage estimation error—if the max error across state-action pairs is small, then there’s an implication that the agent has strong coverage across the state(-action) space. That’s fine, but the term in the bound $-2\\epsilon/(1 - \\gamma)$ implies that $\\epsilon$ likely has to be quite low in practice for the guarantee of improvement to hold, as $1/(1-\\gamma)$ can be quite high for typical discount factors. If exploration is hard, it’s not unreasonable to expect that $\\epsilon$ could be relatively large. This does seem to be a weakness of the result, although my primary concern with it is that it doesn’t seem to be discussed/addressed.\n\n- This is relatively minor, but in several cases runtimes are reported without error bars. \n\n- The performance on continuous control is more ambiguous—I don’t think the results support the conclusion that WPO/SPO are clearly better than other methods. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduces, analyzes, and tests two conceptually simple, yet novel policy optimization algorithms based on taking gradient steps constrained by trust regions defined in terms of the Wasserstein distance (WPO) or the related Sinkhorn divergence (SPO). While previous work has studied similar ideas, this paper is the first to provide analysis of the convergence properties of such methods. Specifically, they show a performance improvement guarantee for WPO given possibly inaccurate advantage estimates, a global convergence guarantee for WPO, and that SPO converges to WPO as the regularizer coefficient $\\lambda\\to\\infty$. Empirically, they apply practical versions of these algorithms to both tabular and discrete and continuous control tasks with function approximation, demonstrating strong performance relative to other on-policy trust region-based policy optimization algorithms.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Requested Changes\n\n- Empirically, results on more challenging continuous control tasks (e.g., Humanoid) seem important to support the claim that WPO/SPO are superior to other methods in this domain. Alternately, more random seeds on existing tasks would also be helpful. \n\n- While the paper is clear overall, more discussion regarding the intuition of the WPO update would be helpful. \n\n- A discussion surrounding $\\epsilon$ in the performance improvement result for WPO would also be welcome. \n\n- I would appreciate error bars / variance measurements be reported for the wall-clock numbers. \n\n- Figure 10 in the Appendix seems to indicate that a KL-based method outperforms WPO for higher values of $N_A$, which doesn’t seem to be defined. Could an explanation of this result be added? \n\n- Could the authors add runtime results for the continuous control experiments? I would think the need to perform additional sampling would drive up computational cost. \n\n- The result regarding exploration in the motivating example at the beginning of the paper is interesting, and a discussion of how WPO/SPO can affect exploration would be very helpful.\n\n- Could the authors add a discussion and/or results regarding the factors governing the “use optimal $\\beta$ for the first 20% of runtime then decay/freeze” strategy? Would similar performance be expected for only computing the optimal $\\beta$ for 10% of the run? For 5? How environment/task-dependent is this? \n\nI am aware that all of these changes together may be tough to fit within the main paper or infeasible, but any added results or discussion in the appendix would be appreciated. Thank you!", "broader_impact_concerns": "I don't have any concerns regarding the ethical implications of this work. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: **Summary of the Paper:**

The paper proposes a new trust region methods for policy optimization. As opposed to TRPO, which uses the KL divergence to define the trust region, this paper suggests using the Wasserstein distance, as well as its regularized variant, to define the trust region. The resulting algorithm are WPO and SPO.
The paper derives the algorithms, shows some theoretical properties regarding their convergence, and empirically compares them with other algorithms including TRPO and two other methods (BPGD and WNPG) that also use Wasserstein distance.



**Evaluation:**

After the initial review, the authors revised their paper, and many initial concerns have been satisfactorily addressed. We have now three reviewers [br9i, xF5g, yvPF] who are Leaning Accept and one [dwsm] who is Leaning Reject.

At the high-level, the concerns are related to the novelty of the algorithm, the significant of theoretical results, and the significance of empirical results. I briefly explain them here, and postpone more discussions, and my reaction to them, to later.

*Regarding the novelty of methods:* Although there are similar algorithms (proposed by Pacchiano et al. 2020, Moskovitz et al. 2021), they are not the same. So even though the method may not be considered unprecedentedly novel, this is not a serious issue, especially because TMLR's main acceptance criteria is not based on novelty.

*Regarding theoretical result:* The theoretical results may not be the strongest. For example, the error of critic is measured in the supremum norm, which is technically simpler than an Lp norm but loses some nuances of how critic actually affects the performance. Nonetheless, I believe that is not a real issue as it still shows some basic and important properties of the algorithm.

*Regarding the empirical evidence:* The learning curves overlap significantly in some domains such as Humanoid-v2 or Ant-v2. This makes it difficult to infer much about the superiority of the algorithm, especially compared to WNPG and BGPG that use Wasserstein penalty as well. Therefore, claims such as "our methods achieve better sample efficiency, faster convergence, and improved final performance" (Item 3 in contributions in Section 1) sounds stronger than justified.


I have a few other comments from my own reading of the paper, which I will describe later. I believe they can be addressed with some work. Altogether, **I consider this paper acceptable after a few clarifications and minor revisions**.



**Detailed Comments:**

I summarize some of the comments of the reviewers, and write down my reaction to them and what I recommend the authors should do about them:

**Concern:** [yvPF,dwsm] The improvement over the baseline are not particularly convincing.

**Reaction and recommendation:**
I agree that the improvement over the baselines, especially other Wasserstein-base algorithms, are not visually significant in some domains. For example, most methods on Ant and Humanoid have a significant overlap in their shaded areas. This brings the question of what the shaded areas in the figures represent. Are they standard deviation across runs? Or standard error? Or some 95% or 95% confidence interval? (I see that in Figures 9 and 10 in the the supplementary material, standard deviation is mentioned, but not in other figures.)

My recommendation is that the authors be more clear about what they are presenting, and ideally do a statistical test to see whether the differences are actually significant or not. I know this is uncommon practice nowadays, but when the results are so close, it is warranted. It may also warrant to tone down the claim of superiority in the paper.



**Concern:** [dwsm, xF5g, yvPF] The theoretical results are not strong enough, for example, its use of the uniform error bound on the critic hides away the issue of exploration.
Moreover, at the limit of $\lambda$, they reduce to the Policy Iteration algorithm (Theorem 5). In particular, they do not shed light on why Wasserstein distance should be preferred to the KL divergence (mentioned in the private  recommendation).

**Reaction and recommendation:**
As far as I can see, the theoretical results do not explain why Wasserstein distance should be preferred, so I agree with reviewer dwsm on this front. It would be great if they did, or if they do but is not clearly apparent, the authors expand on it. This isn't a critical shortcoming for this work though.

On the other hand, I do not agree that when $\lambda$ goes to infinity, the algorithm reduces to the Policy Iteration algorithm. That would be the case when $\beta$ goes to zero, which is basically when we disregard the constraints imposed by the algorithm and "turn it off". So even if in the limit the algorithms become the same as Policy Iteration, that limit is not a good description of what the algorithm is doing and is analyzed by the theorem.

Overall, I think even though the theoretical results are perhaps not groundbreaking, they provide some evidence on the soundness of the algorithm.



**Concern:** [br9i] Comparison with SAC.

**My reaction and recommendation:** Given the similarity of the use of entropy, a comparison can be helpful, though this is optional.


**Concern:** The Wasserstein distance has been used in the policy optimization literature before (Pacchiano et al. 2020, Moskovitz et al. 2021), albeit in different forms. This somewhat reduces the novelty of the paper.

**Reaction and recommendation:**
This does not seem to be a major issue, given that there is enough differences between those algorithms and that the TMLR's criteria is not focused on the novelty anyway.

**Concern:** [xF5g, dwsm] It is not clear how to choose the action metric d.

**Reaction and recommendation:**
The question of how to choose the action metric d is an important one, but I suppose it can be done as a future work. Though I encourage the authors to explore it further for this work as well, for example, compare the effect of various choices of d on the performance.



**My Comments:**

**Major:**
* I am confused about the upper bound in Theorem 4. Can't we choose a large $\beta$ to make the RHS arbitrary large? As far as I see, all terms in the summations of the RHS are positive, so if we choose $\beta$ large enough, we show that the performance at the new iteration is arbitrary better than the performance at the old iteration. The only way I see this does not happen is if $f$ goes to zero faster than $\beta$ increases. I have not verified this myself. Please expand on this.

* One of the motivations of the using SPO over WPO has been the computational cost of computing the Wasserstein distance by solving the Optimal Transport problem compared to the Sinkhorn algorithm (for example, this is alluded in the last paragraph of Section 2). But the results in Table 5 in the supplementary material and wall-clock report in Section 7.3 do not show much of a difference. Please explain this more.



* After Theorem 2, it is mentioned that one can use gradient-based method to find the optimal solution of $\beta$. Is it established that Eq. (9) is convex?



**Minor:**

* What is distribution $\mu$ in Theorem 4? How is related to $\nu$?

* Given that most of the algorithmic development is for finite action spaces, it would be better to be upfront about that, perhaps as early as Section 1 or 2.

* Some of the references do not seem to be the right ones. For example, Mnih et al. 2013 or Silver et al. 2016 are not good examples of Policy-based RL.

* The equation numbers and references are not clickable.

* There are some typos. For example:
- Abstract and Remark 3: close-form --> closed-form
- In a few places, "few" is used instead of "a few". For example, "Recently, few work ..." (Page 1) or "In addition, few recent work ..." (Page 3).



**Requested Changes and Clarifications:**
These are the list of changes and clarification that I'd like the authors to make in their revised paper:

- Clarify the meaning of shaded areas in the revised paper and ideally perform a statistical test to show that whether the suggested methods are better than existing ones, especially for results with overlapping shaded areas. If needed, tone down the claims.
- Clarify my questions above, especially the one regarding the RHS of Eq. (10) in Theorem 4.
- Fix all the typos.
- [Optional] Compare with SAC.
- [Optional] Perform experiments with different choice of action distance d.

Decision: Accept (poster)