Imitation with Neural Density Models

1 Introduction . Imitation Learning ( IL ) algorithms aim to learn optimal behavior by mimicking expert demonstrations . Perhaps the simplest IL method is Behavioral Cloning ( BC ) ( Pomerleau , 1991 ) which ignores the dynamics of the underlying Markov Decision Process ( MDP ) that generated the demonstrations , and treats IL as a supervised learning problem of predicting optimal actions given states . Prior work showed that if the learned policy incurs a small BC loss , the worst case performance gap between the expert and imitator grows quadratically with the number of decision steps ( Ross & Bagnell , 2010 ; Ross et al. , 2011a ) . The crux of their argument is that policies that are `` close '' as measured by BC loss can induce disastrously different distributions over states when deployed in the environment . One family of solutions to mitigating such compounding errors is Interactive IL ( Guo et al. , 2014 ; Ross et al. , 2011b , 2013 ) , which involves running the imitator ’ s policy and collecting corrective actions from an interactive expert . However , interactive expert queries are expensive and seldom available . Another family of approaches ( Fu et al. , 2017 ; Ho & Ermon , 2016 ; Ke et al. , 2020 ; Kim & Park , 2018 ; Kostrikov et al. , 2020 ; Wang et al. , 2017 ) that have gained much traction is to directly minimize a statistical distance between state-action distributions induced by policies of the expert and imitator , i.e the occupancy measures ⇢⇡E and ⇢⇡✓ . As ⇢⇡✓ is an implicit distribution induced by the policy and environment1 , distribution matching with ⇢⇡✓ typically requires likelihood-free methods involving sampling . Sampling from ⇢⇡✓ entails running the imitator policy in the environment , which was not required by BC . While distribution matching IL requires additional access to an environment simulator , it has been shown to drastically improve demonstration efficiency , i.e the number of demonstrations needed to succeed at IL ( Ho & Ermon , 2016 ) . A wide suite of distribution matching IL algorithms use adversarial methods to match ⇢⇡✓ and ⇢⇡E , which requires alternating between reward ( discriminator ) and policy ( generator ) updates ( Fu et al. , 2017 ; Ho & Ermon , 2016 ; Ke et al. , 2020 ; Kim et al. , 2019 ; Kostrikov et al. , 2020 ) . A key drawback to such Adversarial Imitation Learning ( AIL ) methods is that they inherit the instability of alternating min-max optimization ( Miyato et al. , 2018 ; Salimans et al. , 2016 ) which is generally not guaranteed to converge ( Jin et al. , 2019 ) . Furthermore , this instability is exacerbated in the IL setting where generator updates involve high-variance policy optimization and leads to sub-optimal demonstration efficiency . To alleviate this instability , ( Brantley et al. , 2020 ; Reddy et al. , 2017 ; Wang et al. , 2019 ) have proposed to do RL with fixed heuristic rewards . Wang et al . ( 2019 ) , for example , uses a heuristic reward that estimates the 1we assume only samples can be taken from the environment dynamics and its density is unknown 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . support of ⇢⇡E which discourages the imitator from visiting out-of-support states . While having the merit of simplicity , these approaches have no guarantee of recovering the true expert policy . In this work , we propose a new framework for IL via obtaining a density estimate q of the expert ’ s occupancy measure ⇢⇡E followed by Maximum Occupancy Entropy Reinforcement Learning ( MaxOccEntRL ) ( Islam et al. , 2019 ; Lee et al. , 2019 ) . In the MaxOccEntRL step , the density estimate q is used as a fixed reward for RL and the occupancy entropy H ( ⇢⇡✓ ) is simultaneously maximized , leading to the objective max✓ E⇢⇡✓ [ log q ( s , a ) ] +H ( ⇢⇡✓ ) . Intuitively , our approach encourages the imitator to visit high density state-action pairs under ⇢⇡E while maximally exploring the state-action space . There are two main challenges to this approach . First , we require accurate density estimation of ⇢⇡E , which is particularly challenging when the state-action space is high dimensional and the number of expert demonstrations are limited . Second , in contrast to Maximum Entropy RL ( MaxEntRL ) , MaxOccEntRL requires maximizing the entropy of an implicit density ⇢⇡✓ . We address the former challenge leveraging advances in density estimation ( Du & Mordatch , 2018 ; Germain et al. , 2015 ; Song et al. , 2019 ) . For the latter challenge , we derive a non-adversarial model-free RL objective that provably maximizes a lower bound to occupancy entropy . As a byproduct , we also obtain a model-free RL objective that lower bounds reverse Kullback-Lieber ( KL ) divergence between ⇢⇡✓ and ⇢⇡E . The contribution of our work is introducing a novel family of distribution matching IL algorithms , named Neural Density Imitation ( NDI ) , that ( 1 ) optimizes a principled lower bound to the additive inverse of reverse KL , thereby avoiding adversarial optimization and ( 2 ) . advances state-of-the-art demonstration efficiency in IL . 2 Imitation Learning via density estimation . We model an agent ’ s decision making process as a discounted infinite-horizon Markov Decision Process ( MDP ) M = ( S , A , P , P0 , r , ) . Here S , A are state-action spaces , P : S ⇥ A ! ⌦ ( S ) is a transition dynamics where ⌦ ( S ) is the set of probability measures on S , P0 : S ! R is an initial state distribution , r : S ⇥ A ! R is a reward function , and 2 [ 0 , 1 ) is a discount factor . A parameterized policy ⇡✓ : S ! ⌦ ( A ) distills the agent ’ s decision making rule and { st , at } 1t=0 is the stochastic process realized by sampling an initial state from s0 ⇠ P0 ( s ) then running ⇡✓ in the environment , i.e at ⇠ ⇡✓ ( ·|st ) , st+1 ⇠ P ( ·|st , at ) . We denote by p✓ , t : t+k the joint distribution of states { st , st+1 , ... , st+k } , where setting p✓ , t recovers the marginal of st . The ( unnormalized ) occupancy measure of ⇡✓ is defined as ⇢⇡✓ ( s , a ) = P1 t=0 t p✓ , t ( s ) ⇡✓ ( a|s ) . Intuitively , ⇢⇡✓ ( s , a ) quantifies the frequency of visiting the state-action pair ( s , a ) when running ⇡✓ for a long time , with more emphasis on earlier states . We denote policy performance as J ( ⇡✓ , r̄ ) = E⇡✓ [ P1 t=0 t r̄ ( st , at ) ] = E ( s , a ) ⇠⇢⇡✓ [ r̄ ( s , a ) ] where r̄ is a ( potentially ) augmented reward function and E denotes the generalized expectation operator extended to non-normalized densities p̂ : X ! R+ and functions f : X ! Y so that Ep̂ [ f ( x ) ] = P x p̂ ( x ) f ( x ) . The choice of r̄ depends on the RL framework . In standard RL , we simply have r̄ = r , while in Maximum Entropy RL ( MaxEntRL ) ( Haarnoja et al. , 2017 ) , we have r̄ ( s , a ) = r ( s , a ) log ⇡✓ ( a|s ) . We denote the entropy of ⇢⇡✓ ( s , a ) as H ( ⇢⇡✓ ) = E⇢⇡✓ [ log ⇢⇡✓ ( s , a ) ] and overload notation to denote the -discounted causal entropy of policy ⇡✓ as H ( ⇡✓ ) = E⇡✓ [ P1 t=0 t log ⇡✓ ( at|st ) ] = E⇢⇡✓ [ log ⇡✓ ( a|s ) ] . Note that we use a generalized notion of entropy where the domain is extended to non-normalized densities . We can then define the Maximum Occupancy Entropy RL ( MaxOccEntRL ) ( Islam et al. , 2019 ; Lee et al. , 2019 ) objective as J ( ⇡✓ , r̄ = r ) +H ( ⇢⇡✓ ) . Note the key difference between MaxOccEntRL and MaxEntRL : entropy regularization is on the occupancy measure instead of the policy , i.e seeks state diversity instead of action diversity . We will later show in section 2.2 , that a lower bound on this objective reduces to a complete model-free RL objective with an augmented reward r̄ . Let ⇡E , ⇡✓ denote an expert and imitator policy , respectively . Given only demonstrations D = { ( s , a ) i } ki=1 ⇠ ⇡E of state-action pairs sampled from the expert , Imitation Learning ( IL ) aims to learn a policy ⇡✓ which matches the expert , i.e ⇡✓ = ⇡E . Formally , IL can be recast as a distribution matching problem ( Ho & Ermon , 2016 ; Ke et al. , 2020 ) between occupancy measures ⇢⇡✓ and ⇢⇡E : maximize✓ d ( ⇢⇡✓ , ⇢⇡E ) ( 1 ) where d ( p̂ , q̂ ) is a generalized statistical distance defined on the extended domain of ( potentially ) non-normalized probability densities p̂ ( x ) , q̂ ( x ) with the same normalization factor Z > 0 , i.eR x p̂ ( x ) /Z = R x q̂ ( x ) /Z = 1 . For ⇢⇡ and ⇢⇡E , we have Z = 1 1 . As we are only able to take samples from the transition kernel and its density is unknown , ⇢⇡✓ is an implicit distribution2 . Thus , optimizing Eq . 1 typically requires likelihood-free approaches leveraging samples from ⇢⇡✓ , i.e running ⇡✓ in the environment . Current state-of-the-art IL approaches use likelihood-free adversarial methods to approximately optimize Eq . 1 for various choices of d such as reverse Kullback-Liebler ( KL ) divergence ( Fu et al. , 2017 ; Kostrikov et al. , 2020 ) and Jensen-Shannon ( JS ) divergence ( Ho & Ermon , 2016 ) . However , adversarial methods are known to suffer from optimization instability which is exacerbated in the IL setting where one step in the alternating optimization involves RL . We instead derive a non-adversarial objective for IL . In this work , we choose d to be ( generalized ) reverse-KL divergence and leave derivations for alternate choices of d to future work . DKL ( ⇢⇡✓ ||⇢⇡E ) = E⇢⇡✓ [ log ⇢⇡E ( s , a ) log ⇢⇡✓ ( s , a ) ] = J ( ⇡✓ , r̄ = log ⇢⇡E ) +H ( ⇢⇡✓ ) ( 2 ) We see that maximizing negative reverse-KL with respect to ⇡✓ is equivalent to Maximum Occupancy Entropy RL ( MaxOccEntRL ) with log ⇢⇡E as the fixed reward . Intuitively , this objective drives ⇡✓ to visit states that are most likely under ⇢⇡E while maximally spreading out probability mass so that if two state-action pairs are equally likely , the policy visits both . There are two main challenges associated with this approach which we address in the following sections . 1. log ⇢⇡E is unknown and must be estimated from the demonstrations D. Density estimation remains a challenging problem , especially when there are a limited number of samples and the data is high dimensional ( Liu et al. , 2007 ) . Note that simply extracting the conditional ⇡ ( a|s ) from an estimate of the joint ⇢⇡E ( s , a ) is an alternate way to do BC and does not resolve the compounding error problem ( Ross et al. , 2011a ) . 2 . H ( ⇢⇡✓ ) is hard to maximize as ⇢⇡✓ is an implicit density . This challenge is similar to the difficulty of entropy regularizing generators ( Belghazi et al. , 2018 ; Dieng et al. , 2019 ; Mohamed & Lakshminarayanan , 2016 ) for Generative Adversarial Networks ( GANs ) ( Goodfellow et al. , 2014 ) , and most existing approaches ( Dieng et al. , 2019 ; Lee et al. , 2019 ) use adversarial optimization .

Summary: This paper presents an imitation learning approach based on density estimation of the state-action pairs. The underlying idea is to take expert demonstrations and find a distribution over the inherent state occupancy and then follow-up with a reinforcement learning step. As a result, the learner visits the high density state action pairs while also exploring new states during training. The paper makes a number of interesting contributions, among a new cost function for distribution matching IL and a proof showing that it is equivalent to maximizing a non-adversarial model-free RL objective. In contrast to other methods, like AIL, no adversarial optimization is needed; which typically would have resulted in instabilities. Instead the new objective pushes the lower bound in a single direction. The paper also features a list of tradeoffs between different distribution matching IL methods. Experiments on Mujoco tasks show state of the art performance.