Behavior-Guided Reinforcement Learning
Authors: Aldo Pacchiano
Jack Parker-Holder
Yunhao Tang
Anna Choromanska
Krzysztof Choromanski
Michael I. Jordan
Venue:  (2020)
Keywords: Reinforcement Learning
Optimal Transport
Evolution Strategies

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.

Under review as a conference paper at ICLR 2020
BEHAVIOR-GUIDED REINFORCEMENT LEARNING
Anonymous authors Paper under double-blind review
ABSTRACT
We introduce 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, 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 Guidedd Evolution Strategies, which we demonstrate can outperform existing methods in a variety of challenging environments. We also provide an open source demo !T
learn score functions
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 al. 2016 Conti al. 2018). The fundamental question we aim 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 BEM to 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 class of Wasserstein distances (WDs, Villani 2008)). There 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 the same embedding point (for example in the case of the last-state embedding). This precludes the use of likelihood-basee distances such as the KL divergence (Kullback Leibler 195Ip. which we discuss in Section 6
Behavioral Test Functions. Solving the dual formulation of the WD objective yields pair of test functions over the space 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-regularin formulations. This allows us to update the learned test functions in computationally efficient manner
Available at https:::s / thub. com hbehabeaa i orgu dedRI./Boe /BRRI. We emphasize this is not exact replica of the code from our experiments, but demo build intuition and clarify our methods.

---

Under review as a conference paper at ICLR 2020
via stochastic gradient descent (SGD) on 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-efficiene representations of learned behavioral test functions. Finally, having laid the groundwork for comparing trajectories via hehavior-driven trajectory scores, address our core question by introducing two new on-policy RI. algorithms:
Behavior Guided Policy Gradients (BGPG): We propose to replace the KL-bbaed trust region from Schulman et al. (2015) with WD-based trust region in the PPE space.
Behavior Guided Evolution Strategies (BGES): Inspired by the NSR-ES algorithm from Conti ct al. (2018), BGES jointly optimizes for reward and novelty using the WD in the PPE space.
In addition, we also demonstrate 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 SEEAVIOR-GUIDDRS 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 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 propose can be thought of 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 embedding map (behavioral summary) and a cost function. Crucially these choices occur in the semantic space of behaviors as opposed to parameters visitation frequencies?) One of the advantages of choosing 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 (5 we show that in order to get 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 metrid]] 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 (c.g. Imitation Learning). We develop versions of on policy RL algorithms which 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 expert's embedding. Although our framework also has obvious applications 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 significant enhancement.
II we choose an appropriate embedding map our framework handles visitation frequencies as well. 'The emhedding space can he discrete continuous and the metric need not be smooth, and can be for example simple discrete {0, 1} valued criterion
2

---

Under review as a conference paper at ICLR 2020
3 DEFINING BEHAVIOR IN REINFORCEMENT LEARNING
Markov Decision Process (MDP) is a tuple (S,A,PRR) states and actions respectively, such that for s, E S and E A: P(s'la, s) is the probability that the system/agent transitions from s to s given action a and R(s', a, s) is reward obtained by an agent transitioning from to s' via A policy T S > A is (possibly randomized) mapping (parameterized by 0 E Id) from S to A. Let r = {T = sg. ag. Tu.... SH. dH, CH s.t. 8; E S a, E A, E RJ be the set of possible trajectories enriched by sequences of partial rewards under some policy и. The undiscounted reward function R. I + I& (which expectation is to be maximized by optimizing 0) satisfies R(T) = Ei= Ti, where r; = R($++1 8z)
R). Here S and stand for the sets of
3.1 BEHAVIORAL EMBEDDINGS
We start with a Behavioral Embeddng Space (BES) which denote as & and Behavioral Embedding Map (BEM), Φ : - &, mapping trajectories to embeddings in & (Fig. Tb. Importantly, the mapping does not need to be surjective. We will provide examples of BESs and BEMs at the end of the section. Given policy и, we let I denote the distribution induced over the spaces of trajectories I' and by F the corresponding pushforward distribution on & induced by 4. We call Oe the Probabilistic Policy Embedding (PPE) of a policy æ. A policy иt can be fully characterized by the distribution P..
1
Figure Behavioral Embedding Maps (BEMs) map trajectories to points in the behavior embedding space E. Two trajec- tories may map the same point in E.
Additionally, we require the BES & to be equipped with a metric (or cost function) C & x & + R. Given two trajectories Ti, 72 in I, C(Φ(7)), 2(T2)) measures how different these trajectories are 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 Section
1. State-based: the final state & (T) 8H. the visiting frequency of fixed state 25(TT Ei (s = s), the frequency vector of visited states (T((r) Σ/
Σi=0 Cs, (where e, E R|S)
is the one-hot vector corresponding to state s); see also Section 7.2
2. Action-based: the concatenation of actions (T) = [Aa,... (H) see also Section 7.1
3. Reward-based: the total reward (T)
Σ rre reward-to-go vector 46(7)
Σ Tt (Eizo e,) (where E E "++ is one-hot vector corresponding i and with dimensions indexed from 0 to i/); see also Section 7.1 and Section 7.3
For instance, putss 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| < oo) while others are universal.
4 WASSERSTEIN DISTANCE & OPTIMAL TRANSPORT PROBLEM
Let μ4, v be (Radon) probability measures over domains C I"" , C R" and let C: : be a cost function. For 7 > 0, a smoothed Wasserstein Distance is defined as: WD. v) := min
x
x
+ R
sTII((H.) |xxy C(xxy)dT(x, y) + +> (KL(T)))
(I)
where [[ v) is the space of couplings (joint distributions) over X' x y with marginal distributions and v, KLL(II)) denotes the KL divergence between distributions and p with support X' x ) defined
as: (L(TTp) Jxxy (Log
(x,y)) dx y) and & is reference measure over X' V. When de
da
the cost is an l, distance and 7 = 0, WDy 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, first need to consider the dual form of Equation T Under the subspace topology (Bourbaki
3

---

Under review as a conference paper at ICLR 2020
Algorithm 1 Random Features Wasserstein SGD
Input: kernels к, é over X, respectively with corresponding random feature maps @xs Φe, smooth- ing parameter ",, gradient step size number of optimization rounds M, initial dual vectors Po Pö- for 0, .... do 1. Sample (Ia, Ut) ~ 11 X V.
Update (b:) = (P
p + vi
1-exp
(pf (1'))(())-( (p) 3) (x()) -((ss))(
))
( (ee((ve)
Return: PY.-
1966) for X' and V' let C(X) denote the space of continuous functions N' and let C()) denote the space of continuous functions over V. The choice of the subspace topology ensures our discussion encompasses the discrete case.
Let C:: X' x y - IR be cost function, interpreted as the "ground cost" to move unit of mass from y. Define I as the (0, 20) indicator function, where the value 0 denotes set membership. Using Fenchel duality, we can obtain the following dual formulation of the problem in Eq. |
[ {x(x)du(x) - I,. Au())duly) Ec Ec(Au,do): WD, E
max
(2)
where Ec(Aps Av) defined as:
fxxx exp ( (Nals)-Ao(o)-Ctan)) dE(x,
if 0 I((Au, )u) E ((,, s.t. V(x,y) E N' x y u(x) - vly) & c(x,))))
Ec(Ap;Au) An)
(3)
if
If Ai A5 are the functions achieving the maximum in Eq. 2 and is sufficiently small then Ex We will set dC(x, y) x for discrete domains and dE(x,y) = du(x))dv(y) otherwise.
X = = and 2(x, [x)( E, (y)), with equality when 7 0. When for example 7
(D,
for all x X, it is easy to see (xx (² (x)x) for all æ E X. In this case the difference between F. [x*(x) and F,,, [X*(y) equals the WD. In other words, the function X gives higher scores to regions of the space X' where has more mass. This observation
key to the success of our algorithms in guiding optimization towards desired behaviors.
4.2 COMPUTING X AND X
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 the of their ambient spaces. In this paper choose the Gaussian kernel and approximate using random Fourier feature maps (Rahimi & Recht 2008) to increase efficiency. Consequently, the functions À learned by our algorithms have the following form: A(x) (p^) d(x), where is random feature map with m standing for the number of random features and p^ E R.M For the Gaussian kernel, is defined as follows: (0z)
continuous functions
Figure 2: Behavioral embedding func- tions corresponding to two policies ı (green) and na (blue) whose BEMs map trajectories to points in the real line.
cos(Gz b) for Rd. vm
1
where G Rd is Gaussian with iid entries taken from (0, 1), b E R with iid b,s such that b, ~ Onif(0, 2T] and the cos function acts elementwise.
Henceforth, when we refer to optimization over X, we mean optimizing over corresponding dual vectors p² associated with A. We can solve for the optimal dual functions by performing SGD over the dual objective in Eq. 2 Algorithm| is the random features equivalent of Algorithm 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.
4

---

Under review as a conference paper at ICLR 2020
p" pK are the optimal dual vectors and (z1,)1), ... (ax: yk) iid a p& v. then Algorithm can be used to get estimator of WD,, (H, v) follows:
1 k
1 exp ( pº Φe(vr) p. Bi) (4)
(D,(H,))
pp. (x.)) (p", (x(x)) +
5 BEHAVIOR-GUII REINFORCEMENT LEARNING
Here we introduce the framework which allows us to incorporate our behavioral approach to rein- forcement learning into practical on-policy algorithms. Denote by my a policy parameterized by 0 E IRd, The goal of policy optimization algorithms is to find policy maximizing, as function of the policy parameters, the expected total reward (()) == E,wP= [R(r]].
5.1 BEHAVIORAL. TEST FUNCTIONS
IFC 8.C:ExE R a cost function defined over behavior space E, and Iτ, #2 are two policies, then:
T WD, mj Ro E,-P,, [AT((())) (
ET~7,, (T
(5)
where \i, \i are the optimal dual functions. The maps 81 := if o R and 82 X2 o T -> I. define score functions over the space of trajectories. If is close to zero, the score function 8; gives higher scores to trajectories from π, whose behavioral embedding common under m, but rarely appears under 5j for j (Fig. 2
5.2 ALGORITIMMS
We propose to solve a WD-regularized objective to tackle behavior behavior-guided policy optimization. All of our algorithms hinge on trying maximize objective of the form:
F(0)= = (()) + 3WD, (ID, PT)
(6)
where F is base distribution over behavioral embeddings (possibly dependent 0) and B E R
could be positive negative. Although the base distribution 1n could be arbitrary, algorithms
will instantiate p [S] Ures I for some family of policies S (possibly satisfying |S] = 1) we
want the optimization attract 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. in Eq.
2
obtain:
F(0) x Br-Pan (R(T) + (s (r) SE [X(())]
where 81 r R equals o the Behavioral Test Function of policy my and X the optimal dual function of embedding distribution P.. Consequently VgF(0) ~
VgEr~PT, [R(r) + Bs (T)].
We learn a score function $1 over trajectories that can guide our optimization by favoring those trajectories that show desired global behaviors.
Eq. is an approximation to the true objective from Eq. whenever 7 > O. In practice, the entropy regularization requires damping term defined in Equation If ( IP is the joint distribution of choice then F()) L()) 3V for
max Xey (CCE), F~T
E
DA (e((T)))
E drue
[b( +
E Φi ElPe p.)
[A A2)1,
Ang( - C(()))
where A(01, 02) exp
When the embedding space & is not discrete
7
and p = p& for some policy n, we let E(PP No. pg) uniform distribution over x E.
p& do
otherwise ((P))) not pPg.) |テ 1,
All of our methods perform version of alternating SGD optimization: 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
5

---

Under review as a conference paper at ICLR 2020
that make up the optimization objectives 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 I0 in the Appendix.
We consider two distinct approaches to optimizing this objective, by exploring in the action space and backpropagating, in policy gradient methods (Schulman et al. 2013 (017), and by considering black-box optimization problem as in Evolution Strategies (ES, Salimans et al. (2017)). These two different approaches lead to two new algorithms: Behavior- Policy Gradient (BGPG) and Behavior-Guides 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 P below).
Algorithm 2 Sehvaior-Guided Policy Gradient
Input: Initialize stochastic policy To parametrized by 0. B < 0, 1,, M, for ( = 1,..., T do
Run - in the environment to get advantage values (* a) and trajectories (r. (M
2. Update policy and test functions via several alternating gradient steps over the objective:
To(a:(s;)
F(0)
Σ iml
T 1S8, a4)
3A1 (o(T)))
X Pro
T TT
>2 (2(0))))))((()))))-))))))))))))) - -((())
-322(0(T2)) Byexp
Where TI = 50, do, ro: 8H, CH,aH,TH TH. Let Let = 0:-11
for 1,...
Approximate P Pr.., P No via 1 (
1 1==1
[T = P where of iid
P.T
b. Take SGA step 0²
using samples from PRI-1 c. Use samples from P P and Algorithm to update \1, A2-
1-1 +
(e-1)
Set 0: = (M)
Specifically, we maintain a stochastic policy Ty 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 mundom-feature Wasserstein SGD Importantly, this stochastic gradient can be approximated by samples collected from the policy ag- In its simplest form, the V Step 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. lower- variance gradient estimator alternative.
BGPG can be thought of as a variant of Trust Region Policy Optimization with a Wasserstein penalty. As opposed vanilla TRPO, the optimization path of BGPG flows through policy parameter space while encouraging it to follow smooth trajectory through the geometry of the PPE space. We proceed to show that given the right embedding and cost function, we can prove monotonic improvement theorem for BGPG, showing that our methods satisfy at least similar guarantees as TRPO.
For given policy म, we denote as: V Q" and AT a) (s a) (s) the: value function, Q-function and advantage function (see Appendix: Section 10.5). Furthermore, let V(T) be the
expected reward of policy x and Pm = Er<P_ Σ 1(s s) be the visitation measure.
Two distinct policies x and can be related via the equation (see: Sutton al. (1998) V(ã) V(=)+ (=(s)) (aa A (s, a)da) ds and the linear approximations to V around n via: I(R))

---

Under review as a conference paper at ICLR 2020
V(a)+ Ss (J4 (sals)A" d)daa ds (see: Kakade & Langford (2002)). Let S be finite
2t=0 1 ( = 8) and related cost function defined as: Cvv, w) = li w Then WDo(PF P is related to visitation
set. Consider the following embedding t - R defined by ($(7)
frequencies since WDo PPT 2 ELEs lpx (s) (s)] (see Section 10.5 for the proof). These
observations enable us prove an analogue of Theorem from Schulman et (2015). namely:
Theorem 5.1.
(IP p S and maxs,a (s, a)), then V(+) L òc.
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 RI., by considering rollout of a policy, parameterized by 0 as black-box function F. This approach has gained in popularity recently (Salimans et al. [017] 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 simultancously maximizes or minimizes the difference behavior from the base embedding distribution P. We call this method the Shhavio-GGuider Evolution Strategies (BGES) algorithm (see Algorithm below).
Algorithm 3 Rehavior-Guidee Evolution Strategies
Input: learning rate noise standard deviation a, iterations T, BEM 4 B Initialize: Initial policy mo parametrized by f. Behavioral Test Functions > A2- Evaluate policy To to return trajectory Tu and subsequently use the BEM to produce an initial PPE p
for / 1,..., 1 do
1. Sample 61, en independently from (0. I). Evaluate policies [m[ parameterized by {0: + Fex]R=1 return rewards Rg and trajecto-
ries тı for all k.
Use BEM to map trajectories тą to produce empirical PPEs for for all k.
4. Update X and A2 using Algorithm where μ Uh for - and 1 U 1 are the uniform distribution over the set of PPEs from 3 for t and Approximate WD> (1ot 1P plugging in A2 into Eq. for each perturbed policy Tx a.
Update Policy: 0c+1 = 0, + + V ES F, where:
ES F
1
Σ
(1
B) R.) + HWD., att 1 1e
When 3 > 0, and we take _ = P P BGES resembles the NSR-ES algorithm from Conti et al. (2018), an instantiation 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, ie.. 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, A2 for each perturbed policy Tg.
If take < O, and assume P Fe to correspond to embedded trajectories from oracle or expert policy, we can perform imitation learning. Despite not accessing the expert's ''s policy (just the trajectories 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 intro- duced by Lehman & Stanley (2008) and Lehman (0122), who proposed to search directly for novelty,
7

---

Under review as a conference paper at ICLR 2020
rather than simply assuming it would naturally arise in the process of optimizing 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 Fuclidean 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 (Tiang 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 KL divergences, it is possible to take SGD steps using entropy regulay-ee Wasserstein objectives. Computing an estimator of the KI. divergence is hard without 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.
We are not the first to propose using WDs in RL. Zhang et al. (0118) 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-orde WDs. We note that computing WGFs is a nontrivial task. In Zhang 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.22). 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.
WDs for RL:
Distributional RL Distributional RL (DRL. Bellemare et al. (017)) expands on traditional off- policy methods (Mnih et al. 2013) by attempting to learn distribution of the return from given state, rather than just the expected value. These approaches have impressive experimental results (Bellemare et al. 2017 [Dannyy et al. 2018), with a growing body of theory (Rowland al. 2018 Qu et al. 2019] Bellemare et al. 2019 Rowland et al. 2019). Superficially it may seem that learning 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- off-policy whereas our BGES and BGPG algorithms are on-policy, and 2) DRL is typically designed for discrete domains, since Q-I-earnin 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 DRI. algorithm- algorithm itself does not use WDs (Bellemare et al. 2017).
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 basclines 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 provide additional details in the Appendix.
7.1 BEHAVIOR--UIDEE 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 OpenAl Gym and Hopper: Stand, Hooper: Hop and Walker: Stand from the DeepMind Control Suite (Tassal et al. 2018). For the BEM, we use the concatenation-o--actioos (as used already in TRPO). We also
8

---

Under review as a conference paper at ICLR 2020
confirm results from (Schulman et al. 2015) that trust region greatly improves performance, see the black curve (without one) often fails to learn.
30
100
KL
Reward
-500
None
400
1000
10
200
-5500
0
0.5
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Timestess
110
Timesteps X10"
Timestops
100
Timestops x10
Pendulum
Hopper: Stand
(c) Hopper: Hop
(d) Walker: Stand
Figure 3: BGPG vs. TRPO: We compare BGPG and TRPO (KI. divergence) on several continuous control
tasks. As baseline also include results without trust region (B in Algorithm 2}. Plots show the mean std across random seeds. BGPG consistently outperforms other methods.
To illustrate computational benefits of alternating optimization (AO) of in BGPG, 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.
Wall Clock Time:
Normalized Reward
1.0
1.0
1.0
1.0
0.5
0.5
0.5
Partiele
BGPG
0.0
2000
4000
20000 Clock Time
20000
20000 Clock Time
Clock Time
Clock Time
Pendulum
(b) Hopper: Stand
(c) Hopper: Hop
(d) Walker: Stand
Figure 4: The clock-time comparison (in sec) of BGPG (alternating optimization) with particle approximation.
7.2
SIAAVIOR-GUIDEE EVOLUTION STRATEGIES
As 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 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 goal. The deception comes from wall situated in the middle, which means that initially positive rewards from moving directly forward will lead to suboptimal policy.
600
BGES
Reward
-5000
NSR-ES ES
Vanilla-ES
Noisy ToP
6000
1000
Timestops x10"
Timestepe 10*
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.

---

Under review as a conference paper at ICLR 2020
We consider two types of agents a two-dimensional point and much larger quadruped. Details are provided in the Appendix (Section[9). We compare with state-of-the-an on on-policy methods for efficient exploration: NSR-ES from (Conti et al. 2018). which assumes the BEM is deterministic and uses the Fuclidean distance to compare policies, and NoisyNet- Net-TRPO from Fortunato et al. (2018). Results are presented 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-ge 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, -ansen-Shannon (JS), KL and Total Variation (TV) distances), as well as vanilla ES (i.e., with no distance regularizer). Experiments were performed on Swimmer environment from OpenAl 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.
400
Heelinger
Reward
2011
KL TV
Vanil ba
BGES
2011
Timesstpps 107
Figure 6: Escaping Local Maxima. A comparison of BGES with those using different distances PPEs.
7.3 IMITATION LEARNING
As discussed in Section 5.3 we can also utilize the BGES algorithm for imitation learning. by setting 3 < O, and using an expert's tra- jectories for the PPE. For this experiment we use the reward-to-go to
600
Vanilla Es
UGES
Rewarrd
a00
BEM 5S 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.
200
Timnodap x10"
Figure 7: Imitation Learning.
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. g where we see that both the reward-to-ge (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.
Vanilla-ES
-000
BGES (RTC)
Reward
BGFS (8V)
-800
BGES (SF)
1000
Timesteps
Figure 8: Choice of BEM
8 CONCLUSION AND FUTURE WORK
In this paper we proposed a new paradigm for on-policy learning in RL, where policies are em- bedded 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 Wasser- stein distances (WDs) guarantees flexibility in choosing cost funtions between embedded policy trajectories, enables stochastic gradient steps through corresponding regularized objectives (as op- posed to KI 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 population setting. learning the behavioral embedding maps from data, or adapting the degree of repulsion/attraction during optimization (parameter B).
10

---

Under review as a conference paper at ICLR 2020
REFERENCES
Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. In Doina Precup and Yee Whye Tch (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 214-223, International Convention Centre, Sydney, Australia, 06-11 Aug 2017. PMLR.
Sivaraman Balakrishnan, Martin J. Wainwright, and Bin Yu. Statistical guarantees for the em algorithm: From population to sample-based analysis. Ann. Statist. 45(1):77-120, 02 2017. doi: 10. N0.1214/16-AOS111 1335. URL https: //tpi.:///do..e/ www./oi...grrooooooooo worg/100.... 101..1111116-A0000
Mare G. Bellemare. Will Dabney, and Rémi Munos. distributional perspectivee on reinforcement learning. In Proceedings of the 34th International Conference on Machine Learning Volume 70, ICML pp. 449 458. JMI.R.rrg 2017. Utt.tttp:///// http::t/ttll..... http://aa.aacc.oocciiiccaaaaaa
cfm
id=330538 3305428
Marc G. Bellemare, Nicolas Le Roux, Pablo Samuel Castro, and Subhodeep Moitra. Distributional reinforcement learning with linear function approximation. In Kamalika Chaudhuri and Masashi Sugiyama (eds.), Proceedings of Machine Learning Research, volume 89 of Proceedings of Machine Learning Research, pp. 2203-2211. PMLR, 16-18 Apr 2019.
Nicolas Bourbaki. General Topology: Elements of Mathematics. Addison- Wesley, 1966.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint wrXiv::606.01555 2016.
Krzysztof Choromanski, Aldo Pacchiano, Jack Parker-Holdee Yunhao Tang, Deepali Jain, Yuxiang Yang. Atil Iscen, Jasmine Hsu, and Vikas Sindhwani. Provably robust blackbox optimization for reinforcement learning. CoRR, abs/1903.00293, 2019.
Edoardo Conti, Vashisht Madhavan, Felipe Petroski Such, Joel Lehman, Kenneth O. Stanley, and Jeff Clune. Improving exploration in evolution strategies for deep reinforcement learning via a population of novelty-seeking agents. In Advances in Neural Information Processing Systems 31, NeurlPS, pp. 5032-5043, 2018
Will Dabney, Georg Ostrovski, David Silver, and Remi Munos. Implicit quantile networks for distributional reinforcement learning. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1096-1105, Stockholmsmssan. Stockholm Sweden, 10-15 Jul 2018. PMLR.
Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, Yuhuai Wu, and Peter Zhokhov. Openai baselines. https
//github. cww
opena www/opea/aaaoooo//o/////e/eoeaat /bbasel nes 2017.
Meire Fortunato, Mohammad Gheshlaghi Azar, Bilal Piot, Jacob Menick, Matteo Hessel, Ian Osband, Alex Graves, Volodymyr Mnih, Remi Munos, Demis Hassabis, Olivier Pietquin, Charles Blundell, and Shane Legg. Noisy networks for exploration. In International Conference Learning Representations, ICLR, 2018.
Alexander Gajewski, Jeff Clune, Kenneth O. Stanley, and Joel Lehman. Evolvability es: Scalable and direct optimization of evolvability. In Proceedings of the Genetic and Evolutionary Computation Conference, GECCO '19, pp. 107-115, New York, NY, USA, 2019. ACM. ISBN 978-1- 6111-8. doi: 10. 19.114///21707732288866
Aude Genevay, Marco Cuturi, Gabriel Peyré, and Francis Bach. Stochastic optimization for large- scale optimal transport. In Advances in neural information processing systems, pp. 2016.
Ray Jiang, Aldo Pacchiano, Tom Stepleton, Heinrich Jiang, and Silvia Chiappa. Wasserstein fair classification. arXiv preprint 190. 2019.
Sham Kakade and John Langford. Approximately optimal approximate reinforcement learning. 267-274, 2002.
11

---

Under review as a conference paper at ICLR 2020
Diederik P Kingma and Max Welling. arXiv:1312.61 14, 2013.
Auto-encoding variational bayes.
arXiv preprint
S. Kullback and R. A. Leibler. On information and sufficiency. Ann. Math. Statist., 22(1):79- 86, 03 1951. doi:
aoms
1177729694
Seong Jae Lee and Zoran Popovic. Learning behavior styles with inverse reinforcement learning. ACM Trans. Graph., 29(4); 12:::222:2112211 2010. doi: 10. 10.1145/7778765.1778885 URL https
Joel Lehman. Evolution through the Search for Novelty. PhD thesis, University of Central Florida, 2012.
Joel Lehman and Kenneth O. Stanley, Exploiting open-endedness to solve problems through the search for novelty. In Proceedings of the Eleventh International Conference on Artificial Life (Alife XI. MIT Press, 2008
Horia Mania, Aurelia Guy, and Benjamin Recht. Simple random search of static linear policies is competitive for reinforcement learning. Advances in Neural Information Processing Systems 31, NeurlPS, 2018.
Elliot Meyerson, Joel Lehman, and Risto Miikkulainen. Learning behavior characterizations for novelty search. In Proceedings of the 2016 on Genetic and Evolutionary Computation Conference, Denver; CO, USA, July 20 24, 2016, pp. 159-156, 2016. doi: 10.11452908881229089999 URL
https: //dps:
doi. org. /ttpg//c... /1t////..1111/200002020200000000 2908999
Charles A. Micchelli, Yuesheng Xu, and Haizhang Zhang. Universal kemelss. Journal of Ma-
chine Learning Research, 7:2651-2661 2006. URI. http::/
http::/jmlr /jml wwt.....rrrppppr///....
micchel1106a.atmm
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller. Playing atari with deep reinforcement learning. ArXiv, abs/1312.5602,, 2013.
Yurii Nesterov, Introductory Lectures on Convex Optimization: Basic Course. Springer Publishing Company, Incorporated, edition, 2014. ISBN 1461346916, 9781 461346913.
Chao Qu, Shie Mannor, and Huan Xu. Nonlinear distributional gradient cemporal-differees learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), Proceedings of the 36th International Conference Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 5251-5260, Long Beach, Califormia, USA, 09-15 Jun 2019. PMLR.
Ali Rahimi and Benjamin Recht. Random features for large-scale kernel machines. In J. C. Platt, D. Koller, Singer, and S. Roweis (eds.), Advances in Neural Information Processing Systems 20, pp. 1177-1 184. Curran Associates, Inc., 2008. URL httpttt:::://eeeeee http://aaprs.ss/ss www.aaprrssiiiss........p.ss..se nips cc/paper/ scale 3182-zandom- features-eeeeeeeeeeeeee for httpeeeeeaeeeeeeeeeaaaaeeeeeeeee pdf
Pierre H. Richemond and Brendan Maginnis. On wasserstein reinforcement learning and the fokker- planck equation. CoRR, abs/1712.07185, 2017. URL http://arxiv.org/
1712
07185
Mark Rowland, Marc Bellemare, Will Dabney, Remi Munos, and Yee Whye Teh. An analysis of categorical distributional reinforcement learning. In Amos Storkey and Fernando Perez-Cruz (eds.), Proceedings of the Twenty-Firg International Conference on Artificial Intelligence and Statistics, volume 84 of Proceedings of Machine Learning Research, pp. 29 29-37, Playa Blanca, Lanzarote, Canary Islands, 09-11 11 Apr 2018. PMLR.
Mark Rowland, Robert Dadashi, Saurabh Kumar, Remi Munos, Mare G. Bellemare, and Will Dabney, Statistics and samples in distributional reinforcement learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.). Proceedings of the 36th International Conference on Machine Learning, volume 97 of Proceedings of Machine Learning Research, pp. 2228-5536,11 Long Beach, California, USA, 09-15 Jun 2019. PMLR.
12

---

Under review as a conference paper at ICLR 2020
Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever. Evolution strategies scalable alternative to reinforcement learning. CoRR, abs/1703.03864, 2017.
John Schulman, Sergey Levine, Pieter Abbeel, Michael I. Jordan, and Philipp Moritz. Trust region policy optimization. In Proceedings of the 32nd International Conference on Machine Learning, ICML 2015, Lille, France, 6-11 July 2015, pp. 1889-1897, 2015. URL http: /
proceedinng//pppptt///////hhhhhp
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347 2017.
Richard S Sutton, Andrew G Barto, al. Introduction Reinforcement Learning, volume 135, MIT press Cambridge, 1998.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite, arXiv preprint 0rXiv:1801.00690 2018.
Cedric Villani. Optimal transport: Old and new. Springer, 2008.
Ruiyi Zhang, Changyou Chen, Chunyuan Li, and Lawrence Carin. Policy optimization wasserstein gradient flows. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässaan Stockholm, Sweden, July 10-15 15, 2018, pp. 5741-5750, 2018. URL ngs ml r. press/ vv00 8a w00zz//ng/////ggggg zhang html
13

---

0.1; the Under review as a conference paper at ICLR 2020
APPENDIX: BEHAVIOR-GUIDED REINFORCEMENT LEARNING
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 loss function that involves the Behavior Test Functions. Though we could apply vanilla reinforced gradients on both parts, it 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. particular, we note that under random feature approximation
(4), well the action-concater conccttcnnattnnioonn embedding, the Wasserstein distance loss (PP, x P% is differentiable function of 0. To see this more clearly, notice that under Gaussian policy a (s)
the actions (e((s) + ao(s) c are reparametrizable for being standard Gaussian noises. We can directly apply the reparametrization trick to this second objective to obtain gradient estimator with potentially much lower variance. 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 dual formulation of the trust region method and aim to optimize the augmented objective Er- [R(T]] 3WD,, ( ). We apply the concatenation-offaelee 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(1læ: 100; recall the embedding is
1
/ with a 0.1: the number of random features ((T) [d1, [2 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
trust region trade-off constant E {0.1, 1, 10}. The alternate gradient descent carried out with T 100 alternating steps and test function coefficients p RDD are updated with learning rate Op 0.01.
The baseline algorithms 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 e = 0.01.
Across all algorithms, adopt the open source implementation (Dhariwal al. 2017). Hyper- parameters such as number of time steps per update as well as implementation techniques such as state normalization are default the original code base.
The additional experiment results can be found in Figure continuous control benchmarks: Tasks with DM are from DeepMind Contol Suites (Tassa 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).
9 where we show comparison on additional
Cummlttivee
Tinestaps
12
Timestope
1M
Timesteps
(a) Reacher
(b) MountainCar
Acrobot
Figure 9: Additional Experiment on TRPO. We compare No Trust Region with two alternative trust region constructions: Kl-divergence and Wassertein distance (ours).
14

---

Under review as a conference paper at ICLR 2020
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. Alternatielly;
can also adopt 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, -ffatiionn
when using the concatenation-of-uee acttionn 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 3 E {0.1, 1, 10). We adopt the particle method in (Zhang et al. 2018) where for each state there are M = 16 particles. derived based RKHS where 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 approximate the regularized Wasserstein distance.
The gradients
9.2 BGES
Efficient Exploration: To demonstrate the effectiveness of method in exploring deceptive environments, we constructed two new environments using the MuJoCo simulator. For the point environment, we have dimensional state and 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 reach the goal. The quadruped environment is based on Ant from the Open Al 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 and 100 warm-start updates for the WD at each iteration. For point, we use the final state embedding. learning rate 7 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 learning rate of r) = 0.02 and a 0.02. The hyper- parrmeaarsestrss were the same for all ES algorithms. When computing the WD, we used the previous 2 policies, 8,-1 and 0,-2. An ablation study for the point environment for both the choice of embedding and number of prior policies is shown in Fig 12
1000,
.500
100
W: 2531 (5V)
Wi:055 (ST)
-700
M011
-80
1000
S00000
3000000
1500000
3000000
-1002
2000000
ac0C00 600000 800000 1000000
Timostops
TTmostope
(a) Embeddings
(b) Previous Policies
Figure 10: A sensitivity analysis investigating the impact of the embedding and b) the number of previous policies oc-i,i e 1,2,5
15

---

Under review as a conference paper at ICLR 2020
For embeddings, compare the reward-to-ge (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, and 5 do learn the correct behavior.
the SF embedding, and using 2 appears to work best, but both 1
Escaping Local Maxima: We also demonstrated that method leads to faster training even in more standard settings, where exploration not that crucial, but the optimization can be trapped in local maxima. To show 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, in the previous paragraph. We noticed that adding Wasserstein regularizers drastically improved optimization, whereas regularizers based on other distances/divergenciiss namely: Hellinger, Ponsen-Shannon KL and TVV did not have any impact. We considered Swimmer task from OpenAT Gym and to make it challenging, reduced the number perturbations per iteration 80. In that setting our method was the only one that was not trapped in local maxima and managed learn effecttiee 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 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- Repulsion Learning
Input: 8, > 0, M N Initialize: Initial stochastic policies лß. my. parametrized by 88 ob respectively, Behavioral Test Functions xÇ
for f = 1,..., do
Collect M trajectories {rā M Ji=1 from Pay and M trajectories ( from Pap..
Approximate F.,.
P bb via 1 [T 8 ( := P
M
2. Form two distinct surrogate rewards for joint trajectories of agents a and b:
- A((0r)))-MMua))
= 3 exp
((())))))))))))))))))
Rb(TI, R(12) 3A (m))) By exp
Forc E {a, b) use the Reinforce estimator to take gradient steps:
#-1
V log (Tf ( ]
0f 0 +1
E 7 b 1
Re(Ta,
Σ
Where -= . ag, r0, ...
a a r and bb
8 4g
Use samples from P mb and Algorithm to update the Behavioral Test Functions X1, Ab.
Although this was not discussed in the main section of the paper, is possible to use our behavioral approach to simultaneously lear multiple policies exhibiting different behaviors all of which able to solve the same task. This not the main focus of the main paper, but we chose to include these results in attempt provide the readers with better understanding of Behavioral Test functions.
16

---

Under review as a conference paper at ICLR 2020
Algorithm 4 maintains two policies n" and nh. Each policy is optimized by taking policy gradient step (using the Reinforce gradient estimator) in the direction optimizing surrogate rewards R. and Rs that combines the signal from the task's reward function R. and the repulsion score encoded by the behavioral test functions A and bb.
We conducted experiments testing Algorithm 4 on simple Mujoco environment consisting of particle that moves on the plane and whose objective is to Iearn a policy that allows it to reach one of two goals. Each policy outputs velocity vector and stochasticity is achieved by adding Gaussian noise to the mean velocity encoded by neural network with two size 5 hidden layers and ReLu activations. If agent performs action state it moves to state + The reward of an agent after performing action a at state equals -I|a|| * 30 min(d(s Goal,) ) d(s, Goalz))² where d(x, y) denotes the distance between x and y in R The initial state chosen by sampling Gaussian distribution with mean (") 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 & T IR maps trajectories to their mean displacement in the x axis. We use the squared absolute value difference the cost function.
Edwoss Se Krdee
(a) n6
(b) b
(c) X and ー at =
Figure 11: Initial state of policies n", ab and Behavioral Test functions X", bb the Multigoal environment.
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.
(a) n22
(b) b
(c) x and > at 22
(d) TM
(e)
(9) A and b at 118
Figure 12: Evolution of the policies and Behavioral Test Functions throughout optimization.
17

---

Under review as a conference paper at ICLR 2020
Let X' and 3/ be the domains of two measures μ., and Recall that in ? 0, X' y and C(r, = for all z E X. then (x) = (²())= = x for all I E X. In the case 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 p, / have disjoint supports, since the soft constraint eex (x())))A()))Cx.)))
is enforced in
expectation over the product measure there may exist optimal solutions X > that do not satisfy Xr X
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 (T) = Σe Cse,a with Cs.a E RISH++A the indicator vector for the state action pair (8, a). Let C(v, w) = 1|v -wl 1. y = 0 and WD, (p)) = 0 then π = T².
Proof. If WD, (
O, there exists coupling II between p and p& such that:
Eu [llu 2
Consequently;
Bu
Uusa - Vosal
EujenIl [1u,,0 - Us
(s.a)ESXA
(s.0)ESXA
Therefore for all (s, a) E S x A:
P
Buwpy [Me,a] - B [Vs,a] I E Eul [|us, Uo.a -2,,al' ===== =
Where Us a and Us.a denote the (s, a) entries of u and respectively. Notice that for all (s, E Sx.::
P a) p (s. a)
(8)
Since for all S and 1:
Σ
Us.a
U0.0
Us,a
Uosal" P
aEA
aeA
Therefore for all s E S:
E Purt
Σ Usaa
E
Eyaps
Σ Vs.a
Σ 1u
Us.el" PP
0
Consequently F (s) = poth
(s) for all s E S. By Bayes rule, this plus equation
g
yields:
(als) = p (als)
And therefore: श = T
These results can be extended in the following ways:
1. In the case of 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 Φ in Section 3.1 or &(7)
Σ# Csr,ar> when y 0, if
WD,, IP) then [V(T) V(r')) cR for R max R(T) thus implying that small Wasserstein distance between π and T's PPEs implies a small difference in their value functions.
18

---

Under review as a conference paper at ICLR 2020
10.1 RANDOM FEATURES STOCHASTIC GRADIENTS
Let Øx and Φe be two feature maps over X' and } and corresponding to kernels and respectively. For this and the following sections we will make use of the following expression:
G(p",p') f (p")² @x(x)dp(xx()) -8 //1, (p") dely)duly)+
(9)
73 / cxp (pa) Φx(x) - (p')) or(y) C(x,y) du(x))d(y))
xxy
We now show how compute gradients with respect to the random feature maps: Lemma 10.2. The gradient V (I.)) G(p", ) of the objective function from Equation 9 with respect to the parameters (p=)) satisfies:
V CE:)C(p" p²) 3B Br -epp ((D") 6x(xx) (P")) -(xx,)) )
)) (
(xxxx)
)
7
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 the stochastic gradients we use in Algorithm
10.2 BGPG, BGES AND THEIR THEORETICAL GUARANTEES
Here we provide some thcoretical 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 3 0 the algorithms propose for solving the optimization objective in Equation turn into an Alternating Min-Max optimization procedure. Henceforth whenever refer to the problem defined by setting 3 < 0 will call the Min-Max problem.
2. Max-Max optimization. When 3 0 the algorithms we propose for solving the optimization objective in Equation turn into an alternating Max-Max procedure. Henceforth whenever we refer the problem defined by setting
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(O,p"
,p')== "'pp",,''"
Br,TERP &p.r R(TT) (p (x()))))) B(p (((2))) C(((T)) ((2)))
(10)
(p To (P(,))(((T)) - (0t o ( (o"
BB exp
Where No is policy parametrized by 0, n' is target policy, and Th and pa are the feature vectors corresponding to the Behavioral Test functions from Tg and n².
10.3 MAX-MAX PROBLEM: THEORETICAL ANALYSIS
We will analyze our AO algorithm for the Max-Max optimization problem. We show that obtained so- lutions converge to the local maxima of the objective function. Consider the function F(l,, p A Da.))
Aro in 10 We denote by (6* P. A=o _.'') some of its local maxima. Define F(0) F(0, p. i.e. F is F function of for locally optimal values of Awy and p
p. 2r)).
19

---

Under review as a conference paper at ICLR 2020
We will assume that F is locally {-strongly concave and 8-smooth for some fixed Ç. >> 0 in the neighborhood (0", r)) of its optimal value 0". We will also assume that gradicnt of F is Lg Lz-Liischitz with Lipschitz coefficient & in that neighborhood. The following convergence theorem holds:
Our goal in this section is prove the following Theorem:
4 Theorem 10.3. For the entropy coefficient 7, denole: Φ4 = 27 zee, and My = Sea. Denote b+C; un)) Let s(t) = (f(t), (#T (t), (t)) be the solution from iteration and 8, the local maximum considered above. Assume that optimization starts in 80 E N r). If oz < 3. 28, then the evoor at iteration t of the presented AO algorithm for the
6² = max(ø, (ass(,,c, 04) and & = min( 265
Max-Max problem with decaying gradient step size at
3/2
satisfies:
2 30.
[1+2)+
[[||s(t 1 $.115) S [0] S- 1211
+0²
(11)
t t+3 3
12 (t
where a = 2(1 + et)2 ()))()-),r) (O)(
We will need several auxiliary technical results. will use the following notation: (0, 91, 32) F(0, p where F is the objective function from the main body of the paper parameterised by entropy coefficient y 0, 31 = rr and 82 ~^ 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 TO.3 are given in Section 10.3.
We consider the dual optimization problem:
max 8
max A (2)
(0,
max (EC((A).NEE(0))
max
E (agyy.dd ..... wo ())_,,.))))))xxxxxxxxxx
(((())) Ai ]
[(0) + (1((())) + (2()(1)) exp
?
where y is parameter, / Tg. n², o is fixed trajectories' embedding and furthermore:
(=)=) 51)
A2(2) 511 f(2),
such that f(-) is random feature vector. The 7th entry of the feature vector is constructed follows:
[f()], = / cos(2 w w + bz.1), where w N (0 and ba. Uniform(O, 2TT). For the ease of notation denote a (w bai)) We consider stochastic gradient ascent optimization strategy of the following form:
at time f. we receive a single sample (It, Mt, T6.
FGD y6. . ....., Not ") ~ J XXXXXx x w, then we form feature vectors [f(x)) / cos(1, (etx + 6: b and [f(ye)]s =
V 3 cos w + b) and finally update:
g/41 + ITI [8 (xz)) (1 (1-exi) exp
(0)) ((x))+ ((o))))((())--o -(x.))
8211 Il. [3 av) (1-exp exp
((8)) f(ve) ) C(.
Ig (d 1) denotes the projection onto the Euclidean ball B2 2 ( of some given radius T centered at the initial iterate 19--
Let [,f. 35} denote the global optimum of W, (81, 32) computed on the entire data population, i... given access to infinite number of samples ("oracle"). Let (2 B) denote the Euclidean ball of radius r centered at 3. For the sake of the theoretical analysis we assume that (a lower- bound on) the radii of convergence rı, 12 for 31, 82, respectively, is known to the algorithm (this assumption
20

---

Under review as a conference paper at ICLR 2020
is potentially easy to eliminate with more careful choice of the step size in the first iterations). To be more specific, if at any point in time parameter 31 or Bg falls outside the ball (2 (8]) or B2 (r2, respectively, the projection is applied that pushes the parameter of interest to stay in the ball. Also, let v,, (31, f2) and VB, W (B1, (2) denote the gradients of W, with respect to 31 and $2, respectively, computed for single sample. Similarly, Va. W (51, 32) and Va, W., (31, 32) be the gradient of W, with respect to B1 and $2, respectively, computed for the entire data population, i.e. infinite number of samples.
Note that given any initial vector (8) in the ball of radius f centered at Sg we guaranteed that all iterates remain within an ra-ball of 22. This is true for all = 2. The projection is necessary for theoretical analysis but in practice makes little difference. The above is two-step alternated optimization scheme.
Let the population gradient operator, E2),,,82), where 1, be defined 9a(B1. 82) = Pa + 3 3. WV,,W (3),
10.3.1 ASSUMPTIONS
Let W;; (3.) = WW, (3,, (5) and W 2 W, 82). Let Q2, Q2 denote non-empty compact convex sets such B1 E 22 and B2 E f2. The following assumptions made:
Assumption 10.1 (Strong concavity). For d = 1,2, the function W (aaa is Ca-strongly concave near 84 i.e. for all pairs (aa, ba) in the neighborhood of B3 the following holds
Sa | balla.
W (aaa Wi, (haa (Va, W (ba)) aa ba) S
where Ça > O is the strong concavity modulus. Assumption 10.2 (Smoothness). For d = 1,2, the function W (Ba) is Sa-smooth, i.e. for all pairs (aa, ba) the following holds
W Wa,a(aa) W (bba (V W (bd), aa ba) 54 laa ball2.
where da > 0 the smoothness constant.
Assumption 10.3 (Gradient stability (GS) / Lipschitz condition). We assume W., 32) satisfy GS
(da) condition, for all d 1,2, over Euclidean balls B1 E (2(r1,88), E B2 , (8) given as follows
IV ( V Wr,W, 32)1la S dall33 87112. where da and d= (d mod 2) + 1.
Finally, define the bound that considers the expected value of the norm of gradients of our objective function follows: Vof + ozz where of sup W. B2) Br (((1,,8))A) B2 B2(r2, B2)) 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
{at ) and with & < 2 the error at iteration + satisfies the recursion
3
(a²)2 a²
E || 91112 || 35 (1 E 8 Bi
|1
33 12
+
where áb = maxx=112((a), o² 1-20"E420'6 and È min 1.2
26
I-a'*
DatCCa
The recursion in Theorem 10.4 expanded yielding the convergence theorem:
Theorem 10.5. Given the stochastic gradient iterates of Max-Max method with decaying step size
af
3/2
and assuming that & < 25, the error at iteration + satisfies
[2c- (+))+ T
3
2
9
E |||| - 37 12 + 18 --+1-1711 -351 E 11
311 + 1169-133
to
t 3
[2g
(t
where & = (2( and = min
28aCa
d2+
21

---

Under review as a conference paper at ICLR 2020
10.3.3 ANALYSIS
The theoretical analysis we provide below is extension of the analysis in Balakrishnan al. (2017) to the two-step alternated optimization scheme.
Proof of Theorem TO.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 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 91 (B1, (8)) under Assumption 70.1
(ttrngg concavity) and
Assumption 10.2 (smoothness) with constant step size choice 0 < a 5,451
is contractive, i.e.
||9| (5) -3ill S (1-
20d1G1
1131
8ill2
(12)
81 G
for all B1 E B2(r1,87).
Similarly, the gradient operator (2( 922 under Assumption tion 10.2 (smoothness) with constant step size choice 0 <
10.1
strong concavity) and Assump-
5ut is contractive, i.e.
(2) -32122 S (1
2002C2
|18 15112 ll
(13)
82 2
for all B2 E B2(r2, 15).
The next theorem also holds for d = 1, Let r1, 0 and B1 E B2(r1181): E B2(r,,
Theorem 10.7. For some radius ra > 0 (d 1,2) and triplet (ød, Sa, 6a) such that O S da < Sa S 8a suppose that the function W (3a) is Ça-strongly concave and 8m-smo,th, and that the GS (da) condition holds. Then the population gradient operator Ga(31: 32) with step such that 0 < a S. ming=1,2 batta is contractive over a ball (r 37). i.e.
(Ga()3133) 32 s (1 [a B7l12 + ard1 Billz
(14)
where d (d mod 2) + 1, Φ = max dd. and & == min ba+
26aCa
Proof.
(9, - -7112 113 + a V 3. W, (81. (81, 32) -37122
by the triangle inequality (and with d * d) we further get
(31, (2) by the contractivity from Lemma 10.6 and GS condition 2adaGa
|| + a p, W,, aa,,((a))) /2
V Ba (,
V Ba W (34)
11.ua - Sblla a3 1--- 64 + Cd
Proof of Theorem 10.4
Let 3411 (l where 8[+1 = 3 + aVV ( and 35 = B
a' a'V 22 (8 18). where Ve, W22 is the gradient computed with respect to single sample,
31 and A2 are the updates prior to the projection onto a ball (2 (2). Let A = 82+1 B3 and
A+ = -/4 6g. Thus
s | 12 1
= | -B1l 1 -
= (3 - 3d 34t1 + 3/ - 23
22

---

Under review as a conference paper at ICLR 2020
Let Qi Ve, (3i. (5) and Q5 = V #r (3 .85).
Then we have that Bz - Bgg a a'Q
We combine it with Equation 17 and obtain:
1 1 12 - 11
(a' a' + 2(33 -5))
s
(a)))((2))) Q0) 2a' (Q2) ( -8)
= (a')9112112 + 2a (0%. A1)
Let Qi = V., W3, 35) and Q == Vo, W, .8). By the properties of martingales, i.e. iterated expectations and tower property:
+ H & B|| + (a')"E||Q>1 + 2a'B((Qa, A
(15)
Let -= V W B2). By self-consistency, i.e. 32 arg maa W 7, (Ba). and convexity of Da we have that
(02.42) WV 22) =
Combining this with Equation 15 we have
E B ++1 12 <
E|| + (a ( + 20 20' Q3 At)).
Define 0 = By Q'Qa and 0 = 33 + a a'Qa Thus
a' (Qh- - Q2 A1i
= (9) -h - (3) Bi), 84 - -1))
= (g) - B1 -3)) -113/ Bi
by the fact that g 87 + +a' 3 (since Q7 0):
( By 84 81) -1181 BI
by the contractivity of g' from Theorem 10.7
}
(1 -1- a' - 33ll
Σ
11
11
+
Σ
1133
B3 1|2
1182
B3ll2
-1.81
-3113
im
i=d+1
dd-1
2
|ا }}
< (1 1' + a
Σ
|
++
1l
+
Σ
||
1
imd+1
Combining this result with Equation 16 gives
{ (1- (1- A'Il + a
d-1
Σ [ )
2
B ++ '121
(at) + 2E
Σ
|1
l1le +
i=d+1
-|Alle |
{ (1-01))1Aaa + +a'
(d-1
Σ || }
< E| + 60')003 + 2F
Σ
||
A
++1
2 +
k=++1
||l |
23

---

Under review as a conference paper at ICLR 2020
After re-arranging the terms we obtain
2
B S (at) + 20' 23 + 2F Σ
11 +
Σ
||
1=d+1
apply 2ab I a² + +²²
s (a') + (1 2 + a' Σ
(114 12 + (
+ a'o
(1 114
dd-1
= +11 11- 11-2a't + + +
| + 12
+ ta¹EE 6E
|
We obtained
d Σ
2
B1 S (a')
-0aa--a--- + ('61)11411 + o'
|4 + 11
+ a' Ob
Σ
|
Li=d+1
we next re-group the terms as follows
dd
2
B1|
a' Q
Σ
12
$ [1
11--2a' + 0 + + O'
Σ
|
+ (a
Lidd++
and then sum over d from
to
2
d-1
E
12 1
a oE
Σ
| +
112
1=1
2
$ 20' + +
Σ
|
11
OE
Σ
+ (aa)² Σ 03
ind-1
Let oo 02.
Also, note that
2
E
12
a' OE
Σ
11
E
| t 11
a'δE
Σ
|4 12
1=1
d= i=1
and
2
[1 2a' + a'Ø|E
Σ
|
+ aa' OE
Σ
Σ
|
+ (at)202
< [1 - 11-20'8 +
Σ
|
+ a OE
| 11
+ (at))22
Combining these two facts with our previous results yields:
[1 a | 1
2
< ( 22a' + aL$JEE
|
a OF
Σ |
+ (a')202
22-'g + 20'0EE
+ (a))20²
1==1
24

---

Under review as a conference paper at ICLR 2020
Thus:
2
20'5 + 2a Φ
2
A
++1 d 12
<
Σ
1 - aa0
1=1
(a
+ 1 - a²
Since Φ 25
1-2a'5-20'e 1.
3.
I-a's
Proof of Theorem 10.5
To obtain the final theorem we need to expand the recursion from Theorem 10.4
We obtained
E
1
1 20 [
+ (a')
2
1 1 ato o²
1 -ap
a al(25 - -36]
2
(at)²
|| 2
1 - a'
a 0
Ld=1
Recall that we defined o in Theorem 10.4 as
9' =1 1 1 20' + 20's
a¹(2E 3Φ)
1- -'ó
1- a'a
and denote
(a')2 1 - -aa
f
Thus we have
E
112
<1 9²) q²)TE
1 +f'02
d=1
S (1 9¹)
(1 9 10 II
1/
+ 1 12 {
f2
d==1
2
= (1 (1 - DE=
1A4-116
+ (1 )') + l'
dd=1
< (1 - (1) )11
(1 /² 2
A 12
4-2
+ 2
+ (1 -1-- f
d=1
(1- (1- (1 2
- +(1-
- (11-))(1---11111 -1-1-11111--
+ (1 2
We end-up with the following
E
11
<
E
A 2 II
Σ
f II
(1-))
j=it1
25

---

Under review as a conference paper at ICLR 2020
Set q¹ and
3 1+2
g'
a'
2 30 + g'th
3
[2 - 3 + 2) + 3 20
Denote A = 25- 3/ and B 26. Thus
3 2
a²
A(t + 2) +B
and
(a
ft
1
Ba
A(t + 2)[ + 2) + B
E
++1
12
1=11
3
<
2
II
2
1
1
A(i + 2) A(i
2) B
j ++2
==
to²
A(t
2)(A(t + 2) + B
242
1+1
t+2
3
11 2
+ -
1
2
AilAi B
j
A(l
[){ - 2) + B]
j=i+1
Since A > 0 and B >
thus
E
d=11
1+2
++1
1+2
<
2
+ 0
Ai Ai +
B
j
A(L + () - 2) + B)
d==1
jmit1
3
<
E
A
2
2
1
+ 0
(Ai)2
j
[A(t 2)}²
1=2
j=i+1
We can next use the fact that for any a E (1,2):
t+2
II
(1-
2
T + 1
t 3
i=t+1
The bound then becomes
2
2+2
++
t42
<
E
1
II
+o2
2
+ +
(Ai)²
j
[A(t 2)1²
1=2
2=2
Ø=
#+1
(i+) 1
|
+ 2
t
t+3
(At)²
1+3 3
[A(t + 2)
1+2
g
ii+1
E
Σ
L+3 3
(Ai)²
1 3
1
26

---

Under review as a conference paper at ICLR 2020
Note that (i+ 1)
28 for i = 2,3, ..... thus
E
A
+++1
2
(i
1)
E
1
+ 0²
t++3
A (t + 3)
122
i=2
1+2
2
S
E
+
f
1+ 3
A? (t + 3)
12
1+2
1+2
finally note that
Σ
- da < 2(1 + 3) Thus
F2
2
2
9
E
Σ
1471
+ 0²
t+
A2 + +
d=
substituting A = 25
30 gives
2 (
E
+ -
(++
- (t +3)
This leads the final theorem.
Proof of Theorem 10.3
In order prove Theorem 10.3 it suffices to apply Theorem 1O.5 and notice that:
function hv,c R + R defined as follows: lw, w Ae
x for A = 7 - is
2 2e+ 2 -sstrnngyy concave and its gradient is Lipschitz with Lipschitz coefficient 1 7
7
27 e7
(with respect to L2-norm) for C and satisfying: |w
1,
[[Vhv.c(v)lB & 2(1 + e²)² under above conditions.
10.4 MINMAX PROBLEM: THEORETICAL ANALYSIS
In this section aim to obtain similar results for Min-Max Problem 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 dx max (b, Q7) and & min( o+c u.)) where Φy and us are as in Theorem 10.3 Let s(t) (o(t), x p² (t)) be the solution obtained in iteration t and the local optimum. Assume that optimization starls in On E N(O", ,r). th < then the error at iteration + of the alternating optimization algorithm for the Min-Max problem with decaying gradient step size a²
[2-66. satisfies:
3/2
[[]] + - s- E1])s(0) 8-13. 1.1110 1+3
(16)
12 6o..3(1+33)
where 2 + SU) V F()2
We the same technical notation in the previous section, with the only exception that this time p
we denote: (8, B1, 32) = - F(t,, İ= .p x(.)). We consider the MinMax problem the following form
min
max (
(0.
min
max
E( De..g................ .........D...... .o w............... 1)tpp//x..pxxxxxx-----
0
(2(2(z)) + A2( ((((x). (0)()))))
- -C(0) - x)((()))))) + )g(((v))) exp
27

---

Under review as a conference paper at ICLR 2020
where y is 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:
nit mNt..) ~ Tge x v x w X ... x w, then we form feature vectors [f(x)]]: = V cos(2, wit + bb.:1) and [f(e)]: =
at time we receive a single sample (2t, yt, Hie #t.l ...., Rig Ki
B..D 12lell ,...,,
V i cos wit. 3 + b)t"", and finally update:
ge+l - I1 |0 at Vo o= (6, B1 35)
17t1 + 112 B1 + (1ze)
(1 (1- exp
(0)) (B)))
s(x.) -(x))-((a.))) -((x,)
(10t1) ((0))+(88)) f( ((a)) ))
85*1 + II3 5/ aye
e1-exp
exp
II, denotes the projection onto the Euclidean ball B2 ( and IL (d 1, 2) denotes the projection onto the Euclidean ball B2( (2))
Let (3.3) denote the the global optimal solution of the saddle point problem ming ,, W,(G (8, (h.32) computed on the entire data population, i.e. given access an infinite number of samples ("oracle").. As before, we assume that (a lower-bound on) the radii of convergence r2, r3 for 0, 31, 32, respectively, is known to the algorithm and thus the projection applied to control 0, 31, f32 to stay in their respective balls. Also, let Vo ( B1, f2), Va, W (0, 31, (22 and V3, W (0, , 82) denote the gradients of W,, with respect 0,3, and 32, respectively, computed for a single sample. Similarly, V VoW, (0, B1, 82). v (0, 9,, 82) and V3, (,,,, ,1, be the gradient of W/ with respect to 0, B1 and 82, 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 Р centered at 0", we are guaranteed that all iterates remain within an 11 bbll of 6 and given any initial vector 02 (d = 1, 2) in the ball of radius Ia centered 3%. we are guaranteed that all iterates remain within an rall of B3. The projection 2
is necessary for theoretical analysis but in practice makes little difference. The above is a threc-step alternated optimization scheme.
Let the population gradient operator, Ea,0,31,,3 where d = 1,2,3, be defined as
82) = 0 a (e,,
and
1,00,,,80 (2) = &V (0, for d = 2.3.
10.4.1 ASSUMPTIONS
3,,185) (0",3,,,, and W (32) = (3",,,,,,00 Let Q1 S2252 22 denote non- non-empty compact convex sets such E 01 E Q2 and B3 E Ggg The following assumptions are made:
Let W 7 (0) (0, 35), W (3)) = WO
Assumption 10.4 (Strong convexiny/concavity). The function W (0) is Ç1 St-stroy convex near l and the functions W 2 (B1) and W (32) are Ç2- and Ça-strongly concave, respectively, near 3i and 3r, respectively, where Ç1, S2, {3
Assumption 10.5 (Smoothness). The functions W (0). 291),, and (B2) are 61-,82,, and d3- smooth,, respectively, where d1, 82, 83 > O are the smoothness constants.
Assumption 10.6 (Gradient stability (GS) / Lipschitz condition). We assume W., (0,00, (1) satisfy GS (da) condition, for all d 1, 2,3, over Euclidean balls 0 E B2(r1, 0") 31 B2(r2, (1), 32 Ba (r3, 82) given as follows
VoW, (0) VOW, 22 d1
Σ d=1
||3a B7ll,
later we also use alternative notation for gradient V == (7 (5 WWW, (8,,
28

---

Under review as a conference paper at ICLR 2020
for d -12
1V W 7, (Ba) (0,, (1 S & 0- + 115 -5112))
where do and d (d mod 2) + 1.
Finally, as before, define the bound o that considers the expected value of the norm of gradients of our objective function as follows: a of + + oog where of (0, 31, 32)] E 82(r1,0"), E (2 81), S2 E B2 (5)) and for d = 1,2 V W sup
0+1 = sup E BII V W (0, 31, 32)1 8 (2(11. 81 E B2(r2, (t(r,,81), E (2(
10.4.2 MAIN THEOREMS
Theorem 10.9. Given the stochastic gradient iterates of MinMax method with decaying step size
{0 and with & < 5/ the error at iteration t satisfies the recursion
E [[]et+] - -11
112 + ||
11 || -2112]
(at)
(1 D¹) [110* 0- 1012 + | 11 + |1 - 3212
0²
1
2a ②
where @ = max (d)) o² 1- 2- E+4n'o and & = mina-1,2,3
28 C2+Ca
1-2a't
The recursion in Theorem 10.4 expanded yielding the convergence theorem:
Theorem 10.10. Given the stochastic grudient iterates of MinMax method with decaying step size 3/2
a [2 66) (1+2))33 and assuming that Φ < & the error at iteration + 1 satisfies
[|
111 +
||
1T + 11 [311
9
< E |1100 1- + 11 - 11 + 1130 - B311
+
1 + 3
[2 6t) + 3)
where o = max and & = min.4-12,3
26 Ca
which in turn relies Theorem 10. and Lemma 10.11 both of which are stated below. Proofs of the lemma and theorems follow in the subsequent subsections.
Proof of Theorem relies on Theorem 10.9
10.4.3 ANALYSIS
The next result is standard result from convex optimization (Theorem 2.1.14 in Nesterov (2014)) and is used in the proof of Theorem Io. 12 below.
Lemma 10.11. The gradient operator Gi (0, B 2
(2) under strong convexity and smoothness assump- tions with constant step size choice 0 < a S a.tCi is contractive, i.e.
0 (1- 20di51
(0,
110
00
81 + +
for all 0 E B2(r1,0*).
Similarly, the gradient operator (1 .2,32) under strong concavity and smoothness assumptions with constant step size choice 0 < a 82tc2 is contractive, i.e.
10", 81, (5) Bill2 (1
2ad252
|1 - -9112
82 +
for all 31 B2(r,,
And similarly, the gradient operator ,31,3) ,2,0",3i, 32) under strong concavity and smoothness assump- tions with constant step size choice 01 < a TstCa is contractive, i.e.
2adg6a
|| ,8, B2) B2ll2 1
11
17112
83 E3
for all B2 Br(rs, 85)
29

---

Under review as a conference paper at ICLR 2020
The next theorem holds for d 1,2.3. Let 1et,1,72,77 0 and 0 E 0r(r1,0").33 B1)r,,81): B2 (rs. 82). Theorem 10.12. For some radius rı and a triplet (Φ1, S1, 51) such that 0 S 01 < S dj, suppose that the function W (0) is Si-strongly corvex and 61 smooth and that the GS ((1) condition holds.
Then the population gradient operator 9I (0, 82) with step a such that 0 < S ming=1,2,3 52- is contractive over a ball (2 0²), i.e.
(0,,0,,1, 0* S (a) 00 + ac Σ
113 Ball2
d=1
26a5a where 6 max du and & == min 6a+
For some radius ra (d = 2,3) and triples (oa: Sa, 8a) such that S da Sa S ôa. suppose that the function W (Ba is Ça-strongly concave and 8a-smooth and that the GS (0a) condition holds.
Then the population gradient operator Ga B1, (2) with step such that 0 S mind=1,2,3 Ca contractive over a ball (r(r (2)) i.e.
|] $1, B2) B $ (1 fa)l - fall2 + a - + 110 0 112).
where d = (1) mod 2) + 1, max4-1,2,3 ds and & min.- 1,2,33 da+Ca
28
Proof.
|| (0, (1, by the triangle inequality we further get
V0, 0²112
< |10 Vo 00 al V (0, B1, B2) VoW by the contractivity from Lemma 10.11 and GS condition
(1- 2 110 0- + aΦi Σ
51 +c
Ba 82ll.
The proof of the rest of the theorem is analogous to the proof of Theorem 10.7
Proof of Theorem 10.9
Let 01 = 0, 02 = 31, and 0g = Az.
Let 05+1 = la(0 where 0t+1 = 0f - a'V,, W1 W (0), 022 1g).
05+1 = 0g
a (0 02 08), and 0,+1 02 a Vo, W 0,+1 0%), where Vo, is the gra- dient computed with respect to single sample, 01 02. and 0g are the updates prior to the projection. Let A1+1 = - 0 + 07 and for d = 2, 3,
1- A = 0 6% Thus
14 1|4111 s 14 1 1
10 0nl 110 07
(0% 02 0th 022 02-203)
Let Qi = Ve, W: (0,.00,0 05, Ø5 = Vo, (0' 02 0.), and Q = Vo, (0 (0,t1,ttt,
Thus:
| 113 - 1 < (ata'n a d'Q², + 2(0') 03)
(a))((2)) Ø + 20' (2)) [00; 07)
(a99110111 20' Q 4/)
30

---

Under review as a conference paper at ICLR 2020
V,, Let Qi = Vo, (o, 82, 03) Q = VV, (0,+ 8 83), and Q3 == Ve, W, (0]+1 0-211 .8). By
the properties of martingales, i.. iterated expectations and tower property:
12
E| + (a') 11 2a Ot: A:)1
Let Q7 = Va, (0,, 85 05). By self-consistenge i.e. 0j = arg E0X W Wr,a(0a) and convex- ity of S4a we have that
C -3) = (Ve, W,,,, 05, 05), A') = O.
Combining this with the above inqquality yields
++1 d 18 s
E|1411@1 + (a') 3 + 22' Qà: A.)1.
Define Gf == 0 - a Qi and gi = 0f - ai Also, for d = 2,3 define 9 = 02 + Q'Qa and
g 07 + az Q² Thus
a' (Qa Qa. A1)
(g) 9 (0'2 02), 0²- 02)
(9) 0 02 07) |10| 0
by the fact that G 0' + a QQ- 0g (since Q7 = 0):
(g) 0g 02 0 110's 1
by the contractivity of G from Theorem 10.7
}
(1 (1-a's)110) 0ill + 0.
Σ
1000
0,112
Σ
100 0:112
110 02112
-10 0011
1=d+1
d-1
3
|االه2 }}
< (1- d'E)ll ta's
Σ
1l + Σ
1
1=d+1
Thus
B1
<
| + 21II + 2n "B(B(' - Q3. Ai)1
E| + (a Q1 2E {( (1-
d-1
lAille }
(1 a')
A + a'
Σ
||
l12 + Σ
id++1
-14 12 1
{ (1 (1- + a Σ
|llill }
B| + ( 2E
A lle + Σ
1=1
f=d+1
114
After re-arranging the terms we obtain
E||l B 1 |1] < (at) (1 2a ( + 2a' OE
Σ
| |Alt +
Σ
||
||
imd+1
apply 2ab S a² b
s (a')20²+ + (1 20 L + aLØB
(11 ++ 12 + (
+ a' SE
Σ
(1 - 1
Limdt1
dd-1 Σ
3
(af)²0z+ + [1-22' + 20 [] + aO
| 1
+ + + OE
Σ
|A
+=d+1
31

---

Under review as a conference paper at ICLR 2020
We obtained
3
B| S (a + -2a' 2a' [
+a'eE dE
Σ
A +1
11
+ a a'
Σ
1
we next re- group the terms as follows
B a'
Σ
12
S
20' + 2a' 2230)1114111 o' dE
Σ
|
(a')
i=d+1
and then over d from 3
3
d-1
E
++
12
Σ
| +1 1
3
3 Σ
$ [1 2 + 2a 2|
Σ
A 2
+ a éE
Σ
Σ
|
- 2(a²))
od
1=1
d=1i=d+1
d=1
Note that
3
3
3 d-1
E
1 +
12
2a 6E
A 1+1
S
12
- a PE
A ++ 12
and
3 Σ
3
[1 20'E + 20'
|
112
a'
'
+ +af)202
-d=1
d==1 i=d+1
3
< [1-2a'E 2a 21
Σ
|
12
2a' 'E
Σ
|1
+a0)002
d=1
Combining these facts with our previous results yields:
2a |
1
3
3
< [1 - 2a'5 + 2a $1E
|
+ 2a dE
Σ
|| 2
+ 0at)202
3
[1 - 2a'E + 22'
| 2
+ 22²
Thus:
2 + 4a $0
A ++] 11
A 11
20/
1=1
(a²)²
+
1 2xx 1
Since & < 1-2n'E+4a's 1.
31
-00
Proof of Theorem 10.10
To obtain the final theorem need expand the recursion from Theorem 10.4 We obtained
32

---

Under review as a conference paper at ICLR 2020
E
1+11
1
1==1
20¹15 15-
200]
(at)²
EE
| 11
20' o²
20¹0
a' (2f 66]
(at)²
|
+ 1 20'
1 20
Recall that we defined q' in Theorem 10.4 as
2a'5 + 4a'
a a²12E - 66]
9¹ =
1 -2
1 2
and denote
(at)²
M 1 20'0
Thus have
E
Σ 1 11
S (1 q¹)E
Σ
ff
s (1 90) (1 -1--1111E I
| 1-1-1 12
+ -
+ f²
d==1
3
= (1 - (1 - DE
A 1 112
+ g²) 1r 0² + praz
d==1
s (1 - (1 /11) (1 /² 2)
d 113
+ r
+ (1 01 l'oz
d=1
(1 (1-99) d1) q²=1)(1 - 3 Σ
+(1 - d') - ---1-----~-- + (1 /-) +
We end-up with the following
E
E
0 2
II
g)
)) Σ f 2=0
II j=+11
(-
d
d
1=0
Set g¹
and
t+2
d' 2 6Φ + 2q
a'
3
[2 66) (t + 2) + 36
Denote A = 25 66 and B = 3d. Thus
a'
2
A(t + 2) + B
and
(at)
2 Ba
A(t + [)[A(t + 2) + B
3
33

---

Under review as a conference paper at ICLR 2020
2
E
1+1
12
d=1
(1
2
E
11 1=0
+ 2
II
i+2
A(i 2) + 2) B
j+2
d=11
s=i+1
to²
A(t [){t(t + 2) + B]
3
E
2
II 1=2
1
+ 0²
II
1 1-
2
+
Ai[Ai B
]
A(t 2) A(t - 2) + B
1=1
1=2
jmit1
Since A > 0 and B >
0 thus
3
E
A
++1
12
1=1
<
E
21
II
0²
II
1
1-
2
+ +
Ai(Ai BJ
j
A(t ()
-
1=11
1=2
1=2
j=++1
3
1+2
H+1
1+2
2
<
2112
+ 0²
(Ai)
j
[A(t + 2)}²
d=11
1=2
j=it1
We can next use the fact that for any a (1,2): 2):
++1
II 1=T+1
(1-
t+ 3
The bound then becomes
242
++1
t42
E
1
II
+o2
2
II
+
(ii)
[A(t 2)1²
1=2
2=2
Ø=
#+1
g
(i+)
|4
+ 2
t
t+3
(At)²
t+ 3
[A(t + 2)
1+2
g
ii+1
Σ |
L+3 3
(Ai)²
1
1 3
1
34

---

Under review as a conference paper at ICLR 2020
Note that (i+ 1)
28 for i = 2,3, .... thus
E
1 ++1
2
142
1)
E
+ +
t +
A + 3)
12
1=2
2
1+2
1
E
14
+ 0
1 3
A? (t + +3)
1=2
finally note that
1
142
dx S 2(1 + 3) Thus
T2
2
9
E
o²
|+3
A + +3)
-d=1
substituting A
26
60 gives
2 8
E
14
+ 0
1
3
2
(1
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.
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 's definitions here for the reader's convenience: 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 / can't be visited at any other time.
4. S and A are finite sets.
The third assumption is solely avoid having to define time indexed Value function. It can be completely avoided. We chose not to do this in the spirit of notational simplicity. These assumptions be relaxed, most notably We chose to present the finite horizon proof because of the nature of our experimental results.
show similar results for the discounted and infinite horizon case.
Let & id be the identity embedding that & P. In this case p& denotes the distribution of trajectories corresponding to policy æ. We define the value function V - as
V (&t = =}) = ETOrad
R($++1 o2 02, 82)18;
The Q-function Q" S x A - as:
T Σ f=1
Q" a; =
Eywpia
R(s441, ar,
Similarly, the advantage function is defined as:
A (s, a) = -*(( (s)
We denote by V(T) Eyweea Σf R(5141, at, at,ss)) si) the visitation frequency as:
the expected reward of policy n and define
(=()) Eynpi
Σ 1(s s)
The first observation in this section is the following lemma:
35

---

Under review as a conference paper at ICLR 2020
Lemma 10.13. nwo distinct policies π and it be related via the following equation
V(#) V + Σ
p=(s) Σ #(als) (s.
Proof. Notice that A (s,
Pp' (a,s) ( +VT(s)) (T)))) Therefore:
T
p
An (5),
= BBrpp
R($141, 5²) + V (s(+1) V (st)
1=0
R(s 01 st)
Beg [V* (so)]
-V( + V(t)
The result follows.
See Sutton et al. (1998) for alternative proof. We also consider the following linear approximation to V around policy (see: Kakade & Langford (2002)):
L(T)= = + Σ
Pr(s)
Σ
i(als) (s.
SES
AaEA
Where the only difference is that f= substituted by Pr- Consider the following embedding Φ. r R SI defined by (9(()), Σ 1(s = s), and related cost function defined as: C(v, w) = -wwl-
Lemma 10.14. The Wasserstein distance (P m is related to visit frequencies since:
(IP D
Σ (p*(s) p=(()]
sES
Proof. Let be the optimal coupling between pa and p.' Then:
(PP
P$" = E [||u elli]
En,wol [lu, Uyl]
SES
Where u, and denote the S indexed entry of the u and v vectors respectively, Notice that for all E S the following is true:
E Bnmpo' [ue]
Ex- [0a]
I En [lue voll
(*(
(a)(s)
The result follows.
These observations enable us to prove analogue of Theorem from Schulman et al. 2015). namely:
Theorem 10.15. If WDD. p#1
< 8 and e = max,,a (A a)), then V 2 L( dc.
As in Schulman et al. (2015). Theorem 5.1 implies a policy improvement guarantee for BGPG from
Section |.3
36

---

Under review as a conference paper at ICLR 2020
Proof. Notice that:
Σ (x(ab))4(s00)))
V( - - = Σ
(p=(s) p=(())
SES
NEEA
Therefore by Holder inequality:
( - <
Σ
(x)(8) pa(5)]
sup
Σ aEA
A A"a
sES
SESS
SWDo(B?" 121)58
The result follows.
We can leverage the results of Theorem 10.15 embedding &* give monotonically improving sequence of policies. The proof can be concluded by following the logic of Section in Schulman et (2015).
show wasserstein trust regions methods with
37


Reviews:
[{"review_id": "rkxv0JQtYS", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 1, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary\n\nThis 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\".\n\nReview\n\nThe idea of generalizing the trust-region of TRPO by using the WD measure and behavior maps is intriguing. The paper introduces many choices of behavior maps that could lead towards interesting algorithms that merit additional study. I think the connections between the proposed family of algorithms and WD-based TRPO and distributional RL is highly motivating.\n\nThere are, however, a several key issues with the empirical study of the proposed methods that make it challenging to assess the value of introducing another partially understood deep policy gradient algorithm. The results show only behavioral studies of the algorithm, not investigating the effects of the WD based regularizer or the effects of the choice of behavioral map. The results are also significantly limited in their statistical significance, making distinguishing between algorithms difficult in most cases. And the comparison of wall-clock time is particularly difficult to assess due to the uncountably many possible sources of noise when comparing wall-clock time of highly complex algorithms. In the following paragraphs, I will expand upon each of these points.\n\nOne of the key contributions of this paper is the ability to define regularizers based on the definition of the behavior map. The paper introduces many such behavior maps, those over state visitation, actions, and returns; however, the paper does not investigate the effects this novel choice has on the results. It is unclear if introducing each of these behavior maps leads towards different results or if they each induce roughly the same final performance of the agent. It would be highly valuable for me to see a more careful study of the effect of choosing each of these behavior maps on a single simple environment, clarifying that this formulation leads to a family of useful algorithms. As it stands, BGPG may only beat TRPO on these domains because it had more meta-parameters to choose between (BGPG introduces many new meta-parameters: choice of behavior map, kernel for produce RKHS, the meta-parameters of that kernel, the meta-parameters of the behavior map like trajectory length, the entropy regularization term in the WD, etc.). Without a careful study, or intuitive explanation of any of these parameter choices, it is unclear if BGPG won simply through overfitting to the problem.\n\nThe statistical significance of the proposed algorithm is impossible to assess in the given form. The comparisons are made using only five random seeds and the standard error bars are frequently quite large. I refer to Henderson et al. 2017 for further explanation as to why five random seeds is simply too few to provide a meaningful comparison between algorithms. Instead, I'll discuss a few of the results in particular. In figure 3a, I notice that the proposed method has high variance until it plateaus around -300 reward; why? In each of the remaining plots of Figure 3, I notice that the propose method is significantly higher variance than any of its competitors. This greatly leads me to suspect that the proposed method would have much lower average performance if studied across a greater number of random seeds. In the walker domain (Figure 3d), why do BGPG and TRPO both plateau at the same point for many timesteps, then eventually BGPG starts improving again? From Figure 5, the paper somewhat misleadingly states that \"BGES is the only method that drives the agent to the goal in both settings.\" However Figure 5 on the left clearly shows that BGES and NSR-ES have indistinguishable performance, and the error bars indicate that neither significantly outperform NoisyNet-TRPO. In fact, due to the high variance of BGES it is unclear if it would drive the agent towards the goal on average if run over more random seeds. On the right, NoisyNet-TRPO noticeably outperforms BGES and is significantly lower variance. The language used to describe Figure 6 is strong, stating that Figure 6 \"proves that the benefits come here not just from introducing the regularizer, but from its particular form.\" However, I tend to disagree that Figure 6 reliably proves anything. Although it appears that BGES is significantly outperforming other methods, the number of meta-parameters over which it gets to optimize makes it difficult to say if the performance gain is from overfitting to the problem or the form of the regularizer. Additionally, good performance demonstrated across a single problem is hardly proof especially using only five random seeds (Henderson et al. 2017 Figure 5).\n\nAlthough this plays comparatively little role in my scoring of the paper, I feel it is necessary to discuss briefly. Because it is impossible to determine the source of the speed differences between the particle approximator and BGPG, the results in Figure 4 are extremely difficult to interpret. It simply could be that BGPG uses slightly more optimized code than the particle approximator, or it could be that there is in fact a significant performance difference between the algorithms. The only true way to discuss performance differences between algorithms generally is through computational complexity for these reasons. In the case that these statistical algorithms share the same computational complexity, then further analysis including convergence rates would be able to shed light on the speed difference. However, wall-clock time vs. performance has so many confounding factors, it is a fairly meaningless unit of measure. For an empirical investigation of speed, I would suggest giving each algorithm similar number of learning steps or similar number of updates to their weights and compare performance in this way.\n\nAdditional Comments (do not affect score)\n\nI'm curious if this work could be extended to the off-policy case. It does seem like a minor disadvantage that distributional RL is well-defined in both the on-policy and off-policy cases, but the proposed family of methods is not. From my reading of the paper, it appears that there is little preventing this extension. Is this true?\n\n--------\nEdit after reading other reviews and rebuttals.\n\nI appreciate the response noting that the meta-parameters of this algorithm were not tuned. I still feel that understanding the sensitivity to these parameters is important to understand how the choice of these parameters impacts the performance. While it is certainly likely that extensive tuning of these parameters could yield even more improved results, it is also quite possible that the chosen parameters have biased the results. If a more intuitive discussion of the choice of meta-parameters could be included in the paper, I would be more willing to concede this point.\n\nI recognize that 3 or 5 random seeds are a standard in the deep RL community, but I think Henderson et al. 2017 make this point better than I could: this is not a standard that we should hold to. I think that using small domains and smaller networks would allow running across more random seeds and providing a much more careful scientific study. This paper does not convincingly show state of the art performance, a metric that is nearly impossibly to define in a field that moves as quickly as ours, so should not strive to follow the same demonstration study design that SOTA papers follow. A paper with strong theoretical motivations (such as this) should compliment with strong empirical understanding of how that theory translates to practice. Instead, having strong theoretical motivations backed by a demonstration that the algorithm still works is much less convincing (to me). I think a more convincing set of experiments that match well with the intended purpose of the paper (theoretical contribution of behavioral embeddings) would be an investigation to _how_ the optimization affects the learned policies or representations. Instead of performance benchmarks, this would provide understanding for the various effects of the new framework.\n\nAfter the rebuttal phase, I now have an additional new concern about the clarity of the writing and the completeness of the empirical details. Many of the results do not mention most of the important details necessary to replicate or even interpret the results. I had assumed that details were consistent across figures (e.g. figure 3 uses 5 random seeds, so I assumed the same to be true for all other figures that did not specify), but the author response made clear that this is not true. In additional, little to no details are provided in the paper referring to the choices of meta-parameters for any algorithm or the methodology used to chose these meta-parameters. After these details were given in the rebuttal, my concerns were greatly alleviated, but not completely removed. It is unclear if the authors intend to update the paper before the final version, but based on the language of the rebuttal (e.g. indicating it is purely a misunderstanding of the paper), I am forced to assume that the paper will not include these important details.\n\nTo summarize. I think the theoretical underpinnings of these paper are extremely motivating and interesting. I do not want to see these innovations lost. However, the current state of the empirical section of the paper leaves too much open for me to be able to recommend an accept at this time. I do not believe a paper should solve everything in one pass (e.g. careful parameter studies, SOTA demonstrations, real-world demonstrations, etc.), but I do not believe this paper demonstrates any one of these convincingly.", "extra_scores": {"experience_assessment": "I have published one or two papers in this area.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper thoroughly.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "r1l0QRy0FH", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "This work explores two uses of Wasserstein distances (WD) within reinforcement learning: the first is a variant of policy gradient, where WD is used to guide the policy search (instead of alternative such as Trust-region used in TRPO); the second is a variant of evolutionary search where WD is used again to guide the policy updates.\n\nOne of the strengths of the work is to clarify the notion of Behavior embeddings (Sec.3), which I expect can have several uses in RL.   In this paper, the behavioral embeddings are assumed to be given; it would be interesting to discuss/explore learning these embeddings.\n\nSection 4 of the paper reviews key concepts related to WD.  This is much harder to follow for an RL researcher, and would be improved by adding some intuition relating the material presented to the concepts of Sec.3.  Furthermore, this confusion carries out in Sec.5.  For example, what is the best way to think of \\lambda_1 and \\lambda_2?  And the maps s_1 and s_2?  What are necessary/desirable properties of P^\\phi_b?   There are also many steps packed in Alg.2 & Alg.3, which are difficult to unpack.  For example, what are the \\epsilon (step 1., Alg.3), scalars or vectors, how are they sampled?  It would be helpful to have a discussion of the complexity (both data & compute) of both algorithms.\n\nSection 6 presents empirical results for each proposed algorithm.  Corresponding baselines are presented, but I would be interested to see a wider set of baseline methods. The literature is rich with methods in these classes, both variants of TRPO and ES.  It’s necessary to at least pick a representative sample to show and compare (e.g. GAE, SAC).  I am also puzzled by the actual results presented, for example the Hopper reward shown in Fig.3 seems much worse (by orders of magnitude) compared to that reported in the SAC paper (Haarnoja et al. 2018).\n\n\n", "extra_scores": {"experience_assessment": "I have published in this field for several years.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I carefully checked the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}, {"review_id": "Bye86rHeqB", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "\nSummary: 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. \n\n\nOverall, 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. \n\n- 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. \n\n- 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? \n\n- What is the effect of \\beta when positive vs. negative?\n\n- 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.\n\n- It would also be more self-contained if the paper includes experiment settings.\n\n- 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.\n\n\n* Minor comments:\n- The KL definition on page 4 use rho instead of \\xi\n\n\n- are u and v in Eq. 3 functions in \\cal C(X) and \\cal C(Y), respectively?", "extra_scores": {"experience_assessment": "I have read many papers in this area.", "review_assessment:_thoroughness_in_paper_reading": "I read the paper at least twice and used my best judgement in assessing the paper.", "review_assessment:_checking_correctness_of_experiments": "I assessed the sensibility of the experiments.", "review_assessment:_checking_correctness_of_derivations_and_theory": "I assessed the sensibility of the derivations and theory."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The authors introduce the idea of using Wasserstein distances over latent "behavioral spaces" to measure the similarity between two polices, for use in RL algorithms.  Depending on the choice of behavioral embedding, this method produces different regularizers for policy optimization, in some cases recovering known algorithms such as TRPO.  This approach generalizes ideas of similarity used in many common algorithms like TRPO, making these ideas widely applicable to many policy optimization approaches.  The reviewers all agree that the core idea is interesting and would likely be useful to the community.  However, a primary concern that was not sufficiently resolved during the rebuttal period was the experimental evaluation -- both the ability of the experiments to be replicated, as well as whether they provide sufficient insight into how/why the algorithm performs.  Thus, I recommend rejection of this paper at this time.
Decision: Reject