BruceYuan commited on
Commit
693a539
·
verified ·
1 Parent(s): d4872fa

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +215 -3
README.md CHANGED
@@ -1,3 +1,215 @@
1
- ---
2
- license: mit
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ license: mit
3
+ language:
4
+ - en
5
+ library_name: pytorch
6
+ tags:
7
+ - reinforcement learning
8
+ - model based reinforcement learning
9
+ - world model
10
+ - diffusion policy
11
+ - continuous control
12
+ - robotics
13
+ - offline reinforcement learning
14
+ - offline to online reinforcement learning
15
+ pipeline_tag: reinforcement-learning
16
+ datasets:
17
+ - nicklashansen/tdmpc2
18
+ ---
19
+
20
+ # Model Card for MBDPO
21
+
22
+ Official release of MBDPO model checkpoints for the paper
23
+
24
+ [Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization](http://arxiv.org/abs/2605.26282) by
25
+
26
+ Xiaoyuan Cheng\*, Wenxuan Yuan\*, Zhancun Mu, Yuanzhao Zhang, Yiming Yang, Hai Wang, Zhuo Sun<sup>†</sup>, and Che Liu<sup>†</sup>.
27
+
28
+ **Quick links:** [[Paper]](http://arxiv.org/abs/2605.26282) [[Code]](https://github.com/Edmond1Cheng/MBDPO)
29
+
30
+ > **Note:** This repository currently contains the initially uploaded offline checkpoints. More checkpoints and related files are still being updated.
31
+
32
+ ## Model Details
33
+
34
+ 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.
35
+
36
+ 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.
37
+
38
+ ### Model Description
39
+
40
+ - **Developed by:** Xiaoyuan Cheng, Wenxuan Yuan, Zhancun Mu, Yuanzhao Zhang, Yiming Yang, Hai Wang, Zhuo Sun, and Che Liu
41
+ - **Model type:** Model-based reinforcement learning checkpoints with diffusion policy optimization
42
+ - **Framework:** PyTorch
43
+ - **Task type:** Continuous control
44
+ - **License:** MIT
45
+
46
+ ### Model Sources
47
+
48
+ - **Repository:** [https://github.com/Edmond1Cheng/MBDPO](https://github.com/Edmond1Cheng/MBDPO)
49
+ - **Paper:** [http://arxiv.org/abs/2605.26282](http://arxiv.org/abs/2605.26282)
50
+
51
+ ## Uses
52
+
53
+ 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.
54
+
55
+ 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.
56
+
57
+ ### Direct Use
58
+
59
+ Model checkpoints can be loaded and evaluated using the [official implementation](https://github.com/Edmond1Cheng/MBDPO).
60
+
61
+ Example evaluation command:
62
+
63
+ ```bash
64
+ python scripts/evaluate.py \
65
+ task=mt80 \
66
+ checkpoint=/path/to/checkpoint.pt \
67
+ eval_episodes=10
68
+ ````
69
+
70
+ ### Out-of-Scope Use
71
+
72
+ 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.
73
+
74
+ We do not expect checkpoints to generalize reliably to unseen tasks or substantially different environments without fine-tuning or further training.
75
+
76
+ ## How to Get Started with the Models
77
+
78
+ Please first install the official implementation:
79
+
80
+ ```bash
81
+ git clone https://github.com/Edmond1Cheng/MBDPO.git
82
+ cd MBDPO
83
+ ```
84
+
85
+ Create the corresponding Conda environment. For example, for MT80 experiments:
86
+
87
+ ```bash
88
+ conda env create -f conda_envs/mbdpo-mt80.yml
89
+ conda activate mbdpo-mt80
90
+ ```
91
+
92
+ Other environment files are also provided for different experiment suites, such as ManiSkill2 and MyoSuite.
93
+
94
+ After downloading a checkpoint from this repository, run evaluation with:
95
+
96
+ ```bash
97
+ python scripts/evaluate.py \
98
+ task=mt80 \
99
+ checkpoint=/path/to/checkpoint.pt \
100
+ eval_episodes=10
101
+ ```
102
+
103
+ For offline-to-online fine-tuning:
104
+
105
+ ```bash
106
+ python scripts/offline_to_online.py \
107
+ checkpoint=/path/to/checkpoint.pt \
108
+ save_path=/path/to/output_dir \
109
+ off2on_task="walker-run" \
110
+ steps=40000
111
+ ```
112
+
113
+ Please refer to the official repository for detailed installation instructions, configuration files, and experiment scripts.
114
+
115
+ ## Training Details
116
+
117
+ MBDPO supports three main experimental settings:
118
+
119
+ 1. Online training from scratch
120
+ 2. Multi-task offline pretraining
121
+ 3. Offline-to-online fine-tuning
122
+
123
+ ### Training Data
124
+
125
+ For multi-task offline pretraining, MBDPO uses replay buffer data from the open-sourced TD-MPC2 dataset:
126
+
127
+ * [nicklashansen/tdmpc2](https://huggingface.co/datasets/nicklashansen/tdmpc2)
128
+
129
+ The relevant subsets include `mt30` and `mt80`.
130
+
131
+ ### Supported Tasks
132
+
133
+ MBDPO supports 121 continuous control tasks across the following domains:
134
+
135
+ | Domain | Number of Tasks |
136
+ | ---------- | --------------: |
137
+ | DMControl | 39 |
138
+ | MetaWorld | 50 |
139
+ | ManiSkill2 | 5 |
140
+ | MyoSuite | 10 |
141
+ | Locomotion | 7 |
142
+ | Visual RL | 10 |
143
+ | **Total** | **121** |
144
+
145
+ In the DMControl domain, MBDPO follows the TD-MPC2 setting and includes 11 custom tasks.
146
+
147
+ ### Example Training Commands
148
+
149
+ Single-task online training from scratch:
150
+
151
+ ```bash
152
+ python scripts/train.py task=dog-run seed=1 steps=4000000
153
+ ```
154
+
155
+ Multi-task offline pretraining:
156
+
157
+ ```bash
158
+ python scripts/train.py task=mt80 multitask=true
159
+ ```
160
+
161
+ or:
162
+
163
+ ```bash
164
+ python scripts/train.py task=mt30 multitask=true
165
+ ```
166
+
167
+ Parallel online training:
168
+
169
+ ```bash
170
+ python scripts/online_parallel_train.py --config cfgs/online_parallel_config.yaml
171
+ ```
172
+
173
+ ## Checkpoint Status
174
+
175
+ This Hugging Face repository currently contains the initially uploaded **offline checkpoints**.
176
+
177
+ Additional checkpoints and files are still being updated. Please check the repository file list for the latest available models.
178
+
179
+ ## Limitations
180
+
181
+ * Checkpoints are trained and evaluated in simulated environments.
182
+ * Performance may depend on random seeds, simulator versions, hardware, and configuration details.
183
+ * Offline checkpoints may require compatible task definitions and configuration files.
184
+ * Generalization to unseen tasks is not guaranteed.
185
+ * Real-world deployment is out of scope without additional safety testing and validation.
186
+ * This repository may not yet contain all checkpoints used in the paper.
187
+
188
+ ## Citation
189
+
190
+ If you find our work useful, please consider citing the paper as follows:
191
+
192
+ **BibTeX:**
193
+
194
+ ```bibtex
195
+ @misc{cheng2026scalingworldmodelreinforcementlearning,
196
+ title={Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization},
197
+ author={Xiaoyuan Cheng and Wenxuan Yuan and Zhancun Mu and Yuanzhao Zhang and Yiming Yang and Hai Wang and Zhuo Sun and Che Liu},
198
+ year={2026},
199
+ eprint={2605.26282},
200
+ archivePrefix={arXiv},
201
+ primaryClass={cs.LG},
202
+ url={http://arxiv.org/abs/2605.26282}
203
+ }
204
+ ```
205
+
206
+ ## Contact
207
+
208
+ For questions about the paper, please contact:
209
+
210
+ * Xiaoyuan Cheng: [ucesxc4@ucl.ac.uk](mailto:ucesxc4@ucl.ac.uk)
211
+ * Wenxuan Yuan: [YUAN0186@e.ntu.edu.sg](mailto:YUAN0186@e.ntu.edu.sg)
212
+
213
+ For bugs, feature requests, or contributions, please open an issue or pull request in the official GitHub repository:
214
+
215
+ [https://github.com/Edmond1Cheng/MBDPO](https://github.com/Edmond1Cheng/MBDPO)