Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes to learn the invariances relating to MDPs with different observations (e.g. invariances in the underlying dynamics of the systems) using a generative model of the observations in an unsupervised manner. The inference in this generative model is accomplished using a GAN framework with auxiliary losses for the inverse dynamics of the system.
main_review: ## Strengths 

* **Novelty**. As far as I am aware the authors propose a novel method for solving the problem of generalisation to new (but similar) domains in RL. 

* **Relevance**. This is an important problem to solve for the real world deployment of reinforcement learning agents.

* **Empirical results**. The results provided indicate that the authors' method performs well.


## Weaknesses

* **Clarity**. Throughout the paper there were notation and wording choices which hinder clarity, in my opinion. I list them here:
    1. The phrase "ex-ante" is unnecessarily fancy, for lack of a better word. A more commonly used term would be better. 
    2. The statement "... with different forward dynamics, so long as the distinct dynamics maintain *some semblance of similarity*" is very vague. Could it not be replaced with something more concrete? 
    3. Regarding notation in Sec 4.1. 
         A. In some cases G_\theta has two arguments, the observation and action, and in other cases it only has the observation. Which is correct?
         B. (Related to A), should the text after eqn 2, also mention a mapping similar observations *and actions* to close latents?
         C. Perhaps it would be helpful to write \hat{P}_\phi rather than P_\phi in eqn 2, to help distinguish P_\phi from the P in the MDP?
    4. \tilde{z} and D are not introduced before eqn 3. 
    5. D_{dist} and D_{non-dist} in eqn 4 are never explicitly described. 
    6. In eqn 5, it is not clear why J_inv is within the arg max over D, given that it does not depend on D (i.e. D does not appear in eqn 2)?
    7. Section 4.2 would be much more clear if written in an algorithm environment rather than as paragraphs of text. 
    8. In figure 4, should 2. not contain P_\phi rather than D?  
       

* **Experimental evaluation.**
    1. It seems to me that the experimental evaluation in this paper is not very comprehensive, only showing two main experiments and a single ablation. I would expect to see a wider range of experiments in an ICLR paper introducing a new method. Otherwise, how can readers know whether the method is generally applicable? 
    2. Regarding the modifications to DistractingCS, it would be great to see the effects of the various combinations of 2 distractions.
    3. Why was the comparison to PAD not also made in Figure 5? It would be informative to compare how the degree of distraction impacts PAD. 
    4. How many repetitions were made to get the confidence intervals in Table 1? 
    5. The proposed future work of exploring the effect of exploration strategy seems like it should be included in this paper, in order to paint a full picture of the proposed method. 

## Other comments/questions for the authors

1. A relevant but missing citation/discussion in the related work (specifically in **Invariant representation learning**) seems to be World Models (Ha and Schmidhuber, 2018).

2. (nit) Regarding the slight abuse notation for eqn 1, there seems to be enough space to write things correctly. Furthermore, the explanation for the abuse of notation takes more extra space than the correct notation would.

3. Regarding the discussion of data augmentation techniques in the paragraph above Sec 4. While it is true that even extreme data augmentation techniques would fail to cover the test MDP observation space, it is not clear that data augmentation does not ultimately play the same role as *invariance through inference* albeit in a less direct manner. That is, data augmentation should force the policy learned to be invariant to the noise from the data augmentations and thus be more likely to generalise to the test MDP setting.

4. Could the authors elaborate on their choice of the GAN framework as opposed to VAEs or some other deep generative model?

5. How much of a challenge was it to get the adversarial training to converge in the various experiments presented in the text. GANs are well known to be difficult to train and often require many tricks to converge reliably. This was not touched on in the manuscript and seems an important detail for the general applicability of this work. 


## References

David Ha, Jürgen Schmidhuber:
World Models. CoRR abs/1803.10122 (2018)

summary_of_the_review: While the authors' proposed method seems like a promising solution to a challenging and important problem, this paper is let down by a lack of clarity in presentation, and a narrow range of experimental evaluation. 

++++ Post-revision update ++++

The authors have largely rewritten the paper which has resulted in a much clearer presentation. They have also expanded on their experimental results and clarified that the scope of their experimental work is larger than I initially understood. With these two points in mind, I have increased my score from 5 to 6. Unfortunately, I do still feel that the experimental evaluation could be more comprehensive, which is why I have not increased my score further. Additionally, I have not had the time I would like to go through the heavily revised paper, thus I have decreased my confidence.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: 
**Introduction:**

In this paper the authors aim to tackle the problem of learning a model that generalizes well on a test distribution that samples outside of the training data distribution.  They aim to do this not via data augmentation but by using an inference model over the latent space of the input.  Such an approach needs to go beyond in-distribution generalization. If the target domain is not know a priori things like domain transfer, domain randomization and meta-learning techniques may fail.

Out-of-distribution generalization is very important but more challenging:
 * Prior experience will be critical
 * Lack a reward function over the test distribution
 * In robotics for instance, this type of learning is critical

The authors propose to: *"recast out-of-distribution generalization problem into an unsupervised policy adaptation between two MDPs that share a similar latent dynamics and reward structure, but with distinct observations"*

Statement of purpose for the paper: *"Harness probabilistic inference to produce invariance by taking advantage of latent structure"*  This work attempts to distinguish itself from other approaches that *"bake policy invariance"* in at training time.


**Problem Formulation:**

Given an MDP: M = <S, O, A, R, P, γ>, S contains ground truth info not always accessible to agent.  In particular two distinct MDPs: A train MDP M_train and a test MDP M_test.  The agent has no access to the reward function at test time.  The problem is defined: *"Often, we are most interested in the case where shift between M_train and M_test is induced by differences in the observation spaces O_train and O_test. That is, the state and reward structures between the train and test MDPs are quite similar, but the observation spaces between the two MDPs are significantly different."* It is noted that recent work in RL has tried to deal with this via data augmentation but this can lead to instability.

**Invariance Through Inference and Algorithm**

The goal here is to learn an encoder mapping semantically similar states in the train and test distribution to latent vectors.  This is to be accomplished given there is no access to rewards at test time and no access to paired trajectories (ie. matched test/train trajectories). The approach is defined by two objectives:
 1. a distribution matching objective encourages the latent distribution induced by both MDPs to match. 
 2. a GAN-style loss to help ensure that the latent code is not using distracting information that would help distinguish the train and test MDPs.

The authors state that *"This adaptation objective is effectively maximizing the mutual information between the prototypical representation kept from during training, and the adapted representations of the new observations"* and further claim that *"our adaptation objective produce invariance by making statistical inference in the latent space."*  As regards the algorithm they also state that: *"The encoder G takes an observation from target environment, and learns to trick the discriminator, while the discriminator predicts whether the input comes from the encoder or source buffer."*

**Experiments**

The authors us the Deepmind Control Suite (DMC) as a training environment and the Distracting Control Suite (DistractingCS) for a test environment.  *"DistractingCS adds three types of distractions to the DeepMind Control Suite through deviations in background, camera pose, and color."*. Results are presented showing that baseline methods augmented with ITI show much slower performance degradation as the distraction intensity is increased.

The authors also compare ITI to Policy Adaptation during Deployment (PAD).  They compare to ITI on DMC for fixed distraction intensity.
ITI proves to yield stronger results in this case and it is posited that PAD does not encourage latent train and test distributions to match as ITI does: *We suspect this instability is caused by the large deviations in the latent variable distribution as a result of
changes in the target environment. In particular, we posit that the signal from PAD’s inverse dynamics head does not encourage the latent train and test distributions to match, which is a feature specifically baked into Invariance through Inference.*

main_review: **Strengths:**

 * The approach itself is novel and relies on leveraging the model's ability to infer invariance via an adversarial loss on latent representations from source and target distributions. This provides a more principled way to model environment invariance and can possibly scale well in comparison to data driven approaches. 

 * The approach does not rely on rewards from the target distribution.  However, it is unclear whether this is typically the case for the other baselines.

* The results showing decreased degradation for distraction on DMC (Fig. 5) look good and are consistent across baselines and distraction dimensions.  The results against PAD are also significantly in favour of ITI.

* A nice ablation demonstrating the effect of the adversarial objective.

**Weaknesses:**

* The evaluation domain seems to be narrow and focuses on robotics control domains or environments that simulate that.  How well does this approach work on other domains?  Are there any theoretical bounds or existing results that might give us confidence that this approach scales well across any out-of-distribution evaluation scenario?

 * The paper focuses on comparisons to data augmentation approaches to modeling out-of-distribution data.  Are there other approaches that might be considered?  Could combinations of this approach and data augmentation lead to even more robust adaptation to out-of-distribution estimation?

 * A more formal description of an adaptation and the criteria for being out-of-distribution would be very helpful.  In general, in sections 1 and 3 clearer explanation of the problem and definitions would lead to more clarity around the problem that this model proposes to solve.  For instance, what are the failure cases of generalization that this approach aims to solve and what in general are common features of a model with robust generalization?  What if we simply have access to large amounts of more varied data, do the gains of ITI remain strong in such scenarios? 

 * First sentence of the last paragraph of intro, it is stated "this challenging scenario".  Can you be precise and explicit here?

 * The *Algorithm Overview* section seems like it could us more detail. An algorithm box could be useful here.  Do we use training samples only to train the encoder?

* Both the ablations and comparison to PAD are done over a fixed set of hyper-parameters.  It would be helpful if these conclusions were better supported across more variation and/or datasets.

* In the experiments section you emphasize that these methods "expand support for the training distribution", could you contrast these methods with each other?  Do they have access to reward of the target distribution?

* Table 1 doesn't seem to describe what is being compared.  Could you add more description here?

* When discussing the comparison to PAD, it 's stated that the hypothesis is that it does not encourage latent train and test distributions to match as ITI does, could this be described in more detail wrt to PAD itself?  Just a brief description would help lend insight as to why ITI performs better.


summary_of_the_review: 
While I believe that the ideas of this work are novel and very interesting and provide a potential pathway away from data augmentation method for out-of-distribution generalization I have two issues with the paper: 1) the evaluation domain may be too narrow to provide confirmation of the approach with confidence over other methods, 2) more clarity and details of the approach would be helpful in understanding this work, how it contrasts with others, and where it exceeds those. 

Post Rebuttal:

My initial two main concerns with this work: 1) how well the approach generalizes and 2) the clarity of the paper. The updated draft addresses the second concern very well and so I consider that resolved. The sheer volume of updates here is impressive. For the first point, the authors have carried out a large number of new experiments which strengthens the case for control domains but to me the overall claims of the paper seems to be more general than that. That said, I believe this is a worthy contribution given the efforts put forth by the authors and will also increase my score to a 6.  Great work!
review_title: null

final_score: rating: 5: marginally below the acceptance threshold
review_content: summary_of_the_paper: In this paper, the transfer of a reinforcement model (RL) setting from an idealized (training) environment to a more realistic environment with distractors in the observations is considered. Instead of augmenting the training environment with more data so as to make the system more resilient to variations and distractors, the system is adapted at test time to be invariant to the specific distractors found in the environment. Experiments in simulation show the benefits of the proposed approach. Crucially, the agent is not able to access any reward data at test time.
main_review: The paper devotes a lot of space to motivating the setup and justifying its value. I do believe that this setup is relevant in practice in robotics applications. There are indeed many cases in which sim2real faces generalization problems. Unfortunately, the value of this approach is not actually demonstrated in a sim2real application, but in transfer between simulated applications.

While the goal of the paper seems very worthy, and the proposed algorithm seems to provide empirical benefits, the theoretical presentation and the justification of the approach are severely lacking, and I would say incorrect. In particular:
- Why is a distinction between S and O even made? The presentation goes beyond the standard RL presentation making a distinction between the state and the observation, and makes a point of saying that the observation doesn't uniquely identify the state. And that the transition function P is a function of states (and actions) and not observations. To then ignore such distinction and make the transition function only a function of the observations, negating the explanation about partial observability.
- Why is Eq. (5) just adding the two losses? Is this dimensionally correct? How do we know they are balanced? It seems that the relative scaling could be arbitrary by simply changing the scale of the actions.
- What's \tilde{z} in Eq. (3)? This is never defined.
- In Eq. (4),  \bar{G} is said to be the frozen generator. However, this frozen generator is put inside the expectation over "outside of the distribution samples". Isn't the frozen generator used for the training samples, therefore inside the distribution? There's no proper definition of what "distr" and "non-distr" means.
- At the beginning of section 4.1, the case is made for a single generator being used both for training and testing data, and that this generator should ignore the distractors. However, Eq. (4) and the rest of the paper uses different generators for training and testing.
- Figure 4.2 Are you sure this is a D? Shouldn't this be a P_\phi?
- Most importantly, the use of the GAN to match the distributions over z is completely unjustified, see more below.

The need for the GAN: It is unclear what the GAN portion of the architecture is bringing to the table. Since the inverse dynamics head is locked and learned on the training data, making it work on the test data should result in a z that aligns well with the latent space discovered at training time. The paper contains an ablation section showing that removing the GAN makes the system not work at all. How is this possible? The authors hypothesize that "the inverse dynamics head has a tendency to simply memorize dynamics in the training and testing domains without collapsing the commonalities". But the authors had said that, after training the inverse dynamics head on the training data, this head is frozen and remains locked during adaptation to test data. How can it possibly memorize the test dynamics then? And, even though I don't understand how this overfitting can happen, if it did, shouldn't the solution be regularization of the inverse dynamics network? The account of this architecture is inconsistent and the GAN seems completely unwarranted. Very little detail about the whole process is given and there's not supplementary information or code available.

Minor:
Typo: "training environment the Distracting"
The meaning of the values in Table 1 are unclear. The legend has no information and the text says it's the reward difference between before and after adaptation, but then there's a column for before and another for after, so the numbers themselves are not the difference. What is it?
summary_of_the_review: This paper tackles an important problem, but contains technical inconsistencies, lacks rigor and provides very little detail about the proposal as a whole. The use of a GAN, a major part of their architecture, is completely unjustified, and we are told that it's necessary because the system doesn't work if we remove it... but there's no explanation as to why. No code is provided, no supplementary details are provided.

Empirical results seem good, but there's little support as to why this is the case.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes a method for learning invariant latent representation of the observations from MDP processes that share some aspects of their dynamics but differ in states.  This is useful for generlising reinforcement learning agents to wider range of variability of test conditions (i.e. test data can be out of distribution of the train data).  The invariant latent representation is derived through minimisation of mutual information between latent encoding of the in-distribution and the out-of-distribution experiences by "fooling" a GAN-like discriminator tasked at differentiating the two.  
main_review: I quite like the aim of the proposed method, and the results are impressive.  The fundamental principle of using the discriminator to obtain invariant latent representation of an out-of distribution state space of the MDP seems like a very sensible and practical idea.  The only problem with the paper is that is very high level and I can't quite figure out how it all works and comes together.  Even with Figure 3 and Algorithm overview in section 4.2, it's not clear what happens when.  Seems there are different phases to training...but I can't quite understand what happens when.  When is P_\phi fixed and when it is trainable?  What is a_t in Equation 2 when training on MDP_test?  What's z and \tilde{z} in 3?  Do D_{non_dist} and D_{distr} refer to data from MDP_{train} and MDP_{test}?  How big does the source buffer storing training data need to be?  Does it store all training data?  
summary_of_the_review: I don't think it would require a huge amount of work, but a bit more meticulous explanation of the loss functions and the training process is needed.  At the moment it seems that only those intimately familiar with the method and code (that is not provided) would be able to implement it.  The direction seems good, results are very encouraging, just need to flesh out the actual method a bit more, otherwise it feels there isn't enough information to reproduce.
review_title: null

Meta-review: Meta Review for Invariance Through Inference

The motivation of this work is to address the problem of learning a model that generalizes well on a test distribution that samples outside of the training data distribution. Reviewer X3P2 wrote a good summary of the paper:

In this paper, the transfer of a reinforcement model (RL) setting from an idealized (training) environment to a more realistic environment with distractors in the observations is considered. Instead of augmenting the training environment with more data so as to make the system more resilient to variations and distractors, the system is adapted at test time to be invariant to the specific distractors found in the environment. Experiments in simulation show the benefits of the proposed approach. Crucially, the agent is not able to access any reward data at test time.

Reviewers, including myself, recognize the novelty of the approach, in particular appreciate the authors' motivation to provide a more principled way to model environment invariance that can possibly scale well in comparison to data driven approaches. However, the initial round of feedback is generally negative, in particular, most reviewers raise concerns regarding the lack of clarity in presentation, and also have issues with the narrow range of experimental evaluation. Clearly, this is promising work, but possibly had to be rushed for submission.

To the authors' credit, they devoted substantial efforts to completely revamp their paper, addressing many of the issues head on. The resulting updated manuscript is almost a complete rewrite of the paper. All reviewers acknowledge (and praise) the effort from the authors' to improve the paper, and 3 out of 4 reviewers had improved (or maintained) their scores from rejection to a 6. But as the paper is a complete revamp, reviewers did not have the time to assess the entire rewrite of the paper (it's like the need to review a paper from scratch), so the confidence is reduced.

While X3P2 did not change their score, they did lead a discussion amongst myself and other reviewers, and they spent the time to take a detailed look at the completely revised draft. Here are the comments from that discussion, for full transparency:

---

*The paper has been completely revamped, to the point in which the presented technique is actually different (the dynamics loss now includes a new forward term). The changes are overall welcome since it significantly improves the clarity of the presentation. Experiments still show promise.*

*I still have problems with the theoretical aspect of it, though. I think that it is unclear why the proposed system is working and fails to provide the minimal system that works.*

*Equation (4) is dimensionally incorrect. It's summing squared error over actions with squared error over latents. Both of these are arbitrary units that can lead to the forward or backward losses dominating. It's also unclear why both of these losses are necessary and not just the inverse one.*

*Equation (7) is similarly dimensionally incorrect. Again, units are arbitrary and for all we know the joint loss could be ignoring the dynamics loss or the adversarial loss.*

*The use of a GAN and the corresponding loss is unjustified. As the authors acknowledge, if the dynamic loss is very small, then the system should already work. They argue that finding that parameterization without the adversarial loss is challenging. There's a difference between using the right loss and finding the right way to optimize it, but here it seems that the right loss is being modified for optimization purposes. Is the adversarial loss something we really want to minimize or just something that helps find the best dynamics loss? What if we remove the adversarial loss after being close to convergence? What if we use multiple restarts or other techniques to help with the optimization of only the dynamics loss? My take from the theory is that the adversarial loss term shouldn't be needed and that the challenging optimization problem should be addressed (rather than modifying the loss).*

*The new ablation experiments are also confusing: If the dynamics loss is the actual driver, and the adversarial loss only helps with finding a good solution, how come that we get almost equally good results when we remove the dynamic loss? Matching the latent distribution shouldn't be enough to have aligned latents. Maybe there's something about the architecture of F that matches the ground truth, so that matching the distribution aligns the latents. This hints towards the adversarial loss actually playing an important role beyond helping with the optimization problem. This is not supported at all by the theory, since matching distributions should result in arbitrary latents and potentially performance of a random system. In fact, given the problems with units, the joint loss might be dominated by the adversarial loss, which would explain this result.*

*It seems like there's something here, but I think more work is necessary to really understand which pieces are necessary in this system and whether there's some sort of adaptation between the experimental setup and F that would explain why distribution matching results in latent alignment, which is not expected (Zhu et al. 2017). Also, the units problem makes the ablation results even harder to interpret: Maybe the dynamics loss is playing a small role in the joint loss, and that's why removing it completely has a small impact.*

---

After much assessment, while I do find this work to be interesting and potentially highly impactful (since they introduce an alternative approach to data-centric one OOD), the final manuscript's assessment is still borderline (the reviewers all mentioned that while they recognize the improvement, they list issues from preventing their full endorsement), and X3P2 still found several issues with the revision (which I do believe can be addressed in due time). While I'm fully confident that with additional work, this paper could have the potential to be an impactful one, I am currently on the side of not recommending it for acceptance for ICLR 2022.

Note to the PC's, that this is a borderline decision. If the PC's want to flip the decision to an accept, and think the post rebuttal issues are small enough, I'll be fine with that. But in any case, I look forward to seeing a further improved version (of the revamped manuscript) published in a journal or presented at a future conference. Good luck!
Decision: Reject