Paper: Conservative World Models (ICLR-2024)

Reviewer comment (Theory):
The “world model” as described in the paper seems to be potentially mischaracterized. The referenced document, [2209.14935.pdf (arxiv.org)](https://arxiv.org/pdf/2209.14935.pdf), elucidates that "Both SFs and FB lie in between model-free and model-based RL, by predicting features of future states, or summarizing long-term state-state relationships. Like model-based approaches, they decouple the dynamics of the environment from the reward function. Contrary to world models, they require neither planning at test time nor a generative model of states or trajectories." Therefore, the paper’s description might need a revision for accuracy.

Author response:
**Q1.** We have retitled the paper and dropped any reference to world models.

Reviewer comment (Experiments):
The set of baselines employed in the experimental section appears to be inadequate. As depicted in Figure 2 of [2209.14935.pdf (arxiv.org)](https://arxiv.org/pdf/2209.14935.pdf), FB does not have the "sufficient performance gap" as claimed by the authors, necessitating the inclusion of a broader spectrum of baselines for a more comprehensive performance comparison.

Author response:
**Q2:** **The set of baselines employed in the experimental section appears to be inadequate. As depicted in Figure 2 of [2209.14935.pdf (arxiv.org)](https://arxiv.org/pdf/2209.14935.pdf), FB does not have the "sufficient performance gap" as claimed by the authors, necessitating the inclusion of a broader spectrum of baselines for a more comprehensive performance comparison.**

We disagree that FB does not sufficiently outperform other SF-based, zero-shot RL methods. Touati et. al (2022) show FB achieved an aggregate score ~10% higher than the SOTA SF-based method after extensive experimental analysis (Figure 1 (an aggregation of Figure 2) and Table 2 in [1]).  FB arguably enjoys better theoretical properties as it avoids the performance ceiling induced by the representation learning methods required for SF-based methods. That said, our proposals are fully compatible with SF-based methods, and they too will suffer from the same failure mode as FB on low quality datasets (TD learning with $a_{t+1} \sim \pi(s_{t+1})$--see Equation 6 in [1]). As such we shall add a new appendix that derives losses for training Conservative Successor Features—see Revision 5. Focussing new zero-shot RL work on FB but showing compatibility with SF in this way has been shown in other works (under review at ICLR) [2].

Reviewer comment (Writing):
The term "world model" seems to have limited relevance and is infrequently used throughout the paper. Could its significance to the paper’s core content be clarified, or is it a concept that could potentially be omitted without loss of clarity?

Author response:
**Q1.** We have retitled the paper and dropped any reference to world models.

Reviewer comment (Experiments):
To bolster the robustness of the study, it would be beneficial for the authors to incorporate a wider array of zero-shot RL baselines, particularly those utilized in [2209.14935.pdf (arxiv.org)](https://arxiv.org/pdf/2209.14935.pdf). Additionally, considering the 2022 inception of the FB method, it may be pertinent to include more recent and relevant baselines in the analysis.

Author response:
**Q3: Additionally, considering the 2022 inception of the FB method, it may be pertinent to include more recent and relevant baselines in the analysis.**

To the best of our knowledge, we are the first work to build upon FB representations in the context of zero-shot RL. If the reviewer is aware of works we have missed we ask for them to be shared so we can include them.

Reviewer comment (Novelty):
A relatively minor weakness is the choice of the phrase "world model" to describe the approach. World models already are a well-established and distinct approach in this area [1].

Author response:
**Q1.** We’ve changed the title in our revised manuscript and removed all mention of world models.

Reviewer comment (Novelty):
There is a relative lack of novelty. The approach is essentially a reweighing of training data, but similar reweighing strategies have already been proposed and are not compared against [2].

Author response:
**Q3: Second is the relative lack of novelty. The approach is essentially a reweighing of training data, but similar reweighing strategies have already been proposed and are not compared against [2].**

We disagree that our approach is reweighting of the training data and therefore lacks novelty. We believe the reviewer has misunderstood our proposals; we are not reweighting the training data, we are (in the case of VC-FB) regularising Q functions to mitigate OOD value overestimation—data is sampled uniformly from dataset $\mathcal{D}$. The method proposed in the cited paper is not a zero-shot RL method (the policy cannot generalise to arbitrary tasks in an environment with zero downstream planning or learning) and so we do not believe it would constitute a valid baseline.

Reviewer comment (Experiments):
The evaluation setup is questionable. It is odd to have two variants of the approach in comparison to a single baseline, which essentially gives the approach twice the opportunity to outperform the baseline. For simplicity, it might have been better to stick with VC-FB and leave MC-FB to the appendix.

Author response:
**Q4: It's odd to have two variants of the approach in comparison to a single baseline, which essentially gives the approach twice the opportunity to outperform the baseline. For simplicity, it might have been better to stick with VC-FB and leave MC-FB to the appendix. However, it's also unfair that both approaches are given 3 times the training duration to FB. Ablations or additional baselines could have helped to avoid this issue.**

Firstly, we note that we compare against 3 baselines not 1—see Table 5 Appendix C. Secondly, we proposed two variants of our method because it was not clear _a priori_ which method of $z$ sampling would prove beneficial. We do not think this constitutes twice the opportunity to outperform our baselines. As it transpired, both methods outperformed all of our baselines in aggregate on their own merit. Thirdly, and finally, we hold the number of update steps and training data constant for all algorithms which is standard practice. The reason training our proposals takes 3x longer than existing methods is not because the methods were exposed to more data or learning steps, but rather because the $Q$-function regularisation invokes expensive operations. We discuss this in Section 5, and suggest this could be resolved with thoughtful engineering of the codebase. Doing so was left for future work.

Reviewer comment (Experiments):
It is unfair that both approaches are given three times the training duration compared to FB. Ablations or additional baselines could have helped to avoid this issue.

Author response:
**Q4: It's odd to have two variants of the approach in comparison to a single baseline, which essentially gives the approach twice the opportunity to outperform the baseline. For simplicity, it might have been better to stick with VC-FB and leave MC-FB to the appendix. However, it's also unfair that both approaches are given 3 times the training duration to FB. Ablations or additional baselines could have helped to avoid this issue.**

Firstly, we note that we compare against 3 baselines not 1—see Table 5 Appendix C. Secondly, we proposed two variants of our method because it was not clear _a priori_ which method of $z$ sampling would prove beneficial. We do not think this constitutes twice the opportunity to outperform our baselines. As it transpired, both methods outperformed all of our baselines in aggregate on their own merit. Thirdly, and finally, we hold the number of update steps and training data constant for all algorithms which is standard practice. The reason training our proposals takes 3x longer than existing methods is not because the methods were exposed to more data or learning steps, but rather because the $Q$-function regularisation invokes expensive operations. We discuss this in Section 5, and suggest this could be resolved with thoughtful engineering of the codebase. Doing so was left for future work.

Reviewer comment (Evaluation):
The presentation of the results is problematic. The paper repeatedly aggregates over distinct problems, whose scores are not comparable, and presents these results as summed values and percentages. This may not be appropriate.

Author response:
**Q4: Why is it appropriate to aggregate the results across evaluation domains?**

DeepMind Control domains are specifically designed such that task scores always lie in $[0, 1000]$ and so are directly comparable. Aggregating across these domains is standard practice in the community, we refer the reviewer to e.g. [1, 2] for examples. We aggregate in the same way as the work we build upon [3].

Reviewer comment (Evaluation):
Given the claims around the value of conservative world models, it is unexpected that they did not outperform FB the most when the dataset was of a poorer quality.

Author response:
**Q5.** On the full datasets the largest performance gap between FB and our conservative variants was on the lowest quality dataset (Random). We believe this corroborates our hypothesis that the Random-100k dataset is too uninformative to learn good models/polices from, and addressed your question around why relative performance was worst on the lowest quality dataset.

Reviewer comment (Writing):
The relationship between the task structure of the set of MDPs and the reward functions is unclear. The MDP framework outlined in section 2 paragraph 1 does not come with such structure and the structure is only mentioned but not well-defined in the problem formulation section. As a result, it is not clear where the differences in the sets of tasks come from. It might be beneficial to mention the distribution of task vectors earlier and how it relates to reward functions in paragraph 2 (Problem formulation.).

Author response:
**Q1: How is the task distribution $Z$ defined in practice? It is not quite clear to me, for instance, how I would define such an abstract distribution for either of the given environments in section 5 other than by possibly defining a distribution over final states which would have to be as large as the state space? Can you give an example of what this would look like in, e.g. the Cheetah environment?**  

We discuss the formulation of $\mathcal{Z}$ in Appendix B.1.2, but note, and apologise, that a reference to this appendix in the main body was missed. We address will this in Revision 4.3. We believe further misunderstanding was likely caused by our presentation of the FB theory which we address in Revision 2.

Reviewer comment (Novelty):
Although there is abundant literature on conservativeness in RL, the paragraph on this in the related works section only contains a total of 7 references. Contextualization could be stronger, possibly with less detail about the cited methods and more concise statements about previous literature’s commonalities to condense the text.

Author response:
**Q6: Contextualization with prior work: There is an abundance of literature on conservativeness in RL and the paragraph on this in the related works section contains a total of 7 references. This contextualization could be a little stronger. The downside of this lack of contextualization shows somewhat in the experiments as I will outline later. Less detail about the cited methods and concise statements about previous literature’s commonalities might provide a way to condense the text in this section.**  

We agree that the related work section on Offline RL could provide more detail, and as a result will add several more references. Namely: IQL [3], PEVI, [4], COMBO [5] and [6]

Reviewer comment (Evaluation):
The evaluation metrics in the experiments are not fully specified. The text mostly refers to performance without clarifying what is being measured, and in Figure 4, the meaning of a Task Score x is unclear.

Author response:
**Q7: The evaluation metrics are not fully specified. The text mostly talks about performance but it is not clear what is being measured. […] I think a clear definition of what’s being measured and how the measures are aggregated would be useful to provide context for the numbers and a more detailed description of when the method works well.**  

Agreed. See Revision 4.4.  

Finally, we thank the reviewer provided feedback on textual clarity, these will be incorporated in Revision 4.5.

Reviewer comment (Experiments):
The claim that conservativeness is not harmful is only supported on what is referred to as a high-quality dataset, in which the dataset already provides decent coverage and Q-values should be easy to approximate.

Author response:
**Q5: It is a little counterintuitive to me that conservativeness seems to work better when the datasets are of higher quality rather than when datasets provide poor support. Do you have any idea why that is?**  

As mentioned in W2 of the global response, we suspect that the combination of low diversity and small absolute size of the Random dataset makes model/policy inference particularly tricky. We hope the proposed experiments on the full Random and DIAYN datasets (Revision 1) will shed light on this hypothesis.

Reviewer comment (Experiments):
The differences between VC-FB and MC-FB are highlighted in section 3, but the effects of these differences are not explicitly analyzed in the experiments.

Author response:
b) The most recent experiments carried out for 2wo9 provide better evidence of our claims in Section 5 that MC-FB performs worse than VC-FB because of poorer task coverage. (This was the testable hypothesis you referred to from Q5 in our original response to 2wo9). We lay out the new experiment and results at the start of Section 5.

Reviewer comment (Writing):
Would be useful if you would properly define the abbreviations for VC-FB and MC-FB in the text before using them.

Author response:
**Q1: Would be useful if you would properly define the abbreviations for VC-FB and MC-FB in the text before using them.**

Agreed. See Revision 4.1.

Reviewer comment (Theory):
In equation 5, the \( z \) variable is not an argument to the Q function, as it should be.

Author response:
**Q2: In equation 5 the $z$ variable is not an argument to the Q function, as it should be?**

It is implicit via the superscript $\pi_z$, we took this notation directly from Touati et al (2022) but agree it could be clearer. We hope the rewrite of Section 2 proposed in Revision 2 will better clarify notation.

Reviewer comment (Theory):
Broken English in the line before equation 6. What is the connection between (5) and (6)? In (6) the forward representation here is a function of the same \( z \), i.e. \(\langle F(.., z), z \rangle\), however this is not how the Q value is defined in (5). Later you sample \( z \)'s independently from \( s^+ \), hence the backward representation is completely independent of the task \( z \) in the MC-FB. How is this valid?

Author response:
**Q4: Why don’t we use B(s_+) for z inside the forward model of MC-FB?**

FB’s approximation of the successor measure **$M^{\pi_z}(s_0, a_0, s_+) \approx F(s_0, a_0, z)^\top B(s_+)$** is the cumulative time spent in future state $s_+$ starting in state $s_0$, taking action $a_0$ and attempting to solve task $z$. Our goal was to mitigate OOD successor measure overestimation *for all tasks* and so sample $z \sim \mathcal{Z}$ for input to the forward model. Our language justifying MC-FB was therefore imprecise, particularly the phrase: *“Instead, it may prove better to direct updates towards tasks we are likely to encounter*” (Section 3, paragraph 4). We agree with the reviewer that the best way to direct updates would be to use $z = B(s_+)$ as input to the forward model and, although not mentioned by the reviewer, as inputs to the Q functions for policy training (Equation 7). We believe it would no longer make sense to call such a variant _measure-conservative_, but rather _directed-value-conservative_, with all $z$s _directed_ by the backward embedding rather than sampled uniformly. It is straightforward to explore the usefulness of such a change, and we are running experiments on all datasets and domains to do so—see Revision 3. We will update the language used to justify MC-FB—see Revision 4.9.

Reviewer comment (Writing):
In which cases does the \( z \) correspond to the backward representation?

Author response:
**Q3: In which cases does the $z$ correspond to the backward representation?**

As above, this question is valid, and a consequence of ambiguity in our introduction of the FB background theory. We hope to provide much better clarity via Revision 2.

Reviewer comment (Theory):
Follow up to the previous question, why does your MC-FB version treat the backward representation and \( z \) independently, while the VC-FB version does not? Shouldn't the \( z \) in the \( F(...) \) be essentially the output of \( B(s^+) \) as per your description?

Author response:
**Q4: Why don’t we use B(s_+) for z inside the forward model of MC-FB?**

FB’s approximation of the successor measure **$M^{\pi_z}(s_0, a_0, s_+) \approx F(s_0, a_0, z)^\top B(s_+)$** is the cumulative time spent in future state $s_+$ starting in state $s_0$, taking action $a_0$ and attempting to solve task $z$. Our goal was to mitigate OOD successor measure overestimation *for all tasks* and so sample $z \sim \mathcal{Z}$ for input to the forward model. Our language justifying MC-FB was therefore imprecise, particularly the phrase: *“Instead, it may prove better to direct updates towards tasks we are likely to encounter*” (Section 3, paragraph 4). We agree with the reviewer that the best way to direct updates would be to use $z = B(s_+)$ as input to the forward model and, although not mentioned by the reviewer, as inputs to the Q functions for policy training (Equation 7). We believe it would no longer make sense to call such a variant _measure-conservative_, but rather _directed-value-conservative_, with all $z$s _directed_ by the backward embedding rather than sampled uniformly. It is straightforward to explore the usefulness of such a change, and we are running experiments on all datasets and domains to do so—see Revision 3. We will update the language used to justify MC-FB—see Revision 4.9.

Reviewer comment (Writing):
In 4.3, what does "stochastically dominate" mean?

Author response:
**Q6: In 4.3, what does "stochastically dominate" mean?**

We take this term directly from Agarwal et. al (2021) where performance profiles are introduced. Formally, a random variable $M$ stochastically dominates $N$ if $P(M > x) \geq P(N > x) \; \forall x$. In practice the y-axis values of one algorithm’s performance profile are higher than another for all x-axis values. We will add a footnote on this, see Revision 4.2.

Reviewer comment (Evaluation):
Can you explain the failure cases of MC-FB in comparison to VC in figure 5? For the Walker2D environment, the MC variant completely fails for the RND dataset — I suspect that the reason is low task coverage. Also, in some cases there is a big gap between CQL and your method (random Jaco).

Author response:
We have now added experimental evaluation of *directed* VC-FB to our final revision as per proposed revision R3. We report the new Table 2 here for convenience.

| Dataset | Domain | Task | FB | DVC-FB | MC-FB | VC-FB |
| --- | --- | --- | --- | --- | --- | --- |
| All (100k) | all domains | all tasks | 99 | 108 | 136 | 148 |

To recap: VC-FB regularises OOD actions on $F(s, a, z)^\top z$, with $s \sim \mathcal{D}$, and $z \sim \mathcal{Z}$; MC-FB regularises OOD actions on $F(s, a, z)^\top B(s_+)$, with $(s, s_+) \sim \mathcal{D}$ and $z \sim \mathcal{Z}$; and *directed*-VC-FB regularises OOD actions on $F(s, a, B(s_+))^\top B(s_+)$, with $(s, s_+) \sim \mathcal{D}$.

As expected, if $|\mathcal{D}|$ is small, then deriving tasks from $B(s_+)$ with $s_+ \sim \mathcal{D}$ seems to provide poorer task coverage than $z \sim \mathcal{Z}$, and the extent to which $z$s are obtained in this way for each method correlates negatively with performance i.e. VC-FB > MC-FB > DVC-FB. We refer the reviewer to Appendix B.1.6 for implementation details, Appendix C for the full results and section 5 for further discussion. We hope this answers your original question, and sheds further light on discrepancies between VC-FB and MC-FB performance. 

Having responded to each of your concerns either experimentally or in writing, we hope you feel comfortable updating your score. Thanks very much again for your feedback.