File size: 6,365 Bytes
693a539 5cdbdc4 693a539 0c45551 693a539 08ecfef 693a539 25d5319 693a539 25d5319 693a539 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 | ---
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) |