Paper: Reward-Free Curricula for Training Robust World Models (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
> It is questionable whether a single intrinsically-motivated policy could explore regions of the state-action space associated with high rewards

As you point out, our work assumes that an intrinsically motivated agent will sufficiently explore regions of the state-action space that are relevant to learning performant policies for downstream tasks. **We have added discussion of this assumption and potential limitation to the conclusion.** The conclusion now reads:

*In future work, we are excited to scale our approach to even more complex domains with many variable parameters. One limitation of our approach is that it relies upon an intrinsically motivated policy to adequately explore the state-action space across a range of environments. This may pose a challenge for scalability to more complex environments. To scale WAKER further, we plan to make use of function approximation to estimate uncertainty throughout large parameter spaces, as opposed to the discrete buffer used in this work. We also plan to investigate using WAKER for reward-free pretraining, followed by task-specific finetuning
to overcome the challenge of relying upon intrinsically motivated exploration.*

We would like to point on that this assumption is also made in other works that consider reward-free world model learning (such as Plan2Explore). As the results from the Plan2Explore paper demonstrate, it is possible to learn performant policies for downstream tasks such as locomotion by using only intrinsically motivated exploration.

Likewise, the results from our paper show that it is possible to learn performant policies for a range of tasks *and environments* by only using intrinsically motivated exploration. The terrain locomotion results in our paper use the same reward functions as some of the DeepMind Control Suite (DMC) [1] tasks. The rewards received are a similar order of magnitude as those obtained by online RL algorithms in DMC [1], despite the fact that we use only intrinsically motivated exploration and train and evaluate across varying terrain. In the block pushing tasks in our paper, the best learned policies achieve > 95% in terms of average completion (see Table 5 in Appendix B.4), which also demonstrates that the policies learned from intrinsically motivated exploration are performant across a range of environments.

[1] Tassa, Yuval, et al. "Deepmind control suite." *arXiv preprint arXiv:1801.00690* (2018).

Reviewer comment (Theory):
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.

Author response:
> Learning the reward function from the dataset

We agree that the dataset must cover high-reward regions to learn an accurate reward function, and successfully learn a particular task. As mentioned above, we have added a discussion of the limitations of intrinsically motivated exploration to the conclusion.

Assumption 1 states that we can recover the optimal policy *within the world model*. It makes no assumption that the world model (or the learnt reward function) is accurate throughout the state-action space. The assumption that addresses the coverage of the state-action space is that we can find a suitable $\pi_\theta^{expl}$ that will seek out the regions of the state-action space with high uncertainty. If the state-action space is very large, it may be infeasible for this exploration policy to explore all regions with high uncertainty. 

&nbsp;
> Does Assumption 1 imply that the learned latent dynamics will have support in the full state-action distribution?

The world model is defined and implemented in such a way that it will embed any sequence of previous observations and actions into the latent space. Then, conditioned on the current action, it will make a prediction over the next latent state. So the answer is that yes, it is the case that the learned latent dynamics will take any sequence of observations and actions and make a prediction for the next latent state (i.e. it has full support). Note, however, that Assumption 1 does not imply that the prediction will be accurate.

Reviewer comment (Evaluation):
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.

Author response:
> World Model Error Prediction

Yes, our evaluation of world model error is computed in terms of the mean-squared error of pixel-based predictions. **We have revised the error results in the main body of the paper so that they are evaluated under a performant task-specific policy** (rather than a random policy) as you have suggested. The new results for error predictions are qualitatively similar to our original results. The original error results under a uniform random policy are deferred to the appendix. 

You have also suggested that we should compare the accuracy against the underlying physics simulator state. However, the world model only receives image observations as input, and decodes images. Therefore, it does not make any predictions about the physics state of the simulator. So we are not able to make this suggested evaluation.

Reviewer comment (Theory):
**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)

Author response:
## Implementation and learning a Markovian latent space

The **desire to learn a Markovian latent space is the reason why we chose to use a recurrent model** to compute the latent state. The intuition for why we can treat the recurrent latent state as Markovian is as follows. The RNN encodes the *history* of previous observations into the latent state. The history of observations is a sufficient statistic for the *belief* over partially observable variables [7]. It is well known that if we treat the belief as the state of an MDP, then the optimal behaviour in this *Markovian* model is the optimal behaviour for the original non-Markovian problem [8]. This idea is commonly referred to as the *belief MDP* [8]. Thus, by treating output of a recurrent module as the state, the resulting model is Markovian. Therefore, using a recurrent model allows us to fulfil our assumption of a Markovian latent space.

This is a well-established assumption in many works on deep RL. Previous works on Deep RL have established that utilising an RNN enables non-Markovian problems (such as POMDPs or Meta RL) to be treated as Markovian [1, 2]. In the context of world models specifically, this assumption is also common. For example, the seminal paper [3] addresses POMDPs (i.e. non-Markovian problems) with a recurrent world model. The authors of [3] state that “The latent dynamics define a Markov decision process (MDP; Sutton, 1991) that is fully observed because the compact model states are Markovian”. The same assumption is also made in other previous works on world models such as [4, 5, 6].

In light of this clarification, we would appreciate it if you would consider raising your score.

Reviewer comment (Novelty):
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

Author response:
## Latent dynamics discrepancy

Thank you for pointing out this relevant concept, which evaluates the difference between the latent predictions of a pretrained world model and a world model that is fine-tuned to a specific task. In our work, we do not finetune the world model to each specific task (we only pretrain the world model using reward-free exploration). Therefore, we cannot use this evaluation metric in our problem setting. 

We have added a reference to this work to the introduction as we agree it is relevant.

Reviewer comment (Writing):
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.

Author response:
## Wording of Assumption 2

We have changed the wording to say “for which the expected reward is the same as the real environment”.

Reviewer comment (Theory):
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.

Author response:
## Future work: Large space of parameters

If the space of parameters is much larger, we think that our work could be extended by using a function approximator to estimate the uncertainty across the space of parameters (rather than the discrete buffer of uncertainty estimates in our current implementation). We have added this discussion to the conclusion.

Reviewer comment (Writing):
In section 3.1, "Futhermore".

Author response:
## Typos/Writing Suggestions

Thank you for pointing these out, we have made these corrections!

Reviewer comment (Novelty):
**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.

Author response:
We agree that model uncertainty has been used extensively to guide exploration within a single MDP. Uncertainty in the value function is also often used to guide exploration in problems where the reward function is provided. Our main contribution is posing the problem of generating automatic curricula over a space of MDPs without rewards and proposing an algorithm for this problem. Our method actively explores the space of MDPs, in contrast to previous methods for reward-free exploration which do not allow for direct control (and therefore exploration) of the environment configurations during training. In this way, our method differs from these prior methods by directly generating curricula over the space of training environments.

We discuss the relationship between our work and the aforementioned previous works in the related work section. Please let us know if you feel that there are additional existing works that we should discuss!

Reviewer comment (Experiments):
**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.

Author response:
We would like to point out that WAKER results in a statistically significant improvement over domain randomisation (DR) for the in-distribution evaluation results. The results in the table show that WAKER-M outperforms DR for almost every task. Therefore we disagree that similar performance is obtained compared to DR for the in-distribution evaluation. 

For the out-of-distribution (OOD) evaluation, we chose to evaluate difficult OOD environments as we think it is more important to evaluate how well the world model generalises to harder instances of the environments. The results for the new baselines show that the *Hardest Environment Oracle* baseline performs poorly for both the in-distribution evaluation, and the OOD evaluation on stairs terrain. **This demonstrates that simply mastering the most complex environment is insufficient to obtain strong performance on our evaluations.**

Reviewer comment (Experiments):
**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.

Author response:
Thank you for your suggestion for additional baselines! **We have added results for the additional baselines you have suggested - please see the general response above.** 

Thank you for updating your review! We have just updated the paper to include results for the new baselines for three out of four domains. We will ensure that the full baseline results are complete in the final version.

Kind regards,

The authors

Reviewer comment (Experiments):
(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.

Author response:
Thank you for your review! We have added results for **additional baselines** as described in the general response. We will respond to your individual comments below.

&nbsp;
> From the provided empirical evaluation it is hard to say whether the procedure can actually work when facing more diverse classes of environments.

We would like to point out that our paper includes experimental results where we learn a **single world model** for **both the terrain and block pushing environments**. Due to space constraints, these results are in Appendix B.1. These experiments demonstrate that WAKER is effective when learning across very diverse domains.

We have added further discussion of these results in the main paper to more strongly emphasise them.

Reviewer comment (Experiments):
From the provided empirical evaluation it is hard to say whether the WAKER algorithm can actually work when facing more diverse classes of environments.

Author response:
Thank you for your review! We have added results for **additional baselines** as described in the general response. We will respond to your individual comments below.

&nbsp;
> From the provided empirical evaluation it is hard to say whether the procedure can actually work when facing more diverse classes of environments.

We would like to point out that our paper includes experimental results where we learn a **single world model** for **both the terrain and block pushing environments**. Due to space constraints, these results are in Appendix B.1. These experiments demonstrate that WAKER is effective when learning across very diverse domains.

We have added further discussion of these results in the main paper to more strongly emphasise them.

Reviewer comment (Novelty):
(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 becomes 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.

Author response:
> (C1)  Theoretical contribution claims.

Our work extends results for the sub-optimality gap in a single MDP with a known reward function to a **set of MDPs** with an **unknown reward function**. We have modified the wording in the introduction to make it clear that our theoretical contribution extends existing theoretical works. The contributions part of the introduction now reads:

“*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”

Extending these theoretical results to the setting of a space of MDPs with an unknown reward function enables us to analyse robustness (in terms of minimax regret) in the reward-free setting. Thus, this analysis serves to motivate our final practical algorithm (WAKER).

Reviewer comment (Novelty):
(C2) The regret definition in Eq. 1 is 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).

Author response:
> (C2) Problem Formulation

The domain randomisation policy defined by Chen et al., 2022 optimises for the expected value in expectation across a distribution of MDPs. Likewise, the Bayes-optimal policy is a history-dependent policy that optimises the expected value in expectation across a prior distribution of MDPs. Because of the *expectation* over MDPs, these works optimise the *expected* regret.

In our work, we address a different problem setting. We are interested in the *maximum regret* over a set of MDPs. As we reference in the paper (above Equation 1) this objective has been used to define a *robust* policy in a number of prior works (Dennis et al., 2020; Jiang et al., 2021a; Parker-Holder et al., 2022; Rigter et al., 2021). Intuitively, the minimax regret objective results in robust performance because to have low maximum regret, the policy must have a low sub-optimality gap across all environments.

As noted above, domain randomisation (DR) optimises for the expected value assuming that all environments are equally likely. Therefore, DR does not optimise for minimax regret: the optimal DR policy may have very high regret in some environments if this improves overall performance in expectation. To optimise for minimax regret, we must consider the sub-optimality in the worst-case possible environment.

Reviewer comment (Theory):
(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).

Author response:
> (C3) Assumption 1: Planning Oracle for POMDPs
   
While we agree that we cannot expect to solve large POMDPs exactly, we think that assuming the optimal policy in the world model can be found is a suitable assumption for the sake of deriving and motivating our algorithm. Our opinion is echoed by Reviewer VCF1. As we write in the paper, our assumption is justified by the fact that we can generate *unlimited data* within the world model to optimise the policy, so there are no sample efficiency constraints.

Recent works on world models (e.g. Hafner et al. 2021) have shown that optimising policies using large amounts of synthetic data generated by a world model obtains strong performance on difficult POMDPs (such as Atari).

Reviewer comment (Theory):
(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 latent states, looks really strong.

Author response:
> (C4) Assumption that a suitable latent representation can be learnt (Assumption 2)

This assumption comprises two parts: First, that we can learn a Markovian latent space for a single POMDP, and second, that we can learn this latent space for a set of POMDPs. 

We begin by addressing the first part of this assumption. We consider a world model to be a model that utilises a recurrent neural network (RNN) to encode the history of previous observations into the latent representation (Lines 88 - 91). It is well-established in existing works on deep reinforcement learning that utilising an RNN enables non-Markovian problems (such as POMDPs or Meta RL) to be treated as Markovian [1, 2].

In the context of world models specifically, this assumption is also common. For example, the seminal paper [3] addresses POMDPs (i.e. non-Markovian problems) with a recurrent world model. The authors of [3] state that “The latent dynamics define a Markov decision process (MDP; Sutton, 1991) that is fully observed because the compact model states are Markovian”. The same assumption is also made in other previous works on world models such as [4, 5, 6].

The intuition for why this is a reasonable assumption is as follows. The RNN encodes the *history* of previous observations into the latent state. The history of observations is a sufficient statistic for the *belief* over partially observable variables [7]. It is well known that if we treat the belief as the state of an MDP, then the optimal behaviour in this *Markovian* model is the optimal behaviour for the original non-Markovian problem [8]. This idea is commonly referred to as the *belief MDP* [8].

Thus, we have established that we can expect to learn a Markovian latent space for a single POMDP if the latent state contains memory of the history. To extend this to a *set* of POMDPs, we consider the following reasoning. Different POMDPs will generate different histories of states and observations. Therefore, the RNN will encode the histories for each POMDP into different latent states. Therefore, different POMDPs will be encoded to regions of the latent space. 

Our empirical results for training a single world model across both the terrain and block pushing environments in Appendix B.1 demonstrate that this approach can successfully learn powerful world models even when trained across diverse environments.

Reviewer comment (Evaluation):
(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?

Author response:
> (C5) Evaluating Robustness

As discussed in the problem definition, we are interested in evaluating the *maximum regret* across a range of environments. To evaluate the *maximum* regret, knowing the optimal performance for a single test environment is not sufficient. We need to know the optimal performance for *every single possible environment*. This is infeasible, which is why we use the CVaR evaluation criterion as an alternative robustness measure in the paper for the results in Table 1.