Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper studies multi-task generalization in model-based reinforcement learning (MBRL). One of the standard ways to approach the problem is given by inferring a latent variable Z encoding each task and then conditioning the dynamics model on it. The authors propose to encode segments of state-action trajectories into Z vectors and maximize similarity between Zs from the same trajectory. Using the mechanism, the paper then studies the ability of the agents to generalize to unseen variations of the training environments. For example, if the agent was trained on MuJoCo HalfCheetah with body mass 1, it is asked to generalize to body masses 0.5 and 1.5. The paper compares the method against the baselines across 3 axes: in terms of dynamics prediction error, in terms of returns on testing environments, and in terms of separability of the inferred Z for different environments.
main_review: Novelty

To the best of the reviewer’s knowledge, the paper is fairly novel. The authors use CaDM and TMCL, other model-based approaches that aim at generalization, as baselines. Figures 10-13 suggest that TMCL fails to learn separable encodings Z, while Figure 9 suggests that CaDM achieves lower returns. The proposed approach combines the best of both worlds achieving comparable returns and learning Zs that are clustered for the same domain.

Significance

The reviewer is unsure about the effect size of the proposed method. First of all, the experiments were repeated using 3 random seeds only. Combined with the overlapping confidence intervals of returns on a fair amount of test environments (e.g. Fig 5), the significance of the results is unclear. Increasing the number of runs would improve the credibility of the conclusions. For a discussion about the importance of rigorous experimentation in RL, see [1].

On a positive side, the paper follows the setup from previous papers and performs an ablation study of the proposed modification. The reviewer also appreciates qualitative findings about separability of learned Z in Figures 1 and 7 and quantitative evaluations of model prediction errors in training environments in Figures 4 and 8.

Detailed comments and questions

- Figure 1 demonstrates the linear projection of context vectors Z obtained by different methods and serves (at least partially) as a motivation for the paper. Would the conclusion that Zs for TMCL are not separable hold for a non-linear method for visualization (say, t-SNE)?
- In Table 3, C_HalfCheetah row, CaDM has a lower error but the number for the proposed method is bold.
- Conditioning $p(s’|s,a)$ on $Z_{t-k:t-1}$ makes the dynamics non-markovian. Would it be possible to use an RNN-like baseline in this setting?
- Could authors present the plots for prediction errors on *testing* environments?
- What would be a relevant model-free baseline in this setting?

[1] Henderson, Peter, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger. "Deep reinforcement learning that matters." In Proceedings of the AAAI conference on artificial intelligence, vol. 32, no. 1. 2018.

-------------

POST-REBUTTAL UPDATE:

I update the score to 6 and confidence to 4
summary_of_the_review: The reviewer leans towards recommending the submission for rejection. The ultimate objective of an RL system is maximizing the returns and it is unclear from the empirical evaluation to which extent the proposed method improves the performance of the agent on unseen tasks. However, the reviewer has limited familiarity with the related work on causality and it is possible that they missed something. Addressing the outlined concerns might improve the overall score.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper considers the unsupervised dynamics generalization problem. In this problem there are a set of train MDPs and a set of test MDPs, all with the same state and action spaces, but with different dynamics functions. The authors build an approaches that use past transition segments to estimate an environment-specific vector $Z$, which is then intended to act as a contextual input to the dynamics function. The problem the authors identify is that a naive application of inferring $Z$ from past transition segments undermines the ability for Z to be similar in the same environment and dissimilar in different environments. The authors claim that generalization ability is hurt by the inability to preserve this property. The authors propose a loss (equation 6) that pull $Z$'s from the same trajectory together and push away $Z$'s from different trajectories together. However, to avoid not pushing away different trajectories from the same environment, the authors propose to estimate the controlled direct effect (CDE) of the $Z$ on the next state $S_{t+1}$: if $z^i$ and $z^j$ from trajectories $i$ and $j$ have a similar CDE, then the authors declare $i$ and $j$ to be likely from the same environment. This likelihood is used as a weighting coefficient on the loss: the authors use this weighting coefficient as a "soft" indicator variable to distinguish examples from the same environment (positive examples) or examples from different environment (negative examples). Figure 1 shows that a PCA visualization of the learned context vectors using the proposed method separate the context vectors into distinct clusters. Empirical evaluation on standard benchmarks shows that the proposed method does slightly better than the next best baseline, TMCL from (https://arxiv.org/abs/2010.13303).
main_review: Strengths
- novel idea for using CDE for clustering trajectories
- I like the separation of the two objectives (1) and (6). (1) enforces that the z's are useful for predicting dynamics. (6) enforces that the z's within the same trajectory _and_ the same environment are clustered together. The authors show that in practice the z's do end up being predictive of the environment id.
- Figure 1 shows that the learned context vectors are clearly distinguishable

Weaknesses
- the results in Figures 4 and 5 show that TMCL very similarly to the proposed method, which seems to slightly refute the authors claim that "However, because environments are not labelled, the extracted information inevitably contains redundant information unrelated to the dynamics in transition segments and thus fails to maintain a crucial property of Z: Z should be similar in the same environment and dissimilar in different ones. As a result, the learned dynamics prediction function will deviate from the true one, which undermines the generalization ability." Essentially, why is it necessary to "explicitly learn the meaningful dynamics change information" instead of doing it implicitly? The empirical results do not seem to suggest that being explicit is necessarily much more helpful for generalization and performance.
- I was interested to see that in Figure 6, not having the intervention does not seem to hurt performance or generalization that much, and in some cases even improves it. Would the authors be able to conduct an analysis similar to Figure 7 for the case where the interventional prediction is not used and test whether the contexts can still be predictive of the environment? 
summary_of_the_review: It's not clear how much the interventional prediction aspect of this work actually helps (see Figure 6), even though this seems to be the main proposed contribution. It is also not clear what advantage that the proposed method has over TMCL, other than the ability to qualitatively visualize distinct cluster's in the z's. But if this does not hurt performance or generalization (Fig 4 and 5), what would be the reason that the interventional prediction is crucial?

== After Rebuttal ==
The quantitative results seem to still indicate a marginal improvement over the baselines, but perhaps what is really required is to identify environments that the community cares about that would more drastically distinguish the proposed method from those that do not use the proposed intervention prediction method. However, given the current experiments, the authors have adequately addressed my concerns and therefore I have raised my score to a 6.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This paper proposes a new method to tackle the problem of generalization to unseen dynamics for model-based RL. Previous methods learn to predict a vector $Z$ that characterizes a particular environment dynamics from past transitions. However, as the environment id or label is not available, this vector inevitably contains redundant information, which might hurt the generalization of the model. The paper therefore proposes an interventional approach to estimate the probability that two vectors $\hat{z}_i$ and $\hat{z}_j$ belong to the same environment, and then uses a relational head to force similarity between them. As a result, the redundancy in $Z$ is reduced which leads to improved generalization.
main_review: #### **Strengths**
- The paper tackles a very relevant and important problem in the RL community, where generalization to changing dynamics is one of the key aspects to enable RL algorithms for real-world applications.
- To the best of my knowledge, the novelty of the proposed method is significant. While other works have used similar ideas to characterize the environment dynamics using a vector $Z$, the paper presents an interesting idea to reduce the redundant information in $Z$, or in other words, learn a better representation of $Z$. The interventional prediction, while being simple, is elegant and works well in practice. I believe this interventional approach would open some new directions to improve current unsupervised representation learning methods, which use instance-based discrimination objectives (similar to (2) in the paper).
- The improvements of the proposed method are significant compared to the existing baselines. Visualizations such as Figures 1,7,10-13 empirical show that the proposed method actually learns semantical information of $Z$.
- The paper is well written and organized.

#### **Questions**
- What happens to the representation of $Z$ when we omit the interventional module? I would love to see a visualization similar to Figure 1. My wild guess is $Z$ will be spread out?
summary_of_the_review: Overall, I like the idea of the paper and believe it has significant contributions to the RL community.
review_title: null

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: This work tackles the problem of unsupervised dynamics generalisation in model-based reinforcement learning, improving approach introduced by additional of inferring a latent context which conditions the dynamics model and captures the variation in dynamics between environments. They introduce a set of auxiliary losses based on relational intervention and causal reasoning to encourage the inferred context to be the same in trajectories from the same environment (even when the environment identification is unknown) through inferring which environments are the same and weighting the context-similarity loss by this environment similarity. The method shows improved prediction error and test reward on a range of continuous control tasks compared to SOTA baselines from MBRL, Meta-RL and dynamics generalisation MBRL specifically. Visualisations qualitatively show that their method clusters transition segments from the same trajectory together better than previous methods, which could explain their improved performance.
main_review: # Strengths
- The problem statement identified by the paper is clear and compelling, which motivates the paper well. 
- The method is interesting and novel, and seems well-designed to tackle the problem being addressed
- The empirical results and visualisations both support the claims that this method is performing better than previous SOTA methods due to the better clustering of contexts.

# Weaknesses
- The notation in the last paragraph of page 2 is confusing - what's the difference between Z and \hat Z?
- In general, I found the explanation of the method difficult to parse at first - it took me several passes to put all the pieces together, understand the justification for each piece and how it fulfilled it's role. I think that's partly due to slightly confusing notation (see above point) which isn't well defined, and also due to the complexity of the method. In some sense this can't be avoided, but it's still worth mentioning here.
- The biggest weakness of the paper (although I don't think it prevents acceptance at all) is that the method is a fairly incremental improvement to previous methods, and hence isn't highly novel. There's not much that can be done for this, unless additional benchmarks were created and assessed on (which is clearly a lot of work).

# Questions
- In the PCA visualisations, how are the different dots of the same colour generated? Is from different subsequences of transitions from a single episode, or multiple episodes, or both? It would be interesting to see how the context prediction evolves over an episode, and whether it converges over time as more information is gathered within the episode (as we might expect it to).
- The relational encoder take in a fixed number of past transitions, is that correct? How does this work for time-steps near the beginning of the episode, where not enough transitions are available? Have you investigated using a recurrent encoder which processes all the information in the episode?
- What's L_dist in the loss function in Algorithm 1?

# Suggestions for Improvement
- AugWM missing from related work: https://arxiv.org/abs/2104.05632, and seems relevant
- It would be interesting to visualise clustering of masses which are in testing distribution as well, to see how well the relational embedding generalises.
- The justification that their method improves through better clustering is only supported qualitatively through visualisation. It would be great to see a quantitative measure of how good the clustering is in the latent space for different methods, and then seeing if this correlates with evaluation performance or prediction error across the range of methods and environments (or even across hyperparameter options for each method).
- It would be good to see a plot of average prediction error on test environment during training, to see how that changes as the methods train.
- You provide some intuition as to why the ablation with no intervention has lower prediction error but lower performance in 4.3, but you could test this explicitly by plotting the test prediction error and seeing whether it's higher for the no-intervention ablation
- PCA is quite a limited visualisation technique for clustering - have you considered other techniques such as t-SNE or UMAP?
summary_of_the_review: I think the paper is well-motivated, presents and interesting and novel method and demonstrates that method's improved performance on the problem being targeted, and hence deserves acceptance. I think the visualisations and ablations are useful qualitative evidence for the benefit of the method, but additional visualisations, or more quantitative investigation of the clustering hypothesis would be beneficial. Also, the explanation of the method is sometimes difficult to understand. The method is only an incremental improvement on previous work, which is necessary to perform but limits it's impact and novelty.

EDIT: After seeing the revisions and additional experiments, I have raised my confidence from 3 to 4.
review_title: null

Meta-review: Description of paper content:

The authors propose a dynamics model that can generalize to novel environments. The train and test MDPs have the same state and action spaces but different dynamics. Environment specific inference is achieved by estimating latent vectors Z that describe the non-stationary or variable part of the dynamics. These Z-s are inferred from trajectory segments in unlabeled environments. The Z-s are learned contrastively: Z-s from the same trajectory are pulled together, and Z-s from separate trajectories are pushed apart. However, to mitigate the error of distancing Z-s from different trajectories but the same environment, Z-s on trajectories with similar transitions are also pushed together using a soft clustering penalty. These losses are justified based on ideas from Pearl’s causal inference.

Summary of paper discussion:

The reviewers concluded that the contributions are conceptually interesting and “somewhat” novel. The reviewers felt that the empirical performance gains of the method over baselines were demonstrated but not extremely impressive.
Decision: Accept (Poster)