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 introduces an approach for imitation learning based on density estimation. The approach uses the previously introduced idea of minimizing some divergence between policy and expert occupancy measures, state-action distributions induced by these policies. The authors propose to first estimate expert occupancy measure either using an autoregressive density model or an energy based model. Then authors use the Donsker-Varadhan KL representation to compute a log ratio between $p(s_{t+1}|s_t)$ and $p(s_t)$ where $s_{t}$ is a state at time $t$. Then, the expert occupancy and the KL representation are used as RL reward for imitation learning.