Unsupervised Model-based Pre-training for Data-efficient Reinforcement Learning from Pixels
Abstract
Reinforcement learning (RL) aims at autonomously performing complex tasks.
To this end, a reward signal is used to steer the learning process.
While successful in many circumstances, the approach is typically data hungry, requiring large amounts of task-specific interaction between agent and environment to learn efficient behaviors.
To alleviate this, unsupervised RL proposes to collect data through self-supervised interaction to accelerate task-specific adaptation.
However, whether current unsupervised strategies lead to improved generalization capabilities is still unclear, more so when the input observations are high-dimensional.
In this work, we advance the field by closing the performance gap in the Unsupervised RL Benchmark, a collection of tasks to be solved in a data-efficient manner, after interacting with the environment in a self-supervised way.
Our approach uses unsupervised exploration for collecting experience to pre-train a world model. Then, when fine-tuning for downstream tasks, the agent leverages the learned model and a hybrid planner to efficiently adapt for the given tasks, achieving comparable results to task-specific baselines, while using 20x less data.
We extensively evaluate our work, comparing several exploration methods and improving the fine-tuning process by studying the interactions between the learned components.
Furthermore, we investigate the limitations of the pre-trained agent, gaining insights into how these influence the decision process and shedding light on new research directions.
Decision Awareness in Reinforcement Learning Workshop at the 39th International Conference on Machine Learning (ICML) , Bal timore, Maryland, USA, 2022.
Copyright 2022 by the author(s).
1. Introduction
Modern successes of Reinforcement Learning (RL) have been realized through specialized agents that rely on taskspecific rewards.
Such autonomous agents have shown promising results scaling to raw high-dimensional inputs, e.g. playing Atari video games directly from pixels (Silver et al., 2016; Mnih et al., 2016), learning robotic manipulation policies from raw sensory input (Levine et al., 2016), etc.
However, training an agent for each task individually requires a large amount of task-specific environment interactions, incurring huge redundancy.
On the other hand, training agents that can generalize quickly on more than a single task is often desirable towards building intelligent autonomous systems.
Developing algorithms that can efficiently adapt and generalize to new downstream tasks has hence become an active area of research in the RL community.
In computer vision and natural language processing, unsupervised learning strategies have enabled learning representations without supervision, that can quickly be adapted for a variety of tasks (Chen et al., 2020;
Radford et al., 2019).
In a similar fashion, unsupervised RL (URL) methods have fo- cused on improving exploration performance (Pathak et al., 2017; Burda et al., 2019a; Bellemare et al., 2016) or learning generalizable skills (Eysenbach et al., 2019; Hansen et al., 2019).
A key aspect here is to be able to generate useful behaviors in a self-supervised manner, without requiring any reward annotations, and subsequently adapt the learned behaviors to downstream tasks.
Recently, Laskin et al.
(2021) introduced the Unsupervised
RL Benchmark (URLB), where agents are tested on different downstream tasks across three challenging continuous control domains of the DM Control Suite (Tassa et al., 2018).
In this benchmark, an agent is first allowed a taskagnostic pre-training stage, where it can interact with the environment in an unsupervised way, followed by a finetuning stage where, given a limited budget of interactions with the environment, the agent should be able to quickly adapt for a specific task.
The results obtained by Laskin et al.
(2021) suggest that none of the tested unsupervised RL algorithms completely solve the benchmark and that models pre-trained by unsupervised RL are not universally better than random initialization.
Furthermore, a large gap in performance was found between using state-based inputs versus high-dimensional pixel-based inputs, as the latter proved to be more difficult and sample inefficient.
However, a string of recent literature on representation learning and model-based RL has shown that learning from pixel-based inputs can be effective, given an appropriate representation of input (Srinivas et al., 2020; Laskin et al., 2020) and/or learning the environment dynamics (Hafner et al., 2019a;b).
On control tasks (Tassa et al., 2018), these algorithms match state-based efficiency when learning from pixels, in a fully supervised setting.
So the question remains as to how do we translate similar performance gains to unsupervised RL.
In this work, we show that, with the right precautions, unsupervised RL can be employed to nearly solve the URL benchmark from pixels.
We present an unsupervised model-based RL approach that significantly improves dataefficiency when fine-tuned in a low-data regime, by planning and adapting task-specific policies on the synthetic data generated by the unsupervised pre-trained model.
Furthermore, we empirically study the interactions between the multiple modules learned during unsupervised pre-training and analyze the quality of the learned model, identifying appropriate design choices and shedding light on what could prevent the agents from adapting faster.
Our efforts are aimed at understanding and scaling unsupervised RL algorithms to operate efficiently on highdimensional images.
Our contributions can be summarized as: (i) the design of a class of model-based unsupervised
RL approaches, which enable fast adaptation after an unsupervised pre-training stage, (ii) a study of the interplays between the pre-trained modules that allow to improve sample efficiency during the fine-tuning stage, (iii) an analysis of the model learned through unsupervised interaction with the environment, aimed at understanding what aspects could be improved to facilitate fast adaptation, (iv) novel quantitative metrics to evaluate misspecifications in the learned model.
We demonstrate through extensive experimentation that, following our approach, it is possible to bridge the performance gap between state-based and pixel-based inputs, and to achieve the asymptotic performance of supervised RL agents (Figure 1).
2. Preliminaries
The RL setting can be formalized as a Markov Decision
Process (MDP), denoted with the tuple {S , A , T, R, γ } , where S is the set of states, A is the set of actions, T is the state transition dynamics, R is the reward function, and γ is a discount factor.
The objective of an RL agent is to maximize the expected discounted sum of rewards over time for a given task, also called return, and indicated as
[Equation]
− 1) r k .
In continuous action settings, one popular approach to predict the most rewarding actions is to combine a model that learns to output the best action given a certain state, referred to as the actor model, and a model that learns to estimate the expected value of the actor’s actions over time, given a certain state, referred to as the critic model.
Actor-critic algorithms can be combined with the expressiveness of neural network models to solve complex continuous control tasks (Haarnoja et al., 2018;
Lillicrap et al., 2016; Schulman et al., 2017).
In this work, we investigate the problem of fast adaptation for a downstream task, after a phase of unsupervised training and interaction with the environment.
We adopt the URL benchmark, which consists of three control domains, Walker , Quadruped and Jaco , and twelve tasks, four per each domain.
Consistently with URLB (Laskin et al., 2021), our experimental procedure is made of two phases: a pre-training (PT) phase, where the agent can interact with a task-agnostic version of the environment for up to 2 M steps, and a fine-tuning phase (FT), where the agent interacts with the same environment, being provided a task to solve and a limited budget of 100 k steps.
During the PT phase, rewards are removed from the environment.
Sensible information about the environment can be obtained by exploring the domain-dependent dynamics, which will remain unchanged in the downstream tasks.
During FT, the agent receives taskspecific rewards when interacting with the environment.
As the agent has no prior knowledge of the task, it should both understand the task and solve it efficiently, in the limited budget.
Crucially, we focus on the pixel-based setup of URLB, where the environment is perceived by the agent only through images.
In this setting, the performance of several exploration strategies, combined with a state-of-the-art model-free approach (Yarats et al., 2021a), were shown to lack behind the asymptotic performance of an RL agent trained on the downstream task, as reported in (Laskin et al., 2021) and Figure 1.
We believe one of the causes of this is that model-free RL algorithms cannot successfully leverage the information observed about the environment dynamics during PT, as they rely uniquely on actor and critic’s predictions.
To overcome this limitation, we ground our work on a model-based RL agent, whose learned model should allow preserving important information about the environment.
3. Approach
In order to perform well on the URL benchmark, it is important that an 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 understand and master the downstream task.
In this section, we expand on how we addressed these challenges, giving rise to an approach that nearly solves the benchmark by achieving 98.5 ± 4.7% of a supervised RL agent’s overall performance (Figure 1).
An overview of the end-to-end approach is illustrated in Figure 2 and a detailed algorithm is presented in Appendix C for reference.
3.1. Model-based Agent
We build our model-based agent upon DreamerV2 (Hafner et al., 2021), whose agent attempts to learn a world model (Ha & Schmidhuber, 2018; Hafner et al., 2019b; 2021) that allows predicting the outcomes of future actions in the environment.
The environment dynamics is captured into a latent space Z , which allows a compact representation of the high-dimensional inputs of the agent.
The world model consists of the following components:
Encoder:
[Equation]
The model states z t have both a deterministic component, modeled using the recurrent state of a 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 et al., 2019b;a).
In order to plan actions, the agent learns latent actor and critic networks:
[Equation]
The actor is used to generate actions, given the model state, while the critic estimates the expected return for a certain model state, when following the actor’s actions.
Both components are trained online within the world model, by imagining the model state outcomes of the actions produced by the actor, using the model dynamics.
Rewards for imagined trajectories are provided by the reward predictor and combined with the critic predictions to produce a GAE- λ estimate of the returns (Schulman et al., 2015).
The actor maximizes such an estimate of the returns, backpropagating its gradients through the model dynamics.
For the encoder and the decoder networks, we used the same architecture as in Hafner et al.
The hyperparameters for the agent, which we keep fixed across all domains and tasks, can be found in Appendix D.
3.2. Unsupervised Pre-training
In the PT stage, different unsupervised RL strategies can be used to explore the environment and 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 to reduce sample complexity during FT.
As we employ a model-based agent, we use the experience collected to train the agent’s world model, along with the actor and the critic networks.
We note that, during PT, the reward predictor of the world model is either unused or used to predict intrinsic rewards, according to the unsupervised approach employed, as reward information should not be available to the agent.
During FT, the reward predictor is trained to predict the downstream task rewards.
Unsupervised RL methods can be grouped in three categories: knowledge-based, data-based and competencebased (Schmidhuber, 2010; Laskin et al., 2021). We study multiple approaches, focusing primarily on knowledgebased methods as these combine well with the model-based nature of our agent, and implement LBS (Mazzaglia et al., 2021), ICM (Pathak et al., 2017), RND (Burda et al., 2019b), and Plan2Explore (Sekar et al., 2020).
As a data-based approach, we choose APT (Liu & Abbeel, 2021b), and as a competence-based approach, we use DIAYN (Eysenbach et al., 2019).
Finally, we add a Random action baseline, as a maximum entropy approach (Haarnoja et al., 2018).
Details on these methods and how we combined them with our model-based agent are discussed in Appendix B.
3.3. Fine-tuning for Downstream Tasks
During the unsupervised PT phase, the agent collects experience from the environment that is used to train several components: a task-agnostic world model (without the reward predictor), an actor and a critic network.
Moving to the FT phase, the pre-trained weights of these components can be copied into a new instance of the model, aiming to leverage previous experience for faster adaptation.
Since the domain dynamics stays the same between the
PT and FT phases, initializing the world model with the pre-trained one should facilitate adaptation.
However, the reward is changing from pseudo-reward to task reward when changing from the PT phase to FT phase.
Hence, it is not clear if pre-training of the actor and critic can help for the downstream task.
To shed light on this question, we conduct experiments in Section 4 to determine if it useful to transfer the actor and the critic.
Unless specified, for our Default
FT model, which we refer to as (w/ model, w/ actor, w/o critic), we copy the weights of the pre-trained world model and actor but initialize the critic from scratch.
As we train a latent world model, we can exploit modelbased planning to adapt with limited additional environment interaction.
When an accurate model of the environment is available, traditional model-based control approaches, such as Model Predictive Control (MPC) (Williams et al., 2015;
Chua et al., 2018; Richards, 2005), can be used to plan the agent’s action.
Nonetheless, using an actor and a critic has several advantages, such as amortizing the cost of planning by caching previously computed (sub)optimal actions and amortizing the cost of computing long-term returns from a certain state, without the need to imagine outcomes that are far in the future.
We found it useful to adopt a hybrid planning strategy, which exploits the actor and critic’s predictions as well as an evolutionary sampling strategy based on the Cross-Entropy Method (CEM; Rubinstein & Kroese (2004)).
4. Experiments and Analysis
In all the experiments, the results show average normalized returns with error bars showing the standard deviation.
To normalize results in a comparable way for all tasks, we train
Unsupervised a fully-supervised agent with 2M steps per each task.
We use the mean performance of this agent, which we refer to as "oracle", as the reference scores to normalize results in the plots (details in Appendix A).
For all experiments, results are presented with at least three random seeds.
Model-based URLB.
The results of the different exploration approaches are shown in Figure 3.
Results are presented by taking snapshots of the agent at different times during training, i.e. 100k, 500k, 1M and 2M steps, and fine-tuning the pre-trained policies and models for 100k steps.
As opposed to the model-free experiments in URLB, where they fine-tuned the actor and the critic, we found that leveraging a pre-trained world model during fine-tuning dramatically improves the performance.
Simply using random actions for unsupervised exploration already increases performance compared to a supervised agent trained from scratch for 100k steps ( Dreamer@100k ).
This is in contrast to the results in (Laskin et al., 2021), where the improvements when exploiting pre-training were less significant.
Unsupervised exploration strategies generally lead to higher performance in complex tasks, with performance that increases over time in the Walker and Quadruped tasks.
We note that the performance in the Jaco domain is less consistent over time and tends to have higher variance.
We hypothesize that this is due to a greater discrepancy between pseudo-reward and task-reward for this environment.
DIAYN underperforms compared to the other methods, providing some support to the observation in (Laskin et al., 2021), claiming that current competence-based approaches tend to perform worse than the rest.
Exploiting Pre-Trained Modules.
The results of the ablation studies on fine-tuning different sets of pre-trained modules are presented in Figure 4, averaging across all un- supervised RL methods.
Overall, the default configuration, which reuses the weights of the world model and the actor performs best.
Initializing the agent with the pre-trained actor is particularly useful in the Walker and Quadruped domains, but it is harmful in the Jaco tasks.
A possible explanation for this could be that the exploration actor that is transferred from pre-training might have a precise explorative goal when brought to the fine-tuning stage, which may be particularly far from the target state of the task.
In the
Walker and Quadruped tasks, where the rewards are denser, exploring multiple states, even far from the downstream task behavior, provides useful reward information to exploit for the task.
On the other hand, for the Jaco sparser tasks, if the exploration is initialized to reach a point that is too far from the reaching target, it might become arduous to encounter useful rewarding states within the reduced budget of FT.
This hypothesis is supported by the fact that the random action methods and the DIAYN approach, which are believed to explore less than the others (thus, staying closer to the initial agent’s position), perform well, and in some cases even better, than the other approaches (Figure 3).
Initializing the critic with the pre-trained one has little impact in Walker and Quadruped but has been problematic in the Jaco domain, likely because of the more sparse rewards of the Jaco tasks, which are very different from the dense pseudo-rewards used to pre-train the critic.
For this reason, we default to not reusing the critic’s weights.
Finally, the world model is confirmed to be the most valuable component to initialize, as, independently from the other components, initializing the world model brings a significant improvement in performance.
Detailed results per each method are available in Appendix E.
[Figure] Planning.
In order to better exploit the pretrained model during FT, while also leveraging the advan- tages of using an RL-like learning mechanism, we employ a hybrid planning strategy that combines MPC and temporal difference actor-critic learning, as in the recently proposed TD-MPC approach (Hansen et al., 2022).
TD-MPC claims improved performance, thanks to the combination of short-term planning using MPC and long-term predictions, using the actor-critic architecture.
We implement the hybrid planner on top of our models and compare the performance with the standard actor-critic learning in Figure 5.
We observe that using a planner slightly improves the performance in all domains.
Combining the improved planning strategy with the previous insight of not initializing the weights of the actor in the Jaco domain, we obtain our overall best performance (Jaco w/o PT actor + Plan), with the LBS-based model being the overall best-performing method ( 98.5 ± 4.7% ; Figure 1).
Detailed results for each method are available in Appendix E.
4.1. Analysis
First, we showed how we managed to bridge the performance gap in the URL benchmark by leveraging modelbased unsupervised RL, improving the exploitation of the different agent’s components during FT, and leveraging planning to improve data efficiency.
Now, we focus on providing an analysis of the world model learned during the unsupervised pre-training stage, aiming to gain insights on which aspects improve decision awareness in the URLB setting, i.e. how can we best exploit the two stages of training to improve the agent’s decision making.
To shed light on this, we analyze discrepancies in the model’s dynamics and in the reward prediction process that is leveraged during the fine-tuning planning.
Zero-shot Performance.
How useful is the model learned during the unsupervised pre-training stage?
To gain insights into this matter, we perform some additional tests where we provide the FT agent with a pre-trained reward predictor, that we train on the data collected during PT, separately from the agent.
Given such a reward predictor, it should be possible to achieve high performance on the downstream tasks by simply planning within the model, e.g. performing MPC in a zero-shot setting.
This assumes that the model correctly learned the dynamics of the environment and explored rewarding transitions that are relevant to the downstream task.
In Figure 6, we compare the results of performing MPC in a zero-shot setting with the performance of an MPC agent that is allowed the typical 100k steps for fine-tuning.
As for the
MPC method, we employ Model Predictive Path Integral control (MPPI)
(Williams et al., 2015).
Because MPC is particularly expensive to test, we just perform this experiment on top of the models trained with the Plan2Explore URL approach.
We also plot the performance of a non-pre-trained model and of using an actor-critic planning strategy (also provided with the reward predictor since the beginning of fine-tuning), for comparison.
We observe that the performance of MPC (zero-shot) is generally weak.
While it overall performs better than the nonpre-trained model, simply applying MPC leveraging only the pre-trained modules and the reward predictor trained on the PT 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 a major benefit from the fine-tuning stage.
Still, the performance of MPC generally lacks behind the actor-critic performance, suggesting that, especially in a higher-dimensional action space such as the Quadruped one, amortizing the cost of planning with actor-critic seems crucial to achieve higher performance.
Learning the Reward Predictor.
Given that the agent strongly benefits from the 100k fine-tuning steps, we are interested in quantifying how much of this improvement is related to the necessity of learning a good reward function for the downstream task.
In Figure 7, we measure the gap in performance between pre-trained agents that have no knowledge of the reward function at the beginning of finetuning and agents whose reward predictor is initialized from a reward predictor learned on top of the unsupervised pretraining data.
Interestingly, the performance gap is overall small and irrelevant in the Quadruped and Walker domains.
In the Jaco tasks, which have sparser reward functions, an a priori knowledge of the downstream task at the beginning of FT strongly improves performance.
According to our results, it is important that during the 100k steps of fine-tuning the actor is able to quickly obtain information about the downstream task.
This might be easier for dense reward tasks, such as the Walker and Quadruped ones, but trickier in sparse settings like Jaco.
As the initialization of the actor-critic modules also showed to be a compelling issue in our ablation study on the pre-trained modules exploitation (Figure 4), finding more efficient ways to pre-train/fine-tune the actor-critic modules could be an impactful research direction that will facilitate the adoption of unsupervised pre-training for RL.
Latent Dynamics Discrepancy.
A useful measure to assess the uncertainty or inaccuracy of a given model’s dynamics is the model’s misspecification, generally measured as the difference between the dynamics predictions and the real environment dynamics.
When this metric is available, it is also possible to build robust RL strategies, that take the dynamics uncertainty into account while searching for the optimal behavior (Talvitie, 2018).
Dealing with pixel-based inputs, we observe the dynamics of the environment through high-dimensional images, which hinders the possibility to evaluate such a metric, as distances in pixel space can be misleading.
In our approach, we use a model-based RL agent that learns a model of the environment in a compact latent space Z .
In order to quantify the “misspecification" of the learned latent dynamics, we propose a new metric, which we call
Latent Dynamics Discrepancy , that suits the setup of URLB.
We aim to quantify the distance between the predictions of the pre-trained model and the same model after finetuning on a 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 ran fine-tuning experiments where the agent’s decoder weights are frozen, so that the decoder cannot be updated and the model can only improve the posterior and the dynamics.
This ensures that the mapping Z −→ 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 Jensen-Shannon divergence:
[Equation]
where the expectation is taken over the previous model states z t sampled from the fine-tuned posterior q FT ( z t ) , actions a t − 1 sampled from an oracle actor π ∗ ( a t | z t ) , so that we evaluate the metric on optimal trajectories, whose environment’s state distribution corresponds to the stationary distribution induced by the actor s t ∼ d π ∗ ( s t ) .
We usedtrajectories per task in our evaluation.
In Figure 8, we plot the correlation between our metric and the performance ratio between a zero-shot model and a fine-tuned model, where Plan2Explore was used for the 2M steps pre-training phase.
We observed a strong negative
Pearson correlation ( − 0 . 62 ), with a p-value of 0 .
03 < 0 .
05 , asserting that we must reject the null hypothesis, i.e. there exists a correlation between the two factors.
This means that major updates in the model dynamics during fine-tuning 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 link, by either improving the model’s 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.
Unsupervised Rewards and Performance.
How well can the unsupervised RL approaches we employed help improving adaptation further?
To answer this question, we analyze the correlation between the normalized performance of the different agents and the intrinsic rewards they provide [Table] for optimal trajectories obtained by an oracle agent.
A strong negative correlation between the two factors would indicate that the agent will be more interested in seeing the optimal trajectories when its performance is low on the task.
We summarize the results of our analysis in Table 1.
We observe that there is negative correlation between
Plan2Explore (P2E), ICM, LBS’s performance and their intrinsic rewards, while we found ∼ 0 correlation for RND.
In particular, the correlation for LBS, which overall performed best in the benchmark, has a statistical significance, as its p-value is < 0 .
Given such correlation, we believe the intrinsic rewards of LBS might be one of the causes of its outstanding performance.
As LBS searches for transitions of the environment that are difficult to predict for its dynamics, the model likely learns those transitions more accurately, facilitating planning during the fine-tuning stage and eventually leading to higher performance (particularly in the most difficult domain, Quadruped).
It is important that future work would consider learning a less ambiguous dynamics during the unsupervised RL phase, which can be efficiently leveraged by the agent for fine-tuning.
5. Related Work
Our work lies at the intersection of unsupervised RL, modelbased RL, and representation learning for RL.
We discuss below the relevant literature in these three fields.
5.1. Model Based RL
In continuous control, model-based RL combined with powerful search methods has led to impressive results on a wide variety of tasks (Hafner et al., 2019a).
In this work, we used the MPC approach MPPI, which is based on the CrossEntropy Method (CEM Rubinstein & Kroese (2004)).
These methods perform trajectory optimization by fitting a multivariate Gaussian distribution to the imagined future actions allowing them to search the space efficiently.
Alternative search methods such as Monte Carlo Tree search (Coulom, 2006) and Sequential Monte Carlo planning (Piché et al., 2018) could also have been used.
Given that we do not have the reward information during pre-training, we base our model on the Dreamer architecture which reconstructs the future frames to learn a transition model.
This has the advantage of being simple and not requiring the task specification a priori, whereas task and reward awareness is otherwise necessary to learn a model to be self-consistent in latent space (Schrittwieser et al., 2020; Grimm et al., 2020).
5.2. Unsupervised RL
Research in Unsupervised RL spans many fields, from computational accounts of useful intrinsic motivations (Barto, 2004) to empirical evidence for certain intrinsic costs in humans (Kool et al., 2013).
Such intrinsic behavior learning could aid an RL agent to adapt across tasks posed by the environment in a sample-efficient manner.
Oudeyer & Kaplan (2008) classified intrinsic motivation algorithms into three different kinds - knowledge-based, competence-based and data-based models.
Knowledge-based models Schmidhuber (2010) rely on a prediction-error signal to build pseudorewards (Pathak et al., 2017; 2019; Burda et al., 2019a;
Mazzaglia et al., 2021; Burda et al., 2019b; Rajeswar et al., 2021).
Competence models aim at learning a set of diverse and repeatable policies through information-theoretic objectives (Mohamed & Rezende, 2015).
This is achieved by maximizing the mutual information between the trajectory or states and latent skill variables (Eysenbach et al., 2019; Gregor et al., 2016; Liu & Abbeel, 2021a; Frank et al., 2014).
Data-based methods try to increase the diversity of the dataset, often times through explicit maximum entropy objectives or count-based objectives.
Bellemare et al.
Ostrovski et al.
; Liu & Abbeel (2021b); Yarats et al.
An unsupervised RL approach that is closely related to our method is Plan2Explore
(Sekar et al., 2020), which combines Disagreement (Pathak et al., 2019), a knowledgebased exploration approach, with the first iteration of the Dreamer agent (Hafner et al., 2019a), showing improved performance in a few-shot adaptation setting.
We improve upon their work by considering several unsupervised RL strategies (and finding LBS to be the best performing one), better exploiting the pre-trained components of the agent for fast adaptation, and employing a hybrid planning strategy to improve data-efficiency.
5.3. Representation Learning
When inputs are high-dimensional images, it is beneficial to learn compact state representations of the inputs.
Much progress in unsupervised representation learning for RL has been influenced by developments in vision-based unsupervised learning (Chen et al., 2020; Kingma & Welling, 2013).
More recently, a number of works have investigated representation learning for RL (Srinivas et al., 2020;
Laskin et al., 2020; Yarats et al., 2021b).
In our work, we focus on representation learning with autoencoders (Hafner et al., 2019b; Yarats et al., 2019).
Specifically, we ground upon DreamerV2 (Hafner et al., 2021), as predicting ahead in learned latent space allows to efficiently predict thousands of compact state sequences in parallel. Learning a sensible representation is also crucial when aiming to generalize to different domains and/or tasks (van Driessel & FrancoisLavet, 2021; Sasso et al., 2022).
6. Conclusion
Unsupervised RL has been a promising direction aimed at training generalist fast adapting agents.
However, the existing approaches have only yielded sub-optimal results when evaluated on a standardized setup, especially on highdimensional pixel inputs.
We build on existing unsupervised
RL algorithms with a model-based formulation to improve the sample-efficiency and performance on pixel-based inputs.
We empirically show that algorithms that learn a world model through self-supervised exploration can significantly improve their performance, compared to model-free approaches.
Specifically, with our approach we were able to attain near-optimal performance on URLB, a challenging benchmark to evaluate unsupervised RL algorithms.
Furthermore, we propose useful evaluation strategies to assess the quality of the learned model, quantify the uncertainty of the latent dynamics, and assess the intrinsic rewards provided by different unsupervised RL strategies, granting additional understanding of our approach.
Given that our model-based framework relies on the predictions of a learned reward function and, in turn, from the value predictions of the critic, we aim to find faster ways to adapt those functions, given the pre-trained experience.
One idea would be to leverage the flexibility of successor representations and features (Barreto, 2018), which allow learning a task-agnostic estimate of the features expected under the actions of a certain actor.
These predictions could then be exploited for transfer and adaptation.
If we could learn useful successor features during pre-training, we might use them to solve downstream tasks faster during fine-tuning.
We aimed at improving the action selection process during fine-tuning, by employing a hybrid planner that adopts both a classical MPC-based planner and an actor-critic architecture.
However, our strategy overlooks the uncertainty of the model.
In order to account for this, Bayes-adaptive RL strategies could be attempted (Mehta et al., 2021).
Different approaches to learning could also be studied for this particular training setting, where the agent experiences training through two different stages: pre-training and finetuning.
Meta-learning, or learning to learn, strategies for RL could help designing a new approach to tackle the problem end-to-end (Finn et al., 2017; Gupta et al., 2018). Acknowledgements P.M., T
V., and B.D. received funding from the Flemish
Government under the “Onderzoeksprogramma Artificiële
Intelligentie (AI) Vlaanderen” programme.
A. Reference scores
[Table]
Domain [Table] Run [Table]
Walk
[Table] [Table] [Table] Jaco
Reach bottom left
[Table]
[Table]
Reach bottom right
[Table]
[Table]
[Table] left [Table] [Table]
right [Table] 227 ± In Table 2, we report the mean scores for DrQ-v2 (URLB Expert normalization scores in the
, and
Dreamer, which we use to normalize returns in our work, where both supervised baselines have been trained individually on each of the 12 tasks from URLB for 2M steps.
We additionally report mean and standard deviations for the best performing unsupervised baseline from URLB.
which is Disagreement (Pathak et al., 2019), and our best performing method, which employs LBS and the hybrid planner (with no actor initialization for Jaco).
We notice that the LBS + Plan (Ours) scores approach the Dreamer@2M’s scores in several tasks, eventually outperforming them in a few tasks (e.g. Walker Flip, Quadruped Jump).
We believe this merit of LBS + Plan is due both to the exploration pre-training, which may have found more rewarding trajectories than greedy supervised RL optimization, and of the improved planning strategy (Hansen et al., 2022).
B. Unsupervised Reinforcement Learning Strategies
We summarize the unsupervised RL approaches adopted in our work.
For all approaches, rewards have been normalized during training using an exponential moving average with momentum 0 .
95 , with the exception of RND that follows its original implementation reward normalization (Burda et al., 2019b).
ICM.
The Intrinsic Curiosity Module (ICM; Pathak et al.
(2017)) defines intrinsic rewards as the error between states projected in a feature space and a feature dynamics model’s predictions.
We use the Dreamer agent encoder e t = f ϕ ( s t ) to obtain features and train a forward dynamics model g ( e t | e t − 1 , a t − 1 ) to compute rewards as:
[Equation]
As ICM requires environment states to compute rewards, we train a reward predictor to allow estimating rewards in imagination.
Plan2Explore.
The Plan2Explore algorithm (Sekar et al., 2020) is an adaptation of the Disagreement algorithm (Pathak et al., 2019) for latent dynamics models.
An ensemble of forward dynamics models is trained to predict the features embedding e t = f ϕ
( s t ) , given the previous latent state and actions, i.e. g ( e t | z t − 1 , a t − 1 , w k ) , where w k are the parameters of the k-th predictor.
Intrinsic rewards are defined as the variance of the ensemble predictions:
[Equation]
( { g (
[Equation]
| z t
[Equation]
1 , a t
[Equation] 1 , w [Equation] ) |
[Equation]
[1 , ..., K ] } ) .
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 to predict the output of a randomly initialized network n ( s t ) 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.
Intrinsic reward here is defined as:
[Equation]
As RND requires environment states to compute rewards, we train a reward predictor to allow estimating rewards in imagination.
LBS.
In Latent Bayesian Surprise (LBS; Mazzaglia et al. (2021)), they use the KL divergence between the posterior and the prior of a latent dynamics model as a proxy for the information gained with respect to the latent state variable, by observing new states.
Rewards are computed as:
[Equation]
As LBS requires environment states to compute rewards, we train a reward predictor to allow estimating rewards in imagination.
APT.
Active Pre-training (APT; Liu & Abbeel (2021b)) uses a particle-based estimator based on the K nearest-neighbors algorithm (Singh et al., 2003) to estimate entropy for a given state.
We implement APT on top of the deterministic component of the latent states ¯ z t , providing rewards as:
[Equation]
where k are the nearest-neighbors states in latent space.
As APT requires only latent states, it can be computed directly in imagination.
We used k = 12 nearest neighbors.
DIAYN.
Diversity is All you need (DIAYN; Eysenbach et al.
(2019)) maximizes the mutual information between the states and latent skills w .
We implement DIAYN on top of the latent space of the model, writing the mutual information as I ( w t , z t ) = H ( w t ) − H ( w t | z t ) . The entropy H ( w t ) is kept maximal by sampling w t ∼ p ( w t ) from a discrete uniform prior distribution, while H ( w t | z t ) is estimated learning a discriminator q ( w t | z t ) .
Additionally, DIAYN maximizes the entropy of the actor, so we compute intrinsic rewards as:
[Equation]
As DIAYN requires environment states and sampled skills to compute rewards, we train a reward predictor to allow estimating rewards in imagination.
C. Algorithm
Algorithm 1 Model-based Unsupervised RL
Require: Actor θ , Critic ψ , World Model ϕ
Require: Intrinsic reward r int , extrinsic reward r ext
Require: Environment, M , downstream tasks T k , k ∈ [1 , . . .
, M ]
Require: Pre-train steps N PT , fine-tune steps N FT , environment steps/update τ
Require:
Initial model state z 0 , hybrid planner Plan, replay buffers D PT , D FT 1: for t = 0 , . .
. , N PT do 2:
Draw action from the actor, a t ∼ π θ ( a t | z t ) 3:
Apply action to the environment, s t +1 ∼ P ( ·| s t , a t ) 4:
Add transition to replay buffer, D PT ← D PT ∪ ( s t , a t , s t +1 ) 5:
Infer model state, z t +1 ∼ q ( z t +1 | z t , a t , f ϕ ( s t +1 ))
6: if t mod τ = 0 then 7:
Update world model parameters ϕ on the data from the replay buffer D PT 8:
Update actor-critic parameters { θ, ψ } in imagination, maximizing r int 9: end if 10: end for 11: Output pre-trained parameters { ψ PT , θ PT , ϕ PT } 12: for T k ∈ [ T 1 , . .
, T M ] do 13:
Initialize fine-tuning world-model with ϕ PT 14:
(Optional) Initialize fine-tuning actor with θ PT 15: for t = 0 , . .
. , N FT do 16:
Use the planner for selecting action, a t ∼ Plan ( z t ) 17:
Apply action to the environment, s t +1 , r ext t ∼ P ( ·| s t , a t ) 18:
Add transition to replay buffer, D FT ← D FT ∪ ( s t , a t , r ext t , s t +1 ) 19:
Infer model state, z t +1 ∼ q ( z t +1 | z t , a t , f ϕ ( s t +1 ))
20: if t mod τ = 0 then 21:
Update world model parameters ϕ on the data from the replay buffer D FT 22:
Update actor-critic parameters { θ, ψ } in imagination, maximizing r ext 23: end if 24: end for 25:
Evaluate performance on T k 26: end for
D. Hyperparameters
Most of the hyperparameters we used for world-model training are the same as in the original DreamerV2 work (Hafner et al., 2021).
Specific details are as outline here:
[Table]
[Table]
critic
[Table]
[Table]
Imagination [Table] scale
[Table]
TD-MPC
Iterations
Number of samples
Number of elite actions
Mixture coefficient (Actor/CEM) [Table]
Min [Table] fixed)
[Table]
Temperature
[Table]
Momentum
[Table]
Horizon
[Table]
Common
Environment steps/update
MLP number [Table] layers
MLP number of units
[Table]
layers dimension
[Table]
Weight [Table]
Gradient clipping
For the pure MPC-based experiments, we increased the number of MPPI samples from 512 to 1000, the number of elite actions from 64 to 100, and the horizon from 5 to 15.
E. Additional Results
We present complete results, for each unsupervised RL method, for the experiments in Section 4, when using only the actor-critic algorithm, in Figure 9, and when also employing the hybrid planner, in Figure 10.
[Figure] [Figure]