Title: BEHAVIOR-GUIDED REINFORCEMENT LEARNING

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

1. INTRODUCTION:

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

2. MOTIVATING BEHAVIOR-GUIDED REINFORCEMENT LEARNING:

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.

3. DEFINING BEHAVIOR IN REINFORCEMENT LEARNING:

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 θ ∈ R d ) from S to A. Let Γ = {τ = s 0 , a 0 , r 0 , • • • s H , a H , r H s.t. s i ∈ S, a i ∈ A, r i ∈ R} be the set of possible trajectories enriched by sequences of partial rewards under some policy π. The undiscounted reward function R : Γ → R (which expectation is to be maximized by optimizing θ) satisfies R(τ ) = H i=0 r i , where r i = R(s i+1 , a i , s i ). 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 π .

3.1. BEHAVIORAL EMBEDDINGS:

Additionally, we require the BES E to be equipped with a metric (or cost function) H t=0 e st (where e s ∈ R |S| is the one-hot vector corresponding to state s); see also Section 7.2. 2. Action-based: the concatenation of actions Φ 4 (τ ) = [a 0 , ..., a H ]; see also Section 7.1. 3. Reward-based: the total reward Φ 5 (τ ) = H t=0 r t , reward-to-go vector Φ 6 (τ ) = H t=0 r t t i=0 e i (where e i ∈ R H+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.

C : E × E → R. Given two trajectories τ 1 , τ 2 in Γ, C(Φ(τ 1 ), Φ(τ 2 )) measures how different

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.

4. WASSERSTEIN DISTANCE & OPTIMAL TRANSPORT PROBLEM:

Let µ, ν be (Radon) probability measures over domains X ⊆ R m , Y ⊆ R n and let C : X × Y → R be a cost function. For γ > 0, a smoothed Wasserstein Distance is defined as:

EQUATION

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

4.1. WASSERSTEIN DISTANCE: DUAL FORMULATION:

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 (x t , y t ) ∼ µ ν. 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:

EQUATION

where E C (λ µ , λ ν ) is defined as:

EQUATION

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. 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 λ ∈ R m . For the Gaussian kernel, φ is defined as follows:

φ(z) = 1 √ m cos(Gz + b) for z ∈ R d

, where G ∈ R m×d is Gaussian with iid entries taken from N (0, 1), b ∈ R m with iid b i s such that b i ∼ 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 (x 1 , y 1 ),

• • • , (x k , y k ) i.i.d

∼ µ ν, then Algorithm 1 can be used to get an estimator of WD γ (µ, ν) as follows:

5 BEHAVIOR-GUIDED REINFORCEMENT LEARNING 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 θ ∈ R d . 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(τ )].

5.1. BEHAVIORAL TEST FUNCTIONS:

If C : E × E → R is a cost function defined over behavior space E, and π 1 , π 2 are two policies, then:

EQUATION

where λ * 1 , λ * 2 are the optimal dual functions. The maps s 1 := λ * 1 • Φ : Γ → R and s 2 := λ * 2 • Φ : Γ → R define score functions over the space of trajectories. If γ is close to zero, the score function s i gives higher scores to trajectories from π i whose behavioral embedding is common under π i but rarely appears under π j for j = i (Fig. 2 ).

5.2. ALGORITHMS:

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:

EQUATION

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:

EQUATION

where s 1 : Γ → R equals s 1 = λ * 1 •Φ, the Behavioral Test Function of policy π θ and λ * 2 is the optimal dual function of embedding distribution

P Φ b . Consequently ∇ θ F (θ) ≈ ∇ θ E τ ∼Pπ θ [R(τ ) + βs 1 (τ )].

We learn a score function s 1 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.

5.3. BEHAVIOR-GUIDED POLICY GRADIENT (BGPG):

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 } M i=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 (s i , a i ) π θ (a i |s i ) π t-1 (a i |s i ) + βλ 1 (Φ(τ 1 )) -βλ 2 (Φ(τ 2 )) + βγ exp λ 1 (Φ(τ 1 )) -λ 2 (Φ(τ 2 )) -C(Φ(τ 1 )), Φ(τ 2 )) γ

Where

τ 1 = s 0 , a 0 , r 0 , • • • , s H , a H , r H . Let θ (0) t-1 = θ t-1 . for = 1, • • • , L do a. Approximate P πt-1 P π θ via 1 M {τ (t) i } M i=1 1 M {τ θ i } M i=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(s t = 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(s t = s) and related cost function defined as:

C(v, w) = v -w 1 . Then WD 0 (P Φ s π , P Φ s π ) is related to visitation frequencies since WD 0 (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 WD 0 (P Φ s π , P Φ s π ) ≤ δ and = max s,a |A π (s, a)|, then V (π) ≥ L( θ) -δ .

As in Schulman et al. (2015) , Theorem 5.1 implies a policy improvement guarantee for BGPG.

5.4. BEHAVIOR GUIDED EVOLUTION STRATEGIES (BGES):

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 . 4. Update λ 1 and λ 2 using Algorithm 1, where

for t = 1, . . . , T -1 do 1. Sample 1 , • • • , n independently from N (0, I). 2. Evaluate policies {π k t } n k=1 parameterized by {θ t + σ k } n k=1 to return rewards R k and trajecto- ries τ k for all k.

µ = 1 n ∪ n k=1 PΦ π k t-1 and ν = 1 n ∪ n k=1 PΦ π k t

are the uniform distribution over the set of PPEs from 3 for t -1 and t. 5. Approximate WDγ(P Φ π k t , P Φ πt ) plugging in λ 1 , λ 2 into Eq. 4 for each perturbed policy π k 6. Update Policy: θ t+1 = θ t + η∇ ES F , where:

∇ ES F = 1 σ n k=1 [(1 -β)(R k -R t ) + β WD γ (P Φ π k t , 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.

6. RELATED WORK:

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, Under review as a conference paper at ICLR 2020 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 Fokker-Planck 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) .

7. EXPERIMENTS:

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.

7.1. BEHAVIOR-GUIDED POLICY GRADIENT:

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. 

7.2. BEHAVIOR-GUIDED EVOLUTION STRATEGIES:

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

7.3. IMITATION LEARNING:



7.4. HYPERPARAMETER SELECTION:



8. CONCLUSION AND FUTURE WORK:

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

APPENDIX: BEHAVIOR-GUIDED REINFORCEMENT LEARNING 9 FURTHER EXPERIMENTAL DETAILS

9.1. BGPG:

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 WD γ (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( xy 2 2 /σ 2 ) with σ = 0.1; the number of random features is D = 100; recall the embedding is

Φ(τ ) = [a 1 , a 2 ...a H ]

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 ∈ R D 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.

9.2. BGES:

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.

9.3. IMITATION LEARNING::

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.

9.4. REPULSION LEARNING:

Algorithm 4 Behvaior-Guided Repulsion Learning Input: β, η > 0, M ∈ N Initialize: Initial stochastic policies π a 0 , π b 0 , parametrized by θ a 0 , θ b 0 respectively, Behavioral Test Functions λ a 1 , λ b

2 for t = 1, . . . , T do 1. Collect M trajectories {τ a i } M i=1 from P π a t-1 and M trajectories {τ b i } M i=1 from P π b t-1 . Approximate P π a t-1 P π b t-1 via 1 M {τ a i } M i=1 1 M {τ b i } M i=1 := Pπ a t-1 ,π b t-1

2. Form two distinct surrogate rewards for joint trajectories of agents a and b:

EQUATION

3. For c ∈ {a, b} use the Reinforce estimator to take gradient steps:

θ c t = θ c t-1 + η E τ a ,τ b ∼ Pπ a t-1 ,π b t-1 Rc (τ a , τ b ) H-1 i=0 ∇ θ c t-1 log π c t-1 (a c i |s c i )

Where

τ a = s a 0 , a a 0 , r a 0 , • • • , s a H , a a H , r a H and τ b = s b 0 , a b 0 , r b 0 , • • • , s b H , a b H , r b H . 5. Use samples from Pπ a t-1 ,π b t-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 Ra and Rb that combines the signal from the task's reward function R 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, Goal 1 ), d(s, Goal 2 )) 2 where d(x, y) denotes the distance between x and y in R 2 . 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 λ * µ = λ * ν .

10. THEORETICAL RESULTS:

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 e st,at with e s,a ∈ R |S|+|A| the indicator vector for the state action pair (s, a). Let

C(v, w) = v -w p p 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:

E u,v∼Π u -v p p = 0 Consequently: E u,v∼Π   (s,a)∈S×A |u s,a -v s,a | p   = (s,a)∈S×A E u,v∼Π [|u s,a -v s,a | p ] = 0

Therefore for all (s, a) ∈ S × A:

E u∼P Φ π [u s,a ] -E v∼P Φ π [v s,a ] p ≤ E u,v∼Π [|u s,a -v s,a | p ] = 0

Where u s,a and v s,a denote the (s, a) entries of u and v respectively. Notice that for all (s, a) ∈ S ×A:

EQUATION

) Since for all s ∈ S and p ≥ 1:

a∈A u s,a -v s,a p ≤ a∈A |u s,a -v s,a | p

Therefore for all s ∈ S:

E u∼P Φ π a∈A u s,a -E v∼P Φ π a∈A v s,a p ≤ E u,v∼Π a∈A |u s,a -v s,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 e st,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.

10.1. RANDOM FEATURES STOCHASTIC GRADIENTS:

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

10.2. BGPG, BGES AND THEIR THEORETICAL GUARANTEES:

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

10.3. MAX-MAX PROBLEM: THEORETICAL ANALYSIS:

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 neighborhood N (θ * , r) of its optimal value θ * . We will also assume that gradient of F is L 2 -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: φ γ = 1 2γ e 4 γ , and u γ = 8 3γ 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

EQUATION

)

where σ = 2(1 + e 2 γ ) 2 + sup N (θ * ),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 i th entry of the feature vector is constructed as follows:

[f (z)] i = 2 D cos(z w z,i + b z,i ), where w z,i ∼ N 0, I n 1 ρ 2

and b z,i ∼ Uniform(0, 2π). For the ease of notation we denote κ z,i = (w z,i , b z,i ) ∼ ω. We consider stochastic gradient ascent optimization strategy of the following form:

• at time t we receive a single sample

(x t , y t , κ xt,1 t , . . . , κ xt,D t , κ yt,1 t , . . . , κ yt,D t ) ∼ µ × ν × ω × • • • × ω, then we form feature vectors [f (x t )] i = 2 D cos(x t w xt,i t + b xt,i t ) and [f (y t )] i = 2 D cos(y t w yt,i t + b yt,i t )

, and finally update:

β t+1 1 ← Π 1 β t 1 + α t f (x t ) 1 -exp (β t 1 ) f (xt)+(β t 2 ) f (yt)-C(x,y) γ β t+1 2 ← Π 2 β t 2 + α t f (y t ) 1 -exp (β t+1 1 ) f (xt)+(β t 2 ) f (yt)-C(x,y) γ . Π d (d = 1, 2) denotes the projection onto the Euclidean ball B 2 ( r d 2 , β 0 d ) of some given radius r d 2 centered at the initial iterate β 0 d . 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 B 2 (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 r 1 , r 2 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 B 2 (r 1 , β * 1 ) or B 2 (r 2 , β * 2 ), respectively, the projection is applied that pushes the parameter of interest to stay in the ball. Also, let ∇ β1 W 1 γ (β 1 , β 2 ) and ∇ β2 W 1 γ (β 1 , β 2 ) denote the gradients of W γ with respect to β 1 and β 2 , respectively, computed for a single sample. Similarly, ∇ β1 W γ (β 1 , β 2 ) and ∇ β2 W γ (β 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 d in the ball of radius r d 2 centered at β * d , we are guaranteed that all iterates remain within an r d -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, G d (β 1 , β 2 ), where d = 1, 2, be defined as 

G d (β 1 , β 2 ) := β d + α∇ β d W γ (β 1 , β 2 ). 10.3.1 ASSUMPTIONS 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 .

W * γ,d (a d ) -W * γ,d (b d ) -∇ β d W * γ,d (b d ), a d -b d ≤ - ζ d 2 a d -b d 2 2 ,

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 (a d , b d ) the following holds W * γ,d (a d ) -W * γ,d (b d ) -∇ β d W * γ,d (b d ), a d -b d ≥ - δ d 2 a d -b d 2 2 ,

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 ∈ B 2 (r 1 , β * 1 ), β 2 ∈ B 2 (r 2 , β * 2 ) given as follows ∇ β d W * γ,d (β d ) -∇ β d W γ (β 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:

σ = σ 2 1 + σ 2 2 , where σ 2 d = sup{E[ ∇ β d W 1 γ (β 1 , β 2 ) 2 2 ] : β 1 ∈ B 2 (r 1 , β * 1 ), β 2 ∈ B 2 (r 2 , β * 2 )} for d = 1, 2.

10.3.2. MAIN THEOREMS:

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+1 1 -β * 1 2 2 + β t+1 2 -β * 2 2 2 ≤ (1 -q t )E β t 1 -β * 1 2 2 + β t 2 -β * 2 2 2 + (α t ) 2 1 -α t φ σ 2 , where φ = max d=1,2 (φ d ), q t = 1 -1-2α t ξ+2α t φ 1-α t φ

, and ξ = min d=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)+ 3 2 φ and assuming that φ < 2ξ 3 , the error at iteration t + 1 satisfies

E β t+1 1 -β * 1 2 2 + β t+1 2 -β * 2 2 2 ≤ E β 0 1 -β * 1 2 2 + β 0 2 -β * 2 2 2 2 t + 3 3 2 +σ 2 9 [2ξ -3φ] 2 (t + 3)

,

where φ = max d=1,2 (φ d ) and ξ = min d=1,2

2δ d ζ d δ d +ζ d .

10.3.3. ANALYSIS:

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 G 1 (β 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.

G 1 (β 1 , β * 2 ) -β * 1 2 ≤ 1 - 2αδ 1 ζ 1 δ 1 + ζ 1 β 1 -β * 1 2 (12)

for all β 1 ∈ B 2 (r 1 , β * 1 ). Similarly, the gradient operator G 2 (β * 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.

G 2 (β * 1 , β 2 ) -β * 2 2 ≤ 1 - 2αδ 2 ζ 2 δ 2 + ζ 2 β 2 -β * 2 2 (13) for all β 2 ∈ B 2 (r 2 , β * 2 ).

The next theorem also holds for d = 1, 2. Let r 1 , r 2 > 0 and Proof.

β 1 ∈ B 2 (r 1 , β * 1 ), β 2 ∈ B 2 (r 2 , β * 2 ).

G d (β 1 , β 2 ) -β * d 2 = β d + α∇ β d W γ (β 1 , β 2 ) -β * d 2

by the triangle inequality (and with d = d) we further get

≤ β d + α∇ β d W * γ,d (β d ) -β * d 2 + α ∇ β d W γ (β 1 , β 2 ) -∇ β d W * γ,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+1 d = Π d ( βt+1 d )

, where βt+1

1 := β t 1 + α t ∇ β1 W 1 γ (β t 1 , β t 2 ) and βt+1 2 := β t 2 + α t ∇ β2 W 1 γ (β t+1 1 , β t 2 )

, where ∇ β d W 1 γ is the gradient computed with respect to a single sample, β1 and β2 are the updates prior to the projection onto a ball B 2 (

r d 2 , β 0 d ). Let ∆ t+1 d := β t+1 d -β * d and ∆t+1 d := βt+1 d -β * d . Thus ∆ t+1 d 2 2 -∆ t d 2 2 ≤ ∆t+1 d 2 2 -∆ t d 2 2 = βt+1 d -β * d -β t d -β * d = βt+1 d -β t d , βt+1 d + β t d -2β * d . Let Qt 1 := ∇ β1 W 1 γ (β t 1 , β t 2 ) and Qt 2 := ∇ β2 W 1 γ (β t+1 1 , β t 2 ). Then we have that βt+1 d -β t d = α t Qt d .

We combine it with Equation 17 and obtain:

∆ t+1 d 2 2 -∆ t d 2 2 ≤ α t Qt d , α t Qt d + 2(β t d -β * d ) = (α t ) 2 ( Qt d ) Qt d + 2α t ( Qt d ) (β t d -β * d ) = (α t ) 2 Qt d 2 2 + 2α t Qt d , ∆ t d . Let Q t 1 := ∇ β1 W γ (β t 1 , β t 2 ) and Q t 2 := ∇ β2 W γ (β t+1 1 , β t 2 )

. By the properties of martingales, i.e. iterated expectations and tower property: 

E[ ∆ t+1 d 2 2 ] ≤ E[ ∆ t d 2 2 ] + (α t ) 2 E[ Qt d 2 2 ] + 2α t E[ Q t d , ∆ t d ] (15) Let Q * d := ∇ β d W γ (β * 1 , β * 2 ).

Q * d , ∆ t d = ∇ β d W γ (β * 1 , β * 2 ), ∆ t d = 0.

Combining this with Equation 15we have

E[ ∆ t+1 d 2 2 ] ≤ E[ ∆ t d 2 2 ] + (α t ) 2 E[ Qt d 2 2 ] + 2α t E[ Q t d -Q * d , ∆ t d ]. Define G t d := β t d + α t Q t d and G t * d := β * d + α t Q * d . Thus α t Q t d -Q * d , ∆ t d = G t d -G t * d -(β t d -β * d ), β t d -β * d = G t d -G t * d , β t d -β * d -β t d -β * d 2 2

by the fact that

G t * d = β * d + α t Q * d = β * d (since Q * d = 0): = G t d -β * d , β t d -β * d -β t d -β * d 2 2

by the contractivity of G t from Theorem 10.7:

≤ (1 -α t ξ) β t d -β * d + α t φ d-1 i=1 β t+1 i -β * i 2 + 2 i=d+1 β t i -β * i 2 β t d -β * d 2 -β t d -β * d 2 2 ≤ (1 -α t ξ) ∆ t d 2 + α t φ d-1 i=1 ∆ t+1 i 2 + 2 i=d+1 ∆ t i 2 • ∆ t d 2 -∆ t d 2 2

Combining this result with Equation 16gives

E[ ∆ t+1 d 2 2 ] ≤ E[ ∆ t d 2 2 ] + (α t ) 2 E[ Qt d 2 2 ] + 2E (1 -α t ξ) ∆ t d 2 + α t φ d-1 i=1 ∆ t+1 i 2 + 2 i=d+1 ∆ t i 2 • ∆ t d 2 -∆ t d 2 2 ≤ E[ ∆ t d 2 2 ] + (α t ) 2 σ 2 d + 2E (1 -α t ξ) ∆ t d 2 + α t φ d-1 i=1 ∆ t+1 i 2 + 2 i=d+1 ∆ t i 2 • ∆ t d 2 -∆ t d 2 2 .

After re-arranging the terms we obtain

E[ ∆ t+1 d 2 2 ] ≤ (α t ) 2 σ 2 d + (1 -2α t ξ)E[ ∆ t d 2 2 ] + 2α t φE d-1 i=1 ∆ t+1 i 2 + 2 i=d+1 ∆ t i 2 ∆ t d 2 apply 2ab ≤ a 2 + b 2 : ≤ (α t ) 2 σ 2 d + (1 -2α t ξ)E[ ∆ t d 2 2 ] + α t φE d-1 i=1 ∆ t+1 i 2 2 + ∆ t d 2 2 + α t φE 2 i=d+1 ∆ t i 2 2 + ∆ t d 2 2 = (α t ) 2 σ 2 d + E[ ∆ t d 2 2 ] • 1 -2α t ξ + α t φ + α t φE d-1 i=1 ∆ t+1 i 2 2 + α t φE 2 i=d+1 ∆ t i 2 2

We obtained

E[ ∆ t+1 d 2 2 ] ≤ (α t ) 2 σ 2 d + [1 -2α t ξ + α t φ]E[ ∆ t d 2 2 ] + α t φE d-1 i=1 ∆ t+1 i 2 2 + α t φE 2 i=d+1 ∆ t i 2 2

we next re-group the terms as follows

E[ ∆ t+1 d 2 2 ] -α t φE d-1 i=1 ∆ t+1 i 2 2 ≤ [1 -2α t ξ + α t φ]E[ ∆ t d 2 2 ] + α t φE 2 i=d+1 ∆ t i 2 2 + (α t ) 2 σ 2 d and then sum over d from 1 to 2 E 2 d=1 ∆ t+1 d 2 2 -α t φE 2 d=1 d-1 i=1 ∆ t+1 i 2 2 ≤ [1 -2α t ξ + α t φ]E 2 d=1 ∆ t d 2 2 + α t φE 2 d=1 2 i=d+1 ∆ t i 2 2 + (α t ) 2 2 d=1 σ 2 d Let σ = σ 2 1 + σ 2 2 . Also, note that E 2 d=1 ∆ t+1 d 2 2 -α t φE 2 d=1 ∆ t+1 d 2 2 ≤ E 2 d=1 ∆ t+1 d 2 2 -α t φE 2 d=1 d-1 i=1 ∆ t+1 i 2 2

and

[1 -2α t ξ + α t φ]E 2 d=1 ∆ t d 2 2 + α t φE 2 d=1 2 i=d+1 ∆ t i 2 2 + (α t ) 2 σ 2 ≤ [1 -2α t ξ + α t φ]E 2 d=1 ∆ t d 2 2 + α t φE 2 d=1 ∆ t d 2 2 + (α t ) 2 σ 2

Combining these two facts with our previous results yields:

[1 -α t φ]E 2 d=1 ∆ t+1 d 2 2 ≤ [1 -2α t ξ + α t φ]E 2 d=1 ∆ t d 2 2 + α t φE 2 d=1 ∆ t d 2 2 + (α t ) 2 σ 2 = [1 -2α t ξ + 2α t φ]E 2 d=1 ∆ t d 2 2 + (α t ) 2 σ 2

Thus:

E 2 d=1 ∆ t+1 d 2 2 ≤ 1 -2α t ξ + 2α t φ 1 -α t φ E 2 d=1 ∆ t d 2 2 + (α 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+1 d 2 2 ≤ 1 -2α t [ξ -φ] 1 -α t φ E 2 d=1 ∆ t d 2 2 + (α t ) 2 1 -α t φ σ 2 = 1 - α t [2ξ -3φ] 1 -α t φ E 2 d=1 ∆ t d 2 2 + (α t ) 2 1 -α t φ σ 2

Recall that we defined q t in Theorem 10.4 as

q t = 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+1 d 2 2 ≤ (1 -q t )E 2 d=1 ∆ t d 2 2 + f t σ 2 ≤ (1 -q t ) (1 -q t-1 )E 2 d=1 ∆ t-1 d 2 2 + f t-1 σ 2 + f t σ 2 = (1 -q t )(1 -q t-1 )E 2 d=1 ∆ t-1 d 2 2 + (1 -q t )f t-1 σ 2 + f t σ 2 ≤ (1 -q t )(1 -q t-1 ) (1 -q t-2 )E 2 d=1 ∆ t-2 d 2 2 + f t-2 σ 2 + (1 -q t )f t-1 σ 2 + f t σ 2 = (1 -q t )(1 -q t-1 )(1 -q t-2 )E 2 d=1 ∆ t-2 d 2 2 +(1 -q t )(1 -q t-1 )f t-2 σ 2 + (1 -q t )f t-1 σ 2 + f t σ 2

We end-up with the following

E 2 d=1 ∆ t+1 d 2 2 ≤ E 2 d=1 ∆ 0 d 2 2 t i=0 (1 -q i ) + σ 2 t-1 i=0 f i t j=i+1 (1 -q j ) + f t σ 2 . Set q t = 3 2 t+2 and α t = q t 2ξ -3φ + q t φ = 3 2 [2ξ -3φ](t + 2) + 3 2 φ .

Denote A = 2ξ -3φ and B = 3 2 φ. Thus

α t = 3 2 A(t + 2) + B and f t = (α t ) 2 1 -2 3 Bα t = 9 4 A(t + 2)[A(t + 2) + B] . E 2 d=1 ∆ t+1 d 2 2 ≤ E 2 d=1 ∆ 0 d 2 2 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 ∆ 0 d 2 2 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+1 d 2 2 ≤ E 2 d=1 ∆ 0 d 2 2 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 ∆ 0 d 2 2 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+1 d 2 2 ≤ E 2 d=1 ∆ 0 d 2 2 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 ∆ 0 d 2 2 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 ∆ 0 d 2 2 2 t + 3 3 2 + σ 2 t+2 i=2 9 4 (Ai) 2 i + 1 t + 3 3 2 3 2 ≤ 2i for i = 2, 3, . . . , thus E 2 d=1 ∆ t+1 d 2 2 ≤ E 2 d=1 ∆ 0 d 2 2 2 t + 3 3 2 + σ 2 9 4 A 2 (t + 3) 3 2 t+2 i=2 (i + 1) 3 2 i 2 ≤ E 2 d=1 ∆ 0 d 2 2 2 t + 3 3 2 + σ 2 9 2 A 2 (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) 1 2 . Thus ≤ E 2 d=1 ∆ 0 d 2 2 2 t + 3 3 2 + σ 2 9 A 2 (t + 3) substituting A = 2ξ -3φ gives = E 2 d=1 ∆ 0 d 2 2 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: 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

• function h v,C : R d → R defined as follows: h v,C (w) = w v -Ae w v λ for A = γe -C γ is 2 γ e 4 γ -

EQUATION

)

where σ = 2(1 + e 2 γ ) 2 + sup N (θ * ),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)) γ , by the fact that G t * d = θ * d + α t Q * d = θ * d (since Q * d = 0): = G t d -θ * d , θ t d -θ * d -θ t d -θ * d 2 2

by the contractivity of G t from Theorem 10.7:

≤ (1 -α t ξ) θ t d -θ * d + α t φ d-1 i=1 θ t+1 i -θ * i 2 + 3 i=d+1 θ t i -θ * i 2 θ t d -θ * d 2 -θ t d -θ * d 2 2 ≤ (1 -α t ξ) ∆ t d 2 + α t φ d-1 i=1 ∆ t+1 i 2 + 3 i=d+1 ∆ t i 2 • ∆ t d 2 -∆ t d 2 2 Thus E[ ∆ t+1 d 2 2 ] ≤ E[ ∆ t d 2 2 ] + (α t ) 2 E[ Qt d 2 2 ] + 2α t E[ Q t d -Q * d , ∆ t d ] = E[ ∆ t d 2 2 ] + (α t ) 2 E[ Qt d 2 2 ] + 2E (1 -α t ξ) ∆ t d 2 + α t φ d-1 i=1 ∆ t+1 i 2 + 3 i=d+1 ∆ t i 2 • ∆ t d 2 -∆ t d 2 2 ≤ E[ ∆ t d 2 2 ] + (α t ) 2 σ 2 d + 2E (1 -α t ξ) ∆ t d 2 + α t φ d-1 i=1 ∆ t+1 i 2 + 3 i=d+1 ∆ t i 2 • ∆ t d 2 -∆ t d 2 2 .

After re-arranging the terms we obtain Thus:

E[ ∆ t+1 d 2 2 ] ≤ (α t ) 2 σ 2 d + (1 -2α t ξ)E[

E 3 d=1 ∆ t+1 d 2 2 ≤ 1 -2α t ξ + 4α t φ 1 -2α t φ E 3 d=1 ∆ t d 2 2 + (α 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+1 d 2 2 ≤ 1 -2α t [ξ -2φ] 1 -2α t φ E 3 d=1 ∆ t d 2 2 + (α t ) 2 1 -2α t φ σ 2 = 1 - α t [2ξ -6φ] 1 -2α t φ E 3 d=1 ∆ t d 2 2 + (α t ) 2 1 -2α t φ σ 2

Recall that we defined q t in Theorem 10.4 as

q t = 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+1 d 2 2 ≤ (1 -q t )E 3 d=1 ∆ t d 2 2 + f t σ 2 ≤ (1 -q t ) (1 -q t-1 )E 3 d=1 ∆ t-1 d 2 2 + f t-1 σ 2 + f t σ 2 = (1 -q t )(1 -q t-1 )E 3 d=1 ∆ t-1 d 2 2 + (1 -q t )f t-1 σ 2 + f t σ 2

≤ (1 -q t )(1 -q t-1 ) (1 -q t-2 )E 3 d=1 ∆ t-2 d 2 2 + f t-2 σ 2 + (1 -q t )f t-1 σ 2 + f t σ 2

= (1 -q t )(1 -q t-1 )(1 -q t-2 )E +(1 -q t )(1 -q t-1 )f t-2 σ 2 + (1 -q t )f t-1 σ 2 + f t σ 2

We end-up with the following

E 3 d=1 ∆ t+1 d 2 2 ≤ E 3 d=1 ∆ 0 d 2 2 t i=0 (1 -q i ) + σ 2 t-1 i=0 f i t j=i+1

(1 -q j ) + f t σ 2 .

Set q t = 3 2 t+2 and α t = q t 2ξ -6φ + 2q t φ = 3 2 [2ξ -6φ](t + 2) + 3φ

.

Denote A = 2ξ -6φ and B = 3φ. Thus 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.

α t =

10.5. BEHAVIOR GUIDED POLICY GRADIENT AND WASSERSTEIN TRUST REGION:

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 π (s t = s) = E τ ∼P id π T =t R(s +1 , a , s )|s t = s

The Q-function Q π : S × A → R as:

Q π (s t , a t = a) = E τ ∼P id π 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 τ ∼P id π T t=0 R(s t+1 , a t , s t ) the expected reward of policy π and define the visitation frequency as:

ρ π (s) = E τ ∼P id π T t=0 1(s t = 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) = E s ∼P (s |a,s) [R(s , a, s) + V π (s ) -V π (s)]. Therefore:

E τ ∼P id π T t=0 A π (s t , a t ) = E τ ∼P id π T t=0 R(s t+1 , a t , s t ) + V π (s t+1 ) -V π (s t ) = E τ ∼P id π T t=0 R(s t+1 , a t , s t ) -E s0 [V π (s 0 )] = -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(s t = s), and related cost function defined as: C(v, w) = vw 1 .

Lemma 10.14. The Wasserstein distance WD 0 (P Φ s π , P Φ s π ) is related to visit frequencies since:

WD 0 (P Φ s π , P Φ s π ) ≥ s∈S |ρ π (s) -ρ π (s)|

Proof. Let Π be the optimal coupling between P Φ s π and P Φ s π . Then:

WD 0 (P Φ s π , P Φ s π ) = E u,v∼Π [ u -v 1 ] = s∈S E u,v∼Π [|u s -v s |]

Where u s and v s denote the s ∈ S indexed entry of the u and v vectors respectively. Notice that for all s ∈ S the following is true:

E u∼P Φ s π [u s ] ρπ(s) -E v∼P Φ s π [v s ] ρ π (s) ≤ E u,v∼Π [|u s -v s |]

The result follows.

These observations enable us to prove an analogue of Theorem 1 from Schulman et al. (2015) , namely:

Theorem 10.15. If WD 0 (P Φ s π , P Φ s π ) ≤ δ and = max s,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.

:

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

If we choose an appropriate embedding map our framework handles visitation frequencies as well.

The 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

≤ 2i for i = 2,

3, . . . , thus

Figures and Tables:

Figure 1: Behavioral Embedding Maps (BEMs) map trajectories to points in the behavior embedding space E. Two trajectories may map to the same point in E.

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 (τ ) = s H , the visiting frequency of a fixed state Φ s 2 (τ ) = H t=0 1(s t = s), the frequency vector of visited states Φ 3 (τ ) =

Figure 2: Behavioral embedding functions corresponding to two policies π1 (green) and π2 (blue) whose BEMs map trajectories to points in the real line.

3. Use BEM to map trajectories τ k to produce empirical PPEs PΦ π k t for all k.

Figure 3: BGPG vs. TRPO: We compare BGPG and TRPO (KL divergence) on several continuous control tasks. As a baseline we also include results without a trust region (β = 0 in Algorithm 2). Plots show the mean ± std across 5 random seeds. BGPG consistently outperforms other methods.

Figure 4: The clock-time comparison (in sec) of BGPG (alternating optimization) with particle approximation.

Figure 5: Efficient Exploration. On the left we show a visualization of the simulated environment, with the deceptive barrier between the (quadruped) agent and the goal. On the right, we show two plots with the median curve across five seeds, with the IQR shaded for the quadruped and point environment respectively.

Figure 6: Escaping Local Maxima. A comparison of BGES with those using different distances on PPEs.

Figure 7: Imitation Learning.

Figure 8: Choice of BEM

Figure 9: Additional Experiment on TRPO. We compare No Trust Region with two alternative trust region constructions: KL-divergence and Wassertein distance (ours).

(a) Embeddings (b) Previous Policies

Figure 10: A sensitivity analysis investigating a) the impact of the embedding and b) the number of previous policies θt-i, i ∈ 1, 2, 5

Figure 11: Initial state of policies π a , π b and Behavioral Test functions λ a , λ b in the Multigoal environment.

Figure 12: Evolution of the policies and Behavioral Test Functions throughout optimization.

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 (a d , b d ) in the neighborhood of β * d the following holds

Theorem 10.7. For some radius r d > 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 operatorG d (β 1 , β 2 ) with step α such that 0 < α ≤ min d=1,2 2 δ d +ζ d is contractive over a ball B 2 (r d , β * d ), i.e. G d (β 1 , β 2 ) -β * d 2 ≤ (1 -ξα) β d -β * d 2 + αφ β d -β * d 2(14) where d = (d mod 2) + 1, φ := max d=1,2 φ d , and ξ := min d=1,2 2δ d ζ d δ d +ζ d .

By self-consistency, i.e. β * d = arg max β d ∈Ω d W * γ,d (β d ), and convexity of Ω d we have that

Table content: <table/>

≤ (α t ) 2 σ 2 d + [1 -2α t ξ + 2α t φ]E[ ∆ t

Table content: <table><tr><td colspan="2">We obtained</td><td/><td/><td/><td/><td/><td/><td/></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>d-1</td><td>3</td></tr><tr><td colspan="2">E[ ∆ t+1 d</td><td colspan="8">2 2 ] d</td><td>2 2 ] + α t φE</td><td>∆ t+1 i</td><td>2 2 + α t φE</td><td>∆ t i</td><td>2 2</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>i=1</td><td>i=d+1</td></tr><tr><td colspan="8">we next re-group the terms as follows</td><td/></tr><tr><td/><td/><td/><td/><td>d-1</td><td/><td/><td/><td/><td>3</td></tr><tr><td colspan="2">E[ ∆ t+1 d</td><td colspan="3">2 2 ] -α t φE</td><td colspan="2">∆ t+1 i</td><td colspan="3">2 2 ≤ [1 -2α t ξ + 2α t φ]E[ ∆ t d</td><td>2 2 ] + α t φE</td><td>∆ t i</td><td>2 2 + (α t ) 2 σ 2 d</td></tr><tr><td/><td/><td/><td/><td>i=1</td><td/><td/><td/><td/><td>i=d+1</td></tr><tr><td colspan="6">and then sum over d from 1 to 3</td><td/><td/><td/></tr><tr><td/><td>3</td><td/><td/><td/><td>3</td><td colspan="2">d-1</td><td/></tr><tr><td>E</td><td colspan="3">∆ t+1 d</td><td>2 2 -α t φE</td><td/><td/><td colspan="3">∆ t+1 i</td><td>2 2</td></tr><tr><td colspan="2">d=1</td><td/><td/><td/><td>d=1</td><td colspan="2">i=1</td><td/></tr><tr><td/><td/><td/><td/><td>3</td><td/><td/><td/><td/><td>3</td><td>3</td><td>3</td></tr><tr><td colspan="5">≤ [1 -2α t ξ + 2α t φ]E</td><td colspan="2">∆ t d</td><td colspan="3">2 2 + α t φE</td><td>∆ t i</td><td>2 2 + 2(α t ) 2</td><td>σ 2 d</td></tr><tr><td/><td/><td/><td/><td colspan="2">d=1</td><td/><td/><td/><td>d=1</td><td>i=d+1</td><td>d=1</td></tr><tr><td>Note that</td><td/><td/><td/><td/><td/><td/><td/><td/></tr><tr><td>3</td><td/><td/><td/><td/><td>3</td><td/><td/><td/><td>3</td><td>3</td><td>d-1</td></tr><tr><td>E</td><td colspan="2">∆ t+1 d</td><td colspan="2">2 2 -2α t φE</td><td colspan="3">∆ t+1 d</td><td colspan="2">2 2 ≤ E</td><td>∆ t+1 d</td><td>2 2 -α t φE</td><td>∆ t+1 i</td><td>2 2</td></tr><tr><td>d=1</td><td/><td/><td/><td colspan="2">d=1</td><td/><td/><td/><td>d=1</td><td>d=1</td><td>i=1</td></tr><tr><td/><td/><td/><td/><td/><td>3</td><td/><td/><td/><td>3</td><td>3</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td>∆ t d</td><td colspan="2">2 2 + α t φE</td><td>∆ t i</td><td>2 2 + (α t ) 2 σ 2</td></tr><tr><td/><td/><td/><td/><td/><td colspan="2">d=1</td><td/><td/><td>d=1</td><td>i=d+1</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td>3</td><td/><td>3</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td colspan="2">∆ t d</td><td>2 2 + 2α t φE</td><td>∆ t d</td><td>2 2 + (α t ) 2 σ 2</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td colspan="2">d=1</td><td>d=1</td></tr><tr><td colspan="10">Combining these two facts with our previous results yields:</td></tr><tr><td/><td/><td/><td/><td>3</td><td/><td/><td/><td/></tr><tr><td/><td colspan="4">[1 -2α t φ]E</td><td colspan="2">∆ t+1 d</td><td>2 2</td><td/></tr><tr><td/><td/><td/><td/><td>d=1</td><td/><td/><td/><td/></tr><tr><td/><td/><td/><td/><td/><td/><td/><td>3</td><td/><td>3</td></tr><tr><td/><td colspan="5">≤ [1 -2α t ξ + 2α t φ]E</td><td/><td colspan="2">∆ t d</td><td>2 2 + 2α t φE</td><td>∆ t d</td><td>2 2 + (α t ) 2 σ 2</td></tr><tr><td/><td/><td/><td/><td/><td/><td colspan="2">d=1</td><td/><td>d=1</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td>3</td><td/></tr><tr><td/><td colspan="5">= [1 -2α t ξ + 2α t φ]E</td><td/><td colspan="2">∆ t d</td><td>2 2 + (α t ) 2 σ 2</td></tr><tr><td/><td/><td/><td/><td/><td/><td colspan="2">d=1</td><td/></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>d-1</td><td>3</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>∆ t d</td><td>2 2 ] + 2α t φE</td><td>∆ t+1 i</td><td>2 +</td><td>∆ t i 2</td><td>∆ t d 2</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>i=1</td><td>i=d+1</td></tr><tr><td colspan="5">apply 2ab ≤ a 2 + b 2 :</td><td/><td/><td/><td/></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>d-1</td><td>3</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>∆ t+1 i</td><td>2 2 + ∆ t d</td><td>2 2</td><td>+ α t φE</td><td>∆ t i</td><td>2 2 + ∆ t d</td><td>2 2</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>i=1</td><td>i=d+1</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>d-1</td><td>3</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>∆ t+1 i</td><td>2 2 + α t φE</td><td>∆ t i</td><td>2 2</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td/><td>i=1</td><td>i=d+1</td></tr></table>

Peer review (iclr2020):
Summary: 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. Overall, the paper studies an interesting problem in RL. The idea is somewhat interesting. The experiment results also look promising. However, the writing sometimes has unclear descriptions, probably because there are many things packed in this paper. I have some following concerns about it. - The idea of using behavior embedding is new in PO. The description in section 3 is a bit unclear. There are lacks of motivation why in the paper there needs BES/BEM, the interplay of policies, trajectories, and embedding maps, and why non-surjective of BEM is still fine in this case, etc.. In addition, the definition of state-based, action-based, reward-based assume only discrete domains? Besides, they seem not to reappear in other places in the paper. - Does the choice of the functions lambda in 4.2 as a function in RKHS, especially when approximated as a linear function with random Gaussian features, limit the representation power of the embedding space? - What is the effect of _x0008_eta when positive vs. negative? - Experiments: The proposed methods show a lot of potentials, but the description in this section is sometimes unclear. Is TRPO with KL divergence the standard algorithm defined without using BEM or with BEM? Then one can wonder how more ablations can be added, e.g. BGPG with KL divergence, TRPO with WD distances are compared with the proposed algorithm. In addition, how BGPG is compared with other related work that also uses skill/policy embedding, instead of a flat PO approach like TRPO. Similar questions are also applied to the experiments for BGES. - It would also be more self-contained if the paper includes experiment settings. - The paper also comes with theoretical results. The author could also consider mentioning one key result in the main paper, instead, everything is put in the appendix. * Minor comments: - The KL definition on page 4 use rho instead of xi - are u and v in Eq. 3 functions in cal C(X) and cal C(Y), respectively?