State Alignment-based Imitation Learning

1 INTRODUCTION . Learning from demonstrations ( imitation learning , abbr . as IL ) is a basic strategy to train agents for solving complicated tasks . Imitation learning methods can be generally divided into two categories : behavior cloning ( BC ) and inverse reinforcement learning ( IRL ) . Behavior cloning ( Ross et al. , 2011b ) formulates a supervised learning problem to learn a policy that maps states to actions using demonstration trajectories . Inverse reinforcement learning ( Russell , 1998 ; Ng et al. , 2000 ) tries to find a proper reward function that can induce the given demonstration trajectories . GAIL ( Ho & Ermon , 2016 ) and its variants ( Fu et al . ; Qureshi et al. , 2018 ; Xiao et al. , 2019 ) are the recently proposed IRL-based methods , which uses a GAN-based reward to align the distribution of stateaction pairs between the expert and the imitator . Although state-of-the-art BC and IRL methods have demonstrated compelling performance in standard imitation learning settings , e.g . control tasks ( Ho & Ermon , 2016 ; Fu et al . ; Qureshi et al. , 2018 ; Xiao et al. , 2019 ) and video games ( Aytar et al. , 2018b ) , these approaches are developed based on a strong assumption : the expert and the imitator share the same dynamics model ; specifically , they have the same action space , and any feasible state-action pair leads to the same next state in probability for both agents . The assumption brings severe limitation in practical scenarios : Imagine that a robot with a low speed limit navigates through a maze by imitating another robot which moves fast , then , it is impossible for the slow robot to execute the exact actions as the fast robot . However , the demonstration from the fast robot should still be useful because it shows the path to go through the maze . We are interested in the imitation learning problem under a relaxed assumption : Given an imitator that shares the same state space with the expert but their dynamics may be different , we train the imitator to follow the state sequence in expert demonstrations as much as possible . This is a more general formulation since it poses fewer requirements on the experts and makes demonstration collection easier . Due to the dynamics mismatch , the imitator becomes more likely to deviate from the demonstrations compared with the traditional imitation learning setting . Therefore , it is very important that the imitator should be able to resume to the demonstration trajectory by itself . Note that neither BC-based methods nor GAIL-based IRL methods have learned to handle dynamics misalignment and deviation correction . To address the issues , we propose a novel approach with four main features : 1 ) State-based . Compared to the majority of literature in imitation learning , our approach is state-based rather than action-based . Not like BC and IRL that essentially match state-action pairs between the expert and the imitator , we only match states . An inverse model of the imitator dynamics is learned to recover the action ; 2 ) Deviation Correction . A state-based β-VAE ( Higgins et al. , 2017 ) is learned as the prior for the next state to visit . Compared with ordinary behavior cloning , this VAE-based next state predictor can advise the imitator to return to the demonstration trajectory when it deviates . The robustness benefits from VAE ’ s latent stochastic sampling ; 3 ) Global State Alignment . While the VAE can help the agent to correct its trajectory to some extent , the agent may still occasionally enter states that are far away from demonstrations , where the VAE has no clue how to correct it . So we have to add a global constraint to align the states in demonstration and imitation . Inspired by GAIL that uses reward to align the distribution of state-action pairs , we also formulate an IRL problem whose maximal cumulative reward is the Wasserstein Distance between states of demonstration and imitation . Note that we choose not to involve state-action pairs as in GAIL ( Ho & Ermon , 2016 ) , or state-state pairs as in an observation-based GAIL ( Torabi et al. , 2018a ) , because our state-only formulation imposes weaker constraints than the two above options , thus providing more flexibility to handle different agent dynamics ; 4 ) Regularized Policy Update . We combine the prior for next state learned from VAE and the Wasserstein distance-based global constraint from IRL in a unified framework , by imposing a Kullback-Leibler divergence based regularizer to the policy update in the Proximal Policy Optimization algorithm . To empirically justify our ideas , we conduct experiments in two different settings . We first show that our approach can achieve similar or better results on the standard imitation learning setting , which assumes the same dynamics between the expert and the imitator . We then evaluate our approach in the more challenging setting that the dynamics of the expert and the imitator are different . In a number of control tasks , we either change the physics properties of the imitators or cripple them by changing their geometries . Existing approaches either fail or can only achieve very low rewards , but our approach can still exhibit decent performance . Finally , we show that even for imitation across agents of completely different actuators , it is still possible for the state-alignment based method to work . Surprisingly , a point mass and an ant in MuJoCo ( Todorov et al. , 2012 ) can imitate each other to navigate in a maze environment . Our contributions can be summarized as follows : • Propose to use a state alignment based method in the imitation learning problems where the expert ’ s and the imitator ’ s dynamics are different . • Propose a local state alignment method based on β-VAE and a global state alignment method based on Wasserstein distance . • Combine the local alignment and global alignment components into a reinforcement learn- ing framework by a regularized policy update objective . 2 RELATED WORK . Imitation learning is widely used in solving complicated tasks where pure reinforcement learning might suffer from high sample complexity , like robotics control ( Le et al. , 2017 ; Ye & Alterovitz , 2017 ; Pathak et al. , 2018 ) , autonomous vehicle ( Fu et al . ; Pomerleau , 1989 ) , and playing video game ( Hester et al. , 2018 ; Pohlen et al. , 2018 ; Aytar et al. , 2018a ) . Behavioral cloning ( Bain & Sommut , 1999 ) is a straight-forward method to learn a policy in a supervised way . However , behavioral cloning suffers from the problem of compounding errors as shown by ( Ross & Bagnell , 2010 ) , and this can be somewhat alleviated by interactive learning , such as DAGGER ( Ross et al. , 2011b ) . Another important line in imitation learning is inverse reinforcement learning ( Russell , 1998 ; Ng et al. , 2000 ; Abbeel & Ng , 2004 ; Ziebart et al. , 2008 ; Fu et al . ) , which finds a cost function under which the expert is uniquely optimal . Since IRL can be connected to min-max formulations , works like GAIL , SAM ( Ho & Ermon , 2016 ; Blondé & Kalousis , 2018 ) utilize this to directly recover policies . Its connections with GANs ( Goodfellow et al. , 2014 ) also lead to f -divergence minimization ( Ke et al. , 2019 ; Nowozin et al. , 2016 ) and Wasserstein distance minimization ( Xiao et al. , 2019 ) . One can also extend the framework from matching state-action pairs to state distribution matching , such as Torabi et al . ( 2018a ) ; Sun et al . ( 2019 ) ; Schroecker & Isbell ( 2017 ) . Other works ( Aytar et al. , 2018b ; Liu et al. , 2018 ; Peng et al. , 2018 ) also learn from observation alone , by defining reward on state and using IRL to solve the tasks . Works like ( Lee et al. , 2019 ; Lee et al . ) also use state-based reward for exploration . Torabi et al . ( 2018b ) ; Edwards et al . ( 2018 ) will recover actions from observations by learning an inverse model or latent actions . However , our work aims to combine the advantage of global state distribution matching and local state transition alignment , which combines the advantage of BC and IRL through a novel framework . 3 BACKGROUNDS . Variational Autoencoders Kingma & Welling ( 2013 ) ; Rezende et al . ( 2014 ) provides a framework to learn both a probabilistic generative model pθ ( x|z ) as well as an approximated posterior distribution qφ ( z|x ) . β-VAE is a variant VAE that introduces an adjustable hyperparameter β to the original objective : L ( θ , φ ; x , z , β ) = Eqφ ( z|x ) [ log pθ ( x|z ) ] − βDKL ( qφ ( z|x ) ‖p ( z ) ) ( 1 ) Larger β will penalize the total correlation ( Chen et al. , 2018 ) to encourage more disentangled latent representations , while smaller β often results in sharper and more precise reconstructions . Wasserstein distance The Wasserstein distance between two density functions p ( x ) and q ( x ) with support on a compact metric space ( M , d ) has an alternative form due to Kantorovich-Rubenstein duality ( Villani , 2008 ) : W ( p , q ) = sup φ∈L1 Ep ( x ) [ φ ( x ) ] − Eq ( x ) [ φ ( x ) ] ( 2 ) Here , L1 is the set of all 1-Lipschitz functions from M to R. Compared with the prevalent KLdivergence and its extension , the f-divergence family , Wasserstein distance has a number of advantages theoretically and numerically . Please refer to Arjovsky et al . ( 2017 ) and Solomon ( 2018 ) for a detailed discussion . 4 SAIL : STATE ALIGNMENT BASED IMITATION LEARNING 4.1 OVERVIEW Our imitation learning method is based on state alignment from both local and global perspectives . For local alignment , the goal is to follow the transition of the demonstration as much as possible , and allow the return to the demonstration trajectory whenever the imitation deviates . To achieve both goals , we use a βVAE ( Higgins et al. , 2017 ) to generate the next state ( Figure 2 Left ) . For global alignment , we set up an objective to minimize the Wasserstein distance between the states in the current trajectory and the demonstrations ( Figure 2 Right ) . There has to be a framework to naturally combine the local alignment and global alignment components . We resort to the reinforcement learning framework by encoding the local alignment as policy prior and encoding the global alignment as reward over states . Using Proximal Policy Optimization ( PPO ) by Schulman et al . ( 2017 ) as the backbone RL solver , we derive a regularized policy update . To maximally exploit the knowledge from demonstrations and reduce interactions with the environment , we adopt a pre-training stage to produce a good initialization based on the same policy prior induced by the local alignment . Our method is summarized in Algorithm 1 . In the rest parts of this section , we will introduce all the components of our method in details .

Summary: This paper seeks a solution to the problem of performing imitation learning when the dynamics of the demonstrator are different from the dynamics of the imitator. The authors present a novel approach that combines global alignment by minimizing the Wasserstein distance between state occupancies with local alignment via a state-predictive VAE and inverse dynamics model. The experimental results support the claims that the method works for different dynamics and the proposed approach usually outperforms existing imitation learning methods.