{"id":28366936,"url":"https://github.com/thuml/harmonydream","last_synced_at":"2025-07-25T17:41:05.813Z","repository":{"id":242630880,"uuid":"810085617","full_name":"thuml/HarmonyDream","owner":"thuml","description":"Code release for \"HarmonyDream: Task Harmonization Inside World Models\" (ICML 2024), https://arxiv.org/abs/2310.00344","archived":false,"fork":false,"pushed_at":"2024-06-27T05:30:00.000Z","size":2104,"stargazers_count":39,"open_issues_count":0,"forks_count":1,"subscribers_count":5,"default_branch":"main","last_synced_at":"2025-05-29T00:13:40.836Z","etag":null,"topics":["model-based-reinforcement-learning","world-models"],"latest_commit_sha":null,"homepage":"","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/thuml.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":null,"funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2024-06-04T02:45:20.000Z","updated_at":"2025-05-20T17:09:32.000Z","dependencies_parsed_at":"2024-06-27T06:52:58.482Z","dependency_job_id":null,"html_url":"https://github.com/thuml/HarmonyDream","commit_stats":null,"previous_names":["thuml/harmonydream"],"tags_count":0,"template":false,"template_full_name":null,"purl":"pkg:github/thuml/HarmonyDream","repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/thuml%2FHarmonyDream","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/thuml%2FHarmonyDream/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/thuml%2FHarmonyDream/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/thuml%2FHarmonyDream/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/thuml","download_url":"https://codeload.github.com/thuml/HarmonyDream/tar.gz/refs/heads/main","sbom_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/thuml%2FHarmonyDream/sbom","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":261012854,"owners_count":23096924,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["model-based-reinforcement-learning","world-models"],"created_at":"2025-05-29T00:13:37.134Z","updated_at":"2025-07-25T17:41:05.767Z","avatar_url":"https://github.com/thuml.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"# HarmonyDream: Task Harmonization Inside World Models (ICML 2024)\n\nOfficial implementation of [HarmonyDream: Task Harmonization Inside World Model](https://arxiv.org/abs/2310.00344) in PyTorch. Implementation of DreamerV2 and DreamerV3 in PyTorch are also included.\n\nFor **HarmonyDream in JAX and Atari 100K experiments**, please refer to the `dreamerv3-jax` folder and this [README](dreamerv3-jax/README.md).\n\n## Method\n\n**TL;DR**: We find that harmonizing domination among two tasks, observation and reward modeling, in world models can dramatically boost the sample-efficiency of model-based RL and present HarmonyDream to maintain a dynamic equilibrium between them.\n\n\u003cdiv align=\"center\"\u003e\u003cimg align=\"center\" src=\"assets/method.png\" width=\"70%\"\u003e\u003c/div\u003e\n\u003cbr/\u003e\n\u003cdiv align=\"center\"\u003e\u003cimg align=\"center\" src=\"assets/environments.png\" width=\"100%\"\u003e\u003c/div\u003e\n\n## Implementation of HarmonyDream\n\nThe implementation of HarmonyDream only involves a few lines of changes to the base method [here](wmlib-torch/wmlib/agents/dreamerv2.py). These changes can be easily applied to your own method, as listed below:\n1. Add three scalar variables `harmony_s1`, `harmony_s2`, and `harmony_s3`. (line 117-119)\n2. Multiply each loss by the corresponding coefficient and add the regularization term. (line 161-177)\n3. Backward and update both model parameters and si using the new losses. (line 56-57)\n\n\n## Dependencies\n\nThe conda environment can be created by the following command:\n\n```\ncd wmlib-torch\nconda create -n wmlib python=3.9 cudatoolkit=11.3\nconda activate wmlib\npip install wheel==0.38.4 setuptools==57.5.0\npip install -r requirements.txt -f https://download.pytorch.org/whl/cu113/torch_stable.html\n```\n\n**Meta-world**\n\nMeta-world depends on [MuJoCo200](https://www.roboti.us/download.html). You need to install it manually.\nMeta-world itself can be installed using the following command:\n\n```\npip install git+https://github.com/rlworkgroup/metaworld.git@a0009ed9a208ff9864a5c1368c04c273bb20dd06#egg=metaworld\n```\nIf an error (`ERROR: Failed building wheel for mujoco-py`) occurs during the installation, you can try down-grading cython and installing mujoco-py manually before reinstalling Meta-world:\n```\npip install cython==0.29.32\npip install cffi lockfile imageio\npip install mujoco-py==2.0.2.5\n```\n\n**RLBench**\n\nTo enable headless mode, you would need to install xorg\n```\nsudo apt-get install xorg libxcb-randr0-dev libxrender-dev libxkbcommon-dev libxkbcommon-x11-0 libavcodec-dev libavformat-dev libswscale-dev\n```\n\nInstallation and headless mode instructions of RLBench can be found in the [official README](https://github.com/stepjam/RLBench?tab=readme-ov-file#install).\nSince the RLBench environment does not provide dense rewards for the Push Button task, we have implemented a custom reward function in `wmlib-torch/wmlib/envs/rb/push_button.py` following MWM.\n\n**DMCR**\n\nWe adopt the original DMCR implementation provided by [QData](https://github.com/QData/dmc_remastered) and integrated it into our codebase. \nYou need to additionally download the assets of DMCR from [here](https://github.com/QData/dmc_remastered/tree/main/dmc_remastered/assets) and move them to the `wmlib-torch/wmlib/envs/dmcr/assets` folder.\n\n**Natural Background DMC**\n\nWe use the Natural Background DMC environment as in [RePo](https://github.com/zchuning/repo).\nTo train on natural background DMC, download the driving_car videos from Kinetics 400 dataset following these [instructions](https://github.com/Showmax/kinetics-downloader).\nAfterwards, modify the `resource_files` (line 58) in `wmlib-torch/wmlib/envs/natural_bg_dmc.py` to point to the downloaded videos.\n\n**MineDojo**\n\nPlease refer to the [official installation guide](https://github.com/MineDojo/MineDojo/tree/main?tab=readme-ov-file#installation) for the installation of MineDojo.\nWe have enabled the headless mode of MineDojo by setting `os.environ[\"MINEDOJO_HEADLESS\"] = \"1\"` in `wmlib-torch/wmlib/envs/minecraft.py`.\n\nThe MineDojo environment we use is a modified version of the original environment. There are three main differences:\n\n1. **Reward**: We use a manual dense reward function provided by [MineCLIP](https://github.com/MineDojo/MineCLIP/blob/e6c06a0245fac63dceb38bc9bd4fecd033dae735/mineclip/dense_reward/animal_zoo/dense_reward.py), which consists of a navigation reward and an attack reward.\n2. **Action Space**: Since the original action space of MineDojo is too large, we have reduced the action space to 5 discrete actions, which is modified from the pruned action space of [STG-Transformer](https://github.com/zhoubohan0/STG-Transformer/blob/b57312258ee1455d8b39a2f89d092345367f2aec/minecraft/envs/minecraft_nav.py#L65).\nThe action space is defined as follows:\n\n| Index | Description | Details | Num Actions |\n| --- | --- | --- | --- |\n| 0 | Forward and backward | 0:noop, 1:forward, 2:backward | 3 |\n| 1 | Move left and right | 0:noop, 1:left, 2:right | 3 |\n| 2 | Jump, sneak, and sprint | 0:noop, 1:jump, 2:sneak, 3:sprint | 4 |\n| 3 | Camera delta pitch and yaw | 0:noop, 1:pitch -15degree, 2:pitch +15degree, 3:yaw -15degree, 4:yaw +15degree | 5 |\n| 4 | Functional actions | 0:noop, 1:use, 2:attack | 3 |\n3. **Reset Mode**: By default, MineDojo uses a _fast reset_, i.e. resets the environment using MC native command /kill, instead of waiting for re-generating new worlds. Fast reset does not change the terrain of the environment, which results in a loss of diversity.\nRe-generating new worlds on every reset, on the other hand, is too slow for training. As a solution, we still use the fast reset method to reset the environment, but we also re-generate a new world every five episodes. This way, we can maintain the diversity of the environment while keeping the training speed.\n\nDuring our experiments, we found that when running multiple instances of MineDojo in parallel, the environment may fail to reset due to port conflicts. To solve this issue, we let the environment randomly sleep for a short time before retrying to reset if a reset fails.\nFor further details, please refer to the `reset` method in `wmlib-torch/wmlib/envs/minecraft.py`.\n\n## Experiments\n\n### HarmonyDream\n\nChange directory to the `wmlib-torch` folder to train HarmonyDream on different environments.\n\n**Meta-world**\n\n```\npython examples/train_dreamerv2.py --logdir {save path} --configs metaworld --task metaworld_lever_pull --seed 0 --harmony True\n```\n\n**RLBench**\n\n```\npython examples/train_dreamerv2.py --logdir {save path} --configs rlbench plainresnet --task rlbench_reach_target --seed 0 --harmony True\n```\n\n**DMC Remastered**\n\n```\npython examples/train_dreamerv2.py --logdir {save path} --configs dmc_remastered --task dmcr_cheetah_run --seed 0 --harmony True\n```\n\n**Natural Background DMC**\n\n```\npython examples/train_dreamerv2.py --logdir {save path} --configs dmc_natural_background --task dmcnbg_cheetah_run --seed 0 --harmony True\n```\n\n**Minecraft**\n\nExperiments on Minecraft are run using HarmonyDream applied to DreamerV3.\n```\npython examples/train_dreamerv2.py --logdir {save path} --configs minecraft dreamerv3 large eval_hard --task mincraft_hunt_cow --seed 0 --harmony True\n```\n\n### DreamerV2 and DreamerV3 Baselines\n\nOur baseline methods [DreamerV2](https://github.com/danijar/dreamerv2) and [DreamerV3](https://github.com/danijar/dreamerv3) can be easily run using the following commands.\n\n**DreamerV2**\n\n```\npython examples/train_dreamerv2.py --logdir {save path} --configs metaworld --task metaworld_lever_pull --seed 0\n```\n\n**DreamerV3**\n\nDreamerV3 is implemented as an addon to DreamerV2, therefore we use the same entry script as DreamerV2, but with different configurations.\n```\npython examples/train_dreamerv2.py --logdir {save path} --configs minecraft dreamerv3 large eval_hard --task minecraft_hunt_cow --seed 0\n```\n\n## Tips\n\nMixed precision are enabled by default, which is faster but can probably cause numerical instabilities. It is normal to encounter infinite gradient norms, and the training may be interrupted by `nan` values. You can pass `--precision 32` to disable mixed precision.\n\nSee also the tips available in [DreamerV2 repository](https://github.com/danijar/dreamerv2).\n\n## Citation\nIf you find our codebase useful for your research, please cite our paper as:\n\n```\n@inproceedings{ma2024harmonydream,\n  title={HarmonyDream: Task Harmonization Inside World Models}, \n  author={Haoyu Ma and Jialong Wu and Ningya Feng and Chenjun Xiao and Dong Li and Jianye Hao and Jianmin Wang and Mingsheng Long},\n  booktitle={International Conference on Machine Learning},\n  year={2024}\n}\n```\n\n## Contact\n\nIf you have any questions, please contact mhy22@mails.tsinghua.edu.cn (Haoyu Ma) or wujialong0229@gmail.com (Jialong Wu).\n\n## Acknowledgement\n\nWe sincerely appreciate the following github repos for their valuable code base we build upon:\n\n- https://github.com/danijar/dreamerv2\n- https://github.com/danijar/dreamerv3\n- https://github.com/younggyoseo/apv\n- https://github.com/esteveste/dreamerV2-pytorch\n- https://github.com/NM512/dreamerv3-torch\n- https://github.com/QData/dmc_remastered\n- https://github.com/zchuning/repo\n\nThe `wmlib-torch` folder is an addon to the https://github.com/thuml/ContextWM codebase.\n\nThe `dreamerv3-jax` folder is an addon to the https://github.com/danijar/dreamerv3 codebase.\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fthuml%2Fharmonydream","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Fthuml%2Fharmonydream","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Fthuml%2Fharmonydream/lists"}