Efficient Wasserstein and Sinkhorn Policy Optimization
Authors: Jun Song
Chaoyue Zhao
Niao He
Venue: ICLR 2022 Submitted (2022)

Abstract: Trust-region methods based on Kullback-Leibler divergence are pervasively used to stabilize policy optimization in reinforcement learning. In this paper, we examine two natural extensions of policy optimziation 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. Experiments across tabular domains and robotic locomotion 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.

Under review as a conference paper at ICLR 2022
EFFICIENT WASSERSTEIN AND SINKHORN POLICY OP- TIMIZATION
Anonymous authors Paper under double-bling review
ABSTRACT
Trust-region methods based Kullback-Leibler divergence are pervasively used to stabilize policy optimization in reinforcement learning. In this paper, we 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 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. Experiments across tabular domains and robotic locomotion tasks further demonstrate the performance improvement both approaches, more robustness of WPO to sample insufficiency, and faster convergence of SPO, over state-ofant art policy gradient methods.
INTRODUCTION
Policy-hased reinforcement learning (RL) approaches have received remarkable success in many domains, including video games (Mnih et 2013 Mnih et al. 2015). board games (Silver et al. 2016 Heinrich & Silver 2016), robotics (Grudic et al.
[0033 Gu et al. (017). and continuous control tasks (Duan et al. 2016 Schulman al. 2016). One prominent example policy gradient method
2016;; Sutton et al. 1999]] Williams 1992] Mnih al.j 2016 Silver et al. (014)) 4p. The core idea to represent the policy with probability distribution T 8 Pla 0j. such that the action in state is chosen stochastically following the policy Th controlled by parameter 0. Determining the right step size to update the policy is crucial
(Grudic et al. 2003
Peters & Schaal 2006 Lillicrap et al.
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.
To control the size of policy updates, Kullback-Leibler (KL) divergence is commonly adopted to measure the difference between two policies. For example, the seminal work on trust region policy optimization (TRPO) by /Schulman et al. (2015) introduced KL divergence based constraints (trust region constraints) to restrict the size of the policy update; see also Peng et al. (2019); Abdolmaleki al. (2017) introduced a KL-based penalty term the
et al. 2018). Kakade (2001) and (Schllman objective to prevent excessive policy shift.
Though KL-based policy optimization has achieved promising results, remains interesting whether using other metrics to gauge the similarity between policies could bring additional benefits. Recently, few work (Richemond & Maginnis 2017: Zhang et al.) 2018, Moskovitz al. 2020; Pacchiano et al. 2020p has explored the Wasserstein metric to restrict the deviation between consecutive policies. Compared with KL divergence, the Wasserstein metric has several desirable properties. Firstly, is a true symmetric distance measure. Secondly, it allows flexible user-defined costs between actions and less sensitive to ill-posed likelihood ratios. Thirdly but most importantly, the Wasserstein metric has a weaker topology (Arjovsky et al. 2017), thus possibly leading to more robust policy and more stable performance.
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 differ- ent techniques to approximate the policy update under Wasserstein regularization. For example, Richemond & Maginnis (2077) solved the resulting RL. problem using Pokker-Plance equations;

---

Under review as a conference paper at ICLR 2022
Zhang et al. (2018) introduced particle approximation method to estimate the Wasserstein gradient flow. Recently, Moskovitz et al. (2020) instead considered the second-order Taylor expansion of Wasserstein distance based on Wasserstein information matrix to characterize the local behavioral structure of policies. |Pacchiano et al. 2020) tackled behavior-guider policy optimization with smooth Wasserstein regularization by solving approximate dual reformulation defined on reproducing kernel Hilbert spaces. Aside from such approximation, some of these work also limits the policy representation a particular parametric distribution class, As indicated in Tessler et al. (2019), since parametric distributions are not convex in the distribution space, optimizing over such distributions results local movements in the action space and thus leads convergence to a sub-optimal solution. Henceforth, 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 a smooth variant of Waserstein distance by imposing an entropic regularization to the optimal transport problem (Cuturi) 2013). We call them, Wasserstein Policy Optimization (WPO) and Sinkhorn Policy Optimization (SPO), respectively. Instead confining the distribution of policy to a particular distribution class, work 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 exact characterization of the policy updates. We highlight our contributions as follows:
1. Algorithms: We develop close-form expressions of the policy updates for both WPO and SPO based on the corresponding optimal Lagrangian multipliers of the trust region constraints. particular, the optimal Lagrangian multiplier of SPO admits a simple form and can be computed efficiently. A practical on-policy actor-critic algorithm is proposed based on the derived expressions of policy updates and advantage value function estimation.
2. Theory: We theoretically show that WPO guarantees a monotonic performance improvement through the iterations, even with non-optimal Lagrangian multipliers, yielding better and more robust guarantee compared to that using KL divergence. Moreover, we prove that SPO converges to WPO as the entropic regularizer diminishes.
3. Experiments: A comprehensive evaluation with several types of testing environments including tabular domains and robotic locomotion tasks demonstrates the efficiency and effectiveness of WPO and SPO. Compared to state-of-art policy gradients approaches using KL divergence such as TRPO and PPO. our methods achieve better sample efficiency. faster convergence, and improved final performance. Our numerical study indicates that by properly choosing the weight of the entropic regularizer, SPO achieves a better trade-off between convergence and final performance than WPO.
Related work: Wasserstiein-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 metries were recently utilized for representation learning of RL; see e.g., Castro (2020); Agarwal et al. (2020). The most related work to ours are Richemond & Maginnis (2017); Zhang et al. (2018); Moskovitz al. (2020); Pacchiano al. (2020). These work directly use Wasserstein-like distance to measure proximity of policies instead of states. Unlike ours, these work apply Wasserstein distance as an explicit penalty function instead of trust-region constraints. Moreover, they use different strategies to approximate the Wasserstein distance. The only work that exploited Sinkhorn divergence in RI., to our best knowledge, Pacchiano et al.] (2020). In addition, few recent work has also exploited Wasserstein distance for imitation learning; see e.g., Xiao et al. (2019) Dadashi et al. (2021).
Wasserstsein-liee metrics are also pervasively studied in distributionally robust optimization (DRO); see e.g., the survey by Kuhn et al. (019) and references therein. Despite the similarity shared in the duality formulations, the DRO problems are fundamentally different from constrained policy optimization. We also point out that a recent concurrent work by Waan et al. (2021a) studied DRO using the Sinkhorn distance.
2

---

Under review as a conference paper at ICLR 2022
2 BACKGROUND AND NOTATIONS
Markoy Decision Process (MDP): We consider an infinite-horian discounted MDP, defined by the tuple (S, A. P. r, Po: 7), where S is the state space, is the action space, P: : S A x - R the transition probability, r S x A - IR is the reward function, Po S + R is the distribution of the initial state so. and y is the discount factor. We define the return of timestep f as the accumulated discounted reward from t, R
Ex= r( at+k). and the performance of stochastic policy ('r(s,, 4x) where ~ π [s) S1+1 P(st+1|$$, at). As shown p. the expected return of a new policy x' can be expressed in terms
as (T) Eso.a0,5 51
Σ -==
00
in Kakade & Langford (2002)
of the advantage over the old policy т: J(x) = [T)) + E srppp'' (=",(_"(ss where (s a) E[Re s; x represents the advantage function and P. represents the unnormalized discounted visitation frequencies with initial state distribution v, i.e., pT()) Esg (Et 7 P(((s s|
E[R, is s, a: t||
Trust Region Policy Optimization (TRPO): In TRPO (Schulman et al. 2015): the policy π is parameterized as गe 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(x') [TT,, or equivalently, the expected value of the advantage function:
max E Ptmpl (A
gr s.t. Es (dxr 0)] S 6,
where dxz represents the KL divergence and 5 is the threshold of the distance between the new and the old policies. Note that here the expected value of the advantage function is an estimation as the visitation frequency pe is used rather than p. which means the changes in state visitation frequencies caused by the changes in policy are ignored.
Wasserstein Distance: Given two probability distributions of policies т and n on the discrete action space A {a1, 02,,..,, aN J. the Wasserstein distance between the policies is defined as:
dw (w' π) x==
inf (Q, M), QEII(F',))
(2)
where ( denotes the Frobenius inner product. The infimum is taken over all joint distributions Q with marginals T' and r,, and M is the cost matrix with elements M., d(a,, az). where d(a;, a,) is defined as the distance between actions a, and ay-
Sinkhorn Divergence: Sinkhorn divergence (Cuturi, 2013) provides smooth approximation of the Wasserstein distance by adding an entropic regularizer. The Sinkhorn divergence is defined as follows:
ds(n' T a/A)=
inf QETT(=(())
(0 (Q. -
(3)
where h(Q)= Σラ= N Qiy log Qij represents the entropy term, and À >> is regularization parameter. Similarly, use Q denote the joint distribution of I (-|s) and (-1s) with Σ Qi (a,ls)) and ΣØ= Oi = T' r'(alss) The intuition of adding the entropic regularization is: since most elements of the optimal joint distribution Q will be 0 with a high probability, by trading the sparsity with entropy, a smoother and denser coupling between distributions can be achieved (Courty et al.) 2014 2016). Therefore, when the weight of the entropic regularization decreases (i.c., À increases). the sparsity of the divergence increases, and the Sinkhorn divergence converges to the Wasserstein metric, i.e., limx (s'' x{A) = dw (T', T). More critically, Sinkhorn 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 Wassersterin- like metrics in many machine learning domains, including online learning (Cesa-Bianchil Bianchi & Lugosi| 2006), model selection (Juditsky et al. [2008: Rigollet & Tsybakov] (001), ) generative modeling (Genevay et al., 2018 (0188/Pet Petzka et al. 2017: 1 Patrini et al. 2019), dimensionality reduction
(Huang et al. 2021 Lin et al. 2020 Wang al. (022bb.
3

---

Under review as a conference paper at ICLR 2022
3 WASSERSTEIN POLICY OPTIMIZATION
Motivated by TRPO, here we consider a trust region based on the Wasserstein metric. Moreover, we lift the restrictive assumption that policy has to follow a parametric distribution class by allowing all admissible policies. Then, the new policy I' 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
Esppy Ox
E'ED
(4)
where
D { {Tr|Be (dw (T (-1s), (-1+)) <d},
where the Wasserstein distance (:;-,) ) is defined in
In most practical cases, the reward is bounded and correspondingly, the accumulated discounted reward RG is bounded. So without loss of generality, we can make the following assumption:
Assumption 1. Assume A (s, a) is bounded, i.e., suPaEA,EES |A (s, a)l & AT for some A
With Wasserstein metric based trust region constraint, policy update shown in Theorem [ The main idea to form the Lagrangian of the constrained optimization problem presented above, and the detailed proof be found in Appendix A
able derive the closed-form of the
Theorem 1. (Closed-form policy update) Let k, (B. j) argmax ( a 3Mgj ). where M denotes the cost matrix. If Assumption [ holds, then an optimal solution to the WPO problem in (4) is given by:
"t (a,ls)) = E (a,b6)))111)
where = if and fi = 0 otherwise, and 3" is an optimal Lagrangian multipler corresponds to the following dual formulation:
min F(B) 820
min 320
36 + E
EaNp
(14ys))AA"(s,04448))) -BM47085111
(6)
Moreover, we have 8 S = max. (Mkj)- ax) (s,
The exact policy update for WPO in (5) requires computing the optimal Lagrangian multiplier 5 by solving the one-dimensional subproblem (6). A closed form of 3" is not easy to obtain in general, except for special cases of the distance d(z, y) or cost matrix M. In Appendix G we provide the closed form of 3* for the case when d(æ,, v) = 0 if x y and 1 otherwise.
WPO Policy Update:
Based on Theorem I we introduce the following WPO policy updating rule:
Te+1 (a:]s) = FF
Ta(agls)f{
(WPO)
where we choose an arbitrary kTt (Be,j) E argmax No. of (s, B.Maş} and set (B) = 1 and other entries to be 0.
Note that different from (5). we allow 8r to be chosen arbitrarily and time dependently. We show that such policy update always leads to monotonic improvement of the performance even when 3 is not the optimal Lagrangian multiplier. In particular, we propose two efficient strategies update the multiplier 3p:
(i) Time-dependent Be: To improve the computational efficiency, we can simply treat 8, as a time-dependent control parameter, e.g., can choose 3, to be diminishing sequence.
iii) Approximation of optimal By: To improve the convergence, can approximately solve the optimal Lagrangian multiplier based on Sinkhorn divergence. We will discuss this more detail in Section[4
4

---

Under review as a conference paper at ICLR 2022
Next, we provide theoretical justification that WPO policy update is always guaranteed to improve the true performance J monotonically if we have access to the true advantage function. If the advantage function only be evaluated inexactly with limited samples, then extra estimation error will incur. The detailed proof can be found in Appendix B
Theorem 2. (Performance improvement) For any initial state distribution and any Br 2 0, if || - ls I c ffoor some e > O, the WPO policy update with the inaccurate advantage function A", guarantees the following performance improvement bound,
J( 2 ((m)) + B.E. supp 76+1 Σ
j=1 Te(ayls) Ry (B63)) --7
Note that when the estimation error
O. have monotonic performance improvement J (T++1) 2 ((Te) for any 3r 2 0. If the second term in 77 is non-zero, then we have a strict monotonic improvement. Compared to the performance bound when using KL-based trust region 2e ( 2 ((()) 1-7 (see, e.g.. Schulman et al. 2015 Cen et al. (2020), using the Wasserstein metric yields a tighter performance improvement bound and is more robust to the choice of f3.
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 in each step when using Sinkhorn divergence based trust region. The proof follows similar procedure as the Wasserstein policy optimization framework by forming the Lagrangian of the constrained optimization problem. Details are provided in Appendix C
Theorem 3. If Assumption holds, then the optimal solution to the trust region constrained problem 4 with Sinkhorn divergence is:
exp A AT (s, XMa))
(ails)
A 3x exp (s, ax) AMKs) T(a])));
where M denotes the cost matrix and B% is an optimal solution to the following dual formulation:
ming20 (x = min 35 Exmpy Ej=1 ग (agls)( + In(T((a)s))
N In D exp (sA a.) AM,))) Esppp N N Σラ
exp EN exp (5
(9)
Moreover, we have 37 2A
In contrast to the Wasserstein dual formulation (6). the objective in the Sinkhorn dual formulation 9 differentiable in 33 and admits closed-form gradients (shown in Appendix (E. With this gradient information, we can use gradient-basse global optimization algorithms (Wales & Doye 1998 Zhan
al. 2006 Leary 2000) to find global optimal solution 8x
Next, we show that if the entropic regularization parameter À is large enough, then the optimal solution 8 is close approximation the optimal solution 3 the Wasserstein dual formulation. The detailed proof is provided in Appendix
[F
Theorem 4. Define BUB = max 2 B).
The following holds:
1. (B) converges to F(3) uniformly on [0, Sup).
2. lim arg min (x) C arg min F(3)
X 0S
0535880
Although is difficult to obtain the exact value the optimal solution 3" to the Wasserstein dual formulation (6) the above theorem suggests that approximate 8 via 83 by setting up relative large X. In practice, we can also adopt a smooth homotopy approach by setting an increasing sequence X2 for each iteration and letting X co.
5

---

Under review as a conference paper at ICLR 2022
SPO Policy Update:
Based on Theorem
B
we introduce the following SPO policy updating rule:
de A (s, az) M₂)) N exp
(atl(s = ESPO (TT)
=
exp ( +े+ A (s, ax) - Me Tr (agls))
B
(SPO)
Here A 2 and B 2 are some control parameters, The parameter 3, can be either computed via solving the one-dimensional subproblem simply set as a diminishing sequence.
5 A PRACTICAL ALGORITHM
In practice, the advantage value functions are often estimated from sampled trajectories. In this section, we provide a practical on-policy actor-critic algorithm, described in Algorithm that combines WPO/SPO with advantage function estimation.
In each iteration of Algorithm the first step is to collect trajectories, which can be either complete or partial. The difference is whether the return is considered thoroughly to the end of a planning horizon. the trajectory is com- plete, the total return can be directly expressed as the accumulated discounted rewards R, = ST=-1 r If the trajectory is partial, can be estimated by applying the multi- step temporal difference (TD) methods (De) Asis et al. 2017): Re:t+n nn-1 r +
Algorithm 1: On-policy WPO/SPO algorithm
Input: number of iterations K, learning rate a Initialize policy xo and value network Vy, with random parameter tho for k 0, 1, K do
Collect a set of trajectories Dz on policy a For each timestep t in each trajectory, compute total returns G, and estimate advantages AT
Update value:
(* ). Then for the advantage estima- tion, we can use Monte Carlo advantage esti- mation, i.e., An* = R₂ Vos, (st) or General- ized Advantage Estimation (GAE) (Schulman et al. 2016), which provides more explicit control over the birs-variance trade-off. In the neural net to rep-
th+1 + t/k a Σ(G, Vpa (s:))
Update policy:
Th+1 + F(mT) via WPO or SPO with A7
end
value update step,
resent the value function, where e is the parameter that specifies the value net s - V(s). Then, we can update va by using gradient descent, which significantly reduces the computational burden of computing advantage directly.
6 EXPERIMENTS
In this section, we evaluate the proposed WPO and SPO approaches on tabular domains and robotic locomotion tasks as presented in Algorithm We compare the performance of proposed methods with several benchmarks, including TRPO (Schulman et al.) 2015). PPO (Schulman et al.) 2017). and A2C (Mnih et al. 2016 We compare with A2C because is similar to our framework in the sense that both of them are simple on-policy actor-critic methods that utilize the advantage information to perform policy updates. For environments with a discrete state space (e.g.. tabular domains), policy updates are performed for all states at each iteration. For environments with a continuous state space, a random subset of states sampled at each iteration to perform policy updates.
6.1 ABLATION STUDY
In this experiment, we first examine the sensitivity of WPO in terms of different strategies of 3,, We test four settings of 3 value for WPO policy update: (1) Setting I: computing optimal 33 value for all policy update; (2) Setting 2: computing optimal 3 value for first 20% of policy updates and decaying for the remaining; (3) Setting 3: computing optimal 3 value for first 20% of policy updates and fix 3 its last updated value for the remaining;
(4) Setting 4: decaying 3 for all policy updates (c.g..
B = 12) ). In particular, Setting 3 is rooted in the observation that 8 does not change significantly
Wee use the implementations of TRPO, PPO and A2C from OpenAl Baselines (Dhariwll et al. 2017 for
MuluCo tasks and Stable Baselines (Hill al. 2018 for other tasks.

---

Under review as a conference paper at ICLR 2022
throughout all the policy updates, especially in the later stage in 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 is selected for this experiment.
The performance comparisons and average run times are shown in Figure | and Table [ respectively. Figure and [ indicate a tradeoff between computation efficiency and accuracy in terms of different choices of value. Setting 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. We also compare WPO with SPO under different constant and time varying values on the Taxi task. As shown in Figure [] SPO converges faster than WPO. With more weight on the entropic regularization of Sinkhorn divergence (i.e., smaller A), SPO can
Table 1: Run time comparison for different B settings
Runtime
Taxi (s)
CartPole (s)
Setting
1224
130
Setting 2
648
63
Setting 3
630
67
Setting 4 522
44
speed up its convergence more; while as À increases, the convergence becomes slower but the final performance of SPO improves and becomes closer to the final performance of WPO, which verifies the convergence property of Sinkhorn Wasserstein distance shown in Theorem 4 Therefore, the choice of can effectively adjust the trade-off between convergence and final performance. With proper Ả choice, SPO is able attain faster convergence speed with optimum that is only slightly lower than WPO.
More experiments for ablation study is conducted on the Chain (Dearden al. 1998) and CartPole
(Barto et al. 1983) tasks. Results are provided in Appendix H
Taxi
Taxi
Taxi
0
0
-100
-100
-100
Average Return
-200
-200
Averape Return
-300
Averaçe
300
300
WFO Oytimal
WFO Optimal + drcey
-500
WPb
5-0
400
SPD (A logo
-600
WPO Docay 8
constant
-600
30
-500
SPD
-700
-600
0.0
0.5
1.0
1.5
2.0
25 3.0
0.0
0.5
1.0
1.5
2.0
0.0
0.5
1.0
1.5
2.0
Timestaps
x10
Timeslops
x10
Timeslaps
x10
Figure 1: Episode rewards during the training process for the Taxi task with different 8 and À settings, averaged across 3 runs with a random initialization. The shaded area depicts the mean = the standard deviation.
6.2 TABULAR DOMAINS
We evaluate WPO and SPO on a set of tasks including Taxi, Chain (Dearden et al. 1998), and Cliff Walking (Sutton & Barto 2018), which are intentionally designed test the exploration ability of the algorithms. The tabular domain has special environment structure with discrete state space and discrete action space. Thus, we use an array of size SS] x |A| to represent the policy T(ass). 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. Each algorithm is evaluated 5 times with a random initialization. Results are reported Table 2 and Figure 2 The setting of hyperparamaters and network sizes of our algorithms and additional results are provided in Appendix H
Table 2: Trained agents perfor- mance Taxi (averaged over 1000 episodes)
WPO
TRPO
Success (+20)
0.753
0
Fail (-10)
0.232
0
Timesteps (-1)
70.891
200
Avg Return
-58.151
-200
As shown in Figure 2 the performances of WPO, SPO and TRPO are manifestly better than A2C and PPO. Between the trust region based methods, WPO and SPO outperform TRPO in most tasks, except in Chain, where the performances of these three methods are not significantly different. In Taxi and Cliff Walking, SPO converges to the optimum the fastest, while in Taxi, WPO converges to the best optimum, among all methods. We further analyze the performance of the trained agent for each algorithm the Taxi environment. As shown in Table 2 WPO has higher successful drop-off

---

Under review as a conference paper at ICLR 2022
rate and a lower task completion time while the original TRPO reaches the time limit with drop-off rate 0. Therefore, the results in Taxi show that WPO finds a better policy than the original TRPO.
Taxi
Cliff
4000
Chain
Aerage Return
200
3500
-300
-1000
3000
-400
Average
-1500
Averagge
-500
UPD
5PC
2000
2500
-700
-600
TIPO
-2500
2000
H90
800
3.0
0.5
20
25
-300
3.00
0.05
0.10
3:15
0:20
0.25
0.30
1500
0.0
3.5
Timestops
x10'
Timestops
x10
Thr esteps
K10
Figure 2: Episode rewards during the training process for the tabular domain tasks, averaged across runs with a random initialization. The shaded area depicts the mean + the standard deviation.
We also show that compared with the KL. divergence, which is used in traditional TRPO and PPO approaches, the utilization of Wasserstein metric can cope with the inaccurate advantage function estimations caused by the lack of samples. We compare WPO with KL (Algorithm | framework with KL based policy update derived in Peng al. (2019) on the Chain task. We evaluate the performance of these two algorithms under different NA. which denotes the number of samples used to estimate the advantage function at each iteration. As shown Figure 3 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.
Chain
Chain
Chain
4000
2500
3000
2500
Neorgg Recur
3000
2500
2500
2500
2000
WPO
WPD
WP0
1500
KI
1600
1500
1000
100g
1000
0.0
0.1 0.2
0.3 0.1 0.5 0.6 0.7 0.8 38
0.0 0.1 0.2
0.3 0.1 0.5 0.8 0.7 0.8 099
0.0 0.1 0.20222222 0.3 04 0.5 0.6 0.7 0,5 08
Timesteps
X10
Timestape
x10'
Timestops
x10*
(a) N 100
(b) NA 250
(c) Nc = 1000
Figure 3: Episode rewards during the training process for the Chain task, averaged across 3 runs with a random initialization. The shaded area depicts the mean + the standard deviation.
6.3 ROBOTIC LOCOMOTION TASKS
We then integrate deep neural network architecture into MPO and SPO and evaluate their performance on several discrete locomotion tasks, including CartPole (Barto et al. 1983) and Acrobot (Geramifard et al. 2015). We use two separate neural nets to represent the policy and the value. The policy neural net receives state as an input and outputs the categorical distribution of π (als). The performance of WPO and SPO are compared to that of TRPO, PPO and A2C under the same neural net structure. We run each algorithm times with a random initialization.
Figure shows the episode rewards during training process for WPO, SPO and baseline algorithms. As seen in Figure 4 WPO and SPO outperform TRPO, PPO and A2C in most tasks in terms of final performance, except in Acrobot where PPO performs the best. In most cases, SPO converges faster but WPO has a better final performance.
Final Performance:
To train 10a timesteps in the discrete locomotion tasks, the training wall-clock time is around 63s for WPO, 65s for SPO, 59s for TRPO and 70s for PPO. Therefore, WPO has a similar computational efficiency TRPO and PPO.
Training Time:
The performances of WPO and KL are also compared for the discrete locomotion tasks under different NA. As shown in Figure when NA is 500, KL performs better than WPO. However, when NA decreases to 100, WPO sigmificantly outperforms KL. These results indicate that for discrete locomotion tasks, WPO is more robust than KI. when advantage values are inaccurate.
8

---

Under review as a conference paper at ICLR 2022
CariPolo-v1
CarPole-v11
Acrobol-vl
4001
100
300
200
200
100
<
102
25
50
75.
103 125 Epasdds
150
200
50
100 150 Epraade
200
202
-500
90
15
23
25
30
35
Runtime (seconds)
Pantiime (ERcOnTS)
Figure 4: 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.
:001
CartPPloo-11
500
CarPPolowr
-100
Acrobot v1
-100
Acrobot
WFO
IPO
Wecogg Roam
A00
KL
-220
-200
-250
-200
-000
200
200
-990
100
900
-430
WHD
-400
-450
-400
25 50
75 100
125 Equende
190 200
20 50
130
125 150
175 200
100 125
175
Eppecee
(a) A 100
(b) NA 500
(c) NA
100
(d) NA
500
Figure 5: Episode rewards during the training process for the locomotion tasks, averaged across 3 runs with a random initialization. The shaded area depicts the mean + the standard deviation.
6.4 CONTINUOUS ACTION SPACE:
We further extend the evaluation of WPO and SPO environments with a continuous action space by discretizing the action space following Tang & Agrawal (202)). For comparison, we additionally consider Behavior Guided Policy Gradient (BGPG) algorithm from Pacchiano et al.] (2020). Similar results are observed in Figure/oaas the discrete action tasks: WPO and SPO outperform the benchmark algorithms in terms of final performance.
LunarLander-v2
Hopper v2
Walker2dd-2 x2
4000
4000
3500
SPO D0PG
3500
Relur
100
PTO
3000
ADD
2500
2500
-00
2000
2000
1500
1000
1000
-500
570
500
500
1.0
1.5
2.0
2.5
2
4
5
6
4 Timestaps
Timzeleps
x10
x10?
Figure 6: Episode rewards during the training process for continuous action space tasks, averaged across 3 runs with random initialization. The shaded area depicts the mean + the standard deviation.
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 or 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 3 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. For future work, it remains interesting to extend the idea to PPO and natural policy gradients, which penalize the policy update instead of imposing trust region constraint, and extend it to off- off-policy frameworks.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Abbas Abdolmalcki, Jost Tobias Springenberg. Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller. Maximum a posteriori policy optimisation. ArXiv Preprint, pp. arXiv:1806. 886.06920 2018.
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 8th International Conference Learning Representations, 2020.
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein GAN. ArXiv Preprint, pp. arXiv: 100:170110000 10787,, 2017.
Andrew G. Barto, Richard S. Sutton, and Charles W. Anderson. Neuronlike adaptive elements that can solve difficult learning control problems. IEEE Transactions on Systems, Man, and Cybernetics, SMC-13(5):334 846, 1983.
Pablo Samuel Castro. Scalable methods for computing state similarity in deterministic Markov decision processes. In Proceedings of the Thirty-Fourth AAAI Conference on Artificial Intelligence, volume 34, pp. 10069-100777 2020.
Shicong Cen, Chen Cheng, Yuxin Chen, Yuting Wei, and Yuejie Chi. Fast global convergence of natural policy gradient methods with entropy regularization. ArXiv Preprint, pp. arXiv:2007.06558 2020.
Nicolo Cesa- Bianchi 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 Joint European Conference on Machine Learning and Knowledge Discovery in Databases, pp. 274 289. Springer, 2014.
Nicolas Courty, Rémi Flamary, Devis Tuia, and Alain Rakotomamonjy. Optimal transport for domain adaptation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 1853-18833 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, Fernando Hernandez-Garcia, G. Zacharias Holland, and Richard S. Sutton. Multi-step reinforcement learning: A unifying algorithm. ArXiv Preprint, pp. arXiv: 1703.01327, 2017.
Richard Dearden, Nir Friedman, and Stuart Russell. Bayesian Q-learning Proceedings of the Fifieenth National Conference Artificial Intelligence and Tenth Innovative Applications of Artificial Intelligence Conference, pp. 761-768, 1998.
Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, Yuhuai Wu, and Peter Zhokhov. OpenAl baselines. https
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, 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. 1329-1338, 2016.
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite Markov decision processes. In Uncertainty in Artificial Intelligence, volume 4, pp. 162-169, 2004.
10

---

Under review as a conference paper at ICLR 2022
Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with Sinkhorn divergences. In International Conference on Artificial Intelligence and Statistics, pp. 1608- 2018.
Alborz Geramifard, Christoph Dann, Robert H. Klein, William Dabney, and Jonathan P. How. RLPy: f-unctionnnbneed reinforcement learning framework for education and research. Journal of Machine Learning Research, (6(46):1573-1577 2015.
Gregory Z. Grudic, Vijay Kumar, and Lyle H. Ungar. Using policy gradient reinforcement learning autonomous robot controllers. In Proceedings of the 2003 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 406-411, 2003.
Shixiang Gu, Ethan Holly, Timothy Lillicrap, and Sergey Levine. Deep reinforcement learning for robotic manipulation with asynchronous off-policy updates. In Proceedings of the 2017 IEEE International Conference on Robotics and Automation, pp. 3389-33883333 0399,339,2220 2017.
Johannes Heinrich and David Silver. Deep reinforcement learning from self-play in imperfect- information games. ArXiv Preprint. pp. arXiv:1603.01 121, 2016.
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. bttps://gitee com/
hi 1 a/ htt/stabe staslttaaaaeeeeeee lines 2018.
Minhui Huang, Shiqian Ma, and Lifeng Lai. Riemannian block coordinate descent method for computing the projection robust Wasserstein distance. In Proceedings of the 38th International Conference on Machine Learning, pp. 4446-4455, 2021.
Anatoli Juditsky, Philippe Rigollet, and Alexandre B Tsybakov. Learning by mirror averaging. The Annals of Statistics, 36(5):2133-2066 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 Santeerzade---baadee Wasserstein distributionally robust optimization: Theory and applications machine learning. In Operations Research & Management Science in the Age of Analytics, pp. 130-166. INFORMS, 2019.
Robert Leary. Global optimization funneling landscapes. Journal of Global Optimization, 18: 367-383, 2000.
Timothy Lillicrap. Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Proceedings of the 4th International Conference on Learning Representations, 2016.
Tianyi Lin, Chenyou Fan, Nhat Ho, Marco Cuturi, and Michael Jordan. Projection robust Wasserstein distance and Riemannian optimization. ArXiv Preprint, pp. arXiv:2006.07458 2020.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin Riedmiller. Playing Atari with deep reinforcement learning. ArXiv Preprint, pp. arXiv:1312.5602. 2013.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Marc G. Belle- mare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen, Charles Beattie, Amir Sadik, loannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning. Nature, 518(7540):529-535 2015.
11

---

Under review as a conference paper at ICLR 2022
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 Machine Learning, pp. 1928-
1937,2016
Ted Moskovitz, Michael Arbel, Ferenc Huszar, and Arthur Gretton. Efficient Wasserstein natural gradients for reinforcement learning. ArXiv Preprint, pp. arXiv:2010.05380, 2020.
Aldo Pacchiano, Jack Porder-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-7454, 2020.
Giorgio Patrini, Rianne den Berg, Patrick Forre, Marcello Carioni, Samarth Bhargav, Max Welling, Tim Genewein, and Frank Nielsen. Sinkhorn autoencoders. In Uncertainty in Artificial Intelligence, pp- 733-743, 2019.
Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine. Advantage-weighted regression: Simple and scalable off-policy reinforcement learning. ArXiv Preprint, pp. 2019.
Jan Peters and Stefan Schaal. Policy gradient methods for roboties. In Proceedings of the 2006 JEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 2219-2225, 2006.
Henning Petzka, Asja Fischer, and Denis Lukovnicov. On the regularization Wasserstein GANs.
ArXiv Preprint, pp. arXiv:17090 709.08894, 2017.
Pierre H. Richemond and Brendan Maginnis. On Wasserstein reinforcement learning and the Fokker- Planck equation. ArXiv Preprint, pp. arv.a:117112...1185 2017.
Philippe Rigollet and Alexandre Tsybakov. Exponential sereening and optimal rates of sparse estimation. The Annals of Statistics, 39(2)731-777,, 2011.
R. Tyrrell Rockafellar and Roger J. B. Wets. Variational Analysis. Springer, 1998.
Johannes O. Royset. Approximations and solution estimates in optimization. Mathematical Program- ming, 170:479-506, 2018.
John Schulman, Sergey Levine, Philipp Moritz, Michael I. Jordan, and Pieter Abbeel. Trust region policy optimization. In Proceedings the 32nd International Conference Machine Learning. pp. 1889-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:
2017.
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.
David Silver, Aja Huang, Chris J. Maddison, Arthur Guez, Laurent Sifre, George den Driess- che, Julian Schrittwieser, loannis Antonoglou, Vedavyas Panneershelvam, Mare Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy P. Lillicrap, Madeleine Leach, Koray Kavukcuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of Go with deep neural networks and search. Nature, 529(7587):8844455 2016.
Maurice Sion. On general minimax theorems. Pacific Journal of Mathematics, 8(1::171-176, 1958.
Richard S. Sutton and Andrew G. Barto. Reinforcement Learning: An Introduction. MIT Press, 2018.
Richard S. Sutton, David McAllester, Satinder Singh, and Yishay Mansour. Policy gradient methods for reinforcement learning with function approximation. In Advances in Neural Information Processing Systems, pp. 1057-1063, 1999.
12

---

Under review as a conference paper at ICLR 2022
Yunhao Tang and Shipra Agrawal. Discretizing continuous action space for on-policy optimization. In Proceedings of the Thirty-Fourty AAAI Conference on Artificial Intelligence, volume 34, AAAI Press. 2020.
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-56; 1998.
Jie Wang, Rui Gao, and Yao Xie. Sinkhorn distributionally robust optimization. ArXiv Preprint, pp. arXiv:2109.1 1926, 2021a.
Jie Wang, Rui Gao, and Yao Xie. Two-sample test using projected Wasserstein distance. In Proceedings of IEEE International Symposium on Information Theory, volume 21, 202
Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Machine Learning, 8(3-4):229-23 256, 1992.
Huang Xiao, Michael Herman, Joerg Wagner, Sebastian Ziesche, Jalal Etesami, and Thai Hong Linh.
Wasserstein adversarial imitation learning. ArXiv Preprint, pp. arXiv:1906.08 13, 2019.
Lixin Zhan, Jeff Chen, and Wing-Ki Ki Liu. Monte Carlo basin paving: An improved global optimization method. Physical Review E, 73:015701, 2006.
Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin. Policy optimization Wasserstein gradient flows. In Proceedings of the 35th International Conference on Machine Learning, pp. 5737-5746, 2018
13

---

Under review as a conference paper at ICLR 2022
A PROOF OF THEOREM II
Theorem 1. (Closed-form policy update) Let (5) ar. N (s,
.x) 3MRj
where M denotes the cost matrix. If Assumption 7 holds, then optimal solution to the WPO problem in (4) is given by:
x <a(a]s)) = Σ (a,))))))())))
(5)
1=1
and f2 otherwise, and 3 is an optimal Lagrangian multipler corresponds to the following dual formulation:
where
N
min F(B)
min 820
35 + R Snpp
Σ ==1
3M 3Maz(onnd
(6)
820
Moreover, we have 8 s 1- (A A ()
Proof of Theorem 7 First, we denote Q" as the joint distribution of x(-1s) and <'(|ss with Σ Ui <Tayls)) and Σ/= Qi r'(als). Also, let fa(i,)) represent the conditional distribution of (a:ls) under T(alls). Then Qi = ((ag)))).(i..) r'(als)) = EE
Σe= (i. In addition:
dw (T (-))s)))
min Qig
Mi
Mi and min fa(i))
N
N = Σ Af(s.a.)((alss))4
Bexr (1)4
A (s, (a)a))
Thus, the WPO problem in 4 can be reformulated as:
E ExNPE
(10a)
(((,))
N
8.b. Ennng Σ
Σ (1 86,
(10b)
=1
Σ 1(.
Vs E S.j...
(10c)
Note here that (0bb is equivalent to Eanpy min
E Σ (i <
5 because if we have feasible fa(1.) to make (10b) hold,
we must have
Egppy (in) EN Ej=1 Ey=MMMM-M Myyr(ay)s)SA( 5.
Since both the objective function and the constraint linear in fa(i, 10 is convex optimization problem. Also, Slater's condition holds for (I0) as the feasible region has an interior point, which is fi vi. and fi(i,)) o + j. Meanwhile, since AF a) is bounded based on Assumption 1 the objective is bounded above. Therefore, strong duality holds for (0).. At this point we can derive the dual problem of (10) as its equivalent reformulation:
min $20.C
38+ +
ds
SES
j=1
(11)
s.t. A (" ((s,aa)T(a,ls)) - BM, (a,ls(
S 0, (T(ss
As S =1...
1...,
We observe that with a fixed 8, the optimal E3 will be achieved at: max pT(s) (AT az) 3M,))
55 (B)
(12)
-= I....N
14

---

Under review as a conference paper at ICLR 2022
Denote 8 optimal solution to 11I) and f (i. as optimal solution to complimentary slackness, the following equations hold:
10 Due to the
(8*) PT(() (i, - 0,
(A*(s, - 8 Mas((ay)s)
Vs,i,j.
In this case, (i.j) can have non-zero values only when (3 (s,,ai)7((al))) 8 Mayt(ayls) (19))) O, which means (8*) = p. ((()) 5 Msy). Given the expression og(s)
of the optimal 5 12 (i,j) can have non-zero values only when K3 .j), where (B.j) x N... (s, a&) BMR.. Since E (i 1 indicated in 10c).
can choose an arbitrary optimizer kg (0", E (8", .j) to derive an optimal solution f (2,3):
(. = (
if i k=
otherwise,
And then the corresponding optimal solution is, x" (ails) = Σッ a(a, a(a,ss))3((..).
Last, by substituting (5) we can reformulate IIT) into:
(. into the dual problem
/ LES (sr = min 188++
min 38 + 820
Egmpp Σ (sT(8)) 3M (M
820
(13)
The optimal 33 can then be obtained by solving 13
We will further show that 8 <B = max
(Maj)-1(A'6 (s,
In the general case, i.e., $ 2 0. 10a) is non- non-negative because:
N N
Exppp
A (s (1.
(14a)
N x (a;lss) Σ
E Expps
A (s, ((((())))()))))
(14b)
j=1
((al()) (sg))s)))))) (x.
(14c)
N
> E
Expp
T(ayls) ((a)ls)(A"(ssaa) + 3r Mag(0-51))
(14d)
N
E Eenpp
T(a, M37
(14c)
2 0,
(14f)
where (4d) holds since AT aky (1- ,)) ) - 33 May Mi 2 (=s,,
- 33 = (s,))
When 5 max { A= (s.az)=A"'s.0 ), we have that for all E S, Kg (5) =
Maj
Thus, (i, i) 1, vi and (i, 0, Vi # j. The objective value [Oa] will be 0 because
Egppp Σ j=1 (a)
Exmpp IN A as.ssa))((a))
0. The left hand
side of 10b equals to Exppz EN Σjn
Ex-oz EN Mar(a.ls) 0.
Thus, for any 8 > 0. 10b is always satisfied.
Since the objective of the primal Wasserstein trust-region constrained problem in 66 constantly A" ("(s.e)) A (e.0s) }. and is non-negative when 3** <
evaluates to 0 when 3* > max.es N.4j
Mas
15

---

Under review as a conference paper at ICLR 2022
M..j... (=(aaa))=A" we use J.j... Ma; an upper bound for the optimal dual variable 3r.
A (a.0z) 1- (s,aa) Mai
as
B PROOF OF THEOREM 2
Theorem 2. (Performance improvement) For any initial state distribution μ and any Be 2 0, if
|| 1o Se for some > O, the WPO policy update with the inaccurate advantage function A", guarantees the the following performance improvement bound,
J( 2 ((T.) + B,E soupH nr Σ (
j=1 Tr(ayls) (B..))) 1-
2
Proof of Theorem 2
J(m++1) ((Te) = Sa Te Ba (1=" [s)
(15a)
E srup
T (o:ls)
(15b)
E T-- Prupp
(15c)
N
E srupp
Te(ayls) Σ ((.
(15d)
E grupy
Tr (s..gT)(3..)))
(15e)
Ø=1
E T(1
m
+ 3. BT
2c]
(15f)
ppup
j=1
B,E
T
Σ
15 (B.5)) 1 1-
2e
(15g)
srupp
where 15a holds due to the performance difference lemma in Kakade Langford (2002):
15f follows from the definition of kg (Br,j) and the fact that 11 - A 11o0 S e, therefore A (B..)) + c B 2 A (p,))) - 3. M.Mina(a. (8,)))
2 A (s,ag)
B. = as, 2 (s, - c; 15g holds since Egw [ (s, a)) 0.
C PROOF OF THEOREM 3
Theorem 3. If Asssmpptio 7 holds, then the optimal solution to the trust region constrained problem with Sinkhorn divergence is:
exp AT AM4)
(a;ls)
1 (T, an) XMkj) (a,)8));
N
J=1
2k=1
exp
where M denotes the cost matrix and By is an optimal solution to the following dual formulation:
mingzo Fx = mingzo {35 EanpE E (a,)ss) In(T((,ls))
E In exp (AA (s. a.) AM,))) Hampz ΣN Ej &h ΣE exp (s.ae)
exp (MM)) (9)
Moreover, we have 8 < mar
16

---

Under review as a conference paper at ICLR 2022
Invoking the definition of Sinkhorn divergence in (3). the trust region con- strained problem with Sinkhorn divergence can be reformulated as:
Proof of Theorem 3
N 1= (s, 0i Σ j=1 Qill
N
max Q
Bgnpy
(16a)
Egnpp Σ N E
s.t.
MiQ
+ x Qi log Qi <
(16b)
Σ Σn Oij <(ylss), vj = N NeS.S.
N
(16c)
Let B and represent the dual variables of constraints 16b grangian duality of (16) can be derived as:
and
16c
respectively, then the La-
N
N
max min L(Q,Bww) 820,2
= max min Eanpy
A
(s, 0i
Qill
Q $20.00
i=1
j=1 N
N
N
1
Vi {ayls)))ds + 3(6 Exmpy
Mij + 2 log Qiş))
SES j=1
i=1
i=1
==1
(17a)
N ΣΣ
max min E $20,0
Egmpp
A (s.a,) Σ
0iyl
(s)ssss)ds
JOES
(T(s)
N wfr(la ssds + B8 - 3. pT
N
N
1
M + X
2i log Qi,1)
(17b)
SSS
N
= max min
38
Σ
waT(a,
$20,00
.ES
j=1
w
B Qi log Qiyl
Eg pp
(A (s,a.)
BMM
14g
(17c)
pE(s)
N
min max 38
Σ war(la)ssds
$20
JES
J=1
N
N
w (5(()
B A Lez log Qiyl.
E Egppp
(A (s, - 3M
Qi
(17d)
where 17d holds since the Lagrangian function L(Q, 3, w) is concave in Q and linear and w, and we can exchange the max and the min following the Minimax theorem (Sion 1958).
Note that the inner max problem of (7d) is an unconstrained concave problem, and we can obtain the optimal Q by taking the derivatives and setting them to 0. That is,
0L AT a.) 3Mg aQ1
w/
3 (log Qi 1) 0, vi.j=1. L.SSEE N, (18)
(T)()
A
Therefore, we have the optimal Qi5 as:
Aw
2 = =xp (s, ax) - XMr3) exp
1), Vi,j = 1, N,SES.
(19)
Bp((s)
In addition, since MN Qi = =x(a, (s), we have the following hold:
Aug
x(ayls)
exp( Bpp(s)
-1)
(20)
EN exp (x XMa))
17

---

Under review as a conference paper at ICLR 2022
By substituting the left hand side of 20 into 19 we can further reformulate the optimal Qij
cxp ( AM (s 0;) XMyy)
(ayls), Vi,j 1, N, (21)
Si=1 exp A ax) XMaj)
To obtain the optimal dual variables, based on 200 we have the optimal wf as:
3 3 L In(T(ay)s))
A
w/"h - pT(())
In exp
Σ
g
a
AMy))), =j=1,... ,,sE,
A
(22)
By substituting 21 and 22
into 17ch.
obtain the optimal 8 via:
3 In(x(ay(s) In exp A" (8, - AM;)))
min 320
38 Egnpp
T <ay
+ A
N
N
exp [ ( (s, 0z) AM1)) (agls))
+ E
Σ exp ( AM, (8, AM&j)
i=1
j=1
The proof for the upper bound of sinkhorn optimal /3 can be found in Appendix
D UPPER BOUND OF SINKHORN OPTIMAL BETA
In this section, we will derive the upper bound of Sinkhorn optimal 3. First, for given 3, the optimal
Q (B) to the Lagrangian dual 3(Q,3,w) can be expressed in 21). With this, will present the following two lemmas:
Lemma 1. The objective function (6a with respect to Qij (B) decreases as the dual variable increases.
Lemma 2. If Assumption 7 holds, then for every 5 > 0, Q
24 is feasible to 16b for any X.
We provide proofs for Lemma and Lemma 2 in Appendix D.I and Appendix D.2 respectively.
Given the above two lemmas, we are able to prove the following proposition on the upper bound of Sinkhorn optimal 8:
Proposition 1. If B% is the optimal dual solution the Sinkhorn dual formulation 9 then BY S mA for any X.
Proof of Proposition| We will prove it by contradiction. According Lemma Q max is
feasible to 16b). Since 3 is the optimal dual solution, Q73 (Bx) is optimal to 16). If 3 > ma
according to Lemma the objective value in (16a) with respect to 2.4 is smaller than the objective value in 16a) with respect to 3x, which contradicts the fact that Qij (35) is the optimal solution to (66).
D. PROOF OF LEMMA 1
Lemma 1. The objective function (6a with respect Q73 (B) decreases as the dual variable increases.
Proof of Lemma
Let Gx(3) represent the objective function 16a). By substituting the optimal
Q in 21 into
16a
we have:
N exp AT 0:) XMas)
G3 = E EaNpy
AT a.)
1=1 E exp ( A- A5 (s, -MM,)) T(a;
(23a)
exp ( A (s, a:) AM.))
x(az(s) (s)
Σ
A (3
(23b)
Ek=1 exp ( A (s, AMA1)
1=1
18

---

Under review as a conference paper at ICLR 2022
For any B2 > 91 > 0, we have:
(x1 G2
N
(7 exp ( A (s. a.) AMa))
Boppp
Σ T(alls) (8) Σ A as)))))
E exp exp (5 ax) - XMaz)
j=1
i=1
exp
A (s, a:) - AM²)
(24a)
Ek=1 exp
A (8, (z) Maas
32
N
exp (3) (s) (M
Eympp
Σ ((alls)
(8)
Σ A (ss) =
k= exp (,t A (x() >M
j=1
exp
A (s, (p)) (Mass
}.
(24b)
Σ exp ( AT (8, m11) - )M1H),)
where [i] denotes sorted indices that satisty A (T, apt) 2 (s. 2 ... - AT (s) Let
exp (,2 (s, an) - (M
cxp (B. A (s, an)) - (M1
[a(1) N
(25a)
Kk=1 k= exp 1// A (1)) (MH) Σ exp (/=-/ A (s, mw)) - AMay)
exp (( ((}}
A (s) exp (33 (" - (M
8
Σm exp (1)) 2/ (s. exp ( A A (14) - (M11)
exp ( 1// A (s,, mn)) XM
(25b)
k= exp 1 Ar (s, ap1) - {M11)))
For notation brevity, we let m²( = exp ((}} - A ) (s,
By
>
0. w. (1)
=
exp
(1,, - {Mar) > and (i) = E (k))
Then we have
E (0)00s(())
mz (i(()u,(()) (1)
(i)
256
(26a)
Em (k)ww,
Ex (kk
1
m. mi))ws()) (
(26b)
Ck=1 m, (k)
Σx m2 (i))),,()))
m, (1)
(26c)
Since A. 3/ O, (i) decreases as i increases. Thus, (i) decreases as i increases. Since ms(1) (1) 2 m. and m. (N) S mz(k) for all h = 1,...,,,, N, we have (1) = ΣR- (k)w.. 1 Ei (N) Ees (K) EE m. (k)ww,
O, and (NV) = Ei my (k)m, EN m.(k), Σiz (k).. = 0. Since (1) 2 0, go (N) S 0 and
1
ΣS mi (K)
(1) decreases as i increases, there exists an index 1 S k. S N such that (1) 2 0 for i i k. and qi( < 0 for >> ks- Since mi) (i) > 0, we have (i( 2 0 for S k2 and (i) <
for í > kg In addition, we have fa(i) = 0 directly follows from the definition. Thus, EN f.(1) -Ee (f.(1)) - Disk, +1 [ss(i)) 0. Therefore,
N
N AT (1))
(x G2(32) = -Eo-ng Σ {(alss)
(27a)
j=l
N
N AT ap) 1/1(- - Σ (s(,,1))(1.00)
= Esnpg j=1
Σ
((gl())
i=
(27b)
N
E, Earps
Σ
[(a
AT ql. (k.s. ()f((())
Σ
j=1
t +1
(27c)
19

---

Under review as a conference paper at ICLR 2022
N
N2 Σ If.(0) - A (k))))
N
Es
Σ
[fa(i)}
j=1
1=1
i=
(27d)
N
Ax ((.))
Es
Σ
[s-00) - A (3 (8,0
(1(1)))
j=l
(27c)
R. [fa(
E Eg
(s.,)) (s ajk as, +1)))
(27
(27g)
where (2cc) and (27g) hold since (s, (()) is non-increasing increases. Furthermore, at least
one inequality of (27c) and (7g will not hold at equality since Σ ((as)))A" az) = 0. Vs E S and for non-trivial cases, Pr(Aa "A (s. = 0, E S. E A} < which means (r 82 e
S, 01, 02 E A, s.t. A (1) # ( } > 0. Therefore, we have Gx(31) Gx > 0. O
D.2 PROOF OF LEMMA 2
Lemma 2. If A'Assumption holds, then for every 6 > 0, Qij (245 (24"H)) is feasible to 16b for any X.
Proof of Lemma 2 By substinuting the optimal Qi in 21) into (16b), we can reformulate the left
hand side of (Tbb) as follows:
N
N
1 x
oi log Qij]
(28a)
==1==1
Rg
N
Mij
1
Oij A
( -MM + log
π (a,lss)
1}
-k=1 exp ( AT (8, ax) XMas)
(28b)
N
N
1
π (aals)) (s)
= Eampy
2
A (s)
+ x 0ig log
}.
(28c)
Ek=1 exp (g) A (s, (s, ax) - AMa1)
Now we prove that when 3 = x
E Eeupy i= Ex= (0) S
4
and
Exsoy (205) (B) log E exp ( $ A (s,02)
x(a,1s)
hold. For the first part, have:
1 (s.
E Exnot
(29a)
3
N
E
(s,a)))
(s,a.))
(29b)
j=1
Exxps
r' (T
(29c)
< Ex
<((ass)|A"
(29d)
8
A maa
(29e)
B 2
For the second part, the followings hold:
<(ayls)
E
Exyp
Qi log
(30a)
k= exp (s \MaR)
20

---

Under review as a conference paper at ICLR 2022
N
N log
a(ayls)
(30b)
A= exp (s, AM1ss)
==1
1=1
x(alls)
Bxupp
a(azls) log
Σk=1 exp exp (s. (x) AMa))
(30c)
/=1
x(ag18)
E X
Bx
T(azls) log
(30d)
exp (3,
E
a(aals) log
(30e)
X
exp (2ATTs,0,))
3=1
N
Exppy
T(aals)( 8 ((
(30f)
==1
E Egpp
x(c)1s)]44(s.a)))]
(30g)
)=1
A maa
(30h)
3
2'
Therefore, Qif 24
is feasible to (16b) for any A.
E GRADIENT OF THE OBJECTIVE IN THE SINKHORN DUAL FORMULATION
The closed formm gradient of the objective in the Sinkhorn dual formulation p is as follows:
8 FENDS Σ [(a { + + In(=(a)1s))
exp A (s, - -M,)]
3
8
exp
A (s, AMa) x (s, 8
À
Σ exp (s, AMa)
-
N N
(a,ls))
exp (xp() A (s, az) AM1s)
+ B
Σ exp A (8, .x) AMas)
(T(alls) exp ( A" (8. a:) - AM 21" (8, a:)8 Σ Σ exp ( (s, - AMas))
+
Ch exp 3 (T ax) AMky)))
BT(ayls) exp (5 (s, a:) - AM+) x Ex eex 3 ax) AMrz) x (s,
EN exp ( (s, >Ma,)))
F PROOF OF THEOREM 4
Given the upper bound of Wassertein optimal / in Theorem I and the upper bound of Sinkhorn optimal B3 in Proposition
I
we are able to derive the following theorem:
Theorem 4. Define Bus max{" mm mix B). The following holds:
1. F3) converges to F(3) uniformly on [0, Suel.
2. lim arg (B) S arg min F(3).
X x USD5
053580n
Proof of Theorem to show that l SuPos0580s BuB
To show that (3) converges to F(3) uniformly on BuBl, is equivalent
Fx(3) - F(3) = 0. Let (3, = A (23().)) -
21

---

Under review as a conference paper at ICLR 2022
(R,)) (A" a.) BM,y] where kg (B.5) E (5,)) argmax, k= -I... N....NN (s, -3Mbs is an arbitrary optimizer, and eï (B, i, 2 0. First, we have
Fx(3) F(3)
N
N
= 38
E Esppt
Σ T(a
5 π x
3 In((((a]1))) - _ In
Σ
exp ( (T) -
1=1
i=1
8 exp A 1 (s) - AM²) x(agls)
+ Bsnp5
38
Ek=1 exp 1/ A (s, 2Mas)
B
(a3 [A* (p(r)) - 3Marcasvll
(31a)
+ +12 Esmpy Σ
3
Esspp pJ
(aylss)
(aall)
==1
/=1
B exp ( AT (s, a:) - AMij) ((ayls)
+ Bx
A
Σk=1 ex=rexpl (g) (s, ux) AM1j)
T(ayls) 3 In Σ exp (g) At (s, AMay]]
A
+
p
N
B
Σ
T(a, (s){Aa((, (lg( (s, - 3M
(31b)
N
N
$ 2
E Eenpp
Σ
T(az1s)
f 7 +
Kamps
Σ (n(((a))
j=1
j=1
8 In
+
E
T(a
exp 3 (s, (s,
AMa))
- B
x(a, ss) (s, (3.0))(3)))) - 3M
(31c)
In addition,
Eampt
Σ
N
T(a 9 x
N E exp ( A (s, XMMi
1=1
Ex
x ((a)l)) [s)[ss([((( (s, (ky())))) - (Ma
(32a)
3 T(a,ls) x
Injexp A A (s, AMag(ssns) Σ exp (- A
E E
N
E Es-ox p3
T [a)]s)][A(() (s, (a))) - BMa
(32b)
N
3 T Ta, - [mmexp ( (3,,01071)55)) - AMEI(B
X
Eampp
Σ
+ In Σ exp (-
3 (B,i,)))
3
1=1
Es-p3
x (aggs)(AATss (s,a))().)) - BMa Σ
(32c)
22

---

Under review as a conference paper at ICLR 2022
N Σ
N
A
E
Taa 18)
In
Σ exp
(B.1,))))
(32d)
2=1
Therefore.
sup A 7 0535000
lim
(x) - F(3)
(33a)
N
N
2
BUB X
< lim
B Bg
T(azlo)
lim 2
Exm p
Σ xlayls)
X
j=1
1=1
3 E B
N
N
+ lim
sup
Σ
{(a)s))
In
exp
(3,
(33b)
> 05959um
j=1
1=1
N
> (-3-((3...)))
lim
sup
E Expp
T(ay(s) In
exp
(33c)
A X USBSBun
j=1
In addition, vB E [0, Punl and AY, (3, (.) is bounded since
< (3,
A (s, (15(,)) BM - [A (s, ai) - 3Mii
(34)
(A aggg (s, as)) (BBMA (8.5)) - 3Mab)
A (8T - (s, a:) + 1MM47 -BMbs
< 2 max A (8, a) + 9UB max Mij
8,a
< mm BuB max Mij i.j
(35)
Then, Therefore in 33c), the optimal B can be achieved.
Baupe Σn= (ayls)) In M exp (- 1- (B,
is bounded.
Let BA
arg 0580
Expoy En (aal()) In D> exp (- (B, and then
we have:
9 Banpy sup
N x la
lim
Σ
Σ exp ( ( ((3,,,,)))
(36a)
Aoc 05859BB
1=1
==1
8 Eamps
lim
Σ
T(asls) In
exp (- 37 30 (i,)))))))
(36b)
2 700
Define (j) = mino M. (B, i, Then since (B, 0 for $ (g,,) based on its definition, we have a. ()) > 0. On one hand, we have
N
X
lim In >
Σ
exp
(-
Bx 5BA ,i,))
(37a)
N
N
lim ln[ -
Σ
exp (- A
BA 3 (.))
A exp
( ((,)) (37b)
(=1lig177
kl (8x.))
N
N
X xx()))
< lim In[ A
Σ
exp
exp
(37c)
BUB
(=lligky((A.))
E (BxJ)
A
lim In[
Σ
exp
os(5))
[C (8x
(37d)
BUB
(
lim
(xx,
(37e)
23

---

Under review as a conference paper at ICLR 2022
On the other hand, we have
N
A
lim
In
exp
Ba
((5))
(38a)
exp( A
exp (- x
lim In
Σ
(B
E
3r (B*, (,,))) (38b)
P00
k (8x.5)
N
A
lim In t
Σ
exp(
BA (00
(38c)
N
lim In X
(0)
(38d)
lim In[] (Bx, +
(38e)
Therefore, limx
In E/ exp (- (B 1,5)i]
lim 1an->oon/KK3 [n[K/] (3x. f)ll Based that,
we have
8a
A
lim 7
B Bsnps Σ
x
{ayls) In Σ exp
3 c3
(3* (3*,,,)))]
(39a)
j=1
N
lim
ln
exp
B (B*, i,5,,,
(39b)
N
lim
In
exp
3 c (8*
(39c)
7
i=1
= lim
In[|K
(Bx,
(39d)
x
BUB
< lim t
Nln N
(39e)
A
which means lim tea SUP FX(B) F(3)
0.
Furthermore, since
lim suPogagau (3) F(3) 2 0 holds naturally, we have limx-+x suPos (F)
F(B) O. Therefore, (3) converges to F(3) uniformly on (0, BUBl. which also indicates
(3) epi-converges to F(3) on [0, SuBB (Royset 2018 Rockafellar & Wets 1998). By properties
lim arg min FX) arg min F(B)
of epi-convergence, we have that X DSASEE 1998).
Roockfellar & Wets
0SA5AB
G OPTIMAL BETA FOR A SPECIAL DISTANCE
Proposition 2. (1). If the initial point Bo is in ( (s, aT(s,a;)),+oo) +oo), the local optimal B solution is max.,j {AT i(A (s, an,)) ) A" (s))
(2). If the initial point 30 is in [0, min (s - as,,,))) if - f PT
л (ax, (s))ds < O, the local optimal (3 is min (A (8, (s,,a,)) - (8.ag)); otherwise, the local optimal B3 solution is O.
(3). If the initial point Bu is in (min,,j4k, {A as, - A (s, )) m (A ax.) AT (s, a;;)), we construct sets I and 13 as:
for s E S.j (1,2 N) N] Then, 5- - E Ejer (a, < 0, the local optimal B is min (s @k.) ( (s,a,)); otherwise, the local optimal 3 is max {A* (8, ax.)) (s,
if 9o 2 ( ox.) - AT az) then Add j 13 else Add j to 13.
24

---

Under review as a conference paper at ICLR 2022
Proof of Proposition (1). When 8 E 1 (s ok. ) (s (s)))))++20) A" A (s. all k I...N, N, we have (s, aj) 2 AT ax) 3 for all s S.j = =1...N,k= =,,=1.... 11...... N. Thus,
have
az) 2 (s, ak. 3 for all E E 15.j=1.... N. Since ( ok.) - 3 (5 3 for
j E (8".)). for all 8 ES = N. Therefore, can be reformulated as:
min ( [36 + E Egsop 820
j=1
Since 8 2 0, we have the local optimal [ = maxs, (A ox.) (s(s
(2). When 3 E [0, min {A as.)) A" (8, H, we have A az) S A ax.) - 3 for
all E j. = 1. ... N. J # kg- Thus hg K2 (8" j) for all E S.j ........ The inner part of (6) then is:
N Σ f=1,ft
36 + EoNpy {
(aa)ss) 3] (akk (s, ax,)}
N
N
= 3(6 Bspp
Σ x(azlo)) + Esupp
Σ
T(agls)A" (sk
j=1,
j=1
= B(8
/
(1( ग (ax, (s)) + Bsnps
Σ x(a,ls)AA7( (s.
JES
If Joes PT (s) - I (s))ds < 0, we have the local optimal B mins. jxt, {AA ax.) (s a.)). 5- - PS (6) - π x(a ds 2 O, we have the local optimal B =
(3). For an initial point B0 in (min,.jæ4.. {Ar (s. ax.) -s, aj)}. max, {A f (s, ax.) - (s. 0z) (,)
we construct partitions I and 13 of the set 21. .. } in the way described in Proposition 2 for all
A (s.aa) for E S.j I and/ B S A (s, as,, for E S.j E T2. Then the inner part of (6) can be reformulated
E S. Consider 3 in the neightborhood of $g. i.e., B 2 (s,
as:
B8 + E. Eanpy!
Σ JEN
T
+
Σ T(ag)s)(A -}}
JET
= B(8 Eanpt
Σ jET?
T(ay(s)) + E Exppp
Σ
(s, (sa
Σ <(aals) ak.)). jER
jEl!
< 0, we have the local optimal B = min.es.jer? { (A (s, ox.) (3 aj)}. If 8 Esnpy Ejer (ayls)) 2 O, we have the local optimal /
If Esnpt EjEr T(ayls)
max (A* ak.) (s),a)))
H IMPLEMENTATION DETAILS
Visitation Frequencies Estimation: The unnormalized discounted visitation frequencies are needed to compute the global optimal 8*. At the k-th iteration, the visitation frequencies PΣ are estimated using samples of the trajectory set Da. Specifically, we first initialize (5(() = 0, Vs E S. Then for +
each timestep in each trajectory from DR, we update pf as p% (st) + (st) ( //Dul-
The general approach depicted in Algorithm [ allows various policy representations including arrays and neural networks. Let S₁ C represent a subset of states to perform the policy update at the k-th iteration. When array is used, the policy update step is simply Th+i (+1s) = (TT)))))s) Vs E Sc. When a neural network is employed, the policy update step can (F(ma)(+)s) T+(s))))²
Policy Representation:
be achieved by obtaining the gradient descent, i.c.,
SES
Policy Updating Strategy:
State space: For environments with a discrete state space (c.g., tabular domains), the WPO/SPO policy update is performed for all states at each iteration. For the environments
25

---

Under review as a conference paper at ICLR 2022
with continuous state space, random subset of states is sampled each iteration to perform the policy update.
Action space: For environments with a continuous action space, we first discretize the action space following TTan & Agrawal (2020) and then WPO/SPO policy update is performed on the discretized action space.
Hyperparamaters and Additional Results:
Our main experimental results are reported in Section In addition, we provide the setting of hyperparamaters and network sizes of our WPO/SPO algorithms in Table B And we present the numerical results of the final performance comparison among our algorithms and the baseline methods (i.e., TRPO, PPO, A2C) in Table
Table 3: Hyperparamaters and network sizes
Taxi-v3, NChain-v0
CartPole-vl
Acrobot-v|
CliffWalking-vo
0.9
0.95
0.95
Irx
10-2
5 x 3
Traluo
10
10
5 x 10
D
60 (Taxi); (Chain);
2
3
3 (CliftWalking)
size
2D array
64, 64
64, 64
Q/v size
[10, 7,5]
64. 64
664, 64
5, 50, 10
10
Table 4: Averaged rewards last 10% episodes during the training process
Environment
WPO
SPO
TRPO
PPO
A2C
Taxi-v3
45 + 27
-8 + 11
-202 +
381 + 34
-338 + 30
NChain-vO
3549 +: 197
3432 + 131
3522 + 258
3506 + 237
1606 + 10
CliffWalking-vO
-35 15
-25
-159 + 94
-3290 2106
55587 + 1942
CartPole-vl
388 54
370 30
297 65
193 45
Acrohot-v1
-162 +
-185 + 15
248 + 33
100 +
-379 39
Additional Experiments on Ablation Study:
CartPole-v1
Chain
500
WPO:
3500
Average Return
400
WPO: Optimal constant
WPO: Optimal
Average Return
3000
WPO: Decay
300
2500
200
2000
SPO
100
1500
SPO (0002)
SPO
SPO
logt)
0
1000
175 200
0.0
0.1
0.2
0.4
0.5
Episode
Timesteps
x x10
Figure 7: Episode rewards during the training process for different 3 and À settings, averaged across
3 runs with random initialization. The shaded area depicts the mean + the standard deviation.
26


Reviews:
[{"review_id": "oR6-Y-y1gp_", "reviewer": "Reviewer_2JvQ", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "Review of EFFICIENT WASSERSTEIN AND SINKHORN POLICY OPTIMIZATION\n\nIn general, I found the paper to be a nice, if modest, extension of TRPO. In essence, the key idea is trading KL regularization for Wasserstein or Sinkhorn. Experiments show this is, in general, a win, particularly for the case of approximated advantage functions. I found the exposition to be mostly fine with some minor critiques. First, it would be better to have a more detailed treatment of the prior work (TRPO is properly introduced later). Second, separately discussing Wasserstein and Sinkhorn is unnecessary and redundant. They are very closely related and the text (and math) would benefit from making that clearer. \n\n\nAdditional points: \n- The discussion of related work is inadequate. \n- It would be desirable to make a clear statement of what the reasons for and consequences of assumption 1 are. \n- \"Compared to the performance bound when using KL-based trust region\n(see, e.g., Schulman et al. (2015); Cen et al. (2020)), using the Wasserstein metric yields a tighter performance improvement bound and is more robust to the choice of parameters βt.\" It would be nice to show this explicitly. \n- Theorems 1 and 3 are closely related, as are Wasserstein and Sinkhorn. Is is really necessary to break them into separate sections? This choice makes the relation less clear. ", "summary_of_the_paper": "The paper introduces Wasserstein and Sinkhorn policy optimization, with three main contributions. First, the paper derives closed form expressions for updates using Lagrange multipliers. Second, it proves monotonicity of performance improvement. Third, experiments show efficiency and effectiveness relative to TRPO and PPO. \n", "summary_of_the_review": "Overall, this is not the most ambitious paper, but generally a nice contribution.  \n", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "RyT3AUY7MwR", "reviewer": "Reviewer_A8WK", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "The paper is overall well written, the results are well organized and reasonably clear to follow. In terms of novelty, I believe the closed form update parts (Theorem 1 & 3) are reasonably standard primal-dual arguments. However, Theorem 2 is an interesting and significant contribution since it provides theoretical guarantees even when the update $\\beta_t$ is not the optimal solution in (6), which circumvent the difficulty in obtaining closed form solution of (6). Theorem 4 also seems important as it gives theoretical guarantee for us to approximate WPO by SPO with large $\\lambda$. I would like to discuss a few more questions with the authors: \n\nThe upper bound of $\\beta^*_{\\lambda}$ for SPO in Theorem 3 is of order $1/\\delta$, which seems not very good because the bound would be trivial if there is no perturbation to the policy (i.e., $\\delta = 0$). Also, the bound for $\\beta^*$ for WPO in Theorem 1 seems to be independent of $\\lambda$, I would like the authors to discuss some insights behind the different dependency on $\\lambda$. Do we expect the bound for $\\beta^*_{\\lambda}$ is actually also independent of $\\delta$. Moreover, it would be also interesting to understand the optimal dependency of $\\beta^*_{\\lambda}$ as a function of $\\lambda$.\nIn Theorem 4, can we characterize the rate of convergence of $|F_{\\lambda}(\\beta) - F(\\beta)|$? Seems like an upper bound could be $1/\\lambda$, but is this correct (or optimal if it is correct)? Also, it would be good if we could understand the convergence of the optimizers in part 2 of Theorem 4.\n\nSmall typo in Abstract: ‘extensions of policy optimziation’ -> ‘extensions of policy optimization’.\n", "summary_of_the_paper": "This paper studies policy optimization in reinforcement learning with Wasserstein and Sinkhorn trust regions. Compare to the standard TRPO which based on KL-divergence, the proposed WPO and SPO go beyond the parametric policy distribution class. The authors also derive closed form policy update, as well as theoretical performance guarantees for both problems.", "summary_of_the_review": "This paper presents two policy update frameworks (WPO and SPO) which relax the restriction to parametric policy distribution in the standard setting. Theoretical guarantees are provided. The numerical results also suggest that proposed policy optimization methods outperform the standard TRPO and PPO with better performance and faster convergence. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "NXZhzF17sws", "reviewer": "Reviewer_i87P", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "Strength:\n\n1. The paper provides a relatively complete analysis of the Wasserstein and Sinkhorn policy optimization, including closed-form policy updates, performance improvement bound, actor-critic algorithm, and experiments on popular instances.\n\nWeakness:\n\n1. The theoretical advances seem quite limited. The closed-form policy update in Theorem 1 and Theorem 3 should follow from existing duality results on Wasserstein and Sinkhorn optimization (for example, [1] and [2]). It is unfortunate that such connections are not recognized in the paper. Theorem 2 is a simple consequence of the well-known performance difference lemma. Theorem 4 is sort of expected, as many existing results on the relationship between Wasserstein distance and Sinkhorn distance. In addition, it is unclear whether SPO has a performance improvement at each iteration for any fixed entropic regularization parameter.\n\n2. Computational-wise, I am wondering how the policy update in (WPO) and (SPO) can be implemented when the state/action space is large, and the computation of the parameters $\\beta_t$ seems also time-consuming whereas I did not find a technically sound argument for the heuristic choices of $\\beta$. \n\n3. The experiment results are not convincing enough to demonstrate clear advantages of WPO or SPO. \n\n(1) There is no comparison with existing Wasserstein policy optimization (with parametric classes), such as algorithms in Moskovitz et al. (2020) and Pacchinao et al. (2020). Therefore, it is unclear to me whether it is worth considering a non-parametric policy update, given that is the high computational cost for large spaces.\n\n(2) In most results (for example, Figures 2,3,5), SPO is outperformed by WPO in terms of the reward at the end of the timesteps, although it converges to a suboptimal value faster. It might be helpful to consider a varying entropic regularization parameter.\n\nReferences.\n[1] Blanchet, Jose, and Karthyek Murthy. \"Quantifying distributional model risk via optimal transport.\" Mathematics of Operations Research 44.2 (2019): 565-600.\n[2] Wang, Jie, Rui Gao, and Yao Xie. \"Sinkhorn Distributionally Robust Optimization.\" arXiv preprint arXiv:2109.11926 (2021).", "summary_of_the_paper": "This paper considers Wasserstein and Sinkhorn trust region policy optimization (WPO and SPO) for reinforcement learning. Unlike existing works on WPO, it does not assume a parametric policy class. Theoretically, it shows the performance improvement of WPO at every iteration, and that SPO converges to WPO. It also conducts some experiements to illustrate the advantages of the proposed methods.", "summary_of_the_review": "The main feature of this paper is considering a non-parametric policy class for Wasserstein policy optimization. However, the theoretical contribution adds marginal to the existing literature, and the numerical findings are not convincing to demonstrate the advantages of the new framework.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "EGfqdxZRXqu", "reviewer": "Reviewer_PvnA", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "The paper is well-organised and present a new way of incorporating the Wasserstein distance within policy optimisation algorithms. \nThe theoretical results look correct at first glance although I admit didn't check them carefully in the Appendix.\n\nIt would be nice for theorem 1 to be self-contained: some of the variables are defined earlier in the text like M or beta but it would ease the reading to define them in the theorem. The presentation of the theoretical results are a bit hard to follow so adding a few explanatory sentences about the importance of each terms would also be helpful.\n\nCould the authors also comment on the computational complexity of the method?\n\nIn terms of experiments, the domains considered seem to show the benefits of the method. I believe stronger tasks would make the paper stronger, in particular continuous control ones or at least showing results for a few more domains.\n\nRelated Work: \"Wasserstein-like metrics have only been recently studied in the context of reinforcement learning.\" Please note that the Wasserstein metric has been used in RL since at least 2012 with the introduction of bisimulation metrics https://arxiv.org/pdf/1207.4114.pdf, and recent work havent only used it for imitation learning but also for generalization https://arxiv.org/pdf/2101.05265.pdf https://arxiv.org/abs/2102.01514 https://arxiv.org/abs/2006.10742.\nI believe the Wasserstein and Sinkhorn metrics have also been used in the distributional RL literature so it might be nice to discuss the similarity of ideas in both areas.\n\nTypos:\nWassersteim \n*the* Wasserstein metric, *etc.* : please dont use \"etc\" and add all necessary details.", "summary_of_the_paper": "This paper proposes to use two extensions of the TRPO algorithm relying on the Wasserstein distance and the Sinkhorn divergence which dont require to explicitly specify a distribution for the policy. The authors provide a theoretical analysis giving a closed form policy update for their two methods and a performance improvement bound in the case of Wasserstein policy optimisation. They evaluate their methods empirically on tabular domains (Taxi, Chain and Cliff Walking) and on some discrete locomotion tasks (Cartpole, Acrobot). They find that their method outperform TRPO and PPO while being more sample-efficient, and converging faster.", "summary_of_the_review": "This paper provides a new way of optimising the policy distribution relying on the Wasserstein and Sinkhorn distances. The methods is theoretically grounded so I would recommend an accept but I believe the authors would need to evaluate their method on more domains to make the paper stronger.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposes two extensions of the TRPO algorithm in which the trust region is defined using the Wasserstein distance and the Sinkhorn divergence. The proposed methods do not restrict the policy to belong to a parametric distribution class and the authors provide 
closed-form policy updates and a performance improvement bound for the Wasserstein policy optimization.
The authors provide an empirical evaluation of their approaches on tabular domains and some discrete locomotion tasks, comparing the performance with some state-of-the-art policy optimization approaches.

After reading the authors' feedback and interacting with the authors, the reviewers did not reach a consensus: one of the reviewers votes for rejection, while the other three reviewers are slightly positive.
In particular, the reviewer that voted for rejection raised a number of concerns that have been discussed at length with the authors, who were able to clarify some of the issues, but some of the answers did not satisfy the reviewer.
I went through the paper and I found the paper solid from a technical point of view, but I share some of the reviewers' concerns and I think that the authors should better position their contribution with respect to the state of the art. 
Overall, this paper is borderline and I feel it needs still some work to deserve clear acceptance (which I think will be soon).
Decision: Reject