Behavior-Guided Reinforcement Learning

1 INTRODUCTION . One of the key challenges in reinforcement learning ( RL ) is to efficiently incorporate the behavioral characteristics of learned policies into optimization algorithms ( Lee & Popovic , 2010 ; Meyerson et al. , 2016 ; Conti et al. , 2018 ) . The fundamental question we aim to shed light on in this paper is : What is the right measure of similarity between two policies acting on the same underlying MDP and how can we devise algorithms to leverage this information for reinforcement learning ? In simple terms , the main thesis motivating the methods we propose is that : Two policies may perform similar actions at a local level but result in very different global behaviors . We propose to define behaviors via so-called Behavioral Embedding Maps ( BEMs ) , which are functions mapping trajectories ( realizations of policies ) into latent behavioral spaces representing trajectories in a compact way . BEMs enable us to identify policies with their Probabilistic Policy Embeddings ( PPEs ) , which we define as the pushforward distributions over trajectory embeddings as a result of applying a BEM to a policy ’ s trajectories . Importantly , two policies with distinct distributions over trajectories may result in the same probabilistic embedding . PPEs provide us a way to rigorously define dissimilarity between policies . We do this by equipping them with metrics defined on the manifold of probabilistic measures , namely a class of Wasserstein distances ( WDs , Villani ( 2008 ) ) . There are several reasons for choosing WDs : • Flexibility . We can use any cost function between embeddings of trajectories , allowing the distance between PPEs to arise organically from an interpretable distance between embedding points . • Non-injective BEMs . Different trajectories may be mapped to the same embedding point ( for example in the case of the last-state embedding ) . This precludes the use of likelihood-based distances such as the KL divergence ( Kullback & Leibler , 1951 ) , which we discuss in Section 6 . • Behavioral Test Functions . Solving the dual formulation of the WD objective yields a pair of test functions over the space of embeddings that can be used to score trajectories . The behavioral test functions underpin all our algorithms , directing optimization towards desired behaviors . To learn them it suffices to define the BEM and the cost function between points in the PPE space . To mitigate the computational burden of computing WDs , we rely on their entropy-regularized formulations . This allows us to update the learned test functions in a computationally efficient manner 1Available at https : //github.com/behaviorguidedRL/BGRL . We emphasize this is not an exact replica of the code from our experiments , but a demo to build intuition and clarify our methods . via stochastic gradient descent ( SGD ) on a Reproducing Kernel Hilbert Space ( RKHS ) . We develop a novel method for stochastic optimal transport based on random feature maps ( Rahimi & Recht , 2008 ) to produce compact and memory-efficient representations of learned behavioral test functions . Finally , having laid the groundwork for comparing trajectories via behavior-driven trajectory scores , we address our core question by introducing two new on-policy RL algorithms : • Behavior Guided Policy Gradients ( BGPG ) : We propose to replace the KL-based trust region from Schulman et al . ( 2015 ) with a WD-based trust region in the PPE space . • Behavior Guided Evolution Strategies ( BGES ) : Inspired by the NSR-ES algorithm from Conti et al . ( 2018 ) , BGES jointly optimizes for reward and novelty using the WD in the PPE space . In addition , we also demonstrate a way to harness our methodology for imitation learning ( Section 7.3 ) and repulsion learning ( Section 9.4 ) , and we believe there may be many more potential applications in the future . 2 MOTIVATING BEHAVIOR-GUIDED REINFORCEMENT LEARNING . Throughout this paper we prompt the reader to think of a policy as a distribution over its trajectories , induced by the policy ’ s ( possibly stochastic ) map from state to actions and the unknown environment dynamics . We care about summarizing ( or embedding ) trajectories into succinct representations that can be compared with each other ( via a cost/metric ) . These comparisons arise naturally when answering questions such as : Has a given trajectory achieved a certain level of reward ? Has it visited a certain part of the state space ? We think of these summaries or embeddings as characterizing the behavior of the trajectory . We formalize these notions in Section 3 . We show that by identifying policies with the embedding distributions that result of applying the embedding function ( summary ) to their trajectories , and combining this with the provided cost metric , we can induce a topology over the space of policies given by the WD over their embedding distributions . The methods we propose can be thought of as ways to leverage this “ behavior ” geometry for a variety of downstream applications such as policy optimization and imitation learning . This topology emerges naturally from the sole definition of an embedding map ( behavioral summary ) and a cost function . Crucially these choices occur in the semantic space of behaviors as opposed to parameters or visitation frequencies2 . One of the advantages of choosing a Wasserstein geometry is that non-surjective trajectory embedding maps are allowed . This is not possible with a KL induced one ( in non-surjective cases , computing the likelihood ratios in the KL definition is in general intractable ) . In Sections 4 and 5 we show that in order to get a handle on this geometry we can use the dual formulation of the Wasserstein distance to learn functions ( Behavioral Test Functions ) that can provide scores on trajectories which then can be added to the reward signal ( in policy optimization ) or used as a reward ( in Imitation Learning ) . In summary , by defining an embedding map of trajectories into a behavior embedding space equipped with a metric3 , our framework allows us to learn “ reward ” signals ( Behavioral Test Functions ) that can serve to steer policy search algorithms through the “ behavior geometry ” either in conjunction with a task specific reward ( policy optimization ) or on their own ( e.g . Imitation Learning ) . We develop versions of on policy RL algorithms which we call Behavior Guided Policy Gradient ( BGPG ) and Behavior Guided Evolution Strategies ( BGES ) that enhance their baseline versions by the use of learned Behavioral Test Functions . Our experiments in Section 7 show this modification is useful . We also show how to use Behavioral Test Functions in Imitation Learning , where we only need access to an expert ’ s embedding . Although our framework also has obvious applications to safety , ( learning policies that avoid undesirable or dangerous behaviors ) we leave this for future work . We also consider simple heuristics for the embeddings ( inspired by other existing use cases ) , but believe future work on learned embeddings could be a significant enhancement . 2If we choose an appropriate embedding map our framework handles visitation frequencies as well . 3The embedding space can be discrete or continuous and the metric need not be smooth , and can be for example a simple discrete { 0 , 1 } valued criterion 3 DEFINING BEHAVIOR IN REINFORCEMENT LEARNING . A Markov Decision Process ( MDP ) is a tuple ( S , A , P , R ) . Here S and A stand for the sets of states and actions respectively , such that for s , s′ ∈ S and a ∈ A : P ( s′|a , s ) is the probability that the system/agent transitions from s to s′ given action a and R ( s′ , a , s ) is a reward obtained by an agent transitioning from s to s′ via a . A policy πθ : S → A is a ( possibly randomized ) mapping ( parameterized by θ ∈ Rd ) from S to A . Let Γ = { τ = s0 , a0 , r0 , · · · sH , aH , rH s.t . si ∈ S , ai ∈ A , ri ∈ R } be the set of possible trajectories enriched by sequences of partial rewards under some policy π . The undiscounted reward functionR : Γ→ R ( which expectation is to be maximized by optimizing θ ) satisfiesR ( τ ) = ∑H i=0 ri , where ri = R ( si+1 , ai , si ) . 3.1 BEHAVIORAL EMBEDDINGS We start with a Behavioral Embeddng Space ( BES ) which we denote as E and a Behavioral Embedding Map ( BEM ) , Φ : Γ → E , mapping trajectories to embeddings in E ( Fig . 1 ) . Importantly , the mapping does not need to be surjective . We will provide examples of BESs and BEMs at the end of the section . Given a policy π , we let Pπ denote the distribution induced over the spaces of trajectories Γ and by PΦπ the corresponding pushforward distribution on E induced by Φ . We call PΦπ the Probabilistic Policy Embedding ( PPE ) of a policy π . A policy π can be fully characterized by the distribution Pπ . Additionally , we require the BES E to be equipped with a metric ( or cost function ) C : E × E → R. Given two trajectories τ1 , τ2 in Γ , C ( Φ ( τ1 ) , Φ ( τ2 ) ) measures how different these trajectories are in the behavior space . The following are examples of BEMs ( with the corresponding BESs ) categorized into three main types ( we will use examples from all three types in our experiments in Section 7 ) : 1 . State-based : the final state Φ1 ( τ ) = sH , the visiting frequency of a fixed state Φs2 ( τ ) =∑H t=0 1 ( st = s ) , the frequency vector of visited states Φ3 ( τ ) = ∑H t=0 est ( where es ∈ R|S| is the one-hot vector corresponding to state s ) ; see also Section 7.2 . 2 . Action-based : the concatenation of actions Φ4 ( τ ) = [ a0 , ... , aH ] ; see also Section 7.1 . 3 . Reward-based : the total reward Φ5 ( τ ) = ∑H t=0 rt , reward-to-go vector Φ6 ( τ ) =∑H t=0 rt ( ∑t i=0 ei ) ( where ei ∈ RH+1 is a one-hot vector corresponding to i and with dimensions indexed from 0 to H ) ; see also Section 7.1 and Section 7.3 . For instance , PΦ3π is the frequency with which different states are visited under π . Note that some of the above embeddings are only for the tabular case ( |S| , |A| < ∞ ) while others are universal . 4 WASSERSTEIN DISTANCE & OPTIMAL TRANSPORT PROBLEM . Let µ , ν be ( Radon ) probability measures over domains X ⊆ Rm , Y ⊆ Rn and let C : X × Y → R be a cost function . For γ > 0 , a smoothed Wasserstein Distance is defined as : WDγ ( µ , ν ) : = min π∈Π ( µ , ν ) ∫ X×Y C ( x , y ) dπ ( x , y ) + γKL ( π|ξ ) , ( 1 ) where Π ( µ , ν ) is the space of couplings ( joint distributions ) over X ×Y with marginal distributions µ and ν , KL ( ·|· ) denotes the KL divergence between distributions π and ρ with support X × Y defined as : KL ( π|ρ ) = ∫ X×Y ( log ( dπ dξ ( x , y ) ) ) dπ ( x , y ) and ξ is a reference measure over X × Y . When the cost is an ` p distance and γ = 0 , WDγ is also known as the Earth mover ’ s distance and the corresponding optimization problem is known as the optimal transport problem ( OTP ) .

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