Year: 2019

Abstract:
We introduce a new approach for comparing reinforcement learning policies, using Wasserstein distances (WDs) in a newly defined latent behavioral space. We show that by utilizing the dual formulation of the WD, we can learn score functions over trajectories that can be in turn used to lead policy optimization towards (or away from) (un)desired behaviors. Combined with smoothed WDs, the dual formulation allows us to devise efficient algorithms that take stochastic gradient descent steps through WD regularizers. We incorporate these regularizers into two novel on-policy algorithms, Behavior-Guided Policy Gradient and Behavior-Guided Evolution Strategies, which we demonstrate can outperform existing methods in a variety of challenging environments. We also provide an open source demo1.

heading: 1 INTRODUCTION
text: One of the key challenges in reinforcement learning (RL) is to efficiently incorporate the behavioral characteristics of learned policies into optimization algorithms (Lee & Popovic, 2010; Meyerson et al., 2016; Conti et al., 2018). The fundamental question we aim to shed light on in this paper is:
What is the right measure of similarity between two policies acting on the same underlying MDP and how can we devise algorithms to leverage this information for reinforcement learning?
In simple terms, the main thesis motivating the methods we propose is that:
Two policies may perform similar actions at a local level but result in very different global behaviors.
We propose to define behaviors via so-called Behavioral Embedding Maps (BEMs), which are functions mapping trajectories (realizations of policies) into latent behavioral spaces representing trajectories in a compact way. BEMs enable us to identify policies with their Probabilistic Policy Embeddings (PPEs), which we define as the pushforward distributions over trajectory embeddings as a result of applying a BEM to a policy’s trajectories. Importantly, two policies with distinct distributions over trajectories may result in the same probabilistic embedding. PPEs provide us a way to rigorously define dissimilarity between policies. We do this by equipping them with metrics defined on the manifold of probabilistic measures, namely a class of Wasserstein distances (WDs, Villani (2008)). There are several reasons for choosing WDs:
• Flexibility. We can use any cost function between embeddings of trajectories, allowing the distance between PPEs to arise organically from an interpretable distance between embedding points.
• Non-injective BEMs. Different trajectories may be mapped to the same embedding point (for example in the case of the last-state embedding). This precludes the use of likelihood-based distances such as the KL divergence (Kullback & Leibler, 1951), which we discuss in Section 6.
• Behavioral Test Functions. Solving the dual formulation of the WD objective yields a pair of test functions over the space of embeddings that can be used to score trajectories.
The behavioral test functions underpin all our algorithms, directing optimization towards desired behaviors. To learn them it suffices to define the BEM and the cost function between points in the PPE space. To mitigate the computational burden of computing WDs, we rely on their entropy-regularized formulations. This allows us to update the learned test functions in a computationally efficient manner
1Available at https://github.com/behaviorguidedRL/BGRL. We emphasize this is not an exact replica of the code from our experiments, but a demo to build intuition and clarify our methods.
via stochastic gradient descent (SGD) on a Reproducing Kernel Hilbert Space (RKHS). We develop a novel method for stochastic optimal transport based on random feature maps (Rahimi & Recht, 2008) to produce compact and memory-efficient representations of learned behavioral test functions. Finally, having laid the groundwork for comparing trajectories via behavior-driven trajectory scores, we address our core question by introducing two new on-policy RL algorithms:
• Behavior Guided Policy Gradients (BGPG): We propose to replace the KL-based trust region from Schulman et al. (2015) with a WD-based trust region in the PPE space.
• Behavior Guided Evolution Strategies (BGES): Inspired by the NSR-ES algorithm from Conti et al. (2018), BGES jointly optimizes for reward and novelty using the WD in the PPE space.
In addition, we also demonstrate a way to harness our methodology for imitation learning (Section 7.3) and repulsion learning (Section 9.4), and we believe there may be many more potential applications in the future.

heading: 2 MOTIVATING BEHAVIOR-GUIDED REINFORCEMENT LEARNING
text: Throughout this paper we prompt the reader to think of a policy as a distribution over its trajectories, induced by the policy’s (possibly stochastic) map from state to actions and the unknown environment dynamics. We care about summarizing (or embedding) trajectories into succinct representations that can be compared with each other (via a cost/metric). These comparisons arise naturally when answering questions such as: Has a given trajectory achieved a certain level of reward? Has it visited a certain part of the state space? We think of these summaries or embeddings as characterizing the behavior of the trajectory. We formalize these notions in Section 3.
We show that by identifying policies with the embedding distributions that result of applying the embedding function (summary) to their trajectories, and combining this with the provided cost metric, we can induce a topology over the space of policies given by the WD over their embedding distributions. The methods we propose can be thought of as ways to leverage this “behavior” geometry for a variety of downstream applications such as policy optimization and imitation learning.
This topology emerges naturally from the sole definition of an embedding map (behavioral summary) and a cost function. Crucially these choices occur in the semantic space of behaviors as opposed to parameters or visitation frequencies2. One of the advantages of choosing a Wasserstein geometry is that non-surjective trajectory embedding maps are allowed. This is not possible with a KL induced one (in non-surjective cases, computing the likelihood ratios in the KL definition is in general intractable). In Sections 4 and 5 we show that in order to get a handle on this geometry we can use the dual formulation of the Wasserstein distance to learn functions (Behavioral Test Functions) that can provide scores on trajectories which then can be added to the reward signal (in policy optimization) or used as a reward (in Imitation Learning).
In summary, by defining an embedding map of trajectories into a behavior embedding space equipped with a metric3, our framework allows us to learn “reward” signals (Behavioral Test Functions) that can serve to steer policy search algorithms through the “behavior geometry” either in conjunction with a task specific reward (policy optimization) or on their own (e.g. Imitation Learning). We develop versions of on policy RL algorithms which we call Behavior Guided Policy Gradient (BGPG) and Behavior Guided Evolution Strategies (BGES) that enhance their baseline versions by the use of learned Behavioral Test Functions. Our experiments in Section 7 show this modification is useful. We also show how to use Behavioral Test Functions in Imitation Learning, where we only need access to an expert’s embedding. Although our framework also has obvious applications to safety, (learning policies that avoid undesirable or dangerous behaviors) we leave this for future work. We also consider simple heuristics for the embeddings (inspired by other existing use cases), but believe future work on learned embeddings could be a significant enhancement.
2If we choose an appropriate embedding map our framework handles visitation frequencies as well. 3The embedding space can be discrete or continuous and the metric need not be smooth, and can be for
example a simple discrete {0, 1} valued criterion

heading: 3 DEFINING BEHAVIOR IN REINFORCEMENT LEARNING
text: A Markov Decision Process (MDP) is a tuple (S,A,P,R). Here S and A stand for the sets of states and actions respectively, such that for s, s′ ∈ S and a ∈ A: P(s′|a, s) is the probability that the system/agent transitions from s to s′ given action a and R(s′, a, s) is a reward obtained by an agent transitioning from s to s′ via a. A policy πθ : S → A is a (possibly randomized) mapping (parameterized by θ ∈ Rd) from S to A. Let Γ = {τ = s0, a0, r0, · · · sH , aH , rH s.t. si ∈ S, ai ∈ A, ri ∈ R} be the set of possible trajectories enriched by sequences of partial rewards under some policy π. The undiscounted reward functionR : Γ→ R (which expectation is to be maximized by optimizing θ) satisfiesR(τ) = ∑H i=0 ri, where ri = R(si+1, ai, si).
3.1 BEHAVIORAL EMBEDDINGS
We start with a Behavioral Embeddng Space (BES) which we denote as E and a Behavioral Embedding Map (BEM), Φ : Γ → E , mapping trajectories to embeddings in E (Fig. 1). Importantly, the mapping does not need to be surjective. We will provide examples of BESs and BEMs at the end of the section. Given a policy π, we let Pπ denote the distribution induced over the spaces of trajectories Γ and by PΦπ the corresponding pushforward distribution on E induced by Φ. We call PΦπ the Probabilistic Policy Embedding (PPE) of a policy π. A policy π can be fully characterized by the distribution Pπ .
Additionally, we require the BES E to be equipped with a metric (or cost function) C : E × E → R. Given two trajectories τ1, τ2 in Γ, C(Φ(τ1),Φ(τ2)) measures how different these trajectories are in the behavior space. The following are examples of BEMs (with the corresponding BESs) categorized into three main types (we will use examples from all three types in our experiments in Section 7): 1. State-based: the final state Φ1(τ) = sH , the visiting frequency of a fixed state Φs2(τ) =∑H t=0 1(st = s), the frequency vector of visited states Φ3(τ) = ∑H t=0 est (where es ∈ R|S|
is the one-hot vector corresponding to state s); see also Section 7.2. 2. Action-based: the concatenation of actions Φ4(τ) = [a0, ..., aH ]; see also Section 7.1. 3. Reward-based: the total reward Φ5(τ) = ∑H t=0 rt, reward-to-go vector Φ6(τ) =∑H
t=0 rt (∑t i=0 ei ) (where ei ∈ RH+1 is a one-hot vector corresponding to i and with dimensions
indexed from 0 to H); see also Section 7.1 and Section 7.3.
For instance, PΦ3π is the frequency with which different states are visited under π. Note that some of the above embeddings are only for the tabular case (|S|, |A| <∞) while others are universal.

heading: 4 WASSERSTEIN DISTANCE & OPTIMAL TRANSPORT PROBLEM
text: Let µ, ν be (Radon) probability measures over domains X ⊆ Rm,Y ⊆ Rn and let C : X × Y → R be a cost function. For γ > 0, a smoothed Wasserstein Distance is defined as:
WDγ(µ, ν) := min π∈Π(µ,ν) ∫ X×Y C(x,y)dπ(x,y) + γKL(π|ξ), (1)
where Π(µ, ν) is the space of couplings (joint distributions) over X ×Y with marginal distributions µ and ν, KL(·|·) denotes the KL divergence between distributions π and ρ with support X × Y defined as: KL(π|ρ) = ∫ X×Y ( log ( dπ dξ (x,y) )) dπ(x,y) and ξ is a reference measure over X × Y . When the cost is an `p distance and γ = 0, WDγ is also known as the Earth mover’s distance and the corresponding optimization problem is known as the optimal transport problem (OTP).

heading: 4.1 WASSERSTEIN DISTANCE: DUAL FORMULATION
text: We will use smoothed WDs to derive efficient regularizers for RL algorithms. To arrive at this goal, we first need to consider the dual form of Equation 1. Under the subspace topology (Bourbaki,
Algorithm 1 Random Features Wasserstein SGD Input: kernels κ, ` over X ,Y respectively with corresponding random feature maps φκ, φ`, smoothing parameter γ, gradient step size α, number of optimization rounds M , initial dual vectors pµ0 ,p ν 0 . for t = 0, · · · ,M do 1. Sample (xt, yt) ∼ µ ⊗ ν.
2. Update ( pµt pνt ) = (pµt−1 pνt−1 ) + α√ t ( 1− exp ( (pµt−1) >φκ(xt)−(pνt−1) >φ`(xt)−C(xt,yt) γ ))( φκ(xt) −φ`(yt) ) Return: pµM ,p ν M .
1966) for X and Y , let C(X ) denote the space of continuous functions on X and let C(Y) denote the space of continuous functions over Y . The choice of the subspace topology ensures our discussion encompasses the discrete case.
Let C : X × Y → R be a cost function, interpreted as the “ground cost” to move a unit of mass from x to y. Define I as the (0,∞) indicator function, where the value 0 denotes set membership. Using Fenchel duality, we can obtain the following dual formulation of the problem in Eq. 1:
WDγ(µ, ν) = max λµ∈C(X ),λν∈C(Y) ∫ X λµ(x)dµ(x)− ∫ Y λν(y)dν(y)− EC(λµ, λν), (2)
where EC(λµ, λν) is defined as:
EC(λµ, λν) :=
{ γ ∫ X×Y exp ( λµ(x)−λν(y)−C(x,y) γ ) dξ(x,y) if γ > 0
I((λν , λν) ∈ {(u, v) s.t. ∀(x,y) ∈ X × Y u(x)− v(y) ≤ C(x,y)}) if γ = 0. (3)
We will set dξ(x,y) ∝ 1 for discrete domains and dξ(x,y) = dµ(x)dν(y) otherwise. If λ∗µ, λ ∗ ν are the functions achieving the maximum in Eq. 2, and γ is sufficiently small then
WDγ(µ, ν) ≈ Eµ [ λ∗µ(x) ] − Eν [λ∗ν(y)], with equality when γ = 0. When for example γ = 0, X = Y , and C(x, x) = 0 for all x ∈ X , it is easy to see λ∗µ(x) = λ∗ν(x) = λ∗(x) for all x ∈ X . In this case the difference between Eµ [λ∗(x)] and Eµ [λ∗(y)] equals the WD. In other words, the function λ∗ gives higher scores to regions of the space X where µ has more mass. This observation is key to the success of our algorithms in guiding optimization towards desired behaviors.
4.2 COMPUTING λ∗µ AND λ ∗ ν
We combine several techniques to make the optimization of objective from Eq. 2 tractable. First, we replace X and Y with the functions from a RKHS corresponding to universal kernels (Micchelli et al., 2006). This is justified since those function classes are dense in the set of continuous functions of their ambient spaces. In this paper we choose the Gaussian kernel and approximate it using random Fourier feature maps (Rahimi & Recht, 2008) to increase efficiency. Consequently, the functions λ learned by our algorithms have the following form: λ(x) = (pλ)>φ(x), where φ is a random feature map with m standing for the number of random features and pλ ∈
Rm. For the Gaussian kernel, φ is defined as follows: φ(z) = 1√
m cos(Gz + b) for z ∈ Rd,
where G ∈ Rm×d is Gaussian with iid entries taken from N (0, 1), b ∈ Rm with iid bis such that bi ∼ Unif[0, 2π] and the cos function acts elementwise. Henceforth, when we refer to optimization over λ, we mean optimizing over corresponding dual vectors pλ associated with λ. We can solve for the optimal dual functions by performing SGD over the dual objective in Eq. 2. Algorithm 1 is the random features equivalent of Algorithm 3 in Genevay et al. (2016) and will be a prominent subroutine of our methods. An explanation and proof of why this is the right stochastic gradient is in Lemma 10.2 in the Appendix.
If pµ∗ ,pν∗ are the optimal dual vectors and (x1, y1), · · · , (xk, yk) i.i.d∼ µ
⊗ ν, then Algorithm 1 can be
used to get an estimator of WDγ(µ, ν) as follows:
ŴDγ(µ, ν) = 1
k k∑ i=1 〈pµ∗ , φκ(xi)〉 − 〈pν∗ , φ`(yi)〉+ 1 γ exp ( φκ(xi) >pµ∗ − φ`(yi)>pν∗ − C(xi, yi) γ ) (4)

heading: 5 BEHAVIOR-GUIDED REINFORCEMENT LEARNING
text: Here we introduce the framework which allows us to incorporate our behavioral approach to reinforcement learning into practical on-policy algorithms. Denote by πθ a policy parameterized by θ ∈ Rd. The goal of policy optimization algorithms is to find a policy maximizing, as a function of the policy parameters, the expected total reward L(θ) := Eτ∼Pπθ [R(τ)].

heading: 5.1 BEHAVIORAL TEST FUNCTIONS
text: If C : E × E → R is a cost function defined over behavior space E , and π1, π2 are two policies, then:
WDγ(PΦπ1 ,P Φ π2) ≈ Eτ∼Pπ1 [λ ∗ 1(Φ(τ))]− Eτ∼Pπ2 [λ ∗ 2(Φ(τ))] , (5)
where λ∗1, λ ∗ 2 are the optimal dual functions. The maps s1 := λ ∗ 1 ◦ Φ : Γ → R and s2 := λ∗2 ◦ Φ : Γ→ R define score functions over the space of trajectories. If γ is close to zero, the score function si gives higher scores to trajectories from πi whose behavioral embedding is common under πi but rarely appears under πj for j 6= i (Fig. 2).

heading: 5.2 ALGORITHMS
text: We propose to solve a WD-regularized objective to tackle behavior-guided policy optimization. All of our algorithms hinge on trying to maximize an objective of the form:
F (θ) = L(θ) + βWDγ(PΦπθ ,P Φ b ), (6)
where PΦb is a base distribution over behavioral embeddings (possibly dependent on θ) and β ∈ R could be positive or negative. Although the base distribution PΦb could be arbitrary, our algorithms will instantiate PΦb = 1 |S| ∪π′∈S P Φ π′ for some family of policies S (possibly satisfying |S| = 1) we want the optimization to attract to / repel from.
In order to compute approximate gradients for F , we rely on the dual formulation of the WD. After substituting the composition maps resulting from Eq. 5 into Eq. 2, we obtain:
F (θ) ≈ Eτ∼Pπθ [R(τ) + βs1(τ)]− βEφ∼PΦb [λ ∗ 2(φ)] , (7)
where s1 : Γ→ R equals s1 = λ∗1◦Φ, the Behavioral Test Function of policy πθ and λ∗2 is the optimal dual function of embedding distribution PΦb . Consequently ∇θF (θ) ≈ ∇θEτ∼Pπθ [R(τ) + βs1(τ)]. We learn a score function s1 over trajectories that can guide our optimization by favoring those trajectories that show desired global behaviors.
Eq. 7 is an approximation to the true objective from Eq. 2 whenever γ > 0. In practice, the entropy regularization requires a damping term as defined in Equation 3. If ξ(PΦπθ ,P Φ b ) is the joint distribution of choice then F (θ) = L(θ) + βV for
V = max λπθ∈C(E),λb∈C(E) E τ∼Pπθ [λπθ (Φ(τ))]− E φ∼PΦb [λb(φ)] + γ E φ1,φ2∼ξ(PΦπθ ,P Φ b ) [Λ(φ1, φ2)] ,
where Λ(φ1, φ2) = exp ( λπθ (φ1)−λb(φ2)−C(φ1,φ2)
γ
) . When the embedding space E is not discrete
and PΦb = PΦπ for some policy π, we let ξ(PΦπθ ,P Φ b ) = PΦπθ ⊗ PΦπ , otherwise ξ(PΦπθ ,P Φ b ) = 1 |E|2 1, a uniform distribution over E × E . All of our methods perform a version of alternating SGD optimization: we take certain number of SGD steps over the internal dual Wasserstein objective, followed by more SGD steps over the outer objective having fixed the current dual functions. Although in practice the different components
that make up the optimization objectives we consider here could be highly nonconvex, in the cases these functions satisfy some convexity assumptions, we can provide a sharp characterization for the convergence rates of our algorithms. Details are given in Section 10 in the Appendix.
We consider two distinct approaches to optimizing this objective, by exploring in the action space and backpropagating, as in policy gradient methods (Schulman et al., 2015; 2017), and by considering a black-box optimization problem as in Evolution Strategies (ES, Salimans et al. (2017)). These two different approaches lead to two new algorithms: Behavior-Guided Policy Gradient (BGPG) and Behavior-Guided Evolution Strategies (BGES), that we discuss next.

heading: 5.3 BEHAVIOR-GUIDED POLICY GRADIENT (BGPG)
text: Our first algorithm seeks to solve the optimization problem in Section 5.2 with policy gradients. We refer to this method as the Behavior-Guided Policy Gradient (BGPG) algorithm (see Algorithm 2 below).
Algorithm 2 Behvaior-Guided Policy Gradient Input: Initialize stochastic policy π0 parametrized by θ0, β < 0, η > 0, M,L ∈ N for t = 1, . . . , T do
1. Run πt−1 in the environment to get advantage values Aπt−1(s, a) and trajectories {τ (t)i }Mi=1 2. Update policy and test functions via several alternating gradient steps over the objective:
F (θ) = E τ1,τ2∼Pπt−1 ⊗ Pπθ [ H∑ i=1 Aπt−1(si, ai) πθ(ai|si) πt−1(ai|si) + βλ1(Φ(τ1))
− βλ2(Φ(τ2)) + βγ exp ( λ1(Φ(τ1))− λ2(Φ(τ2))− C(Φ(τ1)),Φ(τ2))
γ )] Where τ1 = s0, a0, r0, · · · , sH , aH , rH . Let θ(0)t−1 = θt−1. for ` = 1, · · · , L do
a. Approximate Pπt−1 ⊗ Pπθ via 1M {τ (t) i }Mi=1 ⊗ 1 M {τ θ i }Mi=1 := P̂πt,πθ where τθi
i.i.d∼ Pπθ b. Take SGA step θ(`)t−1 = θ (`−1) t−1 + η∇̂θF̂ (θ (`−1) t−1 ) using samples from P̂πt−1,πθ . c. Use samples from P̂πt−1,πθ and Algorithm 1 to update λ1, λ2.
Set θt = θ (M) t−1 .
Specifically, we maintain a stochastic policy πθ and compute policy gradients as in prior work (Schulman et al., 2015). To optimize the Wasserstein distance WDγ , we approximate the gradient of this term via the random-feature Wasserstein SGD . Importantly, this stochastic gradient can be approximated by samples collected from the policy πθ. In its simplest form, the ∇̂θF̂ in Step b. in Algorithm 2 can be computed by the vanilla policy gradient over the advantage component and using the reinforce estimator through the components involving Behavioral Test Functions acting on trajectories from Pπθ . We explain in Appendix 8.1 a lower-variance gradient estimator alternative.
BGPG can be thought of as a variant of Trust Region Policy Optimization with a Wasserstein penalty. As opposed to vanilla TRPO, the optimization path of BGPG flows through policy parameter space while encouraging it to follow a smooth trajectory through the geometry of the PPE space. We proceed to show that given the right embedding and cost function, we can prove a monotonic improvement theorem for BGPG, showing that our methods satisfy at least similar guarantees as TRPO.
For a given policy π, we denote as: V π, Qπ and Aπ(s, a) = Qπ(s, a)− V π(s) the: value function, Q-function and advantage function (see Appendix: Section 10.5). Furthermore, let V (π) be the expected reward of policy π and ρπ(s) = Eτ∼Pπ [∑T t=0 1(st = s) ] be the visitation measure.
Two distinct policies π and π̃ can be related via the equation (see: Sutton et al. (1998)) V (π̃) = V (π)+ ∫ S ρπ̃(s) (∫ A π̃(a|s)A π(s, a)da ) ds and the linear approximations to V around π via: L(π̃) =
V (π) + ∫ S ρπ(s) (∫ A π̃(a|s)A π(s, a)da ) ds (see: Kakade & Langford (2002)). Let S be a finite
set. Consider the following embedding Φs : Γ→ R|S| defined by (Φ(τ))s = ∑T t=0 1(st = s) and related cost function defined as: C(v,w) = ‖v−w‖1. Then WD0(PΦ s π̃ ,PΦ s
π ) is related to visitation frequencies since WD0(PΦ s π̃ ,PΦ s π ) ≥ ∑ s∈S |ρπ(s)− ρπ̃(s)| (see Section 10.5 for the proof). These observations enable us to prove an analogue of Theorem 1 from Schulman et al. (2015), namely:
Theorem 5.1. If WD0(PΦ s π̃ ,PΦ s π ) ≤ δ and = maxs,a |Aπ(s, a)|, then V (π̃) ≥ L(θ̃)− δ .
As in Schulman et al. (2015), Theorem 5.1 implies a policy improvement guarantee for BGPG.

heading: 5.4 BEHAVIOR GUIDED EVOLUTION STRATEGIES (BGES)
text: ES takes a black-box optimization approach to RL, by considering a rollout of a policy, parameterized by θ as a black-box function F . This approach has gained in popularity recently (Salimans et al., 2017; Mania et al., 2018; Choromanski et al., 2019). If we take this approach to optimizing the objective in Eq. 2, the result is a black-box optimization algorithm which seeks to maximize the reward and simultaneously maximizes or minimizes the difference in behavior from the base embedding distribution PΦb . We call this method the Behavior-Guided Evolution Strategies (BGES) algorithm (see Algorithm 3 below).
Algorithm 3 Behavior-Guided Evolution Strategies Input: learning rate η, noise standard deviation σ, iterations T , BEM Φ, β Initialize: Initial policy π0 parametrized by θ0, Behavioral Test Functions λ1, λ2. Evaluate policy π0 to return trajectory τ0 and subsequently use the BEM to produce an initial PPE P̂Φπ0 . for t = 1, . . . , T − 1 do
1. Sample 1, · · · , n independently from N (0, I). 2. Evaluate policies {πkt }nk=1 parameterized by {θt + σ k}nk=1 to return rewards Rk and trajectories τk for all k. 3. Use BEM to map trajectories τk to produce empirical PPEs P̂Φπkt for all k. 4. Update λ1 and λ2 using Algorithm 1, where µ = 1n ∪
n k=1 P̂Φπkt−1 and ν = 1 n ∪ n k=1 P̂Φπkt are the
uniform distribution over the set of PPEs from 3 for t− 1 and t. 5. Approximate ŴDγ(PΦ
πkt ,PΦπt) plugging in λ1, λ2 into Eq. 4 for each perturbed policy πk
6. Update Policy: θt+1 = θt + η∇ESF , where:
∇ESF = 1
σ n∑ k=1 [(1− β)(Rk −Rt) + βŴDγ(PΦπkt ,P Φ πt)] k
When β > 0, and we take PΦb = PΦπt−1 , BGES resembles the NSR-ES algorithm from Conti et al. (2018), an instantiation of novelty search (Lehman & Stanley, 2008). The positive weight on the WD-term enforces newly constructed policies to be behaviorally different from the previous ones (improving exploration) while the R−term drives the optimization to achieve its main objective, i.e., maximize the reward. The key difference in our approach is the probabilistic embedding map, with WD rather than Euclidean distance. We show in Section 7.2 that BGES outperforms NSR-ES for challenging exploration tasks. The approximation introduced by Step 5 bypasses the need of computing a different pair of behavioral test functions λ1, λ2 for each perturbed policy πk.
If we take β < 0, and assume PΦb = PΦπ to correspond to embedded trajectories from an oracle or expert policy, we can perform imitation learning. Despite not accessing the expert’s policy (just the trajectories it generates), we show in Section 7.3 that this approach dramatically improves learning.

heading: 6 RELATED WORK
text: Our work is related to research in multiple areas in neuroevolution and machine learning: Behavior Characterizations: The idea of directly optimizing for behavioral diversity was introduced by Lehman & Stanley (2008) and Lehman (2012), who proposed to search directly for novelty,
rather than simply assuming it would naturally arise in the process of optimizing an objective function. This approach has been applied to deep RL (Conti et al., 2018) and meta-learning (Gajewski et al., 2019). In all of this work, the policy is represented via a behavioral characterization (BC), typically chosen with knowledge of the environment, for example the final (x,y) coordinate for a locomotion task. Additionally, in most cases these BCs are considered to be deterministic, with Euclidean distances used to compare BCs. In our setting, we move from deterministic BCs to stochastic PPEs, thus requiring the use of metrics capable of comparing probabilistic distributions.
Distance Metrics: WDs have been used in many different applications in machine learning where guarantees based on distributional similarity are required (Jiang et al., 2019; Arjovsky et al., 2017). We make use of WDs in our setting for a variety of reasons. First and foremost, the dual formulation of the WD allows us to recover Behavioral Test Functions, thus providing us with behavior-driven trajectory scores. In contrast to KL divergences, WDs are sensitive to user-defined costs between pairs of samples instead of relying only on likelihood ratios. Furthermore, as opposed to KL divergences, it is possible to take SGD steps using entropy-regularized Wasserstein objectives. Computing an estimator of the KL divergence is hard without a density model. Since in our framework multiple unknown trajectories may map to the same behavioral embedding, the likelihood ratio between two embedding distributions may be ill-defined.
WDs for RL: We are not the first to propose using WDs in RL. Zhang et al. (2018) have recently introduced Wasserstein Gradient Flows (WGFs) for finding efficient RL policies. This approach casts policy optimization as gradient descent flow on the manifold of corresponding probability measures, where geodesic lengths are given as second-order WDs. We note that computing WGFs is a nontrivial task. In Zhang et al. (2018) this is done via particle approximation methods. We show in Section 7 that RL algorithms using these techniques are substantially slower than our methods. The WD has also been employed to replace KL terms in standard Trust Region Policy Optimization (Richemond & Maginnis, 2017). This is a very special case of our more generic framework (cf. Section 5.2). In Richemond & Maginnis (2017) it is suggested to solve the corresponding RL problems via FokkerPlanck equations and diffusion processes, yet no empirical evidence of the feasibility of this approach is provided. We propose general practical algorithms and provide extensive empirical evaluation.
Distributional RL Distributional RL (DRL, Bellemare et al. (2017)) expands on traditional offpolicy methods (Mnih et al., 2013) by attempting to learn a distribution of the return from a given state, rather than just the expected value. These approaches have impressive experimental results (Bellemare et al., 2017; Dabney et al., 2018), with a growing body of theory (Rowland et al., 2018; Qu et al., 2019; Bellemare et al., 2019; Rowland et al., 2019). Superficially it may seem that learning a distribution of returns is similar to our approach to PPEs, when the BEM is a distribution over rewards. Indeed, reward-driven embeddings used in DRL can be thought of as special cases of the general class of BEMs. We note two key differences: 1) DRL methods are off-policy whereas our BGES and BGPG algorithms are on-policy, and 2) DRL is typically designed for discrete domains, since Q-Learning with continuous action spaces is generally much harder. Furthermore, we note that while the WD is used in DRL, it is only for the convergence analysis of the DRL algorithm—the algorithm itself does not use WDs (Bellemare et al., 2017).

heading: 7 EXPERIMENTS
text: Here we seek to test whether our behavior-guided approach to RL translates to performance gains for simulated environments. We individually evaluate our two proposed algorithms, BGPG and BGES, versus their respective baselines for a range of benchmark tasks. While in some cases the results may not be state of the art, we believe the improvement vs. popular RL algorithms (in particular TRPO and ES) are exciting results which could stimulate future work. We also include a study of using our method for imitation learning. For each subsection we provide additional details in the Appendix.

heading: 7.1 BEHAVIOR-GUIDED POLICY GRADIENT
text: Our key question is whether BGPG can outperform baseline TRPO methods using KL divergence. In Fig. 3, we see this is clearly the case for four continuous control tasks: Pendulum from OpenAI Gym and Hopper: Stand, Hooper: Hop and Walker: Stand from the DeepMind Control Suite (Tassa et al., 2018). For the BEM, we use the concatenation-of-actions (as used already in TRPO). We also
confirm results from (Schulman et al., 2015) that a trust region greatly improves performance, as we see the black curve (without one) often fails to learn.
Wall Clock Time: To illustrate computational benefits of alternating optimization (AO) of WD in BGPG, we compare it to the particle approximation (PA) method introduced in Zhang et al. (2018) in Fig. 4. In practice, the WD across different state samples can be optimized in a batched manner using AO (see Appendix for details). We see that AO is substantially faster than PA.

heading: 7.2 BEHAVIOR-GUIDED EVOLUTION STRATEGIES
text: As a novelty-search method, BGES is designed to actively explore the environment by behaving differently for previous policies. With that in mind, we seek to evaluate the ability to solve two key challenges in exploration for RL: deceptive rewards and local maxima.
Deceptive Rewards A common challenge in model-free RL is deceptive rewards. These arise since agents can only learn from data gathered via exploration in the environment. To test BGES in this setting, we created two intentionally deceptive environments where agents may easily be fooled into learning suboptimal policies. In both cases the agent is penalized at each time step for being far away from a goal. The deception comes from a wall situated in the middle, which means that initially positive rewards from moving directly forward will lead to a suboptimal policy.
We consider two types of agents—a two-dimensional point and a much larger quadruped. Details are provided in the Appendix (Section 9). We compare with state-of-the-art on-policy methods for efficient exploration: NSR-ES from (Conti et al., 2018), which assumes the BEM is deterministic and uses the Euclidean distance to compare policies, and NoisyNet-TRPO from Fortunato et al. (2018). Results are presented on Fig. 5. Policies avoiding the wall correspond to rewards: R > −5000 and R > −800 for the quadruped and point respectively. In the prior case an agent needs to first learn how to walk and the presence of the wall is enough to prohibit vanilla ES from even learning forward locomotion. We note that BGES is the only method that drives the agent to the goal in both settings. For the quadruped the BEM is the reward-to-go while for the point we used the final state.
Escaping Local Maxima. In Fig. 6 we compare our methods with methods using regularizers based on other distances or divergences (specifically, Hellinger, Jensen-Shannon (JS), KL and Total Variation (TV) distances), as well as vanilla ES (i.e., with no distance regularizer). Experiments were performed on a Swimmer environment from OpenAI Gym (Brockman et al., 2016), where the number of samples of the ES optimizer was drastically reduced. BGES is the only one that manages to obtain good policies which also proves that the benefits come here not just from introducing the regularizer, but from its particular form.
7.3 IMITATION LEARNING
As discussed in Section 5.3, we can also utilize the BGES algorithm for imitation learning, by setting β < 0, and using an expert’s trajectories for the PPE. For this experiment we use the reward-to-go BEM (Section 5). In Fig. 7, we show that this approach significantly outperforms vanilla ES on the Swimmer task. Although conceptually simple, we believe this could be a powerful approach with potential extensions, for example in designing safer algorithms.
7.4 HYPERPARAMETER SELECTION
Our approach includes several new hyperparameters, such as the kernel for the Behavioral Test Functions and the choice of BEM. For our experiments we did not perform any hyperparameter optimization. We only considered the rbf kernel, and only varied the BEM for BGES. For BGPG we chose the concatenation of actions, since this is the same as used in the KL divergence for
TRPO. For BGES, we demonstrated several different BEMs, and we show an ablation study for the point agent in Fig. 8 where we see that both the reward-to-go (RTG) and Final State (SF) worked, but the vector of all states (SV) did not (for 5 seeds). We leave learned BEMs as exciting future work.

heading: 8 CONCLUSION AND FUTURE WORK
text: In this paper we proposed a new paradigm for on-policy learning in RL, where policies are embedded into expressive latent behavioral spaces and the optimization is conducted by utilizing the repelling/attraction signals in the corresponding probabilistic distribution spaces. The use of Wasserstein distances (WDs) guarantees flexibility in choosing cost funtions between embedded policy trajectories, enables stochastic gradient steps through corresponding regularized objectives (as opposed to KL divergence methods) and provides an elegant method, via their dual formulations, to quantify behaviorial difference of policies through the behavioral test functions. Furthermore, the dual formulations give rise to efficient algorithms optimizing RL objectives regularized with WDs.
We also believe the presented methods shed new light on several other challenging problems of modern RL, including: learning with safety guarantees (a repelling signal can be used to enforce behaviors away from dangerous ones) or anomaly detection for reinforcement learning agents (via the above score functions). We are also excited by the possibility of scaling this approach to a population setting, learning the behavioral embedding maps from data, or adapting the degree of repulsion/attraction during optimization (parameter β).

heading: APPENDIX: BEHAVIOR-GUIDED REINFORCEMENT LEARNING

heading: 9 FURTHER EXPERIMENTAL DETAILS

heading: 9.1 BGPG
text: A Lower-variance Gradient Estimator: As explained in Section 5.2, the BGPG considers an objective which involves two parts: the conventional surrogate loss function for policy optimization (Schulman et al., 2017), and a loss function that involves the Behavior Test Functions. Though we could apply vanilla reinforced gradients on both parts, it is straightforward to notice that the second part can be optimized with reparameterized gradients (Kingma & Welling, 2013), which arguably have lower variance compared to the reinforced gradients. In particular, we note that under random feature approximation (4), as well as the action-concatenation embedding, the Wasserstein distance loss ŴDγ(PΦπθ , P Φ b ) is a differentiable function of θ. To see this more clearly, notice that under a Gaussian policy a ∼ N (µθ(s), σθ(s)2) the actions a = µθ(s) + σθ(s) · are reparametrizable for being standard Gaussian noises. We can directly apply the reparametrization trick to this second objective to obtain a gradient estimator with potentially much lower variance. In our experiments, we applied this lower-variance gradient estimator.
Trust Region Policy Optimization: Though the original TRPO (Schulman et al., 2015) construct the trust region based on KL-divergence, we propose to construct the trust region with WD. For convenience, we adopt a dual formulation of the trust region method and aim to optimize the augmented objective Eτ∼πθ [R(τ)] − βWDγ(PΦπ′ ,PΦπθ ). We apply the concatenation-of-actions embedding and random feature maps to calculate the trust region. We identify several important hyperparameters: the RKHS (for the test function) is produced by RBF kernel k(x, y) = exp(‖x− y‖22/σ2) with σ = 0.1; the number of random features is D = 100; recall the embedding is Φ(τ) = [a1, a2...aH ] where H is the horizon of the trajectory, here we take 10 actions per state and embed them together, this is equivalent to reducing the variance of the gradient estimator by increasing the sample size; the regularized entropy coefficient in the WD definition as γ = 0.1; the trust region trade-off constant β ∈ {0.1, 1, 10}. The alternate gradient descent is carried out with T = 100 alternating steps and test function coefficients p ∈ RD are updated with learning rate αp = 0.01.
The baseline algorithms are: No trust region, and trust region with KL-divergence. The KL-divergence is identified by a maximum KL-divergence threshold per update, which we set to = 0.01.
Across all algorithms, we adopt the open source implementation (Dhariwal et al., 2017). Hyperparameters such as number of time steps per update as well as implementation techniques such as state normalization are default in the original code base.
The additional experiment results can be found in Figure 9 where we show comparison on additional continuous control benchmarks: Tasks with DM are from DeepMind Contol Suites (Tassa et al., 2018). We see that the trust region constructed from the WD consistently outperforms other baselines (importantly, trust region methods are always better than the baseline without trust region, this confirms that trust region methods are critical in stabilizing the updates).
Wasserstein AO vs. Particle Approximation: To calculate the regularized Wasserstein distance, we propose a gradient descent method that iteratively updates the test function. The alternting optimization (AO) scheme consists of updating both the test function and the distribution parameters such that the regularized Wasserstein distance of the trainable distribution against the reference distribution is minimized. Alternatively, we can also adopt a particle approximation method to calculate the Wasserstein distance and update the distribution parameters using an approximate gradient descent method (Zhang et al., 2018).
One major advantage of AO against particle approximation is its ease of parallization. In particular, when using the concatenation-of-actions embedding, the aggregate Wasserstein distance can be decomposeed into an average of a set of Wasserstein distances over states. To calculate this aggregated gradient, AO can easily leverage the matrix multiplication; on the other hand, particle approximation requires that the dual optimal variables of each subproblem be computed, which is not straightforward to parallelize.
We test both methods in the context of trust region policy search, in which we explicitly calculate the Waserstein distance of consecutive policies and enforce the constraints using a line search as in (Schulman et al., 2015). Both methods require the trust region trade-off parameter β ∈ {0.1, 1, 10}. We adopt the particle method in (Zhang et al., 2018) where for each state there are M = 16 particles. The gradients are derived based a RKHS where we adaptively adjust the coefficient of the RBF kernel based on the mean distance between particles. For the AO, we find that it suffices to carry out T ∈ {1, 5, 10} gradient descents to approximate the regularized Wasserstein distance.

heading: 9.2 BGES
text: Efficient Exploration: To demonstrate the effectiveness of our method in exploring deceptive environments, we constructed two new environments using the MuJoCo simulator. For the point environment, we have a 6 dimensional state and 2 dimensional action, with the reward at each timestep calculated as the distance between the agent and the goal. We use a horizon of 50 which is sufficient to reach the goal. The quadruped environment is based on Ant from the Open AI Gym (Brockman et al., 2016), and has a similar reward structure to the point environment but a much larger state space (113) and action space (8). For the quadruped, we use a horizon length of 400.
To leverage the trivially parallelizable nature of ES algorithms, we use the ray library, and distribute the rollouts across 72 workers using AWS. Since we are sampling from an isotropic Gaussian, we are able to pass only the seed to the workers, as in Salimans et al. (2017). However we do need to return trajectory information to the master worker.
For both the point and quadruped agents, we use random features with dimensionality m = 1000, and 100 warm-start updates for the WD at each iteration. For point, we use the final state embedding, learning rate η = 0.1 and σ = 0.01. For the quadruped, we use the reward-to-go embedding, as we found this was needed to learn locomotion, as well as a learning rate of η = 0.02 and σ = 0.02. The hyper-parameters were the same for all ES algorithms. When computing the WD, we used the previous 2 policies, θt−1 and θt−2. An ablation study for the point environment for both the choice of embedding and number of prior policies is shown in Fig 12.
For embeddings, we compare the reward-to-go (RTG), concatenation of states (SV) and final state (SF). In both the RTG and SF case the agent learns to navigate past the wall (> −800). For the number of previous policies, we use the SF embedding, and using 2 appears to work best, but both 1 and 5 do learn the correct behavior.
Escaping Local Maxima: We also demonstrated that our method leads to faster training even in more standard settings, where exploration is not that crucial, but the optimization can be trapped in local maxima. To show it, we compared baseline ES algorithm for ES optimization from Salimans et al. (2017) with its enhancements, where regularizers using different metrics on the space of probabilistic distributions corresponding to policy embeddings were used, as in the previous paragraph. We noticed that adding Wasserstein regularizers drastically improved optimization, whereas regularizers based on other distances/divergencies, namely: Hellinger, Jensen-Shannon, KL and TV did not have any impact. We considered Swimmer task from OpenAI Gym and to make it challenging, reduced the number of perturbations per iteration to 80. In that setting our method was the only one that was not trapped in local maxima and managed to learn effective policies.

heading: 9.3 IMITATION LEARNING:
text: For the Imitation Learning experiment we used the reward-to-go embedding, with learning rate η = 0.1 and σ = 0.01. We use one oracle policy, which achieves > 360 on the environment. The only information provided to the algorithm is the embedded trajectory, used to compute the WD. This has exciting future applications since no additional information about the oracle is required in order to significantly improve learning.

heading: 9.4 REPULSION LEARNING
text: Algorithm 4 Behvaior-Guided Repulsion Learning Input: β, η > 0, M ∈ N Initialize: Initial stochastic policies πa0 , πb0 , parametrized by θa0 , θb0 respectively, Behavioral Test Functions λa1 , λ b 2 for t = 1, . . . , T do 1. Collect M trajectories {τai }Mi=1 from Pπat−1 and M trajectories {τ b i }Mi=1 from Pπbt−1 .
Approximate Pπat−1 ⊗ Pπbt−1 via 1 M {τ a i }Mi=1 ⊗ 1 M {τ b i }Mi=1 := P̂πat−1,πbt−1 2. Form two distinct surrogate rewards for joint trajectories of agents a and b: R̃a(τ1, τ2) = R(τ1) + βλa1(Φ(τ1)) + βγ exp ( λa1(Φ(τ1))− λb2 (Φ(τ2))− C(Φ(τ1)),Φ(τ2))
γ ) R̃b(τ1, τ2) = R(τ2)− βλb2 (Φ(τ2)) + βγ exp ( λa1(Φ(τ1))− λb2 (Φ(τ2))− C(Φ(τ1)),Φ(τ2))
γ ) 3. For c ∈ {a,b} use the Reinforce estimator to take gradient steps:
θct = θ c t−1 + η E
τa,τb∼P̂ πa t−1,π b t−1
[ R̃c(τa, τb) ( H−1∑ i=0 ∇θct−1 log ( πct−1(a c i |sci ) ))]
Where τa = sa0 , a a 0 , r a 0 , · · · , saH , aaH , raH and τb = sb0 , ab0 , rb0 , · · · , sbH , abH , rbH . 5. Use samples from P̂πat−1,πbt−1 and Algorithm 1 to update the Behavioral Test Functions λ a 1 , λ b 2 .
Although this was not discussed in the main section of the paper, it is possible to use our behavioral approach to simultaneously learn multiple policies exhibiting different behaviors all of which are able to solve the same task. This is not the main focus of the main paper, but we chose to include these results in an attempt to provide the readers with a better understanding of Behavioral Test functions.
Algorithm 4 maintains two policies πa and πb. Each policy is optimized by taking a policy gradient step (using the Reinforce gradient estimator) in the direction optimizing surrogate rewards R̃a and R̃b that combines the signal from the task’s reward functionR and the repulsion score encoded by the behavioral test functions λa and λb.
We conducted experiments testing Algorithm 4 on a simple Mujoco environment consisting of a particle that moves on the plane and whose objective is to learn a policy that allows it to reach one of two goals. Each policy outputs a velocity vector and stochasticity is achieved by adding Gaussian noise to the mean velocity encoded by a neural network with two size 5 hidden layers and ReLu activations. If an agent performs action a at state s, it moves to state a + s. The reward of an agent after performing action a at state s equals −‖a‖2 ∗ 30 −min(d(s,Goal1), d(s,Goal2))2 where d(x, y) denotes the distance between x and y in R2. The initial state is chosen by sampling a Gaussian distribution with mean ( 0 0 ) and diagonal variance 0.1. In each iteration step we sample 100 trajectories. In the following pictures we plot the policies’ behavior by plotting 100 trajectories of each. The embedding Φ : Γ→ R maps trajectories τ to their mean displacement in the x−axis. We use the squared absolute value difference as the cost function.
There are two optimal policies, moving the particle to the left goal or moving it to the right goal. We now plot how the policies’ behavior and evolves throughout optimization and how the Behavioral Test Functions guide the optimization by favouring the two policies to be far apart.
Let X and Y be the domains of two measures µ, and ν. Recall that in case γ = 0, X = Y , and C(x, x) = 0 for all x ∈ X , then λ∗µ(x) = λ∗ν(x) = λ∗(x) for all x ∈ X . In the case of regularized Wasserstein distances with γ > 0, this relationship may not hold true even if the cost satisfies the same diagonal assumption. For example when the regularizing measure is the product measure, and µ, ν have disjoint supports, since the soft constraint γ exp ( λµ(x)−λν(y)−C(x,y)
γ
) is enforced in
expectation over the product measure there may exist optimal solutions λ∗µ, λ ∗ ν that do not satisfy λ∗µ = λ ∗ ν .

heading: 10 THEORETICAL RESULTS
text: We start by exploring some properties of the Wasserstein distance and its interaction with some simple classes of embeddings. The first lemma we show has the intention to show conditions under which two policies can be shown to be equal provided the Wasserstein distance between its trajectory embeddings is zero. This result implies that our framework is capable of capturing equality of policies when the embedding space equals the space of trajectories. Lemma 10.1. Let S and A be finite sets, the MDP be episodic (i.e. of finite horizon H), and Φ(τ) = ∑H t=0 est,at with es,a ∈ R|S|+|A| the indicator vector for the state action pair (s, a). Let C(v,w) = ‖v −w‖pp for p ≥ 1. If γ = 0 and WDγ(PΦπ ,PΦπ′) = 0 then π = π′.
Proof. If WDγ(PΦπ ,PΦπ′) = 0, there exists a coupling Π between PΦπ and PΦπ′ such that: Eu,v∼Π [ ‖u− v‖pp ] = 0
Consequently:
Eu,v∼Π  ∑ (s,a)∈S×A |us,a − vs,a|p  = ∑ (s,a)∈S×A Eu,v∼Π [|us,a − vs,a|p] = 0
Therefore for all (s, a) ∈ S ×A:∣∣∣Eu∼PΦπ [us,a]− Ev∼PΦπ′ [vs,a]∣∣∣p ≤ Eu,v∼Π [|us,a − vs,a|p] = 0 Where us,a and vs,a denote the (s, a) entries of u and v respectively. Notice that for all (s, a) ∈ S×A:
PΦπ (s, a) = PΦπ′(s, a) (8) Since for all s ∈ S and p ≥ 1:∣∣∣∣∣∑
a∈A us,a − vs,a
∣∣∣∣∣ p ≤ ∑ a∈A |us,a − vs,a|p
Therefore for all s ∈ S:∣∣∣∣∣Eu∼PΦπ [∑ a∈A us,a ] − Ev∼PΦ π′ [∑ a∈A vs,a ]∣∣∣∣∣ p ≤ Eu,v∼Π [∑ a∈A |us,a − vs,a|p ] = 0
Consequently PΦπ (s) = PΦπ′(s) for all s ∈ S. By Bayes rule, this plus equation 8 yields: PΦπ (a|s) = PΦπ′(a|s) And therefore: π = π′.
These results can be extended in the following ways:
1. In the case of a continuous state space, it is possible to define embeddings using Kernel density estimators. Under the appropriate smoothness conditions on the visitation frequencies, picking an adequate bandwidth and using the appropriate norm to compare different embeddings it is possible to derive similar results to those in Lemma 10.1 for continuous state spaces.
2. For embeddings such as Φ5 in Section 3.1 or Φ(τ) = ∑H t=0 est,at , when γ = 0, if
WDγ(PΦπ ,PΦπ′) ≤ then |V (π) − V (π′)| ≤ R for R = maxτ∈ΓR(τ) thus implying that a small Wasserstein distance between π and π′s PPEs implies a small difference in their value functions.

heading: 10.1 RANDOM FEATURES STOCHASTIC GRADIENTS
text: Let φκ and φ` be two feature maps over X and Y and corresponding to kernels κ and ` respectively. For this and the following sections we will make use of the following expression:
G(pµ,pν) = β ∫ X (pµ) > φκ(x)dµ(x, θ)− β ∫ Y (pν) > φ`(y)dν(y)+ (9)
γβ ∫ X×Y exp ( (pµ) > φκ(x)− (pν)> φ`(y)− C(x,y) γ ) dµ(x)dν(y)
We now show how to compute gradients with respect to the random feature maps:
Lemma 10.2. The gradient∇(pµpν)G(p µ,pν) of the objective function from Equation 9 with respect to the parameters ( pµ
pν
) satisfies:
∇(pµpν)G(p µ,pν) = βE(x,y)∼µ⊗ ν
[( 1− exp ( (pµ)>φκ(x)− (pν)>φ` − C(x,y)
γ
))( φκ(x)
−φ`(y)
)]
Proof. A simple use of the chain rule, taking the gradients inside the expectation, and the fact that pµ and pν are vectors yields the desired result.
The main consequence of this formulation is the stochastic gradients we use in Algorithm 1.

heading: 10.2 BGPG, BGES AND THEIR THEORETICAL GUARANTEES
text: Here we provide some theoretical guarantees for our algorithms. Both proposed methods BGPG and BGES follow the alternating optimization algorithmic template. We start by noting that stripped to their bare bones components our algorithms satisfy two paradigms:
1. Min-Max optimization. When β < 0 the algorithms we propose for solving the optimization objective in Equation , turn into an Alternating Min-Max optimization procedure. Henceforth whenever we refer to the problem defined by setting β < 0 we will call it the Min-Max problem.
2. Max-Max optimization. When β > 0 the algorithms we propose for solving the optimization objective in Equation turn into an alternating Max-Max procedure. Henceforth whenever we refer to the problem defined by setting β < 0 we will call it the Max-Max problem.
Our theoretical results will therefore center around providing guarantees for optimizing an objective of the form:
F (θ,pπθ ,pπ ′ ) = Eτ1,τ2∼Pπθ ⊗ Pπ′ [ R(π1) + 〈pπθ , φκ(Φ(τ1))〉 − β〈pπ ′ , φκ(Φ(τ2))〉+ (10)
γβ exp
( 〈pπθ , φκ(Φ(τ1))〉 − 〈φπ
′ , φκ(Φ(τ2))〉 − C(Φ(τ1),Φ(τ2)) γ
)]
Where πθ is a policy parametrized by θ, π′ is a target policy, and pπθ and pπ ′
are the feature vectors corresponding to the Behavioral Test functions from πθ and π′.

heading: 10.3 MAX-MAX PROBLEM: THEORETICAL ANALYSIS
text: We will analyze our AO algorithm for the Max-Max optimization problem. We show that obtained solutions converge to the local maxima of the objective function. Consider the function F (θ,pλπθ ,pλπ′ ) as in 10. We denote by (θ∗,p λπθ ∗ ,p λπ′ ∗ ) some of its local maxima. Define F̃ (θ) = F (θ,p λπθ ∗ ,p λπ′ ∗ ), i.e. F̃ is F as a function of θ for locally optimal values of pλπθ and pλπ′ .
We will assume that F̃ is locally ζ-strongly concave and δ-smooth for some fixed ζ, δ > 0 in the neighborhoodN(θ∗, r) of its optimal value θ∗. We will also assume that gradient of F̃ is L2-Lipschitz with Lipschitz coefficient φ in that neighborhood. The following convergence theorem holds:
Our goal in this section is to prove the following Theorem:
Theorem 10.3. For the entropy coefficient γ, denote: φγ = 12γ e 4 γ , and uγ = 83γ e 4 γ . Denote φ∗ = max(φ, φγ) and ξ = min( 2δζδ+ζ , uγ). Let s(t) = (θ(t),p λπθ (t),pλπ′ (t)) be the solution from iteration t and s∗ the local maximum considered above. Assume that optimization starts in θ0 ∈ N(θ∗, r). If φ∗ < 2ξ3 , then the error at iteration t + 1 of the presented AO algorithm for the Max-Max problem with decaying gradient step size αt = 3/2
[2ξ−3φ∗](t+2)+ 32φ∗ satisfies:
E[‖s(t+ 1)− s∗‖22] ≤ E[‖s(0)− s∗‖22]( 2
t+ 3 )
3 2 + σ2
9
[2ξ − 3φ∗]2(t+ 3) , (11)
where σ = √
2(1 + e 2 γ )2 + supN(θ∗),r)∇θF̃ (θ)2.
We will need several auxiliary technical results. We will use the following notation: Wγ(θ, β1, β2) = F (θ,pλπθ ,pλπ′ ), where F is the objective function from the main body of the paper parameterised by entropy coefficient γ > 0, β1 = pλπθ and β2 = pλπ′ . We will apply this notation also in the next section regarding the Min-Max Problem. For completeness, the definitions of strong concavity, smoothness and Lipschitz condition from Theorem 10.3 are given in Section 10.3.1.
We consider the dual optimization problem:
max θ max λ1∈C(X ),λ2∈C(Y) Wγ(θ, β1, β2)
= max θ max λ1∈C(X ),λ2∈C(Y) E(x,y,κx,1,...,κx,D,κy,1,...,κy,D)∼µ×ν×ω×···×ω[ L(θ) + λ1(Φ(x)) + λ2(Φ(y))− γ exp ( λ1(Φ(x)) + λ2(Φ(y))− C(Φ(x),Φ(y))
γ
)] ,
where γ is a parameter, µ = πθ, ν = π′, Φ is a fixed trajectories’ embedding and furthermore:
λ1(z) = β > 1 f(z) λ2(z) = β > 2 f(z),
such that f(·) is a random feature vector. The ith entry of the feature vector is constructed as follows: [f(z)]i = √ 2 D cos(z >wz,i + bz,i), where wz,i ∼ N ( 0, In 1ρ2 )
and bz,i ∼ Uniform(0, 2π). For the ease of notation we denote κz,i = (wz,i, bz,i) ∼ ω. We consider stochastic gradient ascent optimization strategy of the following form:
• at time t we receive a single sample (xt, yt, κxt,1t , . . . , κ xt,D t , κ yt,1 t , . . . , κ yt,D t ) ∼ µ × ν × ω × · · · × ω, then we form feature vectors [f(xt)]i = √ 2 D cos(x > t w xt,i t + b
xt,i t ) and [f(yt)]i =√
2 D cos(y > t w yt,i t + b yt,i t ), and finally update: βt+11 ← Π1 [ βt1 + αtf(xt) ( 1− exp ( (βt1) >f(xt)+(β t 2) >f(yt)−C(x,y)
γ ))] βt+12 ← Π2 [ βt2 + αtf(yt) ( 1− exp ( (βt+11 ) >f(xt)+(β t 2) >f(yt)−C(x,y)
γ
))] .
Πd (d = 1, 2) denotes the projection onto the Euclidean ball B2( rd2 , β 0 d) of some given radius rd 2 centered at the initial iterate β0d .
Let {β∗1 , β∗2} denote the global optimum of Wγ(β1, β2) computed on the entire data population, i.e. given access to an infinite number of samples (“oracle”). Let B2(r, β) denote the Euclidean ball of radius r centered at β. For the sake of the theoretical analysis we assume that (a lower-bound on) the radii of convergence r1, r2 for β1, β2, respectively, is known to the algorithm (this assumption
is potentially easy to eliminate with a more careful choice of the step size in the first iterations). To be more specific, if at any point in time parameter β1 or β2 falls outside the ball B2(r1, β∗1) or B2(r2, β ∗ 2), respectively, the projection is applied that pushes the parameter of interest to stay in the ball. Also, let ∇β1W 1γ (β1, β2) and ∇β2W 1γ (β1, β2) denote the gradients of Wγ with respect to β1 and β2, respectively, computed for a single sample. Similarly, ∇β1Wγ(β1, β2) and ∇β2Wγ(β1, β2) be the gradient of Wγ with respect to β1 and β2, respectively, computed for the entire data population, i.e. infinite number of samples.
Note that given any initial vector β0d in the ball of radius rd 2 centered at β ∗ d , we are guaranteed that all iterates remain within an rd-ball of β∗d . This is true for all d = 1, 2. The projection is necessary for theoretical analysis but in practice makes little difference. The above is a two-step alternated optimization scheme.
Let the population gradient operator, Gd(β1, β2), where d = 1, 2, be defined as Gd(β1, β2) := βd + α∇βdWγ(β1, β2).

heading: 10.3.1 ASSUMPTIONS
text: Let W ∗γ,1(β1) = Wγ(β1, β ∗ 2) and W ∗ γ,2(β2) = Wγ(β ∗ 1 , β2). Let Ω1,Ω2 denote non-empty compact convex sets such β1 ∈ Ω1 and β2 ∈ Ω2. The following assumptions are made: Assumption 10.1 (Strong concavity). For d = 1, 2, the function W ∗γ,d(βd) is ζd-strongly concave near β∗d , i.e. for all pairs (ad, bd) in the neighborhood of β ∗ d the following holds
W ∗γ,d(ad)−W ∗γ,d(bd)− 〈 ∇βdW ∗γ,d(bd), ad − bd 〉 ≤ −ζd
2 ‖ad − bd‖22,
where ζd > 0 is the strong concavity modulus. Assumption 10.2 (Smoothness). For d = 1, 2, the function W ∗γ,d(βd) is δd-smooth, i.e. for all pairs (ad, bd) the following holds
W ∗γ,d(ad)−W ∗γ,d(bd)− 〈 ∇βdW ∗γ,d(bd), ad − bd 〉 ≥ −δd
2 ‖ad − bd‖22,
where δd > 0 is the smoothness constant. Assumption 10.3 (Gradient stability (GS) / Lipschitz condition). We assume Wγ(β1, β2) satisfy GS (φd) condition, for all d = 1, 2, over Euclidean balls β1 ∈ B2(r1, β∗1), β2 ∈ B2(r2, β∗2) given as follows ‖∇βdW ∗γ,d(βd)−∇βdWγ(β1, β2)‖2 ≤ φd‖βd̄ − β∗d̄‖2, where φd > 0 and d̄ = (d mod 2) + 1.
Finally, define the bound σ that considers the expected value of the norm of gradients of our objective function as follows: σ = √ σ21 + σ 2 2 , where σ 2 d = sup{E[‖∇βdW 1γ (β1, β2)‖22] : β1 ∈ B2(r1, β ∗ 1), β2 ∈ B2(r2, β∗2)} for d = 1, 2.

heading: 10.3.2 MAIN THEOREMS
text: Theorem 10.4. Given the stochastic gradient iterates of Max-Max method with decaying step size {αt}∞t=0 and with φ < 2ξ 3 the error at iteration t+ 1 satisfies the recursion
E [ ‖βt+11 − β∗1‖22 + ‖β t+1 2 − β∗2‖22 ] ≤ (1− qt)E [ ‖βt1 − β∗1‖22 + ‖βt2 − β∗2‖22 ] + (αt)2
1− αtφ σ2,
where φ = maxd=1,2(φd), qt = 1− 1−2α tξ+2αtφ 1−αtφ , and ξ = mind=1,2 ( 2δdζd δd+ζd ) .
The recursion in Theorem 10.4 is expanded yielding the convergence theorem: Theorem 10.5. Given the stochastic gradient iterates of Max-Max method with decaying step size αt = 3/2
[2ξ−3φ](t+2)+ 32φ and assuming that φ < 2ξ3 , the error at iteration t+ 1 satisfies
E [ ‖βt+11 −β∗1‖22 + ‖β t+1 2 −β∗2‖22 ] ≤ E [ ‖β01−β∗1‖22 + ‖β02−β∗2‖22 ]( 2 t+ 3 ) 3 2 +σ2 9 [2ξ − 3φ]2(t+ 3) ,
where φ = maxd=1,2(φd) and ξ = mind=1,2 (
2δdζd δd+ζd
) .

heading: 10.3.3 ANALYSIS
text: The theoretical analysis we provide below is an extension of the analysis in Balakrishnan et al. (2017) to the two-step alternated optimization scheme.
Proof of Theorem 10.5 relies on Theorem 10.4, which in turn relies on Theorem 10.7 and Lemma 10.6, both of which are stated below. Proofs of the lemma and theorems follow in the subsequent subsections.
The next result is a standard result from convex optimization (Theorem 2.1.14 in Nesterov (2014)) and is used in the proof of Theorem 10.7 below. Lemma 10.6. The gradient operator G1(β1, β∗2) under Assumption 10.1 (strong concavity) and Assumption 10.2 (smoothness) with constant step size choice 0 < α ≤ 2δ1+ζ1 is contractive, i.e.
‖G1(β1, β∗2)− β∗1‖2 ≤ (
1− 2αδ1ζ1 δ1 + ζ1
) ‖β1 − β∗1‖2 (12)
for all β1 ∈ B2(r1, β∗1). Similarly, the gradient operator G2(β∗1 , β2) under Assumption 10.1 (strong concavity) and Assumption 10.2 (smoothness) with constant step size choice 0 < α ≤ 2δ2+ζ2 is contractive, i.e.
‖G2(β∗1 , β2)− β∗2‖2 ≤ (
1− 2αδ2ζ2 δ2 + ζ2
) ‖β2 − β∗2‖2 (13)
for all β2 ∈ B2(r2, β∗2).
The next theorem also holds for d = 1, 2. Let r1, r2 > 0 and β1 ∈ B2(r1, β∗1), β2 ∈ B2(r2, β∗2). Theorem 10.7. For some radius rd > 0 (d = 1, 2) and a triplet (φd, ζd, δd) such that 0 ≤ φd < ζd ≤ δd, suppose that the function W ∗γ,d(βd) is ζd-strongly concave and δd-smooth, and that the GS (φd) condition holds. Then the population gradient operator Gd(β1, β2) with step α such that 0 < α ≤ mind=1,2 2δd+ζd is contractive over a ball B2(rd, β ∗ d), i.e.
‖Gd(β1, β2)− β∗d‖2 ≤ (1− ξα)‖βd − β∗d‖2 + αφ‖βd̄ − β∗d̄‖2 (14)
where d̄ = (d mod 2) + 1, φ := maxd=1,2 φd, and ξ := mind=1,2 2δdζdδd+ζd .
Proof. ‖Gd(β1, β2)− β∗d‖2 = ‖βd + α∇βdWγ(β1, β2)− β∗d‖2 by the triangle inequality (and with d 6= d̄) we further get ≤ ‖βd + α∇βdW ∗γ,d(βd)− β∗d‖2 + α‖∇βdWγ(β1, β2)−∇βdW ∗γ,d(βd)‖2 by the contractivity from Lemma 10.6 and GS condition
≤ (
1− 2αδdζd δd + ζd
) ‖βd − β∗d‖2 + αφd‖βd̄ − β∗d̄‖2.
Proof of Theorem 10.4
Let βt+1d = Πd(β̃ t+1 d ), where β̃ t+1 1 := β t 1 + α t∇β1W 1γ (βt1, βt2) and β̃t+12 := βt2 + αt∇β2W 1γ (βt+11 , βt2), where ∇βdW 1γ is the gradient computed with respect to a single sample, β̃1 and β̃2 are the updates prior to the projection onto a ball B2( rd2 , β 0 d). Let ∆ t+1 d := β t+1 d − β∗d and ∆̃t+1d := β̃ t+1 d − β∗d . Thus
‖∆t+1d ‖ 2 2 − ‖∆td‖22 ≤ ‖∆̃t+1d ‖ 2 2 − ‖∆td‖22
= ‖β̃t+1d − β ∗ d‖ − ‖βtd − β∗d‖ = 〈 β̃t+1d − β t d, β̃ t+1 d + β t d − 2β∗d 〉 .
Let Q̂t1 := ∇β1W 1γ (βt1, βt2) and Q̂t2 := ∇β2W 1γ (βt+11 , βt2). Then we have that β̃ t+1 d − βtd = αtQ̂td. We combine it with Equation 17 and obtain:
‖∆t+1d ‖ 2 2 − ‖∆td‖22 ≤ 〈 αtQ̂td, α tQ̂td + 2(β t d − β∗d) 〉 = (αt)2(Q̂td) >Q̂td + 2α t(Q̂td) >(βtd − β∗d)
= (αt)2‖Q̂td‖22 + 2αt 〈 Q̂td,∆ t d 〉 .
Let Qt1 := ∇β1Wγ(βt1, βt2) and Qt2 := ∇β2Wγ(βt+11 , βt2). By the properties of martingales, i.e. iterated expectations and tower property:
E[‖∆t+1d ‖ 2 2] ≤ E[‖∆td‖22] + (αt)2E[‖Q̂td‖22] + 2αtE[ 〈 Qtd,∆ t d 〉 ] (15)
Let Q∗d := ∇βdWγ(β∗1 , β∗2). By self-consistency, i.e. β∗d = arg maxβd∈ΩdW ∗γ,d(βd), and convexity of Ωd we have that 〈
Q∗d,∆ t d 〉 = 〈 ∇βdWγ(β∗1 , β∗2),∆td 〉 = 0.
Combining this with Equation 15 we have
E[‖∆t+1d ‖ 2 2] ≤ E[‖∆td‖22] + (αt)2E[‖Q̂td‖22] + 2αtE[ 〈 Qtd −Q∗d,∆td 〉 ].
Define Gtd := βtd + αtQtd and Gt∗d := β∗d + αtQ∗d. Thus
αt 〈 Qtd −Q∗d,∆td 〉 = 〈 Gtd − Gt∗d − (βtd − β∗d), βtd − β∗d
〉 = 〈 Gtd − Gt∗d , βtd − β∗d 〉 − ‖βtd − β∗d‖22
by the fact that Gt∗d = β∗d + αtQ∗d = β∗d (since Q∗d = 0): = 〈 Gtd − β∗d , βtd − β∗d 〉 − ‖βtd − β∗d‖22
by the contractivity of Gt from Theorem 10.7:
≤ { (1− αtξ)‖βtd − β∗d‖+ αtφ ( d−1∑ i=1 ‖βt+1i − β ∗ i ‖2 + 2∑ i=d+1 ‖βti − β∗i ‖2 )} ‖βtd − β∗d‖2 − ‖βtd − β∗d‖22
≤ { (1− αtξ)‖∆td‖2 + αtφ ( d−1∑ i=1 ‖∆t+1i ‖2 + 2∑ i=d+1 ‖∆ti‖2 )} · ‖∆td‖2 − ‖∆td‖22
Combining this result with Equation 16 gives
E[‖∆t+1d ‖ 2 2] ≤ E[‖∆td‖22] + (αt)2E[‖Q̂td‖22] + 2E [{ (1− αtξ)‖∆td‖2 + αtφ ( d−1∑ i=1 ‖∆t+1i ‖2 + 2∑ i=d+1 ‖∆ti‖2 )} ·‖∆td‖2 − ‖∆td‖22
] ≤ E[‖∆td‖22] + (αt)2σ2d + 2E [{ (1− αtξ)‖∆td‖2 + αtφ ( d−1∑ i=1 ‖∆t+1i ‖2 + 2∑ i=d+1 ‖∆ti‖2
)} ·‖∆td‖2 − ‖∆td‖22 ] .
After re-arranging the terms we obtain
E[‖∆t+1d ‖ 2 2] ≤ (αt)2σ2d + (1− 2αtξ)E[‖∆td‖22] + 2αtφE [( d−1∑ i=1 ‖∆t+1i ‖2 + 2∑ i=d+1 ‖∆ti‖2 ) ‖∆td‖2 ] apply 2ab ≤ a2 + b2 :
≤ (αt)2σ2d + (1− 2αtξ)E[‖∆td‖22] + αtφE [ d−1∑ i=1 ( ‖∆t+1i ‖ 2 2 + ‖∆td‖22 )] + αtφE [ 2∑ i=d+1 ( ‖∆ti‖22 + ‖∆td‖22 )] = (αt)2σ2d + E[‖∆td‖22] · [ 1− 2αtξ + αtφ ] + αtφE [ d−1∑ i=1 ‖∆t+1i ‖ 2 2 ] + αtφE [ 2∑ i=d+1 ‖∆ti‖22 ]
We obtained
E[‖∆t+1d ‖ 2 2] ≤ (αt)2σ2d + [1− 2αtξ + αtφ]E[‖∆td‖22] + αtφE [ d−1∑ i=1 ‖∆t+1i ‖ 2 2 ] + αtφE [ 2∑ i=d+1 ‖∆ti‖22 ] we next re-group the terms as follows
E[‖∆t+1d ‖ 2 2]− αtφE [ d−1∑ i=1 ‖∆t+1i ‖ 2 2 ] ≤ [1− 2αtξ + αtφ]E[‖∆td‖22] + αtφE [ 2∑ i=d+1 ‖∆ti‖22 ] + (αt)2σ2d
and then sum over d from 1 to 2
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ] − αtφE [ 2∑ d=1 d−1∑ i=1 ‖∆t+1i ‖ 2 2 ]
≤ [1− 2αtξ + αtφ]E [ 2∑ d=1 ‖∆td‖22 ] + αtφE [ 2∑ d=1 2∑ i=d+1 ‖∆ti‖22 ] + (αt)2 2∑ d=1 σ2d
Let σ = √ σ21 + σ 2 2 . Also, note that
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ] − αtφE [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ] ≤ E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ] − αtφE [ 2∑ d=1 d−1∑ i=1 ‖∆t+1i ‖ 2 2 ] and
[1− 2αtξ + αtφ]E [ 2∑ d=1 ‖∆td‖22 ] + αtφE [ 2∑ d=1 2∑ i=d+1 ‖∆ti‖22 ] + (αt)2σ2
≤ [1− 2αtξ + αtφ]E [ 2∑ d=1 ‖∆td‖22 ] + αtφE [ 2∑ d=1 ‖∆td‖22 ] + (αt)2σ2
Combining these two facts with our previous results yields:
[1− αtφ]E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ [1− 2αtξ + αtφ]E [ 2∑ d=1 ‖∆td‖22 ] + αtφE [ 2∑ d=1 ‖∆td‖22 ] + (αt)2σ2
= [1− 2αtξ + 2αtφ]E [ 2∑ d=1 ‖∆td‖22 ] + (αt)2σ2
Thus:
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ] ≤ 1− 2α tξ + 2αtφ 1− αtφ E [ 2∑ d=1 ‖∆td‖22 ]
+ (αt)2
1− αtφ σ2.
Since φ < 2ξ3 , 1−2αtξ+2αtφ 1−αtφ < 1.
Proof of Theorem 10.5
To obtain the final theorem we need to expand the recursion from Theorem 10.4. We obtained
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ 1− 2α t[ξ − φ]
1− αtφ E [ 2∑ d=1 ‖∆td‖22 ] + (αt)2 1− αtφ σ2
= ( 1− α
t[2ξ − 3φ] 1− αtφ
) E [ 2∑ d=1 ‖∆td‖22 ] + (αt)2 1− αtφ σ2
Recall that we defined qt in Theorem 10.4 as
qt = 1− 1− 2α tξ + 2αtφ 1− αtφ = αt[2ξ − 3φ] 1− αtφ
and denote
f t = (αt)2
1− αtφ .
Thus we have
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ] ≤ (1− qt)E [ 2∑ d=1 ‖∆td‖22 ] + f tσ2
≤ (1− qt) { (1− qt−1)E [ 2∑ d=1 ‖∆t−1d ‖ 2 2 ] + f t−1σ2 } + f tσ2
= (1− qt)(1− qt−1)E [ 2∑ d=1 ‖∆t−1d ‖ 2 2 ] + (1− qt)f t−1σ2 + f tσ2
≤ (1− qt)(1− qt−1) { (1− qt−2)E [ 2∑ d=1 ‖∆t−2d ‖ 2 2 ] + f t−2σ2 } + (1− qt)f t−1σ2 + f tσ2
= (1− qt)(1− qt−1)(1− qt−2)E [ 2∑ d=1 ‖∆t−2d ‖ 2 2 ] +(1− qt)(1− qt−1)f t−2σ2 + (1− qt)f t−1σ2 + f tσ2
We end-up with the following
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ] ≤ E [ 2∑ d=1 ‖∆0d‖22 ] t∏ i=0 (1− qi) + σ2 t−1∑ i=0 f i t∏ j=i+1 (1− qj) + f tσ2.
Set qt = 3 2
t+2 and
αt = qt
2ξ − 3φ+ qtφ
= 3 2
[2ξ − 3φ](t+ 2) + 32φ .
Denote A = 2ξ − 3φ and B = 32φ. Thus
αt = 3 2
A(t+ 2) +B
and
f t = (αt)2
1− 23Bαt =
9 4
A(t+ 2)[A(t+ 2) +B] .
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ E [ 2∑ d=1 ‖∆0d‖22 ] t∏ i=0 ( 1− 3 2 i+ 2 ) + σ2 t−1∑ i=0
9 4
A(i+ 2)[A(i+ 2) +B] t∏ j=i+1 ( 1− 3 2 j + 2 )
+σ2 9 4
A(t+ 2)[A(t+ 2) +B]
= E [ 2∑ d=1 ‖∆0d‖22 ] t+2∏ i=2 ( 1− 3 2 i ) + σ2 t+1∑ i=2
9 4
Ai[Ai+B] t+2∏ j=i+1 ( 1− 3 2 j ) + σ2
9 4
A(t+ 2)[A(t+ 2) +B]
Since A > 0 and B > 0 thus
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ E [ 2∑ d=1 ‖∆0d‖22 ] t+2∏ i=2 ( 1− 3 2 i ) + σ2 t+1∑ i=2
9 4
Ai[Ai+B] t+2∏ j=i+1 ( 1− 3 2 j ) + σ2
9 4
A(t+ 2)[A(t+ 2) +B]
≤ E [ 2∑ d=1 ‖∆0d‖22 ] t+2∏ i=2 ( 1− 3 2 i ) + σ2 t+1∑ i=2 9 4 (Ai)2 t+2∏ j=i+1 ( 1− 3 2 j ) + σ2
9 4
[A(t+ 2)]2
We can next use the fact that for any a ∈ (1, 2): t+2∏ i=τ+1 ( 1− a i ) ≤ ( τ + 1 t+ 3 )a .
The bound then becomes
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ E [ 2∑ d=1 ‖∆0d‖22 ] t+2∏ i=2 ( 1− 3 2 i ) + σ2 t+1∑ i=2 9 4 (Ai)2 t+2∏ j=i+1 ( 1− 3 2 j ) + σ2
9 4
[A(t+ 2)]2
≤ E [ 2∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 t+1∑ i=2 9 4 (Ai)2 ( i+ 1 t+ 3 ) 3 2 + σ2 9 4 [A(t+ 2)]2
= E [ 2∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 t+2∑ i=2 9 4 (Ai)2 ( i+ 1 t+ 3 ) 3 2
Note that (i+ 1) 3 2 ≤ 2i for i = 2, 3, . . . , thus
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ E [ 2∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 9 4 A2(t+ 3) 3 2 t+2∑ i=2 (i+ 1) 3 2 i2
≤ E [ 2∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 9 2 A2(t+ 3) 3 2 t+2∑ i=2 1 i 1 2
finally note that t+2∑ i=2 1 i 1 2 ≤ ∫ t+2 1 1 x 1 2 dx ≤ 2(t+ 3) 12 . Thus
≤ E [ 2∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 9 A2(t+ 3)
substituting A = 2ξ − 3φ gives
= E [ 2∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 9 [2ξ − 3φ]2(t+ 3)
This leads us to the final theorem.
Proof of Theorem 10.3
In order to prove Theorem 10.3 it suffices to apply Theorem 10.5 and notice that:
• function hv,C : Rd → R defined as follows: hv,C(w) = w>v − Ae w>v λ for A = γe− C γ is
2 γ e 4 γ -smooth, 2γ -strongly concave and its gradient is Lipschitz with Lipschitz coefficient 1 2γ e 4 γ
(with respect to L2-norm) for C > 0 and v satisfying: |w>v| ≤ 1, • ‖∇hv,C(v)‖22 ≤ 2(1 + e 2 γ )2 under above conditions.

heading: 10.4 MINMAX PROBLEM: THEORETICAL ANALYSIS
text: In this section we aim to obtain similar results for Min-Max Problem as for Max-Max problem. We will use the same notation as in the main body of the paper. We prove the following results: Theorem 10.8. Denote φ∗ = max(φ, φγ) and ξ = min( 2δζδ+ζ , uγ), where φγ and uγ are as in Theorem 10.3. Let s(t) = (θ(t),pλπθ (t),pλπ′ (t)) be the solution obtained in iteration t and s∗ the local optimum. Assume that optimization starts in θ0 ∈ N(θ∗, r). If φ∗ < ξ3 , then the error at iteration t + 1 of the alternating optimization algorithm for the Min-Max problem with decaying gradient step size αt = 3/2[2ξ−6φ∗](t+2)+3φ∗ satisfies:
E[‖s(t+ 1)− s∗‖22] ≤ E[‖s(0)− s∗‖22]( 2
t+ 3 )
3 2 + σ2
9
[2ξ − 6φ∗]2(t+ 3) , (16) where σ = √
2(1 + e 2 γ )2 + supN(θ∗),r)∇θF̃ (θ)2.
We use the same technical notation as in the previous section, with the only exception that this time we denote: Wγ(θ, β1, β2) = −F (θ,pλπθ ,pλπ′ ). We consider the MinMax problem of the following form
min θ max λ1∈C(X ),λ2∈C(Y) Wγ(θ, β1, β2)
= min θ max λ1∈C(X ),λ2∈C(Y) E(x,y,κx,1,...,κx,D,κy,1,...,κy,D)∼µ×ν×ω×···×ω[ −L(θ) + λ1(Φ(x)) + λ2(Φ(y))− γ exp ( λ1(Φ(x)) + λ2(Φ(y))− C(Φ(x),Φ(y))
γ
)] ,
where γ is a parameter and the remaining notation is analogous to the Max-Max case.
We consider mixed stochastic gradient descent/ascend optimization strategy of the following form:
• at time t we receive a single sample (xt, yt, κxt,1t , . . . , κ xt,D t , κ yt,1 t , . . . , κ yt,D t ) ∼ πθt × ν × ω × · · · × ω, then we form feature vectors [f(xt)]i = √ 2 D cos(x > t w xt,i t + b
xt,i t ) and [f(yt)]i =√
2 D cos(y > t w yt,i t + b yt,i t ), and finally update:
θt+1 ← Π1 [θt − αt∇θ=θtWγ(θ, βt1, βt2)]4
βt+11 ← Π2 [ βt1 + αtf(xt) ( 1− exp ( (βt1) >f(xt)+(β t 2) >f(yt)−C(x,y)
γ ))] βt+12 ← Π3 [ βt2 + αtf(yt) ( 1− exp ( (βt+11 ) >f(xt)+(β t 2) >f(yt)−C(x,y)
γ
))] .
Π1 denotes the projection onto the Euclidean ballB2( r12 , θ 0) and Πd (d = 1, 2) denotes the projection onto the Euclidean ball B2( r12 , β 0 d). Let {θ∗, β∗1 , β∗2} denote the the global optimal solution of the saddle point problem minθ maxβ1,β2 Wγ(θ, β1, β2) computed on the entire data population, i.e. given access to an infinite number of samples (“oracle”). As before, we assume that (a lower-bound on) the radii of convergence r1, r2, r3 for θ, β1, β2, respectively, is known to the algorithm and thus the projection is applied to control θ, β1, β2 to stay in their respective balls. Also, let∇θW 1γ (θ, β1, β2),∇β1W 1γ (θ, β1, β2) and ∇β2W 1γ (θ, β1, β2) denote the gradients of Wγ with respect to θ,β1 and β2, respectively, computed for a single sample. Similarly, ∇θWγ(θ, β1, β2), ∇β1Wγ(θ, β1, β2) and ∇β2Wγ(θ, β1, β2) be the gradient of Wγ with respect to θ, β1 and β2, respectively, computed for the entire data population, i.e. infinite number of samples.
Note that given any initial vector θ0 in the ball of radius r12 centered at θ ∗, we are guaranteed that all iterates remain within an r1-ball of θ∗ and given any initial vector β0d (d = 1, 2) in the ball of radius rd 2 centered at β ∗ d , we are guaranteed that all iterates remain within an rd-ball of β ∗ d . The projection is necessary for theoretical analysis but in practice makes little difference. The above is a three-step alternated optimization scheme.
Let the population gradient operator, Gd(θ, β1, β2), where d = 1, 2, 3, be defined as
G1(θ, β1, β2) := θ − α∇θWγ(θ, β1, β2)
and Gd(θ, β1, β2) := βd + α∇βiWγ(θ, β1, β2) for d = 2, 3.

heading: 10.4.1 ASSUMPTIONS
text: Let W ∗γ,1(θ) = Wγ(θ, β ∗ 1 , β ∗ 2), W ∗ γ,2(β1) = Wγ(θ ∗, β1, β ∗ 2) and W ∗ γ,3(β2) = Wγ(θ ∗, β∗1 , β2). Let Ω1,Ω2,Ω3 denote non-empty compact convex sets such θ ∈ Ω1, β1 ∈ Ω2 and β2 ∈ Ω3. The following assumptions are made:
Assumption 10.4 (Strong convexity/concavity). The function W ∗γ,1(θ) is ζ1-strongly convex near θ∗ and the functions W ∗γ,2(β1) and W ∗ γ,3(β2) are ζ2- and ζ3-strongly concave, respectively, near β ∗ 1 and β∗2 , respectively, where ζ1, ζ2, ζ3 > 0. Assumption 10.5 (Smoothness). The functions W ∗γ,1(θ),W ∗γ,2(β1), and W ∗γ,3(β2) are δ1-,δ2-, and δ3-smooth, respectively, where δ1, δ2, δ3 > 0 are the smoothness constants.
Assumption 10.6 (Gradient stability (GS) / Lipschitz condition). We assume Wγ(θ, β1, β2) satisfy GS (φd) condition, for all d = 1, 2, 3, over Euclidean balls θ ∈ B2(r1, θ∗), β1 ∈ B2(r2, β∗1), β2 ∈ B2(r3, β ∗ 2) given as follows
‖∇θW ∗γ,1(θ)−∇θWγ(θ, β1, β2)‖2 ≤ φ1 2∑ d=1 ‖βd − β∗d‖2,
4later we also use alternative notation for gradient∇θ=θtWγ(θ, βt1, βt2) as∇θWγ(θt, βt1, βt2)
and for d = 1, 2
‖∇βdW ∗γ,d+1(βd)−∇βdWγ(θ, β1, β2)‖2 ≤ φd(‖θ − θ∗‖2 + ‖βd̄ − β∗d̄‖2),
where φd > 0 and d̄ = (d mod 2) + 1.
Finally, as before, define the bound σ that considers the expected value of the norm of gradients of our objective function as follows: σ = √ σ21 + σ 2 2 + σ 2 3 , where σ 2 1 = sup{E[‖∇θW 1γ (θ, β1, β2)‖22] : θ ∈ B2(r1, θ∗), β1 ∈ B2(r2, β∗1), β2 ∈ B2(r3, β∗2)} and for d = 1, 2 σ2d+1 = sup{E[‖∇βdW 1γ (θ, β1, β2)‖22] : θ ∈ B2(r1, θ∗), β1 ∈ B2(r2, β∗1), β2 ∈ B2(r3, β∗2)}.

heading: 10.4.2 MAIN THEOREMS
text: Theorem 10.9. Given the stochastic gradient iterates of MinMax method with decaying step size {αt}∞t=0 and with φ < ξ 3 the error at iteration t+ 1 satisfies the recursion
E [ ‖θt+1 − θ∗‖22 + ‖βt+11 − β∗1‖22 + ‖β t+1 2 − β∗2‖22 ] ≤ (1− qt)E [ ‖θt − θ∗‖22 + ‖βt1 − β∗1‖22 + ‖βt2 − β∗2‖22 ] + (αt)2
1− 2αtφ σ2,
where φ = maxd=1,2,3(φd), qt = 1− 1−2α tξ+4αtφ 1−2αtφ , and ξ = mind=1,2,3 ( 2δdζd δd+ζd ) .
The recursion in Theorem 10.4 is expanded yielding the convergence theorem: Theorem 10.10. Given the stochastic gradient iterates of MinMax method with decaying step size αt = 3/2[2ξ−6φ](t+2)+3φ and assuming that φ < ξ 3 , the error at iteration t+ 1 satisfies
E [ ‖θt+1 − θ∗‖22 + ‖βt+11 − β∗1‖22 + ‖β t+1 2 − β∗2‖22 ] ≤ E [ ‖θ0 − θ∗‖22 + ‖β01 − β∗1‖22 + ‖β02 − β∗2‖22 ]( 2 t+ 3 ) 3 2 + σ2 9 [2ξ − 6φ]2(t+ 3) ,
where φ = maxd=1,2,3(φd) and ξ = mind=1,2,3 (
2δdζd δd+ζd
) .
Proof of Theorem 10.10 relies on Theorem 10.9, which in turn relies on Theorem 10.12 and Lemma 10.11, both of which are stated below. Proofs of the lemma and theorems follow in the subsequent subsections.

heading: 10.4.3 ANALYSIS
text: The next result is a standard result from convex optimization (Theorem 2.1.14 in Nesterov (2014)) and is used in the proof of Theorem 10.12 below. Lemma 10.11. The gradient operator G1(θ, β∗1 , β∗2) under strong convexity and smoothness assumptions with constant step size choice 0 < α ≤ 2δ1+ζ1 is contractive, i.e.
‖G1(θ, β∗1 , β∗2)− θ∗‖2 ≤ (
1− 2αδ1ζ1 δ1 + ζ1
) ‖θ − θ∗‖2
for all θ ∈ B2(r1, θ∗). Similarly, the gradient operator G1(θ∗, β1, β∗2) under strong concavity and smoothness assumptions with constant step size choice 0 < α ≤ 2δ2+ζ2 is contractive, i.e.
‖G1(θ∗, β1, β∗2)− β∗1‖2 ≤ (
1− 2αδ2ζ2 δ2 + ζ2
) ‖β1 − β∗1‖2
for all β1 ∈ B2(r2, β∗1). And similarly, the gradient operator G2(θ∗, β∗1 , β2) under strong concavity and smoothness assumptions with constant step size choice 0 < α ≤ 2δ3+ζ3 is contractive, i.e.
‖G2(θ∗, β∗1 , β2)− β∗2‖2 ≤ (
1− 2αδ3ζ3 δ3 + ζ3
) ‖β2 − β∗2‖2
for all β2 ∈ B2(r3, β∗2).
The next theorem holds for d = 1, 2, 3. Let r1, r2, r3 > 0 and θ ∈ B2(r1, θ∗), β1 ∈ B2(r2, β∗1), β2 ∈ B2(r3, β ∗ 2). Theorem 10.12. For some radius r1 and a triplet (φ1, ζ1, δ1) such that 0 ≤ φ1 < ζ1 ≤ δ1, suppose that the function W ∗γ,1(θ) is ζ1-strongly convex and δ1-smooth and that the GS (φ1) condition holds. Then the population gradient operator G1(θ, β1, β2) with step α such that 0 < α ≤ mind=1,2,3 2δd+ζd is contractive over a ball B2(r1, θ∗), i.e.
‖G1(θ, β1, β2)− θ∗‖2 ≤ (1− ξα)‖θ − θ∗‖2 + αφ 2∑ d=1 ‖βd − β∗d‖2
where φ := maxd=1,2,3 φd and ξ := mind=1,2,3 2δdζdδd+ζd .
For some radius rd (d = 2, 3) and a triplet (φd, ζd, δd) such that 0 ≤ φd < ζd ≤ δd, suppose that the function W ∗γ,d(βd−1) is ζd-strongly concave and δd-smooth and that the GS (φd) condition holds. Then the population gradient operator Gd(θ, β1, β2) with step α such that 0 < α ≤ mind=1,2,3 2δd+ζd is contractive over a ball B2(rd, β∗d), i.e.
‖Gd(θ, β1, β2)− β∗d‖2 ≤ (1− ξα)‖βd − β∗d‖2 + αφ(‖βd̄ − β∗d̄‖2 + ‖θ − θ ∗‖2).
where d̄ = ((d− 1) mod 2) + 1, φ := maxd=1,2,3 φd, and ξ := mind=1,2,3 2δdζdδd+ζd .
Proof. ‖G1(θ, β1, β2)− θ∗‖2 = ‖θ − α∇θWγ(θ, β1, β2)− θ∗‖2 by the triangle inequality we further get ≤ ‖θ − α∇θW ∗γ,1 − θ∗‖2 + α‖∇θWγ(θ, β1, β2)−∇θW ∗γ,1‖2 by the contractivity from Lemma 10.11 and GS condition
≤ (
1− 2αδ1ζ1 δ1 + ζ1
) ‖θ − θ∗‖2 + αφ1 2∑ d=1 ‖βd − β∗d‖2.
The proof of the rest of the theorem is analogous to the proof of Theorem 10.7.
Proof of Theorem 10.9
Let θ1 = θ, θ2 = β1, and θ3 = β2.
Let θt+1d = Πd(θ̃ t+1 d ), where θ̃ t+1 1 := θ t 1 − αt∇θ1W 1γ (θt1, θt2, θt3), θ̃t+12 := θt2 + αt∇θ2W 1γ (θt+11 , θt2, θt3), and θ̃ t+1 3 := θ t 3 + α t∇θ3W 1γ (θt+11 , θ t+1 2 , θ t 3), where ∇θdW 1γ is the gradient computed with respect to a single sample, θ̃1, θ̃2, and θ̃3 are the updates prior to the projection. Let ∆t+11 := −θ t+1 1 + θ ∗ 1 and for d = 2, 3, ∆̃ t+1 d := θ̃ t+1 d − θ∗d. Thus
‖∆t+1d ‖ 2 2 − ‖∆td‖22 ≤ ‖∆̃t+1d ‖ 2 2 − ‖∆td‖22
= ‖θ̃t+1d − θ ∗ d‖ − ‖θtd − θ∗d‖ = 〈 θ̃t+1d − θ t d, θ̃ t+1 d + θ t d − 2θ∗d 〉 .
Let Q̂t1 := ∇θ1W 1γ (θt1, θt2, θt3), Q̂t2 := ∇θ2W 1γ (θt+11 , θt2, θt3), and Q̂t3 := ∇θ3W 1γ (θ t+1 1 , θ t+1 2 , θ t 3). Thus: ‖∆t+1d ‖ 2 2 − ‖∆td‖22
≤ 〈 αtQ̂td, α tQ̂td + 2(θ t d − θ∗d) 〉 = (αt)2(Q̂td) >Q̂td + 2α t(Q̂td) >(θtd − θ∗d)
= (αt)2‖Q̂td‖22 + 2αt 〈 Q̂td,∆ t d 〉 .
Let Qt1 := ∇θ1Wγ(θt1, θt2, θt3), Qt2 := ∇θ2Wγ(θt+11 , θt2, θt3), and Qt3 := ∇θ3Wγ(θ t+1 1 , θ t+1 2 , θ t 3). By the properties of martingales, i.e. iterated expectations and tower property:
E[‖∆t+1d ‖ 2 2] ≤ E[‖∆td‖22] + (αt)2E[‖Q̂td‖22] + 2αtE[ 〈 Qtd,∆ t d 〉 ]
Let Q∗d := ∇θdWγ(θ∗1 , θ∗2 , θ∗3). By self-consistency, i.e. θ∗d = arg maxθd∈ΩdW ∗γ,d(θd), and convexity of Ωd we have that 〈
Q∗d,∆ t d 〉 = 〈 ∇θdWγ(θ∗1 , θ∗2 , θ∗3),∆td 〉 = 0.
Combining this with the above inequality yields
E[‖∆t+1d ‖ 2 2] ≤ E[‖∆td‖22] + (αt)2E[‖Q̂td‖22] + 2αtE[ 〈 Qtd −Q∗d,∆td 〉 ].
Define Gt1 := θt1 − αtQt1 and Gt∗1 := θ∗1 − αtQ∗1. Also, for d = 2, 3 define Gtd := θtd + αtQtd and Gt∗d := θ∗d + αtQ∗d. Thus
αt 〈 Qtd −Q∗d,∆td 〉 = 〈 Gtd − Gt∗d − (θtd − θ∗d), θtd − θ∗d
〉 = 〈 Gtd − Gt∗d , θtd − θ∗d 〉 − ‖θtd − θ∗d‖22
by the fact that Gt∗d = θ∗d + αtQ∗d = θ∗d (since Q∗d = 0): = 〈 Gtd − θ∗d, θtd − θ∗d 〉 − ‖θtd − θ∗d‖22
by the contractivity of Gt from Theorem 10.7:
≤ { (1− αtξ)‖θtd − θ∗d‖+ αtφ ( d−1∑ i=1 ‖θt+1i − θ ∗ i ‖2 + 3∑ i=d+1 ‖θti − θ∗i ‖2 )} ‖θtd − θ∗d‖2 − ‖θtd − θ∗d‖22
≤ { (1− αtξ)‖∆td‖2 + αtφ ( d−1∑ i=1 ‖∆t+1i ‖2 + 3∑ i=d+1 ‖∆ti‖2 )} · ‖∆td‖2 − ‖∆td‖22
Thus
E[‖∆t+1d ‖ 2 2] ≤ E[‖∆td‖22] + (αt)2E[‖Q̂td‖22] + 2αtE[ 〈 Qtd −Q∗d,∆td 〉 ]
= E[‖∆td‖22] + (αt)2E[‖Q̂td‖22] + 2E [{ (1− αtξ)‖∆td‖2 + αtφ ( d−1∑ i=1 ‖∆t+1i ‖2 + 3∑ i=d+1 ‖∆ti‖2 )} ·‖∆td‖2 − ‖∆td‖22
] ≤ E[‖∆td‖22] + (αt)2σ2d + 2E [{ (1− αtξ)‖∆td‖2 + αtφ ( d−1∑ i=1 ‖∆t+1i ‖2 + 3∑ i=d+1 ‖∆ti‖2
)} ·‖∆td‖2 − ‖∆td‖22 ] .
After re-arranging the terms we obtain
E[‖∆t+1d ‖ 2 2] ≤ (αt)2σ2d + (1− 2αtξ)E[‖∆td‖22] + 2αtφE [( d−1∑ i=1 ‖∆t+1i ‖2 + 3∑ i=d+1 ‖∆ti‖2 ) ‖∆td‖2 ] apply 2ab ≤ a2 + b2 :
≤ (αt)2σ2d + (1− 2αtξ)E[‖∆td‖22] + αtφE [ d−1∑ i=1 ( ‖∆t+1i ‖ 2 2 + ‖∆td‖22 )] + αtφE [ 3∑ i=d+1 ( ‖∆ti‖22 + ‖∆td‖22 )] = (αt)2σ2d + E[‖∆td‖22] · [ 1− 2αtξ + 2αtφ ] + αtφE [ d−1∑ i=1 ‖∆t+1i ‖ 2 2 ] + αtφE [ 3∑ i=d+1 ‖∆ti‖22 ]
We obtained
E[‖∆t+1d ‖ 2 2] ≤ (αt)2σ2d + [1− 2αtξ + 2αtφ]E[‖∆td‖22] + αtφE [ d−1∑ i=1 ‖∆t+1i ‖ 2 2 ] + αtφE [ 3∑ i=d+1 ‖∆ti‖22 ] we next re-group the terms as follows
E[‖∆t+1d ‖ 2 2]− αtφE [ d−1∑ i=1 ‖∆t+1i ‖ 2 2 ] ≤ [1− 2αtξ + 2αtφ]E[‖∆td‖22] + αtφE [ 3∑ i=d+1 ‖∆ti‖22 ] + (αt)2σ2d
and then sum over d from 1 to 3
E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ] − αtφE [ 3∑ d=1 d−1∑ i=1 ‖∆t+1i ‖ 2 2 ]
≤ [1− 2αtξ + 2αtφ]E [ 3∑ d=1 ‖∆td‖22 ] + αtφE [ 3∑ d=1 3∑ i=d+1 ‖∆ti‖22 ] + 2(αt)2 3∑ d=1 σ2d
Note that
E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ] − 2αtφE [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ] ≤ E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ] − αtφE [ 3∑ d=1 d−1∑ i=1 ‖∆t+1i ‖ 2 2 ] and
[1− 2αtξ + 2αtφ]E [ 3∑ d=1 ‖∆td‖22 ] + αtφE [ 3∑ d=1 3∑ i=d+1 ‖∆ti‖22 ] + (αt)2σ2
≤ [1− 2αtξ + 2αtφ]E [ 3∑ d=1 ‖∆td‖22 ] + 2αtφE [ 3∑ d=1 ‖∆td‖22 ] + (αt)2σ2
Combining these two facts with our previous results yields:
[1− 2αtφ]E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ [1− 2αtξ + 2αtφ]E [ 3∑ d=1 ‖∆td‖22 ] + 2αtφE [ 3∑ d=1 ‖∆td‖22 ] + (αt)2σ2
= [1− 2αtξ + 2αtφ]E [ 3∑ d=1 ‖∆td‖22 ] + (αt)2σ2
Thus:
E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ] ≤ 1− 2α tξ + 4αtφ 1− 2αtφ E [ 3∑ d=1 ‖∆td‖22 ]
+ (αt)2
1− 2αtφ σ2.
Since φ < ξ3 , 1−2αtξ+4αtφ 1−2αtφ < 1.
Proof of Theorem 10.10
To obtain the final theorem we need to expand the recursion from Theorem 10.4. We obtained
E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ 1− 2α t[ξ − 2φ]
1− 2αtφ E [ 3∑ d=1 ‖∆td‖22 ] + (αt)2 1− 2αtφ σ2
= ( 1− α
t[2ξ − 6φ] 1− 2αtφ
) E [ 3∑ d=1 ‖∆td‖22 ] + (αt)2 1− 2αtφ σ2
Recall that we defined qt in Theorem 10.4 as
qt = 1− 1− 2α tξ + 4αtφ 1− 2αtφ = αt[2ξ − 6φ] 1− 2αtφ
and denote
f t = (αt)2
1− 2αtφ .
Thus we have
E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ] ≤ (1− qt)E [ 3∑ d=1 ‖∆td‖22 ] + f tσ2
≤ (1− qt) { (1− qt−1)E [ 3∑ d=1 ‖∆t−1d ‖ 2 2 ] + f t−1σ2 } + f tσ2
= (1− qt)(1− qt−1)E [ 3∑ d=1 ‖∆t−1d ‖ 2 2 ] + (1− qt)f t−1σ2 + f tσ2
≤ (1− qt)(1− qt−1) { (1− qt−2)E [ 3∑ d=1 ‖∆t−2d ‖ 2 2 ] + f t−2σ2 } + (1− qt)f t−1σ2 + f tσ2
= (1− qt)(1− qt−1)(1− qt−2)E [ 3∑ d=1 ‖∆t−2d ‖ 2 2 ] +(1− qt)(1− qt−1)f t−2σ2 + (1− qt)f t−1σ2 + f tσ2
We end-up with the following
E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ] ≤ E [ 3∑ d=1 ‖∆0d‖22 ] t∏ i=0 (1− qi) + σ2 t−1∑ i=0 f i t∏ j=i+1 (1− qj) + f tσ2.
Set qt = 3 2
t+2 and
αt = qt
2ξ − 6φ+ 2qtφ
= 3 2
[2ξ − 6φ](t+ 2) + 3φ .
Denote A = 2ξ − 6φ and B = 3φ. Thus
αt = 3 2
A(t+ 2) +B
and
f t = (αt)2
1− 23Bαt =
9 4
A(t+ 2)[A(t+ 2) +B] .
E [ 2∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ E [ 2∑ d=1 ‖∆0d‖22 ] t∏ i=0 ( 1− 3 2 i+ 2 ) + σ2 t−1∑ i=0
9 4
A(i+ 2)[A(i+ 2) +B] t∏ j=i+1 ( 1− 3 2 j + 2 )
+σ2 9 4
A(t+ 2)[A(t+ 2) +B]
= E [ 2∑ d=1 ‖∆0d‖22 ] t+2∏ i=2 ( 1− 3 2 i ) + σ2 t+1∑ i=2
9 4
Ai[Ai+B] t+2∏ j=i+1 ( 1− 3 2 j ) + σ2
9 4
A(t+ 2)[A(t+ 2) +B]
Since A > 0 and B > 0 thus
E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ E [ 3∑ d=1 ‖∆0d‖22 ] t+2∏ i=2 ( 1− 3 2 i ) + σ2 t+1∑ i=2
9 4
Ai[Ai+B] t+2∏ j=i+1 ( 1− 3 2 j ) + σ2
9 4
A(t+ 2)[A(t+ 2) +B]
≤ E [ 3∑ d=1 ‖∆0d‖22 ] t+2∏ i=2 ( 1− 3 2 i ) + σ2 t+1∑ i=2 9 4 (Ai)2 t+2∏ j=i+1 ( 1− 3 2 j ) + σ2
9 4
[A(t+ 2)]2
We can next use the fact that for any a ∈ (1, 2):
t+2∏ i=τ+1 ( 1− a i ) ≤ ( τ + 1 t+ 3 )a .
The bound then becomes
E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ E [ 3∑ d=1 ‖∆0d‖22 ] t+2∏ i=2 ( 1− 3 2 i ) + σ2 t+1∑ i=2 9 4 (Ai)2 t+2∏ j=i+1 ( 1− 3 2 j ) + σ2
9 4
[A(t+ 2)]2
≤ E [ 3∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 t+1∑ i=2 9 4 (Ai)2 ( i+ 1 t+ 3 ) 3 2 + σ2 9 4 [A(t+ 2)]2
= E [ 3∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 t+2∑ i=2 9 4 (Ai)2 ( i+ 1 t+ 3 ) 3 2
Note that (i+ 1) 3 2 ≤ 2i for i = 2, 3, . . . , thus
E [ 3∑ d=1 ‖∆t+1d ‖ 2 2 ]
≤ E [ 2∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 9 4 A2(t+ 3) 3 2 t+2∑ i=2 (i+ 1) 3 2 i2
≤ E [ 3∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 9 2 A2(t+ 3) 3 2 t+2∑ i=2 1 i 1 2
finally note that t+2∑ i=2 1 i 1 2 ≤ ∫ t+2 1 1 x 1 2 dx ≤ 2(t+ 3) 12 . Thus
≤ E [ 3∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 9 A2(t+ 3)
substituting A = 2ξ − 6φ gives
= E [ 3∑ d=1 ‖∆0d‖22 ]( 2 t+ 3 ) 3 2 + σ2 9 [2ξ − 6φ]2(t+ 3)
This leads us to the final theorem. To obtain Theorem 10.8, we proceed in an analogous way as form Theorem 10.3, but this time applying Theorem 10.10 that we have just proved.

heading: 10.5 BEHAVIOR GUIDED POLICY GRADIENT AND WASSERSTEIN TRUST REGION
text: The chief goal of this section is to prove Theorem 5.1. We restate the section’s definitions here for the reader’s convenience: To ease the discussion we make the following assumptions:
1. Finite horizon T . 2. Undiscounted MDP. 3. States are time indexed. In other words, states visited at time t can’t be visited at any other time. 4. S and A are finite sets. The third assumption is solely to avoid having to define a time indexed Value function. It can be completely avoided. We chose not to do this in the spirit of notational simplicity. These assumptions can be relaxed, most notably we can show similar results for the discounted and infinite horizon case. We chose to present the finite horizon proof because of the nature of our experimental results.
Let Φ = id be the identity embedding so that E = Γ. In this case PΦπ denotes the distribution of trajectories corresponding to policy π. We define the value function V π : S → R as
V π(st = s) = Eτ∼Pidπ [ T∑ `=t R(s`+1, a`, s`)|st = s ] The Q-function Qπ : S ×A → R as:
Qπ(st, at = a) = Eτ∼Pidπ [ T∑ `=t R(s`+1, a`, s`) ] Similarly, the advantage function is defined as:
Aπ(s, a) = Qπ(s, a)− V π(s) We denote by V (π) = Eτ∼Pidπ [∑T t=0R(st+1, at, st) ]
the expected reward of policy π and define the visitation frequency as:
ρπ(s) = Eτ∼Pidπ [ T∑ t=0 1(st = s) ] The first observation in this section is the following lemma:
Lemma 10.13. two distinct policies π and π̃ can be related via the following equation :
V (π̃) = V (π) + ∑ s∈S
( ρπ̃(s) (∑ a∈A π̃(a|s)Aπ(s, a) ))
Proof. Notice that Aπ(s, a) = Es′∼P (s′|a,s) [R(s′, a, s) + V π(s′)− V π(s)]. Therefore:
Eτ∼Pidπ̃ [ T∑ t=0 Aπ(st, at) ] = Eτ∼Pidπ̃ [ T∑ t=0 R(st+1, at, st) + V π(st+1)− V π(st) ]
= Eτ∼Pidπ̃ [ T∑ t=0 R(st+1, at, st) ] − Es0 [V π(s0)]
= −V (π) + V (π̃)
The result follows.
See Sutton et al. (1998) for an alternative proof. We also consider the following linear approximation to V around policy π (see: Kakade & Langford (2002)):
L(π̃) = V (π) + ∑ s∈S
( ρπ(s) (∑ a∈A π̃(a|s)Aπ(s, a) ))
Where the only difference is that ρπ̃ was substituted by ρπ. Consider the following embedding Φs : Γ → R|S| defined by (Φ(τ))s = ∑T t=0 1(st = s), and related cost function defined as: C(v,w) = ‖v −w‖1. Lemma 10.14. The Wasserstein distance WD0(PΦ s π̃ ,PΦ s π ) is related to visit frequencies since:
WD0(PΦ s π̃ ,PΦ s π ) ≥ ∑ s∈S |ρπ(s)− ρπ̃(s)|
Proof. Let Π be the optimal coupling between PΦsπ̃ and PΦ s π . Then:
WD0(PΦ s π̃ ,PΦ s π ) = Eu,v∼Π [‖u− v‖1] = ∑ s∈S Eu,v∼Π [|us − vs|]
Where us and vs denote the s ∈ S indexed entry of the u and v vectors respectively. Notice that for all s ∈ S the following is true:∣∣∣∣∣∣∣Eu∼PΦsπ [us]︸ ︷︷ ︸
ρπ(s)
−Ev∼PΦsπ [vs]︸ ︷︷ ︸ ρπ′ (s) ∣∣∣∣∣∣∣ ≤ Eu,v∼Π [|us − vs|] The result follows.
These observations enable us to prove an analogue of Theorem 1 from Schulman et al. (2015), namely:
Theorem 10.15. If WD0(PΦ s π̃ ,PΦ s π ) ≤ δ and = maxs,a |Aπ(s, a)|, then V (π̃) ≥ L(θ̃)− δ .
As in Schulman et al. (2015), Theorem 5.1 implies a policy improvement guarantee for BGPG from Section 5.3.
Proof. Notice that:
V (π̃)− L(π̃) = ∑ s∈S
( (ρπ̃(s)− ρπ(s)) (∑ a∈A π̃(a|s)Aπ(s, a) )) Therefore by Holder inequality:
|V (π̃)− L(π̃)| ≤ (∑ s∈S |ρπ(s)− ρπ̃(s)| ) ︸ ︷︷ ︸ ≤WD0(PΦ s π̃ ,PΦ s π )≤δ ( sup s∈S ∣∣∣∣∣∑ a∈A π̃(a|s)Aπ(s, a) ∣∣∣∣∣ ) ︸ ︷︷ ︸ ≤
The result follows.
We can leverage the results of Theorem 10.15 to show wasserstein trust regions methods with embedding Φs give a monotonically improving sequence of policies. The proof can be concluded by following the logic of Section 3 in Schulman et al. (2015).

Review summaries:
This paper proposes using Wasserstein Distances to measure the difference between higher-level functions of policies, which this paper terms as "behaviors". For example, one such behavior could be the distribution over final states given the policy, or the distribution over returns given policy. Through the lens of these behavioral embeddings, this paper recovers a few important special cases that are well-known in the literature including WD-based TRPO and distributional RL. This paper shows that the dual formulation of the Wasserstein Distance gives the ability to score individual policies based on a given "behavioral mapping".

this paper proposes a new regularized policy optimization (PO) method which is based on Wasserstein distances. Its idea is to use SGD to optimize the dual form of the WD, then used in two different policy search approaches TRPO and Evolution Strategies. The evaluations are carried out on a variety of control tasks from OpenAI Gym.