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: The paper proposes an imitation learning method that aims to align state distributions rather than state-action distributions to account for cases where the imitator dynamics differ from expert dynamics. They achieve this by two objectives: one local, the other global. The local objective aligns the next state to be close to the expert's next state in each transition by first training a VAE on the expert demonstrations, and using the trained VAE in conjunction with a pretrained inverse dynamics model to compute the action that the imitator needs to imitate. The global objective tries to do a global alignment of states encountered in the imitator and expert trajectories, by minimizing the Wasserstein distance between the two trajectory distributions. The paper claims that using these two objectives results in a method that outperforms existing inverse reinforcement learning and behavior cloning approaches in settings where the imitator and expert dynamics differ.