Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper studies the question of whether pre-trained model-based RL can be effectively fine-tuned to solve downstream tasks. To answer this question, during pre-training, this paper proposed to deploy a teacher-student framework: learn several sing-task teacher models via offline RL and distill these single-task policies into a multi-task student policy. Then, the multi-task student model is fine-tuned with the online target task. Offline data and gradient re-weighting are used to prevent catastrophic forgetting and performance degradation during fine-tuning.
strength_and_weaknesses: ### Strength
- The proposed question is quite interesting.
- The method is easy to follow.
- EfficientZero is a reasonable choice to study this question with offline pre-training.

### Weakness
- Major
    - Could the authors please provide more details about the offline data set used during the pre-training stage? How do you obtain the offline data set? What is the size of this offline data set? How are the quality and size affect the downstream performance?
    - Since the proposed model has an offline pre-training stage, the improvement in sample efficiency seems not significant compared with EfficientZero, as shown in Figure 6. 
    - The BC baseline is too weak, there is no policy improvement stage such as the fine-tuning used for the proposed model.
- Minor
    - Figure 1 and figure 2 seemed not being referred anywhere in the paper.
clarity,_quality,_novelty_and_reproducibility: The paper is well written and easy to follow. 
summary_of_the_review: This paper studies the question of whether pre-trained model-based RL can be effectively fine-tuned to solve downstream tasks.  However, the experiments lack important analysis of the effect of the offline data set.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes XTRA, a framework for finetuning a model-based RL agent that is pretrained on offline datasets to a target task through online interactions. Unlike previous work that uses the same task for offline pretraining and online finetuning, this paper tackles the setting where there are multiple pretraining tasks that are different from the target task. For multi-task offline pretraining, the paper shows that simply learning a single model for all tasks using RL cannot work. Hence, it proposes to first learn a teacher model for each task, and then distill all teacher models into a single student model. For online finetuning, to overcome catastrophic forgetting, the paper proposes to use offline data together with online data. To minimize interference from irrelevant tasks in the offline data, the paper further proposes a re-weighting technique based on task similarity. The experiments focus on the Atari 100k benchmark, with EfficientZero being the backbone and main baseline. There are two settings: (1) pretraining on games that are similar to the target game, and (2) pretraining on diverse games. XTRA outperforms EfficientZero in both settings. Further analysis shows that it is crucial to finetune both the pretrained representation and the dynamics function, and that XTRA has a large advantage over EfficientZero in the early stage of training.
strength_and_weaknesses: - Strengths
    - The paper tackles an important new problem that is of interest to the community. Pretraining on multiple tasks and then finetuning on a different target task is more realistic than pretraining and finetuning on the same task as done in previous work, and has the potential for better data efficiency.
    - The paper is well written and easy to follow. The method is well motivated.
    - The ablation and analysis are quite informative, answering interesting questions such as which components are beneficial to finetune and which stage of training does the proposed model bring the most improvement. 
- Main Weaknesses
    - The baselines are mainly online RL methods which do not have access to offline data. The paper can be strengthened by comparing to methods that also do offline pretraining and online finetuning. For example, it is mentioned in related work that Online Decision Transformer (ODT) pretrains and finetunes on the same task. The paper can be much stronger if it shows that when pretrained on the target task, XTRA acheives similar performance as ODT, and when including other pretraining tasks, it outperforms ODT.
    - Although model-based RL enjoys high sample efficiency in online single-task setting, it seems unclear whether it makes sense or is beneficial to use model-based RL in this paper's setting (as opposed to model-free RL). I would expect different tasks to have very different dynamics, so that it is not only difficult to learn a single model for all pretraining tasks, but it is also difficult to transfer the learned model to the target task. A comparison to model-free online finetuning would help clarify.
- Questions and Minor Issues
    - Can you elaborate on what is catastrophic forgetting in online finetuning? This can make the paper more self-contained and help the reader better understand why you want to retain the offline data.
    - There is not much ablation in the diverse task setting. It would be interesting to see if pretraining and/or retaining offline data helps, and if the task weights are close to zero.
    - I am curious why you use distillation during pretraining but switch back to the EfficientZero loss during finetuning (Eq 1). Is it possible to still use the distillation loss in finetuing?
    - It seems the improvement over EfficientZero is even larger in the diverse task setting than in the similar task setting. Can you comment on that? Also, why is the maze-to-shooter result much worse than the diverse task setting?
    - Figure 5 seems not mentioned in the main text.
clarity,_quality,_novelty_and_reproducibility: - Clarity
    - The paper is well written and organized. It clearly explains the motivation behind its design choices. The illustrations are also of high quality.
- Quality
    - The method makes sense. The experiments convincingly show that the proposed XTRA outperforms online RL that is trained from scratch without access to offline data. It would be nice to have additional comparisons as mentioned in Main Weaknesses.
- Novelty
    - This appears to be the first paper tackling the problem of offline pretraining a world model on multiple tasks and online finetuning on a different task.
- Reproducibility
    - Experiments are run with 5 seeds. Architecture details and hyperparameters are provided in Appendix. The authors "are committed to" releasing the code. So it seems likely to reproduce the results.
summary_of_the_review: I am leaning toward acceptance. The paper makes a first attempt at multi-task offline pretraining and online finetuning using a world model. While additional comparisons to single-task / model-free settings are nice to have, the paper already provides sufficient new knowledge to the community.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper introduces XTRA, a model-based algorithm for reinforcement learning transfer. XTRA extends EfficientZero (SOTA model-based RL on Atari games) to the RL transfer setting in 2 steps:

1. Offline multi-task pretraining: distilling N EfficientZero teachers (pretrained on N tasks via offline RL) into a single EfficientZero student.
2. Online finetuning: where the student keeps optimizing for the mixture of pretraining tasks while also learning the downstream task via online RL.

Besides this pretraining-finetuning setup for model-based RL, the authors propose a scheme to reweight pretraining tasks according to their affinity with the downstream task — this aims to mimic the effects of gradient-surgery (Yu et al., 2020).

They evaluate XTRA on the Atari 100k benchmark (learning with only 100k frames), both when transferring between tasks that are similar (e.g., pretraining and downstream are instances of shooter tasks) and when they are dissimilar (i.e., not from any particular game family).
strength_and_weaknesses: - Strengths:
    - The authors tackle a notoriously hard problem: transfer across different tasks in RL, from images. They also chose one of the most competitive benchmarks to evaluate their method, namely Atari 100k, and include the most competitive baselines. In that sense, the results are impressive and XTRA is indeed successful at leveraging pretraining tasks to more efficiently solve the downstream one (my concerns on this point below).
    - To their credit, the authors also include some ablations studies and analyses (see 2nd part of 4.1) where they discuss task relevance, model size, and, more interestingly, which component of the model-based architecture is more transferrable and when. (Transferring the dynamics yields in the largest boost in the early finetuning phases.) I only wished those questions were explored in more depth, to help guide the community understand the mechanics of transfer in model-based RL on Atari.
- Weaknesses:
    - Since most of the experiments focus on benchmarking, the merits of the paper rest on XTRA’s performance. While it shows promising results, I also wonder if the results are as positive as described in the paper. First, how much more data than EfficientZero does XTRA require when taking pretraining into account (I couldn’t find those numbers in the main text)? Second, the confidence intervals are quite large when transferring to similar tasks (see Fig. 6: Carnival, DemonAttack, Alien, Amidar, Bankheist, MsPacman, Wizard of Wor) — are the reported results significant? It’s also surprising that XTRA does better relative to EfficientZero when downstream tasks are dissimilar from pretraining (see Table 2, confidence intervals) — could the authors comment on this point, and describe how the confidence intervals are computed in this setting?
        
        Depending on the answer to those questions, I wonder if this paper isn’t reporting negative results for cross-task RL transfer.
        
    - A more minor point: I wish the paper included better baselines — not just ablations of the proposed method. For example, how does XTRA perform against on-policy alternatives such as Actor-Mimic or Policy Distillation (e.g., on top of DrQ-v2)? In a similar vein, the authors claim their task-weighting scheme yields the same benefits as gradient surgery (see 3.2) but when is that empirically demonstrated?
    - One question re “Most improvements happen within the first  stages of training.”: Do the authors have any insight on why that’s the case? From Fig. 8 it seems that the dynamics are the reason for this boost (but we’re missing “h+f” on this figure). Is it because the dynamics already have a good initialization (i.e., don’t move too much during finetuning) or is it because the initialization lets them quickly move through parameter space (i.e., large movement at first, but not much after)?
    - Details:
        - What is the definition (formula) of $s^i$? I’m assuming it’s Eq. 2 but please clarify.
        - 3.1, p. 3: “tasks is however a difficult” → remove a.
        - 3.1, p. 4: you could cite Parisi et al.’s “Actor-Mimic”, and Rusu et al.’s “Policy Distillation”.
        - 3.1, p. 4: $(\hat{\pi}, \hat{u}, \hat{z})$ should be $(\hat{\pi}, \hat{v}, \hat{z})$?
        - 3.2, eq. 1: The sum should be $\sum_i^{m}$, right?
        - 4.1, p. 7: “our proposed frames” → “our proposed framework”?
        - Table 2, p. 7: \cite{} → \citep{} for Ye et al., 2021.
        - 4.1, p. 8: “This results indicates that (i) selecting […] and (2) in the […]” → enumeration consistency.
clarity,_quality,_novelty_and_reproducibility: - Clarity: mostly clear; my main qualm is about some details postponed to the Appendix (full explanation for the task reweighting scheme), whereas they could have been included with more careful writing.
- Quality: good; I’m not sure what Fig. 4 brings to the exposition but it doesn’t hurt.
- Novelty: good; the method is not so novel (tweaks on top of EfficientZero) but the transfer learning setting in RL is.
- Reproducibility: ok?; the authors promise to open-source their code but I don’t see a supplementary material.
summary_of_the_review: - Strengths:
    - The authors tackle a hard RL transfer problem, and show promising results.
    - Some ablation studies and analyses — I wish there were more.
- Weaknesses:
    - The reported results are not 100% compelling. I’ve asked the authors to clarify.
    - The paper could include better baselines, some of which need to be discussed as prior work (Actor-Mimic / Policy Distillation).
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper studies the problem of cross-task transferability in model-based reinforcement learning by introducing a framework with offline pretraining of world models and online finetuning (XTRA).
XTRA builds on top of EfficientZero and introduces a novel weighted loss formulation to tackle two common problems: (1) catastrophic forgetting and (2) non-aligned downstream task with the training tasks. The experimental section compares the resulting approach with the original EfficientZero in several ablation studies, which highlight under which conditions the proposed solution shows positive results. 
The authors conduct experiments on 14 visual tasks from the ALE benchmark and explicitly consider a reduced-budget setting.
strength_and_weaknesses: The paper is well-written, and the problem under study is clearly defined. The proposed contribution in the loss formulation looks technically sound and the authors fully convey its rationale. Moreover, the experimental section shows an exemplary structure with several ablations that analyse different aspects of the proposed solution. 

The main criticism is on the empirical evidence of the effect of gradient weighting. In the first experiment, the authors evaluate XTRA in tasks with similar game mechanics. While the performances are generally promising, one of the baselines uses constant weights and shows much higher performance in two games. While I acknowledge we cannot expect to have a solution for everything, I think the large performance gap disconfirms one of the main contributions of this work without properly commenting on it. 

Moreover, even if focusing on model-based reinforcement learning, I would like to see the performance compared with model-free approaches. Even if the training budget is not comparable, reporting their asymptotic performance would give more context to the readers.

Minor errors:
- Missing reference to the original paper on World Models by Ha et al. (NeurIPS 2018).
- In Section 3.1, it is unclear what 'u' refers to when saying about "task specific quantities (pi, u, z)". From the appendix, 'u' is used for immediate reward, but I could not find its definition in the main text.
- Typo in 3.1: "Learning a single RL agent for a diverse set of tasks is however a difficult in practice" instead of "Learning a single RL agent for a diverse set of tasks is however difficult in practice".
- End of paragraph "Scaling model size": the sentence "XTRA is robust to model size and can be further benefit from scaling up the model size in the future" sounds a bit overselling and, in my opinion, it is not supported by sufficient experiments to claim something like that.
- Typo in Related work: "or example, He et al (2020)" instead of "For example, He et al (2020)"
clarity,_quality,_novelty_and_reproducibility: - Clarity and Quality: The paper is easy to read, clear in its statement and technically sound. The experimental results are clearly reported with good use of marking and coloring schemes. The experiments are repeated for multiple runs, and the aggregated metrics described in most of the captions.  
- Novelty: the novelty of this work is difficult to evaluate because in practice, it combines existing techniques (finetuning on both online and offline datasets, gradient weighting based on gradient similarity). However, I acknowledge there is merit in studying the effectiveness of this approach on a diverse set of tasks.
- Reproducibility: the authors report the hyperparameters in the appendix but do not attach any code in the supplementary material, which makes its reproducibility only partially addressed. 
summary_of_the_review: Considering the lack of empirical support for the gradient-weighting mechanism, and the partial novelty in the overall framework, I vote for a weak rejection. By addressing these criticisms, the paper would be solid and mature for publication. I would recommend the authors to advance possible theses of the observed performance and validate them with additional experiments. For example, I would like to understand (1) to what extent the gradient weighting captures the task similarities and if the problem should be addressed by other mechanisms; and (2) if there is any model bias dependent on visual features that is playing a role in this performance gap.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: This paper studies the long standing problem of data efficient RL. Pretraining is still not widely used in training RL agents and this paper takes small steps towards that direction. It proposes a framework for data efficient online RL by first learning a world model from N diverse tasks and then fine tuning this model on a target task with gradient re-weighting and N-task experience replays. To avoid catastrophic forgetting and task interference, two key mechanisms are proposed on top of the EfficientZero based architecture: 1) jointly train on all task data during target fine tuning and 2) adaptively re-weight gradient contributions in unsupervised manner based on similarity to target task.

I believe this paper tackles a critical problem in RL -- how to leverage and transfer a world model for online RL. More work needs to be done in this area and this paper proposes a sound and interesting take on this problem. The paper is well written and relatively easy to follow. The ablation and analysis are informative and makes it clear which components are beneficial to fine tune and what contributions bring performance improvement. They also provide a strong baseline and 25% improvement on a subset of the 100k Atari benchmark. 

I also believe that this paper is good enough (assuming it incorporates all the suggestions made during the review process) for other authors to cite for comparisons (and potentially build upon). There were a lot of concerns that came up during the discussion period about the novelty of this approach but nothing major was articulated -- there are novel elements in this approach with sound experimental validation. One of the concerns was that the target task is not truly out of distribution and is still in the atari world -- but I think this is not a fair criticism. Studying generalization of games that are closely related but not fully OOD is an interesting problem in itself that deserves attention. Another concern was that the empirical improvements was 25% and not high enough -- I think the baselines are strong and the approach is interesting; so this type of an improvement along with the formulation makes this approach sound and interesting for other works to build upon.

justification_for_why_not_higher_score: The paper is technically sound and the approach is quite interesting. I would have recommended it for spotlight or oral if they had studied more environments other than atari. Atari 100k may require a nuanced understanding and the broader ICLR audience might resonate more if there were multiple domains across different modalities or realism constraints.
justification_for_why_not_lower_score: This is an important topic with an interesting solution. The use of learning a single world model and using it to transfer learn on a new target task is an open and timely problem to study. This paper clearly proposes a sound approach and should be accepted into ICLR. I believe several people will reference and build on top of these results for the Atari 100k suite and beyond.
Decision: Accept: poster