Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization

• URL: http://arxiv.org/abs/2605.26282v1
• Date: Mon, 25 May 2026 19:06:51 GMT
• Title: Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization
• Abstract summary: Model-based reinforcement learning can be effectively supported at scale through the use of world models.<n>We identify a bottleneck: a structural misalignment between search and value learning in existing world model approaches.<n>We propose Model-Based Diffusion Policy Optimization (MBDPO) in world models, a framework that unifies search and policy optimization through diffusion policy representations.
• Score: 46.32504081845328
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: Model-based reinforcement learning (RL) can be effectively supported at scale through the use of world models. However, in practice, scaling such approaches remains fundamentally limited. A commonly recognized challenge is model bias and error compounding, which degrade long-horizon predictions. Beyond these issues, we identify a more critical yet underexplored bottleneck: a structural misalignment between search and value learning in existing world model approaches. In particular, policy improvement often relies on value functions induced by a separate, non-search policy, resulting in training inconsistency and ultimately suboptimal learning. To address this limitation, we propose Model-Based Diffusion Policy Optimization (MBDPO) in world models, a framework that unifies search and policy optimization through diffusion policy representations, thereby unlocking the potential of world models for scalable policy learning. Instead of constructing an explicit planner over a learned world model, we reformulate policy optimization as a diffusion process over searched trajectories in latent world models. In this view, we extract an implicit energy function from the collected dataset that anchors the policy, enabling MBDPO to refine the score field for policy optimization while mitigating misalignment. We evaluate MBDPO across a wide range of settings, including multi-task offline pretraining, online learning, and offline-to-online fine-tuning. In the offline regime, we further investigate its scaling behavior by pretraining on large-scale datasets, observing consistent and monotonic performance gains with increasing model capacity.

Related papers

• Model Predictive Control with Differentiable World Models for Offline Reinforcement Learning [17.293286354647716]
  We introduce an inference time adaptation framework inspired by model predictive control (MPC)<n>We show that exploiting inference-time information to optimize the policy parameters yields consistent gains over strong offline RL baselines.
  arXiv  Detail & Related papers  (2026-03-23T18:05:29Z)
• Bootstrap Off-policy with World Model [59.129118672069644]
  We propose BOOM, a framework that tightly integrates planning and off-policy learning through a bootstrap loop.<n>BOOM achieves state-of-the-art results in both training stability and final performance.
  arXiv  Detail & Related papers  (2025-11-01T06:33:04Z)
• Policy-Driven World Model Adaptation for Robust Offline Model-based Reinforcement Learning [6.189693079685375]
  offline model-based RL (MBRL) explicitly learns a world model from a static dataset.<n>We propose a framework that dynamically adapts the world model alongside the policy.<n>We benchmark our algorithm on twelve noisy D4RL MuJoCo tasks and three Tokamak Control tasks, demonstrating its state-of-the-art performance.
  arXiv  Detail & Related papers  (2025-05-19T20:14:33Z)
• Model-Based Offline Reinforcement Learning with Adversarial Data Augmentation [36.9134885948595]
  We introduce Model-based Offline Reinforcement learning with AdversariaL data augmentation.<n>In MORAL, we replace the fixed horizon rollout by employing adversaria data augmentation to execute alternating sampling with ensemble models.<n>Experiments on D4RL benchmark demonstrate that MORAL outperforms other model-based offline RL methods in terms of policy learning and sample efficiency.
  arXiv  Detail & Related papers  (2025-03-26T07:24:34Z)
• Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models [79.2162092822111]
  We systematically evaluate reinforcement learning (RL) and control-based methods on a suite of navigation tasks.<n>We employ a latent dynamics model using the Joint Embedding Predictive Architecture (JEPA) and employ it for planning.<n>Our results show that model-free RL benefits most from large amounts of high-quality data, whereas model-based planning generalizes better to unseen layouts.
  arXiv  Detail & Related papers  (2025-02-20T18:39:41Z)
• Learning from Random Demonstrations: Offline Reinforcement Learning with Importance-Sampled Diffusion Models [19.05224410249602]
  We propose a novel approach for offline reinforcement learning with closed-loop policy evaluation and world-model adaptation. We analyzed the performance of the proposed method and provided an upper bound on the return gap between our method and the real environment under an optimal policy.
  arXiv  Detail & Related papers  (2024-05-30T09:34:31Z)
• World Models via Policy-Guided Trajectory Diffusion [21.89154719069519]
  Existing world models are autoregressive in that they interleave predicting the next state with sampling the next action from the policy. We propose a novel world modelling approach that is not autoregressive and generates entire on-policy trajectories in a single pass through a diffusion model.
  arXiv  Detail & Related papers  (2023-12-13T21:46:09Z)
• Reparameterized Policy Learning for Multimodal Trajectory Optimization [61.13228961771765]
  We investigate the challenge of parametrizing policies for reinforcement learning in high-dimensional continuous action spaces. We propose a principled framework that models the continuous RL policy as a generative model of optimal trajectories. We present a practical model-based RL method, which leverages the multimodal policy parameterization and learned world model.
  arXiv  Detail & Related papers  (2023-07-20T09:05:46Z)
• A Unified Framework for Alternating Offline Model Training and Policy Learning [62.19209005400561]
  In offline model-based reinforcement learning, we learn a dynamic model from historically collected data, and utilize the learned model and fixed datasets for policy learning. We develop an iterative offline MBRL framework, where we maximize a lower bound of the true expected return. With the proposed unified model-policy learning framework, we achieve competitive performance on a wide range of continuous-control offline reinforcement learning datasets.
  arXiv  Detail & Related papers  (2022-10-12T04:58:51Z)
• Model Generation with Provable Coverability for Offline Reinforcement Learning [14.333861814143718]
  offline optimization with dynamics-aware policy provides a new perspective for policy learning and out-of-distribution generalization. But due to the limitation under the offline setting, the learned model could not mimic real dynamics well enough to support reliable out-of-distribution exploration. We propose an algorithm to generate models optimizing their coverage for the real dynamics.
  arXiv  Detail & Related papers  (2022-06-01T08:34:09Z)
• MOPO: Model-based Offline Policy Optimization [183.6449600580806]
  offline reinforcement learning (RL) refers to the problem of learning policies entirely from a large batch of previously collected data. We show that an existing model-based RL algorithm already produces significant gains in the offline setting. We propose to modify the existing model-based RL methods by applying them with rewards artificially penalized by the uncertainty of the dynamics.
  arXiv  Detail & Related papers  (2020-05-27T08:46:41Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.