Venue: NeurIPS 2022 Conference

Reviews:

final_score: rating: 4: Borderline reject: Technically solid paper where reasons to reject, e.g., limited evaluation, outweigh reasons to accept, e.g., good evaluation. Please use sparingly.
review_content: summary: This paper proposes a new framework (called DiAMetR) for meta RL, aiming at improving the ability of meta policy in adapting to varying levels of task distribution shift. This is achieved by leveraging the distributional robust framework to re-formulate the meta RL objective to optimize the worst-case empirical risk under a uncertainty set of test-time task distribution. To be specific, the task distribution is parameterized through modeling the latent representation of the reduced reward functions. Then, the training alternates between 1) the training of meta policy against multiple imagined task distributions, each of which has a different distribution deviation constraint, and 2) the training of the imagined task distributions in an adversarial manner.

The performance of DiAMetR is evaluated in four goal-reaching environments, demonstrating the superiority in adaptation to varying levels of task distribution shift when compared with RL$^2$, VariBAD and HyperX. Ablation study and analysis are also provided.

strengths_and_weaknesses: Pros:
+ The paper is well written and the presentation of the algorithm is very clear.
+ The related work is almost sufficient.
+ The main idea is general and novel in meta RL. Insightful formal analysis is provided.
+ The proposed method is evaluated from multiple perspectives, from performance comparison to visual analysis.

&nbsp;

Cons:

- The proposed method is kind of limited.
  - This is mainly because DiAMetR relies on the assumption that  the target tasks only differ at reward function while the transition dynamics remain unchanged. Personally, I think this cripples the significance of the proposed method a lot.
  - In addition, it is unclear how DiAMetR will perform in dense-reward environments or non-goal reaching environments. It seems that the specific model structure (i.e., structured VAE in appendix is specialized to sparse-reward environments).
- The latent representation and generative model of reward function (Eq.6) implicitly depends on the distribution of state-action pairs. In practice, the distribution of state-action pairs is generated by the policies optimized during the learning process under the original training task distribution.
  - I have no question on the worst-case objective (Eq.3 and Eq.4) when oracle imagined task distribution and task generation are considered. However, since the latent representation and generative model of reward function are trained with samples collected by changing policies under the original task distributional, I worry about the ability of the generalization of such reward function models, in turn the effectiveness and rationality, especially in more possible environments.
- Important experimental details are missing (see my questions below).

&nbsp;

There are some related works which can be included for more complete empirical evaluation and related work survey:

- MetaCURE: Meta Reinforcement Learning with Empowerment-Driven Exploration. ICML 2021
- Towards Effective Context for Meta-Reinforcement Learning: an Approach based on Contrastive Learning. AAAI 2021



questions: 1) How do the authors think about the generalization ability of the reward function models trained with samples collected by changing policies under the original task distributional? (as my comment says in the part of Strengths and Weaknesses)
2) For Line 7 in Algorithm 1, i.e., training the meta policy with imagined task distributions, how are the state-action trajectories collected? I can imagine three options: 1) directly interact with environment since the transition dynamics is the same; 2) interact with a learned transition dynamics model; 3) reuse the experiences in the replay buffer.
3) How are the set of task distribution deviation (i.e., $\{\epsilon\}$) chosen for each environments in the experiments? I did not found the associated content (I also checked the appendix), did I miss it?
4) Do the baselines VariBAD and HyperX also use finetuning with the same number of samples during test-time adaptation?
5) Although the authors provide Figure 8, it would be useful and informative to provide the analysis between the arm (i.e., some $\epsilon$) adaptively chosen by Algorithm 2 and the oracle test distribution.

&nbsp;

Minors:
- The legends in Figure 5,6,8 can be enlarged for better clarity.

&nbsp;


I may improve my rating if my concerns and questions are well addressed.

limitations: The major limitation is the proposed method focuses on reward variations while assuming transition dynamics unchanged. This is considered as one of the future directions in Section 7.

&nbsp;


I think there are three more limitations which are not well discussed in this paper:
1) The first one is the choice of the uncertainty sets (i.e., the multiple $\epsilon$). Intuitively, these hyperparameters are significant to the learning performance of the proposed method, and one may need several attempts or prior knowledge to determine the task-dependent choices.
2) The second one is that the performance of current method in dense-reward environments, or non-goal reaching environments remains unclear. I also notice that the structured VAE is kind of specialized for sparse-reward tasks.
3) The third one is my concern on the sample cost of test-time adaptation (i.e., 250 meta episode). I am a little reluctant to agree with the ‘quick adaptation’, especially when considering other possible environments with horizons longer than e.g. 50, 200 (Table 2 in the Appendix) used in this work.
review_title: null

final_score: rating: 5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.
review_content: summary: The paper proposes a meta-RL algorithm for dealing with test-time distribution shift.
This is achieved by learning a population of policies that have different levels of distributional robustness by an adversarial algorithm.
At test time, a policy is chosen amongst the population by an efficient bandit algorithm.
The paper demonstrates empirically that this way of adapting to test-time distribution achieves better performance than competing methods.
Theoretical result provides a bound on the difference of expected regret of a mis-specified robustness epsilon.

strengths_and_weaknesses: # Strengths

* The method proposed in the paper is novel (as far as I know) and seems to show improvement in comparison to other meta RL methods.
* The presentation of the method is clear, I didn't have many problems following the storyline

# Weaknesses

* Missing details in the experiment section.
* Simplified task spaces where goals are be sampled from uniform distribution and limitation to only shift in reward function as result, not transition dynamics for example.
* It is unclear to me based on the presentation how conservative the policies that are chosen by the bandit are in the end.* “the distributional meta-robustness objective can be modelled as an adversarial game between the policy and the task distribution”. Unfortunately this is the only sentence that hints at a game-theoretic treatment of the topic, but it doesn’t exist in the paper. I like that using an adversarial formulation for inferring the task sampling distribution q allows for a lot of flexibility regarding task distributions, the environments though used in the experiments to validate the method all deal with relatively simple distribution shifts concerning the reward (just uniform distributions of goals), so perhaps the tasks used for evaluation are too simple?
questions: # Major Concerns

* Equation 1 is expected regret, shouldn’t the regret be written as  - expected return of the meta-policy and the actual term in the expectation at the moment is just an MC estimate, this would seem more natural to me, because the outer expectation is taken over different tasks, not environment dynamics, so the expected return of the meta-policy doesn’t even appear in the minimization problem. 

* Fix fonts in figure 5, they are way too small. Also, RL2 and RL2 finetune legend doesn’t match the dashed lines. It’s not clear to me what is diameter oracle here? If it’s the best policy that can be chosen from the set of trained meta-policies, why is then the diameter adapt mean higher than the oracle’s mean? This figure needs better clarification.

* It seems to me that your results hinge on the fact that if you are lucky enough that the test distribution is close to the train distribution, you should be less conservative, otherwise you should be more conservative and therefore sacrifice performance. Doesn’t it make more sense to have policies that are less conservative (train on task distributions of smaller support, slightly overlapping support) and then choose amongst them? In this way there would be no ‘conservativeness tradeoff’. This is not shown in figure 7.

*  Ideally I would also like to see the distribution over policies from the population of meta-policies that get picked for solving tasks. Is it skewed towards more conservative policies? The plot should show that conservativeness grows with epsilon, currently you only have the result that picking the most conserative policy isn’t the best thing. (figure 8).

* l.141 I don’t understand why equation 3 incentivizes exploration as you say it? I understand how it leads to distributional robustness, since you want to minimize the maximum regret, which is the difference to the optimal expected return. But how does this encourage exploration? I don’t see any reason why this should be the case…

* How many adaptation policies are used in the experiments? Experiments need better clarification in general.

# Minor Concerns

* Section 6.2, by test reward distributions you mean q(z | h)? Please put inline math somewhere in this section to make the clear where the samples come from? What projection method was used to generate these heatmaps in 2d? Or is z actually 2-dimensional?

* Figure 8: fix the fonts, same as figure 5. I don’t understand how is diameter mid uncertainty set chosen here exactly, so it’s hard for me to infer anything from this result, what does it imply?

* Equation 2, why is q parametrized here? It should be just a set of distributions that are close in terms of f-divergence to the training distribution p, there is no reason why they should be parameterized.

Requirements to increase my score:

* Better explanation of the experiments section (what do the plots mean).
* Better explanation of section 6.2 (projection method, where do the samples come from).
* Argue clearly why does it make sense to have more and less conservative policies in comparison to less conservative ones with different task-space specialisations.
* Answer all the listed points



limitations: The authors should make it clear how much it costs them to learn this population of policies and how it relates to the test-time distribution shift. I imagine that the number of policies is also going to affect the performance in the end, because it's directly related to the granularity of epsilon (which controls conservativeness).
review_title: null

final_score: rating: 7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.
review_content: summary: This paper identifies and addresses the problem of the distributional shift in test task distributions for meta-RL agents, which often struggle under such circumstances. For this matter, the work leverages the distributional robustness framework to alleviate this effect, by training a population of meta-policies in different levels of distribution shift. In order to do so, the proposed method (DiAMetR) firstly learns a single meta-policy in the train distribution and uses the collected trajectories to learn a probabilistic latent space to parameterize task distributions. Then, each subsequent meta-policy is learned in an adversarial game between itself and a task proposal distribution with some degree of shift. At test-time, a bandit algorithm based on Thompson sampling selects the appropriated meta-policy (level of distribution shift) to interact with the test task. The proposed method is evaluated in 4 different continuous control environments and presents better adaptation performance in the scenarios of test distribution shift when compared to strong baselines. The work also presents ablations to validate the components of the method and a theoretical proposition to better understand a meta-policy performance under distribution shift.
strengths_and_weaknesses: Strengths:

- Addressing distributional shift for meta-RL agents is a very important topic, especially since OOD generalization is a property that the community is actively looking for understanding and improvements;
- The paper is very organized and well written, all ideas are clear and the work is appropriately placed in prior literature;
- The method is novel, technically very sound, and brings a solid contribution to the field of meta-RL towards the mitigation of distributional shift in test task distribution;
- The theoretical framework is clearly stated and the derived proposition is relevant; the experimental setups support the main claims and back up the proposed method. As a strong positive point, there is a helpful, extensive ablation and the raised questions during the experimental setup are very relevant to validating the method.

Weaknesses/Concerns:

- Using a multi-armed bandit for meta-policy selection seems to be very costly (in terms of test-time sample complexity), and the complexity grows with the number of meta-policies. For instance, the example from Section 6.1 shows that  DiAMetR requires 250 meta-episodes (or 500 episodes) to complete adaptation. Prior literature shows evidence for few-episode adaptation, sometimes online adaptation [1, 2]. I understand that these methods are not focused on distributional shift, but in scenarios where test-time sample complexity is also a requirement, this performance seems concerning.
- In the same line as the previous point, it would be interesting to show the results from variBAD and hyperX in the scenario of test time adaptation (Figure 5 and Appendix G). It would be helpful to understand how these other methods would handle this trade-off between performance and adaptation complexity for the distributional shift scenario.
- As noted by the paper, the proposed implementation is limited by meta-RL environments comprising only variations in the reward distribution, missing those ones with variations in the dynamics. This is an important concern, as there is no way to validate if the proposed method is still valid for more complex benchmarks. For instance, it would be interesting to see how it performs in MetaWorld ML10/ML45 [3].

Despite these weaknesses, the understanding is that the method and empirical validation present satisfactory quality for acceptance, given the proposed scope.

Minor issues:

- Figure 2 is not cited and actually does not seem to have any support in the text besides Algorithm 1.
- A diagram with all the modules and training/test stages would be very helpful to improve clarity and presentation. For example, something like Figure 2 in the [VariBAD paper](https://arxiv.org/pdf/1910.08348.pdf).
- In Equation 1, $J(\pi_{t})$ is not defined (Is the minus sign a typo?)
- Algorithm 1, Line 3 cites Equation 7, but this actually refers to the meta-policy selection, is this right?
- In Appendix F, the summary of proof (L575): the paper describes a 3-stage proof, but stage “1)” is repeated twice.

References:

[1] Zintgraf et. al. VariBAD: A Very Good Method for Bayes-Adaptive Deep RL via Meta-Learning. ICLR, 2020.

[2] LC Melo. Transformers are Meta-Reinforcement Learners. ICML, 2022.

[3] Yu et. al. Meta-World: A Benchmark and Evaluation for Multi-Task and Meta Reinforcement Learning. CoRL, 2019.


questions: The main suggestions or questions/concerns to be clarified are pointed out in the Strengths/Weaknesses and Limitations sections.
limitations: First of all, I appreciate that the work raised the limitation regarding handling only tasks with reward variation (as opposed to dynamics variation). It would be interesting to discuss ways to address this limitation in future work, in light of the proposed framework.

Besides this limitation, there is also another point worth discussing:

The distributional shift phenomenon can happen in many different ways. In other words, considering the training distribution, there might have several different test distributions with the same level of shift (given a divergence metric). It is hard to know if a single “adversarial game” per shift level would be able to cover enough possibilities of test task distributions.

For the presented experimental setup, this does not seem to be a problem, as the shift would happen in only one dimension (reward) that varies across very few variables (e.g, radius and angle of the goal). However, in more complex cases, this can be concerning, and adding new adversarial games incurs new meta-policies, which also increases adaptation complexity (which is already a concern as pointed out in the weaknesses). Curious to know your thoughts on this point.

Figure 7d illustrates this concern: it is possible to see that the imagined test distributions are more “skewed” in the west-south directions, but the real test distribution could be the opposite, for example.
review_title: null

final_score: rating: 6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
review_content: summary: In meta RL, the most common assumption is that the training task distribution is the same as the test task distribution, which may not hold in practice. To address the unknown distribution shift in practice, this paper proposes an adaptive method, by learning a set of policies with varying levels of distribution shift. Under a certain level of distribution shift, the policy is learned by leveraging the framework of distributional robustness.  At test time, the selection of policies is formulated as a MAB problem. 
strengths_and_weaknesses: Strengths
1. The distribution shift problem is very common in practice. 
2. The idea of learning a set of policies with varying levels of distribution shift, when the test distribution cannot be observed in advance, is quite reasonable. 
3. The paper is well-organized. 

Weaknesses:
1. The theoretical analysis is only limited to a toy environment.
2. The proposed method only parameters the task distribution by parameterizing distribution over the reward function. However, different tasks usually differ in both reward and transition functions.
3. The test-time meta policies selection is based on MAB. I was wondering if it is possible to improve this MAB method, perhaps by considering some uni-modality properties of this special MAB problem(Is it possible to approve that?). 
questions: 1. I was wondering why the authors did not capture the distribution over reward and transition functions, in their method to parameter the task distribution. It seems not very difficult. 
2. Assuming p_train(z) = N(0, I) does not seem reasonable to me. N(0, I) is prior. Even training the ELBO  only on training task distribution, the posterior is usually different from the prior. 
3. Is it possible to approve some uni-modality property of this special MAB problem? If so, you can use a more efficient method with the unimodal bandit. 
limitations: 1. The method can be extended to capture both variations in reward function and transition function. 
2. The theoretical results are limited to a toy environment.
review_title: null

Meta-review: metareview: The paper introduces a new framework for meta RL and validates it on navigation and goal reaching tasks. The authors addressed several of reviewers’ concerns in the rebuttal and significantly improved the quality of the paper. I think the paper is novel and interesting for the community. It would be interesting to see the performance of the proposed algorithm in more continuous tasks as suggested by reviewer Asdp.
confidence: Less certain
Decision: Accept