Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 7: Good paper, accept
review_content: summary: This paper proposes a new imitation learning algorithm, ICIL, which aims to make imitation learning more robust to irrelevant distractor features (i.e. features that are not causally upstream of the chosen expert action across a range of environments). Starting from BC, ICIL first adds a state encoder that splits states into an action-relevant component $s_t$ and a noise/distractor component $\eta_t$. Next it adds three additional auxiliary losses to prevent distractors from influencing $s_t$: a "dynamics-preserving" loss to ensure that $(s_t, \eta_t)$ are sufficient to predict the next environment observation; an "environment invariance" loss to make $s_t$ marginally independent of the environment; and a mutual information loss that regularises $s_t$ and $\eta_t$ towards marginal independence. Finally, ICIL incorporates an auxiliary loss to match next-step expert state occupancy under a learnt dynamics model, with the intent of decreasing compounding error. Experiments show that ICIL gets higher return than several baselines in variants of standard benchmarks that have been adapted to include additional distractor features.
main_review: In my view, the main strengths of this paper are:

1. Clarity. Writing quality was high on the whole: the problem is motivated in a compelling way, assumptions are clearly stated, and description of the algorithm and experiments makes sense. My only objection is to the assumptions and desiderata that are introduced without enough justification in sections 3-4 (see weaknesses below).
2. The experiments have a wide selection of baselines. Incorporating IRM into the baselines was particularly appreciated—it's hard to find appropriate baselines for a problem like this, so it's good to see some novel baselines that at least try to solve the same problem in different ways. Moreover, the proposed algorithm manages to improve substantially over the baselines in all the examined environments.

However, there are also some weaknesses to the algorithm:

1. The algorithm makes several non-trivial assumptions and has many moving parts, and some of them seemed arbitrary or poorly motivated. For instance, it's assumed that the transitions of decision-relevant variables are the same between "environments", but transitions of noise variables are not. What kinds of domains is this true for in practice? The desiderata in Section 4.1 (invariance across environments, dynamics-preserving $s$/$\eta$, disentangled $s$/$\eta$) are also vague and get introduced without justification:
   1. Why is having equal $s_t$ marginals across environments an adequate notion of "invariance"? What about environments where the demonstrator really does have different $s_t$ marginals because of some difference in environment structure? It seems like $\mathcal L_{\text{inv}}$ objective is going to throw away some relevant information in such cases.
   2. Having "dynamics-preserving" state representations seems too weak—why not just have an invertible mapping from $x$ to $(s,\eta)$? Such a mapping would still be "dynamics-preserving" while also satisfying Assumption 3.1 directly.
   3. What does it mean for $s$ and $\eta$ to be "disentangled"? Why is this desirable? "Disentangled" is a loaded word, and statistical independence (which an MI penalty incentivises) is neither necessary nor sufficient to obtain the kind of disentanglement that you see in carefully hand-designed state representations (for a counterexample, see Theorem 1 of "Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations" by Locatello et al.).
2. I also had several concerns about experiments:
   1. The various losses (there are five of them!) do not seem to have been examined separately in an ablation, so it's not clear which are responsible for the improvement over baselines.
   2. The experiment domains all have manually-inserted distractors that are deliberately constructed to meet the assumptions of the algorithm. The experiments thus do not provide evidence that the assumptions are adequate to handle the sort of distractors present in real-world tasks.
   3. The experiments report return/accuracy, but do not directly evaluate how sensitive the various models are to the inserted distractor features. In some ways this would be more informative than just looking at overall performance, since it would distinguish between gains from reduced reliance on distractors and gains from improved optimisation, etc.

On the whole I feel that this is a well-written paper, but I am currently ambivalent about acceptance because I do not feel the justification and evaluation of the algorithm is adequate given how complex the algorithm is. Clarifying points raised in the weaknesses above would help.
limitations_and_societal_impact: Limitations and societal impact were discussed adequately.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: The authors consider the setting of batch imitation learning with access to expert demonstrations over multiple environments which correspond to interventions on shared, underlying causal structure. They use the framework from [7] to propose a method that disentangles the relevant and irrelevant components of the state for a task with an additional regularization component that minimizes the mutual information between the two components with MINE. For the imitation policy, they propose an energy-based model approach to approximate the occupancy measure of the expert policy as a way to regularize and train the policy. 

main_review: At first glance imitation learning in the batch setting seems like a strange setup, and therefore not so unusual that it is original. However, the authors motivate this setting well by pointing out that in the medical domain, these are reasonable assumptions to make — that you are in an offline setting with access to expert data from different environments with shared latent structure.

As the authors themselves point out, the setup, assumption, and parts of the method are directly from [7] but there are novel components to modify it for the batch and imitation learning setting. The experiments seem reasonable, although the unusual metric on the Open AI Gym tasks make it difficult to ground these results to those reported in other works.  Perhaps a more relevant benchmark would be D4RL, which has an expert dataset. 

The MIMIC III experiment, given the setting, is well chosen. The results there also are clearly quite significant and with respect to reasonable baselines. It would also perhaps be best to report both train and evaluation performance given the emphasis on the focus being generalization.

Overall this seems like a solid contribution and I would lean towards accept.
limitations_and_societal_impact: Yes
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: This paper proposes a method for performing imitation learning from a set of expert demonstrations collected across a set of environment variations, where each variation is an intervention on the noise variables of the environment. The method, ICIL, is demonstrated to successfully control for spurious correlations in a set of OpenAI Gym envronments and a healthcare domain dataset.
main_review: The method itself is well motivated with a clear formalization of the problem setting. The work is well positioned with respect to related works, and the specific problem of learning a generalizable policy from offline data is an important one. 

My main concern relates to how the experiments focus on highly contrived toy cases specifically designed to fit the formal assumptions of their method. It would be beneficial to include experiments that test the method in settings where the examples were not hand-designed by the authors, e.g. OpenAI Procgen, where many levels show differences that can be viewed as spurious noise, e.g. background colors; or a version of Mujoco where the background changes across episodes. These more challenging examples would test how well this method can scale in practice, and how robust this method is to domains that do not perfectly match the structural assumptions around per-environment interventions only occuring on noise variables.

If the authors provide results on a less contrived domain, I will gladly raise my score to at least 7.

### Post-rebuttal update:
The authors' rebuttal addresses my concerns, so I am raising my score to a 7.
limitations_and_societal_impact: - The method makes the assumption that the environments from which demonstrations are collected correspond to interventions only on the noise variables of the environment. It would be valuable for the authors to describe real world scenarios meeting these assumptions, or if such examples are hard to come by, describe how these assumptions are relaxed in example practical domains.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: The paper presents an imitation learning approach that can generalize across demonstrations acquired from multiple environments. The approach projects the observations in an invariant shared latent space across environments. The variability across environments is modeled as noisy perturbations to reflect spurious correlations with the actions. The noisy perturbations and invariant state representations are learned such that the transition dynamics are preserved, while the mutual information between them is minimized. Experiments on OpenAI gym benchmarks and MIMIC III derived dataset suggest performance improvement over other competing methods.
main_review: The paper is well-written for the most part. The problem formulation is interesting and the rationale for the choice of loss functions is well-motivated. Fig. 3 is helpful to explain the ideas of the paper. Experimental results are sound, and interesting for a wider imitation learning community. 

- The problem of generalization in imitation is studied in several different contexts. Environments can vary across pixel distributions (such as sim or real [1]), dynamics [2], reward functions and so on. It would be useful to investigate these cases independently, and in combination with each other. For example, situations where the dynamics and the reward function are the same but pixel distributions are different, does \eta^e converge to a trivial solution ? 

- As the paper addresses dynamics invariance, it would be useful to specify the generalization context clearly (including the title).

- Transfer learning and domain generalization are closely related to the proposed approach. Related work can be better grounded in literature, in particular with [7]. I am also not convinced if the paper is the first to tackle generalizable imitation learning in a batch setting.

- It would be useful to report ablation studies. While the loss functions are well-motivated, they likely have competing objectives which can make the learning unstable. Results reporting the role of each loss function alone and in combination would give furhter insights about the efficacy of different terms.

- Generalizing across different environments here means doing well on average, but the resulting imitation behvaior may not be optimal for a particular environment. Additional experiments to evaluate the imitation performance in a single environment vs multiple environments as a function of learned \eta would be helpful.

[1] DIRL: Domain-Invariant Representation Learning for Sim-to-Real Transfer, 2020

[2] Off-Dynamics Reinforcement Learning: Training for Transfer with Domain Classifiers, 2020-21

limitations_and_societal_impact: see above. No negative societal impact.
review_title: null

Meta-review: This paper proposed a method which exploits principles of causal invariance from causality in order to develop an imitation learning algorithm which is more robust to spurious/distractor features in the expert trajectories. It shows the improved performance of this method in OpenAI Gym tasks and a healthcare dataset. The reviewers were unanimous in recommending acceptance, and no significant argument made by the lone weak accept has made me doubt this paper is worth publishing at NeurIPS. Without much difficulty, I can also recommend acceptance, and hope to see the minor improvements suggested by the reviewers in the final version of the paper.
Decision: Accept (Poster)