Unsupervised Model-based Pre-training for Data-efficient Control from Pixels
Authors: Sai Rajeswar
Pietro Mazzaglia
Tim Verbelen
Alexandre Piché
Bart Dhoedt
Aaron Courville
Alexandre Lacoste
Venue: Submitted to ICLR 2023 (2023)
Keywords: unsupervised reinforcement learning
world models
planning

Abstract: Controlling artificial agents from visual sensory data is an arduous task. Reinforcement learning (RL) algorithms can succeed in this but require large amounts of interactions between the agent and the environment. To alleviate the issue, unsupervised RL proposes to employ self-supervised interaction and learning, for adapting faster to future tasks. Yet, whether current unsupervised strategies improve generalization capabilities is still unclear, especially in visual control settings. In this work, we design an unsupervised RL strategy for data-efficient visual control. First, we show that world models pre-trained with data collected using unsupervised RL can facilitate adaptation for future tasks. Then, we analyze several design choices to adapt faster, effectively reusing the agents' pre-trained components, and planning in imagination, with our hybrid planner, which we dub Dyna-MPC. By combining the findings of a large-scale empirical study, we establish an approach that strongly improves performance on the Unsupervised RL Benchmark, requiring 20$\times$ less data to match the performance of supervised methods. The approach also demonstrates robust performance on the Real-Word RL benchmark, hinting that the approach generalizes to noisy environments.

Under review as a conference paper at ICLR 2023
UNSUPERVISED MODEL-BASED PRE-TRAINING FOR DATA-EFFICIEN CONTROL FROM PIXELS
Anonymous authors Paper under double-blind review
ABSTRACT
Controlling artificial agents from visual sensory data is arduous task. Reinforce- ment learning (RL) algorithms can succeed this but require large amounts interactions between the agent and the environment. To alleviate the issue, unsuper- vised RL proposes to employ self-supervised interaction and learning, for adapting faster to future tasks. Yet, whether current unsupervised strategies improve gener- alization capabilities is still unclear, especially in visual control settings. In this work, we design an unsupervised RI. strategy for data-efficient visual control. First, we show that world models pre-trained with data collected using unsupervised RL can facilitate adaptation for future tasks. Then, we analyze several design choices adapt faster, effectively reusing the agents' pre-trained components, and planning imagination, with our hybrid planner, which we dub Dyna-MPC. By combining the findings of large-scale empirical study, establish approach that strongly improves performance on the Unsupervised RL Benchmark, requiring 20 less data match the performance of supervised methods. The approach also demonstrates robust performance on the Real-Word RL benchmark, hinting that the approach generalizes to noisy environments.
INTRODUCTION
Modern successes of deep reinforcement learning (RL) have shown promising results for control
problems (Levine et al., 2016 OpenAl al. 2019 Lu al. 202T). However, training an agent for each task individually requires a large amount of task-specific environment interactions, incurring huge redundancy and prolonged human supervision. Developing algorithms that can efficiently adapt and generalize to new tasks has hence become active area of research in the RL community.
In computer vision and natural language processing, unsupervised learning has enabled training models without supervision to reduce sample complexity on downstream tasks (Chen et al.) 2020 Radford et al. 2019). In a similar fashion, unsupervised RL (URL) agents aim to learn about the environment without the need for external reward functions, driven by intrinsic motivation (Pathak) ct al. 2017 Burda et al.) 2019a| Bellemare et al. 2016). Any learned models can then be adapted to downstream tasks, aiming to reduce the required amount of interactions with the environment.
Recently, the Unsupervised RI. Benchmark (URLB) (Laskin et al. 2021) established a common protocol compare self-supervised algorithms across several domains and tasks from the DMC Suite (Tasaa et al. 2018). In the benchmark, an agent is allowed task-agnostic pre-training stage, where it interact with the environment in an unsupervised manner, followed by fine-tuning stage where, given limited budget of interactions with the environment, the agent should quickly adapt for a specific task. However, the results obtained by Laskin et al. (202I) suggest that current URL approaches may be insufficient to perform well on the benchmark, especially when the inputs of the agent are pixel-based images.
World models have proven highly effective for solving RL tasks from vision both in simulation (Hafner al. 2021 2019a) and in robotics (Wu et al., 2022), and they are generally data-efficient as they enable learning behavior in imagination (Sutton) 1991). Inspired by previous work on exploration (Sekar et al. 2020 we hypothesize this feature could be key in the unsupervised RL setting, pre-trained world model can leverage previous experience to leam behavior for new tasks in imagination, and in our work, we study how to best exploit this feature. We adopt the URLB setup perform a large-scale study, involving several unsupervised RL methods for pre-training

---

Under review as a conference paper at ICLR 2023
State
Weight
Prsupervioad
Reward
Planser
action
World
Teank agnestic
Model
Tack
Pre-training
Fine-tuning
Figure I: Method overview. Our method considers a pre-training (PT) and fine-tuning (FT) stage. During pre-training, the agent interacts with the environment through unsupervised RL, maximizing an intrinsic reward function, and concurrently training a world model on the data collected. During fine-tuning, the agent exploits its pre-trained components and plans in imagination, to efficiently adapt to different downstream tasks, maximizing the rewards received from the environment.
model-based agents, different fine-tuning strategies, and a new improved algorithm for efficiently planning with world models. The resulting approach, which combines the findings of our study, strongly improves performance on the URL benchmark from pixels, nearly achieving the asymptotic performance of supervised RL agents, trained with 20x more task-specific data, and bridging the gap with low-dimensional state inputs (Laskin et al. 2021
Contributions. This work does not propose novel complex method. Rather, we study the interplay of various existing components and propose a novel final solution that outperforms existing state of the art on URLB by staggering margin. Specifically:
we demonstrate that unsupervised RL combined with world models can be an effective pre-training strategy to enable data-efficient visual control (Section 3.T).
we study the interplays between the agent's 's pre-trained components that improve sample efficiency during fine-tuning (Section 3.2).
we propose novel hybrid planner we call Dyna-MPC, which allows us to effectively combine behaviors learned imagination with planning (Section 3.3).
combining our findings into approach, outperform previous approaches URLB from pixels, nearly solving the benchmark (Section (.I1).
show the approach is resilient to environment perturbations, evaluating it on the Real et al. 2020) (Section (42).
World RL benchmark (Dulac-Amoo
present an extensive analysis of the pre-trained agents, aimed at understanding in-depth
the current findings and limitations (Section 4.3
An extensive empirical evaluation, supported by more than 2k experiments, among main results, analysis and ablations, was used to carefully design our method. We hope that our large-scale evaluation will inform future research towards developing and deploying pre-trained agents that can be adapted with considerably less data to more complex/realistic tasks, as it has happened with unsupervised pre-trained models for vision (Parisi et 2022) and language (Ahn al. 2022
PRELIMINARIES
Reinforcement learning. The RL setting can be formalized as a Markov Decision Process (MDP), denoted with the tuple A, T, R, 7) where S is the of states, A is the set of actions, 7' is the state transition dynamics, R is the reward function, and is discount factor. The objective of an RL agent to maximize the expected discounted sum of rewards over time for given task, also called return, and indicated as G. = Σt=t+1 7 (k-(-1) In continuous-action settings, you can learn an actor, i.e. model predicting the action to take from a certain state, and critic, i.e. model that estimates the expected value the actor's 's actions over time. Actor-critic algorithms can be combined with the expressiveness of neural network models to solve complex continuous control tasks HHarnoja et al. 2018 Lillicrap et al. 2016 Schulman et al." 2017
'The PyTorch code for the experiments will be open-sourced upon publication.
2

---

Under review as a conference paper at ICLR 2023
Unsupervised RL. In this work, we investigate the problem of fast adaptation for downstream task, after a phase of unsupervised training and interaction with the environment. Our training routine, based on the setup of URLB (Laskin al. 202T), is made of two phases: pre-training (PT) phase, where the agent can interact with a task-agnostic version of the environment for up to 2M frames, and fine-tuning phase (FT), where the agent given a task to solve and limited budget of 100k frames. During the PT phase, rewards are removed so that sensible information about the environment should be obtained by exploring the domain-dependent dynamics, which is expected to remain similar or unchanged in the downstream tasks. During FT, the agent receives task-specine rewards when interacting with the environment. As the agent has no prior knowledge of the task, should both understand the task and solve it efficiently, in a limited interaction budget. In this setting, the performance of unsupervised model-free RL Yarats et al. 2022 were shown to insufficient as reported in (Laskin et al. 202T) p. We believe the key reason for this is that model-free RL algorithms can exploit only little part of the information obtained with self--uperrised interaction, as they rely uniquely on actor and critic's predictions.
World models. In this work, we ground upon the Dreamer V2 agent (Hafner et al. 2021 which learns world model (Ha & Schmidhuber 20188 Hafner et al., 2019b) predicting the outcomes of actions in the environment. The dynamics is captured into a latent space 2, providing a compact representation of the high-dimensional inputs. The world model consists of the following components:
Encoder: e 1 (o(e)) p=(21|2111, Q-1).
(0(6)( Posterior: 21 0111
Decoder:
Dynamics:
The model states ž, have both a deterministic component, modeled using the recurrent state of GRU (Chung et al. 2014). and a (discrete) stochastic component. The encoder and decoder are convolutional neural networks (CNNs) and the remaining components are multi-layer perceptrons (MLPs). The world model is trained end-to-end by optimizing an evidence lower bound (ELBO) on the log-likelihood of the data collected in the environment (Hafner al.) 2019h a). For the encoder and the decoder networks, we used the same architecture as in Hafner et al. (2021).
For control, the agent learns latent actor Ta (a, 22) and critic (2)) networks. Both components trained online within the world model, by imagining the model state outcomes of the actions produced provided by a reward predictor, (rela)) 24) trained to predict environment rewards, and they are combined with the critic predictions produce a GAE-À estimate of the returns (Schulman et al. 2016). The actor maximizes estimates of returns, backpropagating gradients through the model dynamics. The hyperparameters for the agent, which we keep fixed across all domains/tasks, can be found in Appendix H
by the actor, using the model dynamics. Rewards for imagined trajectories
UNSUPERVISED MODEL-BASED PRE- TRE-TRAINING FOR DATA-EFFICIENT CONTROL FROM PIXELS
To best exploit self-supervised pre-training for data-efficient adaptation, it is important that the agent: (i) meaningfully interacts with the environment during the PT phase, to discover useful transitions; (ii) successfully reuses the modules learned during PT for fast adaptation; and (iii) efficiently employs the FT phase to quickly understand and master the downstream task. In this section,
experiment-drive driven approach to find which methods components are best at tackling these challenges.
Experimental procedure. We employ the URL benchmark that consists of three control domains, Walker, Quadruped and Jaco, and twelve tasks, four per domain. To evaluate the agents, we take snapshots of the agent at different times during training, i.c. 100k, 500k, 1M, and 2M frames, and fine- tune the agent for 100k frames. In all bar plots, we show average normalized returns downstream tasks with error bars showing the standard deviation. To normalize results in a comparable way for all tasks, we train a sully-supervised agent with 2M frames per task. We use the mean performance of this agent, which we refer to as "oracle", as the reference scores to normalize our results in the plots (details in Appendix A). For all experiments, results are presented with at least three random seeds.
3

---

Under review as a conference paper at ICLR 2023
Drog
Noratred Beturr (h4)
Roturn
Nomalzzd
300e
500c
100k
tock Quadraped
Walker
(a) Aggregated results from Laskin al.
2021 of unsupervised RL pre-train the model-free DrQ agent.
Dreamerg:00m L0Ck
Random actions + Dreamer
Knowledee based + Dreamer
-Np Data-based + Dreamer
Cumprtance Dreamer
Normatzet Suttern M54
Normaltoo
2004
500c
14
24
10ok
3M
3204
500k
1/4
100k
Walker
Quadruped
Jaco
Overall
(b) Our aggregated results of unsupervised RL to pre-train the model-based Dreamer agent.
Figure 2: Unsupervised pre-training. Aggregated performance of different URL techniques for PT, with FT snapshots taken at different times along training.
(a) With the model-free DrQ agent, techniques.
performance slightly improves over time only using knowledge-basse model-based Dreamer agent, performance is higher and, overall, improves for all techniques. We also report Dreamer@ 100k and DrQ@ 100k results, which are obtained 100k FT steps with no PT.
(b) With the
3.1 UNSUPERVISED PRE-TRAINING
In the PT stage, unsupervised RL can be used explore the environment, collecting the data train the components of the agent. The resulting networks are then used to initialize respective components in the agent deployed for the downstream task, aiming reduce sample complexity during FT. The first question we address is thus "What kinds of agents work best with unsupervised tre-training?"
Unsupervised RL methods can be grouped into three categories (Laskin et al. 2021| knowledge- based, which aim to increase the agent's 's knowledge by maximizing error prediction (Pathak et al. 2019 Burda et al., 2019b). data-based, which aim achieve diversity of data Yarrat et al. and competence-based which aim to learn diverse skills (Liu & Abbeel 2021a Eysenbach al. 2019). In Figure 2a we report the results from Laskin et al. (2021), showing that none of these approaches is particularly effective on URLB when combined with the DrQ model-free agent (Yarats et al. 2022). state-of-the-art in RL from pixels, where the data collected with unsupervised RL is used pre-train the agent's actor, critic, and encoder.
2017
2021 Liu & |
To demonstrate that world models can be used to effectively exploit unsupervised RL data collection for fast adaptation, study multiple approaches and use them to pre-train the Dreamer's 's world model and latent actor. As knowledge-ben
methods we employ ICM (Pathak et al. 2017). LBS (Mazzaglia let al. 2021b), Plan2Explore (P2E; (Sekar et al., (0200). ) and RND (Burda al. 2019b). As data-based approach, we choose APT (Liu & Abbeel (2221b), and as competence-based approaches, adopt DIAYN (Eysenbach al. 2019) and APS (Liu & Abbeel 202 Finally, we also test
random actions, as a naive maximum entropy baseline (Haarnoja et al. 2018). Details on these
methods and how we combined them with the Dreamer algorithm are discussed in Appendix B
we show that by leveraging pre-trained world model the overall performance improves over time for all categories, as opposed to the model-free results, where only shooledde--aseeee approaches slightly improve. In particular, data-based and knowledge-birt based methods are more effective in the Walker and Quadruped domains, and random actions and competence-basee based are more effective in the Jaco domain. Detailed results for each method are available in Appendix E
Aggregating results per category, in Figure 2b
3.2 FINETUNING PRE-TRAINED AGENTS
Some of the components learned during the PT phase, such as the world model, can be reused for fast adaptation during FT. However, the reward is changing from pseudo-reward to task reward when
4

---

Under review as a conference paper at ICLR 2023
Pre-training
Pre-training
Normalized Return (8)
100
100
60
80-
Normaiized
20
20
Quadruped
jaco
Quadruped
jace
Fine-tuning pre-trained agents. Comparison of the results when fine-tuning dif- ferent pre-trained components of the agent. Re- sults averaged across all URL methods.
Figure 3:
Figure 4: Learning and planning, Using Dyna- MPC during fine-tuning improves performance, increasing mean scores and reducing variance. Results averaged across all URL methods.
changing from the PT to the FT phase, it is not clear if pre-training of the actor and critic can help the downstream task. shed light on this, we seek to answer: "Which pre-trained components are useful for downstream tasks?".
Here, we test different fine-tuning configurations, where we copy the weights of some the PT components into the agent to fine-tune for the downstream task. We run the tests for the several unsupervised RL methods combined with Dreamer that we presented in Section 3.1 and show aggregated results in Figure (detailed results per each method in Appendix E.
Overall, fine-tuning the PT world model provides the most significant boost in performance, strength- ening the hypothesis that world models very effective with unsupervised RL. Fine-tuning the actor improves performance slightly Walker and remarkably in Quadruped, but is harmful in the Jaco domain. An intuitive explanation is that in the Quadruped and Walker moving tasks, the exploratory behaviors help discovering reward faster. Instead, in the Jaco goal-reaching tasks, the agent needs to reach certain target with sparse rewards. If the PT actor is initialized move far from the target, the agent might struggle to find rewards in the small FT budget. Finally, using PT critic is systematically worse. This can be explained by the discrepancy between intrinsic rewards and task rewards.
3.3 LEARNING AND PLANNING IN IMAGINATION
Knowing model of the environment, traditional model-based control approaches, e.g. model predictive control (MPC) (Williams et al. 2015 Chua al. 2018 Richards 2005), can be used to plan the agent's action. Nonetheless, using actor-critic methods has several ad advantages, such as amortizing the cost of planning by eaching previously computed (sub)optimal actions and computing long-term returns from certain state, without having to predict outcomes that are far in the future. More recent hybrid strategies, such as LOOP (iikchi et al. 2020) and TD-MPC (Hansen et al. 2002), allow combining trajectories sampled from the actor with trajectories sampled from distribution over actions that is iteratively improved. The model and the critic are used to evaluate the trajectories,
Algorithm Dyna-MPC
Require: Actor 0, Critic vb, World Model Ø
1:
μ., initial parameters for sampling actions
2: N, N₂ num trajectories, num policy trajectories
24, H: current model state, planning horizon 4: for each iteration j 1..J do
3:
5: Sample N trajectories of length H from M(u, 22), starting from 2₁
6: Sample N₂ trajectories of length H using the actor πe, starting from 2t
7: Estimate future states, using the model, and returns, using reward and critic predictions
8: Update p and a using MPPI (Williams et al. 2015)
9: end for
10: return at N of))
5

---

Under review as a conference paper at ICLR 2023
improve them, and eventually select the most promising actions, i.e. planning. In this section, answer the question: Can we accelerate downstream task adaptation by leveraging planning?
Dyna-MPC. As we pre-train a world model, could exploit planning in latent space to adapt with limited additional environment interaction. One problem with the above strategies is that they are based upon learning off-policy actor and critic, which in our context would prevent from exploiting the PT model to learn the actor and critic in imagination. In order to enable hybrid planning with the behavior learned in imagination (Hafner et al. 2019a),, we develop a modification of these approaches, which we call Dyna-MPC, that combines the actor and critic learned in imagination with MPPI (Williams et al. 2013) for planning.
As detailed in Algorithm| ! at each time step, we imagine set of latent trajectories using the model, by sampling actions from a time-dependent multivariate gaussian and from the actor policy, trained with Dreamer in imagination. Returns for MPPI are estimated using reward predictions by the model and the critic. MPPI is used to update the parameters of the multivariate gaussian for J iterations. Details on how returns are estimated and the MPPI updates work are given in Appendix C One significant difference with previous approaches is that the policy in Dyna-MPC is learned on-policy in imagination, thus no correction for learning off-policy is required (Sikchi et al. 2020).
Given the insights from the previous section, we use the world models and actors pre-trained with all the different unsupervised strategies considered (see Section 3.1 and test their FT performance with and without planning with DyP. Aggregated scores are reported in Figure 4 and detailed results for each method are available in Appendix E We observe that adopting Dyna-MPC is always beneficial, as it improves the average performance and reduces variance in all domains.
3.4 OUR METHOD: COMBINING THE FINDINGS TOGETIIER
In the large-scale study, we explored several design choices to establish the most adequate approach to tackle the URL benchmark, aiming to provide general recipe for data-efficient adaptation thanks to unsupervised RL. Our approach combines the main findings we presented in the previous sections:
1. learning a model-based agent with data collected using unsupervised RL (Figure 2
2. fine-tuning the PT world model (always) and the pre-trained actor (where beneficial), while learning the critic from scratch (Figure
3. adopting a hybrid planner, as the proposed DyCa-MPC, to leverage both learning and planning in imagination (Figure
An overview of the method is illustrated in Figure and the algorithm is presented in Appendix D We believe the above recipe could be generally applied unsupervised settings, also outside of URLB, with the precaution that one should carefully make two decisions: (a) whether fine-tuning the PT actor meaningful for the downstream task it's better to re-learn it from scratch, (b) what is the best URL strategy to collect data. Both decisions strongly depend on the target domain/task and so it is difficult to assess their implications beforehand. However, adopting unsupervised strategies that specifically focus on interacting with interesting elements of the environment, e.g. objects, or that quickly explore large areas of the environment at the beginning of fine-tuning may help exploring and revisiting crucial states of the environment more easily Parisi al. 2021
For URLB, we already established (a) that the PT actor is effective in Walker and Quadruped tasks, but it is better re-learn the actor from scratch in Jaco, in Section 3.2 To decide which URL strategy to use (b) we present a detailed comparison of the performance of our approach using different exploration strategies. The results in Figure 5 show that the agent using LBS during pre- pre-training performs overall best, as it has the highest interquartile mean (IQM) and mean scores, and the lowest optimality gap. Thus, in the evaluation section, we present Ours (LBSS our approach.

---

Under review as a conference paper at ICLR 2023
Median
IQM
Mean
Optimality Gap
APS DIAYN APT
LBS
RND
P2E
ICM
0.6
0.7 0.8
0.9
0.60
0.75
0.90
0.6
0.7
0.8
0.9
0.1
0.2
0.3
0.4
Normalized Return
Figure 5: Unsupervised RL comparison. Stratified bootstrapped performance of our approach on URLB under different data collection strategies, with 95% Cls.
(Agarwal et al. 2021
DrOg ICM DrQ
Disagreement RND DrQ
DrQ
Dreamen@1000 100k
ProtoRL DrQ
DIAYN DrQ
100
90
80
60
50
40
30
Nome
20
10
Walker
Quadruped
Jaco
Overall
Figure 6: URLB results. We compare approach the results on URLB after 2M frames of
pre-training from [Laskin et al. 2021 (DrQ pre-trained with different unsupervised RL strategies).
4 EVALUATION AND ANALYSIS
4.1
UNSUPERVISED REINFORCEMENT LEARNING BENCHMARK
In Section 3 we presented our approach, which combines the findings from our empirical large-scale study on URLB. In F we compare the results from the original URLB paper with our approach. The performance of our method is superior in all domains. The second strongest method (DrQ with Disagreement) approaches an overall performance of 40% of the respective supervised baseline performance, while our method recovers more than 90% of its supervised counterpart.
4.2 REAL-WORLD REINFORCEMENT LEARNING BENCHMARK
Algorithms developed in simulation struggle to transfer to real-world systems due to series of implicit assumptions that are rarely satisfied in real environments, c.g. URLB assumes the dynamics between PT and FT stay the same. The RWRL benchmark (Dulac- et al. 2020) considers several challenges that are common in real- world systems and implements them on top of DMC tasks.
We employ vision-basee variants of the Walker Walk and Quadruped Walk tasks from the RWRL benchmark. These tasks introduce system delays, stochasticity, and perturbations the robot's model and sensors, which are applied with three degrees of intensity to the original environment, i.e. 'easy', "medium", and 'hard' (details in Appendix We seek answer whether in perturbed settings:
does unsupervised PT enable faster adaptation?
does unsupervised RL provide an advantage over random exploration?
does hybrid planning improve performance, in URLB?
In Figure we present the results of our method, using LBS during PT, with and without planning with Dyna-MPC for FT, and compare to random exploration and training from scratch for 100k. 1M, and 2M frames. Crucially, the PT models are trained in the vanilla task-agnostic version of the environments from the DMC Suite, that the results highlight the extent which models trained in ideal conditions generalize to perturbed settings when fine-tuned low-data regime.
We exceptionally do not the pre-trained actor in the Jaco tasks, as this was shown to lead to hetter nerformance in Section 3.2 (Figure
7

---

Under review as a conference paper at ICLR 2023
Overall, found that fine-tuning PT models offer advantage training from scratch for 100k frames, despite all the variations in the environment. Furthermore, on the Quadruped Easy and method performs better than Dreamer @ 1M and not far from Dreamer @ 2M while using 10x and 20x less task-specific data, respectively. Our method also performs close to Dreamer @ M/2M in the Walker Easy task. Unsupervised RI. for data collection (Ours) outperforms random actions in the "easy" and 'medium" settings, showing that a better PT model yields higher FT performance, even when the dynamies of the downstream task affected by misspecifications and noisy factors. Finally, in contrast with the findings on URLB, adopting the hybrid planner is not generally beneficial. We believe this is because the model's predictions are less certain and precise in this setting and thus cannot inform the short-term planner accurately.
Medium settings,
4.3 EXTENDED ANALYSIS
To better analyze the learned components, conducted range of additional experiments. For conciseness, detailed descriptions of the experimental settings are deferred to Appendix Gla we briefly summarize the takeaways in this section.
Learning rewards online. We verify whether having to discover and learn the reward function during FT' impacts performance. In Figure we compare against agents that (violating the URLB settings) know the task in advance and can pre-train reward predictor during the PT stage. We see that learning the reward predictor does not affect performance significantly for dense-reward tasks, such the Walker and Quadruped tasks. However, in sparser reward tasks, i.e. the Jaco ones, knowing reward information in advance provides an advantage. Efficient strategies to find sparse rewards efficiently represent challenge for future research. More details in Appendix G.1
Zero-shot adaptation. Knowing reward predictor from PT, it could be possible to perform zero-shot control with MPC methods if the model and the reward function allow it. In Figure 9 we show that despite the zero-shot MPC (ZS) offers an advantage over Dreamer@ 100k, the FT phase is crucial to deliver high performance on the downstream tasks, as the agent uses this phase to collect missing information about the environment and the task. Further details in Appendix G.2
Latent dynamics discrepancy (LDD). We propose a novel metric, Latent Dynamics Discrepancy, which evaluates the distance between the latent predictions of the PT model and the same model after FT on a task. In Figure 10 we show the correlation between our metric and the performance ratio between using the PT model and the FT model for planning (see Appendix G.3 for detailed expla- nation). We observed strong negative Pearson correlation (-0.62, p-value: 0.03), highlighting that major updates in the model dynamics during FT played an important role in improving performance.
Unsupervised rewards and performance. We analyze the correlation between the normalized performance of different agents and their intrinsic rewards for optimal trajectories obtained by an oracle agent in T In particular, the correlation for LBS, which overall performs best in URLB, has a statistical significance, as its p-value is < 0.05. We believe this correlation might be one the causes of LBS outstanding performance. Further insights are provided in Appendix G.4
Walker
Real-wordd Quadruped Walk
1000
100
Normalized
20
20
10
10
Vanila
Medium
Hand
Vanila
Medium
Hlnd
Dreamerg110k
Dreameng1M
Dreamerg2M1
Random Actions Dreamerg110kk
Ours (no Dyro MPC10100k
Oursg110kk
Figure 7: RWRL results. We compare method to random exploration and training from scratch on the tasks from the RWRL benchmark. Our models are pre-trained on the vanilla version of the environment for 2M frames and fine-tuned for 100k frames on the perturbated tasks from RWRL.
8

---

Under review as a conference paper at ICLR 2023
Pre-training
120
Normalzzed Return (8)
100
1.0
Dreameng150k
NFC MPC Im
Dreamer
80
60
Normaizzd
80-
40
20
20
Quadruped
jaco
Quadruped
jace
Figure 8: Reward learning ablation. Results
Figure 9: Zero-shot (ZS) vs Fine-tuned (FT). Results of the agent using Plan2Explore.
averaged across all unsupervised RL methods.
RELATED WORK
Model-based control. Dynamics models combined with powerful search methods have led to impressive results on wide variety of tasks such as Atari (Schrittwieser et al. 2020) and continuous control (Hafner et al. 2019a; Janner et al. 2019
Sikchi et al. [021 Lowrey al. 2018). LOOP
(Sikchi et al. 2020) and TD-MPC (Hansen al. 2022) combine temporal difference learning and MPC. The model proposed with TD-MPC is task-oriented and thus requires task to accelerate learning. In our work, we focus on unsupervised model learning, grounding on the Dreamer model (Hafner al., [021], whose supervision comes from predicting the environment's observations. Methods that use no reconstruction could generalize better to visual differences (Mazzaglia et al.)
2021a
Ma al. 2020 but they lose in explainability, as they cannot decode imagined trajectories.
Unsupervised RL. Prior to our work, the large-scale study of curiosity (Burda et al. 2018) provided an insightful analysis of the performance of knowledge-basee our work, we leverage the URLB setting, to provide an analysis of a combination of model- control techniques with unsupervised RL. This allowed us to formulate strategy to adapt pre-trained models visual control tasks in data bata-efficient manner. Closely, (Sekar et al. (2020) combines adapts the Disagreement (Pathak al. 2019) to work with Dreamer (Hafner et al. (0119)..
methods in the reward- setting. In
our work, in addition to analyzing wider choice of unsupervised RL strategies, we show how to better exploit the agent PT components for adaptation, and we propose a hybrid planner to improve data-efficiency.
Transfer learning. In the field of transfer learning, fine-tuning is the most used approach. However, fine-tuning all the pre-trained agent components may not be the most effective strategy. In transfer learning for RL, they have studied this problem, mainly with the objective of transferring from one environment to another (Farebrother et al.. 2018 Sasso et al. 2022 van Driessel & Francois-Lavee analyze which agent's components should be transferred from the unsupervised
2021). Instead, PT stage to the supervised FT stage when the environment's 's dynamics is assumed to stay similar or be the same. Another stream of work has studied successor representations, enable better transfer of the agent's 's actor-critic (Hansen et al. 2020
Barreto et al. 2016).
Gusdrupesljmmp. Walker stand
0,7
0.6
Quedmuped
Pre-training for 2M environment frames
Perfemaace
Walker fll
Quadruped jaco top left
ICM
LBS
P2E
RND
Quadruped walk
Correlation
-0.54
-0.60
-0.34
-0.03
Walker walk
0.3
run
jacá settem left
Jaco top right
p-value
0.07
0.04
0.28
0.91
0.2
right
Latont Dynomi Diporeponcy
Figure 10: LDD and performance correlation. The line and shaded area represent linear re- gression model fit and its confidence intervals.
Table 1: Pearson correlation and p-value between fine-tuned performance across URLB tasks and intrinsic rewards.

---

Under review as a conference paper at ICLR 2023
6 CONCLUSION
In order to accelerate the development and deployment of learning agents for real-world tasks, it is crucial that the employed algorithms can adapt in a data-efficient way for multiple tasks. Our study provides an empirical analysis of several design choices, which allowed us obtain near-optimal performance URLB and that showed robustness to perturbations in the environment, on the RWRL benchmark. We also analyzed several aspects of the learned models, understand what could be improved further in the future to case the adaptation process.
Limitations. In the Jaco reaching tasks, we found that bad initialization of the pre-trained actor can actually harm the agent's performance. While competence- based approaches should address this limitation, by learning variety of skill behaviors, their performance the other domains has been subpar. Future work should aim to find a more general approach to pre-train behavior for fast adaptation or improve the exploration capabilities of competence- baaed approaches.
Another issue we encountered, on the RWRL benchmark, that if the environment introduces too intense perturbations during adaptation, relying on the predictions of the adopted world model becomes problematic, to the extent that exploiting planner is not useful anymore. Developing more resilient models that can be trained in an unsupervised fashion and used for data-efficient planning, even in presence of complex perturbations, will be the focus of future studies.
Reproducibility statement We reported in the main text (Algorithm D the pseudo-code for Dyna- MPC and Appendix the pseudo-code for our end-to-end approach. We also provide instructions on how we implemented our methods (Appendix [) and all the model and training hyperparameters implement and reproduce the results (Table We will release our code and scripts.
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare G. Bellemare. Deep reinforcement learning at the edge of the statistical precipice, 2021. URL https://
wwtp:/rgggbbb///eee 13264
Michael Ahn, Anthony Brohan, Noah Brown, Yevgen Chebotar, Omar Cortes, Byron David, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Daniel Ho. Jasmine Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Eric Jang, Rosario Jauregui Ruano, Kyle Jeffrey, Sally Jesmonth, Nikhil Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Kuang-Huei Lee, Sergey Levine, Yao Lu, Linda Luu, Carolina Parada, Peter Pastor, Jornell Quiambao, Kanishka Rao, Jarek Rettinghouse, Diego Reyes, Pierre Sermanet, Nicolas Sievers, Clayton Tan, Alexander Toshev, Vincent Vanhoucke, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, and Mengyuan Yan. Do as can and not as say: Grounding language robotic affordances. arXiv preprint 12X.v:220..01611, 2022.
Arthur Argenson and Gabriel Dulac-Aroold. Model-hased offline planning, 2020. URI. https
7/axxiv. org/
05556
André Barreto, Will Dabney, Rémi Munos, Jonathan J. Hunt, Tom Schaul, Hado van Hasselt, and David Silver. Successor features for transfer in reinforcement learning, 2016. (tLhts
7/arxiv. org/ abs/ /aeps//666.... 05312
Mare Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. In D. Lee, M. Sugiyama, U. Luxburg, Guyon, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 29, 2016.
Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexei A. Efros. Large-scale study of curiosity-driven learning, 2018. URL
04355
Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexei A. Efros. Largescale study of curiosity-driven learning. ICLR, 2019a.
10

---

Under review as a conference paper at ICLR 2023
Yuri Burda, Harrison Edwards, Amos J. Storkey, and Oleg Klimov. Exploration by random network distillation. ICLR, 2019b.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In Proceedings of the 37th International Conference on Machine Learning, pp. 1597-1607, 2020.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Informa- fion Processing Systems. Curran Associates, Inc., 2018
Junyoung Chung, Caglar Gulcehre, Kyunghyun Cho, and Yoshua Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. In NIPS Workshop on Deep Learning, 2014, 2014.
Gabriel Dulac-Amold, Nir Levine, Daniel J. Mankowitz, Jerry Li, Cosmin Paduraru, Sven Gowal, and Todd Hester. An empirical investigation of the challenges of real-world reinforcement learning, 2020.
Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine. Diversity is all you need: Learning skills without a reward function. In ICLR, 2019.
Jesse Farebrother, Marlos C. Machado, and Michael Bowling. Generalization and regularization in dqn, 2018. URL
hwtprrggaas//111/// 00123
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett (eds.), Advances in Neural Information Processing Systems. Curran Associates, Inc., 2018.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maxi- mum entropy deep reinforcement learning with a stochastic actor. arXiv preprint arXiv:/801.012900 2018.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: behaviors by latent imagination. In ICLR, 2019a.
Learning
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In ICML, pp. 2555-2565, 2019b.
Danijar Hafner, Timothy P Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. In ICLR, 2021.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. 2022.
Steven Hansen, Will Dabney, Andre Barreto, David Warde-Farleyy Tom Van de Wiele, and Volodymyr Mnih. Fast task inference with variational intrinsic successor features. In ICLR, 2020.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model-based policy optimization. ArXiv, abs/1906. (at//9006.8253. 2019.
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Lerrel Pinto, and Pieter Abbeel. URLB: Unsupervised reinforcement learning benchmark. In NewrlPS Dalasets and Benchmarks Track (Round 2), 2021.
Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel. End-to-end training of deep visuomotor policies. J. Mach. Learn. Res., 2016.
Timothy P. Lillicrap. Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra. Continuous control with deep reinforcement learning. In Yoshua Bengio and Yann LeCun (eds.), ICLR, 2016.
Hao Liu and Pieter Abbeel. Aps: Active pretraining with successor features. In ICML, pp. 6736-6747 2021a.
11

---

Under review as a conference paper at ICLR 2023
Hao Liu and Pieter Abbeel. Unsupervised active pre-training for reinforcement learning. ICLR, 2021b.
Kendall Lowrey, Aravind Rajeswaran, Sham Kakade, Emanuel Todorov, and Igor Mordatch. Plan online, learn offline: Efficient learning and exploration via model-based control, 2018. URL
attpps:/arzae org/ htth/bbs/111111 01848
Yao Lu, Karol Hausman, Yevgen Chebotar, Mengyuan Yan, Eric Jang, Alexander Herzog, Ted Xiao, Alex Irpan, Mohi Khansari, Dmitry Kalashnikov, and Sergey Levine. AW-opt: Learning robotic skills with imitation andreinforcement at scale. In 5th Annual Conference on Robot Learning (CoRL), 2021. 2021.
Xiao Ma, Siwei Chen, David Hsu, and Wee Sun Lee. Contrastive variational model-based reinforce- ment learning for complex observations. In Proceedings of the 4th Conference on Robot Learning (CoRL), 2020.
Pietro Mazzaglia, Tim Verbelen, and Bart Dhoedt. Contrastive active inference. In A. Beygelzimer, Y. Dauphin, Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021a.
Pietro Mazzaglia, Ozan Çatal, Tim Verbelen, and B. Dhoedt. Curiosity-driven exploration via latent bayesian surprise. ArXiv, abs/2104.07495, 2021b.
OpenAI, Ilge Akkaya, Marcin Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, Jonas Schneider, Nikolas A. Tezak, Jerry Tworek, Peter Welinder, Lilian Weng, Qiming Yuan, Wojciech Zaremba, and Lei M. Zhang. Solving rubik's cube with robot hand. ArXiv, abs/ abss/991000711.. 2019.
Simone Parisi, Victoria Dean, Deepak Pathak, and Abhinav Gupta. Interesting object, curious agent: Learning task-agnostic exploration, 2021. URL https: arxiv http://awwwwwvo.o.////o// wrtp://///bbb//111 13119
Simone Parisi, Aravind Rajeswaran, Senthil Purushwalkam, and Abhinav Gupta. The unsurprising effectiveness of pre-trained vision models for control, 2022.
Deepak Pathak, Pulkit Agrawal, Alexei Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. ICML, 2017.
Deepak Pathak, Dhiraj Gandhi, and Abhinav Gupta. Self-supervised exploration via disagreement. In ICML, 2019.
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever. Language models are unsupervised multitask learners. 2019.
Arthur George Richards. Robust constrained model predictive control. PhD thesis, Massachusetts Institute of Technology, 2005.
Remo Sasso, Matthia Sabatelli, and Marco A. Wiering, Multi-source transfer learning for deep model-based reinforcement learning. 2022. hwll///csss//arri..o///////////
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7833):604-669, 2020.
John Schulman, Philipp Moritz, Sergey Levine, Michael Jordan, and Pieter Abbeel. High-dimensional continuous control using generalized advantage estimation. In ICLR, 2016.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. ArXiv, abs/ 1737776777777.. 2017.
Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-super world models. ICML, 2020.
Harshit Sikchi, Wenxuan Zhou, and David Held. Learning off-policy with online planning, 2020. URL hhttps://arat atwpg//bs/2000.. 10066
12

---

Under review as a conference paper at ICLR 2023
Harshit Sikchi, Wenxuan Zhou, and David Held. Learning ofllcy with online planning. In 5th Anmual Conference on Robot Learning, 2021. URL htLLitttpp///peeeeeee.etttttt
Id=1GNV9SW95e
Harshinder Singh, V. Hnizdo, Adam Demchuk, and Neeraj Misra. Nearest neighbor estimates of entropy. American Journal of Mathematical and Management Sciences, 23, 02 2003.
Richard S. Sutton. Dyna, an integrated architecture for learning, planning, and reacting. SIGART Bull., 2(4):160 163, jul 1991. ISSN 0163-5719. doi: 10.1 10.114551223441222777 URL https?
7/ 7/doi.
Erik Talvitie. Learning the reward function for misspecified model. In Jennifer Dy and Andreas Krause (eds.), JCML, volume 80 of Proceedings of Machine Learning Research, pp. 4838 4847.
PMLR, 10-15 Jul 2018.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmalcki, Josh Merel, Andrew Lefrancq, Timothy P. Lillicrap, and Martin A. Riedmiller. Deepmind control suite. CoRR, abs/1801. 801.00690, 2018.
Geoffrey van Driessel and Vincent Francois-Lavee. Component transfer learning for deep based on abstract representations, 2021. URL htt.:t/////a....oog//////// http//bs/221111 11525
Grady Williams, Andrew Aldrich, and Evangelos Theodorou. Model Predictive Path Integral Control using Covariance Variable Importance Sampling. arXiv e-prints, art. 1 2015.
Philipp Wu, Alejandro Escontrela, Danijar Hafner, Ken Goldberg, and Pieter Abbeel. Daydreamer: World models for physical robot learning, 2022. URL httpps://axxii......// wwwpxiv.rr//bb//2000 14176
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Reinforcement learning with prototypical representations. 2021.
Denis Yarats, Rob Fergus, Alessandro Lazaric, and Lerrel Pinto. Mastering visual continuous control: Improved data-augmented reinforcement learning. In ICLR, 2022.
13

---

Under review as a conference paper at ICLR 2023
APPENDIX
A
NORMALIZATION SCORES
Pro-trainining for 2M environment frames
Domain
Task
URLB Expert
URLB Disagreement
Dreamer 2M
Ours
Flip
799
346 41
778
938 + 12
Run
796
208 + 47
724
596 + 38
Walker
Stand
984
746 107
909
973 14
Walk
971
549 - 117
965
959 -
Jump
888
389 1 196
753
822 1 33
Run
888
337 1 95
904
642 1 99
Quadruped
Stand
920
512 281
945
28
Walk
866
293 117
947
816 61
Reach bottom left
193
124
223
192 1 23
Reach bottom right
203
115 1:
231
192 19
Jaco
Reach top
38
233
197
17
Reach top right
223
139 22
225
212 + 13
Table 2: Performance of expert baseline and the best method pixel-based URLB from Laskin et al. (2021) and performance our oracle baseline (Dreamer@: @ 2M) and best approach, using LBS for unsupervised data collection, after pre-training for 2M frames and fine-tuning for 100k steps.
In Table we report mean scores for the URLB Expert, used to normalize the scores in URLB paper, and for Dreamer@21 2M, which we use to normalize returns of our methods, where both supervised baselines have been trained individually on each of the 12 tasks from URLB for 2M frames. We additionally report mean and standard deviations for the best performing unsupervised baseline from URLB. which is Disagreement (Pathak et al. 2019), and our method (using LBS for data collection). We notice that scores approach the Dreamer@ 2M's scores in several tasks, eventually outperforming them few tasks (e.g. Walker Flip, Quadruped Jump). We believe this merit is due both to the exploration pre-training, which may have found more rewarding trajectories than greedy supervised RL optimization and of the improved Dyna-MPC planning strategy.
B INTEGRATING UNSUPERVISED RL STRATEGIES
We summarize here the unsupervised RL approaches tested and how we integrated them with the Dreamer algorithm for exploration. For all methods, rewards have been normalized during training using an exponential moving average with momentum 0.95, with the exceptions of RND, which follows its original reward normalization (Burda et 2019b), and APS, whose rewards are not normalized because they are used to regress the skill that is closer to the downstream task during FT.
The Intrinsic Curiosity Module (ICM; Pathak er al. (2077) defines intrinsic rewards as error between states projected in a feature space and a feature dynamics model's predictions. We use the Dreamer agent encoder e = (st) to obtain features and train a forward dynamics model g(erle:-1, a1) to compute rewards as:
ICM.
ICM lgg(er (1-1,00-1) ea
ry
As the rewards for ICM require environment states (going through the encoder to compute prediction error), we train reward predictor to allow estimating rewards imagination.
Plan Pxplorx The Plan2Explore algorithm (Sekar et al. 2020) is adaptation of the Disagree- ment algorithm (Pathak et al. 2019) for latent dynamics models. An ensemble of forward dynamics models is trained to predict the features embedding e, fs()),, given the previous latent state and actions, i.e. g(eal 2 at-1, wx). where wi are the parameters of the k-th predictor. Intrinsic rewards defined as the variance of the ensemble predictions:
P2E
Var( /21 0t-1, (1,..., K]}).
Pt
14

---

Under review as a conference paper at ICLR 2023
Plan2Explore requires only latent states and actions, thus it can be computed directly in imagination. We used an ensemble of 5 models.
RND. Random Network Distillation (RND; Burda et al. (2019b)) learns predict the output of a randomly initialized network r((s:) that projects the states into a more compact random feature space. As the random network is not updated during training. the prediction error should diminish for already visited states. The intrinsic reward here is defined as:
RND & (1g(()) (s)))
As the rewards for RND requires environment states (to encode with the random network), we train a reward predictor to allow estimating rewards in imagination.
LBS. In Latent Bayesian Surprise (LBS; Mazzaglia et al. (001b)), they use the KL divergence between the posterior and the prior of latent dynamics model proxy for the information gained with respect to the latent state variable, by observing new states. Rewards are computed as:
Tt LBS
xk DwLlg(=/l a ai)ll a-11)
As the rewards for LBS requires environment states (to compute the posterior distribution), we train a reward predictor allow estimating rewards in imagination.
APT. Active Pre-training (APT; Liu & Abbeel (2021b)) uses particle-based estimator based on the K nearest- nearest-neighbors algorithm (Siggh et al. 2003) to estimate entropy for given state. We implement APT on top of the deterministic component of the latent states 2,, providing rewards as:
APT rt
log E =
where & are the nearest-neighbor states in latent space. As APT requires only latent states, it can be computed directly in imagination. We used = 12 nearest neighbors.
Diversity is All you need (DIAYN: Eysenbach al. (2019) maximizes the mutual information between the states and latent skills w. We implement DIAYN on top of the latent space of Dreamer, writing the mutual information as / (w,, 2x) H (w:) H/ (w, (zt). The entropy H (w,) is kept maximal by sampling estimated learning discriminator q (z.) We compute intrinsic rewards as:
DIAYN.
Unif(u) w.) from discrete uniform prior distribution, while H (wwllz)
DIAYN & log q(wa]2a)
rt
Additionally, DIAYN maximizes the entropy of the actor, so we add entropy maximization term to Dreamer's s objective (Haarnoja et al.) 2018). As DIAYN requires model states and skills sampled from a uniform distribution to compute rewards, we can directly compute them in imagination. For FT, the skill adapted is the one with the highest expected rewards, considering the states and rewards obtained the initial episodes.
APS. Active Pre-training with Successor features (APS; Liu Abbeel (2021 la maximizes the mutual information between the states and latent skills w. We implement APS on top of the latent space of Dreamer. writing the mutual information as / [w,,, 2x) H (2,) (2, ww).. The entropy term ((z, is estimated using a particle-based estimator on top of the deterministic component of the latent states ε,, as for APT, while the term ((2 (w:) is estimated learning discriminator q(2, (w)). The intrinsic rewards for APS can be written as:
APS
APT + log g(w;, xx))
Px x r
As APS requires model states and uniformly sampled skills to compute rewards, we can directly compute them in imagination. For FT, the skill to adapt is selected using linear regression over the states and rewards obtained in the initial episodes (Liu Abbeel, 2021a).
15

---

Under review as a conference paper at ICLR 2023
C DYNA-MPC
To further improve data efficiency, we chose to use an hybrid planner that combines reinforcement learning and MPC (Hansen al. 2022 Sikchi et al. 2020 Lowrey al.) 2018). Previous works leveraged model-free off-policy algorithms (Hansen et al. 2022 Sikchi et al. 2020 to learn the actor and critic in a more computationally efficient manner. The policy used to act on the environment combines action samples from the actor network with MPC, while the critic and the actor are learned "offline" from previously collected data. This has several benefits but also leads to issue referred to as "actor divergence" (Sikchi et al. (020), which consists of the policy used for data collection being different from the policy that used learn the critic.
In our study, we found that using the PT world model to learn the actor and the critic is crucial to improve data-efficiency during FT (see Figure 3) Thus, we discard the option learning the actor and critic with off-policy deep RL. Instead, we design new hybrid planner, which we call Dyna-MPC. that learns actor and critic functions in the model imagination (Sutton 1991), using the Dreamer algorithm (Hafner al. 2019a), and then combines their predictions with MPPI (Williams et al. 2015) for acting on the environment. By doing so we mitigate the "actor divergence" issue as actor and critic learned on-policy the trajectories generated with the model.
The critic is learned in the model's imagination, computing the expected value of the actor's actions using GAE-À estimates of the returns (Schulman al., 2016 Hafner et 2019a):
(1 - 2) (2++1) + AV4 if << H, V = r + Tt ( (zH)
if 1= H,
where F1 is the reward for state yielded by the reward predictor of the world model, and H is the imagination horizon. When computing returns for MPPI we use the same return estimates.
At each time step, we use MPPI to select the best action. MPPI iteratively fits the parameters of time-dependent multivariate Gaussian distribution with diagonal covariance, updating mean and standard deviation parameters using an importance weighted average of the top-k trajectories with the highest estimated returns. At every step. N trajectories r, {0o,;: 01 an.;} of length are obtained sampling actions from the distributions a, Nu,,, 1) and N₂ trajectories are sampled from the actor network a₂ TT(a; (2)) and their outcomes are predicted using the model. At each MPPI iteration, the distribution parameters are updated as follows:
Σ Ei 1 MM Ω
2r o= max
ΣN LT(T; - 11)
e
E Q
where Ω. = exp(r is temperature parameter, indicates the trajectory is in the top-k, and a clipping factor to avoid too small standard deviations (Hansen al., 2022). To reduce the number of iterations required for convergence, we reuse the I-step shifted mean obtained at the previous timestep (Argenson & Dulac Ammol (020).
16

---

end for Under review as a conference paper at ICLR 2023
D ALGORITHM
Algorithm 2 Unsupervised Model-based Pre- Fre-Trainnng for Data-efficient Control from Pixels
Require: Actor 0, Critic v, World Model &
1:
Intrinsic reward ydm! extrinsic reward yau
2: Environment, M, downstream tasks Te. &o E [I,...,
3: Pre-train frames Ner. fine-tune frames NFT. environment frames/update т
4: Initial model state 2p, hybrid planner Dyna-MPC, replay buffers Dyr, Drr
5:
6: // Pre-training
7: for = .Pp do
8:
Draw action from the actor, a, ~ TT(ael2))
9: Apply action the environment, Sc+1 ~ P(- ax)
10: Add transition to replay buffer, Dpr + Drr U (St, a,, S141)
11:
Infer model state, 21+1 (22++1/2,, at fo(s++))
12:
if I mod T = 0 then
Update world model parameters á on the data from the replay buffer Dpr Update actor-critic parameters {0, o in imagination, maximizing r
14:
end if
16: end for
Output pre-trained parameters {t/rr, OPT, ΦFT}
18:
19: // Fine-tuning
20: for Te E .T TM] do
21: Initialize fine-tuning world-mode with Ørr
22: (Optional) Initialize fine-tuning actor with 6pT
23:
for -= O...., NFT do
24:
Draw action from the actor, a, no (0el2e)
25:
Use the planner for selecting best action, a, ~ Dyna-MPC(z))
Apply action the environment, St+1, r (( [s, ar)
27:
Add transition to replay buffer, DFr + Drr (S,, a, 17 8141)
Infer model state, 24+11 q(2++1/26 Ge, fo(S(41)) if t mod = 0 then
29:
30:
Update world model parameters on the data from the replay buffer DFr Update actor-critic parameters {0, + in imagination, maximizing ext
31:
32:
end
33:
34: Evaluate performance on Ti
35:
end for
17

---

Under review as a conference paper at ICLR 2023
E ADDITIONAL RESULTS
We present complete results, for each unsupervised RL. method, for the large-scale study experiments presented in Section 3
Dreameng
Dentom
APS
Pre-trining for 100k frames
Pre-traiinng for So0k framess
Pro traitrainnng for IN frames
Pro traiiing for 2M frames
Normatzed (84
80
60
60
60
40-
40-
40-
Quadruped
Jaco
Quidruped
jaco
Quadruped
jaco
Walcer
Quadd uped
jaco
Figure 11: Complete results for Section 3.1 (Figure 2b
RestEopions
(a) Model
(b) Model, Actor
(c) Model, Actor, Critic
Figure 12: Complete results for Section 3.2 (Figure
MAA
Preiiiiiing for humes
100
330
40
Dawhepcd
Ada
OustTuedd
(a) FT
(b) FT Dyna-MPC
Figure 13: Complete results for Section 3.3 (Figure
18

---

Under review as a conference paper at ICLR 2023
Bunnkked Bettern (S)
Numolized Return Sa
Nomasteed Rettrn
100
90
34
Quadruped
LBS
10c
Wtrm Reeum t90
200.0
00.5
20.2
2H4
4M
SM
114
2M
25
144
41
2M
3M Overall
Walker
Quadruped
+
P2B
Noraltee Bettr (90)
Butuum (M))
100
alfnnt Retum (60
100
30-
200
S0LE
002
Walker
Quadruped
Overall
Normaloed betum (62)
160.5
Vormalion Retum DM6
200.3
Vemallend Retum
21.3
50.5
250
2M
9M
3M
25
Walker
Quadruped
Ja00
AF5
BormaTzed Return
Horatted
1M
2.5
Walcer
Quadruued
J8C0
Overall
Figure 14: Longer pre-training, Fine-tuning performance of our method when pre-training for longer than 2M steps. Every bar reports mean performance and standard errors.
Can a pre-training stage longer than 2M frames be beneficial? In Figure 14 we report FT results with our full method, every 1M frames up to 5M PT frames. The aggregated results show that, adopting our method, longer PT can increase performance further, especially until 4M steps. The performance in all domains keeps increasing remains steady until 5M steps, with two exceptional cases, Walker for Plan2Explore and Jaco for APS, where performance drops between 4M and 5M steps.
For these experiments, kept the size of the model and all the hyperparameters unvaried with respect to the 2M PT frames experiments but we increased the replay buffer maximum size to 5M frames. Increasing model capacity, and adopting additional precautions, such annealing learning rate, it is possible that the agent could benefit even more from longer pre-training and we aim to analyse this more in details for future work.
19

---

Under review as a conference paper at ICLR 2023
F RWRL SETTINGS
We take the Quadruped and Walker tasks from the RWRL benchmark and replace the low-dimensional sensor inputs with RGB camera inputs. While this removes some of the perturbations planned in the benchmark (Dulac-Arnold et al. 20200, such as noise in the sensors, it introduces the difficulty of a different dynamics in pixel space (due to the other perturbations), compared to the one observed during pre pre-trainnn training in the vanilla simulation environment.
Setting
Easy
Medium
Hard
System Delays
Time Steps
Time Steps
Time Steps
Action
3
Rewards
10
20
40
Action Repetition
3
Gaussian Noise
Std. Deviation
Std. Deviation
Std. Deviation
Action
0.1
0.3
1.0
Perturbation
Max]
Std.
[Mi]
Std.
[Min,Max]
Std.
Quadruped
(shin length)
10.25, 0.3]
0.005
[0.25, 0.8]
0.05
[0.25, 1.4] [Min,Max]
0.1
Perturbation
[Mx]
Std.
[Min,Max]
Std.
Std.
Walker
(thigh length)
[0.225, 0.25]
0.002
[0.225, 0.4]
0.015
[0.15, 0.55]]
0.04
Table 3: Perturbations setting for each challenge of our adapted tasks from the RWRL benchmark, in increasing levels of intensity.
G EXTENDED ANALYSIS
We note that, to run the experiments faster, we did not use Dyna-MPC for the extended analysis. Furthermore, the Jaco tasks used slightly differ from the original ones in URLB, only in that the target to reach cannot move. This allows consistency of the reward function between PT and FT, so that reward predictor can be trained on laweledl PT data. However, because of this change, the performance in Jaco may differ from the other main results (particularly in Figure R and Figure
G.1 LEARNING REWARDS ONLINE
In Figure of the main text, we measure the gap in performance between pre-trained agents that have no knowledge of the reward function at the beginning of fine-tuning and agents whose reward predictor is initialized from reward predictor learned on top of the unsupervised pre-training data (violating the URLB settings). Crucially, the agent during unsupervised PT can learn the reward predictor without affecting neither the model learning or the exploration process. To not affect the model, gradients are stopped between the reward predictor and the rest of the world model. To not affect exploration, the rewards used to train the agent's 's actor and critic remain the intrinsic rewards, for exploration.
G.2 ZERO-SHOT ADAPTATION
Using agents that have access to PT reward predictor, we explore the idea of zero-shot adaptation using MPC, which trying to solve the URLB tasks using only planning and the pre-trained world model and reward predictor. In order to obtain good performance, this assumes that the model correctly learned the dynamics of the environment and explored rewarding transitions that are relevant to the downstream task, during pre-training In FigureDoo of the main text, compare the results of performing MPC in zero-shot setting (ZS) with the performance of an MPC agent that is allowed 100k frames for fine-tuning (FT). As for the MPC method, we employ MPPI (Williams et al., (015). Because these experiments are particularly expensive to run, we just them on the agents trained with the Plan2Explore URL approach.
We observe that the performance of zero-shot MPC is generally weak. While it overall performs better than the non-pre-trained model, simply applying MPC leveraging the pre-trained world model
20

---

Under review as a conference paper at ICLR 2023
and reward predictor trained the pre-training stage data is not sufficient to guarantee satisfactory performance. The fact that exploiting the fine-tuning stage using the same MPC approach generally boosts performance demonstrates that the model has major benefit from the FT stage. Still, the performance of MPC generally lacks behind the actor-critic performance, suggesting that, especially in higher-dimensional action space such as the Quadruped one, amortizing the cost of planning with actor-critic seems crucial to achieve higher performance.
G.3 LATENT DYNAMICS DISCREPANCY
Model misspecification useful measure to assess the uncertainty or inaccuracy of the model dynamics. is computed as the difference between the dynamics predictions and the real environment dynamics. The metric helps build robust RL strategies, that take the dynamics uncertainty into account while searching for the optimal behavior (alvitie 2018). However, with pixel-based inputs the dynamics of the environment are observed through high-dimensional images. And this in- in-turn could hurt the metric evaluation, since the distances in pixel space can be misleading. In approach, use a model-based RL agent that learns the dynamics model in a compact latent space 2.
Our novel metric, Latent Dynamics Discrepancy (LDD),, quantifies the "misspecification" of the learned latent dynamics accordingly. The metric quantifies the distance between the predictions of the pre-trained model and the same model after fine-tuning on downstream task. However, as the decoder of the world model gets updated during fine-tuning, the latent space mapping between model states z and environment states s might drift. For this reason, we freeze the agent's decoder weights, so that the model can only improve the posterior and the dynamics. This ensures that the mapping 2 S remains unchanged and allows to compare the dynamics model after fine-tuning with the one before fine-tuning. In order to measure the distance between the distribution output by the dynamics network, we chose the symmetrical -aasnn-Shnnoon divergence:
LDD E(,,,))) Dis[PFT (2 1 az) ((Pp))llp((ee) Ze+ 1 dz)l]
where the expectation is taken over the previous model states 2, sampled from the fine-tuned posterior (rr(=;), actions 0ę-1 sampled from an oracle actor π (ael2z). so that we evaluate the metric on optimal trajectories, whose environment's state distribution corresponds to the stationary distribution induced by the actor 8t d= (s:). We used 30 trajectories per task in evaluation.
We observe in our experiments that there exists a correlation between the metric and the performance ratio between zero-shot model and fine-tuned model (see Figure 10 in the main paper). The key observation is that major updates in the model dynamics during fine-tuning phase played an important role in improving the agent's performance, compared to the pre-trained model and zero-shot performance. Future research may attempt to reduce such dependency by either improving the model learning process, so that the pre-trained dynamics could have greater accuracy, or the data collection process, proposing URL methods that directly aid to reduce such uncertainty.
G.4 UNSUPERVISED REWARDS AND PERFORMANCE
We further analyzed the correlation between the normalized performance the different exploration agents and their intrinsic rewards for optimal trajectories obtained by oracle agent. A strong negative correlation between the two factors should indicate that the agent is more interested in seeing the optimal trajectories when its performance is low the task.
We observe that there is negative correlation between Plan2Explore (P2E), ICM, LBS's performance and their intrinsic rewards, while found ~0 correlation for RND (see Table in the main text). Out of the methods tested, LBS significantly demonstrated the correlation, as its p-value is < 0.05. This is likely one of the key factors for the high performance of the agent using LBS on the benchmark.
One possible explanation is that LBS searches for transitions of the environment that are difficult to predict for the dynamics, so the model likely learns those transitions more accurately, facilitating planning during the fine-tuning stage. Another potential explanation is that, given the high correlation between intrinsic and extrinsic rewards, the actor initialized by LBS performs better at the beginning of FT, speeding up adaptation.
21

---

Under review as a conference paper at ICLR 2023
H
HYPERPARAMETERS
Most of the hyperparameters we used for world-model training are the same as the original DreamerV2 work (Hafner et al. 2021). Specific details are as outlined here:
Name
Value
World Model
Batch size
50
Sequence length
50
Discrete latent state dimension
32
Discrete latent classes
32
GRU cell dimension
200
KI. free nats
KL balancing
0.8
Adam learning rate
3 10
Slow critic update interval
100
Actor-Critt
Imagination horizon
y parameter
0.99
parameter
0.95
Adam learning rate
8 10-5
Actor entropy loss scale
1 10-4
Dyna-MPC
Iterations
12
Number of samples
512
Top-k
64
Mixture coefficient (Actor/CEM)
0.05
Min std (fixed)
0.1
Temperature
0.5
Momentum
0.1
Planning horizon
5
Common
Environment frames/update MLP number of layers
4
MLP number of units
400
Hidden layers dimension
400
Adam epsilon
1 10-5
Weight decay
10-6
Gradient clipping
100
Table 4: World model, actor-critic, planner (Dyna-MPC) and common hyperparameters.
For the pure MPC-based experiments, we increased the number of MPPI samples from 512 to 1000, the number of top-k from 64 to 100, and the horizon from 5 to 15, to compensate for the absence of actor network's samples and the critic's predictions in the return estimates.
22


Reviews:
[{"review_id": "jK3vuvuxWpC", "reviewer": "Reviewer_xEiB", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper considers the task of unsupervised pretraining of the reinforcement learning (RL) agents.\nIt analyses and compares different design choices for pretraining and fine-tuning the pretrained components; it shows the improvement of using proposed model-based reinforcement learning over existing model-free technique, DrQ.", "summary_of_the_review": "Based on the merits of the evaluation and experimental contribution, as well as the quality of writing, I recommend acceptance. ", "strength_and_weaknesses": "Strengths:\n- The paper is really well written: the narrative is shaped around choosing different components for pre-training and fine-tuning, with the structure of the overall approach neatly summarised in Section 3.4\n- It contains compelling evidence of benefits of using proposed model-based reinforcement learning method over existing state-of-the-art model free technique \n- The analysis includes necessary justification of the architecture: choice of the best possible pretraining model\n\nWeaknesses:\n- The model is engineered out of the existing components (apart from incremental, in a positive sense of this word, contribution of Dyna-MPC); this is not, however, a problem at all in my opinion as we need large-scale studies showing the impact of architectural choices, and this is a good contribution for the community and this very knowledge of how to improve architectural choices is a novel aspect of the paper", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "4: The contributions are significant, and do not exist in prior works."}, "extra_text": {}}, {"review_id": "ixbojj58KDt", "reviewer": "Reviewer_kuHx", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The paper studies the problem of reward free learning, where the setup involves a pre-training stage (when an intrinsic reward is used based on different prior methods) and a fine-tuning stage (when a particular extrinsic reward is provided). The claim is that model-based methods perform better than model-free methods in such a setup. The authors perform certain ablations where different components of the model-based method (in this case the actor, critic, and dynamics model of the Dreamer algorithm) are shown to affect fine-tuning performance differently. Finally, the authors deploy a MPC style planning algorithm instead of using the learnt actor when fine-tuning to new rewards. ", "summary_of_the_review": "The paper lacks a coherent story and it’s hard to parse the clear contributions it makes. The Dyna-MPC algorithm as well as the ablations are not the main contributions in my opinion. As far as the main claim goes, I am not sure why this is surprising, i.e. shouldn’t any model-based method be expected to perform better than model-free method when the reward function is changed? I could not understand why Figure 2 has a flat line for the Drq@100k and Dreamer@100k cases. What does 100k refer to here? Can you provide some commentary on why, in the setup you consider, model-based methods should outperform model-free ones by such a huge margin? I think this is crucial to the reader’s understanding of the contributions and motivation of the paper. ", "strength_and_weaknesses": "Strengths\n\n- Important and relevant problem setup\n- Fairly thorough experiments\n\nWeakness\n \n- Lack of a coherent story\n- Unclear primary contributions\n- Somewhat hard to read through", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "_7DwHafn5I", "reviewer": "Reviewer_zaP8", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The method studies the generalization capabilities of unsupervised RL empirically, and develops a hybrid planner with strong asymptotic performance and high sample efficiency in standard benchmarks. ", "summary_of_the_review": "The empirical findings from the work are insightful for the community, and the proposed method stemming from these findings achieves strong performance in standard benchmarks. Both are valuable contributions to the community. ", "strength_and_weaknesses": "Strength\n* This work offers a large-scale benchmarks for different URL techniques and compare the usefulness of different pre-trained components. \n* The developed method has strong empirical performance. \n\nWeaknesses\n* The main weakness is novelty since most components in this work come from existing works.\n* The paper pinpoints two decisions to make to apply the proposed framework for tasks outside URLB in Section 3.4. Within URLB, these decisions are made given empirical benchmarking results. But more insights or explanations of when and why different techniques work in different settings, and how much alignmenent of the pretraining and downstream task the proposed mechanism can handle, will further strengthen the contribution of this work.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "LRfJjidDzDy", "reviewer": "Reviewer_12hB", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors study model-based methods for tackling the unsupervised reinforcement learning benchmark (URLB). First, the authors show that substituting DrQ (a model-free agent) with DreamerV2 results in increased performance and scaling on URLB. Second, the authors investigate the transfer of specific components (model, actor, critic) from pre-training to fine-tuning, finding that transferring model and actor is beneficial but additionally transferring critic compromises performance. Third, the authors propose Dyna-MPC, a planning procedure that builds upon model predictive path integral (MPPI) control, using the actor as an additional source of sampled trajectories, and the critic and model to score trajectories. This is shown to further improve fine-tuning. These improvements altogether result in dramatic improvement on URLB. Finally, the authors assess transfer under distribution shifts via the real-world reinforcement learning benchmark (RWRL), and find that their proposed model without Dyna-MPC performs best.", "summary_of_the_review": "This work is rather tricky to assess. While its proposed methods do provide dramatically improved empirical results on an existing benchmark, as argued above there are important limitations to both the novelty and significance in this improvement. However, the thoroughness with which the empirical evaluation was conducted is admirable, and there are a couple of genuinely new insights provided on the application of model-based learning to unsupervised RL. I currently lean towards rejection, but look forward to discussion with the authors and other members of the reviewing team to inform my final recommendation.", "strength_and_weaknesses": "Strengths\n\nThe results are particularly strong. The progression of experiments building towards the proposed method is well structured and informative. The set of methods, comparisons, and ablations is extensive.\n\nWeaknesses\n\nThe gains provided by model-based learning are probably significantly inflated due to the simplicity of the three URLB environments, which enables in-imagination training and planning to actually work. In contrast, in Atari, which has far more visual variation per game despite arguably simpler dynamics, model-based methods have still yet to come close to model-free methods, except when the allowed amount of environment interaction or compute is limited. Thus, while impressive, the improvements coming from this work should be taken with a grain of salt, and might not hold for environments with more realistic amounts of variation and complexity.\n\nThe application of model-based learning for unsupervised pre-training for RL has already been known to be effective (e.g. Plan2Explore). I'm not sure why the URLB authors did not benchmark model-based baselines like Plan2Explore, but regardless, the improvements over the DrQ-based runs is expected given Plan2Explore, especially given the environments in Plan2Explore and URLB are both sourced from the DeepMind Control Suite. Regrettably, it seems that the authors formulated their main hypothesis without realizing the above. I recommend rephrasing the contributions in the abstract and intro to convey that the results in this work corroborate and extend upon the findings first provided by Plan2Explore.\n\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject