Venue: NeurIPS 2021 Conference

Reviews:

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: This work studies the partial observability problem of training a policy to adapt to new, unseen out-of-distribution tasks at test time. The authors focus on the setting where the dynamics across all tasks are the same but reward changes. They propose to train an ensemble of latent dynamics models, policies, and reward decoders and train an additional mixture version with dropout. During training, they replace the real reward with an imaginary reward generated by the mixture decoder to create new imaginary tasks. They show that in gridworld and MuJoCo environments, their method generalizes better than other Meta-RL methods. 

main_review: To really evaluate these results requires more than 4 seeds. The variance seems quite high in Figure 3. Given the importance of dropout to prevent memorization of train tasks, another good baseline to check would be RL^2 and/or VariBAD with dropout as well. Perhaps that is all the regularization needed to get better generalization performance? You mention that LDM test time behavior is to first try all the train goal points before exploring for new test goal points. Given that the other recurrence-based baselines also have > 0 reward at test time, do they do something similar?

Overall it seems like a nice idea but I’m not convinced by the results. More analysis of the types of problems that LDM would work well vs. not well for would also be appreciated. All the examples chosen are of tasks with more “continuous” context spaces. I assume LDM fails in tasks with only discrete structure? What types of test tasks would this fail at? It seems like there is some notion of a convex hull over training environments that this captures. Requiring a strict separation of train and test seems like a limitation? Does it need to be strict? Does the method not work if it is not?

I am currently giving this work a borderline reject which I am happy to raise to an accept contingent on some analysis and discussion of limitations and assumptions being made about the structure of the environment and results with 10 seeds. 
limitations_and_societal_impact: Yes.
review_title: null

final_score: rating: 7: Good paper, accept
review_content: summary: The authors attempt to solve a problem in context-based RL, where a model tries to learn and identify the context within which it is behaving. This context can be defined by differences between MDPs, such as differences in the transition dynamics, and the reward function; the agent will make use of such a context in its policy. 

This work builds heavily on the variBAD, as is mentioned in the paper. As its main contribution, it introduces a so-called Latent Dynamics Mixture model where during training, latent context values that are seen over the sampled training tasks are used (additively) to generate new, 'imaginary' tasks, to augment the training. 

This method aims to improve on the problem of generalising to unseen tasks that are generated from an unseen task distribution. 

main_review: ### Originality 
- This paper builds on variBAD, but I believe has sufficient novelty. I haven't seen similar work otherwise. 


### Quality
- I think that experiments that the authors had run, do support their claim that generalisation to unseen tasks from an unseen distribution works better with their models. But some questions do arise. See (Misc. Comments for some questions I have). 

- Are 4 seeds really enough for any meaningful statistical significance?
 
### Clarity
- I think that this paper is written sufficiently well. I would have liked to have seen the equation of the loss they are using, written in terms of their notation. It would help with my understanding of what is going on, and the interaction between the different components of this. 

- That being said, Figure 2 was quite a nice summary of their solution.  

- I appreciate the inclusion of the code, and what appears like sufficient hyperparameter details of how to reproduce this work in the Appendices. 

- Line 40: I think that the language here could be clarified a little. When I first read this, I thought that the policy would be able to learn new 'actions' for its policy. Perhaps more clarity about which environment this is in, and more precise descriptions would be useful. 

- Line 249: Repeated definition of H+ is repeated as in Line 161. Also, I'm not certain H+ is described or used anywhere substantial. 

- I enjoyed Figure 6(b) and (c), as they are nice visualisations the different ways the methods behaved. 

### Significance
- I think that this work is quite interesting. This could be a nice solution for OOD problems in RL as discussed in the conclusion of the paper. 

  But, some questions arise from this work (see Misc. Comments below) that I think should be addressed first. 

### Misc. comments
- Generally speaking, I am confused about the use of the word "Dynamics" in LDM. Is this referring to the dynamics of the MDP? Or is it referring to some dynamics of the latent context variable? 

-  I am also confused as to how exactly this method would help generalise to new tasks. As I understand it, a main novelty is the generation of the imaginary tasks, which are generated as a convex combination of the latent context values from the training tasks. Now, if the claim is that methods that don't use such a generation of imaginary tasks do not generalise well because the latent values are generated only over $M_{train}$, why is there an expectation that a convex sum of such $m$ would result in meaningfully new tasks?

- In fact, from Figure 4, and the description given in the paragraph starting at Line 268, it seems like the dropout plays a key role in generalising to unseen goals. So do the imaginary tasks do what is claimed?

   From the previous submission to ICML 2021, the authors had added Figure 4 to address the following issue:
   
   "We included additional empirical results regarding the tasks generated by LDM in Figure 4. We added a paragraph describing that LDM indeed generates meaningful new tasks that help to solve unseen test tasks."

   I'm not certain the new figure answers this, due to the effects of dropout. 

- What is the significance of encoders for the policy network and latent dynamics network not being shared? If I understand correctly, the policy uses a latent embedding too yes? Couldn't it use the embedding learnt from Encoder-v? Why are they separated? Could there be a comparison where the encoders are shared, as an ablation?

### Conclusion
I think that this work is interesting. The clarifications of the questions above would help my confidence in the work. 
limitations_and_societal_impact: I think that some limitations are mentioned (for example the fact that dropout is important for the imaginary tasks to actually look at unseen goals). I think some more discussion of why this is would have been good. See Main Review for more details about this.
review_title: null

final_score: rating: 6: Marginally above the acceptance threshold
review_content: summary: This paper tackles out-of-distribution generalization in meta-reinforcement learning. The proposed method, latent dynamics mixture (LDM), extends meta-RL methods (e.g. variBAD) that support inference and generation of a reward function by interpolating between inferred latents of training tasks to generate new reward functions. The intuition is that "filling in" areas of the task space not supported by the training task distribution corresponds to shifting towards the test task distribution. 
main_review: ## Originality
The technical contribution is very limited: the delta over variBAD consists entirely of obtaining a convex combination of latent vectors for decoding.

## Quality
The experiments are reasonable, but there is a lack of comparison to methods that also attempt to tackle out-of-distribution generalization, aside from MixReg-based RL2. All other comparisons are made to meta-RL methods that assume iid training and test distributions. Besides quantitative results, informative qualitative analysis that investigates the impact of the proposed method are presented. 

## Clarity
The paper is well-written. 

## Significance
While this is a cleanly executed paper, I am not sure it will have a significant impact. The technical contribution is undisputably incremental, and the resulting improvements marginal. Given that LDM implicitly assumes that test tasks are "interpolations" of training tasks (this is reflected in the design of the experiments) and that the latent space is regularized by the variational inference training procedure, it's also entirely unsurprising that it works. 

## Summary
In essence, my main concern is that the contributions are not sufficient to warrant a NeurIPS publication. However, I look forward to discussing this point with everyone.

## Post-rebuttal update
The authors have clarified via ablations that dropout on the decoder input and separate encoders are important for LDM. They have also expanded their experiments by comparing to meta-learning methods that address out-of-distribution generalization, and on more environments. I have improved my score to a 6.

I still maintain similar concerns to Reviewer 1a3x regarding an implicit "convex hull"/"filling in" assumption underpinning LDM. While the authors have acknowledged this point in their rebuttal, I encourage them to develop this notion further in future revisions. It would make the paper stronger.
limitations_and_societal_impact: This is adequately done.
review_title: null

final_score: rating: 5: Marginally below the acceptance threshold
review_content: summary: This paper focuses on the field of meta-reinforcement learning (meta-RL), where the authors propose a method to augment the tasks for meta-reinforcement by applying a mixture of learned reward functions. Specifically, the paper builds on the framework of variBad [1], which uses a VAE to learn a latent dynamics and reward model by encoding the data and decoding the next states and rewards. The authors propose a method that uses a mixture of the encoded latent variables from different tasks during the training time to obtain new latent encodings corresponding to imaginary tasks, and then applies these latent encodings to the reward decoder to obtain new reward functions to train the policy on these imaginary tasks.

The authors then evaluate the proposed method on grid world tasks and several MuJoCo locomotion tasks, where the authors specifically emphasize the performance of adapting to out-of-distribution tasks. The empirical results demonstrate that the proposed method outperforms several baselines.


References

[1] Zintgraf, L., et al. "VariBAD: a very good method for Bayes-adaptive deep RL via meta-learning." Proceedings of ICLR 2020 (2020).

main_review: Overall I think this paper presents an interesting idea on task augmentation for metal-reinforcement learning. The paper is well-written and the authors include a lot of experiment details and ablation studies in the paper. However, there are also many limitations in this paper that need to be addressed.

Pros:

The main idea is simple and intuitive. Using a mixture of latent codes to augment the task is a natural idea and it is easy to believe that it should work well.

The paper is well written. The main idea is presented in a way that is easy to understand, and the whole paper is structured in an organized way.

The authors provide detailed analysis of the experiments, with many informative behavior visualizations and ablation studies about the proposed method.

Cons:

Several important baseline comparisons are missing. The main advantage of the proposed method is the capability of adapting well to out-of-distribution tasks, where the training tasks and test tasks come from non-overlapping distributions. As the authors mentioned in related works, several prior works such as MQL [1], MIER [2] and FLAP [3] have attempted to address the out-of-distribution problem of test time tasks. Since this paper addresses the same problem, it would be important to compare the proposed method to these baselines.

The scope of experiments in this paper is rather limited. Besides simple grid world environments, the authors only include 3 MuJoCo locomotion environments, which I don’t believe is enough. In order to add more environments, the authors could consider out-of-distribution environments proposed in [1, 2], and also MetaWorld [4]. 

In summary, being an empirical paper, this paper does not include enough experiments to convince me about the performance of the proposed method. Therefore, it is hard to recommend acceptance of this paper at the current state.


References

[1] Fakoor, Rasool, et al. "Meta-q-learning." arXiv preprint arXiv:1910.00125 (2019).

[2] Mendonca, Russell, et al. "Meta-reinforcement learning robust to distributional shift via model identification and experience relabeling." arXiv preprint arXiv:2006.07178 (2020).

[3] Peng, Matt, Banghua Zhu, and Jiantao Jiao. "Linear representation meta-reinforcement learning for instant adaptation." arXiv preprint arXiv:2101.04750 (2021).

[4] Yu, Tianhe, et al. "Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning." Conference on Robot Learning. PMLR, 2020.

limitations_and_societal_impact: See main review for limitations and N/A for social impact.
review_title: null

Meta-review: This paper proposes to learn an ensemble of latent dynamics models and provide a mixture of them as input to a reward decoder, which allows the entire framework to generate imaginary tasks (i.e., imaginary reward functions). By training the agent on the generated tasks, the agent can generalize better to out-of-distribution tasks on grid-world and MuJoCo environments. The idea of generating imaginary reward functions using a mixture of latent dynamics model is novel, though the overall framework builds upon VariBad. Although there was an initial concern around the significance of the results (e.g., lack of baselines, challenging environments), the additional results during the rebuttal period addressed most of the questions raised by the reviewers. However, the question about how the proposed framework can extrapolate so well and how the mixture hyperparameters ($\beta$) affects interpolation/extrapolation performance remains the same, which may need to be discussed in the main paper rather than in the appendix. Given that the main result is still solid and consistent, I'd recommend acceptance for this paper.
Decision: Accept (Poster)