Paper: Harmony World Models: Boosting Sample Efficiency for Model-based Reinforcement Learning (ICLR-2024)

Reviewer comment (Experiments):
There’s no good ablation experiments on different reward blending schemes. The chosen method for blending/weighting the loss terms looks reasonable, but not necessarily the best method or even the most reliable one. There needs to be some other weighting methods tested in the experiment sections. Some simple weighting methods should be compared, which include for example the one in [1] and some other straight-forward heuristic methods.

Author response:
**Q1**: Comparisons to different reward blending schemes

We want to clarify that **the major contribution of our paper is to reveal the overlooked potential of sample efficiency in MBRL through a multi-task perspective, not to propose a specific multi-task method**. Our experiments find that our straightforward method is dramatically effective and sufficient to address the problem revealed in our paper.

We have compared with different multi-task methods and discussed why these methods can hardly make more improvements in the problem of our paper. Please refer to $\underline{\text{Q2 of the global response}}$ for the detailed response.

Reviewer comment (Experiments):
It would also be helpful to show how sensitive the proposed method is. Does it require separate tuning for each task? If it needs to be tuned, how much effort is needed? Is it stable across random seeds?

Author response:
**Q2**: Sensitivity of HarmonyWM

**Our HarmonyWM does not add hyperparameters** to our base method and, therefore, does not need additional tuning for different tasks. The improvement is stable across random seeds, as shown in the confidence intervals reported in our figures.

Reviewer comment (Experiments):
More baselines and environments are needed. The proposed method is only compared against DreamerV2 in a small subset of tasks. What will happen if it is compared to MuZero, TD-MPC and SimPle etc as mentioned in the paper? And what will happen if the algorithm is tested in similar environments such as GO, those Atari environments or locomotion control tasks from image?

Author response:
**Q3**: Additional baselines and environments

Inevitably, the scale of our experiments is restricted due to our limited computational resources. We have conducted experiments on 11 (Metaworld, RLBench, DMCR) and 4 environments (DMC) in the main paper and appendix of the original submission. During rebuttal, we added 4 additional environments (Assembly in Metaworld and Qbert/Seaquest/Boxing in Atari). The number of experimental environments in our work has been on par with or larger than most recent publications on MBRL (please refer to $\underline{\text{response to Q3 of Reviewer 72ZS}}$).

Regarding compared baselines, since our HarmonyWM is based on DreamerV2, it is natural to compare it mainly with DreamerV2 to reveal the significance of harmonizing loss scales among world model learning tasks. We mention MuZero, TD-MPC and SimPLe in Figure 1 to show the high-level spectrum of world model learning. Some model-based methods are tailored to a specific area and therefore unfair to conduct head-to-head comparisons. 

Regarding the environments Reviewer pH7N proposed, it is beyond our computational resources to conduct experiments on GO, and also unfair since MuZero with MCTS is mainly tailored to board games and unable to handle continuous control, while both Dreamer and HarmonyWM are more general architectures. For the Atari domain, SimPLe is a weak baseline and is outperformed by DreamerV3, therefore we perform additional comparisons to DreamerV3. The corresponding results can be found in $\underline{\text{Fig. 23 of Appendix D.10}}$ of our latest revision. For locomotion control tasks from images, we have already presented comparisons to TD-MPC and DreamerV2 on the DMCR domain. Results can be found in $\underline{\text{Fig. 10(b) of main paper}}$ and $\underline{\text{Fig. 18 of Appendix D.5}}$.

Reviewer comment (Evaluation):
The performance of DreamerV2 seems too bad in some environments like Cheetah Run, Walker Run, which is hard to figure out why considering DreamerV2 was applied to far harder problems and succeeded and DreamerV1 has better performance with the same number of training samples. Why does it happen?

Author response:
**Q4**: DreamerV2 performance on DMC Remastered benchmark

Regarding the comment that DreamerV2 performs poorly in environments like Cheetah Run and Walker Run, we need to clarify that we use the **DMC Remastered** benchmark in our main paper, **which is different from and more challenging than the original DMC**, as it contains various random visual distractors, showcased in the $\underline{\text{left of Fig 8 of main paper}}$, and has proven difficult for model-based agents to learn on. DreamerV2 performs well on the original DMC benchmark, as shown in $\underline{\text{Fig. 17 of Appendix D.5}}$, but performs poorly on DMC Remastered as it suffers from the domination of reconstructing irrelevant visual distractors in world model learning.

Reviewer comment (Novelty):
The novelty of the proposed method is limited. The authors build the optimization process of world models as a multi-task optimization and then use uncertainty weighting to harmonize loss scales among tasks. More state-of-the-art methods, e.g., multi-objective optimization or multi-agent learning methods, can be considered and may further improve the performance.

Author response:
**Q1**: Limited novelty of the proposed method

We want to clarify that **the major contribution of our paper is to reveal the overlooked potential of sample efficiency in MBRL through a multi-task perspective**, and to propose a decent way to address this. We acknowledge that our method for balancing different tasks in world models is straightforward, but as pointed out by Reviewer pH7N, **this simplicity serves as an advantage, making it easy to reproduce and apply to existing algorithms**.

Regarding the comment that we use uncertainty weighting to harmonize loss scales, we would like to note that our method differs from uncertainty weighting in several aspects. Uncertainty weighting is derived from maximum likelihood estimation, while our motivation is to harmonize loss scales among tasks. Uncertainty weighting makes assumptions on the distributions behind losses, while our method removes these assumptions and makes it possible to also balance the KL loss. Moreover, we introduce a rectified version to avoid extremely large task weights. Please refer to the $\underline{\text{last paragraph of Sec 3 of main paper in original submission}}$ for a more detailed explanation on this.


**Q2**: Comparison to multi-task or multi-objective baselines

Our work is irrelevant to multi-agent learning methods, as the multi-task perspective is presented for understanding the world model learning part of MBRL, instead of the policy learning part. 

Regarding multi-task learning methods, we have conducted comparison experiments and detailed discussions. Please refer to $\underline{\text{Q2 of the global response}}$.

Reviewer comment (Evaluation):
In Fig.15 & 16, HarmonyWM \( w_d=1 \) learns faster than Harmony WM at the beginning of some tasks. Could the authors give some explanations about this?

Author response:
**Q4**: Ablation on adjusting $w_d$

We observe that HarmonyWM ($w_d=1$) learns faster than HarmonyWM at the beginning of the Hammer task in terms of success rate. We argue that this is because in this task, high success rates do not strictly correspond to high return. The hammer task can be hacked by pushing the nail using the arm instead of the hammer, which results in the task's success with a low reward. When measured in terms of episode returns, HarmonyWM performs comparably with HarmonyWM ($w_d=1$) at the beginning and slightly outperforms HarmonyWM ($w_d=1$) at the end of training.

Reviewer comment (Experiments):
The authors claim "the root cause as the disharmonious interference between two tasks in explicit world model learning: due to *overload of redundant observation signals* ...". Why does, for example, Denoised MDP, which identifies the information that can be safely discarded as noises, serve as a baseline in the experiments?

Author response:
**Q3**: Comparison to Denoised MDP

Although we share a similar point with Denoised MDP to enhance task-centric representations, our approach is entirely irrelevant to it and has the advantage of more simplicity. Our contributions regarding world model optimization are **orthogonal** to Denoised MDP regarding architecture and can be combined with it to further improve performance, which is left for future work.

Nevertheless, we present additional comparisons to Denoised MDP in $\underline{\text{Fig. 22 of Appendix D.9 of latest revision}}$, where Denoised MDP trails behind our HarmonyWM.  Denoised MDP performs information decomposition by changing the MDP transition structure and utilizing the reward as a guide to separate controllable information. However, since Denoised MDP does not modify the weight for the reward modeling task, the observation modeling task can still dominate the learning process. Consequently, the training signals from the reward modeling task may be inadequate to guide decomposition.

Reviewer comment (Experiments):
Is there a sufficient reason why the reward modeling and the observation modeling tasks are scaled to the same constant? Should we maintain the equilibrium between these tasks? Or can we emphasize some tasks at some specific states?

Author response:
**Q5**: Necessity of scaling different tasks to the same constant

As we demonstrated in our analysis ($\underline{\text{Sec 2.3 of main paper}}$), domination of either task in world model learning cannot fully exploit the potential of sample efficiency; thus, we do need to make a proper balance these two tasks in world model learning. We find it is dramatically effective to simply maintain the equilibrium between them, i.e., rescaling losses to the same constant. It is interesting to explore a more fine-grained balancing strategy, for example, as suggested by the reviewer, to emphasize some tasks a bit more at some specific states. This is left for future work.

Reviewer comment (Novelty):
Although the method proposed in this paper is novel, DreamerV3 has already addressed the issue of differing scales in reconstructing inputs and predicting rewards using Symlog Predictions. This paper primarily conducts extensive experiments in comparison with DreamerV2 and lacks more comparisons and discussions with DreamerV3. This significantly diminishes the contributions of this paper.

Author response:
**Q1**: Clarifications and comparison with DreamerV3


Reviewer 72ZS commented that DreamerV3 has already *addressed the issue of differing scales in reconstructing inputs and predicting rewards using Symlog Predictions*. We need to clarify that **using Symlog Predictions does not solve our problem of seeking a balance between reward modeling and observation modeling**. First of all, the Symlog transformation only shrinks extremely large values but is unable to rescale various values into exactly the same magnitude, while our harmonious loss properly addresses this by dynamically approximating the reciprocals of the values. More importantly, the primary reason why $L_r$ has a significantly smaller loss scale is the difference in dimension: as we have stated in our main paper, the observation loss $L_o$ usually aggregates $H\times W\times  C$ dimensions, while the reward loss $L_r$ is derived from only a scalar (a detailed explanation of why $L_o$ is summed over dimensions can be found in $\underline{\text{response to Q3.4 of Reviewer QvDA}}$ ). In summary, **using Symlog Predictions as DreamerV3 only mitigates the problem of differing per-dimension scales** (typically across environment domains) by a static transformation, **while our method aims to dynamically balance the overall loss scales across tasks in world model learning, considering together per-dimension scales, dimensions, and training dynamics**.

In practice, DreamerV3 does not use Symlog predictions in reconstructing visual inputs due to unified scales of visual inputs across domains. Besides, DreamerV3 uses twohot symlog predictions for the reward predictor to replace the MSE loss in DreamerV2. This approach increases the scale of the reward loss, but is insufficient to mitigate the domination of the observation loss. We observe that the reward loss in DreamerV3 is still significantly smaller than the observation loss, especially for visually demanding domains such as RLBench, where the reward loss is still two orders of magnitude smaller. 

Overall, **our contributions are orthogonal to the modifications in DreamerV3** and can be combined with DreamerV3 (which we brand as Harmony DreamerV3) to improve performance further. We have already presented the results of our Harmony DreamerV3 in $\underline{\text{Fig. 9 of main paper in the original submission}}$, and we provide additional results in $\underline{\text{Fig. 20 of Appendix D.7}}$ of our latest revision. We apologize for the limited discussion on DreamerV3, and we have added a discussion section in $\underline{\text{Appendix D.7 of our latest revision}}$.

Reviewer comment (Experiments):
Moreover, as TDMPC is mentioned in the paper as one of the Implicit MBRL methods, I believe it should also be considered as one of the important baselines for comparisons across different benchmarks.

Author response:
**Q2**: Comparison with TD-MPC

We want to highlight that we have already **presented aggregated results of TD-MPC on three Meta-world tasks** in $\underline{\text{Fig. 10(b) of main paper}}$. Moreover, we present **the individual learning curves of TD-MPC compared to our HarmonyWM and DreamerV2 on six Meta-world and DMCR tasks** in $\underline{\text{Fig. 18 of Appendix D.5 in the original submission}}$. These results show that TD-MPC is unable to learn a meaningful policy in some tasks, which further highlights the value of our explicit-implicit method.

Reviewer comment (Experiments):
The experimental environments chosen in the paper are all easy tasks from different domains. I think experiments should be conducted in more challenging environments, such as Hopper-hop and Quadruped-run in DMC, Assembly and Stick-pull in MetaWorld, etc. Moreover, given the current trend in the community to conduct large-scale experiments, having results from only eight environments seems somewhat limited.

Author response:
**Q3**: Difficulties and scales of experimental environments

We **respectfully disagree with the comment that our experimental environments are all easy tasks**. According to the classification standard of Seo et al. [1], ***Sweep-Into*, *Hammer* and *Push* can all be regarded as challenging environments** in the Meta-world domain. The RLBench domain itself is considered more demanding than Meta-world due to its realistic visual observations. It's also worth noting that we use **DMC Remastered, which is different from and more challenging than the original DMC**, as it contains various random visual distractors and has proven difficult for model-based agents to learn on [6]. More details about the environments we use can be found in $\underline{\text{Appendix C.1}}$, and we apologize for not referring to these in our main paper.

Regarding the challenging environments proposed by reviewer 72ZS, we have already reported our results on DMC Quadruped-run in $\underline{\text{Fig. 17 of Appendix D.5}}$, along with several other DMC environments. We apologize for not referring to them in our main paper. **The *Push* task reported in our main paper is of similar difficulty to *Assembly* and *Stick-pull*.** Nevertheless, we provide additional results on *Assembly* in $\underline{\text{Fig. 19 of Appendix D.6}}$ of our latest revision. 

Regarding the comment that our experiments lack variety and that we should perform large-scale experiments, we need to clarify that we perform **our experiments on fifteen instead of eight environments**: six from Meta-world, two from RLBench, three from DMC Remastered, and four from standard DMC (in Appendix). We add another 4 environments (Assembly in Metaworld and Qbert/Seaquest/Boxing in Atari) in our latest revision, which sums to a total of 19. Inevitably, the scale of our experiments is restricted due to our limited computational resources. However, this scale is very common in recent model-based RL publications, such as APV (9 envs) [2], DreamerPro (12 envs) [3], IsoDream (5 envs) [4], Curious Replay (5 envs) [5], IPV (12 envs) [6] and so forth.

[1] Seo et al. Masked world models for visual control. CoRL, 2022.

[2] Seo et al. Reinforcement learning with action-free pre-training from videos. ICML, 2022.

[3] Deng et al. Dreamerpro: Reconstruction-free model-based reinforcement learning with prototypical representations. ICML, 2022.

[4] Pan et al. Iso-Dream: Isolating and Leveraging Noncontrollable Visual Dynamics in World Models. NeurIPS, 2022.

[5] Kauvar et al. Curious Replay for Model-based Adaptation. ICML, 2023.

[6] Wu et al. Pre-training contextualized world models with in-the-wild videos for reinforcement learning. NeurIPS, 2023.

Please let us know whether our response addresses all your concerns. We are happy to discuss more if you have further questions or suggestions to improve our paper.

Reviewer comment (Evaluation):
The paper calls out the following as its contribution: "To the best of our knowledge, our work, for the first time, systematically identifies the multitask essence of world models and analyzes the disharmonious interference between different tasks, which is unexpectedly overlooked by most previous work". While I agree that MBRL can be formulated as a multi-task RL problem (it is clearly a multi-objective problem), the paper does not do a thorough job at analyzing "the disharmonious interference between different tasks". For example, they do not study if and why there is an interference between different tasks. In the multi-task literature, interference often refers to progress on one task hindering the progress of another task. What the authors show is that ad-hoc setting of scalars for the different losses can hurt the performance on the RL task, but they do not show that it hurts the performance on the tasks being directly optimized for. The distinction is important in the multi-task setup (which is what the paper uses).

Author response:
**Q1**: Clarification on the multi-task essence of world models

We apologize for using the misleading word 'interference' to describe the problem in our original submission. In our work, we conduct a dedicated investigation ($\underline{\text{Sec 2.3 of main paper}}$) and demonstrate that observation modeling and reward modeling can benefit from each other, as shown in $\underline{\text{Fig 2 of main paper}}$ and our additional quantitative results in $\underline{\text{Q1 of the global response}}$. More importantly, we also reveal that domination of either task, as in many existing MBRL literature, does not properly exploit these multi-task benefits. Thus, the **main problem our work addresses can be more appropriately named task 'domination'** instead of 'interference,' and we propose a simple yet effective strategy to dynamically balance tasks in world models. We have modified our wording in the text of our $\underline{\text{latest revision}}$, and thank you for pointing out this.

Please also refer to $\underline{\text{Q1 of the global response}}$ for the detailed response.