Diffusion World Model: Future Modeling Beyond Step-by-Step Rollout for Offline Reinforcement Learning

About

We introduce Diffusion World Model (DWM), a conditional diffusion model capable of predicting multistep future states and rewards concurrently. As opposed to traditional one-step dynamics models, DWM offers long-horizon predictions in a single forward pass, eliminating the need for recursive queries. We integrate DWM into model-based value estimation, where the short-term return is simulated by future trajectories sampled from DWM. In the context of offline reinforcement learning, DWM can be viewed as a conservative value regularization through generative modeling. Alternatively, it can be seen as a data source that enables offline Q-learning with synthetic data. Our experiments on the D4RL dataset confirm the robustness of DWM to long-horizon simulation. In terms of absolute performance, DWM significantly surpasses one-step dynamics models with a $44\%$ performance gain, and is comparable to or slightly surpassing their model-free counterparts.

Zihan Ding, Amy Zhang, Yuandong Tian, Qinqing Zheng · 2024

Related benchmarks

Task	Dataset	Result
Offline Reinforcement Learning	D4RL Hopper-M	Normalized Return65	11
Offline Reinforcement Learning	D4RL Walker2d-ME	Normalized Return1.1	11
Offline Reinforcement Learning	D4RL Hopper-ME	Normalized Return1.03	11
Offline Reinforcement Learning	D4RL HalfCheetah-MR	Normalized Return43	11
Offline Reinforcement Learning	D4RL HalfCheetah-ME	Normalized Return75	11
Offline Reinforcement Learning	D4RL HalfCheetah-M	Normalized Return46	11
Offline Reinforcement Learning	D4RL Hopper-MR	Normalized Return0.61	11
Offline Reinforcement Learning	D4RL Walker2d-M	Normalized Return76	11
Offline Reinforcement Learning	D4RL Walker2d-MR	Normalized Return46	11