Imitation with Neural Density Models

Abstract: We propose a new framework for Imitation Learning (IL) via density estimation of the expert's occupancy measure followed by Maximum Occupancy Entropy Reinforcement Learning (RL) using the density as a reward. Our approach maximizes a non-adversarial model-free RL objective that provably lower bounds reverse Kullback–Leibler divergence between occupancy measures of the expert and imitator. We present a practical IL algorithm, Neural Density Imitation (NDI), which obtains state-of-the-art demonstration efficiency on benchmark control tasks. 

Reviews and author responses:

=====POST-REBUTTAL COMMENTS======== 

I thank the authors for the response and the efforts in the updated draft. Most of my concerns were clarified and I still think the paper should be accepted. However, I agree with Reviewer 4 that  additional experiments would be good to better tease out the reasons for this method working.



##########################################################################

Summary:
 

The paper proposes an interesting way to do imitation learning without using an adversarial framework. The proposed approach involves density estimation to learn a surrogate reward function that can be optimized via RL. The approach is motivated by recently proposed maximum occupancy entropy RL and extends this to the imitation learning setting. The authors derive an efficient policy optimization method that outperforms existing approaches for imitation learning.

##########################################################################

Reasons for score: 
 

This paper provides strong theory and strong empirical results validating the theory. GAIL-like methods are notorious for their instability so having non-adversarial IL methods is a significant improvement.
 

##########################################################################

Pros: 
 

1. Non-adversarial approach to IL that seems well suited policy optimization with any RL algorithm.

2. Significant improvement over state-of-the-art IL approaches. Also works with only one demonstration.

3. Nice theoretical results showing the objective lower bounds reverse KL between expert and imitator.

##########################################################################

Cons: 
 

1. Only results on mujoco tasks with state information. Most of these tasks can be solved reasonably well with just a bias toward longer episodes. It would be nice to show that the authors method qualitatively imitates a variety of behaviors rather than just being able to go really fast without falling down. Imitating something like the hopper back flip would be a nice addition (https://openai.com/blog/deep-reinforcement-learning-from-human-preferences/). Some other kind of open ended task would also be interesting where learning from demonstrations actually makes sense.

2. Omits other papers that also perform efficient reward learning, then RL for non-adversarial imitation learning: 
e.g. 
Uchibe. "Model-free deep inverse reinforcement learning by logistic regression." Neural Processing Letters, 2018.
Brown et al. "Better-than-Demonstrator Imitation Learning via Automatically-Ranked Demonstrations." CoRL, 2019. 

3. Quite a few knobs that need to be tuned in terms of hyperparameters. Perhaps I missed it, but it would be nice to have better intuition for how to set these and how imitation behavior changes based on them.


##########################################################################

Questions/Clarifications:

I'm confused about why the authors think this method will work well with hard exploration video games. I'm not sure I buy the claim that optimizing the authors bound will cause an agent to explore new levels.

After equation (9) the authors say this will cause the policy to seek out areas with low randomness and update the policy to be random there, but the alternative also seems equally likely given the objective, e.g., seek out states where there is lots of randomness and try and make the policy less random there.

The notation for the autoregressive section is confusing what does x = (x_1, ..., x_dim(S) + dim(A) ) =  (s,a) mean? What if S and A are inf dimensional? Also, autoregressive models are often used for time series data, but here is looks like q is conditioned on (s,a) pairs from some lexicographic ordering which does not seem to make sense since there may be no correlation between actions at far apart states.

We thank all the reviewers for positive feedback and appreciation of our work! As there were no common major concerns, we provide individual responses to reviewers. Here we provide a meta-summary of our response as well as how we’ve addressed the feedback. We hope the reviewers have had time to read our individual responses and look forward to hearing their updated evaluation! 

**Experiments**: We’re happy the reviewers assessed that we have a good set of experiments (R4) that are complete (R3), convincing(R3), and demonstrate significant (R2) improvement over state-of-the-art (R2, R3).  We hope our individual response to R1 has clarified that the main additional experiment they requested, regarding comparison to baselines that directly apply Donsker-Varadhan representation to bound occupancy divergence, was in fact included in the original submission. To further the completeness of experiments, we have addressed all of the reviewer’s requests for additional experiments which are an additional ablation (R3) study and a minor extension of our current experiment (R1). 

**We would like to emphasize that NDI achieves expert level performance on all Mujoco benchmarks, including Humanoid, which sets a new standard for state-of-the-art demonstration efficiency, as acknowledged by (R2, R3). This is the main empirical contribution of NDI.**

**Theory**: We are also glad that the reviewers have found the theory underlying NDI to be strong (R2), interesting (R1), and novel (R3). We are pleased that R2, R3 has appreciated that the NDI objective is non-adversarial thereby avoiding the instability of alternating min-max optimization (R2). We have further revised the paper to address all conceptual misunderstandings and provided answers to all minor questions in the individual responses below. In particular, we have revised the paper to make it crystal clear what the differences are between NDI (ours) and adversarial IL methods are. This addresses R1 and R4’s questions regarding whether or NDI is truly non-adversarial (R4) and the benefit of applying the Donsker-Varadhan representation to bound the mutual information term instead of directly to bounding the occupancy divergence (R1). 

**We would like to emphasize that  NDI introduces a novel family of non-adversarial IL methods free from the instability of alternating min-max optimization, even when the critic is updated together with the policy.**




**Q. Reference for AIL works poorly in visual environments**

A. Please see [4], [5]


**Q. What does the sentence "We posit that.." mean?**

A. We're referring to our thinking that the SAELBO helps with state-action level exploration. We have added experiments to show how optimizing the SAELBO can be more advantageous for exploration compared to just optimizing policy entropy. Please see the updated Appendix C.1 


**Q. Implementation**

**Q.0 Discriminator regularization?**

A. We attempted spectral, gradient penalty, orthogonal, $L2$, and $L1$ regularization but found it does not lead to better GAIL performance. 


**Q.1. What RL algorithm does GAIL use?**

A.1. It uses SAC like in Ghasemipour et al., 2019. 


**Q.2. SAC replay buffer store on-policy only?**

A.2. No it contains past interactions as we use SAC. 


**Q.3. Using on-policy samples for $q$ marginals?**

A.3. We are trying to approximate on-policy samples by sampling from a recent portion of the replay buffer, but technically no. 


**Q.4. Ratio of model update steps to environment steps**

A.4. If the reviewer is asking about how many RL updates per environment step then all methods perform one gradient step per every environment step. 


**Q. Why does $\lambda_f$ improve KL?**

A. Great question! Setting $\lambda_f > 0$ corresponds to maximizing the SAELBO while $\lambda_f = 0$ corresponds to solely maximizing policy entropy. Maximizing the SAELBO instead of maximizing solely policy entropy improves KL because maximizing the SAELBO can be more effective for occupancy entropy maximization (which is a part of the KL objective). This is because in discrete state-spaces the SAELBO $\mathcal{H}^{f}(\rho_{\pi_{\theta}})$ is a tighter lower bound to occupancy entropy $\mathcal{H}(\rho_{\pi_{\theta}})$ than policy entropy $\mathcal{H}(\pi_{\theta})$, i.e $\mathcal{H}(\pi_{\theta}) \leq \mathcal{H}^{f}(\rho_{\pi_{\theta}}) \leq \mathcal{H}(\rho_{\pi_{\theta}})$, and in continuous state-spaces, where Assumption 1 holds, the SAELBO is still a lower bound while policy entropy alone is neither a lower nor upper bound to occupancy entropy. As an artifact, we found that maximizing the SAELBO ($\lambda_f > 0$) leads to better occupancy distribution matching, i.e imitation performance, than sole policy entropy maximization ($\lambda_f = 0$). We have updated the explanation in Section 6.3 and have added more experiments to support the advantage of SAELBO maximization over policy entropy maximization in Appendix C.1.


**Q. Add reference for "Divergence Minimization Perspective ..."**

A. Thank you for the suggestion. We have added the reference to Section 5. 


**Q. Characterization of optimal policy behavior for MaxOccEntRL**

A. For a general MaxOccEntRL objective of the form $\max_{\pi} E_{\rho_{\pi}}[r(s, a)] + H(\rho_{\pi})$, the optimal RL agent will try it's best to match its own occupancy with a target occupancy $\rho_{target} \propto e^{r(s_t, a_t)}$. As the reviewer said, this is similar to how in MaxEntRL an optimal RL agent tries it's best to match it's own trajectory distribution with a target trajectory distribution $p(\tau) \propto e^{r(\tau)}$. In deterministic environments, an agent can sample trajectories according to any distribution that it wishes, and thus the MaxEntRL optimal agent will sample trajectories exactly according to $\rho_{target}$. However, the environment may not permit sampling state-action pairs according to any distribution, e.g you cannot only sample one particular state action pair $(s^*, a^*)$ unless the dynamics allows you to perpetually stay in $s^*$. Thus it become difficult to analytically characterize the attained optimal occupancy for all MaxOccEntRL problems. 


**Q. Section 2.2: "... Our bound will encourage the agent to continuously seek out…”. You're also increasing entropy, reducing the chances of getting to those states. The initial low variance random policy might/probably won't get anywhere interesting anyways.**

A. Thank you for pointing this out. We agree with the reviewer that the explanation was not convincing. In the revision, we have updated Section 2 to more concisely explain why SAELBO maximization can enable better state-action level exploration than sole policy entropy maximization. Furthermore, we have added experiments in Appendix C.1 to support this claim.  


**Q. SAC already has policy entropy, so do we include it twice?**

A. We do not include an additional policy entropy bonus since SAC already has it (stated in Appendix B). We have clarified this in Section 3 of the main text in the revision.  


**Q. What is $\lambda_{\pi}$**

A. That is the multiplier on the policy entropy (stated under Eq. (11)) which is fixed to be $0.2$ across all experiments. (see Section 6, Architecture)  


**Q. Intuition for poor performance of NDI on Ant**

A. Please see the explanation in Section 6.1. When using MADE, the learned density for the ant task had spurious modes, even with 25 demonstrations, unlike when using EBMs. Spurious modes outside of the data distribution is a problem encountered by many density estimators [7]. 


Thank you very much for your detailed and thoughtful feedback. We would like to clarify misunderstandings and answer your questions. Furthermore, we've updated the revision to address all of your suggestions for improving the paper. 


====**Completed Revisions**====

-- Added Section 4 detailing why NDI is non-adversarial and what trade-offs different distribution matching IL algorithms make. 

-- Included Ass. 1 and a discussion of it in Section 2. 

-- Added experiment in Appendix C.1 to demonstrate effect of SAELBO in aiding exploration

-- Revised explanation of SAELBO in Section 2 and Section 6.3

-- Added clarification regarding checkpointing criterion and GAIL implemenation in Appendix B. 

-- Included suggested reference in Section 5. 


====**Answers to questions**====

**Q. Is our method non-adversarial? Does the mutual information term need to be optimized with adversarial optimization?**

A. Great question. Our method is indeed non-adversarial and the mutual information term does not need to be optimized with adversarial optimization. For example one can simply use gradient descent on samples as in [6] to maximize mutual information lower bounds.
Adversarial IL methods maximize, with respect to the policy $\pi_{\theta}$, an upper bound to the additive inverse of occupancy divergence. The discriminator is updated to minimize the upper bound, thereby tightening the estimated divergence, and the policy is updated to maximize the tightened upper bound. The use of an upperbound innevitably leads to an alternating min-max optimization scheme.
The key insight of NDI is to instead maximize, with respect to the policy $\pi_{\theta}$, a lower bound to the additive inverse of the occupancy divergence. Then, the critic $f$ is updated to maximize the lower bound, thereby tightening the estimated divergence, and the policy $\pi_{\theta}$ is updated to maximize the tightened lower bound. The use of a lowerbound allows NDI to avoid alternating min-max and instead perform alternating max-max.  
While NDI enjoys non-adversarial optimization, it comes at the cost of having to use a non-tight lower bound to the occupancy divergence. On the otherhand, AIL optimizes a tight upper bound at the cost of unstable alternating min-max optimization. Please have a look at Section 4 in the revision for more details. Section 4 explains in detail why our method is non-adversarial and compares the trade-offs between Adversarial IL, Support Matching IL, and NDI. 


**Q. Unique advantage of method and how it's shown in experiments?**

A. Our unique contribution is introducing a novel family of distribution matching IL algorithms 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. (1) is explained in Section 4 of the revision and (2). is shown by results in Table 2. We have made this more clear in the Section 1 of the revision. 


**Q. Move Ass. 1 to main text and add discussion of what happens when it does not hold.**

A. Thank you for this suggestion. We have moved the assumption to the main text (Section 2) as well as a discussion of what happens when the assumption does not hold. More details were added to the beginning of Appendix A as well. Please see revised Section 2 and Appendix A. 


**Q. Why checkpoint on augmented rewards? What about other methods?**

A. The model that attains the highest policy performance with the augmented reward (log density + policy entropy + MI) is the model that maximizes the lower bound to additive inverse of reverse KL. (see Corollary 1) If the question is about why not use the true environment reward, that’s because that reward function is assumed to be unknown in the Imitation Learning setting, and only used for task performance evaluation purposes. All other methods were checkpointed with the method proposed in their respective papers which follows the same principle. In RED we save the model with the highest support matching reward. For ValueDICE, we save the model with the lowest value dice loss from their paper. For GAIL, we save when the discriminator rewards saturate over 40 episodes, similar to how GAN training is stopped. For BC, we save the model with the lowest cross validation loss. We added these details to Appendix B in the revision. 



**Q. Could the author explain the reason why one expert trajectory can provide a good density estimation of expert policy? Typically, we cannot expect one trajectory to cover the state-action space. It is less attractive if the learned policy can reach expert-level performance but stick to the original modality of the expert policy.**

A. Great question. While one trajectory does not contains enough state-actions to cover the state-action space, it contains enough information to learn a density model that positively correlates with forward velocity of the robots (see Section 6.2). Thus, even with one trajectory good task performance can be achieved since the learned density model shares similar characteristics as the true reward function for the considered robotics environments. This means that the imitator can "learn" through interacting with the environment, how to behave in states not covered by the demonstrations. However, the specific way that the imitator accomplishes the task outside of the support of the demonstrations may differ from the expert, thus leading to suboptimal imitation performance, i.e closeness in the expert and imitator occupancies. From the ablation studies in Appendix C.3, it's our belief that having more trajectories allows for better coverage of the state-action space, hence better density estimation on a wider support of the expert occupancy which leads to improved imitation performance. 


**Q. Typo in proof of Lemma 1.**

A. Thanks for point this out, we've fixed it. 


**Q. Regarding theorem 1 Do we need to normalize $p_{\theta, t}$ with  before applying concave property?**

A. You are right. We have fixed it in the revision. 

Thank you very much for your reviews. We are happy you appreciated the novelty of our method and soundness of the experimental results. We would like to answer your questions. Furthermore, we've updated the revision to address all of your suggestions for improving the paper. 


====**Completed Revisions**====

-- Added suggested ablation study to Appendix C.3 regarding the effect of each phase on performance. 

-- Fixed typos in proofs 


====**Answers to questions**====

**Q. Please specify the number of random seeds and the number of evaluation trajectories. What is the mechanism to choose those trajectories?**

A. We use 5 random seeds and 50 evaluation trajectories. These details were added to the revision in the "Pipeline" subsection of Section 6. Regarding the criterion for when to evaluate performance (i.e which trajectories to choose), for NDI, the model that attains the highest policy performance with the augmented reward (log density + policy entropy + MI) was evaluated as it is the model that maximizes the lower bound to KL. (see Corollary 1) In RED we save the model with the highest support matching reward. For ValueDICE, we save the model with the lowest value dice loss from their paper. For GAIL, we save when the discriminator rewards saturate over 40 episodes, similar to how GAN training is stopped. For BC, we save the model with the lowest cross validation loss. We added these details to Appendix B in the revision. 


**Q. Although the sample-efficiency is not the focus of this work, it would be helpful to see the averaged learning curves as in the results in related works, which can give a straight forward intuition on the robustness of the proposed methods.**

A. We had provided environment sample complexity results in Appendix C.5 along with standard deviations which gives information about robustness. We can include a full plot if the reviewer would like to see them.


**Q. Does the improvement come from the better density estimation of the expert's policy? It would be helpful to separate the two phases and conduct an ablation study to show how does each phase affects the performance. It would be nice to see the results of replacing density estimation with oracle [density estimated by a sufficient amount of trajectries]. I didn't see too much difference in terms of final performance between 1 and 25 expert trajectories. Is there a large gap in the sample complexity?**

A. Thank you for your suggestion! We have added the suggested ablation study to Appendix C.3. We first fix the MI reward weight to the optimal value $\lambda_f = 0.005$, then vary the number of demonstrations to isolate the affect of the density estimation phase on overall performance. We found that having more demonstrations for density estimation slightly improves imitation performance (KL) on most mujoco tasks. There's no clear improvement in task performance (Reward), as expert level reward is already attained with one demonstration. Next we fix the EBM density model $q_{\phi}(s, a)$ to be close to "oracle" by training on an ample (25) number of demonstrations then vary the strength of the MI reward $\lambda_f$ in order to isolate the effect of the MaxOccEntRL step on overall performance. We found that, similar to results in Table 3, $\lambda_f$ mainly trades off task and imitation performance. Setting $\lambda_f$ too small drives the imitator to concentrate it's probability mass onto the modes of the expert occupancy, hence achieving good task performance at the expense of imitation performance. Setting $\lambda_f$ too large makes the entropy term dominate the objective leading to poor imitation and task performance. There's a sweet spot value of $\lambda_f = 0.005$ which balances the mode-seeking and mode-covering behavior.

Thank you very much for your reviews. We are happy you appreciated the novelty of our method, soundness of theory, and significance of the experimental results. We would like to answer your questions. Furthermore, we've updated the revision to address your suggestions for improving the paper. 


====**Completed Revisions**====

-- Added suggested references for IRL in Section 5. 

-- Revised explanation of SAELBO in Section 2 and Section 6.3

-- Added experiments in Appendix C.1 to demonstrate effect of SAELBO in aiding exploration. 

-- Fixed typos in description of autoregressive models. 


====**Answers to questions**====

**Q. Would be nice to see imitation of behaviors such as hopper back flip or other open-ended tasks.**

A. Thank you for this suggestion. We will add the hopper back-flip experiments.  


**Q. Include references for papers that perform efficient reward learning, then RL for non-adversarial imitation**

A. Thank you for this suggestion. Both suggested references have been added. 


**Q. Is there some intuition for how to set hyperparameters and how it affects imitation behavior?**

A. There are two main hyperparameters for NDI. The policy entropy multiplier $\lambda_{\pi}$ and MI reward multiplier $\lambda_{f}$. For $\lambda_{\pi}$ we simply follow the protocol outlined in the original SAC [1] paper for tuning it since SAC already includes a policy entropy bonus. In Section 6.3, we had included ablation studies on how $\lambda_{f}$ affects both task and imitation performance. We have also added additional experiments in Appendix C.1 to further demonstrate how $\lambda_f$ helps with exploration in environments where naive policy entropy maximization does not lead to occupancy entropy maximization. 


**Q. I'm confused about why the authors think this method will work well with hard exploration video games. I'm not sure I buy the claim that optimizing the authors bound will cause an agent to explore new levels. After equation (9) the authors say this will cause the policy to seek out areas with low randomness and update the policy to be random there, but the alternative also seems equally likely given the objective, e.g., seek out states where there is lots of randomness and try and make the policy less random there.**

A. Thank you for pointing this out. We agree with the reviewer that the explanation was not convincing. In the revision, we have updated Section 2 to more concisely explain why SAELBO maximization can enable better state-action level exploration than sole policy entropy maximization. Furthermore, we have added experiments in Appendix C.1 to support this claim. We also agree with the reviewer that no strong scientific claim about the effectiveness of SAELBO maximization on hard video games such as Montezuma's revenge can be made at the moment, and we have clarified in the revision that we do not make such claims. 


**Q. The notation for the autoregressive section is confusing. Also, autoregressive models are often used for time series data, but here is looks like q is conditioned on (s,a) pairs from some lexicographic ordering which does not seem to make sense since there may be no correlation between actions at far apart states.**

A. Sorry about the confusion! There were typos in that section which are now fixed in the revision. We meant $x = (x_1, ..., x_{dim(s) + dim(a)}) = (s, a)$ to denote a vector that concatenated one state and one action. $q(x)$ is a distribution over state-action pairs. $i$ in $x_i$ indexes each dimension of a single state-action vector and does not index over time. 


**References**

[1]. Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor, Haarnoja et al. 2018

**References**

[1]. A Divergence Minimization Perspective on Imitation Learning Methods,
Ghasemipour et al. 2019

[2]. Generative Adversarial Imitation Learning, Ho et al. 2016

[3]. Imitation Learning via Off-Policy Distribution Matching, Kostrikov et al. 2020

[4]. Improved techniques for training GANs, Salimans et al. 2016

Thank you very much for your appreciation of our work. We would like to answer your questions and clarify any misunderstandings. Furthermore, we've updated the revision to address all of your suggestions for improving the paper. 


====**Completed Revisions**====

-- Added description of how $f$ is computed in Algorithm 1. 

-- Added Section 4 detailing why NDI is non-adversarial and what trade-offs different distribution matching IL algorithms make. 

-- Included results for $\lambda_f = 0$ in Table 3. 


====**Answers to questions**====

**Q. NDI uses SAC as its underlying RL algorithm so compare with other methods in sample efficient RL.**

A. As mentioned in Section 6.1, the focus of this work is in improving demonstration efficiency, i.e. using fewer demonstrations to get better performance, and not environment sample efficiency, i.e. number of environment interactions needed to reach best IL performance. For completeness, we had included environment sample complexity comparisons in Appendix C.5. In these results (Table 9) we provide comparisons to ValueDICE which, to our knowledge, is the state-of-the-art in environment sample efficient Imitation Learning, out-performing the reviewer's suggested baseline DAC. Moreover, the GAIL baseline in our work is also implemented with SAC. 


**Q. How many additional interactions from the environments does this approach require?**

A. This was provided in Appendix C.5 and referenced at the end of Section 6.1. NDI (ours) roughly requires an order of magnitude less interactions than GAIL. 


**Q. Is the policy entropy term included twice since SAC already has one?**

A. As stated in Appendix B, we do not include an additional policy entropy bonus since SAC already has it. We have made this more clear in the revision by adding this description to Section 3.  


**Q. Algorithm 1 would benefit from explicitly stating how $f$ is computed in practice**

A. Thank you for your suggestion. We have added this to Algorithm 1. 


**Q. What is the benefit of using the representation (3) to compute state density only? The paper could benefit from additional experiments that demonstrate that the proposed formulation is better than simply using (3) for state-action distributions. (Note that the original Eq. (7) has been changed to Eq. (3) in the revision)**

A. Great question. Indeed the Donsker-Varadhan representation of (3) could be used to directly obtain a upper bound to the additive inverse of divergences between state-action occupancies, e.g reverse KL divergence $-D_{\mathrm{KL}}(\rho_{\pi_{\theta}} || \rho_{\pi_{expert}})$. In fact this is precisely what Adversarial Imitation Learning (AIL), such as GAIL and ValueDICE, does as shown in [1]. 
The discriminator $D$ in AIL, analogous to the critic function $f$ from (3), is updated to minimize the upper bound, thereby tightening the estimated divergence, and the policy $\pi_{\theta}$ is updated to maximize the tightened upper bound. So in fact the requested additional experiments are simply comparisons to AIL methods, e.g GAIL [2] and ValueDICE [3], which we have already done (see baselines in Section 6). Specifically, ValueDICE [3] optimizes the same reverse KL as NDI with the representation (3) used to directly upperbound state-action occupancy divergence. Please see [3] for a more detailed explanation. The issue with AIL is that the use of an upperbound innevitably leads to an alternating min-max optimization scheme, which is known to be notoriously unstable [4]. 
The key insight of NDI is to instead maximize, with respect to the policy $\pi_{\theta}$, a lower bound to the additive inverse of the occupancy divergence. Then, the critic $f$ is updated to maximize the lower bound, thereby tightening the estimated divergence, and the policy $\pi_{\theta}$ is updated to maximize the tightened lower bound. The use of a lowerbound allows NDI to avoid alternating min-max and instead perform alternating max-max. Crucially, while both AIL and NDI make use of the representation in (3), we optimize it in different directions with respect to the discriminator/critic, i.e in AIL the discriminator seeks to minimize (3) which upperbounds negative state-action occupancy divergence while in NDI the critic seeks to maximize (3) which lowerbounds Mutual Information between consecutive states (see Theorem 1). 
While NDI enjoys non-adversarial optimization, it comes at the cost of having to use a non-tight lower bound to the occupancy divergence. On the otherhand, AIL optimizes a tight upper bound at the cost of unstable alternating min-max optimization. Please have a look at Section 4 in the revision for more details. Section 4 explains in detail why our method is non-adversarial and compares the trade-offs between Adversarial IL, Support Matching IL, and NDI. 


**Q. Include results for $\lambda_f = 0$ in Table 3.**

A. Thank you for the suggestion. We have added these results to Table 3 in the revision. 


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. 

In overall, I think that this approach is interesting and this direction is under-explored in the context of imitation learning. The paper is well written and easy to follow. However, there are some issues that require additional clarification (see my comments below). I think that experimental evaluation is adequate but limited, the paper requires an additional ablation study in order to justify certain design decisions. Moreover, I believe that since the method is using SAC as its underlying RL algorithm its necessary to perform additional comparisons with other methods in sample efficient RL ([1], [2] or [3]).

In conclusion, I think that at the moment this is a borderline paper. The approach is definitely interesting but right now the experimental evaluation is not solid enough to recommend this paper for acceptance. I highly recommend the authors to revise the paper and I'm looking forward to authors' response.

Comments:
1) How many additional interactions from the environments does this approach require? Some other papers ([1], [2], [3]) plot training performance vs number of additional environment interactions.
2) policy entropy term is included into rewards in algorithm 1 (line 6). Is it the same entropy term as in SAC? If yes, what is the benefit of including it twice?
3) algorithm 1 would benefit from explicitly stating how f is computed in practice;
4) what is the benefit of using the representation (7) to compute state density only? Could the same representation be used to compute $KL(\rho_\pi(s, a)|\rho_{expert}(s,a))$? I think that the paper could significantly benefit from additional experiments that demonstrate that the proposed formulation is better than simply using (7) for state-action distributions.
5) Table 2 demonstrates that smaller values of $\lambda_f$ lead to better performance (with an exception of Ant but the difference is within the standard deviation). The figure table 2 can benefit from including results for $\lambda_f=0$.


References:

[1] Sample Efficient Imitation Learning for Continuous Control, Sasaki et al., ICLR 2019: https://openreview.net/forum?id=BkN5UoAqF7

[2] Discriminator-Actor-Critic, Kostrikov et al., ICLR 2019: https://openreview.net/forum?id=Hk4fpoA5Km

[3] Sample-Efficient Imitation Learning via Generative Adversarial Nets, Blonde et al,. 2018: https://arxiv.org/abs/1809.02064


This work proposes a novel density matching method for learning from demonstration, which achieves state-of-the-art demonstration efficiency. Prior density matching methods utilize the adversarial methods suffers from the instability of optimization. To overcome this issue, this work proposes to separate the imitation process into expert density estimation phase and density matching phase, where a model-free formulation is derived and provably served as the lower bound of reverse KL divergence between $\pi_\theta$ and expert policy $\pi_E$.

This work overcomes the instability issue of min-max optimization
 through transferring the objective to the lower bound leading to 
the model-free objective is attractive and novel as far as I am concerned. 

The experimental results are mostly complete and convincing. 
The improvement in the final performance over state-of-the-art is demonstrated. Although the sample-efficiency is not the focus of this 
work, it would be helpful to see the averaged learning curves as in the results
in related works, which can give a straight forward intuition on the
robustness of the proposed methods. 

Please specify the number of random seeds and, the number of evaluation trajectories [and what is the mechanism to choose those trajectories, such as the last 10 episodes?] of 
the reported results in Table 1 and other tables. 

Questions: 
1. Does the improvement come from the better density estimation of the expert's policy? It would be helpful to separate the two phases and conduct an ablation study to show how does each phase affects the performance. 
It would be nice to see the results of replacing density estimation with
oracle [density estimated by a sufficient amount of trajectries]. I didn't see too much difference in terms of final performance between 1 and 25 expert trajectories. Is there a large gap in the sample complexity?

2. Could the author explain the reason why one expert trajectory can provide a good density estimation of expert policy? Typically, we cannot expect one trajectory to cover the state-action space. It is less attractive
if the learned policy can reach expert-level performance but stick to
the original modality of the expert policy. 


Appendix: 
Lemma 1:
$-\lambda \sum_{x} \hat{p}(x) \log \hat{p}(x)+(1-\lambda) \sum_{x} \hat{q}(x) \log \hat{q}(x) $ $\rightarrow$ $-\left[\lambda \sum_{x} \hat{p}(x) \log \hat{p}(x)+(1-\lambda) \sum_{x} \hat{q}(x) \log \hat{q}(x)\right]$

Regarding theorem 1
Do we need to normalize $p_{\theta}(s)$ with $1/(1-\gamma)$ before applying concave property? 
$\left(\mathcal{H}\left(s_{0}\right)+\sum_{t=1}^{\infty} \gamma^{t} \mathcal{H}\left(s_{t}\right)\right)$ 



## Review
Key points are written in bold font.

* Section 1:
  * You should maybe include a reference to "A Divergence Minimization Perspective on Imitation Learning" (Ghasemipour et al. 2019) since that work studies properties of different divergences used for match the joint state-action distributions (similar to Ke et al. which you cited) as well state-marginal matching using adversarial methods.
* Section 2:
  * For max-ent RL, there is an intuitive connection between RL and probability distributions over the space of trajectories. For example -- when we have deterministic transition dynamics -- max-ent RL's optimal policy is the policy that "samples" trajectories according to the (unnormalized) density $\hat{p}(\tau) = \sum_t r(s_t, a_t)$. I understand the motivation of the addition of this form of entropy term instead, but is there anything that can be said about what the optimal policy for this objective looks like (in a similar vein to the max-ent RL solution I described above)?
  * __Section 2.2, Theorem 1: Assumption 1 should definitely be included in the main text. This is quite a strong assumption (even if it holds in many physical/robotics problems), and the reader should not have to look into the appendix to realize this. Furthermore, I strongly believe that you should include a discussion in the main text of why you needed these assumptions (what goes wrong when these assumptions are not satisfied?). This is important so that 1) you provide intuition to the reader of what can go wrong if this is not satisfied, 2) future work can try to address the limitations.__ For example, if the transition dynamics was not deterministic, the intuition your provide when saying "$I\_{NWJ}$ encourages the agent to seek out states where...", is not necessarily correct. In a non-deterministic transition setting, the mutual information term could seek out states where the dynamics are more determinstic, but I don't see why that would be a good thing necessarily.
  * Section 2.2: "... Our bound will encourage the agent to continuously seek out...": I don't quite agree with this. You're also increasing entropy, which reduces the chances of getting to those states. Also, the initial low variance random policy might/probably won't get anywhere interesting anyways.
* Section 3:
  * You are using SAC as your RL algorithm. SAC itself is introducing entropy, so maybe comment on how that is playing with your entropy term. Is there any concerns one should be aware of?
  * __Top of page 6. You are saying that your method is not adversarial *<<even if you optimize $f$>>*. I believe, despite appendix A.3, this statement is incorrect. Consider the AIRL (Fu et al.) algorithm. According to Ghasemipour et al., 2019, the AIRL objective minimizes $KL(\rho^\pi(s,a)||\rho^{exp}(s,a))$.__
$-KL(\rho^\pi(s,a)||\rho^{exp}(s,a)) = E_{\rho^\pi(s,a)} [\log \frac{\rho^{exp}(s,a)}{\rho^\pi(s,a)}] = \text{RL with the loss }\quad \log \frac{\rho^{exp}(s,a)}{\rho^\pi(s,a)}$
__In works such as AIRL and GAIL there are two things to consider. 1) lower bounds similar to your $I\_{NWJ}$ lower bound are written to estimate quantities of interest (Jensen-Shannon for GAIL, and reverse KL for AIRL). These lower bounds are maximized to obtain tighter estimates of the quantity of interest. 2) Once a good estimate is obtained, an RL step is performed using the estimate. And this process is repeated.
In your work, you are trying to avoid directly estimating $\log \frac{\rho^\pi(s,a)}{\rho^{exp}(s,a)}$, and have decomposed the objective. If we try to write this $KL$ objective in a form that looks closer to your equations, we get:__
$-KL(\rho^\pi(s,a)||\rho^{exp}(s,a))
= E_{\rho^\pi(s,a)}[\log \frac{\rho^{exp}(s,a)}{\rho^\pi(s,a)}] = E_{\rho^\pi(s,a)}[\log \rho^{exp}(s,a) - \log \rho^\pi(s,a)]$
$= E_{\rho^\pi(s,a)}[\log \rho^{exp}(s,a)] - E_{\rho^\pi(s,a)}[\log \rho^\pi(s,a)] \text{(The J and H in your equation 2)}$
$= E_{\rho^\pi(s,a)}[\log \rho^{exp}(s,a)] - E_{\rho^\pi(s,a)}[\log \rho^\pi(a|s)] - E_{\rho^\pi(s,a)}[\log \rho^\pi(s)]$
__In your work, you are estimating $\log \rho^{exp}(s,a)$ directly through generative models, $\log \rho^\pi(a|s)$ is available from the policy, and in your derivations, it seems that you have been able to replace  $- E\_{\rho^\pi(s,a)}[\log \rho^\pi(s)]$ with the mutual information terms (due to Assumption 1 I believe). But this remaining mutual information terms is still being estimated using similar processes/lower bounds as prior work.__
__Hence, it is my view that your method is as "adversarial" as prior work. The reason you are suggesting stability of your optimization process (which I don't think you have quantified) is maybe because you are not actually learning $f(s\_{t+1}, s\_t)$ which is the property that makes those prior works "adversarial".__
* Section 5:
  * __Pipeline: Why did you checkpoint based on your augmented rewards? How were the checkpoints made for the other methods?__
  * Architecture: What is $\lambda_\pi$?
  * In GAIL are you using techniques from the GAN literatures (such as gradient penalty or spectral normalization) to regularize the discriminators? And are you using the original TRPO optimization or a newer technique such as PPO?
  * Do you have intuition for the poor performance or Ant NDI+MADE? For the 25 demos version in the appendix this value is also low.
  * __Section 5.1: Although this does not seem to be a central claim of contribution in your work, I am wondering if your claims of sample-efficiency are well-founded. You are using SAC, which is an off-policy algorithm, whereas GAIL uses on-policy RL (e.g. TRPO). To enable a proper judgement of sample-efficiency, please respond to the following questions: 1) What implementation of GAIL did you use? Does it use TRPO. If so, this might be a contributing factor since there are better on-policy algs now such as PPO. 2) For your method, in your SAC replay buffer are you only storing on-policy samples, or maintaining a large replay buffer or past interactions? 3) When you're estimating the mutual information, you are not using on policy samples for the marginal q distributions, correct? 4) What is the ratio of model update steps to environment steps in the various algorithms?__
  * With regards to the previous point, a more fair comparison might be to run GAIL with SAC and maintaining a full replay buffer, as done in Ghasemipour et al., 2019.
  * __Table 2: $\lambda\_f$ is tuning the knob for the stregth of the mutual information maximization term. Why should this be improving the KL?__
  * __You have a good set of experiments, but ( and I might be wrong) my biggest gripe with your experiments is that it doesn't seem that you have any experiments that would demonstrate a unique advantage of your method. Please clarify the unique advantages of your methodology, and how that is being shown in the experiments.__
* Section 6:
  * I don't think AIL works particularly poorly on visual domains (e.g. "InfoGAIL: Interpretable Imitation Learning from
Visual Demonstrations" or "Third Person Imitation Learning"). Do you have references for this statement?
  * I don't quite understand this statement: "We posit that...". Please clarify.

Meta-review: This is a difficult borderline decision, with the reviewers evenly split in their final recommendation.  Overall, the authors provided good responses to the reviewer questions: this was much appreciated.  The reviewers requested additional ablations and explanations, which the authors provided.

A prevailing concern is that the experimental evaluation, restricted to a few standard MuJoCo environments, does not really demonstrated a distinctive advantage for the proposed approach.  In fact, one of the new ablations added raises concerns about the significance of the paper's main technical contributions: the \lambda_f=0 row added to Table 3 shows very strong reward results, which apparently obviates a key aspect of the proposed approach. 

This work is interesting, and would like to see it published, but the current state of the evaluation does not support the significance of the main contribution.  I think the authors need to expand their empirical evaluation, as suggested, to better highlight the effectiveness of the proposed approach over the \lambda_f=0 baseline.  In the end, I think the authors would be better served by broadening the evaluation, isolating scenarios where the key proposal shows significant (rather than marginal) benefits, and publishing a more compelling version.

Decision: withdrawn-rejected-submissions