# EFFICIENT WASSERSTEIN AND SINKHORN POLICY OP## TIMIZATION

**Anonymous authors**
Paper under double-blind review

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 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. 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.

1 INTRODUCTION

Policy-based reinforcement learning (RL) approaches have received remarkable success in many
domains, including video games (Mnih et al., 2013; Mnih et al., 2015), board games (Silver et al.,
2016; Heinrich & Silver, 2016), robotics (Grudic et al., 2003; Gu et al., 2017), and continuous control
tasks (Duan et al., 2016; Schulman et al., 2016). One prominent example is policy gradient method
(Grudic et al., 2003; Peters & Schaal, 2006; Lillicrap 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 a probability
distribution πθ(a _s) = P_ [a _s; θ], such that the action a in state s is chosen stochastically following_
_|_ _|_
the policy πθ controlled by parameter θ. Determining the right step size to update the policy is 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.

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
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 gauge the similarity between policies could bring additional benefits. Recently,
few work (Richemond & Maginnis, 2017; Zhang et al., 2018; Moskovitz et al., 2020; Pacchiano
et al., 2020) has explored the Wasserstein metric to restrict the deviation between consecutive policies.
Compared with KL divergence, the Wasserstein metric has several desirable properties. Firstly, it is a
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
has a weaker topology (Arjovsky et al., 2017), thus possibly leading to a 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 to different techniques to approximate the policy update under Wasserstein regularization. For example,
Richemond & Maginnis (2017) solved the resulting RL problem using Fokker-Planck equations;


-----

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-guided policy optimization with smooth
Wasserstein regularization by solving an approximate dual reformulation defined on reproducing
kernel Hilbert spaces. Aside from such approximation, some of these work also limits the policy
representation to 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 in local movements in the action space and thus leads to convergence to a sub-optimal solution.
Henceforth, the theoretical performance of policy optimization under the Wasserstein metric remains
elusive in 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 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 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. In 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 metrics 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 et al. (2020); Pacchiano et 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 RL, to our best knowledge, is 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-like metrics are also pervasively studied in distributionally robust optimization (DRO);
see e.g., the survey by Kuhn et al. (2019) 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 Wang et al. (2021a) studied DRO
using the Sinkhorn distance.


-----

BACKGROUND AND NOTATIONS


**Markov Decision Process (MDP): We consider an infinite-horizon discounted MDP, defined by the**
tuple (S, A, P, r, ρ0, γ), where S is the state space, A is the action space, P : S × A × S −→ R is the
transition probability, r : S × A −→ R is the reward function, ρ0 : S −→ R is the distribution of the
initial state s0, and γ is the discount factor. We define the return of timestep t as the accumulated
discounted reward from t, Rt = _k=0_ _[γ][k][r][(][s][t][+][k][, a][t][+][k][)][, and the performance of a stochastic policy]_
_π as J(π) = Es0,a0,s1...[[P][∞]t=0_ _[γ][t][r][(][s][t][, a][t][)]][ where][ a][t][ ∼]_ _[π][(][a][t][|][s][t][)][,][ s][t][+1][ ∼]_ _[P]_ [(][s][t][+1][|][s][t][, a][t][)][. As shown]
in Kakade & Langford (2002), the expected return of a new policy π[′] can be expressed in terms

[P][∞]
of the advantage over the old policy π: J(π[′]) = J(π) + Es∼ρπυ[′] _[,a][∼][π][′]_ [[][A][π][(][s, a][)]][, where][ A][π][(][s, a][) =]
E[Rt|st = s, at = a; π] − E[Rt|st = s; π] represents the advantage function and ρ[π]υ [represents]
the unnormalized discounted visitation frequencies with initial state distribution υ, i.e., ρ[π]υ [(][s][) =]
Es0∼υ[[P][∞]t=0 _[γ][t][P]_ [(][s][t][ =][ s][|][s][0][)]][.]

**Trust Region Policy Optimization (TRPO): In TRPO (Schulman et al., 2015), the policy π is**
parameterized as πθ with parameter vector θ. For notation brevity, we use θ to represent the policy
_πθ. Then, the new policy θ[′]_ is found in each iteration to maximize the expected improvement
_J(π[′]) −_ _J(π), or equivalently, the expected value of the advantage function:_

maxθ[′] Es∼ρθυ[,a][∼][θ][′] [[][A][θ][(][s, a][)]]

(1)
s.t. Es _ρθυ_ [[][d][KL][(][θ][′][, θ][)]][ ≤] _[δ,]_
_∼_

where dKL represents the KL divergence and δ 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 ρ[θ]υ [is used rather than][ ρ]υ[θ][′][, 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 π[′] on the discrete
action space = _a1, a2, . . ., aN_, the Wasserstein distance between the policies is defined as:
_A_ _{_ _}_

_dW (π[′], π) =_ inf (2)
_Q_ Π(π[′],π)[⟨][Q, M] _[⟩][,]_
_∈_

where ⟨·, ·⟩ denotes the Frobenius inner product. The infimum is taken over all joint distributions Q
with marginals π[′] and π, and M is the cost matrix with elements Mij = d(ai, aj), where d(ai, aj) is
defined as the distance between actions ai and aj.

**Sinkhorn Divergence: Sinkhorn divergence (Cuturi, 2013) provides a smooth approximation of**
the Wasserstein distance by adding an entropic regularizer. The Sinkhorn divergence is defined as
follows:

_dS(π[′], π_ _λ) =_ inf (3)
_|_ _Q_ Π(π[′],π)[{⟨][Q, M] _[⟩−]_ _λ[1]_ _[h][(][Q][)][}][,]_
_∈_


where h(Q) = − [P]i[N]=1 _Nj=1_ _[Q][ij][ log][ Q][ij][ represents the entropy term, and][ λ >][ 0][ is a regularization]_

parameter. Similarly, we useP _Q[s]_ to denote the joint distribution of π(·|s) and π[′](·|s) with _i=1_ _[Q]ij[s]_ [=]
_π(aj|s) and_ _j=1_ _[Q]ij[s]_ [=][ π][′][(][a][i][|][s][)][. 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[P][N]
with entropy, a smoother and denser coupling between distributions can be achieved (Courty et al.,

[P][N]
2014; 2016). Therefore, when the weight of the entropic regularization decreases (i.e., λ increases),
the sparsity of the divergence increases, and the Sinkhorn divergence converges to the Wasserstein
metric, i.e., limλ _dS(π[′], π_ _λ) = dW (π[′], π). 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-Bianchi
& Lugosi, 2006), model selection (Juditsky et al., 2008; Rigollet & Tsybakov, 2011), 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

Motivated by TRPO, here we consider a 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 π[′] 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 Es _ρπυ_ _[,a][∼][π][′][(][·|][s][)][[][A][π][(][s, a][)]]_
_π[′]_ _∼_
_∈D_ (4)

where = _π[′]_ Es _ρπυ_ [[][d][W][ (][π][′][(][·|][s][)][, π][(][·|][s][))]][ ≤] _[δ][}][,]_
_D_ _{_ _|_ _∼_


where the Wasserstein distance dW (·, ·) is defined in (2).

In most practical cases, the reward r is bounded and correspondingly, the accumulated discounted
reward Rt is bounded. So without loss of generality, we can make the following assumption:

**Assumption 1. Assume A[π](s, a) is bounded, i.e., supa∈A,s∈S |A[π](s, a)| ≤** _A[max]_ _for some A[max]_ _> 0._

With Wasserstein metric based trust region constraint, we are able to derive the closed-form of the
policy update shown in Theorem 1. The main idea is to form the Lagrangian of the constrained
optimization problem presented above, and the detailed proof can be found in Appendix A.

**Theorem 1. (Closed-form policy update) Let ks[π][(][β, j][) =][ argmax]k=1...N** _[{][A][π][(][s, a][k][)][ −]_ _[βM][kj][}][,]_
_where M denotes the cost matrix. If Assumption 1 holds, then an optimal solution to the WPO_
_problem in (4) is given by:_


_π[∗](ai_ _s) =_
_|_


_π(aj|s)fs[∗][(][i, j][)][,]_ (5)
_j=1_

X


_where fs[∗][(][i, j][) = 1][ if][ i][ =][ k]s[π][(][β][∗][, j][)][ and][ f][ ∗]s_ [(][i, j][) = 0][ otherwise, and][ β][∗] _[is an optimal Lagrangian]_
_multipler corresponds to the following dual formulation:_

_N_

minβ 0 _[F]_ [(][β][) = min]β 0  _υ_ _π(aj|s)[A[π](s, aksπ[(][β,j][)][)][ −]_ _[βM][k]s[π][(][β,j][)][j][]]_ (6)
_≥_ _≥_ _j=1_

 X 

_Moreover, we have β[∗]_ _β := max[βδ][ +]s[ E][s],k,j[∼][ρ][π]=1...N,k=j (Mkj)[−][1](A[π](s, ak)_ _A[π](s, aj))._ _[.]_
_≤_ [¯] _∈S_ _̸_ _−_

The exact policy update for WPO in (5) requires computing the optimal Lagrangian multiplier β[∗] by
solving the one-dimensional subproblem (6). A closed form of β[∗] is not easy to obtain in general,
except for special cases of the distance d(x, y) or cost matrix M . In Appendix G, we provide the
closed form of β[∗] for the case when d(x, y) = 0 if x = y and 1 otherwise.

**WPO Policy Update:** Based on Theorem 1, we introduce the following WPO policy updating rule:


_πt+1(ai|s) = F[WPO](πt) :=_


_πt(aj|s)fs[t][(][i, j][)][,]_ (WPO)
_j=1_

X


where we choose an arbitrary ks[π][t] [(][β][t][, j][)] _∈_ argmaxk=1,...,N _{A[π][t]_ (s, ak) − _βtMkj} and set_
_fs[t][(][k]s[π][t]_ [(][β][t][, j][)][, j][) = 1][ and other entries to be][ 0][.]

Note that different from (5), we allow βt to be chosen arbitrarily and time dependently. We show that
such policy update always leads to a monotonic improvement of the performance even when βt is not
the optimal Lagrangian multiplier. In particular, we propose two efficient strategies to update the
multiplier βt:

(i) Time-dependent βt: To improve the computational efficiency, we can simply treat βt as a
time-dependent control parameter, e.g., we can choose βt to be a diminishing sequence.

(ii) Approximation of optimal βt: To improve the convergence, we can approximately solve the
optimal Lagrangian multiplier based on Sinkhorn divergence. We will discuss this in more
detail in Section 4.


-----

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 can only be evaluated inexactly with limited samples, then an 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 βt ≥** 0, if
_A[ˆ][π]_ _A[π]_ _ϵ for some ϵ > 0, the WPO policy update with the inaccurate advantage function_
_||_ _−_ _||∞_ _≤_
_Aˆ[π], guarantees the following performance improvement bound,_

_N_ 2ϵ
_J(πt+1) ≥_ _J(πt) + βtEs∼ρπtµ_ +1 _j=1[π][t][(][a][j][|][s][)][M]k[ˆ]s[πt]_ [(][β]t[,j][)][j][ −] 1 − _γ [.]_ (7)
X

Note that when the estimation error ϵ = 0, we have a monotonic performance improvement
_J(πt+1)_ _J(πt) for any βt_ 0. If the second term in (7) is non-zero, then we have a strict
monotonic improvement. Compared to the performance bound when using KL-based trust region ≥ _≥_
_J(πt+1) ≥_ _J(πt) −_ 12−ϵγ [(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 βt.

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 a 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 1 holds, then the optimal solution to the trust region constrained problem**
_(4) with Sinkhorn divergence is:_

_N_ exp ( _β[λ]λ[∗]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]_

_πλ[∗][(][a][i][|][s][) =]_ _N_ _π(aj_ _s),_ (8)

_|_

_j=1_ _k=1_ [exp (][ λ]βλ[∗] _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X

_where M denotes the cost matrix andP βλ[∗]_ _[is an optimal solution to the following dual formulation:]_

_N_
minβ≥0 Fλ(β) = minβ≥0 _βδ −_ Es∼ρπυ _j=1_ _[π][(][a][j][|][s][)(][ β]λ_ [+][ β]λ [ln(][π][(][a][j][|][s][))][ −]

_β_ n PN _N_ _β_ exp ( _β[λ]_ _[A][π][(][s,a][i][)][−][λM][ij]_ [)][·][π][(][a][j] _[|][s][)]_
_λ_ [ln[][P]i[N]=1 [exp (][ λ]β _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]) +][ E][s][∼][ρ]υ[π]_ _i=1_ _j=1_ _λ_ _Nk=1_ [exp (][ λ]β _[A][π][(][s,a][k][)][−][λM][kj]_ [)] _. (9)_

P P P o

_Moreover, we have βλ[∗]_ _δ_ _._

_[≤]_ [2][A][max]

In contrast to the Wasserstein dual formulation (6), the objective in the Sinkhorn dual formulation (9)
is differentiable in β and admits closed-form gradients (shown in Appendix E). With this gradient
information, we can use gradient-based global optimization algorithms (Wales & Doye, 1998; Zhan
et al., 2006; Leary, 2000) to find a global optimal solution βλ[∗] [to (9).]

Next, we show that if the entropic regularization parameter λ is large enough, then the optimal
solution βλ[∗] [is a close approximation to the optimal solution][ β][∗] [to the Wasserstein dual formulation.]
The detailed proof is provided in Appendix F.


**Theorem 4. Define βUB = max** [2][A]δ[max]
_{_


_,_ _β[¯]}. The following holds:_


_1. Fλ(β) converges to F_ (β) uniformly on [0, βUB],

_2._ _λlim−→∞_ 0[arg min]≤β≤βUB _Fλ(β) ⊆_ 0arg min≤β≤βUB _F_ (β).

Although it is difficult to obtain the exact value of the optimal solution β[∗] to the Wasserstein dual
formulation (6), the above theorem suggests that we can approximate β[∗] via βλ[∗] [by setting up a]
relative large λ. In practice, we can also adopt a smooth homotopy approach by setting an increasing
sequence λt for each iteration and letting λt →∞.


-----

**SPO Policy Update:** Based on Theorem 3, we introduce the following SPO policy updating rule:

_N_ exp ( _[λ]βt[t]_ _[A][π][t]_ [(][s, a][i][)][ −] _[λ][t][M][ij][)]_

_πt+1(ai_ _s) = F[SPO](πt) :=_ _N_ _πt(aj_ _s),_ (SPO)
_|_ _|_

_j=1_ _k=1_ [exp (][ λ]βt[t] _[A][π][t]_ [(][s, a][k][)][ −] _[λ][t][M][kj][)]_

X

Heresolving the one-dimensional subproblem (9) or simply set as a diminishing sequence. λt ≥ 0 and βt ≥ 0 are some control parameters. The parameterP _βt can be either computed via_

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 1, that
combines WPO/SPO with advantage function estimation.

In each iteration of Algorithm 1, the first step
is to collect trajectories, which can be either

**Algorithm 1: On-policy WPO/SPO algorithm**

complete or partial. The difference is whether
the return is considered thoroughly to the end Input: number of iterations K, learning rate α
of a planning horizon. If the trajectory is com- Initialize policy π0 and value network Vψ0 with
plete, the total return can be directly expressed random parameter ψ0
as the accumulated discounted rewards Rt = **for k = 0, 1, 2 . . . K do**
_Tk=0 −t−1_ _γ[k]rt+k. If the trajectory is partial,_ Collect a set of trajectories Dk on policy πk

For each timestep t in each trajectory,

it can be estimated by applying the multiPstep temporal difference (TD) methods (De compute total returns Gt and estimate
Asis et al., 2017): _R[ˆ]t:t+n =_ _k=0_ _[γ][k][r][t][+][k][ +]_ advantages _A[ˆ][π]t_ _[k]_

Update value:

_γ[n]V (st+n). Then for the advantage estima-_
tion, we can use Monte Carlo advantage esti- _ψk+1_ _ψk_ _α_ _ψk_ (Gt _Vψk_ (st))[2]

[P][n][−][1] _←−_ _−_ _∇_ _−_
mation, i.e., _A[ˆ][π]t_ _[k]_ = Rt _Vψk_ (st) or General- Update policy: X
ized Advantage Estimation (GAE) (Schulman − _πk+1_ F(πk) via WPO or SPO with _A[ˆ][π]t_ _[k]_
et al., 2016), which provides a more explicit **end** _←−_
control over the bias-variance trade-off. In the
value update step, we use a neural net to represent the value function, where ψ is the parameter that specifies the value net s → _V (s). Then, we_
can update ψ 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 1. We compare the performance of our proposed methods
with several benchmarks, including TRPO (Schulman et al., 2015), PPO (Schulman et al., 2017), and
A2C (Mnih et al., 2016)[1]. We compare with A2C because it 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 is 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 βt. We
test four settings of β value for WPO policy update: (1) Setting 1: computing optimal β value for all
policy update; (2) Setting 2: computing optimal β value for first 20% of policy updates and decaying
_β for the remaining; (3) Setting 3: computing optimal β value for first 20% of policy updates and fix_
_β as its last updated value for the remaining; (4) Setting 4: decaying β for all policy updates (e.g.,_
_βt =_ _t[1][2][ ). In particular, Setting 3 is rooted in the observation that][ β][∗]_ [does not change significantly]

1We use the implementations of TRPO, PPO and A2C from OpenAI Baselines (Dhariwal et al., 2017) for
MuJuCo tasks and Stable Baselines (Hill et al., 2018) for other tasks.


-----

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 1 and Table 1 respectively. Figure 1 and Table 1 clearly Table 1: Run time comparison
indicate a tradeoff between computation efficiency and accuracy in for different β settings
terms of different choices of β value. Setting 2 is the most effective
way to balance the tradeoff between performance and run time. For Runtime Taxi (s) CartPole (s)
the rest of experiments, we adopt this setting for both WPO and Setting 1 1224 130
SPO. We also compare WPO with SPO under different constant Setting 2 648 63
and time varying λ values on the Taxi task. As shown in Figure 1, Setting 3 630 67
SPO converges faster than WPO. With more weight on the entropic Setting 4 522 44
regularization of Sinkhorn divergence (i.e., smaller λ), SPO can
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 to Wasserstein distance shown in Theorem 4. Therefore, the
choice of λ can effectively adjust the trade-off between convergence and final performance. With
a proper λ choice, SPO is able to attain a faster convergence speed with an optimum that is only
slightly lower than WPO.

More experiments for ablation study is conducted on the Chain (Dearden et al., 1998) and CartPole
(Barto et al., 1983) tasks. Results are provided in Appendix H.

Figure 1: Episode rewards during the training process for the Taxi task with different β 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), Table 2: Trained agents perforwhich are intentionally designed to test the exploration ability of the mance on Taxi (averaged over
algorithms. The tabular domain has a special environment structure 1000 episodes)
with a discrete state space and a discrete action space. Thus, we use
value function, we use a neural net to smoothly update the values.an array of size |S| × |A| to represent the policy π(a|s). For the Success (+20) WPO0.753 TRPO0
The performance of WPO and SPO are compared to the performance Fail (-10) 0.232 0
of TRPO, PPO and A2C under the same neural net structure. Each Timesteps (-1) 70.891 200
algorithm is evaluated 5 times with a random initialization. Results Avg Return -58.151 -200
are reported in Table 2 and Figure 2. The setting of hyperparamaters
and network sizes of our algorithms and additional results are provided in Appendix H.

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 on the Taxi environment. As shown in Table 2, WPO has a higher successful drop-off


-----

rate and a lower task completion time while the original TRPO reaches the time limit with a drop-off
rate 0. Therefore, the results in Taxi show that WPO finds a better policy than the original TRPO.

Figure 2: Episode rewards during the training process for the tabular domain tasks, averaged across 5
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 1 framework with
KL based policy update derived in Peng et 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 in 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.

(a) NA = 100 (b) NA = 250 (c) NA = 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 s as an input and outputs the categorical distribution of π(a|s). 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 5 times with a random initialization.

**Final Performance: Figure 4 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.

**Training Time: To train 10[5]** 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 as TRPO and PPO.

The performances of WPO and KL are also compared for the discrete locomotion tasks under
different NA. As shown in Figure 5, when NA is 500, KL performs better than WPO. However, when
_NA decreases to 100, WPO significantly outperforms KL. These results indicate that for discrete_
locomotion tasks, WPO is more robust than KL when advantage values are inaccurate.


-----

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.

(a) NA = 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 to environments with a continuous action space
by discretizing the action space following Tang & Agrawal (2020). For comparison, we additionally
consider Behavior Guided Policy Gradient (BGPG) algorithm from Pacchiano et al. (2020). Similar
results are observed in Figure 6 as the discrete action tasks: WPO and SPO outperform the benchmark
algorithms in terms of final performance.

Figure 6: Episode rewards during the training process for continuous action space tasks, averaged
across 3 runs with a random initialization. The shaded area depicts the mean ± the standard deviation.

7 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 β 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-policy frameworks.


-----

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.06920,
2018.

Rishabh Agarwal, Marlos C Machado, Pablo Samuel Castro, and Marc G Bellemare. Contrastive
behavioral similarity embeddings for generalization in reinforcement learning. In Proceedings of
_the 8th International Conference on Learning Representations, 2020._

Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein GAN. ArXiv Preprint, pp.
arXiv:1701.07875, 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):834–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–10076, 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 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, 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.01327,
2017.

Richard Dearden, Nir Friedman, and Stuart Russell. Bayesian Q-learning. In Proceedings of the
_Fifteenth National Conference on 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. OpenAI baselines. https:](https://github.com/openai/baselines)
[//github.com/openai/baselines, 2017.](https://github.com/openai/baselines)

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. 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.


-----

Aude Genevay, Gabriel Peyré, and Marco Cuturi. Learning generative models with Sinkhorn
divergences. In International Conference on Artificial Intelligence and Statistics, pp. 1608–1617,
2018.

Alborz Geramifard, Christoph Dann, Robert H. Klein, William Dabney, and Jonathan P. How. RLPy:
A value-function-based reinforcement learning framework for education and research. Journal of
_Machine Learning Research, 16(46):1573–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/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–3396, 2017._

Johannes Heinrich and David Silver. Deep reinforcement learning from self-play in imperfectinformation games. ArXiv Preprint, pp. arXiv:1603.01121, 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. https://github.com/](https://github.com/hill-a/stable-baselines)
[hill-a/stable-baselines, 2018.](https://github.com/hill-a/stable-baselines)

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–4455, 2021._

Anatoli Juditsky, Philippe Rigollet, and Alexandre B Tsybakov. Learning by mirror averaging. The
_Annals of Statistics, 36(5):2183–2206, 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 Shafieezadeh-Abadeh.
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.

Robert Leary. Global optimization on funneling landscapes. Journal of Global Optimization, 18:
367–383, 2000.

Timothy P. 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 I 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. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Petersen,
Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra,
Shane Legg, and Demis Hassabis. Human-level control through deep reinforcement learning.
_Nature, 518(7540):529–533, 2015._


-----

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. ArXiv Preprint, pp. arXiv:2010.05380, 2020.

Aldo Pacchiano, Jack Parker-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 van 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. arXiv:1910.00177,
2019.

Jan Peters and Stefan Schaal. Policy gradient methods for robotics. In Proceedings of the 2006
_IEEE/RSJ 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 Preprint, pp. arXiv:1709.08894, 2017._

Pierre H. Richemond and Brendan Maginnis. On Wasserstein reinforcement learning and the FokkerPlanck equation. ArXiv Preprint, pp. arXiv:1712.07185, 2017.

Philippe Rigollet and Alexandre Tsybakov. Exponential screening and optimal rates of sparse
estimation. The Annals of Statistics, 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. 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 of the 32nd International Conference on Machine Learning,
pp. 1889–1897, 2015.

John Schulman, Philipp Moritz, Sergey Levine, Michael I. Jordan, and Pieter Abbeel. Highdimensional 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.

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 van den Driessche, Julian Schrittwieser, Ioannis Antonoglou, Vedavyas Panneershelvam, Marc 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 tree search. Nature, 529(7587):484–489, 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._


-----

Yunhao Tang and Shipra Agrawal. Discretizing continuous action space for on-policy optimization.
In Proceedings of the Thirty-Fourth 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 to 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.
arXiv:2109.11926, 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, 2021b._

Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement
learning. Machine Learning, 8(3–4):229–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.08113, 2019.

Lixin Zhan, Jeff Chen, and Wing-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 as Wasserstein
gradient flows. In Proceedings of the 35th International Conference on Machine Learning, pp.
5737–5746, 2018.


-----

A PROOF OF THEOREM 1

**Theorem 1. (Closed-form policy update) Let ks[π][(][β, j][) =][ argmax]k=1...N** _[{][A][π][(][s, a][k][)][ −]_ _[βM][kj][}][,]_
_where M denotes the cost matrix. If Assumption 1 holds, then an optimal solution to the WPO_
_problem in (4) is given by:_


_π[∗](ai_ _s) =_
_|_


_π(aj|s)fs[∗][(][i, j][)][,]_ (5)
_j=1_

X


_where fs[∗][(][i, j][) = 1][ if][ i][ =][ k]s[π][(][β][∗][, j][)][ and][ f][ ∗]s_ [(][i, j][) = 0][ otherwise, and][ β][∗] _[is an optimal Lagrangian]_
_multipler corresponds to the following dual formulation:_

_N_

minβ 0 _[F]_ [(][β][) = min]β 0  _υ_ _π(aj|s)[A[π](s, aksπ[(][β,j][)][)][ −]_ _[βM][k]s[π][(][β,j][)][j][]]_ (6)
_≥_ _≥_ _j=1_

 X 

_Moreover, we have β[∗]_ _≤_ _β[¯] := max[βδ][ +]s[ E]∈S[s],k,j[∼][ρ][π]=1...N,k≠_ _j (Mkj)[−][1](A[π](s, ak) −_ _A[π](s, aj))._ _[.]_

_Proof of Theorem 1.Ni=1_ _[Q]ij[s]_ [=][ π][(][a][j][|] First, we denote[s][)][ and][ P]j[N]=1 _[Q]ij[s]_ _Q[=][s][ π]as the joint distribution of[′][(][a][i][|][s][)][.]_ Also, let fs(i, j) represent the conditional π(·|s) and π[′](·|s) with

distribution ofPN _π[′](ai|s) under π(aj|s). Then Q[s]ij_ [=][ π][(][a][j][|][s][)][f][s][(][i, j][)][,][ π][′][(][a][i][|][s][) =][ P]j[N]=1 _[Q]ij[s]_ [=]
_j=1_ _[π][(][a][j][|][s][)][f][s][(][i, j][)][. In addition:]_

_N_ _N_ _N_ _N_

P

_dW (π[′](_ _s), π(_ _s))_ = min _MijQ[s]ij_ [= min] _Mijπ(aj_ _s)fs(i, j), and_

_·|_ _·|_ _Q[s]ij_ _i=1_ _j=1_ _fs(i,j)_ _i=1_ _j=1_ _|_

X X X X


Ea∼π′(·|s)[A[π](s, a)]


_A[π](s, ai)π[′](ai_ _s) =_
_|_
_i=1_

X


_A[π](s, ai)π(aj_ _s)fs(i, j)._
_|_
_j=1_

X


_i=1_


Thus, the WPO problem in (4) can be reformulated as:


max Es _ρπυ_
_fs(i,j)≥0_ _∼_

_s.t._ Es _ρπυ_
_∼_


_A[π](s, ai)π(aj_ _s)fs(i, j)_ (10a)
_|_
_j=1_

X

_N_

_Mijπ(aj_ _s)fs(i, j)_ _δ,_ (10b)
_|_ _≤_
_j=1_

X


_i=1_

_N_

_i=1_

X


_fs(i, j) = 1,_ _s_ _, j = 1 . . . N._ (10c)
_∀_ _∈S_
_i=1_

X

_N_ _N_
Note here that (10b) is equivalent to Es∼ρπυ [min][f][s][(][i,j][)] _i=1_ _j=1_ _[M][ij][π][(][a][j][|][s][)][f][s][(][i, j][)]_ _≤_
_δ_ because if we have a feasible _fs(i, j)_ to make (10b) hold, we must have
_N_ _N_ P P
Es∼ρπυ [min][f][s][(][i,j][)] _i=1_ _j=1_ _[M][ij][π][(][a][j][|][s][)][f][s][(][i, j][)][ ≤]_ _[δ][.]_

Since both the objective function and the constraint are linear inP P _fs(i, j), (10) is a convex optimization_
problem. Also, Slater’s condition holds for (10) as the feasible region has an interior point, which is
_fs(i, i) = 1_ _i, and fs(i, j) = 0_ _i_ = j. Meanwhile, since A[π](s, a) is bounded based on Assumption
_∀_ _∀_ _̸_
1, the objective is bounded above. Therefore, strong duality holds for (10). At this point we can
derive the dual problem of (10) as its equivalent reformulation:


min _βδ +_ _ζj[s][ds]_
_β≥0,ζj[s]_ Zs∈S _j=1_

X

_s.t._ _A[π](s, ai)π(aj_ _s)_ _βMijπ(aj_ _s)_ _ζj[s]_ _s_ _, i, j = 1 . . . N._
_|_ _−_ _|_ _−_ _ρ[π]υ_ [(][s][)][ ≤] [0][,] _∀_ _∈S_

We observe that with a fixed β, the optimal ζj[s] [will be achieved at:]


(11)


_ζj[s][∗][(][β][) = max]_ _υ_ [(][s][)][π][(][a][j][|][s][)(][A][π][(][s, a][i][)][ −] _[βM][ij][)][.]_ (12)
_i=1...N_ _[ρ][π]_


-----

Denote β[∗] as an optimal solution to (11) and fs[∗][(][i, j][)][ as an optimal solution to (10). Due to the]
complimentary slackness, the following equations hold:

_j_ [(][β][∗][)]
(A[π](s, ai)π(aj _s)_ _β[∗]Mijπ(aj_ _s)_ _s_ [(][i, j][) = 0][,] _s, i, j._
_|_ _−_ _|_ _−_ _[ζ]ρ[s][∗][π]υ_ [(][s][) )][f][ ∗] _∀_

In this case, fs[∗][(][i, j][)][ can have non-zero values only when][ A][π][(][s, a][i][)][π][(][a][j][|][s][)][ −] _[β][∗][M][ij][π][(][a][j][|][s][)][ −]_
_ζj[s][∗][(][β][∗][)]_
_ρ[π]υ_ [(][s][)] = 0, which means ζj[s][∗][(][β][∗][) =][ ρ]υ[π][(][s][)][π][(][a][j][|][s][)(][A][π][(][s, a][i][)][ −] _[β][∗][M][ij][)][. Given the expression]_
of the optimal ζj[s][∗] in (12), fs[∗][(][i, j][)][ can have non-zero values only when][ i][ ∈K]s[π][(][β][∗][, j][)][, where]

_Ks[π][(][β, j][) =][ argmax]k=1...N_ _[A][π][(][s, a][k][)][ −]_ _[βM][kj][. Since][ P][N]i=1_ _[f][ ∗]s_ [(][i, j][) = 1][ as indicated in (10c), we]
can choose an arbitrary optimizer ks[π][(][β][∗][, j][)][ ∈K]s[π][(][β][∗][, j][)][ to derive an optimal solution][ f][ ∗]s [(][i, j][)][:]

1 if i = ks[π][(][β][∗][, j][)]
_fs[∗][(][i, j][) =]_ 0 otherwise,


And then the corresponding optimal solution is, π[∗](ai|s) = _j=1_ _[π][(][a][j][|][s][)][f][ ∗]s_ [(][i, j][)][.]

Last, by substituting ζj[s][∗][(][β][) =][ ρ]υ[π][(][s][)][π][(][a][j][|][s][)(][A][π][(][s, a]ks[π][(][β,j][)][)][ −] _[βM][k]s[π][(][β,j][)][j][)][ into the dual problem]_

[P][N]

(11), we can reformulate (11) into:

_N_ _N_

minβ≥0[{][βδ][ +] Zs∈S _j=1_ _ζj[s][∗][(][β][)][ds][}][ = min]β≥0[{][βδ][ +][ E][s][∼][ρ]υ[π]_ _j=1_ _π(aj|s)[A[π](s, aksπ[(][β,j][)][)][ −]_ _[βM][k]s[π][(][β,j][)][j][]][}][.]_

X X

(13)
The optimal β can then be obtained by solving (13).


We will further show that β[∗] _β := maxs_ _,k,j=1...N,k=j (Mkj)[−][1](A[π](s, ak)_ _A[π](s, aj))._
_≤_ [¯] _∈S_ _̸_ _−_

In the general case, i.e., β ≥ 0, (10a) is non-negative because:


_A[π](s, ai)π(aj|s)fs[∗][(][i, j][)]_ (14a)
_j=1_

X


Es _ρπυ_
_∼_


_i=1_


= Es _ρπυ_
_∼_

= Es _ρπυ_
_∼_

Es _ρπυ_
_≥_ _∼_

= Es _ρπυ_
_∼_


_A[π](s, ai)fs[∗][(][i, j][)]_ (14b)
_i=1_

X


_π(aj_ _s)_
_|_
_j=1_

X


_π(aj|s)A[π](s, aksπ[(][β][∗][,j][)][)]_ (14c)
_j=1_

X

_N_

_π(aj|s){A[π](s, aj) + β[∗]Mksπ[(][β][∗][,j][)][j][}]_ (14d)
_j=1_

X

_N_

_π(aj_ _s)β[∗]Mksπ[(][β][∗][,j][)][j]_ (14e)
_|_
_j=1_

X


_≥_ 0, (14f)

where (14d) holds since A[π](s, aksπ[(][β][∗][,j][)][)][ −] _[β][∗][M][k]s[π][(][β][∗][,j][)][j][ ≥]_ _[A][π][(][s, a][j][)][ −]_ _[β][∗][M][jj][ =][ A][π][(][s, a][j][)][.]_

When β[∗] _> maxs∈S,k,j=1...N,k≠_ _j{_ _[A][π][(][s,a][k]M[)][−]kj[A][π][(][s,a][j]_ [)] _}, we have that for all s ∈S, ks[π][(][β][∗][, j][) =][ j][.]_

Thus, fs[∗][(][i, i][) = 1][,][ ∀][i][ and][ f][ ∗]s [(][i, j][) = 0][,][ ∀][i][ ̸][=][ j][. The objective value (10a) will be][ 0][ because]
_N_ _N_ _N_
Es∼ρπυ _i=1_ _j=1_ _[A][π][(][s, a][i][)][π][(][a][j][|][s][)][f][ ∗]s_ [(][i, j][) =][ E][s][∼][ρ]υ[π] _i=1_ _[A][π][(][s, a][i][)][π][(][a][i][|][s][) = 0][. The left hand]_

_N_ _N_ _N_
side of (10b) equals toP P Es∼ρπυ _i=1_ _j=1_ _[M][ij][π][(][a][j][|][s]P[)][f][ ∗]s_ [(][i, j][) =][ E][s][∼][ρ]υ[π] _i=1_ _[M][ii][π][(][a][i][|][s][) = 0][.]_
Thus, for any δ > 0, (10b) is always satisfied.
P P P

Since the objective of the primal Wasserstein trust-region constrained problem in (6) constantly
evaluates to 0 when β[∗] _> maxs∈S,k,j=1...N,k≠_ _j{_ _[A][π][(][s,a][k]M[)][−]kj[A][π][(][s,a][j]_ [)] _}, and is non-negative when β[∗]_ _≤_


-----

maxs∈S,k,j=1...N,k≠ _j{_ _[A][π][(][s,a][k]M[)][−]kj[A][π][(][s,a][j]_ [)] _}, we can use maxs∈S,k,j=1...N,k≠_ _j{_ _[A][π][(][s,a][k]M[)][−]kj[A][π][(][s,a][j]_ [)] _} as_

an upper bound for the optimal dual variable β[∗].

B PROOF OF THEOREM 2

**Theorem 2. (Performance improvement) For any initial state distribution µ and any βt ≥** 0, if
_A[ˆ][π]_ _A[π]_ _ϵ for some ϵ > 0, the WPO policy update with the inaccurate advantage function_
_||_ _−_ _||∞_ _≤_
_Aˆ[π], guarantees the following performance improvement bound,_


_N_ 2ϵ
_J(πt+1) ≥_ _J(πt) + βtEs∼ρπtµ_ +1 _j=1[π][t][(][a][j][|][s][)][M]k[ˆ]s[πt]_ [(][β]t[,j][)][j][ −] 1 − _γ [.]_ (7)
X

_Proof of Theorem 2._

_J(πt+1) −_ _J(πt) = Es∼ρπtµ_ +1 Ea∼πt+1 [A[π][t] (s, a)] (15a)


= Es∼ρπtµ +1

= Es∼ρπtµ +1

= Es∼ρπtµ +1

= Es∼ρπtµ +1

Es _ρπtµ_ +1
_≥_ _∼_


_πt+1(ai_ _s)A[π][t]_ (s, ai) (15b)
_|_
_i=1_

X


_πt(aj|s) f[ˆ]s[t][(][i, j][)][A][π][t]_ [(][s, a][i][)] (15c)
_j=1_

X


_i=1_


_fˆs[t][(][i, j][)][A][π][t]_ [(][s, a][i][)] (15d)
_i=1_

X


_πt(aj_ _s)_
_|_
_j=1_

X


_πt(aj_ _s)A[π][t]_ (s, akˆs[πt] [(][β]t[,j][)][)] (15e)
_|_
_j=1_

X

_N_

_πt(aj_ _s)[A[π][t]_ (s, aj) + βtMkˆs[πt] [(][β]t[,j][)][j][ −] [2][ϵ][]] (15f)
_|_
_j=1_

X


_N_

2ϵ
_πt(aj|s)Mkˆs[πt]_ [(][β]t[,j][)][j][ −] 1 _γ [,]_ (15g)
_j=1_ _−_

X


= βtEs∼ρπtµ +1


where (15a) holds due to the performance difference lemma in Kakade & Langford (2002);
(15f) follows from the definition of _k[ˆ]s[π][t]_ [(][β][t][, j][)][ and the fact that][ ||][ ˆ]A[π][t] _A[π][t]_ _ϵ, therefore_
_−_ _||∞_ _≤_

[A[π][t] (s, akˆs[πt] [(][β]t[,j][)][) +][ ϵ][]][ −] _[β][t][M]k[ˆ]s[πt]_ [(][β]t[,j][)][j][ ≥] _Aˆ[π][t]_ (s, akˆs[πt] [(][β]t[,j][)][)][ −] _[β][t][M]k[ˆ]s[πt]_ [(][β]t[,j][)][j][ ≥] _Aˆ[π][t]_ (s, aj)
_−_

_βtMjj = A[ˆ][π][t]_ (s, aj) _A[π][t]_ (s, aj) _ϵ; (15g) holds since Ea_ _π[A[π](s, a)] = 0._
_≥_ _−_ _∼_

C PROOF OF THEOREM 3

**Theorem 3. If Assumption 1 holds, then the optimal solution to the trust region constrained problem**
_(4) with Sinkhorn divergence is:_

_N_ exp ( _β[λ]λ[∗]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]_

_πλ[∗][(][a][i][|][s][) =]_ _N_ _π(aj_ _s),_ (8)

_|_

_j=1_ _k=1_ [exp (][ λ]βλ[∗] _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X

_where M denotes the cost matrix andP βλ[∗]_ _[is an optimal solution to the following dual formulation:]_


_N_
minβ≥0 Fλ(β) = minβ≥0 _βδ −_ Es∼ρπυ _j=1_ _[π][(][a][j][|][s][)(][ β]λ_ [+][ β]λ [ln(][π][(][a][j][|][s][))][ −]

_β_ n PN _N_ _β_ exp ( _β[λ]_ _[A][π][(][s,a][i][)][−][λM][ij]_ [)][·][π][(][a][j] _[|][s][)]_
_λ_ [ln[][P]i[N]=1 [exp (][ λ]β _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]) +][ E][s][∼][ρ]υ[π]_ _i=1_ _j=1_ _λ_ _Nk=1_ [exp (][ λ]β _[A][π][(][s,a][k][)][−][λM][kj]_ [)]

P P P

_Moreover, we have βλ[∗]_ _δ_ _._

_[≤]_ [2][A][max]


_. (9)_


-----

_Proof of Theorem 3. Invoking the definition of Sinkhorn divergence in (3), the trust region con-_
strained problem with Sinkhorn divergence can be reformulated as:

_N_ _N_
maxQ Es∼ρπυ [[] _i=1[A][π][(][s, a][i][)]_ _j=1[Q]ij[s]_ []] (16a)
XN _N_ X

_s.t._ Es∼ρπυ [[] _i=1_ _j=1[M][ij][Q]ij[s]_ [+ 1]λ _[Q]ij[s]_ [log][ Q]ij[s] []][ ≤] _[δ]_ (16b)

_N_ X X

_i=1[Q]ij[s]_ [=][ π][(][a][j][|][s][)][,] _∀j = 1, . . ., N, s ∈S._ (16c)

X


Let β and ω represent the dual variables of constraints (16b) and (16c) respectively, then the Lagrangian duality of (16) can be derived as:


max min min _υ_ [[] _A[π](s, ai)_
_Q_ _β≥0,ω_ _[L][(][Q, β, ω][) = max]Q_ _β≥0,ω_ [E][s][∼][ρ][π] _i=1_

X


_Q[s]ij[]]_
_j=1_

X


_N_

_MijQ[s]ij_ [+ 1] _ij_ [log][ Q]ij[s] [])]

_λ_ _[Q][s]_

_j=1_

X

(17a)

_ωj[s]_

_ij[ρ]υ[π][(][s][)][ds]_
_ρ[π]υ_ [(][s][)] _[Q][s]_


_ωj[s][(]_
_j=1_

X


_i=1_ _Q[s]ij_ _[−]_ _[π][(][a][j][|][s][))][ds][ +][ β][(][δ][ −]_ [E][s][∼][ρ]υ[π] [[] _i=1_

X X


_s∈S_


= max min _υ_ [[] _A[π](s, ai)_
_Q_ _β≥0,ω_ [E][s][∼][ρ][π] _i=1_

X


_Q[s]ij[] +]_
_j=1_

X


_s∈S_


_j=1_


_i=1_


_N_

_MijQ[s]ij_ [+ 1] _ij_ [log][ Q]ij[s] [])] (17b)

_λ_ _[Q][s]_

_j=1_

X


_ωj[s][π][(][a][j][|][s][)][ds][ +][ βδ][ −]_ _[β][E][s][∼][ρ][π]υ_ [[]
_j=1_ _i=1_

X X


_s∈S_


_ωj[s][π][(][a][j][|][s][)][ds]_
_j=1_

X


= max min
_Q_ _β_ 0,ω _[βδ][ −]_
_≥_


_s∈S_


_N_ _ωj[s]_

(A[π](s, ai) _βMij +_ _ij_ _ij_ [log][ Q]ij[s] []] (17c)
_j=1_ _−_ _ρ[π]υ_ [(][s][))][Q][s] _[−]_ _[β]λ_ _[Q][s]_

X


+ Es _ρ[π]υ_ [[]
_∼_


_i=1_


_ωj[s][π][(][a][j][|][s][)][ds]_
_j=1_

X


min
_β_ 0,ω [max]Q _[βδ][ −]_
_≥_


_s∈S_


_N_ _ωj[s]_

(A[π](s, ai) _βMij +_ _ij_ _ij_ [log][ Q]ij[s] []][,] (17d)
_j=1_ _−_ _ρ[π]υ_ [(][s][))][Q][s] _[−]_ _[β]λ_ _[Q][s]_

X


+ Es _ρ[π]υ_ [[]
_∼_


_i=1_


where (17d) holds since the Lagrangian function L(Q, β, ω) is concave in Q and linear in β and ω,
and we can exchange the max and the min following the Minimax theorem (Sion, 1958).

Note that the inner max problem of (17d) is an unconstrained concave problem, and we can obtain
the optimal Q by taking the derivatives and setting them to 0. That is,

_∂L_ _ωj[s]_

= A[π](s, ai) _βMij +_ _ij_ [+ 1) = 0][,][ ∀][i, j][ = 1][,][ · · ·][, N, s][ ∈S][.] (18)
_∂Q[s]ij_ _−_ _ρ[π]υ_ [(][s][)][ −] _[β]λ_ [(log][ Q][s]


Therefore, we have the optimal Q[s]ij[∗] [as:]

_j_

_Q[s]ij[∗]_ [= exp (] _[λ]_ (19)

_β [A][π][(][s, a][i][)][ −]_ _[λM][ij][) exp (][ λω]βρ[π]υ_ [(][s][s][)][ −] [1)][,][ ∀][i, j][ = 1][,][ · · ·][, N, s][ ∈S][.]

In addition, since _i=1_ _[Q]ij[s][∗]_ [=][ π][(][a][j][|][s][)][, we have the following hold:]

_j_ _π(aj_ _s)_

[P][N]exp ( _βρ[λω][π]υ_ [(][s][s][)][ −] [1) =] _Ni=1_ [exp (][ λ]β _[A][π][(]|[s, a][i][)][ −]_ _[λM][ij][)]_ _._ (20)

P


-----

By substituting the left hand side of (20) into (19), we can further reformulate the optimal Q[s]ij[∗] [as:]

exp ( _β[λ]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]_
_Q[s]ij[∗]_ [=] _N_ _π(aj_ _s),_ _i, j = 1,_ _, N, s_ _._ (21)

_|_ _∀_ _· · ·_ _∈S_

_k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

To obtain the optimal dual variables, based on (20), we have the optimalP _ω[∗]_ as:


_N_

_ωj[s][∗]_ = ρ[π]υ [(][s][)][{] _[β]_ exp ( _[λ]_

_λ_ [+][ β]λ [ln(][π][(][a][j][|][s][))][ −] _[β]λ_ [ln[] _β [A][π][(][s, a][i][)][ −]_ _[λM][ij][)]][}][,][ ∀][j][ = 1][,][ · · ·][, N, s][ ∈S]_

_i=1_

X

(22)
By substituting (21) and (22) into (17d), we can obtain the optimal β[∗] via:


_N_

_π(aj_ _s)_
_|_ _{_ _[β]λ_ [+][ β]λ [ln(][π][(][a][j][|][s][))][ −] _[β]λ_ [ln[]
_j=1_

X


_N_

exp ( _[λ]_

_β [A][π][(][s, a][i][)][ −]_ _[λM][ij][)]][}]_

_i=1_

X


min _βδ_ Es _ρπυ_
_β_ 0 _−_ _∼_
_≥_


_N_ _N_ _β_ exp ( _β[λ]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)][ ·][ π][(][a][j][|][s][)]_

+ Es∼ρπυ _λ_ _N_ _._

_i=1_ _j=1_ _k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X X

The proof for the upper bound of sinkhorn optimalP _β can be found in Appendix D._

D UPPER BOUND OF SINKHORN OPTIMAL BETA

In this section, we will derive the upper bound of Sinkhorn optimal β. First, for a given β, the optimal
_Q[s]ij[∗][(][β][)][ to the Lagrangian dual][ L][(][Q, β, ω][)][ can be expressed in (21). With this, we will present the]_
following two lemmas:
**Lemma 1. The objective function (16a) with respect to Q[s]ij[∗][(][β][)][ decreases as the dual variable][ β]**
_increases._


**Lemma 2. If Assumption 1 holds, then for every δ > 0, Q[s]ij[∗][(][ 2][A]δ[max]**


) is feasible to (16b) for any λ.


We provide proofs for Lemma 1 and Lemma 2 in Appendix D.1 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 β:
**Proposition 1. If βλ[∗]** _[is the optimal dual solution to the Sinkhorn dual formulation (9), then][ β]λ[∗]_
2A[max] _[≤]_

_δ_ _for any λ._

_Proof of Proposition 1. We will prove it by contradiction. According to Lemma 2, Q[s]ij[∗][(][ 2][A]δ[max]_ ) is

feasible to (16b). Since βλ[∗] [is the optimal dual solution,][ Q]ij[s][∗][(][β]λ[∗][)][ is optimal to (16). If][ β]λ[∗] _[>][ 2][A]δ[max]_,

according to Lemma 1, the objective value in (16a) with respect to [2][A]δ[max] is smaller than the objective

value in (16a) with respect to βλ[∗][, which contradicts the fact that][ Q]ij[s][∗][(][β]λ[∗][)][ is the optimal solution to]
(16).

D.1 PROOF OF LEMMA 1

**Lemma 1. The objective function (16a) with respect to Q[s]ij[∗][(][β][)][ decreases as the dual variable][ β]**
_increases._

_Proof of Lemma 1. Let Gλ(β) represent the objective function (16a). By substituting the optimal_
_Q[s]ij[∗]_ [in (21) into (16a), we have:]


_N_ _N_ exp ( _β[λ]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]_

_A[π](s, ai)_ _N_ _π(aj_ _s)]_ (23a)

_|_

_i=1_ _j=1_ _k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X X

_N_ _N_ P exp ( _β[λ]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]_

_π(aj_ _s)_ _A[π](s, ai)_ _N_ ]. (23b)
_|_
_j=1_ _i=1_ _k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X X

P


_Gλ(β) = Es_ _ρπυ_ [[]
_∼_

= Es _ρπυ_ [[]
_∼_


-----

For any β2 > β1 > 0, we have:

_Gλ(β1)_ _Gλ(β2)_
_−_

= Es∼ρπυ _N_ _π(aj|s)_ _N_ _A[π](s, ai){_ _Nexp (_ _β[λ]1_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]_

_j=1_ _i=1_ _k=1_ [exp (][ λ]β1 _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X X

exp ( _β[λ]2_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]_ P

_N_ (24a)

_−_ _}_

_k=1_ [exp (][ λ]β2 _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

= EPs∼ρπυ _N_ _π(aj|s)_ _N_ _A[π](s, a[i]){_ _Nexp (_ _β[λ]1_ _[A][π][(][s, a][[][i][]][)][ −]_ _[λM][[][i][]][j][)]_

_j=1_ _i=1_ _k=1_ [exp (][ λ]β1 _[A][π][(][s, a][[][k][]][)][ −]_ _[λM][[][k][]][j][)]_

X X

exp ( _β[λ]2_ _[A][π][(][s, a][[][i][]][)][ −]_ _[λM][[][i][]][j][)]_ P

_N_ _,_ (24b)

_−_ _}_

_k=1_ [exp (][ λ]β2 _[A][π][(][s, a][[][k][]][)][ −]_ _[λM][[][k][]][j][)]_

where [i] denotes sorted indices that satisfyP _A[π](s, a[1])_ _A[π](s, a[2])_ _A[π](s, a[N_ ]). Let
_≥_ _≥· · · ≥_

exp ( _β[λ]1_ _[A][π][(][s, a][[][i][]][)][ −]_ _[λM][[][i][]][j][)]_ exp ( _β[λ]2_ _[A][π][(][s, a][[][i][]][)][ −]_ _[λM][[][i][]][j][)]_
_fs(i) =_ _N_ _N_ (25a)

_−_

_k=1_ [exp (][ λ]β1 _[A][π][(][s, a][[][k][]][)][ −]_ _[λM][[][k][]][j][)]_ _k=1_ [exp (][ λ]β2 _[A][π][(][s, a][[][k][]][)][ −]_ _[λM][[][k][]][j][)]_
P exp (( _β[λ]1_ _β2_ [)][A][π][(][s, a][[][i][]][)) exp (][ λ]β2 _[A]P[π][(][s, a][[][i][]][)][ −]_ _[λM][[][i][]][j][)]_

= _N_ _[−]_ _[λ]_

_k=1_ [exp ((][ λ]β1 _[−]_ _β[λ]2_ [)][A][π][(][s, a][[][k][]][)) exp (][ λ]β2 _[A][π][(][s, a][[][k][]][)][ −]_ _[λM][[][k][]][j][)]_
P exp ( _β[λ]2_ _[A][π][(][s, a][[][i][]][)][ −]_ _[λM][[][i][]][j][)]_

_N_ _._ (25b)

_−_

_k=1_ [exp (][ λ]β2 _[A][π][(][s, a][[][k][]][)][ −]_ _[λM][[][k][]][j][)]_

For notation brevity,P we let ms(i) = exp (( _β[λ]1_ _β2_ [)][A][π][(][s, a][[][i][]][))] _>_ 0, _ws(i)_ =

exp ( _β[λ]2_ _[A][π][(][s, a][[][i][]][)][ −]_ _[λM][[][i][]][j][)][ >]m[ 0]s([ and]i)ws[ q](i[s])[(][i][) =]_ PNk=1 _w[m]s[s]1([(][k]i)[)][−][w][s][(][k][λ][)][ −]_ PNk=1 _[m]1[s][(][i][)][w][s][(][k][)]_ [. Then we have]

(25b) = _N_ _N_ (26a)
_−_
_k=1_ _[m][s][(][k][)][w][s][(][k][)]_ _k=1_ _[w][s][(][k][)]_

1 1
= mPs(i)ws(i)( _N_ P _N_ ) (26b)
_−_
_k=1_ _[m][s][(][k][)][w][s][(][k][)]_ _k=1_ _[m][s][(][i][)][w][s][(][k][)]_

= ms(i)ws(i)qs(i). (26c)

P P

Since _βλ1_ _βλ2_ _[>][ 0][,][ m][s][(][i][)][ decreases as][ i][ increases. Thus,][ q][s][(][i][)][ decreases as][ i][ increases. Since]_

_ms(1)_ _[−]ms(k) and ms(N_ ) _ms(k) for all k = 1, . . ., N_, we have qs(1) = _N_ 1
_≥_ _≤_ _k=1_ _[m][s][(][k][)][w][s][(][k][)][ −]_

1 1 1 1
_N_ _N_ _N_ PN
_k=1_ _[m][s][(1)][w][s][(][k][)][ ≥]_ _k=1_ _[m][s][(][k][)][w][s][(][k][)][ −]_ _k=1_ _[m][s][(][k][)][w][s][(][k][)][ = 0][, and][ q][s][(][N]_ [) =] _k=1_ _[m][s][(][k][)][w][s][(][k][)][ −]_

1 1 1

PN P _N_ P _N_ P

_k=1_ _[m][s][(][N]_ [)][w][s][(][k][)][ ≤] _k=1_ _[m][s][(][k][)][w][s][(][k][)][ −]_ _k=1_ _[m][s][(][k][)][w][s][(][k][)][ = 0][. Since][ q][s][(1)][ ≥]_ [0][,][ q][s][(][N] [)][ ≤] [0][ and]

_qPs(i) decreases as i increases, there exists an indexP_ P 1 _ks_ _N such that qs(i)_ 0 for i _ks_
and qs(i) < 0 for i > ks. Since ms(i), ws(i) > 0, we have ≤ _≤ fs(i) ≥_ 0 for i ≤ _k ≥s and fs(i ≤) < 0_
for i > ks. In addition, we have _i=1_ _[f][s][(][i][) = 0][ directly follows from the definition. Thus,]_
_N_
_i=1_ _[f][s][(][i][) =][ P]i[k]=1[s]_ _[|][f][s][(][i][)][| −]_ [P]i[N]=ks+1 _[|][f][s][(][i][)][|][ = 0][. Therefore,]_

_N_ [P][N] _N_

P

_Gλ(β1)_ _Gλ(β2) = Es_ _ρ[π]υ_ _π(aj_ _s)_ _A[π](s, a[i])fs(i)_ (27a)
_−_ _∼_ _|_

_j=1_ _i=1_

X X


_ks_

_A[π](s, a[i])_ _fs(i)_
_|_ _| −_
_i=1_

X

_ks_


_A[π](s, a[i])_ _fs(i)_
_|_ _|}_
_i=ks+1_

X


= Es _ρπυ_
_∼_

Es _ρπυ_
_≥_ _∼_


_π(aj_ _s)_
_|_ _{_
_j=1_

X

_N_

_π(aj_ _s)_
_|_ _{_
_j=1_

X


(27b)


_A[π](s, a[ks])_ _fs(i)_
_|_ _| −_
_i=1_

X


_A[π](s, a[ks+1])_ _fs(i)_
_|_ _|}_
_i=ks+1_

X

(27c)


-----

_ks_

_fs(i)_ _A[π](s, a[ks+1])_
_|_ _| −_
_i=1_

X

_ks_

_fs(i)_ _A[π](s, a[ks+1])_
_|_ _| −_
_i=1_

X

_ks_


_π(aj_ _s)_ _A[π](s, a[ks])_
_|_ _{_
_j=1_

X

_N_

_π(aj_ _s)_ _A[π](s, a[ks])_
_|_ _{_
_j=1_

X


= Es _ρπυ_
_∼_

= Es _ρπυ_
_∼_

= Es _ρ[π]υ_
_∼_


_fs(i)_
_|_ _|}_
_i=ks+1_

X

(27d)

_ks_

_fs(i)_
_|_ _|}_
_i=1_

X

(27e)


_π(aj_ _s)(A[π](s, a[ks])_ _A[π](s, a[ks+1]))_
_|_ _−_
_j=1_

X


_fs(i)_ (27f)
_|_ _|_
_i=1_

X


_≥_ 0. (27g)
where (27c) and (27g) hold since A[π](s, a[i]) is non-increasing as i increases. Furthermore, at least
one inequality of (27c) and (27g) will not hold at equality since _i=1_ _[π][(][a][i][|][s][)][A][π][(][s, a][i][) = 0][,][ ∀][s][ ∈S][,]_
and for non-trivial cases, Pr _A[π](s, a) = 0,_ _s_ _,_ _a_ _< 1, which means Pr_ _s1, s2_
_, a1, a2_ _, s.t. A[π](s1, a1){_ = A[π](s2, a2) _∀ > ∈S 0. Therefore, we have ∀_ _∈A}_ _Gλ(β1)_ _Gλ(β2{∃) > 0._ _∈_
_S_ _∈A_ _̸_ _}_ [P][N] _−_

D.2 PROOF OF LEMMA 2


**Lemma 2. If Assumption 1 holds, then for every δ > 0, Q[s]ij[∗][(][ 2][A]δ[max]**


) is feasible to (16b) for any λ.


_Proof of Lemma 2. By substituting the optimal Q[s]ij[∗]_ [in (21) into (16b), we can reformulate the left]
hand side of (16b) as follows:


_N_

_MijQ[s]ij[∗]_ [+ 1] _ij_ [log][ Q]ij[s][∗][]] (28a)

_λ_ _[Q][s][∗]_

_j=1_

X


Es _ρ[π]υ_ [[]
_∼_


_i=1_


_N_

_π(aj_ _s)_

_MijQ[s]ij[∗]_ [+ 1]λ _[Q]ij[s][∗][[]_ _β [λ]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][ + log]_ _N_ _|_ ]}
_j=1_ _k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X

(28b)

P


= Es _ρπυ_ _[{]_
_∼_

= Es _ρ[π]υ_ _[{]_
_∼_


_i=1_

_N_

_i=1_

X


_N_ _N_

1 _π(aj_ _s)_

= Es∼ρ[π]υ _[{]_ _β [Q]ij[s][∗][A][π][(][s, a][i][) + 1]λ_ _[Q]ij[s][∗]_ [log] _N_ _|_ _}._ (28c)

_i=1_ _j=1_ _k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X X

Now we prove that when β = 2Aδ[max], Es∼Pρπυ _[{][P][N]i=1_ _Nj=1_ _β1_ _[Q]ij[s][∗][(][β][)][A][π][(][s, a][i][)][}]_ _≤_ 2δ [and]

_π(aj_ _s)_

Es∼ρπυ _[{][ 1]λ_ _[Q]ij[s][∗][(][β][) log]_ PNk=1 [exp (][ λ]β _[A][π][(]|[s,a][k]N[)][−][λM]N_ _[kj]_ [)] _[} ≤]_ 2[δ] [hold. For the first part, we have:]P


Es _ρπυ_ _[{]_
_∼_


1

_β [Q]ij[s][∗][A][π][(][s, a][i][)][}]_ (29a)


_i=1_


_j=1_


= [1] _υ_ _[{]_

_β_ [E][s][∼][ρ][π]

= [1] _υ_ _[{]_

_β_ [E][s][∼][ρ][π]

_υ_ _[{]_

_≤_ _β[1]_ [E][s][∼][ρ][π]



[ _Q[s]ij[∗][]][A][π][(][s, a][i][)][}]_ (29b)

_j=1_

X

_π[′](ai_ _s)A[π](s, ai)_ (29c)
_|_ _}_

_π[′](ai_ _s)_ _A[π](s, ai)_ (29d)
_|_ _|_ _|}_


_i=1_

_N_

_i=1_

X

_N_

_i=1_

X


_≤_ _[A][max]β_

For the second part, the followings hold:


= _[δ]_ (29e)

2 _[.]_


_λ1_ _[Q]ij[s][∗]_ [log] _N_ _π(aj|s)_ _}_ (30a)
_k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_
P


Es _ρπυ_ _[{]_
_∼_


_i=1_


_j=1_


-----

_N_ _N_

1 _π(aj_ _s)_

= Es∼ρπυ _[{]_ _λ_ [(] _Q[s]ij[∗][) log]_ _N_ _|_ _}_ (30b)

_j=1_ _i=1_ _k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X X

_N_

P _π(aj_ _s)_

= λ[1] [E][s][∼][ρ]υ[π] _[{]_ _π(aj|s) log_ _N_ _|_ _}_ (30c)

_j=1_ _k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X

_N_

P _π(aj_ _s)_

_υ_ _[{]_ _π(aj_ _s) log_ _|_ (30d)

_≤_ _λ[1]_ [E][s][∼][ρ][π] _j=1_ _|_ exp ( _β[λ]_ _[A][π][(][s, a][j][))]_ _[}]_

X


_υ_ _[{]_

_≤_ _λ[1]_ [E][s][∼][ρ][π]

= [1] _υ_ _[{]_

_λ_ [E][s][∼][ρ][π]


1

(30e)

exp ( _β[λ]_ _[A][π][(][s, a][j][))]_ _[}]_


_π(aj_ _s) log_
_|_
_j=1_

X


_N_

_π(aj_ _s)(_ (30f)
_|_ _−_ _β [λ]_ _[A][π][(][s, a][j][))][}]_
_j=1_

X


_N_

_υ_ _[{]_ _π(aj_ _s)_ _A[π](s, aj)_ (30g)

_≤_ _β[1]_ [E][s][∼][ρ][π] _j=1_ _|_ _|_ _|}_

X


_≤_ _[A][max]β_


= _[δ]_ (30h)

2 _[.]_


Therefore, Q[s]ij[∗][(][ 2][A]δ[max]


) is feasible to (16b) for any λ.


E GRADIENT OF THE OBJECTIVE IN THE SINKHORN DUAL FORMULATION

The closed-form gradient of the objective in the Sinkhorn dual formulation (9) is as follows:


_N_ 1

_π(aj|s)_ _λ_ [+ 1]λ [ln(][π][(][a][j][|][s][))][ −] _λ[1]_ [ln[]
_j=1_

X n


_N_

exp ( _[λ]_

_β [A][π][(][s, a][i][)][ −]_ _[λM][ij][)]]_

_i=1_

X


_δ −_ Es∼ρ[π]υ


_N_

1

_−_ _[β]λ_ _N_ _×_ [exp ( _β [λ]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)][ × −][λA][π][(][s, a][i][)][β][−][2][]]_

_[·]_ _i=1_ [exp (][ λ]β _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]_ Xi=1 o
P _N_ _N_ _π(aj_ _s)_ exp ( _β[λ]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)]_

+ Es∼ρ[π]υ _λ_ _|_ _N_

_i=1_ _j=1_ _k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X X n

exp ( _β[λ]_ _[A][π][(][s, a]P[i][)][ −]_ _[λM][ij][)][ × −][λA][π][(][s, a][i][)][β][−][2][ ×][ P]k[N]=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

+ _[βπ][(][a][j][|][s][)]_

_λ_ _·_ ([P][N]k=1 [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][))][2]_

exp ( _β[λ]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)][ ×][ P]k[N]=1[[exp (][ λ]β_ _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)][ × −][λA][π][(][s, a][k][)][β][−][2][]]_

_._

_−_ _[βπ][(]λ[a][j][|][s][)]_ _·_ ([P][N]k=1 [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][))][2]_

o

F PROOF OF THEOREM 4

Given the upper bound of Wassertein optimal β in Theorem 1 and the upper bound of Sinkhorn
optimal β in Proposition 1, we are able to derive the following theorem:


**Theorem 4. Define βUB = max** [2][A]δ[max]
_{_


_,_ _β[¯]}. The following holds:_


_1. Fλ(β) converges to F_ (β) uniformly on [0, βUB],

_2._ _λlim−→∞_ 0[arg min]≤β≤βUB _Fλ(β) ⊆_ 0arg min≤β≤βUB _F_ (β).

_Proof of Theorem 4. To show that Fλ(β) converges to F_ (β) uniformly on [0, βUB], it is equivalent
to show that limλ−→∞ [sup]0≤β≤βUB _Fλ(β) −_ _F_ (β) = 0. Let ϵ[π]s [(][β, i, j][) =][ A][π][(][s, a]ks[π][(][β,j][)][)][ −]


-----

_βMksπ[(][β,j][)][j][ −]_ [[][A][π][(][s, a][i][)][ −] _[βM][ij][]][ where][ k]s[π][(][β, j][)][ ∈K]s[π][(][β, j][) =][ argmax]k=1...N_ _[A][π][(][s, a][k][)][ −]_ _[βM][kj]_
is an arbitrary optimizer, and ϵ[π]s [(][β, i, j][)][ ≥] [0][. First, we have]

_Fλ(β)_ _F_ (β)
_−_

_N_ _N_

= _βδ_ Es _ρ[π]υ_ _π(aj_ _s)_ exp ( _[λ]_
_−_ _∼_ _|_ _{_ _[β]λ_ [+][ β]λ [ln(][π][(][a][j][|][s][))][ −] _[β]λ_ [ln[] _β [A][π][(][s, a][i][)][ −]_ _[λM][ij][)]][}]_

_j=1_ _i=1_

X X

_N_ _N_ _β_ exp ( _β[λ]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)][ ·][ π][(][a][j][|][s][)]_

+ Es∼ρπυ _λ_ _N_ _−_ _βδ_

_i=1_ _j=1_ _k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X X

_N_

P

_−_ Es∼ρ[π]υ _π(aj|s)[A[π](s, aksπ[(][β,j][)][)][ −]_ _[βM][k]s[π][(][β,j][)][j][]]_ (31a)

_j=1_

X


_λ_ [E][s][∼][ρ]υ[π]

_[β]_


_λ_ [E][s][∼][ρ]υ[π]

_[β]_


_π(aj_ _s)_
_|_
_j=1_

X


_π(aj_ _s) ln(π(aj_ _s))_
_|_ _|_
_j=1_

X


_N_ _N_ _β_ exp ( _β[λ]_ _[A][π][(][s, a][i][)][ −]_ _[λM][ij][)][ ·][ π][(][a][j][|][s][)]_

+ Es∼ρ[π]υ _λ_ _N_

_i=1_ _j=1_ _k=1_ [exp (][ λ]β _[A][π][(][s, a][k][)][ −]_ _[λM][kj][)]_

X X

_N_ _N_

P

+ Es _ρ[π]υ_ _π(aj_ _s)_ _[β]_ exp ( _[λ]_
_∼_ _|_ _λ_ [ln[] _β [A][π][(][s, a][i][)][ −]_ _[λM][ij][)]]_

_j=1_ _i=1_

X X

_N_

_−_ Es∼ρ[π]υ _π(aj|s)[A[π](s, aksπ[(][β,j][)][)][ −]_ _[βM][k]s[π][(][β,j][)][j][]]_ (31b)

_j=1_

X


_λ_ [E][s][∼][ρ]υ[π]

_[β]_


_λ_ [E][s][∼][ρ]υ[π]

_[β]_


_≤_ 2


_π(aj_ _s)_
_|_
_j=1_

X


_π(aj_ _s) ln(π(aj_ _s))_
_|_ _|_
_j=1_

X


_N_ _N_

+ Es _ρ[π]υ_ _π(aj_ _s)_ _[β]_ exp ( _[λ]_
_∼_ _|_ _λ_ [ln[] _β [A][π][(][s, a][i][)][ −]_ _[λM][ij][)]]_

_j=1_ _i=1_

X X

_N_

_−_ Es∼ρπυ _π(aj|s)[A[π](s, aksπ[(][β,j][)][)][ −]_ _[βM][k]s[π][(][β,j][)][j][]]_ _._ (31c)

_j=1_

X


In addition,


_N_

_π(aj_ _s)_ _[β]_
_|_ _λ_ [ln[]
_j=1_

X


_N_

exp ( _[λ]_

_β [A][π][(][s, a][i][)][ −]_ _[λM][ij][)]]_

_i=1_

X


Es _ρ[π]υ_
_∼_

_j_

Es _ρπυ_
_−_ _∼_


_π(aj|s)[A[π](s, aksπ[(][β,j][)][)][ −]_ _[βM][k]s[π][(][β,j][)][j][]]_ (32a)
_j=1_

X


_N_

_π(aj_ _s)_ _[β]_ _s_ [(][β,j][)][)][ −] _[λM][k]s[π][(][β,j][)][j][)]_
_|_ _λ_ [ln[exp (] _β [λ]_ _[A][π][(][s, a][k][π]_
_j=1_

X


_N_

exp (− _β [λ]_ _[ϵ]s[π][(][β, i, j][))]]_
_i=1_

X


Es _ρ[π]υ_
_∼_

_j_

Es _ρπυ_
_−_ _∼_


_π(aj|s)[A[π](s, aksπ[(][β,j][)][)][ −]_ _[βM][k]s[π][(][β,j][)][j][]]_ (32b)
_j=1_

X


_N_ _N_

_π(aj_ _s)_ _[β]_ _s_ [(][β,j][)][)][ −] _[λM][k]s[π][(][β,j][)][j][)] + ln[]_ exp ( _s_ [(][β, i, j][))]][}]
_|_ _λ_ _[{][ln[exp (]_ _β [λ]_ _[A][π][(][s, a][k][π]_ _−_ _β [λ]_ _[ϵ][π]_
_j=1_ _i=1_

X X


Es _ρ[π]υ_
_∼_

_j_

Es _ρπυ_
_−_ _∼_


_π(aj|s)[A[π](s, aksπ[(][β,j][)][)][ −]_ _[βM][k]s[π][(][β,j][)][j][]]_ (32c)
_j=1_

X


-----

_N_

_π(aj_ _s)_ _[β]_
_|_ _λ_ [ln[]
_j=1_ _i_

X


_N_

exp (− _β [λ]_ _[ϵ]s[π][(][β, i, j][))]]_ _._ (32d)
_i=1_

X


= Es _ρ[π]υ_
_∼_

Therefore,


lim sup _Fλ(β)_ _F_ (β) (33a)
_λ−→∞_ 0≤β≤βUB _−_

_N_ _N_

2βUB _βUB_
lim Es _ρ[π]υ_ _π(aj_ _s)_ + lim Es _ρ[π]υ_ _π(aj_ _s) ln(π(aj_ _s))_
_≤_ _λ_ _λ_ _∼_ _|_ _λ_ _λ_ _∼_ _|_ _|_
_−→∞_ _j=1_ _−→∞_ _j=1_

X X

_N_ _N_

_β_

+ lim sup Es _ρ[π]υ_ _π(aj_ _s) ln[_ exp ( _s_ [(][β, i, j][))]] (33b)
_λ−→∞_ 0≤β≤βUB _λ_ _∼_ _j=1_ _|_ _i=1_ _−_ _β [λ]_ _[ϵ][π]_

X X

_N_ _N_

_β_

= lim sup Es _ρ[π]υ_ _π(aj_ _s) ln[_ exp ( _s_ [(][β, i, j][))]] _._ (33c)
_λ−→∞_ 0≤β≤βUB _λ_ _∼_ _j=1_ _|_ _i=1_ _−_ _β [λ]_ _[ϵ][π]_

X X

In addition, ∀β ∈ [0, βUB] and ∀λ, ϵ[π]s [(][β, i, j][)][ is bounded since]
_ϵ[π]s_ [(][β, i, j][)] = _A[π](s, aks[π][(][β,j][)][)][ −]_ _[βM][k]s[π][(][β,j][)][j][ −]_ [[][A][π][(][s, a][i][)][ −] _[βM][ij][]]_ (34)

= (A[π](s, aks[π][(][β,j][)][)][ −] _[A][π][(][s, a][i][))][−][(][βM][k]s[π][(][β,j][)][j][ −]_ _[βM][ij][)]_

_≤_ _A[π](s, aks[π][(][β,j][)][)][ −]_ _[A][π][(][s, a][i][)]_ + _βMks[π][(][β,j][)][j][ −]_ _[βM][ij]_

2 max
_≤_ _s,a_ _[A][π][(][s, a][) +][ β][UB][ max]i,j_ _[M][ij]_

2A[max] + βUB max (35)
_≤_ _i,j_ _[M][ij][ <][ ∞][.]_

Then, Es∼ρ[π]υ _Nj=1_ _[π][(][a][j][|][s][) ln[][P]i[N]=1_ [exp (][−] _β[λ]_ _[ϵ]s[π][(][β, i, j][))]]_ is bounded.

Therefore in (33c), the optimal _β_ can be achieved. Let _β[λ]_ =

P

arg max0≤β≤βUB _βλ_ Es∼ρ[π]υ _Nj=1_ _[π][(][a][j][|][s][) ln[][P]i[N]=1_ [exp (][−] _β[λ]_ _[ϵ]s[π][(][β, i, j][))]]_, and then we have:
P

_N_ _N_

_β_

lim sup Es _ρ[π]υ_ _π(aj_ _s) ln[_ exp ( _s_ [(][β, i, j][))]] (36a)
_λ−→∞_ 0≤β≤βUB _λ_ _∼_ _j=1_ _|_ _i=1_ _−_ _β [λ]_ _[ϵ][π]_

X X

_N_ _N_

_β[λ]_
= lim Es _ρ[π]υ_ _π(aj_ _s) ln[_ exp ( _s_ [(][β][λ][, i, j][))]] _._ (36b)
_λ_ _λ_ _∼_ _|_ _−_ _β[λ][λ][ ϵ][π]_
_−→∞_ _j=1_ _i=1_

X X

Define σs(j) = min0≤β≤βUB mini=1...N,i/∈Ks[π][(][β,j][)][ ϵ]s[π][(][β, i, j][)][. Then since][ ϵ]s[π][(][β, i, j][)][ >][ 0][ for][ i /]∈
_Ks[π][(][β, j][)][ based on its definition, we have][ σ][s][(][j][)][ >][ 0][. On one hand, we have]_

_N_

lim exp ( _s_ [(][β][λ][, i, j][))]] (37a)
_λ_ _−_ _β[λ][λ][ ϵ][π]_
_−→∞_ [ln[] _i=1_

X


exp ( _s_ [(][β][λ][, i, j][)) +]
_−_ _β[λ][λ][ ϵ][π]_


exp ( _s_ [(][β][λ][, i, j][))]][ (37b)]
_−_ _β[λ][λ][ ϵ][π]_


= lim
_λ−→∞_ [ln[]i=1|i/∈KXs[π][(][β][λ][,j][)]

_N_

lim
_≤_ _λ_
_−→∞_ [ln[]i=1|i/∈KXs[π][(][β][λ][,j][)]


_i=1|i∈Ks[π][(][β][λ][,j][)]_


exp (
_−_ _β[λ]UB_


exp (0)] (37c)


_σs(j)) +_


_−→∞_ _i=1|i/∈Ks[π][(][β][λ][,j][)]_ _i=1|i∈Ks[π][(][β][λ][,j][)]_

_N_

= lim exp ( _σs(j)) +_ _s_ [(][β][λ][, j][)][|][]] (37d)
_λ−→∞_ [ln[]i=1|i/∈KXs[π][(][β][λ][,j][)] _−_ _β[λ]UB_ _|K[π]_

= lim _s_ [(][β][λ][, j][)][|][]][.] (37e)
_λ−→∞_ [ln[][|K][π]


-----

On the other hand, we have

_N_

lim exp ( _s_ [(][β][λ][, i, j][))]] (38a)
_λ_ _−_ _β[λ][λ][ ϵ][π]_
_−→∞_ [ln[] _i=1_

X


exp ( _s_ [(][β][λ][, i, j][)) +]
_−_ _β[λ][λ][ ϵ][π]_


exp ( _s_ [(][β][λ][, i, j][))]][ (38b)]
_−_ _β[λ][λ][ ϵ][π]_


= lim
_λ−→∞_ [ln[]i=1|i/∈KXs[π][(][β][λ][,j][)]

_N_

lim
_≥_ _λ_
_−→∞_ [ln[]i=1|i∈KXs[π][(][β][λ][,j][)]


_i=1|i∈Ks[π][(][β][λ][,j][)]_


exp ( _s_ [(][β][λ][, i, j][))]] (38c)
_−_ _β[λ][λ][ ϵ][π]_

exp (0)] (38d)


= lim exp (0)] (38d)
_λ−→∞_ [ln[]i=1|i∈KXs[π][(][β][λ][,j][)]

= lim _s_ [(][β][λ][, j][)][|][]][.] (38e)
_λ−→∞_ [ln[][|K][π]

Therefore, limλ−→∞ ln[[P][N]i=1 [exp (][−] _β[λ][λ][ ϵ]s[π][(][β][λ][, i, j][))]]_ = limλ−→∞ [ln[][|K]s[π][(][β][λ][, j][)][|][]][. Based on that,]

we have

_N_ _N_

_β[λ]_
lim Es _ρ[π]υ_ _π(aj_ _s) ln[_ exp ( _s_ [(][β][λ][, i, j][))]] (39a)
_λ_ _λ_ _∼_ _|_ _−_ _β[λ][λ][ ϵ][π]_
_−→∞_ _j=1_ _i=1_

X X

_N_ _N_

_β[λ]_
lim ln[ exp ( _s_ [(][β][λ][, i, j][))]] (39b)
_≤_ _λ_ _λ_ _−_ _β[λ][λ][ ϵ][π]_
_−→∞_ _j=1_ _i=1_

X X


_β[λ]_
lim
_≤_ _λ_ _λ_
_−→∞_

_β[λ]_
= lim
_λ−→∞_ _λ_


_N_ _N_

ln[ exp ( _s_ [(][β][λ][, i, j][))]] (39c)

_−_ _β[λ][λ][ ϵ][π]_

_j=1_ _i=1_

X X

_N_

ln[|Ks[π][(][β][λ][, j][)][|][]] (39d)
_j=1_

X


_βUB_
lim (39e)
_≤_ _λ_ _λ [N][ ln][ N][ = 0][,]_
_−→∞_

which means limλ−→∞ [sup]0≤β≤βUB _Fλ(β) −_ _F_ (β) _≤_ 0. Furthermore, since

limF (βλ)−→| = 0∞ [sup]. Therefore,0≤β≤βUB _[|][F]λ F[(][β]λ[)]([ −]β)[F] converges to[(][β][)][| ≥]_ [0][ holds naturally, we have] F (β) uniformly on[ lim] [0, βλ−→UB∞], which also indicates[sup]0≤β≤βUB _[|][F]λ[(][β][)][ −]_
_Fλ(β) epi-converges to F_ (β) on [0, βUB] (Royset, 2018; Rockafellar & Wets, 1998). By properties

of epi-convergence, we have that _λlim−→∞_ 0[arg min]≤β≤βUB _Fλ(β) ⊆_ 0arg min≤β≤βUB _F_ (β) (Rockafellar & Wets,
1998).


G OPTIMAL BETA FOR A SPECIAL DISTANCE

**Proposition 2. (1). If the initial point β0 is in [maxs,j{A[π](s, aks** ) − _A[π](s, aj)}, +∞), the local_
_optimal β solution is maxs,j_ _A[π](s, aks_ ) _A[π](s, aj)_ _._
_{_ _−_ _}_
_(2). If the initial point β0 is in [0, mins,j≠_ _ks_ _{A[π](s, aks_ ) − _A[π](s, aj)}]: if δ −_ _s∈S_ _[ρ][π][(][s][)(1][ −]_

_π(aks_ _|s))ds < 0, the local optimal β is mins,j≠_ _ks_ _{A[π](s, aks_ ) − _A[π](s, aj)}; otherwise, the localR_
_optimal β solution is 0._

_(3)._ _If the initial point β0 is in (mins,j=ks_ _A[π](s, aks_ ) _A[π](s, aj)_ _, maxs,j_ _A[π](s, aks_ )
_̸_ _{_ _−_ _}_ _{_ _−_
_A[π](s, aj)}), we construct sets Is[1]_ _[and][ I]s[2]_ _[as:]_

_Then, iffor s ∈S δ_ _−, jE ∈{s∼ρπ1, 2 . . . Nj∈Is[2]_ _[π]}[(] :[a][j] if[|][s][)] β[ <]0 ≥[ 0][, the local optimal]A[π](s, aks_ ) − _A[ β][π]([ is]s, a[ min]j) then[s][∈S][,j] Add[∈][I]s[2]_ _j to Is[1]_ **[else][ Add][ j][ to][ I]s[2][.]**

_otherwise, the local optimal β is maxs_ _,j_ _Is1_ _[{][A][π][(][s, a][k][s]_ [)][−][A][π][(][s, a][j][)][}][;]
_∈S_ _∈_

[P] _[{][A][π][(][s, a][k][s]_ [)][ −] _[A][π][(][s, a][j][)][}][.]_


-----

_Proof of Proposition 2. (1)._ When β [maxs,j _A[π](s, aks_ ) _A[π](s, aj)_ _, +_ ), we have
_∈_ _{_ _−_ _}_ _∞_
_A[π](s, aj)_ _A[π](s, aks_ ) _β for all s_, j = 1 . . . N . Since A[π](s, aks ) _β_ _A[π](s, ak)_ _β for_
_≥_ _−_ _∈S_ _−_ _≥_ _−_
all k = 1 . . . N, we have A[π](s, aj) _A[π](s, ak)_ _β for all s_, j = 1 . . . N, k = 1 . . . N . Thus,
_≥_ _−_ _∈S_
_j ∈Ks[π][(][β][∗][, j][)][, for all][ s][ ∈S][,][ j][ = 1][ . . . N]_ [. Therefore, (6) can be reformulated as:]


minβ 0[{][βδ][ +][ E][s][∼][ρ]υ[π]
_≥_


_π(aj_ _s)A[π](s, aj)_ _._
_|_ _}_
_j=1_

X


Since δ 0, we have the local optimal β = maxs,j _A[π](s, aks_ ) _A[π](s, aj)_ .
_≥_ _{_ _−_ _}_

(2). When β ∈ [0, mins,j≠ _ks_ _{A[π](s, aks_ ) − _A[π](s, aj)}], we have A[π](s, aj) ≤_ _A[π](s, aks_ ) − _β for_
(6) then is:all s ∈S, j = 1 . . . N, j ̸= ks. Thus ks ∈Ks[π][(][β][∗][, j][)][ for all][ s][ ∈S][,][ j][ = 1][ . . . N] [. The inner part of]

_N_

_βδ + Es_ _ρπυ_ _[{]_ _π(aj_ _s)(A[π](s, aks_ ) _β) + π(aks_ _s)A[π](s, aks_ )
_∼_ _|_ _−_ _|_ _}_

_j=1X,j≠_ _ks_


= β(δ Es _ρπυ_
_−_ _∼_


_π(aj_ _s)) + Es_ _ρπυ_
_|_ _∼_
_j=1X,j≠_ _ks_


_π(aj_ _s)A[π](s, aks_ )
_|_
_j=1_

X


_ρ[π]υ_ [(][s][)(1][ −] _[π][(][a][k]s_ _[|][s][))][ds][) +][ E][s][∼][ρ][π]υ_
_s∈S_


_π(aj_ _s)A[π](s, aks_ ).
_|_
_j=1_

X


= β(δ −


If δ − _s∈S_ _[ρ]υ[π][(][s][)(1][ −]_ _[π][(][a][k]s_ _[|][s][))][ds <][ 0][, we have the local optimal][ β][ = min][s,j][̸][=][k]s_ _[{][A][π][(][s, a][k]s_ [)][ −]

_A[π](s, aRj)}. If δ −_ _s∈S_ _[ρ]υ[π][(][s][)(1][ −]_ _[π][(][a][k]s_ _[|][s][))][ds][ ≥]_ [0][, we have the local optimal][ β][ = 0][.]

(3). For an initial pointR _β0 in (mins,j=ks_ _A[π](s, aks_ ) _A[π](s, aj)_ _, maxs,j_ _A[π](s, aks_ ) _A[π](s, aj)_ ),
_̸_ _{_ _−_ _}_ _{_ _−_ _}_
we construct partitions Is[1] [and][ I]s[2] [of the set][ {][1][,][ 2][ . . . N] _[}][ in the way described in Proposition 2 for all]_
_s_ . Consider β in the neightborhood of β0, i.e., β _A[π](s, aks_ ) _A[π](s, aj) for s_ _, j_ _Is[1]_
_∈S_ _≥_ _−_ _∈S_ _∈_
and β ≤ _A[π](s, aks_ ) − _A[π](s, aj) for s ∈S, j ∈_ _Is[2][. Then the inner part of (6) can be reformulated]_
as:


_βδ + Es_ _ρπυ_ _[{]_ _π(aj_ _s)A[π](s, aj) +_
_∼_ _|_

_j_ _Is[1]_

X∈


_π(aj_ _s)(A[π](s, aks_ ) _β)_
_|_ _−_ _}_
_j_ _Is[2]_

X∈


= β(δ Es _ρ[π]υ_ _π(aj_ _s)) + Es_ _ρ[π]υ_ _[{]_ _π(aj_ _s)A[π](s, aj) +_ _π(aj_ _s)A[π](s, aks_ ) _._
_−_ _∼_ _|_ _∼_ _|_ _|_ _}_

_j_ _Is[2]_ _j_ _Is[1]_ _j_ _Is[2]_

X∈ X∈ X∈

If δ − Es∼ρπυ _j∈Is[2]_ _[π][(][a][j][|][s][)][ <][ 0][, we have the local optimal][ β][ = min][s][∈S][,j][∈][I]s[2]_

_A[π](s, aj)}._ PIf δ − Es∼ρπυ _j∈Is[2]_ _[π][(][a][j][|][s][)]_ _≥_ 0, we have the local optimal[{][A][π][(][s, a] β[k][s] [)][ −]=
maxs _,j_ _Is1_
_∈S_ _∈_ P

_[{][A][π][(][s, a][k][s]_ [)][ −] _[A][π][(][s, a][j][)][}][.]_

H IMPLEMENTATION DETAILS

**Visitation Frequencies Estimation: The unnormalized discounted visitation frequencies are needed**
to compute the global optimal β[∗]. At the k-th iteration, the visitation frequencies ρ[π]k [are estimated]
using samples of the trajectory set Dk. Specifically, we first initialize ρ[π]k [(][s][) = 0][,][ ∀][s][ ∈] _[S][. Then for]_
each timestep t in each trajectory from Dk, we update ρ[π]k [as][ ρ]k[π][(][s][t][)][ ←]− _ρ[π]k_ [(][s][t][) +][ γ][t][/][|D][k][|][.]

**Policy Representation:** The general approach depicted in Algorithm 1 allows various policy
representations including arrays and neural networks. Let _k_ represent a subset of states to
perform the policy update at the k-th iteration. When an array is used, the policy update step is simply S _⊆S_
_πk+1(·|s) = F(πk)(·|s), ∀s ∈Sk. When a neural network is employed, the policy update step can_

be achieved by obtaining the gradient descent, i.e., (F(πk)(·|s) − _πk(·|s))[2]._

_[∇]_ _sX∈Sk_

**Policy Updating Strategy:**

-  State space: For environments with a discrete state space (e.g., tabular domains), the
WPO/SPO policy update is performed for all states at each iteration. For the environments


-----

with a continuous state space, a random subset of states is sampled at each iteration to
perform the policy update.

-  Action space: For environments with a continuous action space, we first discretize the action
space following Tang & 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 6. In addition, we provide the setting of
hyperparamaters and network sizes of our WPO/SPO algorithms in Table 3. 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 4.

Table 3: Hyperparamaters and network sizes

Taxi-v3, NChain-v0 CartPole-v1 Acrobot-v1
CliffWalking-v0

_γ_ 0.9 0.95 0.95

_lrπ_ \ 10[−][2] 5 10[−][3]
_×_

_lrvalue_ 10[−][2] 10[−][2] 5 × 10[−][3]

_k_ 60 (Taxi); 1 (Chain); 2 3
_|D_ _|_
3 (CliffWalking)

_π size_ 2D array [64, 64] [64, 64]

Q/v size [10, 7, 5] [64, 64] [64, 64]


_λ_ 5, 50, 10 10 10

Table 4: Averaged rewards over last 10% episodes during the training process

Environment WPO SPO TRPO PPO A2C


Taxi-v3 _−45 ± 27_ _−87 ± 11_ _−202 ± 3_ _−381 ± 34_ _−338 ± 30_

NChain-v0 3549 ± 197 3432 ± 131 3522 ± 258 3506 ± 237 1606 ± 10

CliffWalking-v0 _−35 ± 15_ _−25 ± 1_ _−159 ± 94_ _−3290 ± 2106_ _−5587 ± 1942_

CartPole-v1 388 ± 54 370 ± 30 297 ± 65 193 ± 45 267 ± 61

Acrobot-v1 _−162 ± 8_ _−185 ± 15_ _−248 ± 33_ _−103 ± 5_ _−379 ± 39_

**Additional Experiments on Ablation Study:**

Figure 7: Episode rewards during the training process for different β and λ settings, averaged across
3 runs with a random initialization. The shaded area depicts the mean ± the standard deviation.


-----

