Primal Wasserstein Imitation Learning

1 INTRODUCTION . Reinforcement Learning ( RL ) has solved a number of difficult tasks whether in games ( Tesauro , 1995 ; Mnih et al. , 2015 ; Silver et al. , 2016 ) or robotics ( Abbeel & Ng , 2004 ; Andrychowicz et al. , 2020 ) . However , RL relies on the existence of a reward function , that can be either hard to specify or too sparse to be used in practice . Imitation Learning ( IL ) is a paradigm that applies to these environments with hard to specify rewards : we seek to solve a task by learning a policy from a fixed number of demonstrations generated by an expert . IL methods can typically be folded into two paradigms : Behavioral Cloning , or BC ( Pomerleau , 1991 ; Bagnell et al. , 2007 ; Ross & Bagnell , 2010 ) and Inverse Reinforcement Learning , or IRL ( Russell , 1998 ; Ng et al. , 2000 ) . In BC , we seek to recover the expert ’ s behavior by directly learning a policy that matches the expert behavior in some sense . In IRL , we assume that the demonstrations come from an agent that acts optimally with respect to an unknown reward function that we seek to recover , to subsequently train an agent on it . Although IRL methods introduce an intermediary problem ( i.e . recovering the environment ’ s reward ) they are less sensitive to distributional shift ( Pomerleau , 1991 ) , they generalize to environments with different dynamics ( Piot et al. , 2013 ) , and they can recover a near-optimal agent from suboptimal demonstrations ( Brown et al. , 2019 ; Jacq et al. , 2019 ) . However , IRL methods are usually based on an iterative process alternating between reward estimation and RL , which might result in poor sample-efficiency . Earlier IRL methods ( Ng et al. , 2000 ; Abbeel & Ng , 2004 ; Ziebart et al. , 2008 ) require multiple calls to a Markov decision process solver ( Puterman , 2014 ) , whereas recent adversarial IL approaches ( Finn et al. , 2016 ; Ho & Ermon , 2016 ; Fu et al. , 2018 ) interleave the learning of the reward function with the learning process of the agent . Adversarial IL methods are based on an adversarial training paradigm similar to Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) , where the learned reward function can be thought of as the confusion of a discriminator that learns to differentiate expert transitions from non expert ones . These methods are well suited to the IL problem since they implicitly minimize an f -divergence between the state-action distribution of an expert and the state-action distribution of the learning agent ( Ghasemipour et al. , 2019 ; Ke et al. , 2019 ) . However the interaction between a generator ( the policy ) and the discriminator ( the reward function ) makes it a minmax optimization problem , and therefore comes with practical challenges that might include training instability , sensitivity to hyperparameters and poor sample efficiency . ∗Correspondence to Robert Dadashi : dadashi @ google.com . In this work , we use the Wasserstein distance as a measure between the state-action distributions of the expert and of the agent . Contrary to f -divergences , the Wasserstein distance is a true distance , it is smooth and it is based on the geometry of the metric space it operates on . The Wasserstein distance has gained popularity in GAN approaches ( Arjovsky et al. , 2017 ) through its dual formulation which comes with challenges ( see Section 5 ) . Our approach is novel in the fact that we consider the problem of minimizing the Wasserstein distance through its primal formulation . Crucially , the primal formulation prevents the minmax optimization problem , and requires little fine tuning . We introduce a reward function computed offline based on an upper bound of the primal form of the Wasserstein distance . As the Wasserstein distance requires a distance between state-action pairs , we show that it can be hand-defined for locomotion tasks , and that it can be learned from pixels for a hand manipulation task . The inferred reward function is non-stationary , like adversarial IL methods , but it is not re-evaluated as the agent interacts with the environment , therefore the reward function we define is computed offline . We present a true distance to compare the behavior of the expert and the behavior of the agent , rather than using the common proxy of performance with respect to the true return of the task we consider ( as it is unknown in general ) . Our method recovers expert behaviour comparably to existing state-of-the-art methods while being based on significantly fewer hyperparameters ; it operates even in the extreme low data regime of demonstrations , and is the first method that makes Humanoid run with a single ( subsampled ) demonstration . 2 BACKGROUND AND NOTATIONS . Markov decision processes . We describe environments as episodic Markov Decision Processes ( MDP ) with finite time horizon ( Sutton & Barto , 2018 ) ( S , A , P , r , γ , ρ0 , T ) , where S is the state space , A is the action space , P is the transition kernel , r is the reward function , γ is the discount factor , ρ0 is the initial state distribution and T is the time horizon . We will denote the dimensionality of S and A as |S| and |A| respectively . A policy π is a mapping from states to distributions over actions ; we denote the space of all policies by Π . In RL , the goal is to learn a policy π∗ that maximizes the expected sum of discounted rewards it encounters , that is , the expected return . Depending on the context , we might use the concept of a cost c rather than a reward r ( Puterman , 2014 ) , which essentially moves the goal of the policy from maximizing its return to minimizing its cumlative cost . State action distributions . Suppose a policy π visits the successive states and actions s1 , a1 , . . . , sT , aT during an episode , we define the empirical state-action distribution ρ̂π as : ρ̂π = 1 T ∑T t=1 δst , at , where δst , at is a Dirac distribution centered on ( st , at ) . Similarly , suppose we have a set of expert demonstrations D = { se , ae } of size D , then the associated empirical expert state-action distribution ρ̂e is defined as : ρ̂e = 1D ∑ ( s , a ) ∈D δs , a . Wasserstein distance . Suppose we have the metric space ( M , d ) where M is a set and d is a metric on M . Suppose we have µ and ν two distributions on M with finite moments , the p-th order Wasserstein distance ( Villani , 2008 ) is defined asWpp ( µ , ν ) = infθ∈Θ ( µ , ν ) ∫ M×M d ( x , y ) pdθ ( x , y ) , where Θ ( µ , ν ) is the set of all couplings between µ and ν . In the remainder , we only consider distributions with finite support . A coupling between two distributions of support cardinal T and D is a doubly stochastic matrix of size T ×D . We note Θ the set of all doubly stochastic matrices of size T ×D : Θ = { θ ∈ RT×D+ | ∀j ∈ [ 1 : D ] , T∑ i′=1 θ [ i′ , j ] = 1 D , ∀i ∈ [ 1 : T ] , D∑ j′=1 θ [ i , j′ ] = 1 T } . The Wasserstein distance between distributions of state-action pairs requires the definition of a metric d in the space ( S , A ) . Defining a metric in an MDP is non trivial ( Ferns et al. , 2004 ; Mahadevan & Maggioni , 2007 ) ; we show an example where the metric is learned from demonstrations in Section 4.4 . For now , we assume the existence of a metric d : ( S , A ) × ( S , A ) 7→ R+ . 3 METHOD . We present the theoretical motivation of our approach : the minimization of the Wasserstein distance between the state-action distributions of the agent and the expert . We introduce a reward based on an upper-bound of the primal form of the Wasserstein distance inferred from a relaxation of the optimal coupling condition , and present the resulting algorithm : Primal Wasserstein Imitation Learning ( PWIL ) . 3.1 WASSERSTEIN DISTANCE MINIMIZATION . Central to our approach is the minimization of the Wasserstein distance between the state-action distribution of the policy we seek to train ρ̂π and the state-action distribution of the expert ρ̂e . In other words , we aim at optimizing the following problem : inf π∈Π Wpp ( ρ̂π , ρ̂e ) = inf π∈Π inf θ∈Θ T∑ i=1 D∑ j=1 d ( ( sπi , a π i ) , ( s e j , a e j ) ) pθ [ i , j ] . ( 1 ) In the rest of the paper , we only consider the 1-Wasserstein ( p = 1 in Equation 1 ) and leave the extensive study of the influence of the order p for future work . We can interpret the Wasserstein distance using the earth ’ s movers analogy ( Villani , 2008 ) . Consider that the state-action pairs of the expert are D holes of mass D−1 and that the state-action pairs of the policy are piles of dirt of mass T−1 . A coupling θ is a transport strategy between the piles of dirt and the holes , where θ [ i , j ] stands for how much of the pile of dirt i should be moved towards the hole j . The optimal coupling is the one that minimizes the distance that the earth mover travels to put all piles of dirt to holes . Note that to compute the optimal coupling , we need knowledge of the locations of all piles of dirt . In the context of RL , this means having access to the full trajectory generated by π . From now on , we write θ∗π as the optimal coupling for the policy π , that we inject in Equation ( 1 ) : θ∗π = arg min θ∈Θ T∑ i=1 D∑ j=1 d ( ( sπi , a π i ) , ( s e j , a e j ) ) θ [ i , j ] inf π∈Π W1 ( ρ̂π , ρ̂e ) = inf π∈Π T∑ i=1 c∗i , π with c ∗ i , π = D∑ j=1 d ( ( sπi , a π i ) , ( s e j , a e j ) ) θ ∗ π [ i , j ] . ( 2 ) In Equation ( 2 ) , we have introduced c∗i , π , which we interpret as a cost to minimize using RL . As c∗i , π depends on the optimal coupling θ ∗ π , we can only define c ∗ i , π at the very end of an episode . This can be problematic if an agent learns in an online manner or in large time-horizon tasks . Thus , we introduce an upper bound to the Wasserstein distance that yields a cost we can compute online , based on a suboptimal coupling strategy .

Summary: This paper proposes to use Wasserstein distance in the primal form for imitation learning. Compared with its dual form and f-divergence minimization variants, it avoids the unstable minimax optimization. In order to compute the Wasserstein distance in primal form, they also propose a greedy approximation. Their experiments demonstrate that this method has a better performance compared with baseline methods.