HarmonyDream: Task Harmonization Inside World Models

TL;DR

HarmonyDream investigates the dual-task nature of world models in model-based RL, identifying that domination by either observation or reward modeling can hamper sample efficiency. It proposes a lightweight harmonization mechanism, using learnable loss scales with a rectified harmonic loss, to dynamically balance the two tasks during world-model learning. Empirically, HarmonyDream improves DreamerV2 by 10–69% on visual robotic tasks and achieves a new state-of-the-art 136.5% mean human performance on Atari 100K, while also generalizing to DreamerV3 and DreamerPro. The work advances understanding of the multi-task essence of world models and offers a simple, robust method to exploit it for improved sample efficiency across diverse domains.

Abstract

Model-based reinforcement learning (MBRL) holds the promise of sample-efficient learning by utilizing a world model, which models how the environment works and typically encompasses components for two tasks: observation modeling and reward modeling. In this paper, through a dedicated empirical investigation, we gain a deeper understanding of the role each task plays in world models and uncover the overlooked potential of sample-efficient MBRL by mitigating the domination of either observation or reward modeling. Our key insight is that while prevalent approaches of explicit MBRL attempt to restore abundant details of the environment via observation models, it is difficult due to the environment's complexity and limited model capacity. On the other hand, reward models, while dominating implicit MBRL and adept at learning compact task-centric dynamics, are inadequate for sample-efficient learning without richer learning signals. Motivated by these insights and discoveries, we propose a simple yet effective approach, HarmonyDream, which automatically adjusts loss coefficients to maintain task harmonization, i.e. a dynamic equilibrium between the two tasks in world model learning. Our experiments show that the base MBRL method equipped with HarmonyDream gains 10%-69% absolute performance boosts on visual robotic tasks and sets a new state-of-the-art result on the Atari 100K benchmark. Code is available at https://github.com/thuml/HarmonyDream.

Figures (27)

• Figure 1: A multi-task view of world models. (Left) World models typically consist of components for two tasks: observation modeling and reward modeling. (Right) A spectrum of world model learning in MBRL. Explicit MBRL learns models dominated by observation modeling, while implicit MBRL relies solely on reward modeling. In the spirit of implicit MBRL, our proposed HarmonyDream enables explicit MBRL to maintain a dynamic equilibrium between them to unleash the multi-task benefits of world model learning, thus boosting the sample efficiency of MBRL.
• Figure 2: Analysis experiments revealing the (b) imbalanced nature of world model learning and potential multi-task benefits yet to be properly exploited. Simply adjusting the coefficient of reward loss leads to (a) dramatically boosted sample efficiency of DreamerV2 agents and (c) better representations with lower environment state regression errors.
• Figure 3: Analysis of world models learned with different reward loss coefficients. Rewards are labeled at the bottom right corner, with predictions marked as correct or incorrect. Dominating observation modeling in world models incurs spurious correlations between actions, observations, and rewards, which can be dissolved by properly emphasizing reward modeling.
• Figure 4: Overview of HarmonyDream. (Left) Built upon Dreamer, we introduce lightweight harmonizers to maintain a dynamic equilibrium between tasks. (Right) Comparison between the original harmonious loss (Eq. (\ref{['eq:harmony_loss']})) and the rectified one (Eq. (\ref{['eq:harmony_loss_with_base']})). The latter prevents an extremely large loss weight.
• Figure 5: Visual control domains for evaluation, including robotic manipulation (a-b), distracted locomotion (c), and video games (d-e).

Theorems & Definitions (2)

• Definition 2.1
• Proposition 3.1