Stabilized Likelihood-based Imitation Learning via Denoising Continuous Normalizing Flow

1 INTRODUCTION . Imitation learning ( IL ) ( Abbeel & Ng , 2004 ; Ho & Ermon , 2016 ) aims to learn sequential decisionmaking policies directly from expert demonstrations , without access to reward signals from the environment . State-of-the-art ( SOTA ) IL approaches have primarily followed one of two paradigms : behavior cloning ( BC ) ( Pomerleau , 1991 ) and generative adversarial imitation learning ( GAIL ) ( Ho & Ermon , 2016 ) . While both BC and GAIL have been studied extensively , each has crucial limitations . On the one hand , BC approaches employ supervised learning , which requires a large amount of expert demonstrations to avoid compounding errors due to covariate shifts ( Ross et al. , 2011 ; Ross & Bagnell , 2010 ) . On the other hand , GAIL approaches ( Jena & Sycara , 2020 ; Fei et al. , 2020 ; Li et al. , 2017 ; Hausman et al. , 2017 ; Fu et al. , 2017 ; Ke et al. , 2019 ; Ghasemipour et al. , 2019 ; Zhang et al. , 2020 ; Arjovsky et al. , 2017 ) connect IL with generative adversarial networks ( GAN ) ( Goodfellow et al. , 2014 ) , but adversarial training processes are intrinsically unstable ( Jena & Sycara , 2020 ; Arjovsky et al. , 2017 ) and prone to mode collapse ( especially when learning from multi-mode expert demonstrations ) ( Ghasemipour et al. , 2019 ; Fei et al. , 2020 ; Ke et al. , 2019 ) . To see this , consider Fig . 1 , which shows an example of different IL algorithms on the Reacher task ( Todorov et al. , 2012 ; Hausman et al. , 2017 ) with four targets ( Fig . 1a ) . While the expert tends to visit all targets equally ( Fig . 1b ) , the policy learned by GAIL is mode collapsed , primarily visiting the green one ( Fig . 1e ) . SLIL : Overview and Contributions . In this work , we are motivated to design an IL methodology with a stabilized training process and mode distribution preservation from expert demonstrations . We propose SLIL – Stabilized Likelihood-based Imitation Learning – a novel IL approach that directly maximizes the likelihood of observing the expert demonstrations . SLIL is a stable twostage optimization framework , with the 1st stage focusing on accurately estimating the expert state distribution , and the 2nd stage training the learner policy to match the expert ’ s policy and state distribution . Fig . 1c shows a quick view of our results : SLIL leads to a learner policy preserving the mode visitation distribution of the expert demonstrations . Our key contributions are as follows : • We develop the first likelihood-based IL methodology that tackles the training instability issue of GAIL . SLIL ’ s two-stage optimization framework is based on a tight lower bound on the joint policy and state distribution training objective from the likelihood-based IL formulation ( Sec 4.1 ) . • We propose the denoising continuous normalizing flow ( DCNF ) algorithm to accurately estimate the expert state distribution in the 1st stage of SLIL to preserve expert modes . DCNF maps the expert state to the Gaussian noise distribution with a neural ODE , and trains the ODE using perturbed expert states ( Sec 4.2 ) . • Our evaluation on ten different physics-based control tasks reveals that SLIL obtains superior results compared with SOTA baselines in terms of learner policy performances , training stability , and mode distribution preservation ( Sec 5 ) .1 2 RELATED WORK . Behavior Cloning ( BC ) ( Pomerleau , 1991 ; Bohg et al. , 2020 ) approaches for IL learn the expert policy via maximizing expert trajectory likelihood in the demonstration data , i.e , maximizing expert action likelihood . Though effective with abundant demonstrations , BC suffers from the covariate shift problem with limited data ( Ross et al. , 2011 ; Ross & Bagnell , 2010 ) . Brantley et al . ( 2019 ) aims to address this via first pre-training an ensemble of BC policies , and then using reinforcement learning ( RL ) ( Sutton & Barto , 2018 ) to train a learner policy whose cost is proportional to ensemble policies ’ prediction variance . In SLIL , we instead strive to solve the covariate shift problem via expert state-action distribution matching , learning the expert state distribution explicitly . Generative Adversarial Imitation Learning ( GAIL ) ( Ho & Ermon , 2016 ) employs the GAN framework ( Goodfellow et al. , 2014 ) to minimize the discrepancy between expert and learner stateaction distributions . It jointly trains a generator ( i.e , learner policy ) to imitate the expert behaviors , and a discriminator ( i.e , reward signal ) to distinguish the state-action pair distributions between the expert and the learner . In GAIL , the discrepancy between the behavior distributions of the expert and the learner is measured by JS divergence . Using the variational lower bound of an f -divergence measure , several studies ( Ke et al. , 2019 ; Ghasemipour et al. , 2019 ; Fu et al. , 2017 ; Nowozin et al. , 2016 ; Arumugam et al. , 2019 ; Zhang et al. , 2020 ) have extended GAIL from JS to a pre-defined f - divergence , e.g , KL ( Fu et al. , 2017 ) , Reverse KL ( Ke et al. , 2019 ) , Total Variation ( Ross et al. , 2011 ) . However , all these approaches adversarially train the generator and the discriminator as a minimax game to reach an equilibrium , which can lead to training instability and mode collapse ( Schroecker et al. , 2019 ; Ghasemipour et al. , 2019 ; Fei et al. , 2020 ; Ke et al. , 2019 ) . 1For reproducibility , the code for our experiments is available at https : //www.dropbox.com/sh/ buhqre6hfwdmvwu/AAAkc7SSjqhmtfalqqj5I7Iba ? dl=0 . Several recent works try to avoid adversarial training in IL ( Dadashi et al. , 2020 ; Liu et al. , 2020 ; Kim et al. , 2020b ; Rhinehart et al. , 2018 ; Schroecker et al. , 2019 ) . To do so , Primal Wasserstein IL ( PWIL ) ( Dadashi et al. , 2020 ) considers the primal form of Wasserstein distance to match learner ’ s and expert ’ s state-action distributions . Neural Density Imitation ( NDI ) ( Kim et al. , 2020b ) estimates expert ’ s occupancy measure using which as a reward for maximum occupancy entropy reinforcement learning . Energy-Based IL ( EBIL ) ( Liu et al. , 2020 ) stems from Max-Entropy IRL ( Ziebart et al. , 2008 ) and estimates a surrogate reward function with score matching from expert demonstrations . Imitative Models ( IM ) ( Rhinehart et al. , 2018 ) learns a flow model that assigns high likelihoods to expert-like trajectories for test time goal-directed planning . GPRIL ( Schroecker et al. , 2019 ) applies masked autoregressive flows ( Papamakarios et al. , 2017 ) to learn predecessor state-action distribution in each training iteration , which adds complexity to learner policy training and requires training stabilization . Unlike these methods , we maximize expert state-action likelihood and strive to not only avoid adversarial training , but solve the mode collapse and unstable training problems as well . 3 LIL PROBLEM FORMULATION Notations . We denote S as a set of states , A as a set of actions , P : S⇥ A⇥S 7 ! [ 0 , 1 ] as the transition probability distribution , r : S⇥A 7 ! R as the reward function , ⇢0 : S 7 ! R as the distribution of the initial state s0 , and 2 [ 0 , 1 ] as the discount factor . An agent makes decisions following a policy ⇡ : S ⇥ A 7 ! [ 0 , 1 ] , which specifies a probability distribution of choosing an action a 2 A at a state s 2 S. With s0 ⇠ ⇢0 , then over time t , at ⇠ ⇡ ( at|st ) and st+1 ⇠ P ( st+1|st , at ) according to the policy ⇡ . We will denote the expert policy as ⇡E , and the learner policy as ⇡ . We denote P⇡ ( s , a ) as the probability of observing a state-action pair ( s , a ) when executing the learner policy ⇡ , and denote PE ( s , a ) to represent P⇡E ( s , a ) for brevity . Moreover , for a function h ( s , a ) of interest , we use an expectation over a policy ⇡ to denote an expectation with respect to the trajectories it generates , i.e , E⇡ [ h ( s , a ) ] , E [ P1 t=0 th ( st , at ) ] . Likelihood-based Imitation Learning ( LIL ) . To avoid the instability from adversarial training ( e.g , in GAIL ) , we introduce a likelihoodbased imitation learning paradigm to learn the policy ⇡ by directly maximizing the likelihood of the state-action pairs from the expert demonstration data DE . The LIL objective can be formally modeled as a bilevel optimization problem ( Sinha et al. , 2017 ) : max ⇡ E ( s , a ) 2DE [ logP⇡ ( s , a ) ] , s.t . P⇡ = argmax P E ( s , a ) 2D⇡ [ logP ( s , a ) ] , ( 1 ) with P⇡ ( s , a ) as the probability of observing an expert state-action pair ( s , a ) 2 DE when executing the learner policy ⇡ . Notice that instead of maximizing trajectory likelihood as in BC , LIL maximizes the likelihood of observed state-action pairs . The variable function P⇡ ( s , a ) can be learned from the lower-level optimization to maximize the log-likelihood of observing state-action pairs ( s , a ) 2 D⇡ collected by interacting with the environment using the learner policy ⇡ ( i.e , the upper-level variable function ) . Moreover , since P⇡ ( s , a ) = ⇡ ( a|s ) P⇡ ( s ) , with ⇡ ( a|s ) as the learner policy and P⇡ ( s ) as the state distribution when executing ⇡ , the LIL objective can be rewritten as max ⇡ E ( s , a ) 2DE [ log ⇡ ( a|s ) ] + Es2DE [ logP⇡ ( s ) ] , s.t . P⇡ = argmaxP Es2D⇡ [ logP ( s ) ] . ( 2 ) Challenges . The LIL solution framework alternates between policy learning ( i.e , the upper-level optimization to update ⇡ ) and state distribution estimation ( i.e , the lower-level optimization to learn P⇡ ) , as shown in Fig . 2 . In each iteration , the learner policy ⇡ is used to interact with the environment to collect D⇡ . A normalizing flow is trained to estimate the state distribution P⇡ of policy ⇡ using D⇡ ( i.e , lower-level optimization ) . Then , the policy ⇡ is updated based on the state distribution P⇡ using expert data DE ( i.e , upper-level optimization ) . However , two key challenges prevent a straightforward application of LIL . C # 1 Training stability : It is challenging to properly grow ⇡ and P⇡ jointly to speed up and stabilize the training process in the bilevel optimization framework of Eq. ( 2 ) . C # 2 Manifold hypothesis : Normalizing flows are only valid when the data distribution and the target noise distribution have the same dimensions , which conflicts with the observation that real-world data usually concentrate on low dimensional manifolds in a high dimensional space ( a.k.a . the ambient space ) ( Papamakarios et al. , 2017 ; Kim et al. , 2020a ; Song & Ermon , 2019 ; Ho et al. , 2020 ; Belkin & Niyogi , 2003 ) . Given an expert state distribution supported on a low dimensional manifold , the LIL framework will fail to estimate the ground-truth state distribution P⇡ and fail to learn an expert-like learner policy ⇡ efficiently . To address these challenges , we propose the Stabilized Likelihood-based Imitation Learning ( SLIL ) framework shown in Fig . 3 . SLIL relaxes the LIL problem to a two-stage optimization problem for improved training stability , and learns expert state distribution with denoising continuous normalizing flow ( DCNF ) to overcome the manifold hypothesis challenge . We formalize this methodology next .

Summary: The paper present a novel method for imitation learning aiming to stabilize training and mode coverage compared to existing methods. The proposed method (SLIL) extends behavioral cloning by an additional objective of maximizing a state-only reward that is defined by a density estimate (trained once beforehand) of the expert's state distribution. For density estimation SLIL uses a continuous normalizing flow where the expert states are corrupted by Gaussian noise, where the noise-level is decreased during training. The noise is introduced to tackle the "Manifold Hypothesis Challenge", that is, the inability of normalizing flows to accurate approximate distributions that lie on a lower-dimensional manifold.   The main contributions are: 1. A novel imitation learning method that combines the BC loss with a RL loss for reaching expert states 2. A slightly modified approach to tackle the manifold hypothesis challenge (Kim et al. [2020] used Gaussian noise where the noise-level was sampled from a uniform distribution, independently for each data set, rather than using the same noise-level for all data and annealing it during optimization)