Reward-Free Curricula for Training Robust World Models
Authors: Marc Rigter
Minqi Jiang
Ingmar Posner
Venue: ICLR 2024 poster (2024)
Keywords: curricula
model-based reinforcement learning
world models
reward-free reinforcement learning
robustness
unsupervised environment design
Primary area: reinforcement learning
TL;DR: Curricula for reward-free learning of world models across variable environments.

Abstract: There has been a recent surge of interest in developing generally-capable agents that can adapt to new tasks without additional training in the environment. Learning world models from reward-free exploration is a promising approach, and enables policies to be trained using imagined experience for new tasks. However, achieving a general agent requires robustness across different environments. In this work, we address the novel problem of generating curricula in the reward-free setting to train robust world models. We consider robustness in terms of minimax regret over all environment instantiations and show that the minimax regret can be connected to minimising the maximum error in the world model across environment instances. This result informs our algorithm, WAKER: Weighted Acquisition of Knowledge across Environments for Robustness. WAKER selects environments for data collection based on the estimated error of the world model for each environment. Our experiments demonstrate that WAKER outperforms naı̈ve domain randomisation, resulting in improved robustness, efficiency, and generalisation.

Published as a conference paper at ICLR 2024
REWARD-FREE CURRICULA FOR TRAINING ROBUST WORLD MODELS
Mare Rigter University of Oxford Marcrigter&gnatr com
Minqi Jiang University College London
Ingmar Posner University of Oxford
ABSTRACT
There has been a recent surge of interest in developing generally-capable agents that can adapt to new tasks without additional training the environment. Learn- ing world models from reward-free exploration is a promising approach, and en- ables policies to be trained using imagined experience for new tasks. However. achieving a general agent requires robustmess across different enviromments. In this work, we address the novel problem of generating curricula in the reward-free setting to train robust world models. We consider robustness in terms of mini- max regret over all environment instantiations and show that the minimax regret can be connected to minimising the maximum error in the world model across environment instances. This result informs our algorithm, WAKER: Weighted Ac- quisition of Knowledge across Environments for Robustness. WAKER selects en- vironments for data collection based on the estimated error of the world model for each environment. Our experiments demonstrate that WAKER outperforms several baselines, resulting in improved robustness, efficiency, and generalisation.
INTRODUCTION
Deep reinforcement learning (RL) has been successful on number of challenging domains, such as Go (Silver et al., 2016), and Starcraft (Vinyals et al., 2019). While these domains are difficult, they are narrow in the sense that they require solving a single pre-specified task. More recently, there has been surge of interest in developing generally-cupable agents that can master many tasks and quickly adapt to new tasks without any additional training the environment (Brohan et al., 2022; Mendonca et al., 2021; Reed et al., 2022: Sekar et al., 2020; Stooke al., 2021; Jiang et al., 2022b).
Motivated by the goal of developing generalist agents that are not specialised for a single task, we consider the reward-free setting. In this setting, the agent first accumulates useful information about the environment in an initial, reward-free -free exploration phase. Afterwards, the agent is presented with specific tasks corresponding to arbitrary reward functions, and must quickly adapt to these tasks by utilising the information previously acquired during the reward-free exploration phase. By separating the learning of useful environment representations into initial pre-training phase, reward-free learning provides a powerful strategy for data-efficient RL.
A promising line of work in the reward- free setting involves learning world models (Ha Schmid- huber, 2018), form of model-based RL (Sutton, 1991), where the agent learns a predictive model of the environment. In the reward-free setting, the world model is trained without access a reward function, and instead, is trained using data collected by suitable exploration policy (Sekar et al., 2020). Once world model has been trained for an environment, it possible to train policies en- tirely within the world model (i.e. "in imagination") for new tasks corresponding to specific reward functions within that environment (Sekar et al., 2020; Xu et al., 2022; Rajeswar et al., 2023).
However, to realise the vision of a general agent, it is not only necessary for the agent to be able to learn multiple tasks a single environment: the agent must also be robust to different environ- ments. To enable this, one approach is to apply domain randomisation (DR) (Tobin et al., 2017) to sample different environments uniformly at random to gather a more diverse dataset. However, the amount of data required to learn suitable world model might vary by environment. Unsupervised Environment Design (UED) (Dennis et al., 2020) aims to generate curricula that present the optimal environments to the agent each point of training, with the goal of maximising the robustness of

---

Published as a conference paper at ICLR 2024
the final agent across wide range of environments. However, existing UED approaches require task-specific reward function during exploration (Eimer et al., 2021: Matüien et al., 2019; Portelas et al., 2020; Dennis et al., 2020; Jiang et al., 2021a; Mehta et al., 2020; Purker-Holder et al., 2022; Wang et al., 2019), and therefore cannot be applied in the reward-free setting that we consider. In this work we address the novel problem of generating curricula for training robust agents without access to reward function exploration, we aim to learn a world model that robust downstream tasks and environments.
functinnnction during exploration. To distil the knowledge obtained during reward-free
We first analyse the problem of learning a robust world model in the reward-free setting. We then operationalise these insights in the form of novel algorithms for robust, reward-free world model learning. Inspired by past works on UED with known reward functions (Dennis et al., 2020; Jiang et al., 2021a; Parker-Holder et al., 2022), we consider robustness in terms of minimax regret (Sav- age. 1951). To our knowledge, WAKER is the first work to address automatic curriculum learning for environment selection without access to a reward function. We make the following contributions: a) We formally define the problem of learning a robust world model in the reward-free setting, in terms of minimax regret optimality, b) We extend existing theoretical results for MDPs to prove that this problem is equivalent to minimising the maximum expected error of the world model across all environments under a suitable exploration policy, and finally c) We introduce WAKER, an algo- rithm for actively sampling environments for exploration during reward-fres training based on the estimated error of the world model in each environment. We introduce pixel-based continuous con- trol domains for benchmarking generalisation in the reward-free setting, We evaluate WAKER on these domains, by pairing it with both an caariniicall-mmoiivate exploration policy. Our results show that WAKER outperforms several baselines, producing more performant and robust policies that generalise better to out-of out-of-distributtin (OOD) environments.
rewarffunction.
exploration policy and a random
2 PRELIMINARIES
A reward-free Markov Decision Process (MDP) is defined by M - A, where S is the set of states and A is the set of actions. T S x A + ASS) is the transition function, where A ASS) denotes the set of possible distributions over S. For some reward function, R: S A - [0. 1]. write M to denote the corresponding (standard) MDP (Puterman, 2014) with reward function R and discount factor 7- A reward-free Partially Observable Markov Decision Process (POMDP) (Kaelbling al., 1998) is defined by P = (S, A, O, T, T), where O is the set of observations, and S - A(O) is the observation function. history is sequence of observations and actions, h on, ao..... O2, ar, where E O and a, E A. We use H to denote the set of all possible histories. Analogous to the MDP case, PR denotes a POMDP with reward function R and discount factor y.
We assume that there are many possible instantiations of the environment. To model an un- derspecified environment, we consider a reward-free free Underspecified POMDP (UPOMDP): U {S, A, O. To, T. }} (Dennis et al., 2020). In contrast to a POMDP, the UPOMDP additionally in- cludes a set of free parameters of the environment, θ. Furthermore, the transition function depends on the setting of the environment parameters, i.e. T Sx A x θ - A(S). For each episode, the en- vironment parameters are set to a specific value θ Therefore, for each episode the environment can be represented by a standard POMDP P,- = {S, A, O, Te, J, ), where (a a) = (o a, 0).
World Models Model-based RL algorithms use experience gathered by an agent to learn a model of the environment (Sutton, 1991; Janner et al., 2019). When the observations are high- -igh-dimensional, it is beneficial to learn a compact latent representation of the state, and predict the environment dynamics in this latent space. Furthermore, in partially observable environments where the optimal action depends on the history of observations and actions, recurrent neural networks can be used to encode the history into a Markovian representation (Schmidhuber, 1990; Karl al., 2017). In this work, we consider a world model to be model that utilises a recurrent module to predict environment dynamics in Markovian latent space (Ha & Schmidhuber, 2018; Hafner al., 2021).
Let the environment be some reward-free POMDP, P. A world model, W, can be thought of as consisting of two parts: W = [g, T). The first part is the representation model q H + Z, which encodes the history into a compact Markovian latent representation 2 E Z. The second part is the latent transition dynamics model. Î. Z x A + A(Z). which predicts the dynamics in latent space. Because the latent dynamics are Markovian, can think of the world model is approximating the
2

---

Published as a conference paper at ICLR 2024
original reward-free POMDP P with seward-fres MDP in latent space: M = (Z, A, In this work, we consider policies of the form: क : Z + (A). This corresponds to policies that are Markovian in latent space, and history-dependent in the original environment.
Minimax Regret robust optimisation, there a set of possible scenarios, each defined by pa- rameters f E è The goal is to find a solution that achieves strong performance across all scenarios. In the context of reinforcement learning, we can think of each scenario as a possible instantiation of the environment, Pa. For some reward function, the expected value of policy in Pd is (Tz, o 7're ı, Pf), where rt the rewards received by executing in P.R.
Minimax regret (Savage, 1951) a commonly used objective for robust policy optimisation in RL (Chen al., 2022; Dennis al., 2020; Jiang al., 2021a; Parker-Holdee et al., 2022; Rigter et al., 2021). To define the minimax regret objective, we begin by defining the optimal policy for given environment and reward function, ~े,r arg max, VTT, P/t). We refer each P as "task" when clear from context, "environment." The regret for arbitrary policy in environment PR is
REGRET(T, Pf) (TT PB) (TT,
The minimax regret objective finds the policy with the lowest regret across all environments:
Trogret arg min max REGRET(T, P.P).
6EB
Minimax regret aims to find policy that is near-optimal all environments, and therefore robust.
APPROACH
W
train
D
The minimax regret objective defines how op- timise policy to be robust to different environ- ments when the task is known. However, our aim in this work is to train world model such that poli- cies derived from the world model for future tasks robust to different environments. In this section, we present our approach for gathering data to train a robust world model achieve this aim.
imagine
Ty
Pa
01
(a) WAKER overview.
In Section 3.1,
we outline how we learn a single
Pe
Z
world model for many possible environments. In
Za
Section 3.2 we define the Reward-Free Minimax Re- gret objective, which connects minimax regret to world model training by assuming that when a re- ward function is provided, the optimal policy in the world model for that reward function can be found. We then show that we can optimise an up- bound on this objective by minimising the max- imum expected latent dynamics error the world model across all environments, under suitable ploration policy. This informs our algorithm for se- lecting environments sample data from to train the world model, WAKER: Weighted Acquisition of Knowledge across Environments for Robustness (Section 3.4). WAKER biases sampling towards en- vironments where the world model predicted have the greatest errors (Figure la).
P
Ze
Top
(b) Illustration of world model latent space for UPOMDP.
Figure 1: a) WAKER uses error estimates for each environment to choose the next environment to sample data from, Pe. A trajectory ra is collected by rolling exploration policy d in the lected environment. To is added the data buffer which is used train the world model, W. Imagined trajectories in W are used to update the error estimates. b) In the world model, each envi- ronment is encoded to subset, Ze. of the latent space Z by representation model g.
3.1 WORLD MODELS FOR UNDERSPECIFIED POMDPS
We utilise single world model, W = {g. F) to model the reward-free UPOMDP. Consider any parameter setting the UPOMDP, E O, and the corresponding reward-free POMDP, Pg. For any history in Ps. h E Ha. the representation model encodes this into subset of the latent space,
i.e. q Ho Za. where Zo C Z. We can then the latent dynamics model 7 to predict the
3

---

Published as a conference paper at ICLR 2024
latent dynamics in Ze, corresponding to the dynamies of Pg for any E Q. Thus, think of the world model as representing the set reward POMDPs in the UPOMDP by a set reward -ree MDPs, each with their own latent state space: MM {Za. A, A,T, 7}. This is illustrated in Figure 1b.
Using a single world model across all environment parameter settings is natural approach as it a) utilises the recurrent module to infer the parameter setting (which may be partially observable), and b) enables generalisation between similar parameter settings. Furthermore, it is sufficient to train single generalist policy over the entire latent space Z to obtain policy for all environments.
3.2 REWARD-FREE MINIMAX REGRET: PROBLEM DEFINITION
As discussed in Section 2, we can define robust policy optimisation via the minimax regret objective. However, this definition cannot be directly applied to setting where we do not know the reward function during exploration, and our goal is to train a world model. In this section, we present the first contribution of this work, the Reward-Free Free Minimax Regret problem, which adapts the minimax regret objective to the setting of reward-free world model training that we address.
Consider some world model, W, which discussed in Section 3.1 represents the possible environ- ments by a set of reward-free MDPs in latent space, {MoJoge Assume that after training W we are given some reward function, R. We define Tō.n to be the optimal world model policy for that reward function R and parameter setting 0, i.e.
क R = arg max (T.
This is the optimal policy according to the MDP defined in the latent space of the world model for parameter setting and reward function R. and does not necessarily correspond to the optimal policy in the real environment. From here onwards, will make the following assumption.
Assumption Consider some world model W that defines a set of MDPs in latent space {Ma]eee. functiinnntion Assume that given any reward function R, and parameter setting 0 E θ we can can find T.
Assumption reasonable because we can generate unlimited synthetic training data in the world model for any parameter setting, and we can use this data to find policy that is near-optimal in the world model using RL. In practice, we cannot expect to find the exact optimal policy within the world model, however Assumption 1 enables an analysis of our problem setting. We now define the Heward-Freee Minimax Regret problem.
Problem 1 (Reward-Free Minimax Regret) Consider some UPOMDP, U, with parameter set θ. For world model W, and E O, let MD be the latent-space MDP defined by W, that repre- sents real environment Pe with reward function R. Define the optimal world model policy as To arg max, MB) Find the world model, W*, that minimises the maximum regret of the optimal world model policy across all parameter settings and reward functions:
W arg min max REGRET BEGREET(Rİ.R PB)) 8.R
(4)
Problem 1 differs from the standard minimax regret objective in Equation 2 in two ways. First, Problem optimises a world model (not a policy) under the assumption that the optimal world model policy later be recovered (i.e. Assumption 1). Second, the maximum is over all possible reward functions in addition to parameter settings. This makes Problem 1 suitable for the reward- free setting: do need access specific reward function when training the world model.
3.3 THEORETICAL MOTIVATION
Before presenting our algorithm for Problem 1, we first provide the motivation for our approach. We make the assumption that the world model learns a suitable representation model, q, which encodes any sequence of observations and actions in the UPOMDP into a Markovian latent state.
Assumption Consider the representation model learnt by the world model, q Hy + Ze. for all O E θ. Assume that given the representation model q, there exists a latent transition dynamics func- tion, T, foror which the expected reward the same as the real environment: V(z, P.) (T.M)) where MB (Zo A, R. T, 7), for any policy a, reward function R2 and parameter setting 0.
4

---

Published as a conference paper at ICLR 2024
Assumption 2 states that the representation model successfully encodes any sequence of observa- tions and actions into a Markovian latent state. Therefore, there exists a dynamics function T defined over the latent space that exactly models the real environment. Assumption allows us to reason about the inaccuracy of the world model solely in terms of the difference between the learnt latent dynamics function, T. and the (unknown) exact latent dynamics, T. For this reason, from here on- wards we will solely refer to the latent dynamics function T when discussing the world model, under the implicit assumption that a suitable representation model g is learnt according to Assumption 2.
Using Assumption 2 we bound the sub-optimality of the optimal world model policy for any parameter setting 0 according the difference between the learnt latent dynamics function, T, and the true latent dynamics, in latent MDP Ma. This is stated formally in Proposition 1.
Let T be the learnt latent dynamics in the world model. Assume the existence of representation model g that adheres to Assumption and let T be the true latent dynamics according to Assumption 2. Then, for any parameter setting 0 and reward function R. the regret of the optimal world model policy is bounded according to:
Proposition 1
2y
REGRET (To.R: Pb)) S (1-
E L d n Ma) (T((
+ E 2a in Ma) (TV a), ((H,,,)))
where d(x,M denotes the state-action distribution of a in MDP M, and TV(P, Q) is the total variation distance belween distributions P and Q.
Proof Sketch: This can be proven by applying version of the Simulation Lemma (Kearns & Singh, 2002) twice. The full proof provided in Appendix B.
Proposition tells us that the optimal world model policy will have low regret if T is accurate under the latent state-action distribution of both nे, and #े. in Mo. However, during data collection we do not have access to the reward function. Therefore, we do not know these distributions as the state- action distribution induced by both Tÿ. and क.. depends upon the reward function. To alleviate this issue, we define an exploration policy, xe expl that maximises the expected error (in terms of total variation distance) of the latent dynamics model:
expl "g
arg max E.
2,a
This allows us to write an upper bound on the regret that has no dependence on the reward function:
47
REGRET PB) <
- 71- e. anl MMa) (T) for all R.
E
(6)
Therefore, we can upper bound the objective of Problem by the maximum expected TV error in the latent dynamics function over all parameter settings:
max REGRET 8.R
Pb) I max
(1 7)2 z.and( MM) TV(T()2))
4-
(7)
We now formally state the Minimax World Model Error problem, which proposes to optimise this upper bound approximation the Reward-Free Minimax Regret objective in Problem 1.
Problem 2 (Minimax World Model Error) Consider some UPOMDP, U, with parameter set θ
and world model latent dynamics function T. Let T be the true latent dynamics function according to Assumption 2. Define the world model error for some parameter setting as:
WORLDMODELERROO (T,0) Bz Maa) (T)
Find the world model that minimises the maximum world model error across all parameter settings:
T arg min max WORLDMODELEEROR(I,
f
DEG
Problem 2 optimises an upper bound on our original objective in Problem (see Equation 7). Prob- lem 2 finds world model that has low prediction error across all environments under exploration
5

---

Published as a conference paper at ICLR 2024
policy that seeks out the maximum error. This ensures that for any future reward function, the op- timal world model policy will be near-optimal for all environments. In the next section, we present our algorithm for selecting environments to gather data for world model training, with the aim of minimising the maximum world model error across all environments, per Problem 2.
3.4 WEIGHTED ACQUISITION OF KNOWLEDGE ACROSS ENVIRONMENTS FOR ROBUSTNESS
Overview In this section, we address how to select environments to collect data for world model training. so that the world model approximately solves Problem 2. We cannot directly evaluate the world model eπror in Equation 8, as it requires knowing the true latent dynamics function, T. Therefore, following works on offline RL (Lu et al., 2022: Yu et al., 2020; Kidambi et al., 2020), we use the disagreement between an ensemble of neural networks as an estimate of the total variation distance between the learnt and true latent transition dynamies functions in Equation 8. This enables us to generate estimate of the world model error for each environment. Then, when sampling environments, our algorithm (WAKER) biases sampling towards the environments that are estimated to have the highest error. By gathering more data for the environments with the highest estimated error, WAKER improves the world model on those environments, and therefore WAKER reduces the maximum world model error across environments required by Problem 2.
WAKER is presented in Algorithm 1, and illus- trated in Figure 1a. We train a single exploration policy over the entire latent space to approximately optimise the exploration objective in Equation 5 across all environments. Therefore, we refer to the exploration policy simply as nrpp,, dropping the
Algorithm Weighted Acquisition of Knowledge across Environments for Robustness (WAKER)
1: Inputs: UPOMDP with free parameters B; Boltzmann temperature Imagination horizon h;;
2: Initialise: data huller D) error huffer Davei world model
W = {g. T) exploration policy ooo
3:
while training world model do
dependence on 0. We maintain buffer Devor of the error estimate for each parameter setting that we have collected data for. To choose the envi- ronment for the next episode of exploration, with probability PDR we sample 0 using domain ran- domisation (Line 5). This ensures that we will eventually sample all environments. With prob- ability 1 PDR. we sample 6 from a Boltzmann distribution, where the input to the Boltzmann dis- tribution is the error estimate for each environment in Demor (Line 7). Once the environment has been selected we sample a trajectory, Tp. by rolling out " in the environment (Line 8). We add Ty to the data buffer D, and perform supervised learning on D (Line 10) to update the world model.
4:
i ~ -Uo.a < PDR or (l then
5:
bomeinnarrdonissttoonll
6:
else
7:
~ (ol ixzee)))))))()
8:
TO + rollout of "P in Pa Collect rea traj. for 0
9:
Add TO to
Train W
1I:
a Dime + Imagine(D, Deme. W. W
12:
fumction Imagine(D. Dena. W. n/))
13:
for I...., K do
14:
D.sample()
Sample real trajectory Encode latent states
Z. (x1) 74| g(rs)
to rollout प, for steps from Z-, in W
17:
6p < Eérorr estimate for 0 via Eq. 10 on Fg
18:
Update Deme with dg for
19:
Train nu on Fo
20: return a aul Pame
Imag ine updates expl and Demor using imagined rollouts (Line 11). For each imagined rollout, we first sample real trajectory To E D (Line 14). We encode Tp into latent states Zr, (Line 15). We then generate an imagined trajectory, Fr. by rolling out cupl using I starting from an initial latent state E Zr, (Line 16). Thus, Fa corresponds to an imagined trajectory in the environment with parame- ter setting 8, as illustrated in Figure 1b. We wish to estimate the world model error for environment parameter setting& from this imagined rollout. Following previous works (Mendonca et al., 2021; Rigter et al., 2023; Sekar et al., 2020), we learn an ensemble of N latent dynamics models: (T; L.N. Like (Yu et al., 2020; Kidambi et al., 2020; Lu et al., 2022), we use the disagreement between the ensemble means as an approximation to the world model TV error for parameters 6:
WORLDMODELERROKK
E 2.a snd (Mo) Var [f a)l
(10)
We compute the eтror estimate using the latent states and actions in Fo in Line 17. We use this to update our estimate of the world model error for environment 0 in Demor (Line 18). Optionally, we may also use the imagined rollout to update the exploration policy in Line 19. For the world model architecture, we use DreamerV2 (Hafner et al., 2021). Implementation details are in Appendix E.
Error Estimate Update Function We consider two possibilities for updating the error estimate for each f in D Demor on Line 18: 1) in WAKER-M, MM, D Demor maintains a smoothed average of the magnitude of the error estimate for each 0: 2) In WAKER-R, we update Denoe to maintain smoothed average

---

Published as a conference paper at ICLR 2024
of the rate of reduction in the error estimate for each 8. Thus, WAKER-M biases sampling towards environments that have highest error, while WAKER-R biases sampling towards environments that have the highest rate of reduction in error. More details are in Appendix E.
Exploration Policy We must learn an exploration policy лh that approximately maximises the world model error according to Equation 5. By default, we use Plan2Explore (Sekar et al., 2020), and train the exploration policy to maximise the approximation in Equation 10. To test whether our approach is agnostic to the exploration policy used, we also consider random exploration policy.
Zero-Shot Task Adaptation Once the world model has been trained, we use to derive a task- specific policy without any further data collection. For reward function R. we find single task policy St that is defined over the entire latent space Z, and therefore all environments. We use the same approach as Sekar et al. (2020): we label the data in D with the associated rewards and use this data to train reward predictor. Then, we train the task policy in the world model to optimise the expected reward according to the reward predictor. All task policies are trained in this manner.
4 EXPERIMENTS
We seek to answer the following questions: a) Does WAKER enable more robust policies to be trained in the world model? b) Does the per- formance of WAKER depend on the exploration policy used? c) Does WAKER lead to stronger generalisation to -fiit--distiiuttion environ- ments? The code for our experiments is available at github. coltub.com/mrrc---- rigter/waker.
Methods Compared To answer question a), compare WAKER- M and WAKER-R with the following baselines: Domain Randomi- sation (DR), samples uniformly from the default environment distri- bution; Gradual Expansion (GE) gradually increases the range of en- vironments sampled from the default distribution; Hardest Environ- ment Oracle (HE-Oracle) samples only the single most complex envi- ronment; Re-weighting Oracle (RW-Oracle re-whissts the environ- ment distribution to focus predominantly on the most complex envi- ronments. Note that the HE-Oracle and RW-Oracle baselines require expert domain knowledge. Detailed descriptions of the baselines can be found in Appendix G.
Figure Example training environments. Rows: Ter- rain Walker, Terrain Hopper, Clean Up. Car Clean Up.
For both WAKER variants, we set PDR 0.2 and perform limited tuning of the Boltzmann temper- ature 1). More details are in Appendix F.3. To investigate question b), we pair WAKER and DR with two different exploration policies: Plan2Explore (P2B) (Sekar et al., 2020) or random exploration policy. For the other baselines we always use the P2E exploration policy.
Domains and Tasks All domains use image observations. For Terrain Walker and Terrain Hopper we simulate the Walker and Hopper robots from the DMControl Suite (Tassa et al., 2018) on pro- cedurally generated terrain. For each environment, there two parameters (amplitude and length scale) that control the terrain generation. For each domain, we evaluate number of downstream tasks (walk, run, stand, flip, walk-back. hop, hop-back. stand). The Clean Up and Car Clean Up domains are based on SafetyGym (Ray et al., 2019) and consist of blocks that can be pushed and ei- ther a point mass (Clean Up) or car robot (Car Clean Up). There are three environment parameters: the environment size, the number of blocks, and the block colours. The downstream tasks (sort, push, sort-reverse) each correspond to different goal locations for each colour of block. Examples of training environments are shown in Figure 2, and more details are in Appendix F.I. We also perform experiments where we train single world model for both the Clean Up and Terrain Walker domains. Due to space constraints we defer these results to Appendix C.I.
Our aim is to train the world model such that the policies obtained are robust, as mea- sured by minimax regret. However, we cannot directly evaluate regret it requires knowing the true optimal performance. Therefore, following (Jiang et al., 2021a; Rigter et al., 2023) we evaluate conditional value at risk (CVaR) (Rockafellar et al.. 2000) to measure robustness. For confidence level a, CVaR, is the average performance on the worst a-fraction of runs. We evaluate CVaRg.1 by evaluating each policy on 100 environments sampled uniformly at random, and reporting the aver-
Evaluation
7

---

Published as a conference paper at ICLR 2024
Explaration Putkyy
PinzEEppiee
Expleratinn
WAKER-M
WAKER-R
DR
GE
BE-0nade
RW-Orack
WAKERM
WAKER.R
Cless
Upp
Seil
Sem--Ro:
0.701 2:0.09
BL595 0.00
D:543 x 0.67
0.397=0011 5-312
1.030
0.68
0:436 x: 0.05
143
5:13 3,83 5:
0337
0.00 000
1400
01881
5000 0.00
0.0
10000 1:00
0:000:000
Cur Claa
Upp
Stil
Sem- Sant
0,914 - 08.01
0,994 0:0.4
0:350
0.815 2:0.1
16080
06655::33
2:14
1:018
0.641
= 009
1L31 1041
0041
2:3m 0/03
1:00M 0.67
15
1440
0.33 = 01
0
x 0.0
0:57
= 067
Wilk
$160 15.
805.3
0 42.0
9
8 15.5
741.2
456 545.3
= 15.3 1916
325
343
36
33:3 253
Run
$12.5 199
505.6
+ 15.1
2799
+ 18.1
500
1 17.4
18.
365.5
152
[20
-042--9
114.1
Alg
Malker
Stued
041.3 + 12.3
945.4
= 16.6
$26.5
17.5
658.5- + 526.3
56.5
423.4
22.0
585
= 31.1
58155
= 08.5
591
3: 65.5
Wilk Beck
762-5
+ 31.8
722
+ 13.5
6
10.1
700
23.1 418
54.3
712.2-1
18.
S0
15.1
511.5
25.5
311
Figure 3: Robustness
Sansia
llepper
Stuee
Back.
659.6 68.3
3347
243
286.1
5690
= 76.4 061.9 51.5
4 21.6
2013
25.
025.
$1. 5072 = 59.7 5566
16.
22:3 1240
= 92.1
41-3
24:14
9.
2577
22
180
+ 102
113
evaluation aggregated
Table 1: Robustness evaluation:
CV of policies evaluated on 100 randomly sam-
Cls.
pled environments.
Espiaain TOEES
Kandam Lapiestin
MAKER-M WAKLE-H
Hi--Orace
Cuean Black
ILSSM ALUN
11315 101
0.00 uiG
Cur Clou 1p Batra Block
0.08 0713
om 0.98 om 8.980 0.08 0.60 D... 0.005
8 0.00
0.10 + bo4
n.18 n0
136
Hemms Wallem Staire
MO.
6226 136 622.
11.9 1071
182
162
66A.1 21.
431.1 226
4181
20.3
4145
Jemas Hopper Steep
32211 = 184 2901 145 1566 21.3 292. 31.1 28.5 226 2
19
212-444
104 2-0
19655
Teruir Hoppor Stalis
3982
375
123
543
Table 2: Out-of-distribution evaluation: average performance on OOD environments.
Figure 4: OOD evalu- ation aggregated Cls.
Here, present the average performance across tasks for each domain. Full results for each task are in Table 4 in Appendix C.2.
age of the worst 10 runs. We also report the average performance over all 100 runs in Appendix C.4. To answer question c), we additionally evaluate the policies on out-of-distribution (OOD) environ- ments. For the terrain domains, the OOD environments are terrain with length scale 25% shorter than seen in training (Steep), and terrain containing stairs (Stairs). For the clean up domains, the OOD environments contain one more block than was ever seen in training (Extra Block).
Results Presentation Tables 1-2 we present results for task policies obtained from the final world model at the end of six days of training. For each exploration policy, we highlight results within 2% of the best score (provided that non-trivial performance is obtained), and + indicates the S.D. over seeds. Learning curves for the performance of policies obtained from snapshots of the world model are in Appendix D.2. To assess statistical significance, we present 95% confidence intervals of the probability that algorithm X obtains improved performance over algorithm Y, computed using the aggregated results across all tasks with the rliable (Agarwal et al., 2021) framework (Figures 3-4).
Results Table presents the robustness evaluation results. For both exploration policies. WAKER- M outperforms DR across almost all tasks. For the Plan2Explore exploration policy, WAKER-R also outperforms DR. Figure 3 shows that these improvements over DR are statistically significant, as the lower bound on the probability of improvement is greater than 0.5. Both WAKER variants result in significant improvements over DR when using Plan2Explore, suggesting that WAKER more effective when combined with sophisticated exploration policy. Between WAKER-M. -M and WAKER-R, WAKER-M (which prioritises the most uncertain environments) obtains stronger per- formance. This is expected from our analysis in Section 3.3, which shows that minimising the maximum world model error across environments improves robustness. Figure 5 illustrates that WAKER focuses sampling on more complex environments: larger environments with more blocks in the clean up domains, and steeper terrain with shorter length scale and greater amplitude in the terrain domains. Plots of how the sampling evolves during training are in Appendix C.3. Regard- less of the environment selection method, Plan2Explore leads to stronger performance than random exploration, verifying previous findings (Sekar al., 2020). The results for average performance in Appendix C.4 show that WAKER achieves improved or similar average performance compared to DR. This demonstrates that WAKER improves robustness at no cost to average performance. The results for training single world model for both the Clean Up and Terrain Walker environments in Appendix C.1 demonstrate that even when training across highly diverse environments, WAKER also achieves more robust performance in comparison to DR.
The robustness results in Table show that GE obtains very similar performance to DR. This unsurprising, GE does not bias sampling towards more difficult environments, and only modifies the DR distribution by gradually increasing the range of environments sampled. HE-Oracle obtains poor performance, demonstrating that focussing on the most challenging environment alone is insuf- ficient to obtain a robust world model. This is expected from the analysis in Section 3.3 which shows that to obtain robustness need the world model to have low error across all environments (not just the most complex one). For Terrain Walker and Terrain Hopper, RW-Oracle is a strong baseline, and obtains similar performance to WAKER. However, for Clean Up and Car Clean Up RW-Oracle
8

---

Published as a conference paper at ICLR 2024
obtains significantly worse performance than WAKER. This is likely because by focussing sampling environments with four blocks of any colour, RW-Oracle does not sample diversely enough to ob- tain good robustness. This demonstrates that even with domain knowledge, handcrafting suitable curriculum is challenging.
To assess the quality of the world models, we evaluate the error between the transitions predicted by the world model and real transitions. To evaluate robustness, we compute the error on the worst 10% of trajectories generated by performant task-policy on randomly sampled environments. In Figure 6, we observe that WAKER leads to lower prediction errors on the worst 10% of trajectories compared DR. This verifies that by biasing sampling towards environments with higher error estimates, WAKER reduces the worst-case errors in the world model. This suggests that improved world model training leads to the improved policy robustness that we observe in Table 1, as expected from the upper bound in Equation 7.
Table 2 and Figure 4 present the evaluation on OOD environ- ments, averaged across tasks. Full results for each task are in Appendix C.2. WAKER-M results in a considerable perfor- mance improvement over DR for both exploration policies, and WAKER-R R significantly improves performance when the exploration policy is Plan2Explore. Thus, WAKER-M again obtains better performance than WAKER-R, but both vari- ants of algorithm outperform DR. GE obtains similar slightly worse performance than DR for the OOD environ- ments. For the steep terrain OOD environments, HE--Orace performs quite well as it focuses on sampling the steepest pos- sible in nn-distribution terrain. However, HE-Oracle does not perform well on the stairs OOD environments, demonstrating that sampling a range of environments necessary for strong OOD generalisation. RW-Oracle performs well on the OOD environments across all domains. However, RW-Oracle has the significant drawback that expert domain knowledge is re- quired. These results demonstrate that by actively sampling more uncertain environments for exploration, WAKER leads to world models that are able to generalise more broadly to environments never seen during training, without requiring any expert domain knowledge.
04
30
(a) Clean Up (b) Terrain Walker
Figure 5: parameters sampled by WAKER-M
Heatmaps of environment
P2F.
prediction HES
Image prudiction
20.015
70.055
0.010
0.010
0.000
Wisccrutt PPt
NARK PDB
0.000
PM
(a) Car Clean Up (b) Terrain Hopper
Figure 6: Evaluation of CVaRa.1 a of the world model image prediction error.
CONCLUSION
In this work, we have proposed the first method for automatic curriculum learning for environment selection in the reward-free setting. Due to space constraints, a discussion of related work can be found in Appendix A. We formalised this problem in the context of learning a robust world model. Building on prior works in robust RI., we considered robustness in terms of minimax regret, and we derived a connection between the maximum regret and the maximum error of the world model dynamics across environments. We operationalised this insight in the form of WAKER, which trains a world model in the reward-free setting by selectively sampling the environment settings that induce the highest latent dynamics error. In several pixel-based continuous control domains, we demonstrated that compared to other baselines that do not require expert domain knowledge, WAKER drastically improves the zero-shot task adaptation capabilities of the world model in terms of robustness. Policies trained for downstream tasks inside the learned world model exhibit sig- nificantly improved generalisation to -ff--ii-iitiitiionn environments that were never encountered during training. WAKER therefore represents a meaningful step towards developing more generally- capable agents. In future work, we are excited scale our approach to even more complex domains with many variable parameters. One limitation of our approach is that it relies upon an intrinsi- cally motivated policy to adequately explore the state-action space across a range of environments. This may pose challenge for scalability to more complex environments. To scale WAKER further, we plan to make use of function approximation to estimate uncertainty throughout large parame- ter spaces, opposed to the discrete buffer used in this work. We also plan to investigate using WAKER for reward-free pretraining, followed by speciice finetuning to overcome the challenge of relying upon intrinsically motivated exploration.

---

Published as a conference paper at ICLR 2024
Acknowledgements This work was supported by a Programme Grant from the Engineering and Physical Sciences Research Council (EP/V000748/I) and a gift from Amazon Web Services. Ad- ditionally, this project made of the Tier 2 HPC facility JADE2, funded by the Engineering and Physical Sciences Research Council (FP/TT02205/1).
REFERENCES
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare Bellemare. Deep reinforcement learning the edge of the statistical precipice. Advances in Neural Informa- fion Process Systems, 34:29304-2993 2021.
Igge Akkaya, Marcin Andrychowicz, Maciek Chociej, Mateusz Litwin, Bob McGrew, Arthur Petron, Alex Paino, Matthias Plappert, Glenn Powell, Raphael Ribas, al. Solving rubik's cube with robot hand. arXiv preprint arXiv:1910.071I: 2019.
Minoru Asada, Shoichi Noda, Sukoya Tawaratsumida, and Koh Hosoda. Purposive behavior acqui- sition for real robot by vision-based reinforcement learning. Machine learning, 23:279-303, 1996.
Philip Ball, Jack Parker-Holder. Aldo Pacchiano, Krzysztof Choromanski, and Stephen Roberts. Ready policy one: World building through active learning. In International Conference on Ma- chine Learning, pp. 591-601. PMLR, 2020.
Mare Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos. Unifying count-based exploration and intrinsic motivation. Advances in Neural Information Pro- cesy Systems, 29, 2016.
Homanga Bharadhwaj, Mohammad Babaeizadeh, Dumitru Erhan, and Sergey Levine. Informa- tion prioritization through empowerment in visual model-based RL. International Conference on Learning Representations, 2022.
Anthony Brohan, Yevgen Chebotar, Chelsea Finn, Karol Hausman, Alexander Herzog, Daniel Ho, Julian Ibarz, Alex Irpan, Eric Jang, Ryan Julian, et al. Do as I can, not as I say: Grounding language in robotic affordances. In Conference on Robot Learning, pp. 287-318. PMLR, 2022.
Xiaoyu Chen, Jiachen Hu, Chi Jin, Lihong Li, and Liwei Wang. Understanding domain randomiza- tion for sim-to-real transfer. International Conference Learning Representations, 2022.
David A Cohn, Zoubin Ghahramani, and Michael I Jordan. Active learning with statistical models. Journal of artificial intelligence research, 4:129-145, 1996.
Edward L Deci and Richard M Ryan. Intrinsic motivation and self-eeternnnation in human behav- ior. Springer Science & Business Media, 1985.
Michael Dennis, Natasha Jaques, Eugene Vinitsky, Alexandre Bayen, Stuart Russell, Andrew Critch, and Sergey Levine. Emergent complexity and zero-shot transfer via unsupervised environment design. Advances in Neural Information Process Systems, 33:13049-11 -1009913061, 2020.
Theresa Eimer, André Biedenkapp. Frank Hutter, and Marius I Lindauer. Self-paced context evalu- ation for contextual reinforcement learning. In International Conference on Machine Learning, pp. 2948-2958, PMLR, 2021.
Carlos Florensa, David Held, Markus Wulfmeier, Michael Zhang, and Pieter Abbeel. Reverse cur- riculum generation for reinforcement learning. In Conference on Robot Learning, pp. 482-495. PMLR, 2017.
Carlos Florensa, David Held, Xinyang Geng, and Pieter Abbeel. Automatic goal generation for reinforcement learning agents. International Conference on Machine Learning, pp. 1515- 1528. PMLR, 2018.
Alex Graves, Mare G Bellemare, Jacob Menick, Remi Munos, and Koray Kavukcuoglu. Automated curriculum learning for neural networks. In international conference on machine learning. pp- 1311-1320. PMLR, 2017.
10

---

Published as a conference paper at ICLR 2024
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Advances in Neural Information Process Systems, 31, 2018.
Nick Haber, Damian Mrowca, Stephanie Wang. Fei-Fei and Daniel L Yamins. Learning to play with intrinsically-motivaeed, self-aware agents. Advances in Neural Information Process Systems, 31, 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning. pp. 2555-2565, PMLR, 2019.
Danijar Hafner, Timothy Lillicrap. Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. International Conference on Learning Representations, 2020.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with discrete world models. International Conference on Learning Representations, 2021.
Nick Jakobi. Evolutionary robotics and the radical envelope-of-noise hypothesis. hypothesssssssssssssssssss..s... Adaptive behavior, 6(2):325-36) -388,-666--- 1997.
Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine. When to trust your model: Model- based policy optimization. Advances in Neural Information Process Systems, 32, 2019.
Minqi Jiang, Michael Dennis, Jack Parker-Holder, Jakob Foerster, Edward Grefenstette, and Tim Rocktäschel. Replay-guided adversarial environment design. Advances in Neural Information Processing Systems, 34:1884-1897, 2021a.
Minqi Jiang, Edward Grefenstette, and Tim Rocktäschel. Prioritized level replay. In International Conference on Machine Learning, pp. 49500 PMLR, 2021b.
Minqi Jiang, Michael D Dennis, Jack Parker-Holder, Andrei Lupu, Heinrich Kuttler, Edward Grefenstette, Tim Rocktäschel, and Jakob Nicolaus Foerster. Grounding aleatoric uncertainty for unsupervised environment design. In Advances in Neural Information Processing Systems, 2022a.
Minqi Jiang, Tim Rocktäschel, and Edward Grefenstette. General intelligence requires rethinking exploration. arXiv preprint arXiv:2211.00 2022b.
Leslie Pack Kaelbling, Michael I. Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligncee 101(1-2):99-134 1998.
Maximilian Karl, Maximilian Soelch, Justin Bayer, and Patrick Van der Smagt. Deep variational bayes filters: Unsupervised learning of state space models from raw data. International Confer- ence on Learning Representations, 2017.
Michael Kearns and Satinder Singh. Near-optimal reinforcement learning in polynomial time. Ma- chine learning, 49:209-232, 2002.
Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims. MOREL: Model- based offline reinforcement learning. Advances in Neural Information Process Systems, 33: 21810-21823 2020.
Alexander S Klyubin, Daniel Polani, and Chrystopher L Nehaniv. Empowerment: universal agent-centric measure of control. In IEEE Congress on Evolutionary Computation, volume 1, pp. 128-135. IEEE, 2005.
Michael Laskin, Denis Yarats, Hao Liu, Kimin Lee, Albert Zhan, Kevin Lu, Catherine Cang, Ler- rel Pinto, and Pieter Abbeel. URLB: Unsupervised reinforcement learning benchmark. arXiv preprint arXiv:21 10.iv:2110.15//1 15191, 2021.
Cong Lu, Philip J Ball, Jack Parker-Holder Michael A Osborne, and Stephen J Roberts. Revisit- ing design choices in offline model-basee reinforcement learning. International Conference on Learning Representations, 2022.
11

---

Published as a conference paper at ICLR 2024
Tambet Matiisnn, Avital Oliver, Taco Cohen, and John Schulman. Seacher-studend curriculum learn- ing. IEEE Transactions on Neural Networks and Learning Systems, 31(9):3732-3744 2019.
Bhairav Mehta, Manfred Diaz, Florian Golemo, Christopher J Pal, and Liam Paull. Active domain randomization. Conference on Robot Learning, pp- 1162-1176. PMLR, 2020.
Russell Mendonca, Oleh Rybkin, Kostas Daniilidis, Danijar Hafner, and Deepak Pathak. Discover- ing and achieving goals via world models. Advances in Neural Information Processing Systems, 14:24379-24437 2021.
Sanmit Narvekar, Bei Peng, Matteo Leonetti, Jivko Sinapov, Matthew E Taylor, and Peter Stone. Curriculum learning for reinforcement learning domains: A framework and survey. The Journal of Machine Learning Research, 21(1):7382-7431, 2020.
Georg Ostrovski, Marc G Bellemare, Aäron Oord, and Rémi Munos. Count-based exploration with neural density models. In International Conference on Machine Learning, pp. 2721-2730. PMLR, 2017.
Jack Parker-Holder, Minqi Jiang, Michael Dennis, Mikayel Samvelyan, Jakob Foerster, Edward Grefenstette, and Tim Rocktäschel. Evolving curricula with regret-based environment design. In International Conference Machine Learning, pp. 17473-17498, PMLR, 2022.
Deepak Pathak, Pulkit Agrawal, Alexei A Efros, and Trevor Darrell. Curiosity-driven exploration by self-supervised prediction. In International Conference on Machine Learning, pp. 2778-2787. PMLR, 2017.
Ken Perlin. Improving noise. In Proceedings of the Annual Conference on Computer Graphics and Interactive Techniques, pp. 681-682, 2002.
Rémy Portelas, Cédric Colas, Katja Hofmann, and Pierre- Oudeyer. Teacher algorithms for curriculum learning of deep RL in continuously parameterized environments. In Conference Robot Learming, pp. 835-853, PMLR, 2020.
Rémy Portelas, Cédric Colas, Lilian Weng, Katja Hofmann, and Pierre-yves Oudeyer. Automatic curriculum learning for deep A short survey. In IJCAI 1020-nnternational Joint Conference on Artificial Intelligence, 2021.
Martin L Puterman. Markov decision processes: Discrete stochastic dynamic programming. John Wiley & Sons, 2014.
Sai Rajeswar, Pictro Mazzaglia, Tim Verbelen, Alexandre Piché, Bart Dhoedt, Aaron Courville, and Alexandre Lacoste. Mastering the unsupervised reinforcement learning benchmark from pixels. International Conference Machine Learning, 2023.
Alex Ray, Joshua Achiam, and Dario Amodei. Benchmarking safe exploration in deep reinforcement learning. arXiv preprint arXiv: 1910.0/708, 7(1):2, 2019.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. Transactions on Machine Learning Research, 2022.
Mare Rigter, Bruno Lacerda, and Nick Hawes. Minimax regret optimisation for robust planning in uncertain Markov decision processes. In Proceedings of the AAAI Conference Artificial Intelligence, volume 35, pp. 1990--19999 1993,, 2021.
Marc Rigter, Bruno Lacerda, and Nick Hawes. RAMBO-RI: Robust adversarial model-based offline reinforcement learning. Advances in Neural Information Processing Systems, 2022.
Marc Rigter, Bruno Lacerda, and Nick Hawes. One risk to rule them all: Addressing distributional shift in offline reinforcement learning via risk-aversion. arXiv preprint arXiv.2212.00124 2023.
Tyrrell Rockafellar, Stanislav Uryasev, et al. Optimization of conditional aalue-at-risk. Journal of risk, 2:21-42, 2000.
12

---

Published as a conference paper at ICLR 2024
Leonard Savage. The theory of statistical decision. Journal the American Statistical Association, 46(253):55 67, 1951.
Jürgen Schmidhuber. Reinforcement learning in Markovian and non-Markovian environments, Ad- vances in Neural Information Process Systems, 3, 1990.
Jürgen Schmidhuber. A possibility for implementing curiosity and boredom in model-building neu- ral controllers. In Proceedings of the International Conference on Simulation of Adaptive Behav- ior, pp. 222-222-22222 227, 1991.
Jürgen Schmidhuber. Formal theory creativity, fun, and intrinsic motivation (1990-2010), IEEE Transactions Autonomous Mental Development, 2(3):230-247) 2010.
Julian Schrittwieser, Toannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Gracpel, et al. Mastering Atari, Go, chess and shogi by planning with learned model. Nature, 588(7839):604 2020.
Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak. Planning to explore via self-supervised world models. In International Conference Machine Learning, pp. 8583-8592. PMLR, 2020.
Pranav Shyam, Wojciech Jaskowski, and Faustino Gomez. Model-based active exploration. In International Conference Machine Learning, pp. 5779-5788, PMLR, 2019.
David Silver, Aja Huang, Chris Maddison, Arthur Guez, Laurent Sifre, George Van Den Driessche, Julian Schritwiisser, Ioannis Antonoglou, Veda Panneershelvam, Marc Lanctot, et al. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484 489, 2016.
Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, Max Jaderberg, Michael Mathieu, et al. Open-ended learning leads to generally capable agents, arXiv preprint urXiv:2107. 12808, 2021.
Richard S Sutton. Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 264):160-163, 1991.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Bud- den, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. Deepmind control suite. arXiv preprint arXiv: .0690, 2018.
Josh Tobin, Rachel Fong. Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Do- main randomization for transferring deep neural networks from simulation to the real world. In IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 23-30, IEEE, 2017.
Oriol Vinyals, Igor Babuschkin, Wojeiech M Czamecki, Michaël Mathieu, Andrew Dudzik, Juny- oung Chung, David H Choi, Richard Powell, Timo Ewalds, Petko Georgiev, et al. Grandmaster level in starcraft II using multi-agent reinforcement learning. Nature, 575(7782):350-332) 2019.
Rui Wang, Joel Lehman, Jeff Clune, and Kenneth O Stanley. Paired open-ended trailblazer (POET): Endlessly generating increasingly complex and diverse learning environments and their solutions. arXiv preprini arXiv: 1:XXv:9901.077::0 2019.
Yingchen Xu, Jack Parker-Holder, Aldo Pacchiano, Philip Ball, Oleh Rybkin, Roberts, Tim Rocktäschel, and Edward Grefenstette. Learning general world models in a handful reward-free deployments. Advances in Neural Information Processing Systems, 35:26820-2682 2022.
Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma. MOPO: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33:14129 14111129-1412,, 2020.
13

---

Published as a conference paper at ICLR 2024
A RELATED WORK
Our work focuses on training more robust world-models (Ha & Schmidhuber, 2018; Hafner et al., 2020; 2019; 2021; Janner et al., 2019; Rigter et al., 2022; Schrittwieser et al., 2020) in the reward- free setting. Past works improve data collection by training an exploration policy within the world model (i.e. in imagination) (Sekar et al., 2020; Ball et al., 2020; Mendonca et al., 2021; Xu et al., 2022). Unlike prior methods, WAKER augments action-based exploration by directly exploring over the space of enviromments via autocurriculum. To our knowledge, this is the first work to address curricula for reward-free training of a world model.
WAKER can thus be viewed as an automatic curriculum learning (ACL) (Graves et al., 2017: Porte- las et al., 2021; Narvekar al., 2020) method for training a world model. ACL methods optimise the order in which tasks or environments are presented during training, to improve the agent's 's final performance (Asada al., 1996; Florensa et al., 2017: 2018; Mendonca al., 2021). Existing ACL approaches for environment selection require task-specity evaluation metrics to assess the perfor- mance or learning potential of the agent each environment, and select appropriate environments to sample from (Portelas et al., 2020; Eimer et al., 2021; Dennis et al., 2020; Jiang et al., 2021a; Matüisnn et al., 2019; Mehta et al., 2020; Parker-Holder et al., 2022; Wang et al., 2019). Therefore, existing ACI, approaches are not applicable to the reward-free setting that we address.
Domain randomisation (DR) (Jakobi, 1997; Tobin et al., 2017) samples environments uniformly, and can be viewed as the simplest approach to ACL. Active DR (Mehta et al., 2020; Akkaya et al., 2019) extends DR by upweighting environments that induce divergent behaviors relative to a reference behaviour. However, Active DR requires access to task-specific reward function to generate the reference behaviour and therefore is also not applicable to our reward-free setting.
Unsupervised environment design (UED) (Dennis et al., 2020) refers to class of methods that generate curricula to produce more robust policies. UED typically frames curriculum learning as a zero-sum game between teacher and student, where the teacher actively selects environments to optimise an adversarial objective. When this objective is the student's 's regret in each environment, the student provably follows a minimax regret policy at the Nash equilibria of the resulting game, making minimax regret UED methods a principled approach for learning robust policies (Dennis al., 2020: Jiang et al., 2021a). simple and effective form of UED Prioritised Level Replay (PLR) (Jiang et al., 2021b;a), which selectively revisits environments that have higher estimated "learning potential"' as measured by the temporal difference (TD) error. Extensions of PLR improve on its random search (Parker-Holder et al., 2022) and address issues with UED in stochastic settings (Jiang al., 2022a). However, all existing UED methods require known reward function during curriculum learning. To our knowledge, ours is the first work to address reward free UED.
WAKER uses ideas from intrinsic motivation (Deci & Ryan, 1985; Schmidhuber, 2010) as a bridge between UED and the reward-free setting. By providing agents with intrinsic rewards, such meth- ods enable agents to efficiently explore complex environments in the absence of extrinsic rewards. Intrinsic motivation generalises active learning (Cohn et al., 1996) to sequential decision-making, by rewarding agents for visiting states that: a) have high uncertainty (Haber et al., 2018; Pathak al., 2017: Schmidhuber, 1991; Sekar et al., 2020; Shyam al., 2019), b) can be maximally influ- enced (Bharadhwaj et al.. 2022: Klyubin et al., 2005), or c) have rarely been visited (Bellemare et al., 2016: Ostrovski et al., 2017). Our work extends uncertainty-bast based intrinsic motivation beyond states within single environment to the space of environment metries to be used in place of typical URD objectives that require a task-speeific reward function. WAKER thus highlights the fundamental connection between autocurricula and exploration.
allowing novelty-related
B PROOF OF PROPOSITION 1
We begin by stating and proving a version of the classic Simulation Lemma from (Kearns & Singh, 2002).
Lemma 1 (Simulation Lemma (Kearns & Singh, 2002)) Consider two infinite horizon MDPs,
M = (S, A, A,R,T.Y) and M {S, A, R, T.T, with reward function R S A + [0. 1].
14

---

Published as a conference paper at ICLR 2024
Consider any stochastic policy S (A). Then:
-(E,) $ (1 (1-7) 7)2
2y
E a, MR) (T)
(11)
where d(T,, MR) is the state-action distribution of x in MR, and TV(P, is the total variation distance between two distributions P and Q.
Proof: To simplify the notation of the proof, we will use the notation V ( and V (RR,
Using the Bellman equation for a7 and V we have:
(T7())) (=00s) ((so) R(su: an) + y (00.a0)) (T((')
B Bay (00)
R(s0, + +oB
E B (T((') EN'T((OO VF()))
> (oo
E BXnt(so (7((s)) E E V=((s)
E Ba.T(ess, V=(5') - (s))))
B(o) BBrnT(oo (77(s)) E (00 (s')
+ BBagnz(so) (g) V (*((') Ess (7('')
(12)
We define (s1ls))) SaeA {(a]so)T(s1 so, (s)/s, a)da. This allows us to rewrite the first term from the last line as:
BBoona(so) B (yrr((es))))) (=7(s) Ey (( (=0))) = ,B P (T*(s)) - (si))
We detine Pi (s1, = ( (s PT and Po (32(50) loes laes (s addads. We again apply the Bellman equation:
y sy PT PT (1100)
V (s1) v*T(s)]
ME - (T((0))
(lBuy (BrnA(ssaaa (*(s') E. Ee-Ton (6)
B (:le)
E B r7 (8) E, EynT(ana
7B PT >
7/3
7BaINt(())s))
E Bx (3(')) BeNNTenna (+))
22
by by (1150)
+ 5 a1 ( B. V - ET V (s')
E,
(13)
Combining Equations 12 and we have
(T0) (*0))) = (.aasa((s)) E Bo VT((') - (3
E Bsnt ar
V (ss)) - EVWT(Gn, 1 (=((s)
22 2g (TT(s)) V (822) (14)
15

---

Published as a conference paper at ICLR 2024
Repeatedly applying the reasoning in Equations 12-14 we have
(~0) V (s0) + PT
B Pian( (7())) ((
s,a~
== 7 E s. Ma)
(15)
E ((.T) VT (st) - ByrTT(s,)
where d( MR) the state-action distribution of in MR. By the definition of the reward function (R(s a) [0. we have that (*(s) E 10,1/(1 - -)] We utilise the inequality that:
[E (x) 2 max, [(=))IV((P.Q) Then, taking the absolute value of Equa-
tion 15 and applying the inequality we have:
(~0()) (50o)
- - s.
?
E (a (*))s ENNT(s,)) V*()s)
(16)
27
<
() a))
(1 7)2 s.
Now that we have proven Lemma 1, we return to our original purpose of proving Proposition 1. We begin by restating Proposition I:
Let T be the learnt latent dynamics in the world model. Assume the existence of representation model q that adheres to Assumption 2, and let T be the true latent dynamics according to Assumption Then, for any parameter setting 8 and reward function R. the regret of the optimal world model policy is bounded according
Proposition 1
2y
REGRET (oo T:) PO
E
(1 7)3
dia Ma)
+ E
x.anddRo MMM) TVV (T)
where d(T,M denotes the state-action distribution of π in MDP M. and TV(P, Q) is the total variation distance between distributions P and Q.
Now, let us consider the latent MDP learnt by our world model Mo = (Zn, A, A,,,>) for some pa- rameters 0 E O, as well as the latent space MDP that exactly models the true environment dynamics A, Recall that for some reward function, R, the optimal world model policy is defined to be
according to Assumption 2, .Ma-
0 arg arg V (T, MB)
The regret of the optimal world model policy for parameter setting 0 and reward function R is REGRET (Tr Pp) = V P.)) V RB)
(17)
By Assumption 2 we have that V(T, P.P) = (T, ) for all æ. This allows us to write the regret as the performance difference in M rather than P.:
REGRET(#A P() = V :T
(18)
By the definition of the optimal world model policy, we have that V ( MB (Tā MB) 20
(19)
Adding together Equations 18 and 19 we have REGRET(F2.A.
P S V R MB) V MB) + V(Tē,R,
MB) ((T.R.
< V ( ( M3) ( MB)
V My)
Then, applying the Lemma 1 to both terms on the right-hand side we have
27
REGRET(R3 (..R. PB) <
(1 7)2
20a n Mo)
a.a Mo) Mo)
16

---

Published as a conference paper at ICLR 2024
C KEY ADDITIONAL RESULTS
In this section, we present the most important additional results that were omitted from the main paper due to space constraints.
C.I TRAINING A SINGLE WORLD MODEL FOR TWO DOMAINS
In this subsection, we present results for training single world model on both the Terrain Walker and Clean Up domains. WAKER must choose which domain to sample from (either Terrain Walker or Clean Up) well sample the environment parameters for that domain. The domain randomi- sation (DR) baseline chooses either domain with 50% probability, and then randomly samples the parameters for that domain.
To handle the varying dimensionality of the action spaces between the two domains, we pad the action space of the Clean Up domain with additional actions that are unused in that domain. To ensure that WAKER receives well-calibrated ensemble-based uncertainty estimates between these two significantly different domains, use the following approach. We first normalise the uncer- tainty estimates for Terrain Walker and Clean Up separately, and then concatenate the uncertainty estimates create the buffer of uncertainty estimates. We found that this was necessary because the scale of the ensemble-based uncertainty estimates can differ between the two significantly differ- ent domains. Obtaining well-calibrated uncertainty estimates without requiring this normalisation process is an orthogonal area of research that is outside the scope of this paper.
The results Table show that WAKER outperforms DR on both the robustness evaluation and the out-of- -itto--diistibution evaluation when training single world model on both domains. This highlights the potential for WAKER to be used to train very general world models, to enable agents capable of solving wide range of tasks in wide range of domains.
Table 3: Results for training a single world model for two domains (Clean Up + Terrain Walker) over five seeds. Evaluated after 8M total steps of reward-free training. For each episode, WAKER chooses the domain as well as the environment parameters for that domain. Domain randomisation (DR) randomly samples the domain and the parameters. Due resource constraints. evaluate subset of the tasks.
(b) Oui-of-distribution average performance.
(a) Robustness Evaluation 1C).
Exploration Policy:
Plan2Explore
Exploration Policy:
PlanoExplore
Eny. Sampling:
WAKER-M
DR
Env. Sampling:
WAKER-M
DR
Clean Up: Sort
0.333 1 0.03 0,180
Sort
0,348 + 0.080 0.089 = 0.07
Extra Block Push
0.440 + 0.03 0.294
Clean Up Push
0.436 + 0.141
0.262
0.08
Terrain Walker: Walk
420.5 1 33.4 382.6 34.3
Terrain
Walk
538.5
42.1
499.
1 43.6
Steep
Run
154.9 1 149.6 4 13.9
Walker
Run
195.4 1
192.9 16.6
Terrain Walker:
Walk
417.1 + 55.1
412.8 +3.9
Stairs Run
161.8 + 17.6 160.9 + 10.2
C.2 OUT DISTRIBUTION EVALUATION: FULL TASK RESULTS
In Table in the main paper, we presented the the tasks for each domain. In Table 4 we present the full results for each task.
results in terms of averages across
17

---

Published as a conference paper at ICLR 2024
Explccation Poacy
PlanZExplore
Kandom Exploration
Envirommnt Sampling
WAKER-M
WAKER-R
DR
GR
HE-Orade
RW-Orack
WAKER-M
WAKER-R
DR
Extra Block
Clean Up
Scet
Soct-Rev.
UP99 0734
0.534 0.05
0:437
0.4%
000
0.05
0.300
0.295
+ 028
1 0.06
0.300 0:03
0,04
0458
0441
001
0.03 0.506
0,30 + 0.00
0.03
0.140
0144 000
0.0
0100=
0,110
0005
0.04
6071
0.066 0.05
+ 0m
Push
0,42
0.0
0529
0
0.220
008
0,100
0.141
Cur Clean
Soct
0.749
0.49
0638
000
0.61
0.400
0.04
0415
Up
Soct-Rov
804
0.724
0.72
0.04
0.47
0.00
0.674
0.503
0465
0.374
Block
Pash
0.778
0.04
0.754
0.05
0.74
0.05
0.859
0.04
0.42
0.00
0,740
0.04
0,633
0.00
0,582
0.07
0,466
0.012
Walk
6114 208
5818
51.9
S127
+ S12
S19
29.5
6110
14.9
664.4
21
212.1 212
190.5
16.1
1851
26.5
Terrain
Run
228.6.1:22
213.2
17.2
156.8
9.1
197.
1: 7.6
222.6
1 10.5
252.3
5,6
119,3
1 9.9
1034
10.0
105.2
15.4
Walker
Flip
246.3111
950.3
8.0
929.1
1 23.8
902.0
18.4
978.5
2.4
953.2
15.8
891.5 1 14.8
817.0
59.0
885.4
29.6
Steep
Stand
94008
9949=21.2
951.5
612.5
81.4
Wilk--Rcc
554.
538.1
213.7
203.3
33.5
Wall
68666
0484-
611.7
0226
691.2
27.2
234
27.1
217.3-
219.9
17.0
Terrain
Run
241:2 16.4
235.6
14.1
216
24.3
2259
225.0
94
264.6
8.5
128.2
1 6.0
114.5
119.5 1.7
Walker
Plip
13..9
8.8
927.6
8.4
9293 = 7.1
957.1
= 8.
953.6
2: 6.7
RXL.0
-: 15.1
840
29.7
822.3
10L.
Stairs
Stnnd
972.8
6.2.
971.8
970.2
7
971.2+69
960.7
+ 7.0
969
10.2
607.7
41
665.8
72.2
668.1
72.1
Walks Back
1.5.1 548.1
368.7
34.6
433.2
1 16.5
44.6
32.6
561.9
19.6
312.6
5.4
257.1
21,3
242.6
1 24.8
Terrinn
133
112.9
13,5
14.6
85.7
8,4
1430
16.6
1169
17.9
5.1 2.6
3.9
13.3
137
Hopger
Back.
133.61135
12.3
19.7
12.3
2
Sterp
695.4 29.4
592.6
74.6
20.6
+4.0 12.6
Terrin
Hop
16.5
29.3
25017
249
33.6
2253
2633=
17.4 13.3
25.6 10.2
Hopper
Hop-Back.
1
209.1
23.5
188.5
25.7
164.3
20.4
117.5
177.5
47.3
33.6 1:
19.5 12.2
34.3
Stairs
Stand
720.8 4: 47.6 -9.7
772.2
7728
= 30.8
665.8
54,3
746.7
4:
47.5
51.5 52.6
49.1
29.4
39.2 17.1
Table 4: Out-of d-dstrribbtttoinn evaluation: average performance on OOD environments for each domain and task.
For Terrain Hopper with a random exploration policy, all sampling methods fail to learn a meaningful policy so none are highlighted.
C.3 ILLUSTRATION OF CURRICULA
The plots in Figures 7 and 8 illustrate the parameters sampled throughout world model training for the Terrain Walker and Clean Up domains.
Domain randomisation (DR) samples parameters uniformly throughout training. WAKER-R biases sampling towards environments where the uncertainty estimate is decreasing the fastest. For Terrain Walker, we see in Figure that initially (0 2M steps), WAKER-R samples all parameters equally in a similar fashion to DR. This indicates that during this period, the uncertainty estimate for all envi- romment parameters is decreasing approximately equally. From 2M 5M steps, WAKER-R selects terrain with higher amplitude, and shorter length scale (more complex and undulating terrain). This indicates that during this period, the uncertainty for the complex terrain is decreasing more quickly and therefore is sampled more often. Thus, we observe that for Terrain Walker WAKER-R initially focuses equally on all environments, before switching to sampling more complex terrain when simple domains have converged and therefore no longer exhibit gradient in uncertainty.
WAKER-M biases sampling towards the domains with the highest magnitude of uncertainty. We do not expect WAKER-M to initially focus on the simplest environments, it is always biased towards sampling complex environments where uncertainty is high. In Figure 7, we see that for Terrain Walker WAKER-M initially (0- 0.3M steps) samples the parameters similar DR. From 0.3M steps onwards, WAKER-M more frequently samples complex terrain with high amplitude and short length scale. Thus, WAKER-M consistently samples complex environments with higher uncertainty. Likewise, in Figure 8 we observe that for Clean Up WAKER-M more frequently samples environ- ments with larger arena size and more blocks. Thus, in Clean Up WAKER-M also consistently samples complex environments with higher uncertainty.
For WAKER-R in Clean Up, we observe that in the early stages (0-IM steps) sampling is heavily focused on the most complex environments with larger arenas and more blocks. This indicates that uncertainty is being reduced more quickly on the more complex environments during this stage of training. As training progresses, WAKER-R R gradually places less emphasis on the most complex environments, and samples the environments more evenly.
18

---

Published as a conference paper at ICLR 2024
Length Scale
Amplicudo
0.75
0.50
0.25
0.0
Environment Steps
1e6
Environment Steps
1e6
(a) Terrain Amplitude
(b) Terrain Length Scale
Figure Plots of the parameters sampled for Terrain Walker domain.
DR P2E
WAKER-M + PZE
WAKER-R P2E
Arona Size
Blue Blocks
Groan Blocks
Environment Steps
1e6
Environment Steps
1e6
Environment Steps
1e6
(a) Arena Size
(b) Blue Blocks
(c) Green Blocks
Figure 8: Plots of the parameters sampled for Clean Up domain.
C.4 AVERAGE PERFORMANCE RESULTS
In Table 1 in the main results in Section we presented the results for the robustness evaluation, where we averaged performance over the 10 worst runs on 100 randomly sampled environments. In Table 5 and Figure we present the results from averaging over all 100 episodes. Note that we do not necessarily expect WAKER to improve over DR for this evaluation, as the DR baseline directly optimises for this exact training distribution.
Figure shows that WAKER-M achieves statistically significant improvement in average perfor- mance over DR when using either exploration policy. For WAKER-R, R, the average performance is not significantly different when compared to DR for either exploration policy.
Exploration Poacy
PlanzExplore
Randoms Exploration
Eaviroment Sampling
WAKER-M
WAKER-R
DR
GE
HR-Orack
RW-Oraele
WAKER-M
WAKKR-R
DR
Soct
BUT3
000
012 0U32 0.01
0384
0.44
0342
0002
2
Push
0967
Up
Sect-Rev,
Push
0.991 1
0,990 = n.0
0.988
03999
1: 0.0
0363 0.02
0973
0:01
0.65
0981
0.01 0.828
0.30
0.12
0.04 0.059
0.04
0.911
0.02
001
0.933
001
0199
0002
0569
0.03
Wilk
909 + 12.1
915.4+
10.3
9004
22.7
0058
13.9
6529 75.5 8737
26.
385.8
34.3
3903
4-.10
386.7
74.5
Walker
Terrai
Run
Stnnd
Flip
6.5
973
9704-1-99
388.7
9-6.
965.3
16.11144 14.7
391.5
967.1
970.3 8.2
13.1
403.2
9739
968.0 + 9.0 927.2
13.0
7.1 982.3
261.9 21.4 365.4
25.4 9624
01 973.7
7.4
18.1
12.8
149.4
933.6
1:2.1.222..
10.7
8.0
669
924.3
138.1
119
66.
13.7
680.7
149.6
924.
84.7
15.2
Walks Back.
861.2
575.2
474.1
Terain
Пopper
Hap-
Hap
485. 150
483.7
143
400LS
4563
153
110
25.2
20.7
7290
31.6
794.4
790.3
680.5
771.
Table 5: Average performance evaluation: Average performance 100 randomly sampled environments.
For each exploration policy, highlight results within 2% the best and + indicates the S.D. 5 seeds. For Terrain Hopper with Random Exploration policy, all methods fail learn meaningful policy so none are highlighted.
DR.
DR
PZE
DR Rand.
DR Randd.
Alg
1.0
0.9
0.7
0.6
0.4
Alg
WAKER.
X
PZE
R
PZE
WA Rand.
R Rand.
Figure 9: Average performance aggregated Cls.
19

---

Published as a conference paper at ICLR 2024
D FURTHER RESULTS
Here, we present additional results that expand upon the results already presented in the main paper.
D.1 ADDITIONAL. WORLD MODEL IMAGE PREDICTION ERRORS
To compare the quality of the world model predictions, we compute the errors between the im- ages predicted by the world model, and the actual next image. For each world model, we collect 200 trajectories in each of 200 randomly sampled environments. For the results in the main paper (Figure 6), the trajectories are collected under performant task policies (for the hop task for Terrain Hopper, and for the sort task for Car Clean Up). For the results Figures 10 and 11 the trajectories are collected under uniform random policy.
Along each trajectory, we compute the next image predicted by the world model by decoding the mean of the next latent state prediction, using the decoder learned by DreamerV2. We compute the mean squared error between the image prediction and the actual next image. Then, we average the errors along the trajectory compute the error for the trajectory. We repeat this process along all trajectories to compute the error for each trajectory. Then, we compute CVa of these error evaluations by taking the average over the worst 10% of trajectories. These values are plotted for each domain in Figures 6, 10, and 11.
We observe that the CVaR of the image prediction errors generally lower for world models trained with WAKER than DR. The difference in performance is especially large for the Clean Up and Car Clean Up domains. This mirrors the large difference in policy performance between WAKER and DR on the Car Clean Up and Clean Up domains. This verifies that WAKER is suc- cessfully able to learn world models that are more robust to different environments than DR. For Ter- rain Walker and Terrain Hopper, the difference in the error evaluations between methods is smaller, we still observe that WAKER produces smaller errors than DR when using the Plan2Explore exploration policy.
Note that in Figures and 11 the error evaluation performed using trajectories generated by uniform random policy. Therefore, we observe that the world models trained using data collected by a random policy tend to have lower errors relative to world models trained using Plan2Explore as the exploration policy.
0.0100
Image prediction
MSE (CVaRo.1)
0.0075
0.0050
0.0025
0.0000
WAKER-M + P2E
WAKER-R + PZE
WAKER.M + Random
DR P2E
WAKER-R + Random
DR + Random
Figure 10: Clean Up image prediction emors for uniform random policy. Error indicate standard devia- tions across seeds.
20

---

Published as a conference paper at ICLR 2024
0.0100
Image prediction
MSE (CVaRe.)
0.0050
0.0000
WAKER-M + PZE
WAKER-R + PZE
DR + P2E
WAKER-M + Random
WAKER-R + Random
DR + Random
Figure 11: Terrain Walker image prediction emors for uniform random policy. Error bars indicate standard deviations across seeds.
D.2 PERFORMANCE ON SNAPSHOTS OF WORLD MODEL TRAINED WITH VARIABLE AMOUNTS DATA
In the main results presented in Section 4, we report the performance of policies obtained from the final world model the end six days of training. In this section, we present the performance of policies obtained from snapshots of the world model trained with variable amounts of data. These results are presented Figures 12, 13 and 14. In each of the plots, the r-axis indicates the number of steps of data collected within the environment that was used to train the world model. The values on the y-axis indicate the performance of policies obtained from imagined training within each snapshot of the world model.
As expected, we observe that as the amount of data used to train the world model increases, the performance of policies trained within the world model increases. We observe that for most levels of world model training data, and for both exploration policies, the WAKER variants improve upon DR in the robustness and OOD evaluations in Figures 12 and 13. The magnitude of the improvement tends to become larger as the amount of training data from the environment increases. This suggests that WAKER might be especially effective when scaled to larger amounts of training data.
The average performance of policies trained in snapshots of the world model is presented in Fig- ure 14. For Terrain Walker and Terrain Hopper, the average performance between WAKER and DR for the same exploration policy is similar across all levels of environment data. For Clean Up and Car Cleanup, we observe that both WAKER variants improve the average performance relative to DR when using Plan2Explore as the exploration policy, for almost all levels of data. For the random exploration policy, we observe that the average performance initially improves more slowly when using WAKER, but eventually surpasses the performance of DR. This initially slower improvement in average performance with WAKER is likely due to WAKER putting increased emphasis on col- lecting data from more complex environments.
21

---

Published as a conference paper at ICLR 2024
WAKER-M
WAKER-R
PlanZExplore
WAKER-M Random Exploration
WAKER-R Random Exploration
DR + Random Exploration
Clean Up: Sort
Clean Sort-iteeerse
Clean Push
0.8
0.8
0.0
0.0
0.4
0.2
0.0
steps
1ed
st0ps
1ed
Steps
1et
Car Clean Up: Sort
Car Clean Up: Sort Reversed
Car Clean Up: Push
1.0
1.0
0.0
0.6
0.4
0.2
0.2
0.2
0.0
100
30aps
100
STaps
J0b
Terrain Walker: Walk
Terrain Walker: Run
Terrain Walker Flip
800
800
600
300
800
400
200
100
0
steps
166
300ps
16c
500.00
1eb
Torrain Walker Stand
Terrain Walker Walk- BackaBd
Terroin Hopper: Hop
1.00
300
Rowa/orrrd/Co (Ciaat)
Tital Roward (Clat))
800
300
400
600
200
300
200
S10p3
50055
510p5
Terrain Hooger: Hop: Backnerdd
Terrein Hogper::
800
1hna (CAWO)
300
Iuaa Nonare (DWAN)
600
200
200
Figure 12: Robustness evaluation. Plots show the CVaRa.1 metric: the of worst 10 100 randomly sampled environments. The policies trained zero-shot on snapshots of world model trained on amount of data labelled on æ-axis, Results are averaged across seeds, and shaded regions are standard deviations across sceds.
22

---

Published as a conference paper at ICLR 2024
WAKE Plantexplore
WAKER-R + Man2Explore
DR + Plan2Explore
WAKER-M Random Exploration
WAKER-R + Random Exploration
DR + Random Exploration
Clean Up Extra Block (OOD): Sert
Clean Up Extra Block 1O0D): Sort-Reversed
Clean Up Extra Block (000): Push
3ax. Camp (Mean)
0.
0.3
0.2
0.0
0.0
sbaps
staps
stopa
Car Cleen Up Extra Block (OOD): Sort
Car Clean Up Extra Block (000): Sort Reverse
Cor Clean Up Extra Block (OOD): Push
0.8
0.8
a.6
0.6
Comphetion
6.6
0.4
9.4
0.4
0.2
0.2
0.2
0.0
0.0
0.0
50055
50005
S15g5
Terrain Walker Steep (OOD): Walk
Terrain Walker Stairs (0OO): Walk
250
Terrain Walker Steep (OOD): Run
800
500
200
2014 Noward
150
200
Terrain Walker (COD): Run
Tarrain Walker Steep (GOD):
Terrain Walker Stairs (GCD):
250
1000
230
800
150
100
400
200
Steep Srand
(COD): Stand
Steep
1000
600
"oral Rewwrnt (toom)
800
80
800
600
600
400
800
200
200
1ub
Terain Walter Stairs Walk-Backward
Terrain
Terrain Steep (OOD): Hop Backeard
820
Thei newarrs (Meen)
150
R00
400
160
T0tal
200
155
105
165
Terrain Hopper Steep (00D): Stand
Terrain Hopper Stairs 100D):
Terrain Hopper Stairs (COD): Hlop-Beckward
Tonn Rewert (flesnt
800
250
200
100
30
stegs
1a4
Teb
steps
Jeb
Figure 13: Out-of-distribution evaluation. Plots show the average performance on out-of-diisributtinn environ- ments. The policies are trained zero-shot on snapshots of world model trained on amount of data labelled on a-axis. Results are averaged across 5 seeds, and shaded regions standard deviations across seeds.
23

---

Published as a conference paper at ICLR 2024
WAKER-M
WAKER-R
PlanZExplors
WAKER-M Random Exploration
WAKER-R Random Exploration
DR + Random Exploration
Clean Up: Sort
Clean Sort-itveesee
Clean Push
Compattoon
Completroe
0.0
0.5
0-5
steps
led
steps
led
steps
let
Car Clean Up: Sort
Car Clean Up: Sort Reversed
Car Clean Up: Push
1.0
1.0
1.0
Compkree
0.8
Tast: Comparer
0.8
0.6
2.05
0.4
38856
1.0
30aps
100
Srapa
10b
Terrain Walker: Walk
Terrain Walker: Run
Terrain Walker: Flip
100
1000
800
300
800
800
600
800
200
steps
166
3t0ps
16c
500.00
1eb
Torrain Walker Stand
Terrain Walker Walk- Backkard
Terroin Hopper: Hop
300
800
800
400
honard
600
600
300
400
200
400
200
100
S10g3
5t055
510p5
Terrain Hooger: Hop: Backkerdd
Terrein Hogper: Stand
(Meemt
Teaal Roward O'eant
300
800
Figure 14: Average performance. Plots show the average performance 100 evaluations environments sampled uniformly at random. The policies are trained zero-shot on snapshots the world model trained on amount of data labelled on r-axis. Results are averaged across seeds, and shaded regions are standard deviations across sceds.
24

---

Published as a conference paper at ICLR 2024
E WAKER IMPLEMENTATION DETAILS
Here, we provide additional details implementation of the WAKER algorithm that were omitted from the main paper due to space constraints. The code for our experiments is available at github. haabbb/ommmarrrigee/waakke
Error Estimate Update and Smoothing In Line 17 of Algorithm we compute the error estimate dg according to the average ensemble disagreement over an imagined trajectory. Each 8e value is noisy, so we maintain smoothed average ö9 of the d0 values. Specifically, we use an exponential moving average (EMA). For each parameter setting 0, when we receive a new de value, we update the EMA according to: 8g - ogdg + (1 a0)6, where we set a 0.9999.
For WAKER-M, the error buffer Dener contains the smoothed average of the óg values for each parameter setting: Denro {ōg 60, .... These values are used as input to the Boltzmann distribution used to sample the environment parameters in Line of Algorithm 1.
For WAKER-R, R, we compute the reduction of the d0 values for each parameter setting between each interval of 10,000 environment steps: As 80 - new Because the error estimates change very slowly, we perform a further smoothing of the Ae values using an exponential moving average: Ag + مم + (1 مم)مo, where we set aд = 0.95. For WAKER-R, the error estimate buffer contains these As values: Deror = (Do. Σe, ...,...,, and these values determine the environ- ment sampling distribution in Line of Algorithm
Error Estimate Normalisation We normalize the error values Line 7 of Algorithm I to reduce the sensitivity of our approach to the scale of the error estimates, and reduce the need for hyperpa- rameter tuning between domains. For WAKER-R, we divide each Ag value by the mean absolute value of Ao across all parameter settings. The rationale for this form of normalisation is that if the rate of reduction of error is similar across all environments, then WAKER-R will sample the environments with approximately equal probability.
For WAKER-M, for each 8g value, subtract the mean of dg across all parameter settings, and divide by the standard deviation. This means that regardless of the scale of the error estimates, WAKER-M M will always favour the environments with the highest error estimates, motivated by Problem 2.
World Model Training For the world model, we use the official open-source implementation of D2eamerV2 (Hafner et al., 2021) at https:://ithali com/ dani jar / dreame rv2. For the world model training we use the default hyperparameters from DreamerV2, with the default batch size of trajectories with 50 steps each. For the ensemble of latent dynamics functions, we use 10 fully-connected neural networks, following the implementation of Plan2Explore (Sekar et al., 2020) in the official DreamerV2 repository. We perform one update for every eight environment steps added to the data buffer.
In our implementation Plan2Explore, each member of the ensemble is trained to predict both the deterministic and stochastic part of the next latent state. This is slightly different to the Plan2Explore implementation in the official Dreamer-v2 codebase, where the ensemble only predicts part of the latent state. We made this modification because the derivation of our algorithm indicates that should be concerned with the uncertainty over the prediction of the entire latent state.
Policy Training For the exploration and task policies, we use the actor-critic implementation from the official DreamerV2 repository. During each update, we sample batch trajectories from the data buffer of 50 steps each. For each latent state corresponding we generate an imagined trajectory starting from that latent state, with a horizon of 15 steps. We then update both the task and exploration policies using dynamics gradients (i.e. backpropagation through the world model) computed on these imagined trajectories. We perform one update to each of the task and exploration policies for every eight environment steps.
data sample in the batch,
25

---

Published as a conference paper at ICLR 2024
F EXPERIMENT DETAILS
F.1 DOMAINS AND TASKS
Terrain Walker and Terrain Hopper We simulate the Walker and Hopper robots from the Deep- Mind Control Suite (Tassa et al., 2018). The observations are images of shape 64 x 64 x 3. For each episode, we generate terrain using Perlin noise (Perlin. 2002), a standard technique for proce- dural terrain generation. The terrain generation is controlled by parameters, the amplitude, with values in [0, 1J. and the length scale, with values in [0.2, 2). Domain randomisation samples both of these parameters uniformly from their respective ranges. For examples of each domain and the terrain generated, see Figures 15 and 16.
Figure 15: Examples of training environments for the Terrain Walker domain.
Figure 16: Examples of training environments for the Terrain Hopper domain.
We evaluate 5 different downstream tasks for Terrain Walker. The tasks walk, stand, flip are from the URLB benchmark (Laskin et al., 2021). We also include walk-backward, which uses the same reward function walk, except that the robot is rewarded for moving backwards rather than forwards.
We evaluate 3 different downstream tasks for Terrain Hopper: hop, stand and hop-backward. Hop and stand are from the Deepmind Control Suite Tassa et al. (2018). Hop-backward is the same as hop, except that the robot is rewarded for moving backwards rather than forwards.
Terrain Walker and Terrain Hopper: Out-of-Distributie Environments We use two different types dittr--distribution environments for the terrain environments. The first is Steep, where the during training. This results
length scale of the terrain is 0.15. This is 25% shorter than terrain with sharper peaks than seen training. The second Stairs, where the terrain contains stairs, in contrast the undulating terrain shown in Figure 17 for the Walker robot. The out-of-distribttion environments for Terrain Hopper the same, except that we use the Hopper robot.
in training. The out-of-distribution environments
(a) Terrain Walker Steep
(b) Terrain Walker Stairs
Figure 17: Examples of -itt----iitribttion (OOD) environments for the Terrai Walker domain. The OOD environments for Terrain Hopper the same, except that the Hopper robot.
Clean Up and Car Clean Up These domains based SafetyGym (Ray al., 2019) and consists of a robot and blocks that can be pushed. The observations are images from a top-down view with dimension 64 x 64 x For Clean Up, the robot is point mass robot. For Car Clean Up the robot is differential drive car. In both cases, the action space is two-dimensional.
For each environment, there three factors that vary: the size of the environment, the number of blocks, and the colour of the blocks. For the default domain randomisation sampling distribution, the size of the environment is first sampled uniformly from zize E 1, 4). The number of
26

---

Published as conference paper at ICLR 2024
blocks is then sampled uniformly from {0, ..... size). The number of green blue blocks is then also sampled uniformly. Examples the training environments generated for the Clean Up domain are in Figure 18.
There are three different tasks for both Clean Up and Car Clean Up: sort, push, and sort-reversed. The tasks vary by the goal location for each colour of block. For sort, each block must be moved to the goal location the corresponding colour. For sort-reverse, each block must be moved to the goal location of the opposite colour. For push, all blocks must be pushed the blue goal location, irrespective of the colour of the block.
We define the task completion be the number of blocks in the environment that in the correct goal region divided by the number of blocks in the environment. If there are no blocks in the environment, then the task completion 1. The reward function for each task is defined follows: The agent receives dense reward for moving any block closer to the desired goal region, and the agent also receives a reward each time step that proportional to the task completion.
In the main results we report the task completion at the end of each episode, as it is easier to interpret. We observe that the results for the total reward directly correlate to those for task completion.
Examples of training environments for the Clean Up domain. The environments differ by their size, the number blocks, and the colour of each block. The shaded regions indicate the goal locations.
Figure 18:
Figure 19: Examples of training environments for the Car Clean Up domain.
Clean Up and Car Clean Up: Out-of-Distriuutio Environments For the out-of-dine uuttof-diitribution environments, we place one more block in the environment than was ever seen during training. We set the size of the environment to size = 4, and there are blocks in the environment. The task completion and reward function is defined the same as for the training environments. Examples of the
environments are in Figure 20.
Figure 20: Examples of disttribitiittionnn (OOD) environments for Clean Up Extra Block domain. For Car
Clean Up, the OOD environments are the same except that we use the car robot.
F.2 LENGTH OF TRAINING RUNS
Due to resource limitations, each training run is limited six days run time. This corresponds to total of 7.4 x 10 environment steps for the Terrain Walker and Terrain Hopper domains, and a total of 8.2 x 10 environment steps for the Clean Up and Car Clean up domains. The results reported in Section are for task policies trained in imagination in the final world models at the end of these training runs.
F.3 HYPERPARAMETER TUNING
We the default parameters for DreamerV2 (Hafner al., 2021) for training the world model.
For WAKER, there are two hyperparameters: the probability of sampling uniformly from the de- fault environment distribution, PDR, and the temperature parameter for the Boltzmann environment distribution, 77. In our experiments, we set PDR = 0.2 for all experiments and did not tune this value. We performed limited hyperparameter tuning of the Boltzmann temperature parameter, 7). We
27

---

Published as a conference paper at ICLR 2024
WAKER-M M + Plan2Explore and WAKER-R Plan2Explore for each of three different values of E {0.5, 1.0, 1.5). For each algorithm and 7 value, we ran two seeds for 5e6 environment steps on the Clean Up domain. At the end of 5e6 environment steps, the best performance for each algorithm for CVa on the sort task. We then use this value for WAKER-M and WAKER-R across all experiments, when using both the Plan2Explore and the random exploration policies.
chose the value of 1j that obtained
The hyperparameters used in our experiments for WAKER are summarised in Table 6.
Table Summary hyperparameters used by WAKER.
WAKER-M WAKER-R
PDR
0.2
0.2
0.5
F.4 CONFIDENCE INTERVAL DETAILS
Figures 3, 4 and present 95% confidence intervals of the probability of improvement, computed using the rliable framework (Agarwal et al., 2021). To compute these values, we first normalise the results for each algorithm and task between [0, 1] by dividing by the highest value obtained by any algorithm. We then input the normalised scores in the rliable package to compute the confidence intervals.
confidence interval where the lower bound on the probability of improvement is greater than 0.5 indicates that the algorithm a statistically significant improvement over the baseline.
F.5 COMPUTATIONAL RESOURCES
Each world model training run takes days on an NVIDIA V100 GPU. In our experiments, we train 120 world models in total, resulting in our experiments using approximately 720 GPU days.
G BASELINE METHOD DETAILS
In our experiments, we compare the following methods. Note that two of the baselines (HE-Oracle and RW-Oracle) require expert domain knowledge. The other methods (WAKER, GE, and DR) do not require domain knowledge.
Domain Randomisation (DR) DR samples environments uniformly from the default environment distribution (as described in Appendix F.1).
Hardest Environment Oracle (HE-Oracle) For this baseline, the most complex instance of the environment is always sampled. For the block pushing tasks, the most complex environment is the largest arena, containing two blocks of each colour. For the terrain tasks, the most complex environment is the terrain with the highest possible amplitude (1) and the shortest possible length scale (0.2).
Re-weighting Oracle (RW-Oracle) RW-Oracle re-weights the environment distribution focus predominantly on the most complex environments. 20% of the time RW-Oracle samples from the default domain randomisation distribution. The remaining 80% of the time RW-Oracle samples uniformly from the most complex environments. For the terrain environments, the most complex environments are those where the amplitude within (0.8, and the length scale within [0.2, 0.4). For the block pushing environments, the most complex environments are those where the size of the arena is four and there are four blocks any colour.
Gradual Expansion (GE) 20% of the time, main randomisation distribution. The remaining 80% the time, GE samples from the default dis- tribution, but the default distribution is restricted to only include environments that have been so far. Thus, GE utilises the default domain randomisation distribution to sample new environments gradually increase the range of environments seen during training.
environment is sampled from the default do-
28


Reviews:
[{"review_id": "xfFV6zxLLI", "reviewer": "Reviewer_j9LV", "summary": "This paper tackles the problem of learning world models through reward-free interactions that are robust to environments variations. First, it formulates the learning problem in terms of minimization of the minimax regret. Then, it proves that the problem formulated is equivalent to minimizing the expected error of the world model in the class of environments under an exploration policy. Hence, it designs an algorithm, called WAKER, to actively sample the environments on which the world model is trained on. Finally, it provides an empirical evaluation of the proposed approach in challenging domains, especially showing its superiority against naive domain randomisation.", "questions": "Can the authors address my comments above?", "limitations": "", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- (Relevance) The paper tackles a very relevant problem, which is learning a robust model of the world from reward-free interactions with a class of environments;\n- (Originality) The paper contributes some nice ideas, such as considering the model error in expectation under the worst-case policy.", "weaknesses": "- (Presentation) The first part of the paper seems to set a very ambitious goal, which is to tackle the general problem of learning a model of a class of environments that allows to approximately solve any RL task within the class. I would rather present a narrower scope of extending domain randomisation to reward-free settings through a model-based approach;\n- (Strong assumptions) The set of assumptions does not look completely reasonable. Especially assuming the knowledge of a decoder that maps every environment to an MDP on latent states and a planning oracle for a setting that is essentially POMDP;\n- (Empirical evaluation) Despite the very general problem formulation, the proposed algorithm is evaluated in very structured domains, where the variation between environments and tasks is fairly limited. Moreover, the set of baselines basically counts different variations of the proposed algorithm, without a real external competitor;\n- (Limited theoretical contribution) The paper sells the theoretical characterization of the problem as an important contributions, but there is not much in terms of novel theoretical results.\n\nThis paper proposes an interesting model-based algorithm for domain randomisation in reward-free settings, which is arguably very relevant given the recent advancements in reward-free RL (in a single MDP) and robust RL/domain randomization. However, while the WAKER algorithm looks like a reasonable approach, from the provided empirical evaluation it is hard to say whether the procedure can actually work when facing more diverse classes of environments. Moreover, I think the presentation is not great for this work: The wording of the first sections seems to imply a substantial theoretical contribution, while the actual contribution is essentially an heuristic algorithm. For these reasons, I am currently providing a slightly negative evaluation. I am open to raise my score if the authors can convince me that the proposed algorithm can actually work with reasonable efficiency and the presentation can be reworked to narrow the scope to \"model-based domain randomisation with active training\". I report more detailed comments below.\n\n(C1) The provided formulation of the problem is reasonable and clear, but the paper is overselling that contribution in my opinion. Everything that is said in the first sections is already known for a single environment (especially that the sub-optimality gap of a planning policy can be upper bounded with the approximation error on the transition model). The paper claims to extend that for multiple environments, but the extension become almost straightforward when the paper assumes the existence of a decoder from any of the environment to a single MDP on a latent state representation.\n\n(C2) The regret definition in Eq. 1 is also reasonable, even though it is mostly a generalization gap rather than an online learning regret, but not really novel (e.g., Chen et al., Understanding domain randomization for sim-to-real transfer, 2022). The $\\pi^\\star_{\\text{regret}}$ has also been formulated before under the name of \"domain-randomization policy\" (Chen et al., 2022) or \"Bayes-optimal policy\" (e.g., Ghavamzadeh et al., Bayesian reinforcement learning: A survey, 2015).\n\n(C3) The paper essentially assumes a planning oracle for POMDPs (Assumption 1) as the original problem is a POMDP. This is not uncommon, but it is hardly reasonable, as the original problem is arguably computationally hard (Lusena et al., Nonapproximability results for partially observable Markov decision processes, 2001).\n\n(C4) How can the representation model q be implemented? The requirement defined in Assumption 2, i.e., that a set of partially observable environments can be captured into a single MDP over a latent states, looks really strong.\n\n(C5) The paper says the \"aim is to train the world model such that the policies obtained are robust, as measured by minimax regret. However, we cannot directly evaluate regret as it requires knowing the true optimal performance.\" Is instead evaluating the regret against a world model trained for the test environment a reasonable \"ideal\" target?\n\n(C6) The notion of policy is never defined, which makes unclear whether the paper refers to Markovian policies or general history-dependent policies.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "R8n84jTf7f", "reviewer": "Reviewer_VCF1", "summary": "This work considers the problem of pretraining a world model on reward-free data using a learned exploratory policy, and subsequently using the learned model to train task-specific policies (using a provided reward function) without any additional model learning. The key difference between this work and prior work is the strong focus on learning models that are *robust*: designing curricula for learning world models in the presence of multiple environment variations such that downstream policies learned using the model are more likely to generalize to unseen environment variations. The paper starts by treating the problem formally, then introduces a practical algorithm -- WAKER -- motivated by theoretical insights, and then compares the \"robustness\" of learned world models on OOD environment variations in a number of simulated environments.", "questions": "I'd like the authors to address my questions and suggestions provided in the *weaknesses* section above. Addressing those comments are most likely to change my opinion. I have one additional clarification questions, though:\n- The reward-free MDP introduced in Sec 2 is defined as a tuple that exclude the reward function R (as opposed to the standard MDP) but still includes a discount factor. What is the role of the discount factor in the reward-free setting, if any? Is this just for notational convenience in Sec 3?\n\n\nLastly, typo:\n- P4, below Eq. 3: according the MDP\n\n**Post-rebuttal:** I believe that the authors have addressed my biggest concerns, and I raise my score (5 -> 6) and confidence (3 -> 4) accordingly.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The paper is very well written, and provides sufficient background for readers unfamiliar with the problem to appreciate their theoretical result (Proposition 1). The theoretical treatment is not particularly surprising in itself, but helps motivate the proposed method which I appreciate. Assumptions are clearly stated (perfect model and latent policy exists), which seem pretty reasonable to make for the sake of analysis.\n- The proposed method is intuitive and appears to be fairly straightforward to implement. I consider this a strength.\n- Experiments are conducted on a reasonably large amount of tasks, which makes the results fairly trustworthy overall despite similar results for methods on many of the individual tasks. The proposed method achieves the most significant gains in the OOD evaluations, as one would expect based on the derivation of the algorithm.", "weaknesses": "I have three concerns at the moment:\n\n- **Contributions.** While I really do consider simplicity a strength, the proposed method is undeniably incremental. At surface level, the main algorithmic contribution is automatic curricula by sampling of environments for which the current model error is large, which -- again, at surface level -- has been explored in numerous prior works as a means of guiding e.g. exploration, curricula for which rewards are available, gradient updates, or as model regularization (I will omit references to specific works here to avoid propagating my own biases, since I believe that the authors are aware of the literature, but can provide some during the discussion period if requested). This is in my opinion not a big problem given differences in problem setting as well as some of the more intricate algorithm details and theoretical contributions, but seems worth bringing up.\n- **Experimental design.** It is evident from the empirical results that performance between WAKER and domain randomization is fairly similar for in-domain environment variations, and that most of benefit from the WAKER curricula is in the OOD test environments (Figure 4, Figure 13). However, this appears to mainly be because (1) the task variations considered have a natural order of difficulty (e.g. number of objects or terrain amplitude), (2) WAKER strongly favors sampling more difficult tasks (e.g. max allowed number of objects), and (3) the OOD variations considered here are formulated such that they are most similar to the task variations favored by WAKER (e.g. max possible number of training objects + 1). Therefore, it is not clear whether the models learned by WAKER really are robust or simply converge on a data distribution that is more favorable in these particular test environments.\n- **Baselines.** The authors claim that there are no applicable baselines besides a naive domain randomization baseline that uniformly samples environment variations. However, there are several other manually designed curricula that would be very useful for validating the necessity for automatic curricula: (1) a baseline that strongly favors sampling the most difficult environments, (2) a baseline that is trained **only** on the most difficult environment, and (3) a baseline that performs automatic domain randomization by initially sampling from a small range of environments and gradually increases the range of sampled values (similar to ADR from https://arxiv.org/abs/1910.07113). Baselines 1 and 2 require some domain knowledge in that they need to know which environments are likely to be difficult (could be used as an oracle), but baseline 3 could easily be implemented without any domain knowledge by growing number of environments based on a linear schedule. It seems disingenuous to claim that there are no existing curriculum methods that would be applicable to this problem setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "FVLgpFGObd", "reviewer": "Reviewer_hsNK", "summary": "The work presents an approach for unsupervised environment design in a reward-free setting, using world models and based on the idea of minimax regret. By adopting WAKER, the agent samples environments based on the expected world model error, which is approximated using an ensemble of dynamics models.", "questions": "* In [1], the authors present the \"latent dynamics discrepancy\" metrics to evaluate the world model's misspecification when trained in reward-free settings and show how it correlates with downstream task performance. I think the idea behind this metric is relevant to the latent dynamics error used here and might be worth discussing in related work or exploring the way it is used in [1] for future work\n* In Assumption 2, the authors talk about \"modelling the real environment exactly\", why what matters for the model is actually that the expected rewards are consistent along the distribution of state-actions visited by the actor. I suggest finding a better wording for that.\n* The approach seems to work well in cases where the set of parameters $\\theta$ for the environment's design is discrete/small. I recommend discussing how to extend the work for continuous sets of parameters, as this may require further approximations.\n\nTypos/writing suggestions:\n- in Preliminaries, there is a $\\Delta(X)$ which I think should be $\\Delta(S)$\n- In section 3.1, \"Futhermore\"\n\n[1] Rajeswar et al, 2023, Mastering the Unsupervised Reinforcement Learning Benchmark from Pixels\n\n**Post rebuttal** : the authors have clarified their implementation choices and addressed my concerns. They also extended their evaluation, according to some other reviewers' suggestions. Now, I feel more confident about recommending acceptance for the paper and I have updated my score to reflect this.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "* **Presentation**: the paper is well-written and the presentation is clear and clean. The metrics plotted in the Results section are sensible and the empirical results well reflect the story of the work.\n* **Method**: the idea behind this work is innovative and well-motivated. Given the absence of previous work in this specific setting, I think the method will be useful for future work. The results also corroborate the usefulness of the approach.", "weaknesses": "* **Breaking assumptions in implementation**: I think the approach is well-motivated theoretically. However, in practice, the authors adopted a world model implementation that learns the dynamics using a recurrent module. By exploiting memory, the model is not encouraged to learn a Markovian latent state space. Given this is one of the assumptions behind the work, I think the authors should have opted for a memoryless model (e.g. learning from sequences of data but without keeping a memory) or they should make clear that their implementation is possibly breaking the assumptions (despite providing the expected improvements in practice)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rBEZ40nuyt", "reviewer": "Reviewer_3ipu", "summary": "This work addresses the challenge of learning a robust World Model in a task-agnostic setting, i.e., without the supervision of any reward distribution. For this matter, it formulates the problem as a minimax regret objective and proves that, under assumptions of learnability of the optimal policy (under reward supervision) and existence of a latent dynamics that models the real environment, the regret of a optimal “World Model” policy is bounded and depends on the state-action distribution and the total variation between the learnt latent and the true dynamics. Based on that, the paper introduces WAKER, an agent that actively samples environments to generate a reward-free curricula in order to learn a robust World Model. Experiments show that WAKER outperforms domain randomization in robustness and out-of-distribution settings.", "questions": "- Following my second concern, does Assumption 1 imply that the learned latent dynamics will have support in the full state-action distribution? \n\n- Following my third concern, is the image prediction error indeed computed on pixel-based mean-squared error?\n\n\n====================== **POST-REBUTTAL** =======================\n\nI would like to thank the authors for putting efforts on the rebuttal discussion and addressing concerns. I believe my concerns were partially addressed, and, as a consequence, I am leaning towards acceptance of the paper (5 -> 6). Overall, as described in my initial review, my main concerns are related to the assumptions that the work is built upon. Although I still believe they are strong and limit the applicability of the method in harder scenarios, I also understand that these assumptions are present in previous work and the authors updated the work to highlight these as potential limitations, improving transparency.\n\nIn more details:\n\n- For my first point: the authors addressed the concern as possible, and showed average returns to provide evidence that, for the particular tasks in the paper, this was not a problem.\n\n- For my second point: similarly, they provided discussion on the limitations and clarified my question. \n\n- For my third point: I think it was partially addressed. The paper now provides evidence under a performant policy (rather than random policy) which substantially supports the validity of the experiments. I still have my concerns of using MSE as a predictive metric (and how much semantic value it brings). Nevertheless, I understand that it would indeed require an inverse mapping between image and true physics simulator state, which does not seem easy to obtain as mentioned.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 4, "strengths": "- The work introduces the setting of generating task-agnostic curricula for learning robust world models, arguing that this direction may help learning generally-capable agents, which is an interesting perspective with good motivation. Furthermore, the paper is well written and didactic.\n- The paper mathematically formalizes learning robust World Models under a framework of minimax regret and proves a regret bound for optimal policies under learned World Models. For this, it clearly states the assumptions and presents the proof. The theoretical development is sound and clear. \n    - Furthermore, the derivation starting from the theoretical framework and arriving into the proposed learning agent (WAKER) is clever and well presented.\n\n- The experimental setting shows the effectiveness of WAKER in robust and OOD evaluation settings in comparison with the random sampling approach (DR). It also ablates the exploration policy and shows an interesting and interpretable illustration of the generated curricula.", "weaknesses": "- The proposed learning objective (and, therefore, learning agent) relies on a strong assumption that seems to be overlooked (or at least not discussed with appropriate depth): the “capacity” of the exploration policy to cover the full state-action distribution. Equivalently, the relaxation of the bound by introducing an exploration policy is a big assumption and limits the generality of the proposed agent. It is questionable whether a single intrinsically-motivated policy could explore regions of the state-action space associated with high rewards for locomotion gaits (as in the walker scenario) with terrain variability, for instance.\n\n- Regarding “Zero-Shot Task Adaptation”: The presented methodology for learning task-specific policies relies on the acquired dataset “D” to train the reward predictor. Nevertheless, this raises a similar problem to the previous concern: if D does not contain the state-action pairs associated with high rewards for a particular task, then the learned policy might be potentially very suboptimal. If the work relies on Assumption 1 to disregard this problem, then it should be clearly stated in the paper, but I believe that this would make Assumption 1 unreasonable.\n\n- In Figure 6 and 10, I understood that the prediction error is computed as the mean squared error on the pixel values. If this is indeed the case, it is questionable if this is the best metric to evaluate the prediction error. Furthermore, collecting the test set by sampling a random policy will omit part of the state space, and likely very relevant states (i.e., associated with high rewards in an optimal policy for locomotion). Therefore, the results are not quite convincing. I recommend trying to use the underlying physics of the simulator as a “true latent” dynamics and compute the distance from it.  Another possibility is to learn a separated embedding space (via metric learning or autoencoding) and use it as the ground truth representation. Additionally, I also recommend learning optimal policies for several different environments and use the history of policies generated during training to sample from the space. This would give a more complete covering of the state-action distribution.\n\n\n**Summary of the Review**:\n\nOverall, I believe the paper provides solid contributions, for introducing a new setting, and also providing a sound theoretical framework and cleverly translates it to a learning algorithm. Nevertheless, the introduction setting must consider the challengings of covering the state-action distribution correctly, which was overlooked in the paper. Furthermore, the experiments regarding model prediction error can be improved to be convincing.  From my perspective, concerns 1 and 2 should be addressed by clarifying these limitations in the paper. Concern 3 can be addressed by the suggestions in the experimental setting. *I will recommend acceptance if these concerns were addressed.*", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)