Learning General World Models in a Handful of Reward-Free Deployments
Abstract
Building generally capable agents is a grand challenge for deep reinforcement learning (RL).
To approach this challenge practically, we outline two key desiderata: 1) to facilitate generalization, exploration should be task agnostic; 2) to facilitate scalability, exploration policies should collect large quantities of data without costly centralized retraining.
Combining these two properties, we introduce the reward-free deployment efﬁciency setting, a new paradigm for RL research.
We then present CASCADE , a novel approach for self-supervised exploration in this new setting.
CASCADE seeks to learn a world model by collecting data with a population of agents, using an information theoretic objective inspired by Bayesian
Active Learning.
CASCADE achieves this by speciﬁcally maximizing the diversity of trajectories sampled by the population through a novel cascading objective.
We provide theoretical intuition for CASCADE which we show in a tabular setting improves upon na¨ıve approaches that do not account for population diversity.
We then demonstrate that CASCADE collects diverse task-agnostic datasets and learns agents that generalize zero-shot to novel, unseen downstream tasks on Atari,
MiniGrid, Crafter and the DM Control Suite.
Code and videos are available at https://ycxuyingchen.github.io/cascade/
Introduction
Reinforcement learning (RL, [105]) has achieved a number of impressive feats over the past decade, with successes in games [69, 13, 100], robotics [48, 77] and the emergence of real world applications [11, 25].
Indeed, now that RL has successfully mastered a host of individual tasks, the community has begun to focus on the grand challenge of building generally capable agents [90, 109, 68, 5].
In this work, we take steps towards building generalist agents at scale , where we outline two key desiderata.
First, for agents to become generalists that can adapt to novel tasks, we eschew the notion of restricting agent learning to task-speciﬁc reward functions and focus on the reward-free problem setting instead [83, 28], whereby agents must discover novel skills and behaviors without supervision.
2 Consider the problem of learning to control robotic arms, where we may already have some expert ofﬂine data to learn from.
In many cases this data will cover only a subset of the entire range of possible behaviors.
Therefore, to learn additional general skills, it is imperative to collect additional novel and diverse data, and to do so without a pre-speciﬁed reward function.
Second, to ensure scalability, we should have access to a large ﬂeet of robots that we can deploy to gather this data for a large number of timesteps [60], without costly and lengthy centralized
Indeed, designing reward functions to learn behaviors can have unintended consequences [74, 6].
[Figure]
retraining during this crucial phase [67].
This has recently been referred to as deployment efﬁciency , falling between the typical online/ofﬂine RL dichotomy.
Limiting deployments not only reduces the overhead in retraining exploration policies, but also limits the potential costs and risks present when deploying new policies [108, 53], an important consideration in many real world settings, such as robotics [48], education [65] and healthcare [29].
Combining these two desiderata, we introduce the reward-free deployment efﬁciency setting, a new paradigm for deep RL research.
To tackle this new problem we train a population of exploration policies to collect large quantities of useful data via world models [20, 34, 39].
World models allow agents to plan and/or train policies without interacting with the true environment.
They have already been shown to be highly effective for deployment efﬁciency [67], ofﬂine RL [118, 49, 8, 88] and self-supervised exploration [99, 97].
Furthermore, world models offer the potential for increasing agent generalization capabilities [7, 37, 68, 16, 8, 21, 10, 58], one of the frontiers of RL research [80, 51].
However, since existing self-supervised methods for learning world models are designed to collect only a few transitions with a single exploration policy, they likely produce a homogenous dataset when deployed at scale, which does not optimally improve the model.
Instead, drawing analogies from Bayesian Active Learning [42, 52], we introduce a new information theoretic objective that maximizes the information gain from an entire dataset collected by a population of exploration agents (see
Figure 2).
We call our method C oordinated A ctive S ample
C ollection vi a D iverse E xplorers or CASCADE (Figure 1).
We provide theoretical justiﬁcation for CASCADE , which emphasizes the importance of collecting data with diverse agents.
In addition, we provide a rigorous empirical evaluation across four challenging domains that shows CASCADE can discover a rich dataset from a handful of deployments.
We see that CASCADE produces general exploration strategies that are equally adept at both “deep” exploration problems and diverse behavior discovery.
This makes it possible to train agents capable of zero-shot transfer when rewards are provided at test time in a variety of different settings.
To summarize, our contributions are as follows: 1) We introduce a novel problem setting, RewardFree Deployment Efﬁciency , designed to train generalist agents in a scalable fashion; 2) We propose
CASCADE , a theoretically motivated model-based RL agent designed to gather diverse, highly informative data, inspired by Bayesian Active Learning; 3) We provide analysis that shows CASCADE theoretically improves sample efﬁciency over other na¨ıve methods that do not ensure sample diversity, and demonstrate that CASCADE is capable of improved zero-shot transfer in four distinct settings, ranging from procedurally generated worlds to continuous control from pixels.
Problem Statement
Reinforcement learning (RL) considers training an agent to solve a Markov Decision Process (MDP), represented as a tuple M = {S , A , P, R, ρ, γ } , where s ∈ S and a ∈ A are the set of states and actions respectively, P ( s (cid:48) | s, a ) is a probability distribution over next states given a previous state and action, R ( s, a, s (cid:48) )
→ r is a reward function mapping a transition to a scalar reward, ρ is an initial state distribution and γ is a discount factor.
A policy π acting in the environment produces a trajectory τ = { s 1 , a 1 , . . .
, s H , a H } for an episode with horizon H .
Since actions in the trajectory are sampled from a policy, we can then deﬁne the RL problem as ﬁnding a policy π that maximizes expected returns in the environment, i.e. π (cid:63) =
arg max π E τ ∼ π
[ R ( τ )] .
We seek to learn policies that can transfer to any MDP within a family of MDPs.
This can be formalized as a Contextual MDP [51], where observations, dynamics and rewards can vary given a context.
In this paper we consider settings where only the reward varies, thus, if the test-time context is unknown at training time we must collect data that sufﬁciently covers the space of possible reward functions .
Finally, to facilitate scalability, we operate in the deployment efﬁcient paradigm [67], whereby policy learning and exploration are completely separate, and during a given deployment , we gather a large quantity of data without further policy retraining (c.f. online approaches like DER [112], which take multiple gradient steps per exploration timestep in the real environment).
Taken together, we consider the reward-free deployment efﬁciency problem.
This differs from previous work as follows: 1) unlike previous deployment efﬁciency work, our exploration is task agnostic; 2) unlike previous reward-free RL work, we cannot update our exploration policy π EXP during deployment.
Thus, the focus of our work is on how to train π EXP ofﬂine such that it gathers heterogeneous and informative data which facilitate zero-shot transfer to unknown tasks.
In this paper we make use of model-based RL (MBRL), where the goal is to learn a model of the environment (or world model [96]) and then use it to subsequently train policies to solve downstream tasks.
To do this, the world model needs to approximate both P and R .
Typically, the model will be a neural network, parameterized by ψ , hence we denote the approximate dynamics and reward functions as P ψ and R ψ , which produces a new “imaginary” MDP, M ψ = ( S , A , P ψ , R ψ , ρ ) .
We focus on Dyna-style MBRL [104], whereby we train a policy ( π θ parameterized by θ ) with model-free
RL solely using “imagined” transitions inside M ψ .
Furthermore, we can train the policy on a single
GPU with parallelized rollouts since the simulator is a neural network [54].
The general form of all methods in this paper is shown in Algorithm 1, with the key difference being step 5: We aim to update π EXP in the new imaginary MDP M ψ such that it continues to collect a large, diverse quantity of reward-free data.
Note that π EXP need not be a single policy, but could also refer to a collection of policies that we can deploy (either in parallel or in series), such that π ∈ π EXP .
We focus on learning world models from high dimensional sensory inputs such as pixels [34, 76, 47], where at each timestep we are given access to an observation o t rather than a state s t .
A series of recent works have shown tremendous success by mapping the observation to a compact latent state z t [39, 38, 40].
In this paper we will make use of the model from DreamerV2 [40], which has been shown to produce highly effective policies in a variety of high dimensional environments.
The primary component of DreamerV2 is a Recurrent State Space Model (RSSM) that uses a learned latent state to predict the image reconstruction, reward r t and discount factor γ t .
Aside from the reward head, all components of the model are trained jointly, in similar fashion to variational encoders (VAEs, [50, 92]).
For zero-shot evaluation, we follow [97] and only train the reward head at test time when provided with labels for our pre-collected data, which is then used to train a behavior policy ofﬁne.
Thus, it is critical that our dataset is sufﬁciently diverse to enable learning novel, unseen behaviors.
Coordinated Active Sample Collection
The aim of this work is to train a population of B exploration policies { π ( i )EXP } Bi =1 such that they collectively acquire data which maximally improves the accuracy of a world model .
To achieve this, we take inspiration from the information theoretic approach in Plan2Explore [97], but crucially focus on maximizing information gain over entire trajectories rather than per state-action, and hence drop the conditional dependence on state and action (see App.
C.1 for why this distinction is important):
[Equation]
where d π M ψ is the distribution of states visited by the policy π in the imaginary MDP M ψ .
This objective produces π EXP , a policy whose visitation distribution has a high entropy when computed over model samples, but has low entropy for each individual MDP model (i.e., high epistemic/reducable uncertainty).
We think of each model M ψ as sampled from a posterior distribution over models given the data.
A good exploration policy has low entropy on individual models but large entropy across models, i.e. it is intent in visiting regions of the space where there is large uncertainty about the model transitions.
To make this objective more general, we represent the trajectory data collected by a policy with a “summary” embedding space [79, 72].
Let Φ : Γ → Ω be a summary function mapping trajectories into this embedding space.
P Φ π [ M ψ ] denotes the embedding distribution generated by policy π in imaginary MDP M ψ .
We can now write the objective from Eq.
1 as follows:
[Equation]
This more general framework allows us to consider multiple representations for trajectories, in a similar fashion to behavioral characterizations in Quality Diversity algorithms [85].
For the rest of this discussion, we will use the ﬁnal state embedding as our summary representation, whereby Φ( τ ) =
h H , since in the case of the RSSM, the ﬁnal latent is a compact representation of the entire trajectory collected by the policy, analogous to the ﬁnal hidden state in an RNN [94, 103].
3.1 A Cascading Objective with Diverse Explorers
We now consider a population-based version of
Equation.
2, using
B agents:
[Equation]
where π B = π (1) , · · · , π ( B ) and (cid:81) Bi
=1 P Φ π ( i )
[ M ψ ] is the product measure of the policies’ embedding distributions in M ψ .
By deﬁnition, the conditional entropy factorizes as:
[Equation]
It is now possible to show that maximum information gain is achieved with a diverse set of agents:
Lemma 1.
When all models M ψ in the support of the model posterior are deterministic and tabular, and the space of policies Π consists only of deterministic policies, there always exists a solution { π
[Equation]
EXP } Bi =1 satisfying π ( i )EXP (cid:54) =
π ( j ) EXP ∀ i (cid:54) =
j .
Moreover, there exists a family of tabular MDP models, such that the maximum cannot be achieved by setting π EXP ( i )
=
π for a ﬁxed π .
The proof of Lemma 1 is in Appendix C.1.1.
Since the mutual information objective 3 is submodular, a greedy algorithm yields a (1 − 1 e ) approximation of the optimum (where e is Euler’s number)
Leveraging this insight, let us assume that we already have a set of policies π (1) , · · · , π ( i − 1) ; we then select the next policy π ( i ) based on the following greedy objective:
[Equation]
  Which can be factorized in similar fashion to Equation 4 (See Appendix C).
3.2 A Tractable Objective for Deep RL
Inspired by [97], we make a couple of approximations to derive a tractable objective for { π EXP } Bi
= in the deep RL setting.
First, we assume that the ﬁnal state embedding distributions are Gaussian with means that depend on the policies and sampled worlds, and variances that depend on the worlds, i.e. P Φ π [ M ψ = w ] = N ( µ ( w, π ) , Σ( w )) .
In this case, H ( P Φ π [ M ψ ] |M ψ = w ) =
ρ ( w ) , and
Eq.
5 reduces to solving π ( i ) =
arg max ˜ π
( i ) ∈
Π H (cid:16)(cid:81) ij =1 P Φ˜ π ( j ) [ M ψ ] (cid:12)(cid:12)(cid:12) ˜
π ( j ) =
π ( j ) ∀ j ≤ i − 1 (cid:17) for a policy that maximizes the resulting joint entropy of the embedding distribution when added to the policy population.
This produces the following surrogate objective, maximizing a quadratic cascading disagreement:
[Equation]
where D ( i − 1) is a dataset of imagined trajectories sampled from policies π (1) , · · · , π ( i − 1) in the model, and PopDiv is short for Population Diversity.
Finally, following [97, 9], we also add a per state information gain component to each policy’s reward to encourage a richer landscape for data acquisition: InfoGain( π ) =
E τ ∼ P π [ M ψ ] (cid:104)(cid:80) ( s,a ) ∈
τ σ ( s, a ) (cid:105) where σ ( · , · ) is the variance across the ensemble latent state predictions (for details see App. B).
Taken together, these objectives form our approach, which we call C oordinated A ctive S ample
C ollection vi a D iverse E xplorers or CASCADE .
CASCADE trains agents to optimize: 1) a diversity term ( PopDiv ) that takes into account the behaviors of the other agents in the population; 2) an information gain term ( InfoGain ) that encourages an individual agent to sample states that maximally improve the model:
[Equation]
where λ is a weighting hyperparameter that trades off whether we favor individual model information gain or population diversity.
Finally, we train the B agents in parallel using (policy) gradient descent over θ , which makes it possible to achieve the same wall clock time as training a single agent [23].
3.3 Theoretical Motivation
We now seek to provide a tabular analogue to CASCADE which provides a theoretical grounding for our approach.
In App.
D we outline the pseudo-code of CASCADE - TS , a greedy Thompson Sampling algorithm [4] that produces the i -th exploration policy in a tabular enviornment using “imaginary” data gathered by running policies π (1) , · · · , π ( i − 1) in the model .
We can then show the following:
Lemma 2.
For the class of Binary Tree MDPs, the CASCADE - TS algorithm satisﬁes,
[Equation]
SinglePolicyBatch) where T ( (cid:15), · ) are the expected number of rounds of deploying a population of B policies necessary to learn the true model up to (cid:15) accuracy; SinglePolicyBatch plays a ﬁxed policy B times in each round; Sequential does not have a population, and instead interleaves updates and executions of a single policy B times within each round.
The proof is in App.
D.
Indeed, we see that CASCADE - TS achieves provable efﬁciency gains over a na¨ıve sampling approach that does not ensure diversity in its deployed agents.
This provable gain is achieved by discouraging the i -th policy away from imaginary state-action pairs sampled by the previous i − 1 policies by using imaginary counts.
Now returning to CASCADE , we can see the importance of leveraging the imaginary data gathered in the model by the previous i − 1 policies when training the i -th exploration policy.
Concretely, encouraging policy π ( i ) to induce high disagreement with the embeddings produced by { π (1) , · · · , π ( i − 1) } (i.e. the PopDiv term) is analogous to the imaginary count bonus term of CASCADE - TS in Line 10 of
Alg. 2, which avoids redundant data collection during deployment.
Experiments
In our experiments we test whether CASCADE facilitates the learning of generalist agents by evaluating their zero-shot transfer to unseen tasks, given a limited number of reward-free deployments.
In all experiments, agents cannot train online in the environment, and instead execute a ﬁxed exploration policy to gather new data during each deployment.
We test against three baselines.
The ﬁrst is random exploration, which is used in the majority of RL publications.
Our next baseline is Plan2Explore [97] ( P 2 E ), which trains a single exploration policy that optimizes for expected information gain inside a world model.
We then include a population-based version of P 2 E , which we call Population
Plan2Explore ( PP 2 E ).
PP 2 E trains a set of randomly initialized agents that independently seek to maximize expected information gain.
All methods make use of a DreamerV2 world model [40] and use the same hyperparameters for model and agent training (more details in App. B).
We test all methods in two separate settings: Firstly Exploring Worlds (Sec 4.1), where we seek to collect data in challenging exploration environments.
We test how well the explorers cover the state space and discover sparse rewards (also known as “deep exploration” [78, 75]).
Secondly, we consider Exploring Behaviors (Sec 4.2) where we consider collecting data in continuous control environments.
In both settings we test zero-shot transfer as follows: 1) we provide reward labels to the model; 2) we train a separate reward head; 3) we train a task speciﬁc behavior policy and test it in the environment.
This tests whether our model is general enough to facilitate learning downstream policies for previously unknown tasks [97, 55, 116].
4.1 Exploring Worlds
In many cases we may not know a priori which states in an environment are highly rewarding.
Thus, if we wish to use our world model to subsequently train agents to solve tasks, it is crucial that we accurately model as much of state space as possible.
Here we consider exploring worlds in three different environments: MiniGrid [19], Atari [12] and Crafter [37].
In all three settings, we start with a random deployment before beginning reward-free exploration.
We begin with MiniGrid, a set of sparse-reward, partially observable navigation environments commonly used to test state-of-the-art exploration methods [89, 17, 31, 121].
MiniGrid environments are procedurally generated, which provides an additional generalization challenge [46, 93].
We consider the canonical FourRooms and the more challenging MultiRoom .
In each case we are limited to 5 deployments of 200k transitions, for a total of 1M environment interactions.
In order to gauge the effectiveness of the reward-free exploration, we use two metrics: 1) state coverage, where we evaluate the percentage of the states visited by the exploration policies after each deployment on the same set of held out test levels; and 2) rewarding episodes, where we show the number of rewarding episodes collected during the training process.
Note that these two objectives are not useful in isolation, as solved levels may only include a narrow set of observations, while state coverage alone could be maximized without performing deep, goal-seeking exploration in the environment.
However, in combination they provide a proxy for exploration.
We show the results in Figure 3, where we see the CASCADE exploration policies cover the state space far more effectively than all baselines.
In FourRooms , both CASCADE and PP 2 E ﬁnd the goal an equal number of times, while CASCADE ﬁnds almost double the number of rewarding episodes in the MultiRoom environment.
We also test zero-shot performance, where CASCADE achieves almost 100% success rate.
Interestingly, in both settings the Random policy does cover a large proportion of the state space, but ﬁnds fewer rewarding episodes and subsequently does not have a sufﬁcient quantity to train a policy to solve the task (and accordingly achieves 0% success rate).
Next we consider four Atari games: Montezuma’s Revenge , Frostbite , Hero and Freeway , often used as barometers for exploration capabilities [27].
We use 15 reward-free deployments, each consisting of 200k steps, and use the ﬁnal model to train a behavior policy with task reward.
We test these agents in a zero-shot fashion, with results shown in the ﬁrst three columns of Figure 4.
We see that CASCADE gets statistically signiﬁcantly stronger zero-shot returns compared to the baselines.
Furthermore, the baseline methods display no clear trend, with random exploration sometimes offering the strongest performance.
On the other hand, we see that CASCADE performs consistently well, providing further evidence of it being a strong general exploration strategy.
See App.
A.2 for more detailed analysis.
[Figure]
Finally, we consider the Crafter environment [37], a procedurally generated grid world based on the game of Minecraft.
This more complex setting requires the agent to master a variety of compositional skills to fully explore all possible behaviors.
Online P 2 E represents the state of the art reward-free agent, achieving a “Crafter Score” of 2.1.
Making the problem more challenging, we deploy onlydistinct exploration policies, each collecting 50k timesteps, which leads to weaker performance for
P 2 E vs. the reported result in [40].
However, as we see in Figure 5, adding a population of diverse explorers recovers the majority of this performance.
We ﬁnd this result encouraging as the behaviors required to achieve a higher Crafter score are non-trivial.
Moreover, as shown in Figure 6, CASCADE achieves the highest average success rate and unlocks the most number of unique Crafter skills among all baselines, including Making a Stone Pickaxe.
We believe this opens the door to further gains with more active tuning of the behavioral representation used in CASCADE .
Random
P2E
PP2E
CASCADE
IQM 0.
Random
P2E
PP2E
P(CASCADE > Y) 0.
[Figure]
Training Score
Random
P2E
PP2E
[Figure] [Figure]
4.2 Exploring Behaviors
Next we wish to test is whether our agents sufﬁciently explore a range of behaviors in a continuous control setting, using the walker environment from the DM Control Suite [111] (DMC).
In these experiments, we operate from pixels and assume access to an initial dataset of 200k transitions.
Then, we conduct 14 further reward-free deployments, each collecting 200k transitions.
We consider the possibility of improving generality from arbitrary ofﬂine data, using common datasets from the ofﬂine RL literature: random , medium and expert .
Using random data represents learning from scratch.
In Figure 7 we show the behaviors deployed by CASCADE and PP 2 E in a single deployment. The CASCADE explorers all display diverse, useful behaviors, while PP 2 E agents are homogenous.
We then test the generality of the world model after each deployment by training a separate agent from scratch for each of the four individual tasks: stand , walk , run , flip proposed in URLB [56].
This demonstrates whether the world model is capable of producing an imaginary MDP that facilitates learning speciﬁc behaviors.
To compare the performance of CASCADE and the baselines, in Figurewe show the aggregated statistics from the RLiable library [3] for a given number of deployments.
To do this, we combine the results across all datasets ( random , medium , expert ) and all downstream tasks ( stand , walk , run , flip ) and compute the following: 1) the Inter-Quartile Mean (IQM) of the normalized scores (the robust statistic recommended in [3]); 2) the Probability of Improvement, the likelihood of one method outperforming another on a new, unseen task.
As we see, CASCADE shows clear and consistent gains over the next best baseline ( PP 2 E ) throughout 15 deployments, showing [Figure] statistical signiﬁcance under the Mann-Whitney U test [66].
Note that here the single agent P2E performs poorly, likely due to only collecting data with a single mode of exploratory behavior. See
Figure 9 in the Appendix for a full breakdown of the performance.
4.3 Discussion and Limitations
We have shown that CASCADE is the most effective method in training generally capable agents across a wide range of environments.
Indeed, CASCADE provides the best of both worlds: deep exploration by seeking the frontier of the current knowledge (via InfoGain), while covering the space of possible behaviors (via PopDiv).
Each of our baselines lacked in one of these two key properties.
Random exploration fared poorly when exploring worlds, which requires deep exploration.
However, it can be a highly effective method for exploring behaviors.
This is likely why boosting entropy in the action space is so effective for robotics tasks [36], but less so in deep exploration tasks [75].
Furthermore, the information theoretic baselines were effective in Exploring Worlds, but the disappointing performance for P 2 E in DMC is likely caused by a lack of data diversity, while the
PP 2 E populations were not sufﬁciently heterogeneous to match the coverage of CASCADE .
In terms of limitations, the main weakness in current methods is the absolute level of performance.
Despite outperforming baselines in DMC, the policies learned by CASCADE are far from optimal.
This is likely due to the challenge of the setting (reward-free deployment efﬁciency) and we believe that future work in the community will be crucial to further improve these methods.
Finally, CASCADE does require training additional agents that may provide additional computational expense.
In this setting we assume the time between deployments is arbitrary, but this may not be the case in all settings and future work could consider ways to make this process more efﬁcient and scalable.
Related Work
At the core of our paper is the notion of training generalist agents, which have recently been of increased interest in RL [90, 68, 109, 110].
In particular, we focus on the reward-free paradigm, since in many cases it may not be possible to know a priori all possible tasks we may wish our agent to solve [45].
With this in mind, there has been a surge of interest in methods that can learn in a self-supervised fashion, to build representations that lead to fast (or even zero-shot) adaptation to future tasks.
These methods range from competence-based (i.e. learning “skills”)
[1, 33, 28, 59, 62, 101] to data-based that aim to maximize data diversity [63, 117].
In this work we take inspiration from knowledge-based approaches, which typically seek to learn a world model [9, 83, 84, 30, 107, 15, 98].
In particular we build on Plan2Explore [97] which trains an exploration agent by maximizing an objective resembling information gain inside the world model.
Plan2Explore was the ﬁrst work that showed it is possible to transfer a model learned without rewards in a zero-shot fashion.
However, it differed from our work in that Plan2Explore solely considers the online RL setting, deploying a new exploration policy every timestep.
Other recent work examined the effectiveness of reward-free RL for collecting datasets for subsequent ofﬂine RL [116, 55].
Once again however, this work trains the exploration agent online.
By contrast, we only deploy a handful of exploration agents and collect a larger quantity of data with each one.
Our setting resembles a reward-free version of the deployment efﬁciency paradigm, ﬁrst introduced in [67].
By contrast, their method focuses on policy optimization in the supervised RL paradigm, whereas we focus on exploration and zero-shot generalization in the unsupervised RL setting.
Follow up work has remained focused on supervised policy optimization [102], whereas we are focused on training generalist agents in the reward-free paradigm.
Recent work has also shown the importance of frequent policy updates in ensuring diverse exploration [95], further supporting the need for a diverse population of explorers when operating in the deployment-efﬁcient setting.
To address the challenges of this novel problem setting, we draw inspiration from a variety of ﬁelds.
In particular, several works have considered exploration with a population of agents [106, 26, 64], but these have only considered task-speciﬁc online RL.
Our choice of a diverse population of agents takes inspiration from the ﬁeld of Quality Diversity (QD, [85, 24]), where it has been shown that diversity can boost exploration in RL [22].
Recently, [61] showed it is possible to train diverse populations inside a world model to acquire new skills.
As far as we are aware, we are the ﬁrst to consider a cascading objective for training diverse populations, which is typically done either synchronously [82] or sequentially [122, 119].
Also related are works that speciﬁcally train policies that are able to maximally cover the state space of a given MDP [2, 70].
However, these works focus on the more theoretically amenable (and less general) low-rank MDP setting, and hence do not demonstrate any experiments on larger scale problems.
Furthermore, we aim to achieve state-space coverage by encouraging policy diversity , which is distinct from the approaches taken in these works.
We make use of an information theoretic objective, which has been shown to be highly effective for exploration in RL [71, 43, 41].
Most similar to our work [120] consider the reward free RL setting, and show that maximizing entropy can lead to improved exploration that can facilitate transfer to arbitrary reward functions.
By contrast we consider the limited deployment setting and use model-based RL.
Finally, our work relates to collective intelligence [35, 115], since each member of the population contributes a small piece to the puzzle that produces a general world model.
Conclusion
We introduced the reward-free deployment efﬁciency paradigm—an important problem setting for learning generally capable agents in a scalable fashion.
To address the challenges of this setting, we proposed CASCADE , a theoretically motivated approach that leverages a diverse population of selfsupervised exploration agents.
We showed in a variety of experiments that CASCADE produces general exploration strategies , adept at both deep exploration and gathering sufﬁciently hetereogeneous datasets that facilitate the learning of diverse downstream behaviors.
We believe this work provides a new approach for collecting large, diverse datasets in a self-supervised fashion, inducing a potentially open-ended process for training general world models at scale.
For future work, we could consider alternative means for computing diversity, for example using optimization landscapes [81] or explicit behavioral diversity [28, 82].
These could also be combined with representations that directly incorporate the behaviors of the policy [14] to further improve the design of the “summary” embedding space.
Furthermore, we would be interested in considering transformer models [113] for trajectory modelling [18, 44] that ought to scale more gracefully with larger batches of data [87].
Recent work has shown that this class of models can compose unseen complex multi-task behaviors at test time from ﬁxed ofﬂine datasets [32], as well as leverage signiﬁcantly large and diverse ofﬂine datasets to improve behavioral modeling [91], highly relevant to both aforementioned desiderata for the reward-free deployment efﬁciency paradigm.
Acknowledgments
The authors would like to thank Cong Lu for help with Ofﬂine DreamerV2 code, and Minqi Jiang and Robert Kirk for useful discussions.
This work was funded by Meta AI.
Appendix A Additional Experimental Results A.1 DMC
Here we show additional experimental results, beginning with a full breakdown of the DMC results for each dataset.
In Figure 9 we see the zero-shot transfer performance for models trained with random , medium and expert initial datasets
[Figure]
and 14 subsequent deployments.
Interestingly, we see that the medium dataset proves to be the most effective for achieving high performance, this may be due to being more diverse than the expert dataset, which is a narrow distribution of high performing episodes
[Figure]
A.2 Atari Tasks
Next we show detailed results in four Atari games: Montezuma’s Revenge, Frostbite, Hero and
Freeway in Fig.
We see that generally CASCADE discovers higher rewarding episodes (almost double the next best baseline in Montezuma’s Revenge)
[Figure]
and also discovers more rewarding episodes on average.
This translates to stronger zero-shot performance than all baselines in all four environments.
On the other hand, the other baselines achieve worse zero-shot performance because they either (1) ﬁnd an occasional high reward episode (Max Episode Return), but not in a sufﬁcient quantity (Rewarding Episodes)
[Figure]
or (2) collect abundant low rewarding episodes that are less helpful for learning good behavior polices.
For instance, the random baseline is able to ﬁnd Rewarding
Episodes more frequently in Frostbite
[Figure]
but these are relatively low quality, hence the lower curve when assessing its Max Episode Return
[Figure]
To understand why CASCADE performs well, in Figure 11 we plot trajectories from Montezuma’s
Revenge.
We see that the inclusion of a diversity-inducing objective does indeed lead to more diverse behaviors in the environment, with each policy exploring different regions of the room.
In contrast, for
PP 2 E , agents 3,4,5,6,7,9 all exhibit very similar behavior, likely indicating that random initialization alone is not sufﬁcient to induce diversity when all agents optimize for the same objective.
CASCADE
[Figure]
B Implementation Details
DreamerV2 consists of an image encoder that uses a Convolutional Neural Network (CNN, [57]), a
Recurrent State-Space Model (RSSM [39]) that learns the dynamics, and predictors for the image, reward, and discount factor.
The RSSM uses a sequence of deterministic recurrent states h t .
At each step, it computes a posterior state z t conditioned on the current image x t , as well as a prior state ˆ z t without the current image.
During world model training, the concatenation of h t and z t is used to reconstruct the current image x t , and predict the reward r t and discount factor γ t .
Once the world model is trained, it can be used to roll out “imaginary” trajectories where the model state is instead the concatenation of the deterministic state h t and prior stochastic state ˆ z t .
RSSM
[Equation]
Recurrent model
[Equation]
predictor
[Equation]
Image predictor
[Equation]
Discount
[Equation]
Our implementation was built on top of the ofﬁcial DreamerV2 repository https://github.com/ danijar/dreamerv2 .
We used the default hyperparameter values in the DreamerV2 repository.
Table 1 lists the additional hyperparameters used in our experiments.
We also incorporate a latent disagreement ensemble, following Plan2Explore [97].
This involves training, alongside the RSSM model, an MLP ensemble with 10 members (each having 4 hidden layers, 400 units per layer, and different parameter initializations).
These one-step models take action a t and latents z t , h t as inputs, and try to predict the next stochastic latent z t +1 .
The variance across these outputs during imagined rollouts is then used as the reward signal that forms the Information
Gain ( InfoGain ) objective.
Note that this ensemble is otherwise unused, and is solely trained for the purposes of the exploration objective
[Table]
[Table]
C Deriving the Objective C.1 An Information Theoretic Perspective
Inspired by [97], we derive a data acquisition objective based on the mutual information between the collected data and parameters of the world model M ψ .
M ψ can be either stochastic or deterministic, and can also be an ensemble of empirical models.
Crucially however, we eschew the reliance on one-step information gain (as performed in [97]), and instead aim to maximize diversity directly in the space of trajectories.
To do this, let Φ : Γ → Ω be an summary function mapping trajectories into an embedding ‘summary’ space.
For any model M ψ we use the notation P π [ M ψ ] to denote the distribution over trajectories generated by policy π in M ψ and use M ψ ∼ P M ψ to denote the posterior distribution over models.
The later depends on the data collected so far.
In this work we consider different embedding functions Φ ,
Final State Embedding .
In this setting we deﬁne Φ( τ ) =
h H where h H corresponds to the H − th (and last) state embedding in the trajectory τ .
Since we use an RNN world model, the use of this summary embedding is analogous to the ﬁnal encoder embedding in seq2seq language models [103].
Visitation Frequencies .
In the tabular setting if we deﬁne Φ( τ ) as a discounted count statistic for states in τ , we recover P Φ π [ M ψ ] =
d π M ψ where d π M ψ corresponds to the discounted visitation distribution of policy π in model M ψ .
The authors of [97] study a per state-action mutual information objective that informs the construction of a greedy mutual information maximizing policy.
Implicitly, this per-step objective assumes the model to factor in independent transition operators pertaining to each state-action pair.
This is certainly not the case when using powerful function approximators such as neural networks.
In this case, the dynamics model corresponding to the state transitions of a speciﬁc state action pair is correlated with the state transition model of other state-actions.
This is not captured accurately by the objective in [97].
Instead, we consider the following choice for an exploratory policy in the single policy setting studied in [97]:
[Equation]
π EXP is a policy whose distribution over summaries Φ( τ ) with τ ∼ P π
[ M ψ ] and M ψ ∼ P M ψ has large entropy, but such that the average entropy of the summaries in every individual model is small.
The term H ( P Φ π [ M ψ ]) captures the total uncertainty (epistemic + aleatoric) while subtracting the conditional entropy removes the aleatoric uncertainty resulting from noise within the model and the policy.
To gain intuition about this objective, consider the case where Φ equals the ﬁnal state embedding .
For simplicity, we consider the scenario when all models M ψ ∼ P and all policies π ∼ Π are such that for any realization of M ψ = m ψ , the conditional entropy H ( P Φ π [ M ψ ] |M ψ = m ψ ) =
σ ( m ψ ) is a function of the world m ψ and not of the policy.
In this case, the policy π EXP is the entropy maximizing policy:
[Equation]
The scenario above is realized for example when all models M ψ ∼ P M ψ are deterministic and all policies in Π are deterministic.
In this case, for every realization of M ψ = m ψ , the conditional entropy H ( P Φ π [ M ψ ] |M ψ = m ψ )
= 0 .
The assumption H ( P Φ π [ M ψ ] |M ψ = m ψ ) =
σ ( w ) is also realized when the distribution Φ( τ ) ∼ P Φ π [ m ψ ] is approximated as a Gaussian distribution
P Φ π
[ m ψ ] ≈ N ( µ ( m ψ , π ) , Σ( m ψ )) with policy dependent mean µ ( m ψ , π ) and policy independent covariance Σ( m ψ ) .
In these cases π EXP is the policy that produces the most diverse distribution over ﬁnal states across the posterior distribution over models P M ψ .
Since access to H ( P Φ π [ M ψ ]) may not be possible, in our experiments we make further approximations inspired by [97].
Under the same gaussian approximation P Φ π [ m ψ ] ≈ N ( µ ( m ψ , π ) , Σ( m ψ )) , optimizing max π H ( P Φ π [ M ψ ]) is achieved by ﬁnding the policy π that makes the ensemble means µ ( m ψ , π ) as far apart as possible.
A suitable surrogate for this objective is to maximize the empirical variance over means,
[Equation]
where |M ψ | denotes the number of samples m ψ , and µ (cid:48) ( π )
=
[Equation]
M
[Equation]
| (cid:80) m ψ µ ( m ψ , π ) .
We now consider a “batch” version of Eq.
9 involving a population of B agents:
[Equation] [Equation]
where (cid:81) B i =1 P Φ π ( i )
[ M ψ ] is the product measure of the candidate policies embedding distributions over the model M ψ .
By deﬁnition the conditional entropy factors as,
[Equation]
The objective in Eq.
10 has a similar interpretation as in the single policy case.
We are hoping to ﬁnd a set of policies whose average conditional entropies are small (Eq. 4), but whose total entropies are large.
Intuitively, we see this objective is more amenable to our population of policies.
Concretely, by considering distributions directly over the space of trajectories, we ensure that each agent does not ‘double count’ the explored states under Eq.
10, whereas applying the same principal to a one-step information gain objective would simply ensure diversity conditioned on that state and action, and doesn’t explicitly ensure diversity in the visited states by the population.
C.1.1 Proof of Lemma
In the proof of Lemma 1 we will make use of the following supporting result,
Lemma 3.
Let p ∈
(0 , 1] and p 1 , p 2 > 0 satisfying p 1 + p 2 = p then,
[Equation]
Proof.
Let p 1 = αp with α (cid:54) = 0 , 1 , then,
[Equation]
Lemma 3 implies the following result,
Lemma 4.
Let p ∈
[0 , 1] K be a vector satisfying (cid:80) Ki
=1 p i = 1 and let C be a partition of [ K ] such that |C| ≤ K − 1 .
For all C ∈ C denote p ( C ) = (cid:80) i ∈
C p i .
The following inequality holds, (cid:88)
[Equation]
If p i >
the inequality is strict,
[Equation]
Proof.
W.l.o.g we call C 1 the partition set containing 1 and assume p 1 > 0 .
If p i > 0 for only one element of C 1 , it must be the case that p 1 log(1 / p 1 ) =
p C 1 log(1 / p C 1 ) .
We now assume that | C 1 | > 1 and that C 1 has at least two indices i (cid:54) =
j such that p i , p j > 0 .
Let l denote the size of C 1 and w.l.o.g.
let’s say C
= { 1 , 2 , · · · , l } .
Lemma 3 implies the following inequalities,
[Equation]
Applying this reasoning to each element in the partition C yields the desired result.
Lemma 1.
When all models M ψ in the support of the model posterior are deterministic and tabular, and the space of policies Π consists only of deterministic policies, there always exists a solution { π ( i )EXP } Bi
=1 satisfying π ( i )EXP (cid:54) =
π ( j ) EXP ∀ i (cid:54) =
j .
Moreover, there exists a family of tabular MDP models, such that the maximum cannot be achieved by setting π EXP ( i )
=
π for a ﬁxed π .
Proof.
First let’s observe that when the models and policies are all deterministic, the conditional entropies satisfy,
[Equation]
We assume the set of policies Π is of size at least B .
Otherwise, the result cannot be true.
The ﬁrst result follows immediately by Lemma 4.
For any ﬁxed π , the product distribution (cid:81) Bi
=1 P Φ π
[ M ψ ] is a distribution over ‘diagonal’ tuples of the form ( b, · · · , b ) (cid:124) (cid:123)(cid:122) (cid:125) size B where b is an embedding.
Consider an arbitrary set of policies π (2) , · · · , π ( B ) satisfying π ( i ) (cid:54) =
π ( j ) (cid:54) =
π for all i, j ∈ { 2 , · · · , K } .
Call π (1) =
π .
The distribution over embeddings induced by the product distribution (cid:81) Bi
=1 P Φ π ( i )
[ M ψ ] is made of tuples of the form ( b 1 , · · · , b B ) .
Call p ( b 1 , · · · , b B ) the probability under the product measure (cid:81)
Bi
=1 P Φ π ( i )
[ M ψ ] of the tuple ( b 1 , · · · , b B ) .
Notice the ‘projection measure’ onto the ﬁrst coordinate satisﬁes
[Equation]
and that (cid:80) b 2 , ··· ,b K p ( b 1 , b 2 , · · · , b K ) = p 1 ( b 1 ) .
This induces a partition over the outcome space ( b 1 , · · · , b K ) corresponding to C ( b 1 )
= { ( b 1 , b 2 , · · · , b K ) }
b 2 , ··· ,b K .
A direct use of Lemma 4 implies the entropy of the product distribution over ﬁner grained tuples is larger than the entropy over the diagonal of the product distribution having all coordinates equal to each other.
In fact this result also informs when the inequality will be strict.
This happens when (for example the measure over M ψ is the counting measure) there exists two worlds M (1) ψ and
M (2) ψ such that the embedding tuples ( b (1)1 , · · · , b (1) B ) and ( b (2)1 , · · · , b (2) B ) satisfy b (1) =
b (2)1 and ( b (1)2 , · · · , b (1) B ) (cid:54) =
( b (2)2 , · · · , b ( B ) B ) .
We will use this observation to prove the second claim.
Consider a family of MDPs formed of depth
L binary trees.
In this family, the paths leading to the L − 1 layer are the same, but the connections from layer L − 1 to layer L are unknown.
The ‘posterior’ distribution is assumed to be uniform over all plausible trees.
Layer L − 1 has size 2 L − 1 and layer L has size 2 L .
W.l.o.g.
we assume the set of nodes in layer L is labeled as [1 , 2 , · · · , 2 L ] .
The distribution over models is supported over the set of partitions of size 2 L − 1 of [1 , · · · , 2 L ] where each partition set has size 2 .
We assume 2 L ≥ B so that there are at least B distinct policies.
Let π be a ﬁxed policy.
It is enough to show there exist two worlds M (1) ψ and M (2) ψ such that policy π ends in the same state for both M (1) ψ and M (2) ψ
but that there exist distinct policies π (2) , · · · , π ( B ) such that their endpoints ( b (1)2 , · · · , b (1) B ) and ( b (2)2 , · · · , b (2) B ) in worlds M (1) ψ and M (2) ψ satisfy ( b (1)2 , · · · , b (1) B ) (cid:54) =
( b (2)2 , · · · , b (2) B ) .
The latter always holds because among the set of models that maintain the same endpoint for π , there is a pair of models that has different endpoints for π (2) for any arbitrary π (2) (cid:54)
=
π .
This sufﬁces to show the entropy of the ensemble of distinct policies is strictly larger than the entropy of the ‘diagonal choice’ ( π, · · · , π )
[Equation]
C.1.2 InfoCascade
Since the mutual information objective in Eq.
3 is submodular, a simple greedy algorithm yields a (1 − 1 e ) approximation of the optimum [73].
This is the same observation that gives rise to the greedy algorithm underlying other batch exploration objectives, such as in BatchBALD [52].
Let’s start by assuming we have candidate policies π (1) , · · · , π ( i − 1) .
InfoCascade selects π ( i ) based on the following greedy objective,
[Equation]
Equation 4 implies
H
[Equation] [Equation]
and therefore if we approximate P Φ π [ m ψ ] by a Gaussian P Φ π [ m ψ ] ≈ N ( µ ( m ψ , π ) , Σ( m ψ )) , the conditional entropy becomes a policy independent term.
In this case ﬁnding π ( i ) boils down to solving for the policy that maximizes H (cid:16)(cid:81)
[Equation]
] (cid:12)(cid:12)(cid:12) π (1) , · · · , π ( i − 1) (cid:17) .
Using the same approximations described for the single policy objective, a suitable surrogate for this objective is to maximize max π Var( π | π (1) , · · · , π ( i − 1) ) , the empirical variance over means with respect to the policies found so far { π ( j ) } i − 1 j =1 ,
Var( π | π (1) , · · · , π ( i − 1) )
=
[Equation]
Where µ (cid:48) ( π, π (1) , · · · , π ( i − 1) )
=
[Equation]
D Theoretical RL Intuition
The problem we study in this work can be thought of as the “batch” version of the Reward Free
Exploration formalism [45].
In this setting, the learner interacts with an MDP in two phases: 1) a training phase, where it is allowed to collect data from the environment; 2) a test phase, where the learner is presented with an arbitrary task (parametrized by a reward function unseen during training) and it must produce a near optimal policy.
When faced with this problem, the learner is required to design a careful exploration strategy that permits them to build an accurate model in all areas of the state-space that can be reached with sufﬁcient probability.
There have been multiple follow up works that have also studied this problem in the context of Linear MDPs [114] and neural function approximation [86].
Nonetheless, there has been limited focus on the batch setting, where the learner is required to collect data via parallel data gathering policies in each training phase.
Due to the adaptive nature of the algorithm (every single environment query uses all information collected so far), there is an inevitable drop in performance when moving from the sequential ( fully online) setting to the batch (deployment efﬁcient) setting, as measured by the total number of environment interactions (in our case, T × B where T is the number of training rounds and B the number of parallel collection policies).
To illustrate how CASCADE mitigates this loss in sampleefﬁciency, we outline the pseudo-code of CASCADE - TS , a greedy Thompson Sampling algorithm (see [4]) that produces the i − th batch exploration policy in a tabular enviornment by incorporating fake count data from rolling out policies π (1) , · · · , π ( i − 1) in the model.
In the algorithms we present in the main body, we incorporate data gathered in the model by previous i − 1 selected policies into the batch when building the i − th exploration policy; this is directly related to the approach behind
CASCADE - TS .
Concretely, encouraging policy π ( i ) to induce a high variance among the embeddings produced by { π (1) , · · · , π ( i − 1) } can be thought as adding a bonus to encourage π ( i ) to visit regions of the space that have a low embedding visitation count by the previous policies in { π (1) , · · · , π ( i − 1) } .
[Table]
[Table] [Table]
[Table] [Table]
[Table] [Table]
Solve
, π ( i ) k +1 =
max π E τ ∼ (cid:101) P ( i ) ,π k + (cid:34) H (cid:88)
h
= b ( i ) k
+1 ( s h , a h ) (cid:35)
[Table]
Collect
[Table]
[Table]
14: end for 15: end for
Algorithm 2 works by sampling a model from a model posterior, solving for an optimal uncertainty seeking policy and updating the model with ‘fake’ data corresponding to trajectories collected in the model from this uncertainty seeking policy.
After producing the i − th exploration policy in the batch, the recomputed uncertainty bonuses of the model are reduced in the areas of the state space most visited by the i − th policy.
This will encourage subsequent uncertainty seeking policies (i.e., π ( i +1)
k +1 , π ( i +2)
k
+1 , . . .
π ( B ) k +1 ) to visit parts of the space not yet explored by previous policies.
Why Thompson Sampling?
The reader may wonder why Algorithm 2 samples a model (cid:101) P ( i )
k +1 from a TS posterior instead of using the empirical model resulting from fake and true data.
The answer is that the raw empirical model may ascribe a probability of zero to certain transitions, which means we may miss exploring parts of the true state-action space.
Having a TS prior that ascribes non-zero probabilities to all possible transitions ﬁxes this potential issue.
D.1 Proof of Lemma
Let’s consider the class of deterministic MDPs with S = |S| states and A = |A| actions.
It is clear that playing the same deterministic policy multiple times does not provide us with any more information than playing it once.
We will then show that running Thompson TS, the learner will end up with a nonzero probability of producing at least two distinct policies within a batch of size B > 2 .
We will further assume the posterior is aware of the deterministic nature of the MDP family so that any sample MDP model (cid:101) P ( i ) k +1 is deterministic.
For i > 1 , the model (cid:101) P ( i )
k +1 is one whose transitions are consistent with the fake data generated by π (1) k +1 , · · · , π ( i − 1)
k +1 (and the true data in D k ).
A deterministic MDPs can be encoded as a set of triplets ( s, a, s (cid:48) ) .
We say a model { ( s, a, s (cid:48) ) } s ∈S ,a ∈A is (cid:15) − accurate if at least a fraction of 1 − (cid:15) of the triplets are correct.
We now illustrate how the CASCADE - TS algorithm evolves in this setting.
To simplify our task, we will further restrict our attention to the family of deterministic MDPs made of binary trees of height
L .
The number of nodes in any of such trees equals 2 L +1 − 1 .We will deﬁne the counts to be a large value ( 2 L ) when no data has been collected of a given state action pair.
When data has been collected, we deﬁne the bonus terms to be of order at most 1 .
To ﬁgure out the connectivity of any such trees it is necessary to try out 2 L − 1 distinct sequences of L actions (the nature and connectivity of the remaining leaf nodes can be inferred once all the other ones have been decoded).
To build an (cid:15) − accurate model, it is enough to know the connectivity structure of 1 − (cid:15) proportion of paths corresponding to (1 − (cid:15) )
(2 L − 1) leaf nodes.
To prove the result of Lemma 2, we observe that any batch strategy can be simulated by a fully sequential learner, so it immediately follows that T ( (cid:15), Sequential) ≤ T ( (cid:15), AnyBatchStrategy) .
To show that CASCADE - TS will be better than SinglePolicyBatch in this tree example, as long as B is smaller than the remaining number of paths the learner has not tried out, it will produce a set of B policies different from any policy played so far and also different from each other.
To see why the second inequality is true, it is enough to see that complete paths (real or imagined) have a total reward score of at most L , since the counts are always at most one for each edge that is present.
Nonetheless, the counts over any path (sequence of L actions) that has not been visited nor has been imagined to be visited will be at least 2 L .
Thus, every step in the sequential batch construction mechanism will produce a new policy (sequence of L actions) not existent neither in the true nor the fake data so far.
This ﬁnalizes the proof.
Notice that the SinglePolicyBatch strategy will have tried only K unique paths after K batches have been collected, whereas the CASCADE - TS strategy will have tried BK , resulting in a potentially substantial improvement in sample efﬁciency.