| --- |
| license: mit |
| language: |
| - en |
| library_name: pytorch |
| tags: |
| - model based reinforcement learning |
| - world model |
| - diffusion policy |
| - continuous control |
| - robotics |
| - offline reinforcement learning |
| - offline to online reinforcement learning |
| pipeline_tag: reinforcement-learning |
| datasets: |
| - nicklashansen/tdmpc2 |
| --- |
| |
| # Model Card for MBDPO |
|
|
| Official release of MBDPO model checkpoints for the paper |
|
|
| [Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization](http://arxiv.org/abs/2605.26282) by |
|
|
| Xiaoyuan Cheng\*, Wenxuan Yuan\*, Zhancun Mu, Yuanzhao Zhang, Yiming Yang, Hai Wang, Zhuo Sun<sup>†</sup>, and Che Liu<sup>†</sup>. |
|
|
| **Quick links:** [[Website]](https://wenxuan52.github.io/mbdpo-page/) [[Paper]](http://arxiv.org/abs/2605.26282) [[Code]](https://github.com/Edmond1Cheng/MBDPO) |
|
|
| > **Note:** Due to limited data storage capacity, data loss occurred when saving ckpt. Therefore, the current repository only provides offline trained ckpt and offline-to-online pretrained ckpt. |
|
|
| ## Model Details |
|
|
| We release model checkpoints for **MBDPO**, a model-based reinforcement learning framework that unifies search and policy optimization through diffusion policy optimization inside a learned latent world model. |
|
|
| MBDPO reformulates policy optimization as a diffusion process over imagined trajectories. The diffusion score field is corrected by model-based returns and anchored to the behavior distribution through an implicit energy function. This design removes the need for an explicit planner on top of the world model and addresses the structural mismatch between search and value learning in prior world-model reinforcement learning methods. |
|
|
| ### Model Description |
|
|
| - **Developed by:** Xiaoyuan Cheng\*, Wenxuan Yuan\*, Zhancun Mu, Yuanzhao Zhang, Yiming Yang, Hai Wang, Zhuo Sun, and Che Liu |
| - **Model type:** Model-based reinforcement learning checkpoints with diffusion policy optimization |
| - **Framework:** PyTorch |
| - **Task type:** Continuous control |
| - **License:** MIT |
|
|
| ### Model Sources |
|
|
| - **Repository:** [https://github.com/Edmond1Cheng/MBDPO](https://github.com/Edmond1Cheng/MBDPO) |
| - **Paper:** [http://arxiv.org/abs/2605.26282](http://arxiv.org/abs/2605.26282) |
|
|
| ## Uses |
|
|
| These checkpoints are intended for researchers interested in model-based reinforcement learning, world models, diffusion policies, offline reinforcement learning, and offline-to-online fine-tuning. |
|
|
| They can be used for reproducing MBDPO results, evaluating pretrained agents, analyzing learned world models and policies, and initializing offline-to-online fine-tuning experiments. |
|
|
| ### Direct Use |
|
|
| Model checkpoints can be loaded and evaluated using the [official implementation](https://github.com/Edmond1Cheng/MBDPO). |
|
|
| Example evaluation command: |
|
|
| ```bash |
| python scripts/evaluate.py \ |
| task=mt80 \ |
| checkpoint=/path/to/checkpoint.pt \ |
| eval_episodes=10 |
| ```` |
|
|
| ### Out-of-Scope Use |
|
|
| These checkpoints are research artifacts trained and evaluated in simulated continuous control environments. They are not intended for direct deployment in real-world robotics systems or safety-critical applications without additional validation. |
|
|
| We do not expect checkpoints to generalize reliably to unseen tasks or substantially different environments without fine-tuning or further training. |
|
|
| ## How to Get Started with the Models |
|
|
| Please first install the official implementation: |
|
|
| ```bash |
| git clone https://github.com/Edmond1Cheng/MBDPO.git |
| cd MBDPO |
| ``` |
|
|
| Create the corresponding Conda environment. For example, for MT80 experiments: |
|
|
| ```bash |
| conda env create -f conda_envs/mbdpo-mt80.yml |
| conda activate mbdpo-mt80 |
| ``` |
|
|
| Other environment files are also provided for different experiment suites, such as ManiSkill2 and MyoSuite. |
|
|
| After downloading a checkpoint from this repository, run evaluation with: |
|
|
| ```bash |
| python scripts/evaluate.py \ |
| task=mt80 \ |
| checkpoint=/path/to/checkpoint.pt \ |
| eval_episodes=10 |
| ``` |
|
|
| For offline-to-online fine-tuning: |
|
|
| ```bash |
| python scripts/offline_to_online.py \ |
| checkpoint=/path/to/checkpoint.pt \ |
| save_path=/path/to/output_dir \ |
| off2on_task="walker-run" \ |
| steps=40000 |
| ``` |
|
|
| Please refer to the official repository for detailed installation instructions, configuration files, and experiment scripts. |
|
|
| ## Training Details |
|
|
| MBDPO supports three main experimental settings: |
|
|
| 1. Online training from scratch |
| 2. Multi-task offline pretraining |
| 3. Offline-to-online fine-tuning |
|
|
| ### Training Data |
|
|
| For multi-task offline pretraining, MBDPO uses replay buffer data from the open-sourced TD-MPC2 dataset: |
|
|
| * [nicklashansen/tdmpc2](https://huggingface.co/datasets/nicklashansen/tdmpc2) |
|
|
| The relevant subsets include `mt30` and `mt80`. |
|
|
| ### Supported Tasks |
|
|
| MBDPO supports 121 continuous control tasks across the following domains: |
|
|
| | Domain | Number of Tasks | |
| | ---------- | --------------: | |
| | DMControl | 39 | |
| | MetaWorld | 50 | |
| | ManiSkill2 | 5 | |
| | MyoSuite | 10 | |
| | Locomotion | 7 | |
| | Visual RL | 10 | |
| | **Total** | **121** | |
|
|
| In the DMControl domain, MBDPO follows the TD-MPC2 setting and includes 11 custom tasks. |
|
|
| <!-- |
| ## Checkpoint Status |
|
|
| This Hugging Face repository currently contains the initially uploaded **offline checkpoints**. |
|
|
| Additional checkpoints and files are still being updated. Please check the repository file list for the latest available models. --> |
|
|
| ## Citation |
|
|
| If you find our work useful, please consider citing the paper as follows: |
|
|
| **BibTeX:** |
|
|
| ```bibtex |
| @misc{cheng2026scalingworldmodelreinforcementlearning, |
| title={Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization}, |
| author={Xiaoyuan Cheng and Wenxuan Yuan and Zhancun Mu and Yuanzhao Zhang and Yiming Yang and Hai Wang and Zhuo Sun and Che Liu}, |
| year={2026}, |
| eprint={2605.26282}, |
| archivePrefix={arXiv}, |
| primaryClass={cs.LG}, |
| url={http://arxiv.org/abs/2605.26282} |
| } |
| ``` |
|
|
| ## Contact |
|
|
| For questions about the paper, please contact: |
|
|
| * Xiaoyuan Cheng: [ucesxc4@ucl.ac.uk](mailto:ucesxc4@ucl.ac.uk) |
| * Wenxuan Yuan: [YUAN0186@e.ntu.edu.sg](mailto:YUAN0186@e.ntu.edu.sg) |
|
|
| For bugs, feature requests, or contributions, please open an issue or pull request in the official GitHub repository: |
|
|
| [https://github.com/Edmond1Cheng/MBDPO](https://github.com/Edmond1Cheng/MBDPO) |
