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 proposes an approach to the problem of state-action based imitation learning, in which an agent aims to solve a specific task given state-action expert trajectories. Their objective SLIL combines (or interpolates between) generative modeling and behavioral cloning in order to avoid mode collapse of the agent policy (common issue of adversarial approaches to IL like GAIL) and to mitigate distributional shift which BC typically suffers from.   Contributions: - An extension of continuous normalizing flows, coined denoising continuous normalising flows (DCNF), which is claimed to alleviate the manifold hypothesis, and is trained via maximum likelihood estimation. - A strong imitation learning algorithm consisting in estimating the expert's state distribution, and learning a policy that recovers the expert policy and that has an occupancy distribution close to the distribution induced by generative model learned via the DCNF.