Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization

# Scaling World-Model Reinforcement Learning Through Diffusion Policy Optimization
Source: [https://arxiv.org/html/2605.26282](https://arxiv.org/html/2605.26282)
Xiaoyuan Cheng1,Wenxuan Yuan2,∗Zhancun Mu3Yuanzhao Zhang4 Yiming Yang1Hai Wang1Zhuo Sun5,†Che Liu6,† 1Dynamic Systems Lab, University College London 2College of Computing and Data Science, Nanyang Technological University 3School of Intelligence Science and Technology, Peking University4Santa Fe Institute 5School of Statistics and Data Science, Shanghai University of Finance and Economics 6Department of Computing, Imperial College London Code:[https://github\.com/Edmond1Cheng/MBDPO](https://github.com/Edmond1Cheng/MBDPO) Hugging Face:[https://huggingface\.co/BruceYuan/MBDPO](https://huggingface.co/BruceYuan/MBDPO)Core contributors: Xiaoyuan Cheng:[ucesxc4@ucl\.ac\.uk](https://arxiv.org/html/2605.26282v1/[email protected])and Wenxuan Yuan:[YUAN0186@e\.ntu\.edu\.sg](https://arxiv.org/html/2605.26282v1/[email protected])\.†Corresonding authors: Zhuo Sun:[zhuosunreid@outlook\.com](https://arxiv.org/html/2605.26282v1/[email protected]), Che Liu:[cl522@ic\.ac\.uk](https://arxiv.org/html/2605.26282v1/[email protected])\.

###### 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\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x1.png)Figure 1:Overview of offline and online performance\.\(Left\) Our method \(MBDPO\) significantly outperforms TD\-MPC2\(Hansenet al\.,[2024](https://arxiv.org/html/2605.26282#bib.bib11)\)in multi\-task offline pretraining, exhibiting a clear monotonic scaling behavior as model parameters increase from1\.71\.7M to340340M\. \(Right\) In the online\-from\-scratch setting, MBDPO consistently achieves superior or competitive results across44benchmarks with121121tasks\.## 1Introduction

Model\-based reinforcement learning \(RL\) has long been regarded as a promising paradigm for sample\-efficient planning in general cases\. By learning an explicit dynamical system, agents can dream and plan without multiple interactions with real environments\(Sutton,[1996](https://arxiv.org/html/2605.26282#bib.bib1); Suttonet al\.,[1998](https://arxiv.org/html/2605.26282#bib.bib2)\)\. Recent advances in world models have further expanded this paradigm\(Garridoet al\.,[2024](https://arxiv.org/html/2605.26282#bib.bib3); Zhouet al\.,[2024](https://arxiv.org/html/2605.26282#bib.bib4); Hafneret al\.,[2025a](https://arxiv.org/html/2605.26282#bib.bib5); Psenkaet al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib40)\), enabling large models to capture rich environment dynamics across diverse tasks\. Instead of serving solely as local frame predictors\(SVet al\.,[2023](https://arxiv.org/html/2605.26282#bib.bib6); Thrunet al\.,[1990](https://arxiv.org/html/2605.26282#bib.bib7); Guanet al\.,[2023](https://arxiv.org/html/2605.26282#bib.bib8)\), modern world models aim to learn generalizable representations of environment dynamics that can be reused across downstream control problems\. At the same time, the emergence of large\-scale datasets has created new opportunities for pretraining world models at scale\(Wuet al\.,[2023](https://arxiv.org/html/2605.26282#bib.bib9); Hafneret al\.,[2025b](https://arxiv.org/html/2605.26282#bib.bib10)\)\. Analogous to foundation models in vision and language, pretrained world models offer the potential to unlock scalable model\-based RL among multi\-tasks\(Hansenet al\.,[2024](https://arxiv.org/html/2605.26282#bib.bib11)\)\.

Despite recent advances, model\-based RL methods built upon world models still struggle to meaningfully outperform their model\-free counterparts\(Van Hasseltet al\.,[2019](https://arxiv.org/html/2605.26282#bib.bib12); Fujimotoet al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib13); Changet al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib14)\)\. Two explanations have been proposed\. The most common explanation attributes this shortfall to limited model accuracy and compounding prediction errors\(M\. Moerlandet al\.,[2023](https://arxiv.org/html/2605.26282#bib.bib17); Lambertet al\.,[2022](https://arxiv.org/html/2605.26282#bib.bib15)\)\. When a learned model is used to simulate long\-horizon trajectories, even small one\-step prediction errors accumulate over time, resulting in increasingly unreliable state estimates\. Motivated by this observation, prior work has focused on mitigating policy degradation induced by imperfect model dynamics\. Representative approaches include improving model accuracy\(Nagabandiet al\.,[2018](https://arxiv.org/html/2605.26282#bib.bib18); Hafneret al\.,[2019](https://arxiv.org/html/2605.26282#bib.bib19)\), incorporating uncertainty\-aware modeling\(Deisenroth and Rasmussen,[2011](https://arxiv.org/html/2605.26282#bib.bib20); Curiet al\.,[2020](https://arxiv.org/html/2605.26282#bib.bib21); Fuet al\.,[2022](https://arxiv.org/html/2605.26282#bib.bib22)\), employing truncated or branched rollouts\(Janneret al\.,[2019](https://arxiv.org/html/2605.26282#bib.bib16); Parket al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib23)\), and enhancing long\-horizon prediction\(Maet al\.,[2024](https://arxiv.org/html/2605.26282#bib.bib24); Farebrotheret al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib25)\)\. Underlying these efforts is an implicit assumption: that sufficiently accurate world models will naturally translate into better planning and, ultimately, better control performance\.

However, recent work challenges this point\(Changet al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib14)\), arguing that model accuracy alone does not dictate planning effectiveness\. Even highly accurate models can fail to yield strong policies, causing RL within world models to struggle in both performance and scalability\. The core obstacle lies in the misalignment between policy search and value learning\. Due to the substantial computational cost of search, the value function is typically trained using rollouts generated by a learned non\-search policy network\(Hansenet al\.,[2022](https://arxiv.org/html/2605.26282#bib.bib26),[2024](https://arxiv.org/html/2605.26282#bib.bib11),[2025](https://arxiv.org/html/2605.26282#bib.bib27); Wanget al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib67); Zhanet al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib28); Shimizu and Tomizuka,[2024](https://arxiv.org/html/2605.26282#bib.bib53)\)\. In contrast, policy improvement relies on trajectories produced by a search procedure \(i\.e\., model predictive path integral\(Williamset al\.,[2015](https://arxiv.org/html/2605.26282#bib.bib54)\)\) over a learned world model\. This mismatch introduces a distributional discrepancy: the value function is trained on state\-action pairs generated by a non\-search policy, whereas policy execution and improvement are governed by trajectories produced via search over the world model\. As a result, the value function is often queried on out\-of\-distribution state\-action pairs where its estimates are unreliable and typically overoptimistic; the search procedure then exploits these errors, leading to systematic overestimation bias and ultimately degrading both performance and scalability\(Fujimotoet al\.,[2019](https://arxiv.org/html/2605.26282#bib.bib29)\)\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x2.png)Figure 2:Core framework of MBDPO\.The target policy distribution is progressively shaped by a sequence of stepwise transition kernels, fromπϕN\\pi\_\{\\phi\}^\{N\}toπϕ0\\pi\_\{\\phi\}^\{0\}, which transforms a Gaussian prior𝒩​\(0,I\)\\mathcal\{N\}\(0,I\)into the optimal Gibbs policyπϕ∗\\pi\_\{\\phi\}^\{\*\}through multi\-step refinement\. Crucially, each transition kernelπϕτ−1​\(at:t\+Hτ−1\|at:t\+Hτ,zt\)\\pi^\{\\tau\-1\}\_\{\\phi\}\(a\_\{t:t\+H\}^\{\\tau\-1\}\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)is governed by the score functionϕ^\\hat\{\\phi\}, estimated entirely within the learned world model: Monte Carlo action sequence samples are rolled out through the latent dynamicsℱ\\mathcal\{F\}, evaluated by their energy\-regularized cumulative returnG~\\tilde\{G\}, and reweighted to recover a score that points toward higher\-return actions\. In this way, the world model redefines and corrects the score field, transforming the standard generative denoising process into model\-based policy optimization\. Furthermore, the implicit energyE∝−log⁡βE\\propto\-\\log\\betaanchors the policy to the non\-search policy distribution within a KL trust region, thereby unifying search and policy optimization\.Motivated by the bottlenecks of prior approaches, we investigate how to overcome these two limitations and unlock the scalability potential of model\-based RL within world models\. In this work, we introduce a novel perspective grounded in diffusion models and rethink how to make model\-based RL in world models scalable\. From this perspective, we demonstrate that world models enable direct policy optimization, where score matching is achieved through imagined trajectories\. Our contributions are threefold:

\(1\) We revisit the foundation of policy optimization and show, from the perspectives of both value iteration and policy improvement, that the misalignment between search and value learning can lead to suboptimality in existing world model approaches\.

\(2\) We propose model\-based diffusion policy optimization \(MBDPO\), a unified framework that formulates search and policy optimization as a diffusion process over imagined trajectories of world models, thereby unlocking the potential of world models for scalable policy learning \(see Figure[2](https://arxiv.org/html/2605.26282#S1.F2)\)\. In this formulation, the buffer dataset induces an implicit energy function that anchors policy updates to the behavior distribution\. This mechanism enables score correction beyond the suboptimality in previous methods\. By doing so, MBDPO unifies search and policy optimization, eliminating their structural misalignment\.

\(3\) We develop a scalable diffusion RL algorithm jointly trained with the world model and demonstrate consistent, monotonic performance improvements as model capacity increases across diverse settings, including multi\-task offline pretraining, online learning, and offline\-to\-online fine\-tuning\. Beyond these quantitative results, we visualize our controlled latent trajectories to evaluate the underlying representation, finding that the learned policy consistently exhibits structured behaviors that align with physical intuition\.

Related Work\.We review prior work in two areas: world models for RL, and diffusion policy\.

World Models\.The concept of world models originates from Jay W\. Forrester, who introduced it for constructing mental representations of complex systems\(Forrester,[1971](https://arxiv.org/html/2605.26282#bib.bib51)\)\. It was later adapted to machine learning, enabling agents to learn and act within learned internal environments\(Ha and Schmidhuber,[2018](https://arxiv.org/html/2605.26282#bib.bib31)\)\. Later, the Dreamer series\(Hafneret al\.,[2019](https://arxiv.org/html/2605.26282#bib.bib19),[2020](https://arxiv.org/html/2605.26282#bib.bib33),[2025a](https://arxiv.org/html/2605.26282#bib.bib5),[2025b](https://arxiv.org/html/2605.26282#bib.bib10)\)transformed latent world modeling from a proof\-of\-concept into a scalable and robust RL paradigm\. However, due to model bias, training the policy purely within the learned world model can lead to distributional shift\(Janneret al\.,[2019](https://arxiv.org/html/2605.26282#bib.bib16)\)\. Another line of research, TD\-MPC\(Hansenet al\.,[2022](https://arxiv.org/html/2605.26282#bib.bib26),[2024](https://arxiv.org/html/2605.26282#bib.bib11),[2025](https://arxiv.org/html/2605.26282#bib.bib27)\), combines model\-based and model\-free RL\. However, the existence of misalignment between search and value hinders scalability and performance\. Recent approaches, including Vision\-Language\-Action models with world models\(Zhenet al\.,[2024](https://arxiv.org/html/2605.26282#bib.bib35); Cenet al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib37); Teamet al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib39); Intelligenceet al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib62)\)and World Action Models \(WAMs\)\(Hanet al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib38); Yeet al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib34); Liet al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib36)\), have been proposed to generate action chunks and state representations, driven by the rapid development of large pre\-trained vision models\. Compared to RL, the WAM family does not perform explicit policy optimization, which limits its ability to achieve long\-horizon control and robust generalization\.

Diffusion Policy\.Diffusion policies can be broadly categorized into model\-free and model\-based approaches\. Model\-free diffusion policies do not rely on learned dynamics; instead, they directly learn the policy distribution from a fixed dataset\(Janneret al\.,[2022](https://arxiv.org/html/2605.26282#bib.bib41); Chiet al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib42)\)or a learnedQQ\-function\(Psenkaet al\.,[2023](https://arxiv.org/html/2605.26282#bib.bib43); Donget al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib44); Chenget al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib45); Li and Levine,[2026](https://arxiv.org/html/2605.26282#bib.bib52)\)\. In contrast, model\-based diffusion techniques leverage environment dynamics to estimate the score function without relying on offline data\(Panet al\.,[2024](https://arxiv.org/html/2605.26282#bib.bib46); Xueet al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib47); Chenget al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib48)\)\. Crucially, computing this score function requires access to the transition dynamics, making world models an ideal substrate; they provide the simulated environment necessary\. However, in many practical scenarios, ground\-truth dynamics are unavailable, limiting their direct applicability\. This bottleneck, in turn, highlights the critical necessity of learned world models\. By approximating a diverse spectrum of dynamics, world models hold the key to unlocking the full potential of model\-based diffusion for scalable policy learning\. Despite this immense promise, effective model\-based diffusion policies within learned world models remain underexplored\.

## 2Preliminaries

Model\-Based RL\.We consider a Markov Decision Process \(MDP\) defined by the tupleℳ=\(𝒮,𝒜,𝒯,r,ρ,γ\)\\mathcal\{M\}=\(\\mathcal\{S\},\\mathcal\{A\},\\mathcal\{T\},r,\\rho,\\gamma\)\(Puterman,[1990](https://arxiv.org/html/2605.26282#bib.bib30)\), where𝒮\\mathcal\{S\}and𝒜\\mathcal\{A\}denote the state and action spaces, respectively\.ρ\\rhois the initial state distribution,𝒯:𝒮×𝒜→𝒮\\mathcal\{T\}:\\mathcal\{S\}\\times\\mathcal\{A\}\\to\\mathcal\{S\}represents the transition dynamics,r:𝒮×𝒜→ℝr:\\mathcal\{S\}\\times\\mathcal\{A\}\\to\\mathbb\{R\}is the reward function in the state space, andγ∈\(0,1\)\\gamma\\in\(0,1\)is the discount factor\. For a policyπ​\(a\|s\)\\pi\(a\|s\), the action\-value function \(orQQ\-function\)Qπ​\(s,a\)Q^\{\\pi\}\(s,a\)is defined as the expected discounted return starting from statessand actionaa:Qπ\(s,a\)≔𝔼π\[∑t=0∞γtr\(st,at\)\|s0=s,a0=a\]Q^\{\\pi\}\(s,a\)\\coloneqq\\mathbb\{E\}\_\{\\pi\}\\left\[\\sum\_\{t=0\}^\{\\infty\}\\gamma^\{t\}r\(s\_\{t\},a\_\{t\}\)\\;\\middle\|\\;s\_\{0\}=s,a\_\{0\}=a\\right\]\. Our objective is to find a policyπ\\pithat maximizes the expected discounted return:𝔼s∼dπ,a∼π​\[Qπ​\(s,a\)\]\.\\mathbb\{E\}\_\{s\\sim d^\{\\pi\},a\\sim\\pi\}\[Q^\{\\pi\}\(s,a\)\]\.In this paper, we jointly optimize the policyπ\\piand a world model that approximates the latent MDP dynamics\.

Latent World Model\.Drawing on recent advances\(Hansenet al\.,[2024](https://arxiv.org/html/2605.26282#bib.bib11),[2025](https://arxiv.org/html/2605.26282#bib.bib27); Hafneret al\.,[2025a](https://arxiv.org/html/2605.26282#bib.bib5)\), we employ a latent world model to capture compact representations essential for complex decision\-making tasks\. Our world model comprises the following components:

Encoderℰ:𝒮×ℰe​n​v→𝒵,\\displaystyle\\mathcal\{E\}\\colon\\mathcal\{S\}\\times\{\\color\[rgb\]\{\.5,\.5,\.5\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{\.5,\.5,\.5\}\\pgfsys@color@gray@stroke\{\.5\}\\pgfsys@color@gray@fill\{\.5\}\\mathcal\{E\}\_\{env\}\}\\to\\mathcal\{Z\},\\quadzt=ℰ​\(st,e\);\\displaystyle z\_\{t\}=\\mathcal\{E\}\(s\_\{t\},\{\\color\[rgb\]\{\.5,\.5,\.5\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{\.5,\.5,\.5\}\\pgfsys@color@gray@stroke\{\.5\}\\pgfsys@color@gray@fill\{\.5\}e\}\);\(1\)Latent Dynamicsℱ:𝒵×𝒜×ℰe​n​v→𝒵,\\displaystyle\\mathcal\{F\}\\colon\\mathcal\{Z\}\\times\\mathcal\{A\}\\times\{\\color\[rgb\]\{\.5,\.5,\.5\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{\.5,\.5,\.5\}\\pgfsys@color@gray@stroke\{\.5\}\\pgfsys@color@gray@fill\{\.5\}\\mathcal\{E\}\_\{env\}\}\\to\\mathcal\{Z\},\\quadzt\+1=ℱ​\(zt,at,e\);\\displaystyle z\_\{t\+1\}=\\mathcal\{F\}\(z\_\{t\},a\_\{t\},\{\\color\[rgb\]\{\.5,\.5,\.5\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{\.5,\.5,\.5\}\\pgfsys@color@gray@stroke\{\.5\}\\pgfsys@color@gray@fill\{\.5\}e\}\);Latent Rewardℛ:𝒵×𝒜×ℰe​n​v→ℝ,\\displaystyle\\mathcal\{R\}\\colon\\mathcal\{Z\}\\times\\mathcal\{A\}\\times\{\\color\[rgb\]\{\.5,\.5,\.5\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{\.5,\.5,\.5\}\\pgfsys@color@gray@stroke\{\.5\}\\pgfsys@color@gray@fill\{\.5\}\\mathcal\{E\}\_\{env\}\}\\to\\mathbb\{R\},\\quadr^t=ℛ​\(zt,at,e\);\\displaystyle\\hat\{r\}\_\{t\}=\\mathcal\{R\}\(z\_\{t\},a\_\{t\},\{\\color\[rgb\]\{\.5,\.5,\.5\}\\definecolor\[named\]\{pgfstrokecolor\}\{rgb\}\{\.5,\.5,\.5\}\\pgfsys@color@gray@stroke\{\.5\}\\pgfsys@color@gray@fill\{\.5\}e\}\);wherezt∈𝒵z\_\{t\}\\in\\mathcal\{Z\}denotes the latent state ande∈ℰe​n​ve\\in\\mathcal\{E\}\_\{env\}is a learnable task embedding designed for multi\-task generalization\. For brevity, weomiteein subsequent discussions unless its inclusion is explicitly required\. Rather than operating in the raw observation space, our RL framework is built entirely upon the learned latent world model\.

Diffusion Policy\.Instead of generating a single\-step action, we optimize an action sequenceat:t\+H≔\{at,…,at\+H\}a\_\{t:t\+H\}\\coloneqq\\\{a\_\{t\},\\dots,a\_\{t\+H\}\\\}and execute it in a receding\-horizon fashion, similar to MPC\. We adopt the Denoising Diffusion Probabilistic Model \(DDPM\) framework\(Hoet al\.,[2020](https://arxiv.org/html/2605.26282#bib.bib49)\)to represent our generative policy\. Letτ∈\{1,…,N\}\\tau\\in\\\{1,\\dots,N\\\}denote the diffusion timestep\. Since each step of the reverse\-time transition is determined by the score functionϕ​\(zt,at:t\+Hτ,τ\)≔∇at:t\+Hτlog⁡πϕτ​\(at:t\+Hτ\|zt\)\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)\\coloneqq\\nabla\_\{a\_\{t:t\+H\}^\{\\tau\}\}\\log\\pi^\{\\tau\}\_\{\\phi\}\(a\_\{t:t\+H\}^\{\\tau\}\|z\_\{t\}\), the evolution of the policy is directly governed byϕ\\phi\. Accordingly, the reverse diffusion process is characterized by the reverse\-time transition kernelπϕτ−1​\(at:t\+Hτ−1\|at:t\+Hτ,zt\)\\pi\_\{\\phi\}^\{\\tau\-1\}\(a\_\{t:t\+H\}^\{\\tau\-1\}\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\), where the marginal distribution of the clean action sequence is

πϕ0​\(at:t\+H0\|zt\)=∫at:t\+H1:NπϕN​\(at:t\+HN\|zt\)​∏τ=N1πϕτ−1​\(at:t\+Hτ−1\|at:t\+Hτ,zt\)​d​at:t\+H1:N,\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)=\\int\_\{a\_\{t:t\+H\}^\{1:N\}\}\\pi^\{N\}\_\{\\phi\}\(a\_\{t:t\+H\}^\{N\}\|z\_\{t\}\)\\prod\_\{\\tau=N\}^\{1\}\\pi\_\{\\phi\}^\{\\tau\-1\}\(a\_\{t:t\+H\}^\{\\tau\-1\}\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)da\_\{t:t\+H\}^\{1:N\},\(2\)which describes a multi\-step refinement process that transforms from a Gaussian priorπN​\(at:t\+HN\|zt\)∼𝒩​\(0,I\)\\pi^\{N\}\(a\_\{t:t\+H\}^\{N\}\|z\_\{t\}\)\\sim\\mathcal\{N\}\(0,I\)\. Specifically, the stepwise transition kernel is guided by the score function:

πϕτ−1​\(at:t\+Hτ−1\|at:t\+Hτ,zt\)=𝒩​\(at:t\+Hτ−1;1ατ​\(at:t\+Hτ\+\(1−ατ\)​ϕ​\(zt,at:t\+Hτ,τ\)\),σ2​\(τ\)​I\),\\small\\pi\_\{\\phi\}^\{\\tau\-1\}\(a\_\{t:t\+H\}^\{\\tau\-1\}\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)=\\mathcal\{N\}\\left\(a\_\{t:t\+H\}^\{\\tau\-1\};\\frac\{1\}\{\\sqrt\{\\alpha^\{\\tau\}\}\}\\left\(a\_\{t:t\+H\}^\{\\tau\}\+\(1\-\\alpha^\{\\tau\}\)\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)\\right\),\\sigma^\{2\}\(\\tau\)I\\right\),\(3\)whereατ\\alpha^\{\\tau\}is the scalar factor of noise injection at theτ\\taustep, andσ2​\(τ\)\\sigma^\{2\}\(\\tau\)denotes the step\-dependent variance\. In contrast to conventional diffusion policies in the model\-free setting, our approach leverages the world model to estimate the score functionϕ\\phivia imagined trajectories, thereby unifying the search and policy optimization\. The next section establishes the link between the diffusion refinement in \([2](https://arxiv.org/html/2605.26282#S2.E2)\) and policy optimization within the world model framework\.

Given the unified latent representation within world models, we focus on the following central question:

What are the fundamental bottlenecks in scaling model\-based RL within latent world models, and how can we architect scalable algorithms to overcome them?

In the subsequent section, we analyze the scalability limitations inherent in current model\-based frameworks and propose a principled approach to resolve these challenges\.

## 3Method

In this section, we analyze why existing model\-based RL approaches in world models do not correspond to a monotonic policy improvement, and thus fail to guarantee convergence to the optimal policy\. To address this issue, we propose a corrected and scalable learning paradigm that reconsiders the interplay between the world model and policy optimization\.

### 3\.1Misalignment Between Search and Value Learning

We briefly revisit RL strategies built upon world models, such as those proposed in\(Hansenet al\.,[2022](https://arxiv.org/html/2605.26282#bib.bib26),[2025](https://arxiv.org/html/2605.26282#bib.bib27); Zhanet al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib28); Linet al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib50)\)\. In these representative approaches, a sampling\-based planner operates within the learned world model to generate high\-value actions for environment interaction\. However, this design inherently introduces a misalignment between policy improvement and value learning\(Changet al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib14)\)\. Specifically, while the value function is trained on state\-action pairs generated by a non\-search policyβ\\beta111Here,β\\betadenotes the data\-collection policy associated with the training samples\. In offline settings,β\\betais the behavior policy that generated the fixed offline dataset\. In online settings,β\\betarefers to the implicit replay\-buffer policy, i\.e\., the mixture of historical policies that collected the transitions stored in the replay buffer\., the executed policyπ\\pi222Here, we useπ\\pito denote the search policy used in other methods, contrasting it with our diffusion policyπϕ\\pi\_\{\\phi\}\.relies on the sampling\-based planner over the world model, resulting in a significant distributional discrepancy \(see more details in Appendix[C\.1\.1](https://arxiv.org/html/2605.26282#A3.SS1.SSS1)\)\. This misalignment hinders performance and limits scalability\. In particular, we elucidate why such a procedure may fail to guarantee optimality, suggesting that it does not constitute true policy optimization in the conventional sense\.

\{insightmwithframe\}

\[Value Iteration Gap\] Letβ\\betadenote the learned, non\-search policy, and letπ\\pibe the policy induced by search\. SupposeQ^\\hat\{Q\}is the learned value function trained on data collected underβ\\beta\. Then, the discrepancy between the Bellman operators underπ\\piandβ\\beta\(denoted asTπT^\{\\pi\}andTβT^\{\\beta\}, respectively; see Definition[B\.1](https://arxiv.org/html/2605.26282#A2.Thmtheorem1)\) introduces the following error when updatingQ^\\hat\{Q\}:

‖Tπ​Q^−Tβ​Q^‖∞≤γ​‖Q^‖∞​2​DKLmax​\(π∥β\),\\\|T^\{\\pi\}\\hat\{Q\}\-T^\{\\beta\}\\hat\{Q\}\\\|\_\{\\infty\}\\leq\\gamma\\\|\\hat\{Q\}\\\|\_\{\\infty\}\\sqrt\{2D\_\{\\text\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\)\},\(4\)whereDKLmax\(π∥β\)≔supz∈𝒵DKL\(π\(⋅\|z\)∥β\(⋅\|z\)\)D\_\{\\text\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\)\\coloneqq\\sup\_\{z\\in\\mathcal\{Z\}\}D\_\{\\text\{KL\}\}\(\\pi\(\\cdot\|z\)\\\|\\beta\(\\cdot\|z\)\), see the detailed proof in Appendix[C\.1\.2](https://arxiv.org/html/2605.26282#A3.SS1.SSS2)\.

While search aims to improve the policy, it can induce a vicious cycle of error amplification\. SinceQ^\\hat\{Q\}is trained underβ\\beta, it is unreliable in out\-of\-distribution regions\. The search policyπ\\pitends to select actions that maximizeQ^\\hat\{Q\}, which often correspond to spurious overestimated values, i\.e\.,Q^​\(z,a\)\>Qβ​\(z,a\)\\hat\{Q\}\(z,a\)\>Q^\{\\beta\}\(z,a\)for some\(z,a\)∈𝒵×𝒜\(z,a\)\\in\\mathcal\{Z\\times A\}\. The Bellman updateTπ​Q^T^\{\\pi\}\\hat\{Q\}then bootstraps on these erroneous peaks, reinforcing rather than correcting them\. As a result, the iterative updates can amplify the estimation error ofQ^\\hat\{Q\}as shown in \([4](https://arxiv.org/html/2605.26282#S3.E4)\)\.

\{insightmwithframe\}

\[Policy Improvement Gap\] LetJ​\(π\)J\(\\pi\)andJ​\(β\)J\(\\beta\)denote the standard discounted returns induced by policiesπ\\piandβ\\beta, respectively\. Furthermore, letJ^​\(β\)≔𝔼z∼dβ,a∼β​\[Q^​\(z,a\)\]\\hat\{J\}\(\\beta\)\\coloneqq\\mathbb\{E\}\_\{z\\sim d^\{\\beta\},a\\sim\\beta\}\[\\hat\{Q\}\(z,a\)\]andJ^β​\(π\)≔𝔼z∼dβ,a∼π​\[Q^​\(z,a\)\]\\hat\{J\}\_\{\\beta\}\(\\pi\)\\coloneqq\\mathbb\{E\}\_\{z\\sim d^\{\\beta\},a\\sim\\pi\}\[\\hat\{Q\}\(z,a\)\]represent the estimatedQQ\-based surrogate objectives evaluated under the discounted state occupancy distributiondβd^\{\\beta\}using an approximate value functionQ^\\hat\{Q\}\. Then, the performance gap satisfies the following bound:

J​\(π\)−J​\(β\)⏟true improvement≥11−γ​\(J^β​\(π\)−J^​\(β\)\)⏟estimated improvement−C1​DKLmax​\(π∥β\)−C2​‖Qβ−Q^‖∞⏟value iteration gap in​Problem​[3\.1](https://arxiv.org/html/2605.26282#S3.SS1),\\underbrace\{J\(\\pi\)\-J\(\\beta\)\}\_\{\\text\{true improvement\}\}\\geq\\frac\{1\}\{1\-\\gamma\}\\underbrace\{\\big\(\\hat\{J\}\_\{\\beta\}\(\\pi\)\-\\hat\{J\}\(\\beta\)\\big\)\}\_\{\\text\{estimated improvement\}\}\-C\_\{1\}D\_\{\\text\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\)\-\\underbrace\{C\_\{2\}\\\|Q^\{\\beta\}\-\\hat\{Q\}\\\|\_\{\\infty\}\}\_\{\\text\{value iteration gap in\}~\\text\{Problem\}~\\ref\{prob\_1\}\},\(5\)whereC1C\_\{1\}andC2C\_\{2\}are positive constants depending on the discount factorγ\\gammaand the boundedness of the value functions, see the analysis in Appendix[C\.2](https://arxiv.org/html/2605.26282#A3.SS2)\. The above result in \([5](https://arxiv.org/html/2605.26282#S3.E5)\) highlights two key failure modes induced by the search policy\. First, the policy improvement is inherently local under the metric of KL divergence: once the search policyπ\\pideviates significantly fromβ\\beta, the learned value functionQ^\\hat\{Q\}no longer provides a reliable estimate for evaluatingπ\\pi, leading to biased performance estimation\. In particular, the approximation gap‖Qβ−Q^‖∞\\\|Q^\{\\beta\}\-\\hat\{Q\}\\\|\_\{\\infty\}can grow as updates are performed underπ\\pi\(as shown in Problem[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)\), since the Bellman operatorTπT^\{\\pi\}propagates and reinforces errors in regions whereQ^\\hat\{Q\}is inaccurate\. Empirically, this phenomenon is characterized by a severe inflation in cross TD error and substantial action drift relative to the base policy, as illustrated in Figure[3](https://arxiv.org/html/2605.26282#S3.F3)\. Together, these theoretical and empirical effects culminate in a self\-reinforcing cycle of distribution shift and error amplification\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x3.png)Figure 3:\(Left\) Cross temporal difference \(TD\) error comparison between MBDPO and TD\-MPC2 across two benchmark suites, DMControl and MetaWorld \(detailed subtask curves are provided in Figure[21](https://arxiv.org/html/2605.26282#A1.F21)\)\. \(Right\) Averaged relative action drift during training\. \(Right 1\) Exponential moving average \(EMA\) of the mean drift across 8 tasks; larger values indicate greater deviation\. \(Right 2\) Action drift distributions across different training intervals, measured by the average log\-likelihood ratio1H​∑h=0Hlog⁡π​\(ah∣zh\)β​\(ah∣zh\)\\frac\{1\}\{H\}\\sum\_\{h=0\}^\{H\}\\log\\frac\{\\pi\(a\_\{h\}\\mid z\_\{h\}\)\}\{\\beta\(a\_\{h\}\\mid z\_\{h\}\)\}\. Compared to TD\-MPC2, the diffusion policy optimization framework significantly reduces action drift, with the contrastive variant \(η=0\.1\\eta=0\.1\) yielding the closest alignment with the policy network\. The reported evaluation averages results across 8 distinct online tasks \(individual subtask distributions are detailed in Figure[22](https://arxiv.org/html/2605.26282#A1.F22)\)\.At the core, this issue arises from the interplay between value approximation error and unconstrained policy improvement within world models: the former introduces spurious high\-value estimates, while the latter allows the search policyπ\\pito deviate significantly fromβ\\beta, exacerbating distribution shift\. To bridge this gap, we unify search and policy optimization into a single algorithm, achieving a monotonic policy improvement within the world model\.

### 3\.2Model\-Based Diffusion Policy Optimization \(MBDPO\) in World Models

In this subsection, we define the policy learning objective and detail the core component of our framework: MBDPO within world models\.

Policy Objective\.In RL built upon world models, policy optimization is framed as seeking a policyπϕ\\pi\_\{\\phi\}through the following formulation:

arg⁡maxπϕ⁡𝔼at:t\+H∼πϕ\(⋅\|zt\)​\[G​\(zh,at:t\+H\)\],\\arg\\max\_\{\\pi\_\{\\phi\}\}\\mathbb\{E\}\_\{a\_\{t:t\+H\}\\sim\\pi\_\{\\phi\}\(\\cdot\|z\_\{t\}\)\}\\big\[G\(z\_\{h\},a\_\{t:t\+H\}\)\\big\],\(6\)where

G​\(zh,at:t\+H\)≔𝔼zh\+1∼ℱ\(zh,ah\),at:t\+H∼πϕ\(⋅\|zt\),​\[∑h=tt\+H−1γh−t​ℛ​\(zh,ah\)\+γH​Q^​\(zt\+H,at\+H\)\]\.G\(z\_\{h\},a\_\{t:t\+H\}\)\\coloneqq\\mathbb\{E\}\_\{z\_\{h\+1\}\\sim\\mathcal\{F\}\(z\_\{h\},a\_\{h\}\),\\ a\_\{t:t\+H\}\\sim\\pi\_\{\\phi\}\(\\cdot\|z\_\{t\}\),\}\\big\[\\sum\_\{h=t\}^\{t\+H\-1\}\\gamma^\{h\-t\}\\mathcal\{R\}\(z\_\{h\},a\_\{h\}\)\+\\gamma^\{H\}\\hat\{Q\}\(z\_\{t\+H\},a\_\{t\+H\}\)\\big\]\.\(7\)In prior work\(Hansenet al\.,[2022](https://arxiv.org/html/2605.26282#bib.bib26),[2024](https://arxiv.org/html/2605.26282#bib.bib11),[2025](https://arxiv.org/html/2605.26282#bib.bib27); Zhanet al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib28); Psenkaet al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib40)\), the action sequenceat:t\+Ha\_\{t:t\+H\}is typically obtained via model predictive path integral \(MPPI\)\. However, this approach can be problematic: MPPI is initialized with a Gaussian prior and relies on trajectories generated by the world model, which can lead to a large KL divergence from the non\-search policyβ\\beta\. Consequently, both the Bellman update and policy improvement may enter a detrimental feedback loop, as illustrated in Problems[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)and[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)\. Essentially, this paradigm does not constitute true policy optimization; it merely performs a heuristic search that lacks the guarantees and distributional consistency required for a monotonic improvement\. To address this, it is imperative to establish a principled policy optimization objective that simultaneously maximizes expected returns and constrains the search policy to the non\-search base policyβ\\beta:

maxπϕ𝔼at:t\+H∼πϕ\(⋅\|zt\)​\[G​\(zt,at:t\+H\)\]s\.t\.DKL\(πϕ\(⋅\|zt\)∥β\(⋅\|zt\)\)≤ε⏟KL constraint,zh\+1=ℱ\(zh,ah\),∀h∈\{t,…,t\+H−1\}\.⏟latent dynamics constraint\\begin\{split\}\\max\_\{\\pi\_\{\\phi\}\}\\quad&\\mathbb\{E\}\_\{a\_\{t:t\+H\}\\sim\\pi\_\{\\phi\}\(\\cdot\|z\_\{t\}\)\}\\big\[G\(z\_\{t\},a\_\{t:t\+H\}\)\\big\]\\\\ \\text\{s\.t\.\}\\quad&\\underbrace\{D\_\{\\text\{KL\}\}\(\\pi\_\{\\phi\}\(\\cdot\|z\_\{t\}\)\\\|\\beta\(\\cdot\|z\_\{t\}\)\)\\leq\\varepsilon\}\_\{\\text\{KL constraint\}\},\\\\ &\\underbrace\{z\_\{h\+1\}=\\mathcal\{F\}\(z\_\{h\},a\_\{h\}\),\\quad\\forall h\\in\\\{t,\\dots,t\+H\-1\\\}\.\}\_\{\\text\{latent dynamics constraint\}\}\\end\{split\}\(8\)Here, the KL divergence enforces policy regularization to prevent suboptimality, while the stepwise transitions encapsulate the rollouts within the learned latent world modelℱ\\mathcal\{F\}over a planning horizonHH\. The formulation in \([8](https://arxiv.org/html/2605.26282#S3.E8)\) anchors the policy optimization on the distributionβ\\beta, thereby mitigating the failure modes characterized in Problems[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)and[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)\. Because the exact density ofβ\\betais intractable, we introduce the following implicit energy function\.

Implicit Energy Function\.To this end, the implicit energy functionEθE\_\{\\theta\}is learned from the buffer dataset via a contrastive objective:

minEθ⁡ℒ​\(Eθ\)=−𝔼​\[log⁡exp⁡\(−Eθ​\(z,a\)\)exp⁡\(−Eθ​\(z,a\)\)\+∑jexp⁡\(−Eθ​\(z,aj−\)\)\]\.\\min\_\{E\_\{\\theta\}\}\\mathcal\{L\}\(E\_\{\\theta\}\)=\-\\mathbb\{E\}\\Big\[\\log\\frac\{\\exp\(\-E\_\{\\theta\}\(z,a\)\)\}\{\\exp\(\-E\_\{\\theta\}\(z,a\)\)\+\\sum\_\{j\}\\exp\(\-E\_\{\\theta\}\(z,a\_\{j\}^\{\-\}\)\)\}\\Big\]\.\(9\)As the dataset is large\-scale and diverse, the negative samplesaj−a\_\{j\}^\{\-\}randomly drawn from other latent states within the same batch can be closely approximated by a uniform distribution over the entire action space\. Consequently, the summation in the denominator asymptotically tracks the unweighted partition function, thereby aligning the learned energy function with the log behavior policy density up to a state\-dependent constant:

−E​\(z,a\)≈log⁡β​\(a\|z\)\+C​\(z\)\.\\displaystyle\-E\(z,a\)\\approx\\log\\beta\(a\|z\)\+C\(z\)\.Based on the learned energy, we can reparameterize the reward function by absorbing the implicit energy function asℛ~​\(zh,ah\)=ℛ​\(zh,ah\)−ηγh−t​E​\(zh,ah\)\\tilde\{\\mathcal\{R\}\}\(z\_\{h\},a\_\{h\}\)=\\mathcal\{R\}\(z\_\{h\},a\_\{h\}\)\-\\frac\{\\eta\}\{\\gamma^\{h\-t\}\}E\(z\_\{h\},a\_\{h\}\), whereη\\etais the regularized factor\.

Diffusion Policy\.Instead of relying on MPPI merely for policy search, we adopt MBDPO to unify search and policy optimization\. Unlike the conventional model\-free setting, the score functionϕ\\phiin \([3](https://arxiv.org/html/2605.26282#S2.E3)\) can be estimated directly from trajectories imagined from the learned world model, without requiring a forward diffusion process:

ϕ​\(zt,at:t\+Hτ,τ\)=−at:t\+Hτ1−α¯τ\+α¯τ1−α¯τ​𝔼at:t\+H0\|τ∼𝒩​\(1α¯τ​at:t\+Hτ,1−α¯τα¯τ​I\)​\[w​\(at:t\+H0\|τ\)⋅at:t\+H0\|τ\]\.\\displaystyle\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)=\-\\frac\{a^\{\\tau\}\_\{t:t\+H\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\+\\frac\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\\mathbb\{E\}\_\{a\_\{t:t\+H\}^\{0\|\\tau\}\\sim\\mathcal\{N\}\(\\frac\{1\}\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}a\_\{t:t\+H\}^\{\\tau\},\\frac\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\{\\bar\{\\alpha\}^\{\\tau\}\}I\)\}\\big\[w\(a^\{0\|\\tau\}\_\{t:t\+H\}\)\\cdot a\_\{t:t\+H\}^\{0\|\\tau\}\\big\]\.\(10\)Here,at:t\+H0\|τa\_\{t:t\+H\}^\{0\|\\tau\}is a random variable that follows the sample posterior distribution𝒩​\(1α¯τ​at:t\+Hτ,1−α¯τα¯τ​I\)\\mathcal\{N\}\(\\frac\{1\}\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}a\_\{t:t\+H\}^\{\\tau\},\\frac\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\{\\bar\{\\alpha\}^\{\\tau\}\}I\), andα¯τ=∏k=1ταk\\bar\{\\alpha\}^\{\\tau\}=\\prod\_\{k=1\}^\{\\tau\}\\alpha^\{k\}\. The importance weightw​\(at:t\+H0\|τ\)w\(a\_\{t:t\+H\}^\{0\|\\tau\}\)is defined as:

w​\(at:t\+H0\|τ\)=exp⁡\(G~​\(zt,at:t\+H0\|τ\)κ\)𝔼at:t\+H0\|τ∼𝒩​\(1α¯τ​at:t\+Hτ,1−α¯τα¯τ​I\)​\[exp⁡\(G~​\(zt,at:t\+H0\|τ\)κ\)\],w\(a\_\{t:t\+H\}^\{0\|\\tau\}\)=\\frac\{\\exp\(\\frac\{\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\|\\tau\}\)\}\{\\kappa\}\)\}\{\\mathbb\{E\}\_\{a\_\{t:t\+H\}^\{0\|\\tau\}\\sim\\mathcal\{N\}\(\\frac\{1\}\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}a\_\{t:t\+H\}^\{\\tau\},\\frac\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\{\\bar\{\\alpha\}^\{\\tau\}\}I\)\}\[\\exp\(\\frac\{\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\|\\tau\}\)\}\{\\kappa\}\)\]\},\(11\)whereκ\\kappais the temperature factor, with

G~​\(zt,at:t\+H0\|τ\)≔∑h=tt\+H−1γh−t​ℛ~​\(zh,ah0\|τ\)\+γH​Q^​\(zt\+H0\|τ,at\+H0\|τ\)−η​E​\(zt\+H0\|τ,at\+H0\|τ\)\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\|\\tau\}\)\\coloneqq\\sum\_\{h=t\}^\{t\+H\-1\}\\gamma^\{h\-t\}\\tilde\{\\mathcal\{R\}\}\(z\_\{h\},a^\{0\|\\tau\}\_\{h\}\)\+\\gamma^\{H\}\\hat\{Q\}\(z^\{0\|\\tau\}\_\{t\+H\},a^\{0\|\\tau\}\_\{t\+H\}\)\-\\eta E\(z^\{0\|\\tau\}\_\{t\+H\},a^\{0\|\\tau\}\_\{t\+H\}\)representing the cumulative return regularized with the implicit energy\. By substituting the score function into the reverse\-time transition kernel in \([3](https://arxiv.org/html/2605.26282#S2.E3)\), the action sequenceat:t\+Hτ−1a\_\{t:t\+H\}^\{\\tau\-1\}atτ−1\\tau\-1is

at:t\+Hτ−1=1ατ​\(at:t\+Hτ\+\(1−ατ\)​ϕ​\(zt,at:t\+Hτ,τ\)\)\+σ​\(τ\)​ϵ,ϵ∼𝒩​\(0,I\)a\_\{t:t\+H\}^\{\\tau\-1\}=\\frac\{1\}\{\\sqrt\{\\alpha^\{\\tau\}\}\}\\left\(a\_\{t:t\+H\}^\{\\tau\}\+\(1\-\\alpha^\{\\tau\}\)\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)\\right\)\+\\sigma\(\\tau\)\\epsilon,\\quad\\epsilon\\sim\\mathcal\{N\}\(0,I\)\(12\)Lemma[C\.1](https://arxiv.org/html/2605.26282#A3.Thmtheorem1)in the Appendix presents the detailed derivation of equations \([10](https://arxiv.org/html/2605.26282#S3.E10)\)–\([12](https://arxiv.org/html/2605.26282#S3.E12)\)\. To generate the optimized action sequence, we start from a Gaussian priorat:t\+HN∼𝒩​\(0,I\)a\_\{t:t\+H\}^\{N\}\\sim\\mathcal\{N\}\(0,I\)and iteratively apply the reverse\-time transition kernel defined in \([2](https://arxiv.org/html/2605.26282#S2.E2)\) and \([3](https://arxiv.org/html/2605.26282#S2.E3)\)\. At each denoising step, the score in \([10](https://arxiv.org/html/2605.26282#S3.E10)\) shapes the transition kernel based on the returns evaluated from imagined trajectories within the world model\. In this sense, the world model acts as a score\-field corrector for policy optimization\. Once the exact score function is matched in \([10](https://arxiv.org/html/2605.26282#S3.E10)\), the optimal policy for the KL\-constrained objective in \([8](https://arxiv.org/html/2605.26282#S3.E8)\) is derived, as stated in the following theorem\.

\{theoremwithframe\}

\[Equivalence of Score Matching and Policy Optimization\] Consider the policy refinement task in \([2](https://arxiv.org/html/2605.26282#S2.E2)\) and the energy minimization objective in \([9](https://arxiv.org/html/2605.26282#S3.E9)\)\. For a given temperatureκ\>0\\kappa\>0and any latent statezt∈𝒵z\_\{t\}\\in\\mathcal\{Z\}, letπϕ∗\\pi^\{\*\}\_\{\\phi\}be the solution to the KL\-constrained policy optimization problem formulated in \([8](https://arxiv.org/html/2605.26282#S3.E8)\)\. The optimum is attained if and only if the score matching condition in \([10](https://arxiv.org/html/2605.26282#S3.E10)\) is satisfied\. Specifically, under the transition dynamics of the world model, the optimal policyπϕ∗\\pi^\{\*\}\_\{\\phi\}can be approximated by the following Gibbs distribution:

πϕ∗​\(at:t\+H\|zt\)∝\(∏h=tt\+Hβ​\(ah\|zh\)\)​exp⁡\(G​\(zt,at:t\+H\)κ\)​∏h=tt\+H−1𝟙\{zh\+1=ℱ​\(zh,ah\)\},\\displaystyle\\pi^\{\*\}\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)\\propto\\left\(\\prod\_\{h=t\}^\{t\+H\}\\beta\(a\_\{h\}\|z\_\{h\}\)\\right\)\\exp\\left\(\\frac\{G\(z\_\{t\},a\_\{t:t\+H\}\)\}\{\\kappa\}\\right\)\\prod\_\{h=t\}^\{t\+H\-1\}\\mathbb\{1\}\_\{\\\{z\_\{h\+1\}=\\mathcal\{F\}\(z\_\{h\},a\_\{h\}\)\\\}\},\(13\)where𝟙\{⋅\}\\mathbb\{1\}\_\{\\\{\\cdot\\\}\}denotes the indicator function enforcing the latent dynamics, and the partition function is omitted for brevity\. A detailed proof is provided in Appendix[C\.4](https://arxiv.org/html/2605.26282#A3.SS4)\. Theorem[3\.2](https://arxiv.org/html/2605.26282#S3.SS2)provides a unified perspective that bridges search and policy optimization: rather than explicitly parameterizing a policy network, optimization can be implicitly executed through iterative rollouts and score\-based sampling within the world model\.

Why Diffusion Policy?The diffusion policy intrinsically bridges score matching within the world model and true policy optimization, unlocking the scaling potential of world models through a dual mechanism\. First, the world model continuously bootstraps trajectories by generating imagined rollouts to facilitate score matching; crucially, as established in Theorem[3\.2](https://arxiv.org/html/2605.26282#S3.SS2), this score matching process is mathematically equivalent to performing true policy optimization\. Second, while a standard KL\-constrained objective in \([8](https://arxiv.org/html/2605.26282#S3.E8)\) is typically intractable due to the intractable density functionβ\\beta, this framework bypasses the intractable integration by training an implicit energy function that is naturally absorbed into imagined trajectories within world models\.

Since the optimal policy follows the Gibbs distribution in \([13](https://arxiv.org/html/2605.26282#S3.E13)\), MBDPO directly addresses the bottlenecks identified in Problems[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)and[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)\. Specifically, it orchestrates both value iteration and policy improvement by ensuring that policy updates remain strictly bounded within a local trust region, thereby overcoming the suboptimality bottlenecks inherent in previous methods\.

### 3\.3Practical Algorithm

For the practical implementation of MBDPO, we adopt a joint\-training paradigm that synchronizes world model learning with diffusion policy optimization \(see Algorithm[1](https://arxiv.org/html/2605.26282#algorithm1)in Appendix[D](https://arxiv.org/html/2605.26282#A4)\)\. We maintain a replay bufferℬ\\mathcal\{B\}to store trajectory segments, supporting various training configurations including online, offline, and offline\-to\-online settings\. The framework is jointly trained through an alternating optimization scheme that iteratively updates \(i\) the world model and value iteration, and \(ii\) the diffusion policy optimization\.

\(i\) World Model and Value Iteration:Utilizing trajectory segments sampled from the bufferℬ\\mathcal\{B\}, we optimize the core components of the world model defined in \([1](https://arxiv.org/html/2605.26282#S2.E1)\), namely theencoderℰ\\mathcal\{E\},latent dynamicsℱ\\mathcal\{F\}, andreward predictorℛ\\mathcal\{R\}\. Simultaneously, thevalue functionQ^\\hat\{Q\}is updated via one\-step temporal difference \(TD\) learning, while theimplicit energyEθE\_\{\\theta\}is refined according to the objective in \([9](https://arxiv.org/html/2605.26282#S3.E9)\)\. The joint minimizer forℰ,ℱ,ℛ\\mathcal\{E,F,R\},Q^\\hat\{Q\}andEθE\_\{\\theta\}can be written as

min⁡𝔼τ∼ℬ​\[∑h=tt\+Hγh−t​\(‖ℱ​\(zh,ah\)−sg​\(ℰ​\(sh\+1\)\)‖22\+CE​\(r^h,r​\(sh,ah\)\)\+TD​\(Q^​\(zh,ah\)\)\+ℒ​\(Eθ\)\)\],\\min\\mathbb\{E\}\_\{\\tau\\sim\\mathcal\{B\}\}\\bigg\[\\sum\_\{h=t\}^\{t\+H\}\\gamma^\{h\-t\}\\bigg\(\\\|\\mathcal\{F\}\(z\_\{h\},a\_\{h\}\)\-\\text\{sg\}\(\\mathcal\{E\}\(s\_\{h\+1\}\)\)\\\|\_\{2\}^\{2\}\+\\text\{CE\}\(\\hat\{r\}\_\{h\},r\(s\_\{h\},a\_\{h\}\)\)\+\\text\{TD\}\(\\hat\{Q\}\(z\_\{h\},a\_\{h\}\)\)\+\\mathcal\{L\}\(E\_\{\\theta\}\)\\bigg\)\\bigg\],\(14\)wheresg​\(⋅\)\\text\{sg\}\(\\cdot\)denotes thestop\-gradoperator andCE​\(⋅,⋅\)\\text\{CE\}\(\\cdot,\\cdot\)represents the cross\-entropy loss, utilized to capture sparse or discrete rewards across diverse environments\. Following Lillicrap et al\.\(lillicrap2016continuous\), the one\-step TD loss is defined asTD​\(Q^​\(zh,ah\)\)≔CE​\(Q^​\(zh,ah\),rh\+γ​Q¯​\(zh\+1,ah\+1\)\)\\text\{TD\}\(\\hat\{Q\}\(z\_\{h\},a\_\{h\}\)\)\\coloneqq\\text\{CE\}\\big\(\\hat\{Q\}\(z\_\{h\},a\_\{h\}\),r\_\{h\}\+\\gamma\\bar\{Q\}\(z\_\{h\+1\},a\_\{h\+1\}\)\\big\), whereQ¯\\bar\{Q\}is the target network maintained as an exponential moving average \(EMA\) ofQ^\\hat\{Q\}\. The energy lossℒ​\(Eθ\)\\mathcal\{L\}\(E\_\{\\theta\}\)is optimized according to \([9](https://arxiv.org/html/2605.26282#S3.E9)\), using actions sampled from other latent states as negative pairs\. Other training configurations, such as normalization and encoder architectures, strictly follow the TD\-MPC framework to highlight the significance of scalable policy learning within world models\.

\(ii\) Diffusion Policy Optimization\.As shown in \([10](https://arxiv.org/html/2605.26282#S3.E10)\) and Theorem[3\.2](https://arxiv.org/html/2605.26282#S3.SS2), the policy optimization relies on the imagined samples from the learned world models under a weighted expectation \(red shadow of \([15](https://arxiv.org/html/2605.26282#S3.E15)\)\)\. Since the expectation cannot be computed analytically, we approximate it using Monte Carlo sampling\. MBDPO employs a weighted Monte Carlo estimator based on a proposal distributionq​\(at:t\+H0\|τ\|zt\)q\(a\_\{t:t\+H\}^\{0\|\\tau\}\|z\_\{t\}\), chosen as a Gaussian𝒩​\(1α¯τ​at:t\+Hτ,1−α¯τα¯τ​I\)\\mathcal\{N\}\(\\frac\{1\}\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}a\_\{t:t\+H\}^\{\\tau\},\\frac\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\{\\bar\{\\alpha\}^\{\\tau\}\}I\)according to \([11](https://arxiv.org/html/2605.26282#S3.E11)\)\. The asymptotically unbiased expectation can thus be approximated usingTTMonte Carlo samples\{at:t\+H0\|τ,\(i\)\}i=1T\\\{a\_\{t:t\+H\}^\{0\|\\tau,\(i\)\}\\\}\_\{i=1\}^\{T\}drawn fromq​\(at:t\+H0\|τ\|zt\)q\(a\_\{t:t\+H\}^\{0\|\\tau\}\|z\_\{t\}\):

ϕ​\(zt,at:t\+Hτ,τ\)≈−at:t\+Hτ1−α¯τ\+α¯τ1−α¯τ⋅∑i=1Tw​\(at:t\+H0\|τ,\(i\)\)⋅at:t\+H0\|τ,\(i\)\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)\\approx\-\\frac\{a^\{\\tau\}\_\{t:t\+H\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\+\\frac\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\\cdot\\sum\_\{i=1\}^\{T\}\\hbox\{\\pagecolor\{red\!10\}$w\(a^\{0\|\\tau,\(i\)\}\_\{t:t\+H\}\)$\}\\cdot a\_\{t:t\+H\}^\{0\|\\tau,\(i\)\}\(15\)withw​\(at:t\+H0\|τ,\(i\)\)=exp⁡\(G~​\(zt,at:t\+H0\|τ,\(i\)\)κ\)∑i=1Texp⁡\(G~​\(zt,at:t\+H0\|τ,\(i\)\)κ\)w\(a^\{0\|\\tau,\(i\)\}\_\{t:t\+H\}\)=\\frac\{\\exp\(\\frac\{\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\|\\tau,\(i\)\}\)\}\{\\kappa\}\)\}\{\\sum\_\{i=1\}^\{T\}\\exp\(\\frac\{\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\|\\tau,\(i\)\}\)\}\{\\kappa\}\)\}for samplesat:t\+H0\|τ,\(i\)∼q​\(at:t\+H0\|τ\|zt\)a\_\{t:t\+H\}^\{0\|\\tau,\(i\)\}\\sim q\(a\_\{t:t\+H\}^\{0\|\\tau\}\|z\_\{t\}\)\. These weights rely on imagined trajectories rolled out by the learned latent dynamicsℱ\\mathcal\{F\}to evaluate the cumulative returnG~\\tilde\{G\}, as shown in Figure[2](https://arxiv.org/html/2605.26282#S1.F2)\. To avoid the heavy computational overhead of online rollouts during deployment, this Monte Carlo estimation is exclusively employed during score network training, such that

min⁡𝔼​\[‖ϕ^​\(zt,at:t\+Hτ,τ\)−ϕ​\(zt,at:t\+Hτ,τ\)‖22\],\\min\\mathbb\{E\}\\big\[\\\|\\hat\{\\phi\}\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)\-\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)\\\|\_\{2\}^\{2\}\\big\],\(16\)whereϕ^​\(zt,at:t\+Hτ,τ\)\\hat\{\\phi\}\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)is the learned score function\. During inference, the score network serves as an amortized solution for efficient policy generation, pluggingϕ^\\hat\{\\phi\}into \([12](https://arxiv.org/html/2605.26282#S3.E12)\) for iterative sampling of the target action sequence\. This framework benefits from provable theoretical guarantees: as established in\(Chenet al\.,[2022](https://arxiv.org/html/2605.26282#bib.bib70); Chenget al\.,[2026](https://arxiv.org/html/2605.26282#bib.bib45)\), the error relative to the target Gibbs distribution in \([13](https://arxiv.org/html/2605.26282#S3.E13)\) decays at a polynomial ratePoly​\(N,T,Lϕ\)\\text\{Poly\}\(N,T,L\_\{\\phi\}\), whereN,T,N,T,andLϕL\_\{\\phi\}denote the diffusion steps, Monte Carlo sample size, and the score function’s Lipschitz constant, respectively\.

## 4Experiments

Benchmarks and Tasks\.We evaluate MBDPO across four major benchmarks encompassing 121 diverse control tasks: DMControl\(Tassaet al\.,[2018](https://arxiv.org/html/2605.26282#bib.bib56)\), MetaWorld\(Yuet al\.,[2020](https://arxiv.org/html/2605.26282#bib.bib57)\), ManiSkill2\(Guet al\.,[2023](https://arxiv.org/html/2605.26282#bib.bib58)\), and MyoSuite\(Caggianoet al\.,[2022](https://arxiv.org/html/2605.26282#bib.bib59)\)\. Visual demonstrations for each benchmark are provided in Figures[12](https://arxiv.org/html/2605.26282#A1.F12)–[15](https://arxiv.org/html/2605.26282#A1.F15)\. Our evaluation includes challenging variations and subsets such as Visual RL \(image\-based DMControl\), Locomotion, and PickYCB\. These tasks demand mastery over high\-dimensional state\-action spaces, sparse rewards, and multi\-object manipulation\. Furthermore, they include physiologically accurate musculoskeletal control and complex locomotion with Dog and Humanoid embodiments, covering a wide spectrum of difficulty levels\.

Baselines\.We compare MBDPO against several state\-of\-the\-art and data\-efficient RL algorithms: \(1\) Soft Actor\-Critic \(SAC\)\(Haarnojaet al\.,[2018](https://arxiv.org/html/2605.26282#bib.bib60)\), a representative off\-policy baseline; \(2\) DreamerV3\(Hafneret al\.,[2023](https://arxiv.org/html/2605.26282#bib.bib61)\), a leading world\-model\-based approach; \(3\) TD\-MPC\(Hansenet al\.,[2022](https://arxiv.org/html/2605.26282#bib.bib26)\), the first foundational work of the TD\-MPC series; and \(4\) TD\-MPC2\(Hansenet al\.,[2024](https://arxiv.org/html/2605.26282#bib.bib11)\), which shares its algorithmic core with the most recent advancements in the field\(Hansenet al\.,[2025](https://arxiv.org/html/2605.26282#bib.bib27)\)\.

Result Analysis\.Our experimental evaluation focuses on five key dimensions: \(1\) a comparative performance evaluation of MBDPO against leading model\-free and model\-based RL baselines; \(2\) an investigation into the scaling behaviors of MBDPO; \(3\) an empirical verification of whether MBDPO mitigates the bottlenecks identified in Problems[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)and[3\.1](https://arxiv.org/html/2605.26282#S3.SS1); \(4\) an assessment of the algorithm’s generalization across offline pre\-training, online, and offline\-to\-online settings; and \(5\) a comprehensive evaluation of the computational efficiency of MBDPO\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x4.png)Figure 4:Aggregate performance across four benchmarks in the online setting: DMControl, MetaWorld, ManiSkill2, and MyoSuite\.Detailed learning curves for each subtask are provided in Figures[16](https://arxiv.org/html/2605.26282#A1.F16)–[20](https://arxiv.org/html/2605.26282#A1.F20)and Appendix[A\.2](https://arxiv.org/html/2605.26282#A1.SS2)\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x5.png)Figure 5:Performance comparison between TD\-MPC2 and MBDPO across 10 visual control tasks\.Results are averaged over55random seeds for each task\.1\. MBDPO consistently achieves better results across all evaluated benchmarks, demonstrating significant performance gains over both model\-free and model\-based baselines\.As illustrated in the aggregate performance in Figure[1](https://arxiv.org/html/2605.26282#S0.F1), our algorithm consistently attains higher rewards in online settings\. Notably, in high\-dimensional visual control tasks, MBDPO significantly elevates the performance bottleneck compared to prior methods\. Another compelling observation is that across nearly all tasks, MBDPO exhibits accelerated policy improvement with substantially lower variance during the early training stages \(see Figures[4](https://arxiv.org/html/2605.26282#S4.F4)and[5](https://arxiv.org/html/2605.26282#S4.F5)\)\. We attribute this advantage to our principled policy optimization framework, which, unlike pure search methods based on a sampling\-based planner, provides a more stable basis for diffusion policy refinement from the onset of training \(seamlessly improved with world models and value functions\)\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x6.png)Figure 6:Massively multi\-task world models in the offline pretraining setting\.Normalized score as a function of model size on the two 80\-task and 30\-task datasets\. MBDPO shows sharper scaling behavior with model capacity\.2\. MBDPO unlocks the potential of world models for policy learning, enabling effective multi\-task offline pretraining and demonstrating a monotonic scaling curve as model capacity scales from1\.71\.7M to340340M\.As illustrated in Figure[6](https://arxiv.org/html/2605.26282#S4.F6), several key scaling properties emerge\. Notably, MBDPO demonstrates a more pronounced advantage at smaller model scales\. We attribute this performance to our principled policy optimization, which effectively regularizes the policy \(e\.g\., via the implicit KL constraint with the behavioral policy\), ensuring stable performance even when the world model is not highly accurate\. As model capacity increases, the performance gap between MBDPO and TD\-MPC2 gradually narrows\. This is likely because the enhanced accuracy of a larger world model significantly benefits the trajectory\-based policy search used in TD\-MPC2\. In contrast, MBDPO focuses on direct diffusion policy optimization, which consistently evolves alongside the world model\. Interestingly, on the 30\-task dataset with340340M parameters, the performance of both methods appears to decrease\. We hypothesize that at this scale, the dataset capacity has been potentially exhausted by the model, leading to saturated performance\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x7.png)Figure 7:Overview of offline\-to\-online \(O2O\) performance\.\(Left\) Comparison between MBDPO and TD\-MPC baselines\. \(Right\) Comparison between training from scratch and O2O fine\-tuning\. Results show that fine\-tuning a generalist agent yields superior performance with significantly less data, highlighting the high sample efficiency and transferability of our framework\. Note: The suboptimal results in “Hopper Hop” can be attributed to a severe deviation in the representation’s data distribution \(as illustrated by the latent embeddings in Figure[23](https://arxiv.org/html/2605.26282#A1.F23)\), which subsequently impedes effective policy optimization\.3\. MBDPO mitigates the fundamental bottlenecks \(Problems[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)and[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)\) induced by the misalignment between search and value learning\.We empirically validate this by monitoring both the cross TD error and the average relative policy drift throughout training, as illustrated in Figure[3](https://arxiv.org/html/2605.26282#S3.F3)\. Crucially, MBDPO exhibits a significantly lower and more stable cross TD error compared to TD\-MPC2\. This directly demonstrates that our framework effectively minimizes the value iteration error described in \([4](https://arxiv.org/html/2605.26282#S3.E4)\) \(Problem[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)\), preventing the vicious cycle of value error amplification in out\-of\-distribution regions\. Concurrently, we evaluate the policy drift, quantified as the log\-ratio between the search\-based distributionπ\\piand the current policy networkβ\\beta\. As training progresses, our diffusion policy regularized by the implicit energy function exhibits minimal volatility and smoothly converges toward the optimized distribution \(approaching zero after11M steps\)\. In sharp contrast, heuristic search via MPPI maintains high volatility and severe policy drift throughout the entire training stage\. This persistent misalignment confirms that unconstrained heuristic search fails to achieve true, stable policy optimization, leading to the distribution shift analyzed in \([5](https://arxiv.org/html/2605.26282#S3.E5)\) \(Problem[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)\)\. This phenomenon is particularly magnified in offline settings \(Figure[6](https://arxiv.org/html/2605.26282#S4.F6)\), where MBDPO maintains a highly stable optimization path even when deployed with limited\-capacity world models\. By constraining the KL divergence through the implicit energy, MBDPO suppresses both the value iteration gap and policy distributional drift, achieving a principled policy optimization that is highly congruent with our theoretical analysis in Theorem[3\.2](https://arxiv.org/html/2605.26282#S3.SS2)\.

Table 1:Profiling training efficiency results for online\-from\-scratch and multi\-task offline pretraining\.\(a\) Online training time is measured per environment step \(milliseconds\) on a single NVIDIA A800\-80GB, see the task\-specific results in Table[2](https://arxiv.org/html/2605.26282#A1.T2)\. \(b\) Offline training time is reported in hours for 10M training steps on a single NVIDIA\-H200\. Values are reported as mean±\\pmstd\.\(a\) Online\-from\-scratch training timeDomain\#TD\-MPC2MBDPOMetaWorld501\.08±0\.161\.08\\pm 0\.161\.08±0\.151\.08\\pm 0\.15DMControl390\.78±0\.710\.78\\pm 0\.710\.76±0\.710\.76\\pm 0\.71ManiSkill256\.50±1\.836\.50\\pm 1\.836\.91±2\.516\.91\\pm 2\.51MyoSuite101\.58±0\.221\.58\\pm 0\.221\.54±0\.191\.54\\pm 0\.19Visual RL103\.55±0\.513\.55\\pm 0\.513\.04±0\.373\.04\\pm 0\.37
\(b\) Multi\-task offline training timeParameters30 tasks80 tasks80/301\.7M34\.8±0\.634\.8\\pm 0\.639\.9±0\.439\.9\\pm 0\.41\.15×1\.15\\times6M35\.0±0\.535\.0\\pm 0\.540\.5±0\.640\.5\\pm 0\.61\.16×1\.16\\times21M39\.8±0\.339\.8\\pm 0\.342\.6±0\.442\.6\\pm 0\.41\.07×1\.07\\times54M52\.6±0\.152\.6\\pm 0\.169\.6±0\.769\.6\\pm 0\.71\.32×1\.32\\times340M202\.9±0\.9202\.9\\pm 0\.9206\.8±0\.8206\.8\\pm 0\.81\.02×1\.02\\times

4\. MBDPO is highly versatile, providing unified support for offline pre\-training, online learning, and offline\-to\-online \(O2O\) transfer settings\.Across these diverse paradigms, MBDPO consistently achieves superior performance, as evidenced in Figures[4](https://arxiv.org/html/2605.26282#S4.F4),[6](https://arxiv.org/html/2605.26282#S4.F6), and[7](https://arxiv.org/html/2605.26282#S4.F7)\. Notably, our method demonstrates significant advantages in the offline regime, where stable optimization is more critical than in other settings\. To evaluate generalization, we extended MBDPO to an O2O fine\-tuning scenario—pre\-training on 70 tasks followed by fine\-tuning on 10 unseen tasks with only40​k40kinteractions per task\. We found that MBDPO significantly outperforms training from scratch, verifying that our pre\-trained generalist agent can rapidly adapt to become an expert in novel environments\. This enhanced efficiency in policy optimization underscores MBDPO’s effectiveness across varied scenarios\. Specifically, the comparison between from\-scratch and fine\-tuning curves in Figure[7](https://arxiv.org/html/2605.26282#S4.F7)\(right\) demonstrates that our framework can be seamlessly extended to new tasks with minimal data requirements\.

5\. MBDPO achieves efficient training during both the online\-from\-scratch and multi\-task offline paradigms\.According to Table[1](https://arxiv.org/html/2605.26282#S4.T1), MBDPO incurs negligible overhead in online\-from\-scratch training\. Across121121tasks from five domains, its per\-step training time is comparable to TD\-MPC2, and is faster on DMControl, MyoSuite, and Visual RL, reducing the Visual RL average from3\.553\.55ms to3\.043\.04ms\. For multi\-task offline training, MBDPO scales efficiently from 30 to 80 tasks, with only a1\.02×1\.02\\times–1\.32×1\.32\\timesincrease in total training time\. Notably, the 340M model shows nearly unchanged cost \(202\.9202\.9vs\.206\.8206\.8hours\), indicating strong scalability to broader multi\-task settings\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x8.png)Figure 8:Ablation study of the factorη\\eta\.Experiments are conducted on the 80\-task multi\-task setting with a 21M parameters model, whereη\\etais varied within the range\[0,5\]\[0,5\]\.Ablation Study 1: Regularized Factorη\\eta\.Figure[8](https://arxiv.org/html/2605.26282#S4.F8)demonstrates the importance of the implicit energy functionEEin diffusion policy optimization\. In the absence of this energy function, the policy lacks an effective KL constraint, leading to suboptimal performance as identified in Problems[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)and[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)\. While the performance remains stable forη\\etawithin the range\[0\.3,1\]\[0\.3,1\], scalingη\\etaup to55causes the performance to drop below 70\. This degradation occurs because an excessively largeη\\etainduces an overly dominant behavioral cloning effect, which ultimately restricts the policy’s exploitation of higher\-value regions within the learned world model\. An interesting phenomenon is that even when the regularized factor is set toη=0\\eta=0, our method still outperforms TD\-MPC2 \(75\.29 vs\. 57\.1\)\. This reflects that diffusion policy optimization, by its nature, can lead to a better policy than heuristic search\.

Ablation Study 2: Diffusion Denoise TimestepsNNand Monte Carlo SamplesTT\.We conduct ablation studies on several challenging tasks from DMControl and MetaWorld, with results illustrated in Figure[9](https://arxiv.org/html/2605.26282#S4.F9)and Figure[10](https://arxiv.org/html/2605.26282#S4.F10)\. The findings are straightforward: increasing the number of diffusion timesteps and Monte Carlo samples consistently improves performance\. This aligns with our theoretical error analysis—specifically, the polynomially decaying errorPoly​\(N,T,Lϕ\)\\text\{Poly\}\(N,T,L\_\{\\phi\}\)—confirming that the distributional error decreases polynomially with both parameters\. Conversely, because sampling is executed via a discrete timestep, inference runtime scales linearly with the number of diffusion timesteps\. However, due to the parallelizable nature of Monte Carlo sampling, increasingTThas only a marginal impact on execution time\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x9.png)Figure 9:The ablation study of Monte Carlo samples in the policy with 3 random seeds\.\(Left\) Training runtime comparison under different sample numbers\. \(Right\) Episode reward versus training steps and sample numbers\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x10.png)Figure 10:The ablation study of diffusion denoise timesteps in the policy with 3 random seeds\.\(Left\) Training runtime comparison under different numbers of diffusion timesteps\. \(Right\) Episode reward versus training steps for different diffusion timesteps\.
## 5What We Learned?

1\. Scalability of Diffusion Policy Built Upon World Model\.Our primary insight is that diffusion models act as an intrinsic interface bridging world models and policy optimization\. By leveraging imagined trajectories, the diffusion policy can bootstrap its behaviors to achieve continual value iteration and policy improvement\. Crucially, we find that the joint learning of the world model and RL policy under a simple loss function enforces a stronger capture of underlying causality within the latent space\. This emergent causality is distinctly mirrored in our latent visualizations, where the policy naturally uncovers periodic closed\-loop manifolds for cyclical tasks and goal\-oriented trajectories for sequential manipulations\. Remarkably, even when evaluated under controlled and intervened latent rollouts, the diffusion policy consistently maintains these physically aligned behaviors\. This empirical evidence suggests that the joint optimization in MBDPO helps structure the latent space in a way that better adheres to underlying physical constraints, making the policy less prone to rendering physically implausible decisions when encountering distribution shifts\. Looking forward, as the world model achieves sufficient fidelity, this paradigm opens up a promising avenue toward fully self\-improving autonomous agents\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x11.png)Figure 11:Visualizing latent trajectories via a locally linear embedding\.We plot the latent state trajectories across single and multiple episodes in various simulated environments\. The colorbar gradient indicates the temporal progression of the trajectories, while R and SR denote the accumulated reward and success rate, respectively\. \(a\) For cyclical tasks such as “Cheetah Run Front”, “Reacher Hard”, and “Cup Spin”, MBDPO with the diffusion policy generates distinctly closed\-loop structures, effectively capturing the physical and repetitive nature of these skills\. For other MetaWorld manipulation tasks, MBDPO’s latent states delineate smooth and directed state\-progressing trajectories that represent a clear causal transition from initial state to target goal\. \(b\) In contrast, under the same controlled latent trajectory visualization, TD\-MPC2 exhibits noisier, more disorganized, and less physically aligned behaviors\. This indicates that its sampled policy is not physically meaningful, failing to adhere to the underlying causal dynamics necessary for robust decision\-making\. The controlled latent trajectories for8080tasks are visualized in Appendix Figures[24](https://arxiv.org/html/2605.26282#A1.F24)\-[27](https://arxiv.org/html/2605.26282#A1.F27)\.2\. Rethinking Policy Learning in World Models\.Our second higher\-level insight compels a fundamental rethinking of how RL should be formulated within world models to truly unlock their future large\-scale scaling potential\. As mathematically exposed in Problems[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)and[3\.1](https://arxiv.org/html/2605.26282#S3.SS1), simply scaling up naive bootstrapping or unconstrained rollout search \(such as MPPI, CEM\(Pinneriet al\.,[2021](https://arxiv.org/html/2605.26282#bib.bib75)\)\) is inherently problematic; instead of compounding performance gains, it merely aggregates value overestimation and action drift \(Figure[3](https://arxiv.org/html/2605.26282#S3.F3)\)\. What we truly learned is that the path toward scaling policy learning lies not in designing more complex heuristic search pipelines, but in establishing a simple but principled foundation that inherently links world model with foundational RL paradigms\. Crucially, the KL\-constrained objective in \([8](https://arxiv.org/html/2605.26282#S3.E8)\) serves as a direct mathematical bridge\. On one hand, its implicit energy function suppresses out\-of\-distribution value peaks, capturing the essential conservatism of offline RL\(Levineet al\.,[2020](https://arxiv.org/html/2605.26282#bib.bib72)\)\. On the other hand, it bounds the policy shift within a KL trust region to ensure policy improvement, echoing the core principle of TRPO and PPO\(Schulmanet al\.,[2015](https://arxiv.org/html/2605.26282#bib.bib69),[2017](https://arxiv.org/html/2605.26282#bib.bib73)\)\. Ultimately, this approach reveals that these seemingly separate methods actually point back to the same fundamental challenges of RL\.

3\. Representation Matters\.Our third insight underscores that the quality of the learned latent representation serves as the bedrock for policy optimization, a dependency that manifests through three critical phenomena\. First, a robust, generalist representation unlocks exceptional transferability and data efficiency\. As demonstrated in our offline\-to\-online \(O2O\) transfer experiments, pre\-training MBDPO on 70 diverse tasks constructs an adaptable latent space; this allows the generalist agent to rapidly fine\-tune and become an expert on 10 completely unseen tasks with only40​k40kinteractions per task, significantly outperforming training from scratch \(Figure[7](https://arxiv.org/html/2605.26282#S4.F7), right\)\. Second, conversely, a corrupted or misaligned representation directly and severely cripples policy optimization\. As isolated in the suboptimal performance of “Hopper Hop”, a severe distributional deviation in the latent embedding \(Figure[23](https://arxiv.org/html/2605.26282#A1.F23)\) imposes a fundamental bottleneck that completely impedes the diffusion policy from searching or sampling effectively\. Third, during large\-scale pre\-training, we observe that the controlled latent trajectories evolve to align with physical intuition \(see Figure[11](https://arxiv.org/html/2605.26282#S5.F11)\)\. This empirical alignment proves that scaling the representation is not merely about preserving training statistics, but about capturing the underlying structure of the world\. Ultimately, these findings reveal that policy versatility across offline, online, and O2O paradigms is intrinsically a representation problem, where a physically meaningful latent space is the premier prerequisite for scalable policy learning\.

## 6Conclusion

In this work, we addressed a critical yet underexplored bottleneck in model\-based RL: the structural misalignment between policy search and value learning\. To overcome this limitation, we introduced Model\-Based Diffusion Policy Optimization \(MBDPO\)\. In this framework, the diffusion policy intrinsically links the world model and policy optimization, unlocking the potential of world models for policy learning\. By leveraging an implicit energy function to refine the score field, MBDPO anchors policy updates within a trust region, effectively mitigating training inconsistencies and compounding errors\. Our extensive evaluations across multi\-task offline pretraining, online learning, and offline\-to\-online fine\-tuning demonstrate that MBDPO consistently outperforms leading baselines\. Notably, MBDPO exhibits strong scalability, yielding monotonic performance gains as model capacity increases\. Beyond quantitative improvements, visualizations of the controlled latent trajectories generated by MBDPO align naturally with physical intuition\. Ultimately, MBDPO establishes an effective new paradigm for policy learning within world models\.

Limitations and Future Directions\.While MBDPO demonstrates strong scaling properties and performance, several avenues remain for future exploration\. First, due to bounded computational resources, we have not yet scaled our framework to extremely large parameter sizes or massive foundation\-model scales; investigating its full scaling limits remains a direct next step\. Second, our evaluations are currently focused on simulated environments\. Moving forward, deploying MBDPO on real\-world robots will be critical to testing its robust transferability and practical viability\. Finally, an exciting prospective direction is to integrate MBDPO with even more powerful, pretrained representations \(e\.g\., visual or multimodal foundation models as encoders\)\. Leveraging such richer visual priors could improve the latent representation and unlock truly large\-scale RL within world models\.

## References

- \[1\]\(2013\)Stability theory of differential equations\.Courier Corporation\.Cited by:[§B\.2](https://arxiv.org/html/2605.26282#A2.SS2.p2.4)\.
- \[2\]V\. Caggiano, H\. Wang, G\. Durandau, M\. Sartori, and V\. Kumar\(2022\)MyoSuite–a contact\-rich simulation suite for musculoskeletal motor control\.arXiv preprint arXiv:2205\.13600\.Cited by:[§A\.1](https://arxiv.org/html/2605.26282#A1.SS1.p1.1),[§4](https://arxiv.org/html/2605.26282#S4.p1.1)\.
- \[3\]B\. Calli, A\. Singh, A\. Walsman, S\. Srinivasa, P\. Abbeel, and A\. M\. Dollar\(2015\)The ycb object and model set: towards common benchmarks for manipulation research\.In2015 international conference on advanced robotics \(ICAR\),pp\. 510–517\.Cited by:[Figure 18](https://arxiv.org/html/2605.26282#A1.F18),[Figure 18](https://arxiv.org/html/2605.26282#A1.F18.4.2.2)\.
- \[4\]J\. Cen, C\. Yu, H\. Yuan, Y\. Jiang, S\. Huang, J\. Guo, X\. Li, Y\. Song, H\. Luo, F\. Wang,et al\.\(2025\)Worldvla: towards autoregressive action world model\.arXiv preprint arXiv:2506\.21539\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[5\]W\. Chang, M\. Henaff, B\. Amos, G\. Dudek, and S\. Fujimoto\(2026\)The surprising difficulty of search in model\-based reinforcement learning\.arXiv preprint arXiv:2601\.21306\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1),[§1](https://arxiv.org/html/2605.26282#S1.p3.1),[§3\.1](https://arxiv.org/html/2605.26282#S3.SS1.p1.2)\.
- \[6\]S\. Chen, S\. Chewi, J\. Li, Y\. Li, A\. Salim, and A\. R\. Zhang\(2022\)Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions\.arXiv preprint arXiv:2209\.11215\.Cited by:[§3\.3](https://arxiv.org/html/2605.26282#S3.SS3.p3.14)\.
- \[7\]X\. Cheng, X\. Tang, and Y\. Yang\(2025\)Safe and stable control via lyapunov\-guided diffusion models\.arXiv preprint arXiv:2509\.25375\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p10.1)\.
- \[8\]X\. Cheng, W\. Yuan, B\. Li, Y\. Xu, Y\. Yang, H\. Liang, B\. Peng, R\. Loftin, Z\. Sun, and Y\. Hu\(2026\)How does the lagrangian guide safe reinforcement learning through diffusion models?\.arXiv preprint arXiv:2602\.02924\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p10.1),[§3\.3](https://arxiv.org/html/2605.26282#S3.SS3.p3.14)\.
- \[9\]C\. Chi, Z\. Xu, S\. Feng, E\. Cousineau, Y\. Du, B\. Burchfiel, R\. Tedrake, and S\. Song\(2025\)Diffusion policy: visuomotor policy learning via action diffusion\.The International Journal of Robotics Research44\(10\-11\),pp\. 1684–1704\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p10.1)\.
- \[10\]S\. Curi, F\. Berkenkamp, and A\. Krause\(2020\)Efficient model\-based reinforcement learning through optimistic policy search and planning\.Advances in Neural Information Processing Systems33,pp\. 14156–14170\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[11\]M\. Deisenroth and C\. E\. Rasmussen\(2011\)PILCO: a model\-based and data\-efficient approach to policy search\.InProceedings of the 28th International Conference on machine learning \(ICML\-11\),pp\. 465–472\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[12\]X\. Dong, J\. Cheng, and X\. S\. Zhang\(2025\)Maximum entropy reinforcement learning with diffusion policy\.arXiv preprint arXiv:2502\.11612\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p10.1)\.
- \[13\]J\. Farebrother, M\. Pirotta, A\. Tirinzoni, R\. Munos, A\. Lazaric, and A\. Touati\(2025\)Temporal difference flows\.arXiv preprint arXiv:2503\.09817\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[14\]J\. W\. Forrester\(1971\)Counterintuitive behavior of social systems\.Theory and decision2\(2\),pp\. 109–140\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[15\]H\. Fu, S\. Yu, M\. Littman, and G\. Konidaris\(2022\)Model\-based lifelong reinforcement learning with bayesian exploration\.Advances in Neural Information Processing Systems35,pp\. 32369–32382\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[16\]S\. Fujimoto, P\. D’Oro, A\. Zhang, Y\. Tian, and M\. Rabbat\(2025\)Towards general\-purpose model\-free reinforcement learning\.arXiv preprint arXiv:2501\.16142\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[17\]S\. Fujimoto, D\. Meger, and D\. Precup\(2019\)Off\-policy deep reinforcement learning without exploration\.InInternational conference on machine learning,pp\. 2052–2062\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p3.1)\.
- \[18\]Q\. Garrido, M\. Assran, N\. Ballas, A\. Bardes, L\. Najman, and Y\. LeCun\(2024\)Learning and leveraging world models in visual representation learning\.arXiv preprint arXiv:2403\.00504\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p1.1)\.
- \[19\]J\. Gu, F\. Xiang, X\. Li, Z\. Ling, X\. Liu, T\. Mu, Y\. Tang, S\. Tao, X\. Wei, Y\. Yao,et al\.\(2023\)Maniskill2: a unified benchmark for generalizable manipulation skills\.arXiv preprint arXiv:2302\.04659\.Cited by:[§A\.1](https://arxiv.org/html/2605.26282#A1.SS1.p1.1),[§4](https://arxiv.org/html/2605.26282#S4.p1.1)\.
- \[20\]L\. Guan, K\. Valmeekam, S\. Sreedharan, and S\. Kambhampati\(2023\)Leveraging pre\-trained large language models to construct and utilize world models for model\-based task planning\.Advances in Neural Information Processing Systems36,pp\. 79081–79094\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p1.1)\.
- \[21\]D\. Ha and J\. Schmidhuber\(2018\)World models\.arXiv preprint arXiv:1803\.101222\(3\),pp\. 440\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[22\]T\. Haarnoja, A\. Zhou, K\. Hartikainen, G\. Tucker, S\. Ha, J\. Tan, V\. Kumar, H\. Zhu, A\. Gupta, P\. Abbeel,et al\.\(2018\)Soft actor\-critic algorithms and applications\.arXiv preprint arXiv:1812\.05905\.Cited by:[§4](https://arxiv.org/html/2605.26282#S4.p2.1)\.
- \[23\]D\. Hafner, T\. Lillicrap, I\. Fischer, R\. Villegas, D\. Ha, H\. Lee, and J\. Davidson\(2019\)Learning latent dynamics for planning from pixels\.InInternational conference on machine learning,pp\. 2555–2565\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1),[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[24\]D\. Hafner, T\. Lillicrap, M\. Norouzi, and J\. Ba\(2020\)Mastering atari with discrete world models\.arXiv preprint arXiv:2010\.02193\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[25\]D\. Hafner, J\. Pasukonis, J\. Ba, and T\. Lillicrap\(2023\)Mastering diverse domains through world models\.arXiv preprint arXiv:2301\.04104\.Cited by:[§4](https://arxiv.org/html/2605.26282#S4.p2.1)\.
- \[26\]D\. Hafner, J\. Pasukonis, J\. Ba, and T\. Lillicrap\(2025\)Mastering diverse control tasks through world models\.Nature640\(8059\),pp\. 647–653\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p1.1),[§1](https://arxiv.org/html/2605.26282#S1.p9.1),[§2](https://arxiv.org/html/2605.26282#S2.p2.4)\.
- \[27\]D\. Hafner, W\. Yan, and T\. Lillicrap\(2025\)Training agents inside of scalable world models\.arXiv preprint arXiv:2509\.24527\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p1.1),[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[28\]J\. Han, M\. Tian, J\. Zhu, F\. He, H\. Zhang, S\. Guo, D\. Zhu, H\. Tang, P\. Xu, Y\. Guo,et al\.\(2025\)Percept\-wam: perception\-enhanced world\-awareness\-action model for robust end\-to\-end autonomous driving\.arXiv preprint arXiv:2511\.19221\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[29\]N\. Hansen, H\. Su, and X\. Wang\(2024\)Td\-mpc2: scalable, robust world models for continuous control\.arXiv preprint arXiv:2310\.16828\.Cited by:[Figure 1](https://arxiv.org/html/2605.26282#S0.F1),[Figure 1](https://arxiv.org/html/2605.26282#S0.F1.8.4.4),[§1](https://arxiv.org/html/2605.26282#S1.p1.1),[§1](https://arxiv.org/html/2605.26282#S1.p3.1),[§1](https://arxiv.org/html/2605.26282#S1.p9.1),[§2](https://arxiv.org/html/2605.26282#S2.p2.4),[§3\.2](https://arxiv.org/html/2605.26282#S3.SS2.p2.4),[§4](https://arxiv.org/html/2605.26282#S4.p2.1)\.
- \[30\]N\. Hansen, H\. Su, and X\. Wang\(2025\)Learning massively multitask world models for continuous control\.arXiv preprint arXiv:2511\.19584\.Cited by:[§C\.1](https://arxiv.org/html/2605.26282#A3.SS1.p1.1),[§1](https://arxiv.org/html/2605.26282#S1.p3.1),[§1](https://arxiv.org/html/2605.26282#S1.p9.1),[§2](https://arxiv.org/html/2605.26282#S2.p2.4),[§3\.1](https://arxiv.org/html/2605.26282#S3.SS1.p1.2),[§3\.2](https://arxiv.org/html/2605.26282#S3.SS2.p2.4),[§4](https://arxiv.org/html/2605.26282#S4.p2.1)\.
- \[31\]N\. Hansen, X\. Wang, and H\. Su\(2022\)Temporal difference learning for model predictive control\.arXiv preprint arXiv:2203\.04955\.Cited by:[§C\.1](https://arxiv.org/html/2605.26282#A3.SS1.p1.1),[§1](https://arxiv.org/html/2605.26282#S1.p3.1),[§1](https://arxiv.org/html/2605.26282#S1.p9.1),[§3\.1](https://arxiv.org/html/2605.26282#S3.SS1.p1.2),[§3\.2](https://arxiv.org/html/2605.26282#S3.SS2.p2.4),[§4](https://arxiv.org/html/2605.26282#S4.p2.1)\.
- \[32\]J\. Ho, A\. Jain, and P\. Abbeel\(2020\)Denoising diffusion probabilistic models\.Advances in neural information processing systems33,pp\. 6840–6851\.Cited by:[§B\.1](https://arxiv.org/html/2605.26282#A2.SS1.p1.2),[§2](https://arxiv.org/html/2605.26282#S2.p3.5)\.
- \[33\]P\. Intelligence, A\. A\. Bo, R\. Aniceto, A\. Balakrishna, G\. Balke, K\. Black, G\. Bokinsky, S\. Cao, T\. Charbonnier,et al\.\(2026\)π0\.7\\pi\_\{0\.7\}: A Steerable Generalist Robotic Foundation Model with Emergent Capabilities\.arXiv preprint arXiv:2604\.15483\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[34\]M\. Janner, Y\. Du, J\. B\. Tenenbaum, and S\. Levine\(2022\)Planning with diffusion for flexible behavior synthesis\.arXiv preprint arXiv:2205\.09991\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p10.1)\.
- \[35\]M\. Janner, J\. Fu, M\. Zhang, and S\. Levine\(2019\)When to trust your model: model\-based policy optimization\.Advances in neural information processing systems32\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1),[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[36\]N\. Lambert, K\. Pister, and R\. Calandra\(2022\)Investigating compounding prediction errors in learned dynamics models\.arXiv preprint arXiv:2203\.09637\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[37\]S\. Levine, A\. Kumar, G\. Tucker, and J\. Fu\(2020\)Offline reinforcement learning: tutorial, review, and perspectives on open problems\.arXiv preprint arXiv:2005\.01643\.Cited by:[§5](https://arxiv.org/html/2605.26282#S5.p2.1)\.
- \[38\]L\. Li, Q\. Zhang, Y\. Luo, S\. Yang, R\. Wang, F\. Han, M\. Yu, Z\. Gao, N\. Xue, X\. Zhu,et al\.\(2026\)Causal world modeling for robot control\.arXiv preprint arXiv:2601\.21998\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[39\]Q\. Li and S\. Levine\(2026\)Q\-learning with adjoint matching\.arXiv preprint arXiv:2601\.14234\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p10.1)\.
- \[40\]H\. Lin, P\. Wang, J\. Schneider, and G\. Shi\(2025\)TD\-m\(pc\)2: improving temporal difference mpc through policy constraint\.External Links:2502\.03550,[Link](https://arxiv.org/abs/2502.03550)Cited by:[§C\.1](https://arxiv.org/html/2605.26282#A3.SS1.p1.1),[§3\.1](https://arxiv.org/html/2605.26282#S3.SS1.p1.2)\.
- \[41\]T\. M\. Moerland, J\. Broekens, A\. Plaat, and C\. M\. Jonker\(2023\)Model\-based reinforcement learning: a survey\.Foundations and Trends in Machine Learning16\(1\),pp\. 1–118\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[42\]M\. Ma, T\. Ni, C\. Gehring, P\. D’Oro, and P\. Bacon\(2024\)Do transformer world models give better policy gradients?\.arXiv preprint arXiv:2402\.05290\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[43\]S\. Mahadevan\(1996\)Average reward reinforcement learning: foundations, algorithms, and empirical results\.Machine learning22\(1\),pp\. 159–195\.Cited by:[§C\.2](https://arxiv.org/html/2605.26282#A3.SS2.2.p2.2)\.
- \[44\]A\. Nagabandi, G\. Kahn, R\. S\. Fearing, and S\. Levine\(2018\)Neural network dynamics for model\-based deep reinforcement learning with model\-free fine\-tuning\.In2018 IEEE international conference on robotics and automation \(ICRA\),pp\. 7559–7566\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[45\]C\. Pan, Z\. Yi, G\. Shi, and G\. Qu\(2024\)Model\-based diffusion for trajectory optimization\.Advances in Neural Information Processing Systems37,pp\. 57914–57943\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p10.1)\.
- \[46\]S\. Park, K\. Frans, D\. Mann, B\. Eysenbach, A\. Kumar, and S\. Levine\(2025\)Horizon reduction makes rl scalable\.arXiv preprint arXiv:2506\.04168\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[47\]C\. Pinneri, S\. Sawant, S\. Blaes, J\. Achterhold, J\. Stueckler, M\. Rolinek, and G\. Martius\(2021\)Sample\-efficient cross\-entropy method for real\-time planning\.InConference on Robot Learning,pp\. 1049–1065\.Cited by:[§5](https://arxiv.org/html/2605.26282#S5.p2.1)\.
- \[48\]M\. Psenka, A\. Escontrela, P\. Abbeel, and Y\. Ma\(2023\)Learning a diffusion model policy from rewards via q\-score matching\.arXiv preprint arXiv:2312\.11752\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p10.1)\.
- \[49\]M\. Psenka, M\. Rabbat, A\. Krishnapriyan, Y\. LeCun, and A\. Bar\(2026\)Parallel stochastic gradient\-based planning for world models\.arXiv preprint arXiv:2602\.00475\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p1.1),[§3\.2](https://arxiv.org/html/2605.26282#S3.SS2.p2.4)\.
- \[50\]M\. L\. Puterman\(1990\)Markov decision processes\.Handbooks in operations research and management science2,pp\. 331–434\.Cited by:[§2](https://arxiv.org/html/2605.26282#S2.p1.16)\.
- \[51\]J\. Schulman, S\. Levine, P\. Abbeel, M\. Jordan, and P\. Moritz\(2015\)Trust region policy optimization\.InInternational conference on machine learning,pp\. 1889–1897\.Cited by:[§C\.2](https://arxiv.org/html/2605.26282#A3.SS2.3.p3.2),[§C\.2](https://arxiv.org/html/2605.26282#A3.SS2.SSS0.Px1.p2.6),[§5](https://arxiv.org/html/2605.26282#S5.p2.1)\.
- \[52\]J\. Schulman, F\. Wolski, P\. Dhariwal, A\. Radford, and O\. Klimov\(2017\)Proximal policy optimization algorithms\.arXiv preprint arXiv:1707\.06347\.Cited by:[§5](https://arxiv.org/html/2605.26282#S5.p2.1)\.
- \[53\]Y\. Shimizu and M\. Tomizuka\(2024\)Bisimulation metric for model predictive control\.arXiv preprint arXiv:2410\.04553\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p3.1)\.
- \[54\]Y\. Song, J\. Sohl\-Dickstein, D\. P\. Kingma, A\. Kumar, S\. Ermon, and B\. Poole\(2020\)Score\-based generative modeling through stochastic differential equations\.arXiv preprint arXiv:2011\.13456\.Cited by:[§B\.1](https://arxiv.org/html/2605.26282#A2.SS1.p1.2),[§C\.4](https://arxiv.org/html/2605.26282#A3.SS4.12.p12.3)\.
- \[55\]L\. K\. R\. Sutton\(1996\)Model\-based reinforcement learning with an approximate, learned model\.InProceedings of the ninth Yale workshop on adaptive and learning systems,Vol\.1996,pp\. 101–105\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p1.1)\.
- \[56\]R\. S\. Sutton, A\. G\. Barto,et al\.\(1998\)Reinforcement learning: an introduction\.Vol\.1,MIT press Cambridge\.Cited by:[§B\.2](https://arxiv.org/html/2605.26282#A2.SS2.p1.1),[§1](https://arxiv.org/html/2605.26282#S1.p1.1)\.
- \[57\]J\. SV, S\. Jalagam, Y\. LeCun, and V\. Sobal\(2023\)Gradient\-based planning with world models\.arXiv preprint arXiv:2312\.17227\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p1.1)\.
- \[58\]Y\. Tassa, Y\. Doron, A\. Muldal, T\. Erez, Y\. Li, D\. d\. L\. Casas, D\. Budden, A\. Abdolmaleki, J\. Merel, A\. Lefrancq,et al\.\(2018\)Deepmind control suite\.arXiv preprint arXiv:1801\.00690\.Cited by:[§A\.1](https://arxiv.org/html/2605.26282#A1.SS1.p1.1),[§4](https://arxiv.org/html/2605.26282#S4.p1.1)\.
- \[59\]G\. Team, B\. Wang, C\. Ni, G\. Huang, G\. Zhao, H\. Li, J\. Li, J\. Lv, J\. Liu, L\. Feng,et al\.\(2026\)GigaBrain\-0\.5 m\*: a vla that learns from world model\-based reinforcement learning\.arXiv preprint arXiv:2602\.12099\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[60\]S\. Thrun, K\. Möller, and A\. Linden\(1990\)Planning with an adaptive world model\.Advances in neural information processing systems3\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p1.1)\.
- \[61\]L\. Van der Maaten and G\. Hinton\(2008\)Visualizing data using t\-sne\.\.Journal of machine learning research9\(11\)\.Cited by:[Figure 23](https://arxiv.org/html/2605.26282#A1.F23.1.1),[Figure 23](https://arxiv.org/html/2605.26282#A1.F23.2.1)\.
- \[62\]H\. P\. Van Hasselt, M\. Hessel, and J\. Aslanides\(2019\)When to use parametric models in reinforcement learning?\.Advances in Neural Information Processing Systems32\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p2.1)\.
- \[63\]Y\. Wang, H\. Guo, S\. Wang, L\. Qian, and X\. Lan\(2025\)Bootstrapped model predictive control\.arXiv preprint arXiv:2503\.18871\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p3.1)\.
- \[64\]G\. Williams, A\. Aldrich, and E\. A\. Theodorou\(2017\)Model predictive path integral control: from theory to parallel computation\.Journal of Guidance, Control, and Dynamics40\(2\),pp\. 344–357\.Cited by:[2nd item](https://arxiv.org/html/2605.26282#A3.I1.i2.p1.1)\.
- \[65\]G\. Williams, A\. Aldrich, and E\. Theodorou\(2015\)Model predictive path integral control using covariance variable importance sampling\.arXiv preprint arXiv:1509\.01149\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p3.1)\.
- \[66\]P\. Wu, A\. Escontrela, D\. Hafner, P\. Abbeel, and K\. Goldberg\(2023\)Daydreamer: world models for physical robot learning\.InConference on robot learning,pp\. 2226–2240\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p1.1)\.
- \[67\]H\. Xue, C\. Pan, Z\. Yi, G\. Qu, and G\. Shi\(2025\)Full\-order sampling\-based mpc for torque\-level locomotion control via diffusion\-style annealing\.In2025 IEEE International Conference on Robotics and Automation \(ICRA\),pp\. 4974–4981\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p10.1)\.
- \[68\]S\. Ye, Y\. Ge, K\. Zheng, S\. Gao, S\. Yu, G\. Kurian, S\. Indupuru, Y\. L\. Tan, C\. Zhu, J\. Xiang, A\. Malik, K\. Lee, W\. Liang, N\. Ranawaka, J\. Gu, Y\. Xu, G\. Wang, F\. Hu, A\. Narayan, J\. Bjorck, J\. Wang, G\. Kim, D\. Niu, R\. Zheng, Y\. Xie, J\. Wu, Q\. Wang, R\. Julian, D\. Xu, Y\. Du, Y\. Chebotar, S\. Reed, J\. Kautz, Y\. Zhu, L\. "\. Fan, and J\. Jang\(2026\)World action models are zero\-shot policies\.External Links:2602\.15922,[Link](https://arxiv.org/abs/2602.15922)Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[69\]T\. Yu, D\. Quillen, Z\. He, R\. Julian, K\. Hausman, C\. Finn, and S\. Levine\(2020\)Meta\-world: a benchmark and evaluation for multi\-task and meta reinforcement learning\.InConference on robot learning,pp\. 1094–1100\.Cited by:[§A\.1](https://arxiv.org/html/2605.26282#A1.SS1.p1.1),[§4](https://arxiv.org/html/2605.26282#S4.p1.1)\.
- \[70\]G\. Zhan, L\. Wang, X\. Zhang, J\. Gao, M\. Tomizuka, and S\. E\. Li\(2025\)Bootstrap off\-policy with world model\.arXiv preprint arXiv:2511\.00423\.Cited by:[§C\.1](https://arxiv.org/html/2605.26282#A3.SS1.p1.1),[§1](https://arxiv.org/html/2605.26282#S1.p3.1),[§3\.1](https://arxiv.org/html/2605.26282#S3.SS1.p1.2),[§3\.2](https://arxiv.org/html/2605.26282#S3.SS2.p2.4)\.
- \[71\]H\. Zhen, X\. Qiu, P\. Chen, J\. Yang, X\. Yan, Y\. Du, Y\. Hong, and C\. Gan\(2024\)3d\-vla: a 3d vision\-language\-action generative world model\.arXiv preprint arXiv:2403\.09631\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p9.1)\.
- \[72\]G\. Zhou, H\. Pan, Y\. LeCun, and L\. Pinto\(2024\)Dino\-wm: world models on pre\-trained visual features enable zero\-shot planning\.arXiv preprint arXiv:2411\.04983\.Cited by:[§1](https://arxiv.org/html/2605.26282#S1.p1.1)\.

## Notation

NotationMeaningaaactioneelearnable task embeddingssstatetttime steprrreward functionzzlatent state𝒜\\mathcal\{A\}space of actionEEimplicit energy functionℰ\\mathcal\{E\}encoder of world modelℰe​n​v\\mathcal\{E\}\_\{env\}space of learnable task embeddingℱ\\mathcal\{F\}latent dynamics of world modelGGcumulative return over anHH\-step action rolloutG~\\tilde\{G\}energy\-regularized cumulative return over anHH\-step action rolloutJJstandard discounted returnsℳ\\mathcal\{M\}Markov decision process𝒮\\mathcal\{S\}space of state𝒯\\mathcal\{T\}transition dynamicsQQQQvalue functionQ^\\hat\{Q\}value function learned from buffer datasetℛ\\mathcal\{R\}latent reward functionℛ~\\tilde\{\\mathcal\{R\}\}energy\-regularized latent reward functionT□T^\{\\Box\}Bellman update𝒵\\mathcal\{Z\}space of latent stateα,α¯\\alpha,\\bar\{\\alpha\}scalar factor of noise injectionβ\\betabehavioral policy distributionκ\\kappatemperature factorπ\\pipolicyη\\etaenergy regularized factorρ\\rhoinitial distributionσ2​\(τ\)\\sigma^\{2\}\(\\tau\)variance induced by the reverse diffusion transition kernelπϕ\\pi\_\{\\phi\}policy generated by score functionϕ\\phidenoted symbol of score functionγ\\gammadiscount factor𝟙\{⋅\}\\mathbb\{1\}\_\{\\\{\\cdot\\\}\}indicator function

## A Reader’s Guide to the Appendix

Appendices can often feel like a dense maze of mathematical proofs and exhaustive plots\. To respect your time and make navigating this supplementary material as seamless as possible, we have organized the content into distinct, self\-contained modules\. Depending on your background and primary interests, here is a quick roadmap to help you find exactly what you are looking for:

- •For visual intuitions and empirical rigor:If you want to see how MBDPO behaves across different environments, head to Section[A\.2](https://arxiv.org/html/2605.26282#A1.SS2)for exhaustive per\-task learning curves and action drift metrics\. To visualize what the model actually learns, Section[A\.5](https://arxiv.org/html/2605.26282#A1.SS5)provides 3D manifolds showing how our method forms clean, closed\-loop latent trajectories compared to baselines\. Exhaustive compute profiling and timing evaluations are also provided to ensure full transparency in Section[A\.6](https://arxiv.org/html/2605.26282#A1.SS6)\.
- •For the theoretical foundation:If you are interested in the mathematical mechanisms, jump straight to Section[C](https://arxiv.org/html/2605.26282#A3)\. We break down the intrinsic bottlenecks of existing model\-based policy search paradigms \(Section[C\.1](https://arxiv.org/html/2605.26282#A3.SS1)\) and provide step\-by\-step proofs for the exact score function derivations and Theorem[3\.2](https://arxiv.org/html/2605.26282#S3.SS2)\(Section[C\.3](https://arxiv.org/html/2605.26282#A3.SS3)\)\.
- •For reproduction:If you are looking to implement MBDPO or understand its computational footprint, Section[D](https://arxiv.org/html/2605.26282#A4)provides the complete pseudocode \(Algorithm[1](https://arxiv.org/html/2605.26282#algorithm1)\) and hyperparameter settings alongside practical details for online, offline, and fine\-tuning variants\.

Please feel free to bypass the sections outside your immediate scope and jump directly to the modules most relevant to your expertise\.

## Appendix AMore Experimental Results

### A\.1Training Environments

Our evaluation covers 121 diverse control challenges from four standard benchmarks: DMControl, MetaWorld, ManiSkill2, and MyoSuite\[[58](https://arxiv.org/html/2605.26282#bib.bib56),[69](https://arxiv.org/html/2605.26282#bib.bib57),[19](https://arxiv.org/html/2605.26282#bib.bib58),[2](https://arxiv.org/html/2605.26282#bib.bib59)\]\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x12.png)Figure 12:Demonstration of DMControl tasks\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x13.png)Figure 13:Demonstration of MetaWorld tasks\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x14.png)Figure 14:Demonstration of ManiSkill2 tasks\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x15.png)Figure 15:Demonstration of MyoSuite tasks\.
### A\.2Additional Results of Online Training from Scratch

We evaluate MBDPO across 121 tasks from 4 major benchmarks \(Figures[16](https://arxiv.org/html/2605.26282#A1.F16)\-[20](https://arxiv.org/html/2605.26282#A1.F20)\), with each task averaged over 3 random seeds\. The aggregate results in Figure[1](https://arxiv.org/html/2605.26282#S0.F1)show that MBDPO consistently achieves state\-of\-the\-art performance across all environments\. Compared to both model\-free and model\-based baselines, our framework demonstrates superior final rewards, higher sample efficiency, and greater training stability of policy optimization\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x16.png)Figure 16:Single\-task DMControl results\.Episode return as a function of environment steps\. The first44M environment steps are shown for each task\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x17.png)Figure 17:Single\-task MetaWorld results\.Success rate \(%\) as a function of environment steps\. MBDPO achieves the best averaged performance over all MetaWorld tasks, while outperforming other methods on hard tasks such asPick Place WallandShelf Place\. DreamerV3 often fails to converge\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x18.png)Figure 18:Single\-task ManiSkill2 results\.Success rate \(%\) as a function of environment steps on 5 object manipulation tasks from ManiSkill2\. Pick YCB is the hardest task and involves manipulation of all 74 objects from the YCB\[[3](https://arxiv.org/html/2605.26282#bib.bib63)\]\. We report44M environment steps for each task\. MBDPO achieves a success rate above75%75\\%on the Pick YCB task, whereas other methods fail to learn within the given budget\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x19.png)Figure 19:High\-dimensional locomotion results\.Episode return as a function of environment steps on all 7 “Locomotion” benchmark tasks\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x20.png)Figure 20:Single\-task MyoSuite results\.Success rate \(%\) as a function of environment steps\. This task domain includes high\-dimensional contact\-rich musculoskeletal motor control with a physiologically accurate robot hand\. Goals are randomized in tasks designated as “Hard”\. MBDPO achieves comparable or better performance than existing methods on all tasks from this benchmark\.
### A\.3Additional Results of Action Drift During Training

![Refer to caption](https://arxiv.org/html/2605.26282v1/x21.png)Figure 21:Per\-task cross TD\-error during training\.EMA\-smoothed cross TD\-errors are reported for representative online tasks\. The cross TD\-error measures the temporal\-difference error incurred when evaluating trajectories under a policy different from the one used for value\-function training, thereby reflecting the distribution mismatch between policy improvement and value learning\. Across most tasks, MBDPO exhibits lower cross TD\-error than TD\-MPC2, suggesting that diffusion policy optimization reduces out\-of\-distribution value queries and leads to more stable value iteration during training\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x22.png)Figure 22:Per\-task action drift during training\.EMA\-smoothed action differences are reported for each of the88online tasks\. The Averaged Policy Network \(Grey\) curve reports the mean action drift of the policy networks from TD\-MPC2 and the two MBDPO variants\. Across these tasks, MBDPO exhibits smaller drift than TD\-MPC2, indicating improved temporal stability under diffusion policy optimization\.
### A\.4Latent Task Embedding

Figure[23](https://arxiv.org/html/2605.26282#A1.F23)below shows the task embeddingℰe​n​v\\mathcal\{E\}\_\{env\}for7070\-task pretraining with1010unseen tasks for fine\-tuning\. Each point corresponds to one task embedding projected into a two\-dimensional space using t\-SNE for visualization\. The red circles denote MetaWorld tasks, while the blue squares denote DMControl tasks\. As shown in the figure, the learned task embeddings exhibit a clear domain\-level separation: MetaWorld manipulation tasks form a compact cluster in the upper region, whereas DMControl locomotion and control tasks are mainly grouped in the lower region\. This separation indicates thatℰe​n​v\\mathcal\{E\}\_\{env\}captures task\-specific environment information and organizes tasks according to their underlying dynamics, observation distributions, and action semantics\.

Most tasks within the same benchmark are embedded close to semantically or dynamically related tasks\. For example, MetaWorld tasks involving similar manipulation primitives, such as pushing, pulling, opening, closing, and object placement, tend to be located near each other\. Similarly, DMControl tasks with related morphology or control structure, such as Walker, Cheetah, Cartpole, and Reacher variants, also form local clusters\. This suggests that the learned task embedding is not merely distinguishing task identities, but also reflects meaningful similarities across tasks\. Notably, the task “Hopper Hop” is highlighted by the red circle\. Although it belongs to the DMControl suite, its embedding is relatively isolated from the main DMControl cluster\. This indicates that “Hopper Hop” has a substantially different data distribution from the other pretraining tasks\. Such an outlying embedding may lead to less effective sharing of representations and policies during multi\-task pretraining, thereby making policy optimization more challenging for this task\. Overall, the visualization supports our hypothesis that the learned task embeddingℰe​n​v\\mathcal\{E\}\_\{env\}provides a structured representation of task distributions, while also revealing tasks whose distributions deviate significantly from the majority of the pretraining set\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x23.png)Figure 23:T\-SNE visualization\[[61](https://arxiv.org/html/2605.26282#bib.bib71)\]of task embeddings learned by MBDPO trained on8080tasks from DMControl and MetaWorld\.Detailed labels are shown for clarity\. According to our hypothesis, the task embedding of Hopper Hop can be attributed to a severe deviation in the representation’s data distribution \(see the red circle\), which subsequently impedes effective policy optimization\.
### A\.5Latent Trajectories Comparison

We provide additional visualizations of latent state trajectories for all8080tasks under both MBDPO and TD\-MPC2\. Specifically, Figures[24](https://arxiv.org/html/2605.26282#A1.F24)and[25](https://arxiv.org/html/2605.26282#A1.F25)compare single\-episode trajectories, while Figures[26](https://arxiv.org/html/2605.26282#A1.F26)and[27](https://arxiv.org/html/2605.26282#A1.F27)compare multi\-episode trajectories\. Across both settings, MBDPO produces more structured and temporally coherent latent manifolds, including closed\-loop patterns for cyclic control tasks and smooth goal\-directed transitions for manipulation tasks\. In contrast, TD\-MPC2 often exhibits noisier and less organized trajectories, suggesting weaker alignment between the sampled policy and the underlying task dynamics\.

![Refer to caption](https://arxiv.org/html/2605.26282v1/x24.png)Figure 24:Single\-episode latent trajectory visualization for MBDPO\.We visualize the latent state trajectories produced by MBDPO with the diffusion policy over one episode for all8080tasks using a locally linear embedding\. The color gradient indicates temporal progression along the trajectory\. For cyclical control tasks, the latent trajectories often form closed\-loop structures that reflect the repetitive physical dynamics of the skills \(DMControl tasks\)\. For goal\-directed manipulation tasks \(MetaWorld tasks\), the trajectories typically evolve smoothly from the initial state toward the target state, suggesting that the diffusion policy induces structured and physically meaningful latent transitions\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x25.png)Figure 25:Single\-episode latent trajectory visualization for TD\-MPC2\.We visualize the latent state trajectories produced by TD\-MPC2 over one episode for all8080tasks using the same locally linear embedding protocol\. The color gradient indicates temporal progression along the trajectory\. Compared with MBDPO, TD\-MPC2 produces trajectories that are often more scattered, noisy, and less aligned with the physical or goal\-directed structure of the tasks\. This suggests that its sampled policy may induce less coherent latent dynamics during decision\-making\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x26.png)Figure 26:Multi\-episode latent trajectory visualization for MBDPO\.We visualize latent state trajectories produced by MBDPO with the diffusion policy over multiple episodes for all8080tasks using a locally linear embedding\. The color gradient denotes temporal progression within each episode\. Across repeated rollouts, MBDPO maintains consistent and structured trajectory manifolds, with closed\-loop patterns for cyclic tasks and smooth directed transitions for goal\-conditioned manipulation tasks\. These results indicate that the diffusion policy captures stable task\-specific latent dynamics across episodes\.![Refer to caption](https://arxiv.org/html/2605.26282v1/x27.png)Figure 27:Multi\-episode latent trajectory visualization for TD\-MPC2\.We visualize latent state trajectories produced by TD\-MPC2 over multiple episodes for all8080tasks using the same locally linear embedding protocol\. The color gradient denotes temporal progression within each episode\. Compared with MBDPO, TD\-MPC2 exhibits less consistent trajectory geometry across rollouts, with noisier, more dispersed, and more irregular latent paths\. This indicates that its sampled policy produces less stable latent dynamics and weaker alignment with the underlying causal structure of the environments\.
### A\.6Comprehensive Training and Evaluation Time

Table 2:Profiling results of online\-task training time per environment step across domains\(milliseconds per env step\)\. Values are reported as mean±\\pmstd\. Each experiment runs a single task on one NVIDIA A800\-SXM4\-80GB GPU\.TaskTD\-MPC2MBDPOTaskTD\-MPC2MBDPOMetaWorldmw\-assembly1\.35±0\.0121\.35\\pm 0\.0121\.35±0\.0081\.35\\pm 0\.008mw\-basketball1\.07±0\.0191\.07\\pm 0\.0191\.09±0\.0221\.09\\pm 0\.022mw\-button\-press\-topdown0\.87±0\.0070\.87\\pm 0\.0070\.89±0\.0060\.89\\pm 0\.006mw\-button\-press\-topdown\-wall0\.90±0\.0160\.90\\pm 0\.0160\.89±0\.0060\.89\\pm 0\.006mw\-button\-press0\.89±0\.0170\.89\\pm 0\.0170\.92±0\.0280\.92\\pm 0\.028mw\-button\-press\-wall0\.94±0\.0540\.94\\pm 0\.0540\.91±0\.0130\.91\\pm 0\.013mw\-coffee\-button0\.99±0\.0150\.99\\pm 0\.0151\.00±0\.0041\.00\\pm 0\.004mw\-coffee\-pull1\.26±0\.0261\.26\\pm 0\.0261\.26±0\.0071\.26\\pm 0\.007mw\-coffee\-push1\.22±0\.0071\.22\\pm 0\.0071\.21±0\.0061\.21\\pm 0\.006mw\-dial\-turn0\.94±0\.0240\.94\\pm 0\.0240\.92±0\.0070\.92\\pm 0\.007mw\-disassemble1\.36±0\.0181\.36\\pm 0\.0181\.34±0\.0201\.34\\pm 0\.020mw\-door\-open1\.09±0\.0131\.09\\pm 0\.0131\.08±0\.0151\.08\\pm 0\.015mw\-door\-close1\.01±0\.0101\.01\\pm 0\.0101\.01±0\.0161\.01\\pm 0\.016mw\-drawer\-close1\.09±0\.0411\.09\\pm 0\.0411\.08±0\.0411\.08\\pm 0\.041mw\-drawer\-open1\.02±0\.0071\.02\\pm 0\.0071\.03±0\.0121\.03\\pm 0\.012mw\-faucet\-open0\.91±0\.0220\.91\\pm 0\.0220\.90±0\.0290\.90\\pm 0\.029mw\-faucet\-close0\.90±0\.0050\.90\\pm 0\.0050\.89±0\.0070\.89\\pm 0\.007mw\-hammer1\.33±0\.0161\.33\\pm 0\.0161\.33±0\.0161\.33\\pm 0\.016mw\-handle\-press\-side0\.88±0\.0350\.88\\pm 0\.0350\.87±0\.0180\.87\\pm 0\.018mw\-handle\-press0\.93±0\.0380\.93\\pm 0\.0380\.92±0\.0370\.92\\pm 0\.037mw\-handle\-pull\-side1\.00±0\.0111\.00\\pm 0\.0111\.00±0\.0361\.00\\pm 0\.036mw\-handle\-pull0\.99±0\.0090\.99\\pm 0\.0090\.98±0\.0070\.98\\pm 0\.007mw\-lever\-pull1\.00±0\.0141\.00\\pm 0\.0140\.99±0\.0100\.99\\pm 0\.010mw\-peg\-insert\-side1\.14±0\.0051\.14\\pm 0\.0051\.14±0\.0161\.14\\pm 0\.016mw\-peg\-unplug\-side1\.22±0\.0791\.22\\pm 0\.0791\.20±0\.0791\.20\\pm 0\.079mw\-pick\-out\-of\-hole1\.08±0\.0191\.08\\pm 0\.0191\.07±0\.0481\.07\\pm 0\.048mw\-pick\-place1\.07±0\.0141\.07\\pm 0\.0141\.10±0\.0321\.10\\pm 0\.032mw\-pick\-place\-wall1\.19±0\.0101\.19\\pm 0\.0101\.18±0\.0061\.18\\pm 0\.006mw\-plate\-slide0\.90±0\.0110\.90\\pm 0\.0110\.92±0\.0130\.92\\pm 0\.013mw\-plate\-slide\-side0\.92±0\.0080\.92\\pm 0\.0080\.91±0\.0130\.91\\pm 0\.013mw\-plate\-slide\-back0\.97±0\.0110\.97\\pm 0\.0110\.94±0\.0130\.94\\pm 0\.013mw\-plate\-slide\-back\-side0\.96±0\.0110\.96\\pm 0\.0110\.95±0\.0090\.95\\pm 0\.009mw\-push\-back1\.10±0\.0081\.10\\pm 0\.0081\.10±0\.0111\.10\\pm 0\.011mw\-push1\.11±0\.0101\.11\\pm 0\.0101\.11±0\.0081\.11\\pm 0\.008mw\-push\-wall1\.17±0\.0101\.17\\pm 0\.0101\.15±0\.0141\.15\\pm 0\.014mw\-reach0\.90±0\.0110\.90\\pm 0\.0110\.89±0\.0110\.89\\pm 0\.011mw\-reach\-wall0\.91±0\.0050\.91\\pm 0\.0050\.92±0\.0250\.92\\pm 0\.025mw\-shelf\-place1\.21±0\.0401\.21\\pm 0\.0401\.19±0\.0381\.19\\pm 0\.038mw\-soccer1\.10±0\.0141\.10\\pm 0\.0141\.10±0\.0091\.10\\pm 0\.009mw\-stick\-push1\.35±0\.0151\.35\\pm 0\.0151\.34±0\.0101\.34\\pm 0\.010mw\-stick\-pull1\.43±0\.0171\.43\\pm 0\.0171\.39±0\.0091\.39\\pm 0\.009mw\-sweep\-into1\.19±0\.0301\.19\\pm 0\.0301\.20±0\.0231\.20\\pm 0\.023mw\-sweep1\.12±0\.0241\.12\\pm 0\.0241\.11±0\.0041\.11\\pm 0\.004mw\-window\-open1\.23±0\.0181\.23\\pm 0\.0181\.22±0\.0201\.22\\pm 0\.020mw\-window\-close1\.24±0\.0311\.24\\pm 0\.0311\.23±0\.0081\.23\\pm 0\.008mw\-bin\-picking1\.12±0\.0151\.12\\pm 0\.0151\.19±0\.2291\.19\\pm 0\.229mw\-box\-close1\.44±0\.0921\.44\\pm 0\.0921\.41±0\.0871\.41\\pm 0\.087mw\-door\-lock1\.02±0\.0671\.02\\pm 0\.0671\.00±0\.0761\.00\\pm 0\.076mw\-door\-unlock1\.00±0\.0421\.00\\pm 0\.0421\.00±0\.0411\.00\\pm 0\.041mw\-hand\-insert1\.14±0\.0311\.14\\pm 0\.0311\.11±0\.0261\.11\\pm 0\.026DMControlacrobot\-swingup0\.40±0\.0030\.40\\pm 0\.0030\.38±0\.0020\.38\\pm 0\.002cartpole\-balance0\.51±0\.0020\.51\\pm 0\.0020\.50±0\.0120\.50\\pm 0\.012cartpole\-balance\-sparse0\.40±0\.0040\.40\\pm 0\.0040\.37±0\.0020\.37\\pm 0\.002cartpole\-swingup0\.51±0\.0130\.51\\pm 0\.0130\.51±0\.0010\.51\\pm 0\.001cartpole\-swingup\-sparse0\.38±0\.0020\.38\\pm 0\.0020\.37±0\.0010\.37\\pm 0\.001cheetah\-jump0\.50±0\.0070\.50\\pm 0\.0070\.50±0\.0090\.50\\pm 0\.009cheetah\-run0\.40±0\.0050\.40\\pm 0\.0050\.38±0\.0080\.38\\pm 0\.008cheetah\-run\-back0\.45±0\.0030\.45\\pm 0\.0030\.45±0\.0080\.45\\pm 0\.008cheetah\-run\-backwards0\.40±0\.0080\.40\\pm 0\.0080\.38±0\.0030\.38\\pm 0\.003cheetah\-run\-front0\.47±0\.0070\.47\\pm 0\.0070\.44±0\.0040\.44\\pm 0\.004cup\-catch0\.40±0\.0110\.40\\pm 0\.0110\.39±0\.0020\.39\\pm 0\.002cup\-spin0\.68±0\.0550\.68\\pm 0\.0550\.67±0\.0500\.67\\pm 0\.050dog\-run2\.76±0\.1472\.76\\pm 0\.1472\.70±0\.1372\.70\\pm 0\.137dog\-trot2\.72±0\.0972\.72\\pm 0\.0972\.84±0\.1182\.84\\pm 0\.118dog\-stand2\.69±0\.0692\.69\\pm 0\.0692\.62±0\.0902\.62\\pm 0\.090dog\-walk2\.85±0\.1512\.85\\pm 0\.1512\.74±0\.1382\.74\\pm 0\.138finger\-spin0\.37±0\.0090\.37\\pm 0\.0090\.33±0\.0020\.33\\pm 0\.002finger\-turn\-easy0\.47±0\.0050\.47\\pm 0\.0050\.45±0\.0030\.45\\pm 0\.003finger\-turn\-hard0\.47±0\.0050\.47\\pm 0\.0050\.46±0\.0050\.46\\pm 0\.005fish\-swim0\.51±0\.0030\.51\\pm 0\.0030\.51±0\.0030\.51\\pm 0\.003hopper\-hop0\.45±0\.0190\.45\\pm 0\.0190\.41±0\.0150\.41\\pm 0\.015hopper\-hop\-backwards0\.53±0\.0430\.53\\pm 0\.0430\.41±0\.0160\.41\\pm 0\.016hopper\-stand0\.49±0\.0140\.49\\pm 0\.0140\.46±0\.0150\.46\\pm 0\.015humanoid\-run1\.03±0\.0131\.03\\pm 0\.0131\.02±0\.0121\.02\\pm 0\.012humanoid\-stand1\.07±0\.0101\.07\\pm 0\.0101\.09±0\.0091\.09\\pm 0\.009humanoid\-walk1\.06±0\.0111\.06\\pm 0\.0111\.03±0\.0121\.03\\pm 0\.012pendulum\-spin0\.35±0\.0110\.35\\pm 0\.0110\.33±0\.0010\.33\\pm 0\.001pendulum\-swingup0\.31±0\.0020\.31\\pm 0\.0020\.30±0\.0010\.30\\pm 0\.001quadruped\-run1\.11±0\.0201\.11\\pm 0\.0201\.09±0\.0191\.09\\pm 0\.019quadruped\-walk1\.09±0\.0191\.09\\pm 0\.0191\.09±0\.0221\.09\\pm 0\.022reacher\-easy0\.37±0\.0030\.37\\pm 0\.0030\.34±0\.0040\.34\\pm 0\.004reacher\-hard0\.36±0\.0030\.36\\pm 0\.0030\.34±0\.0020\.34\\pm 0\.002reacher\-three\-easy0\.38±0\.0070\.38\\pm 0\.0070\.36±0\.0030\.36\\pm 0\.003reacher\-three\-hard0\.38±0\.0060\.38\\pm 0\.0060\.35±0\.0030\.35\\pm 0\.003walker\-run0\.61±0\.0060\.61\\pm 0\.0060\.58±0\.0060\.58\\pm 0\.006walker\-run\-backwards0\.61±0\.0050\.61\\pm 0\.0050\.59±0\.0070\.59\\pm 0\.007walker\-stand0\.57±0\.0100\.57\\pm 0\.0100\.55±0\.0210\.55\\pm 0\.021walker\-walk0\.61±0\.0060\.61\\pm 0\.0060\.58±0\.0030\.58\\pm 0\.003walker\-walk\-backwards0\.60±0\.0070\.60\\pm 0\.0070\.59±0\.0070\.59\\pm 0\.007ManiSkill2lift\-cube4\.98±0\.3014\.98\\pm 0\.3014\.85±0\.2494\.85\\pm 0\.249pick\-cube4\.84±0\.3894\.84\\pm 0\.3895\.04±0\.2515\.04\\pm 0\.251stack\-cube5\.74±0\.2295\.74\\pm 0\.2295\.85±0\.3515\.85\\pm 0\.351pick\-ycb8\.24±1\.9218\.24\\pm 1\.9217\.99±1\.8907\.99\\pm 1\.890turn\-faucet8\.69±2\.7738\.69\\pm 2\.77310\.81±5\.65210\.81\\pm 5\.652MyoSuitemyo\-reach1\.59±0\.0621\.59\\pm 0\.0621\.52±0\.0291\.52\\pm 0\.029myo\-reach\-hard1\.59±0\.0411\.59\\pm 0\.0411\.55±0\.0241\.55\\pm 0\.024myo\-pose1\.49±0\.0391\.49\\pm 0\.0391\.48±0\.0491\.48\\pm 0\.049myo\-pose\-hard1\.65±0\.0461\.65\\pm 0\.0461\.59±0\.0871\.59\\pm 0\.087myo\-obj\-hold1\.68±0\.0581\.68\\pm 0\.0581\.61±0\.0361\.61\\pm 0\.036myo\-obj\-hold\-hard1\.58±0\.0431\.58\\pm 0\.0431\.61±0\.0401\.61\\pm 0\.040myo\-key\-turn1\.84±0\.0811\.84\\pm 0\.0811\.70±0\.0851\.70\\pm 0\.085myo\-key\-turn\-hard1\.90±0\.0971\.90\\pm 0\.0971\.86±0\.0901\.86\\pm 0\.090myo\-pen\-twirl1\.23±0\.0221\.23\\pm 0\.0221\.22±0\.0391\.22\\pm 0\.039myo\-pen\-twirl\-hard1\.24±0\.0371\.24\\pm 0\.0371\.24±0\.0191\.24\\pm 0\.019Visual RLacrobot\-swingup2\.99±0\.0232\.99\\pm 0\.0232\.54±0\.0702\.54\\pm 0\.070cheetah\-run3\.20±0\.1533\.20\\pm 0\.1533\.16±0\.5533\.16\\pm 0\.553finger\-spin4\.18±0\.2124\.18\\pm 0\.2122\.98±0\.4382\.98\\pm 0\.438finger\-turn\-easy4\.10±0\.0864\.10\\pm 0\.0862\.96±0\.0502\.96\\pm 0\.050finger\-turn\-hard3\.38±0\.2793\.38\\pm 0\.2793\.04±0\.3503\.04\\pm 0\.350quadruped\-walk4\.42±0\.1014\.42\\pm 0\.1013\.72±0\.0293\.72\\pm 0\.029reacher\-easy3\.48±0\.2803\.48\\pm 0\.2802\.80±0\.0762\.80\\pm 0\.076reacher\-hard2\.91±0\.0272\.91\\pm 0\.0272\.53±0\.0112\.53\\pm 0\.011walker\-run3\.45±0\.1043\.45\\pm 0\.1043\.43±0\.4393\.43\\pm 0\.439walker\-walk3\.44±0\.0893\.44\\pm 0\.0893\.24±0\.0483\.24\\pm 0\.048

Table 3:Profiling results of evaluation time for online tasks\(seconds per evaluation phase, 10 episodes in total\) across domains\. Values are reported as mean±\\pmstd\. All experiments are evaluated with a single task running on a single NVIDIA A800\-SXM4\-80GB GPU\.TaskTD\-MPC2MBDPOTaskTD\-MPC2MBDPOMetaWorldmw\-assembly10\.67±0\.0310\.67\\pm 0\.0318\.89±0\.0218\.89\\pm 0\.02mw\-basketball10\.31±0\.0110\.31\\pm 0\.0118\.76±0\.0418\.76\\pm 0\.04mw\-button\-press\-topdown10\.18±0\.0110\.18\\pm 0\.0118\.77±0\.0218\.77\\pm 0\.02mw\-button\-press\-topdown\-wall10\.32±0\.0110\.32\\pm 0\.0118\.73±0\.0818\.73\\pm 0\.08mw\-button\-press10\.06±0\.0110\.06\\pm 0\.0118\.39±0\.0218\.39\\pm 0\.02mw\-button\-press\-wall10\.16±0\.0110\.16\\pm 0\.0118\.51±0\.0218\.51\\pm 0\.02mw\-coffee\-button10\.33±0\.0210\.33\\pm 0\.0218\.48±0\.0318\.48\\pm 0\.03mw\-coffee\-pull10\.59±0\.0410\.59\\pm 0\.0418\.73±0\.1618\.73\\pm 0\.16mw\-coffee\-push10\.44±0\.0110\.44\\pm 0\.0118\.79±0\.0218\.79\\pm 0\.02mw\-dial\-turn10\.25±0\.0110\.25\\pm 0\.0118\.31±0\.0318\.31\\pm 0\.03mw\-disassemble10\.48±0\.0110\.48\\pm 0\.0118\.68±0\.0918\.68\\pm 0\.09mw\-door\-open10\.29±0\.0110\.29\\pm 0\.0118\.69±0\.0118\.69\\pm 0\.01mw\-door\-close10\.26±0\.0110\.26\\pm 0\.0119\.02±0\.0419\.02\\pm 0\.04mw\-drawer\-close10\.36±0\.0310\.36\\pm 0\.0318\.60±0\.0218\.60\\pm 0\.02mw\-drawer\-open10\.32±0\.0510\.32\\pm 0\.0518\.67±0\.0818\.67\\pm 0\.08mw\-faucet\-open10\.36±0\.0610\.36\\pm 0\.0618\.38±0\.1218\.38\\pm 0\.12mw\-faucet\-close10\.14±0\.0210\.14\\pm 0\.0218\.15±0\.0418\.15\\pm 0\.04mw\-hammer10\.62±0\.0310\.62\\pm 0\.0318\.87±0\.0418\.87\\pm 0\.04mw\-handle\-press\-side10\.08±0\.0010\.08\\pm 0\.0018\.53±0\.0918\.53\\pm 0\.09mw\-handle\-press10\.15±0\.0210\.15\\pm 0\.0218\.61±0\.0418\.61\\pm 0\.04mw\-handle\-pull\-side10\.21±0\.0010\.21\\pm 0\.0018\.40±0\.0218\.40\\pm 0\.02mw\-handle\-pull10\.25±0\.0310\.25\\pm 0\.0318\.52±0\.0418\.52\\pm 0\.04mw\-lever\-pull10\.31±0\.0210\.31\\pm 0\.0218\.73±0\.1218\.73\\pm 0\.12mw\-peg\-insert\-side10\.38±0\.0510\.38\\pm 0\.0518\.79±0\.0418\.79\\pm 0\.04mw\-peg\-unplug\-side10\.42±0\.0310\.42\\pm 0\.0318\.73±0\.0418\.73\\pm 0\.04mw\-pick\-out\-of\-hole10\.66±0\.0110\.66\\pm 0\.0118\.73±0\.0218\.73\\pm 0\.02mw\-pick\-place10\.33±0\.0110\.33\\pm 0\.0118\.50±0\.0718\.50\\pm 0\.07mw\-pick\-place\-wall10\.45±0\.0210\.45\\pm 0\.0218\.70±0\.0318\.70\\pm 0\.03mw\-plate\-slide10\.27±0\.0110\.27\\pm 0\.0118\.42±0\.0318\.42\\pm 0\.03mw\-plate\-slide\-side10\.15±0\.0110\.15\\pm 0\.0118\.57±0\.3018\.57\\pm 0\.30mw\-plate\-slide\-back10\.36±0\.0410\.36\\pm 0\.0418\.67±0\.0318\.67\\pm 0\.03mw\-plate\-slide\-back\-side10\.42±0\.0210\.42\\pm 0\.0219\.25±0\.0919\.25\\pm 0\.09mw\-push\-back10\.32±0\.0110\.32\\pm 0\.0119\.04±0\.0319\.04\\pm 0\.03mw\-push10\.37±0\.0210\.37\\pm 0\.0218\.64±0\.0418\.64\\pm 0\.04mw\-push\-wall10\.62±0\.0510\.62\\pm 0\.0518\.62±0\.0218\.62\\pm 0\.02mw\-reach10\.16±0\.0110\.16\\pm 0\.0118\.60±0\.0418\.60\\pm 0\.04mw\-reach\-wall10\.14±0\.0310\.14\\pm 0\.0318\.42±0\.0418\.42\\pm 0\.04mw\-shelf\-place10\.42±0\.0210\.42\\pm 0\.0218\.64±0\.0218\.64\\pm 0\.02mw\-soccer10\.45±0\.0510\.45\\pm 0\.0519\.08±0\.0519\.08\\pm 0\.05mw\-stick\-push10\.57±0\.0510\.57\\pm 0\.0519\.04±0\.0319\.04\\pm 0\.03mw\-stick\-pull10\.65±0\.0210\.65\\pm 0\.0219\.19±0\.0319\.19\\pm 0\.03mw\-sweep\-into10\.64±0\.0510\.64\\pm 0\.0518\.53±0\.0318\.53\\pm 0\.03mw\-sweep10\.32±0\.0110\.32\\pm 0\.0118\.71±0\.0218\.71\\pm 0\.02mw\-window\-open10\.48±0\.0210\.48\\pm 0\.0218\.48±0\.0218\.48\\pm 0\.02mw\-window\-close10\.41±0\.0110\.41\\pm 0\.0118\.59±0\.0318\.59\\pm 0\.03mw\-bin\-picking10\.36±0\.0110\.36\\pm 0\.0119\.01±0\.0319\.01\\pm 0\.03mw\-box\-close10\.51±0\.0210\.51\\pm 0\.0218\.44±0\.0818\.44\\pm 0\.08mw\-door\-lock10\.47±0\.0510\.47\\pm 0\.0518\.65±0\.0618\.65\\pm 0\.06mw\-door\-unlock10\.45±0\.0210\.45\\pm 0\.0218\.86±0\.0418\.86\\pm 0\.04mw\-hand\-insert10\.73±0\.0410\.73\\pm 0\.0418\.89±0\.0418\.89\\pm 0\.04DMControlacrobot\-swingup48\.96±0\.0448\.96\\pm 0\.0491\.41±0\.0591\.41\\pm 0\.05cartpole\-balance49\.54±0\.0249\.54\\pm 0\.0293\.11±0\.4993\.11\\pm 0\.49cartpole\-balance\-sparse48\.59±0\.0248\.59\\pm 0\.0289\.92±0\.3189\.92\\pm 0\.31cartpole\-swingup49\.14±0\.0249\.14\\pm 0\.0295\.71±0\.2195\.71\\pm 0\.21cartpole\-swingup\-sparse48\.53±0\.0248\.53\\pm 0\.0290\.39±0\.6490\.39\\pm 0\.64cheetah\-jump48\.64±0\.0248\.64\\pm 0\.0290\.12±0\.4290\.12\\pm 0\.42cheetah\-run48\.04±0\.0548\.04\\pm 0\.0588\.86±0\.0688\.86\\pm 0\.06cheetah\-run\-back48\.40±0\.0748\.40\\pm 0\.0789\.86±0\.0589\.86\\pm 0\.05cheetah\-run\-backwards48\.28±0\.1948\.28\\pm 0\.1989\.26±0\.2089\.26\\pm 0\.20cheetah\-run\-front48\.57±0\.1848\.57\\pm 0\.1893\.52±0\.3193\.52\\pm 0\.31cup\-catch48\.53±0\.0648\.53\\pm 0\.0689\.35±0\.3389\.35\\pm 0\.33cup\-spin49\.28±0\.0849\.28\\pm 0\.0888\.98±0\.0888\.98\\pm 0\.08dog\-run81\.49±0\.4981\.49\\pm 0\.49116\.13±0\.88116\.13\\pm 0\.88dog\-trot82\.57±0\.4382\.57\\pm 0\.43112\.60±0\.74112\.60\\pm 0\.74dog\-stand80\.86±0\.4780\.86\\pm 0\.47117\.07±0\.78117\.07\\pm 0\.78dog\-walk76\.16±0\.1476\.16\\pm 0\.14102\.96±0\.42102\.96\\pm 0\.42finger\-spin48\.39±0\.0848\.39\\pm 0\.0888\.62±0\.5188\.62\\pm 0\.51finger\-turn\-easy48\.70±0\.0748\.70\\pm 0\.0790\.84±0\.3890\.84\\pm 0\.38finger\-turn\-hard49\.44±0\.4749\.44\\pm 0\.4790\.07±0\.1890\.07\\pm 0\.18fish\-swim50\.22±0\.4050\.22\\pm 0\.4090\.99±0\.7590\.99\\pm 0\.75hopper\-hop48\.45±0\.0548\.45\\pm 0\.0590\.15±0\.0890\.15\\pm 0\.08hopper\-hop\-backwards48\.56±0\.0348\.56\\pm 0\.0390\.67±0\.0690\.67\\pm 0\.06hopper\-stand48\.93±0\.0448\.93\\pm 0\.0489\.40±0\.0789\.40\\pm 0\.07humanoid\-run65\.93±0\.1665\.93\\pm 0\.1693\.01±0\.3193\.01\\pm 0\.31humanoid\-stand66\.68±0\.4266\.68\\pm 0\.4293\.37±0\.9893\.37\\pm 0\.98humanoid\-walk65\.95±0\.1965\.95\\pm 0\.1993\.96±0\.1593\.96\\pm 0\.15pendulum\-spin48\.23±0\.0448\.23\\pm 0\.0489\.49±0\.4089\.49\\pm 0\.40pendulum\-swingup48\.29±0\.0948\.29\\pm 0\.0989\.57±0\.9389\.57\\pm 0\.93quadruped\-run51\.07±0\.1051\.07\\pm 0\.1092\.25±0\.3492\.25\\pm 0\.34quadruped\-walk50\.48±0\.0650\.48\\pm 0\.0694\.07±0\.1494\.07\\pm 0\.14reacher\-easy48\.58±0\.0548\.58\\pm 0\.0589\.32±0\.3989\.32\\pm 0\.39reacher\-hard48\.50±0\.1948\.50\\pm 0\.1989\.07±0\.3889\.07\\pm 0\.38reacher\-three\-easy48\.62±0\.0448\.62\\pm 0\.0489\.35±0\.0489\.35\\pm 0\.04reacher\-three\-hard48\.29±0\.0848\.29\\pm 0\.0888\.68±0\.4288\.68\\pm 0\.42walker\-run49\.95±0\.0349\.95\\pm 0\.0390\.58±0\.1790\.58\\pm 0\.17walker\-run\-backwards49\.51±0\.0349\.51\\pm 0\.0389\.00±0\.1289\.00\\pm 0\.12walker\-stand49\.05±0\.0349\.05\\pm 0\.0392\.53±0\.1392\.53\\pm 0\.13walker\-walk49\.31±0\.0849\.31\\pm 0\.0890\.07±0\.0890\.07\\pm 0\.08walker\-walk\-backwards49\.61±0\.0449\.61\\pm 0\.0490\.30±0\.0690\.30\\pm 0\.06ManiSkill2lift\-cube15\.27±0\.3215\.27\\pm 0\.3224\.99±1\.1924\.99\\pm 1\.19pick\-cube15\.62±0\.2715\.62\\pm 0\.2726\.34±0\.7626\.34\\pm 0\.76stack\-cube16\.08±0\.3016\.08\\pm 0\.3027\.02±0\.6527\.02\\pm 0\.65pick\-ycb18\.83±1\.0118\.83\\pm 1\.0131\.01±2\.2031\.01\\pm 2\.20turn\-faucet18\.84±1\.1118\.84\\pm 1\.1125\.06±0\.5525\.06\\pm 0\.55MyoSuitemyo\-reach13\.90±0\.0313\.90\\pm 0\.0319\.22±0\.0219\.22\\pm 0\.02myo\-reach\-hard14\.00±0\.0414\.00\\pm 0\.0419\.45±0\.0919\.45\\pm 0\.09myo\-pose13\.90±0\.0713\.90\\pm 0\.0719\.00±0\.0219\.00\\pm 0\.02myo\-pose\-hard13\.95±0\.0413\.95\\pm 0\.0419\.08±0\.0319\.08\\pm 0\.03myo\-obj\-hold13\.88±0\.0513\.88\\pm 0\.0519\.79±0\.0219\.79\\pm 0\.02myo\-obj\-hold\-hard14\.01±0\.0514\.01\\pm 0\.0519\.34±0\.0519\.34\\pm 0\.05myo\-key\-turn14\.27±0\.0814\.27\\pm 0\.0818\.84±0\.0418\.84\\pm 0\.04myo\-key\-turn\-hard14\.10±0\.0314\.10\\pm 0\.0319\.22±0\.0519\.22\\pm 0\.05myo\-pen\-twirl13\.66±0\.0513\.66\\pm 0\.0518\.81±0\.0318\.81\\pm 0\.03myo\-pen\-twirl\-hard13\.67±0\.0713\.67\\pm 0\.0718\.72±0\.0218\.72\\pm 0\.02Visual RLacrobot\-swingup66\.51±0\.4266\.51\\pm 0\.42108\.95±0\.30108\.95\\pm 0\.30cheetah\-run66\.29±0\.3766\.29\\pm 0\.37108\.64±0\.31108\.64\\pm 0\.31finger\-spin66\.35±0\.4766\.35\\pm 0\.47108\.86±0\.28108\.86\\pm 0\.28finger\-turn\-easy66\.36±0\.4666\.36\\pm 0\.46113\.39±0\.09113\.39\\pm 0\.09finger\-turn\-hard65\.37±2\.5965\.37\\pm 2\.59105\.12±0\.20105\.12\\pm 0\.20quadruped\-walk66\.66±1\.6966\.66\\pm 1\.69106\.39±0\.67106\.39\\pm 0\.67reacher\-easy65\.29±2\.5965\.29\\pm 2\.59104\.03±0\.94104\.03\\pm 0\.94reacher\-hard65\.41±2\.6365\.41\\pm 2\.63104\.97±0\.48104\.97\\pm 0\.48walker\-run66\.02±0\.1566\.02\\pm 0\.15104\.89±0\.63104\.89\\pm 0\.63walker\-walk66\.13±0\.2366\.13\\pm 0\.23105\.64±0\.20105\.64\\pm 0\.20

## Appendix BTechnical Definition

In this section, we provide further technical details and formal definitions to supplement the descriptions in the main text\.

### B\.1Diffusion Policy based on DDPM

We briefly review the fundamental mechanism of Denoising Diffusion Probabilistic Models \(DDPM\)\[[32](https://arxiv.org/html/2605.26282#bib.bib49),[54](https://arxiv.org/html/2605.26282#bib.bib64)\]\. In our setting, the diffusion process based on DDPM is used for sampling and generating target policies\. Letτ∈\{1,…,N\}\\tau\\in\\\{1,\\dots,N\\\}denote the diffusion timestep andα¯τ∈\(0,1\]\\bar\{\\alpha\}^\{\\tau\}\\in\(0,1\]be the cumulative noise schedule\. The forward and reverse diffusion processes for the action sequence are defined as:

\(Forward\)p​\(at:t\+Hτ\|at:t\+H0,zt\)=𝒩​\(at:t\+Hτ;α¯τ​at:t\+H0,\(1−α¯τ\)​I\),\\displaystyle p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)=\\mathcal\{N\}\\left\(a\_\{t:t\+H\}^\{\\tau\};\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}a\_\{t:t\+H\}^\{0\},\(1\-\\bar\{\\alpha\}^\{\\tau\}\)I\\right\),\(17\)\(Reverse\)at:t\+Hτ−1=1ατ​\(at:t\+Hτ\+\(1−ατ\)​ϕ​\(zt,at:t\+Hτ,τ\)\)\+σ​\(τ\)​ϵ,ϵ∼𝒩​\(0,I\),\\displaystyle a\_\{t:t\+H\}^\{\\tau\-1\}=\\frac\{1\}\{\\sqrt\{\\alpha^\{\\tau\}\}\}\\left\(a\_\{t:t\+H\}^\{\\tau\}\+\(1\-\\alpha^\{\\tau\}\)\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)\\right\)\+\\sigma\(\\tau\)\\epsilon,\\quad\\epsilon\\sim\\mathcal\{N\}\(0,I\),\(18\)wherep​\(at:t\+Hτ\|at:t\+H0,zt\)p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)is the transition kernel in the forward process, andϕ​\(zt,at:t\+Hτ,τ\)=∇at:t\+Hτlog⁡πτ​\(at:t\+Hτ\|zt\)\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)=\\nabla\_\{a\_\{t:t\+H\}^\{\\tau\}\}\\log\\pi^\{\\tau\}\(a\_\{t:t\+H\}^\{\\tau\}\|z\_\{t\}\)denotes the score function at diffusion stepτ\\tau\. In our framework, as policy optimization is performed within the latent world model, the explicit forward diffusion process is not strictly required for training\. Nevertheless, the forward process remains essential for computing the transition kernel in the reverse process, as formulated in \([3](https://arxiv.org/html/2605.26282#S2.E3)\)\.

### B\.2Bellman Operator

To formally analyze value estimation within the world model, we define the Bellman operator as follows\[[56](https://arxiv.org/html/2605.26282#bib.bib2)\]:

###### Definition B\.1\(Bellman Operator\)\.

Consider the space ofQQ\-functions𝒬\\mathcal\{Q\}equipped with theL∞L\_\{\\infty\}norm333Formally,𝒬=\{Q\|Q:𝒵×𝒜→ℝ\}\\mathcal\{Q\}=\\\{Q\|Q:\\mathcal\{Z\}\\times\\mathcal\{A\}\\to\\mathbb\{R\}\\\}is the set of all bounded real\-valued functions defined on the latent state\-action space\. TheL∞L\_\{\\infty\}norm is defined as‖Q‖∞=sup\(z,a\)∈𝒵×𝒜\|Q​\(z,a\)\|\\\|Q\\\|\_\{\\infty\}=\\sup\_\{\(z,a\)\\in\\mathcal\{Z\}\\times\\mathcal\{A\}\}\|Q\(z,a\)\|\.\. For a given policyπ\\pi, the Bellman evaluation operatorTπ:𝒬→𝒬T^\{\\pi\}:\\mathcal\{Q\}\\to\\mathcal\{Q\}is defined as:

\(Tπ​Q\)​\(zt,at\)=ℛ​\(zt,at\)\+γ​𝔼zt\+1∼ℱ​\(zt,at\),at\+1∼π​\[Q​\(zt\+1,at\+1\)\],\(T^\{\\pi\}Q\)\(z\_\{t\},a\_\{t\}\)=\\mathcal\{R\}\(z\_\{t\},a\_\{t\}\)\+\\gamma\\mathbb\{E\}\_\{z\_\{t\+1\}\\sim\\mathcal\{F\}\(z\_\{t\},a\_\{t\}\),a\_\{t\+1\}\\sim\\pi\}\[Q\(z\_\{t\+1\},a\_\{t\+1\}\)\],\(19\)whereℛ​\(zt,at\)\\mathcal\{R\}\(z\_\{t\},a\_\{t\}\)denotes the latent reward function,ℱ\\mathcal\{F\}represents the latent dynamics within the learned world model, andγ∈\(0,1\)\\gamma\\in\(0,1\)is the discount factor\.

This operator serves exclusively for policy evaluation and value function update\. By repeatedly applyingTπT^\{\\pi\}, theQQ\-function converges toQπQ^\{\\pi\}\(according to the Banach fixed\-point theorem\[[1](https://arxiv.org/html/2605.26282#bib.bib65)\]\), the true value of the current policyπ\\pi\. This evaluation stage is crucial as it provides an accurate value estimate, which MBDPO then uses to guide the policy optimization process toward higher\-reward regions\.

## Appendix CTheoretical Analysis

In this section, we provide a theoretical analysis of the proposed framework\. We first examine the misalignment between search and value learning, highlighting how discrepancies between the search policy and the learned value function can affect optimization\. We then analyze MBDPO from a theoretical perspective, characterizing its learning dynamics and discussing why it can mitigate such misalignment\. Finally, we present an error analysis of the overall algorithm, identifying the key sources of approximation error and their impact on the resulting policy\.

### C\.1Bottlenecks in the Existing Paradigm

Before presenting our main results, we first provide a detailed mathematical analysis of policy search in the TD\-MPC series\[[31](https://arxiv.org/html/2605.26282#bib.bib26),[30](https://arxiv.org/html/2605.26282#bib.bib27),[70](https://arxiv.org/html/2605.26282#bib.bib28),[40](https://arxiv.org/html/2605.26282#bib.bib50)\]\.

#### C\.1\.1How Does Policy Search Work in the TD\-MPC Series?

The policy search procedure can be decomposed into two steps:

- •Step 1: Initialize with a Prior Policy\.The search distribution is initialized with an isotropic Gaussian prior policyp0\(⋅\|z\)∼𝒩\(μ\(z\),σ2\)p\_\{0\}\(\\cdot\|z\)\\sim\\mathcal\{N\}\(\\mu\(z\),\\sigma^\{2\}\)\(usually parameterized by a policy network\), where the meanμ\\muis estimated by maximizing the objective in \([6](https://arxiv.org/html/2605.26282#S3.E6)\)\.
- •Step 2: Policy Search via Model Predictive Path Integral Planning \(MPPI\)\.Starting from the initialized prior policy𝒩​\(μ​\(z\),σ2\)\\mathcal\{N\}\(\\mu\(z\),\\sigma^\{2\}\), the algorithm iteratively samples candidate action sequences and evaluates them using model\-based rollouts\[[64](https://arxiv.org/html/2605.26282#bib.bib66)\]\. These samples are then reweighted according to their estimated returns \(abstracted in \([20](https://arxiv.org/html/2605.26282#A3.E20)\)\), p0​\(at:t\+H\|zt\)⏟Gaussian prior→sample\{at:t\+H\(i\)\}i=1T⏟candidate actions→evaluated byH\-step​Q^\{Q^\(i\)\}i=1T⏟estimated values→reweightπ​\(at:t\+H\|zt\)⏟search target,\\displaystyle\\underbrace\{p\_\{0\}\(a\_\{t:t\+H\}\|z\_\{t\}\)\}\_\{\\text\{Gaussian prior\}\}\\xrightarrow\{\\text\{sample\}\}\\underbrace\{\\\{a\_\{t:t\+H\}^\{\(i\)\}\\\}\_\{i=1\}^\{T\}\}\_\{\\text\{candidate actions\}\}\\xrightarrow\{\\text\{evaluated by $H$\-step \}\\hat\{Q\}\}\\underbrace\{\\\{\\hat\{Q\}^\{\(i\)\}\\\}\_\{i=1\}^\{T\}\}\_\{\\text\{estimated values\}\}\\xrightarrow\{\\text\{reweight\}\}\\underbrace\{\\pi\(a\_\{t:t\+H\}\|z\_\{t\}\)\}\_\{\\text\{search target\}\},\(20\)whereTTrepresents the sample number, which yields an improved policy target through importance\-weighted aggregation\.

This procedure highlights a key limitation of policy search in the TD\-MPC series\. The resulting target policy is induced by the search distribution rather than obtained through an explicit policy optimization problem constrained around an anchored policy distribution\. As a result, the searched target can become effectively uncontrolled: its update is primarily determined by the return estimates used for importance weighting\. When the learned value function or model\-based return estimates are inaccurate, erroneously highQQ\-values may assign excessive weights to suboptimal action sequences\. Consequently, the aggregated policy target can be biased toward poor actions, causing policy learning to imitate a bad search\-induced target rather than ensuring a monotonic improvement over the prior policy\. In this sense, the search step does not constitute a principled policy optimization procedure; it constructs a heuristic target from weighted samples of learned world models, without explicitly optimizing a regularized policy objective or enforcing an improvement guarantee\.

#### C\.1\.2Analysis of Bottleneck[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)

We start from the theoretical analysis from the Bellman update\. In the setting of a latent world model, theQ^\\hat\{Q\}value function is learned from the \(behavioral\) policyβ\\beta\(parameterized with a policy network\)\.

###### Proof\.

By Definition[B\.1](https://arxiv.org/html/2605.26282#A2.Thmtheorem1), for any latent state\-action pair\(zt,at\)\(z\_\{t\},a\_\{t\}\), we have

\(Tπ​Q^\)​\(zt,at\)=ℛ​\(zt,at\)\+γ​𝔼zt\+1∼ℱ​\(zt,at\)​𝔼at\+1∼π\(⋅\|zt\+1\)​\[Q^​\(zt\+1,at\+1\)\],\(T^\{\\pi\}\\hat\{Q\}\)\(z\_\{t\},a\_\{t\}\)=\\mathcal\{R\}\(z\_\{t\},a\_\{t\}\)\+\\gamma\\mathbb\{E\}\_\{z\_\{t\+1\}\\sim\\mathcal\{F\}\(z\_\{t\},a\_\{t\}\)\}\\mathbb\{E\}\_\{a\_\{t\+1\}\\sim\\pi\(\\cdot\|z\_\{t\+1\}\)\}\[\\hat\{Q\}\(z\_\{t\+1\},a\_\{t\+1\}\)\],\(21\)and

\(Tβ​Q^\)​\(zt,at\)=ℛ​\(zt,at\)\+γ​𝔼zt\+1∼ℱ​\(zt,at\)​𝔼at\+1∼β\(⋅\|zt\+1\)​\[Q^​\(zt\+1,at\+1\)\]\.\(T^\{\\beta\}\\hat\{Q\}\)\(z\_\{t\},a\_\{t\}\)=\\mathcal\{R\}\(z\_\{t\},a\_\{t\}\)\+\\gamma\\mathbb\{E\}\_\{z\_\{t\+1\}\\sim\\mathcal\{F\}\(z\_\{t\},a\_\{t\}\)\}\\mathbb\{E\}\_\{a\_\{t\+1\}\\sim\\beta\(\\cdot\|z\_\{t\+1\}\)\}\[\\hat\{Q\}\(z\_\{t\+1\},a\_\{t\+1\}\)\]\.\(22\)Therefore,

\|\(Tπ​Q^\)​\(zt,at\)−\(Tβ​Q^\)​\(zt,at\)\|\\displaystyle\\left\|\(T^\{\\pi\}\\hat\{Q\}\)\(z\_\{t\},a\_\{t\}\)\-\(T^\{\\beta\}\\hat\{Q\}\)\(z\_\{t\},a\_\{t\}\)\\right\|=γ\|𝔼zt\+1∼ℱ​\(zt,at\)\[∑at\+1\(π\(at\+1\|zt\+1\)−β\(at\+1\|zt\+1\)\)Q^\(zt\+1,at\+1\)\]\|\\displaystyle=\\gamma\\left\|\\mathbb\{E\}\_\{z\_\{t\+1\}\\sim\\mathcal\{F\}\(z\_\{t\},a\_\{t\}\)\}\\left\[\\sum\_\{a\_\{t\+1\}\}\\left\(\\pi\(a\_\{t\+1\}\|z\_\{t\+1\}\)\-\\beta\(a\_\{t\+1\}\|z\_\{t\+1\}\)\\right\)\\hat\{Q\}\(z\_\{t\+1\},a\_\{t\+1\}\)\\right\]\\right\|≤γ𝔼zt\+1∼ℱ​\(zt,at\)\[\|∑at\+1\(π\(at\+1\|zt\+1\)−β\(at\+1\|zt\+1\)\)Q^\(zt\+1,at\+1\)\|\]\\displaystyle\\leq\\gamma\\mathbb\{E\}\_\{z\_\{t\+1\}\\sim\\mathcal\{F\}\(z\_\{t\},a\_\{t\}\)\}\\left\[\\left\|\\sum\_\{a\_\{t\+1\}\}\\left\(\\pi\(a\_\{t\+1\}\|z\_\{t\+1\}\)\-\\beta\(a\_\{t\+1\}\|z\_\{t\+1\}\)\\right\)\\hat\{Q\}\(z\_\{t\+1\},a\_\{t\+1\}\)\\right\|\\right\]≤γ∥Q^∥∞𝔼zt\+1∼ℱ​\(zt,at\)\[∫at\+1\|π\(at\+1\|zt\+1\)−β\(at\+1\|zt\+1\)\|dat\+1\]\\displaystyle\\leq\\gamma\\\|\\hat\{Q\}\\\|\_\{\\infty\}\\mathbb\{E\}\_\{z\_\{t\+1\}\\sim\\mathcal\{F\}\(z\_\{t\},a\_\{t\}\)\}\\left\[\\int\_\{a\_\{t\+1\}\}\\left\|\\pi\(a\_\{t\+1\}\|z\_\{t\+1\}\)\-\\beta\(a\_\{t\+1\}\|z\_\{t\+1\}\)\\right\|da\_\{t\+1\}\\right\]=γ∥Q^∥∞𝔼zt\+1∼ℱ​\(zt,at\)\[∥π\(⋅\|zt\+1\)−β\(⋅\|zt\+1\)∥1\]\.\\displaystyle=\\gamma\\\|\\hat\{Q\}\\\|\_\{\\infty\}\\mathbb\{E\}\_\{z\_\{t\+1\}\\sim\\mathcal\{F\}\(z\_\{t\},a\_\{t\}\)\}\\left\[\\\|\\pi\(\\cdot\|z\_\{t\+1\}\)\-\\beta\(\\cdot\|z\_\{t\+1\}\)\\\|\_\{1\}\\right\]\.\(23\)By Pinsker’s inequality,

∥π\(⋅\|zt\+1\)−β\(⋅\|zt\+1\)∥1≤2DKL\(π\(⋅\|zt\+1\)∥β\(⋅\|zt\+1\)\)\.\\\|\\pi\(\\cdot\|z\_\{t\+1\}\)\-\\beta\(\\cdot\|z\_\{t\+1\}\)\\\|\_\{1\}\\leq\\sqrt\{2D\_\{\\mathrm\{KL\}\}\\left\(\\pi\(\\cdot\|z\_\{t\+1\}\)\\\|\\beta\(\\cdot\|z\_\{t\+1\}\)\\right\)\}\.\(24\)Thus,

\|\(Tπ​Q^\)​\(zt,at\)−\(Tβ​Q^\)​\(zt,at\)\|≤γ​‖Q^‖∞​𝔼zt\+1∼ℱ​\(zt,at\)​\[2DKL\(π\(⋅\|zt\+1\)∥β\(⋅\|zt\+1\)\)\]\.\\left\|\(T^\{\\pi\}\\hat\{Q\}\)\(z\_\{t\},a\_\{t\}\)\-\(T^\{\\beta\}\\hat\{Q\}\)\(z\_\{t\},a\_\{t\}\)\\right\|\\leq\\gamma\\\|\\hat\{Q\}\\\|\_\{\\infty\}\\mathbb\{E\}\_\{z\_\{t\+1\}\\sim\\mathcal\{F\}\(z\_\{t\},a\_\{t\}\)\}\\left\[\\sqrt\{2D\_\{\\mathrm\{KL\}\}\\left\(\\pi\(\\cdot\|z\_\{t\+1\}\)\\\|\\beta\(\\cdot\|z\_\{t\+1\}\)\\right\)\}\\right\]\.\(25\)Taking the supremum over all latent state\-action pairs\(zt,at\)∈𝒵×𝒜\(z\_\{t\},a\_\{t\}\)\\in\\mathcal\{Z\}\\times\\mathcal\{A\}gives

‖Tπ​Q^−Tβ​Q^‖∞≤γ​‖Q^‖∞​supzt,at𝔼zt\+1∼ℱ​\(zt,at\)​\[2DKL\(π\(⋅\|zt\+1\)∥β\(⋅\|zt\+1\)\)\]≤γ​‖Q^‖∞​2​DKLmax​\(π∥β\),\\begin\{split\}\\\|T^\{\\pi\}\\hat\{Q\}\-T^\{\\beta\}\\hat\{Q\}\\\|\_\{\\infty\}&\\leq\\gamma\\\|\\hat\{Q\}\\\|\_\{\\infty\}\\sup\_\{z\_\{t\},a\_\{t\}\}\\mathbb\{E\}\_\{z\_\{t\+1\}\\sim\\mathcal\{F\}\(z\_\{t\},a\_\{t\}\)\}\\left\[\\sqrt\{2D\_\{\\mathrm\{KL\}\}\\left\(\\pi\(\\cdot\|z\_\{t\+1\}\)\\\|\\beta\(\\cdot\|z\_\{t\+1\}\)\\right\)\}\\right\]\\\\ &\\leq\\gamma\\\|\\hat\{Q\}\\\|\_\{\\infty\}\\sqrt\{2D\_\{\\text\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\)\},\\end\{split\}\(26\)withDKLmax\(π∥β\)≔supz∈𝒵DKL\(π\(⋅\|z\)∥β\(⋅\|z\)\)D\_\{\\text\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\)\\coloneqq\\sup\_\{z\\in\\mathcal\{Z\}\}D\_\{\\text\{KL\}\}\(\\pi\(\\cdot\|z\)\\\|\\beta\(\\cdot\|z\)\), which completes the proof\. ∎

### C\.2Analysis of Bottleneck[3\.1](https://arxiv.org/html/2605.26282#S3.SS1)

###### Proof\.

We consider the latent MDP induced by the world model transitionzt\+1=ℱ​\(zt,at\)z\_\{t\+1\}=\\mathcal\{F\}\(z\_\{t\},a\_\{t\}\)\. Equivalently, the transition kernel is deterministic:

Pℱ​\(z′\|z,a\)=𝟙\{z′=ℱ​\(z,a\)\},P\_\{\\mathcal\{F\}\}\(z^\{\\prime\}\|z,a\)=\\mathbb\{1\}\_\{\\\{z^\{\\prime\}=\\mathcal\{F\}\(z,a\)\\\}\},\(27\)where𝟙\{⋅\}\\mathbb\{1\}\_\{\\\{\\cdot\\\}\}is the indicator function\. Letdπd^\{\\pi\}anddβd^\{\\beta\}be the discounted latent\-state occupancy distributions under policiesπ\\piandβ\\beta, respectively\.

By the performance difference lemma\[[43](https://arxiv.org/html/2605.26282#bib.bib68)\], the true improvement ofπ\\pioverβ\\betacan be written as

J​\(π\)−J​\(β\)=11−γ​𝔼z∼dπ,a∼π\(⋅\|z\)​\[Aβ​\(z,a\)\],J\(\\pi\)\-J\(\\beta\)=\\frac\{1\}\{1\-\\gamma\}\\mathbb\{E\}\_\{z\\sim d^\{\\pi\},a\\sim\\pi\(\\cdot\|z\)\}\\left\[A^\{\\beta\}\(z,a\)\\right\],\(28\)where

Aβ​\(z,a\)=Qβ​\(z,a\)−Vβ​\(z\),Vβ​\(z\)=𝔼a∼β\(⋅\|z\)​\[Qβ​\(z,a\)\]\.A^\{\\beta\}\(z,a\)=Q^\{\\beta\}\(z,a\)\-V^\{\\beta\}\(z\),\\qquad V^\{\\beta\}\(z\)=\\mathbb\{E\}\_\{a\\sim\\beta\(\\cdot\|z\)\}\[Q^\{\\beta\}\(z,a\)\]\.\(29\)Equivalently,

J​\(π\)−J​\(β\)\\displaystyle J\(\\pi\)\-J\(\\beta\)=11−γ​𝔼z∼dπ​\[𝔼a∼π\(⋅\|z\)​Qβ​\(z,a\)−𝔼a∼β\(⋅\|z\)​Qβ​\(z,a\)\]\.\\displaystyle=\\frac\{1\}\{1\-\\gamma\}\\mathbb\{E\}\_\{z\\sim d^\{\\pi\}\}\\left\[\\mathbb\{E\}\_\{a\\sim\\pi\(\\cdot\|z\)\}Q^\{\\beta\}\(z,a\)\-\\mathbb\{E\}\_\{a\\sim\\beta\(\\cdot\|z\)\}Q^\{\\beta\}\(z,a\)\\right\]\.\(30\)
We now decompose the right\-hand side into a surrogate term evaluated underdβd^\{\\beta\}and a distribution\-shift term:

J​\(π\)−J​\(β\)\\displaystyle J\(\\pi\)\-J\(\\beta\)=11−γ​𝔼z∼dβ​\[𝔼a∼π\(⋅\|z\)​Qβ​\(z,a\)−𝔼a∼β\(⋅\|z\)​Qβ​\(z,a\)\]\+Δocc,\\displaystyle=\\frac\{1\}\{1\-\\gamma\}\\mathbb\{E\}\_\{z\\sim d^\{\\beta\}\}\\left\[\\mathbb\{E\}\_\{a\\sim\\pi\(\\cdot\|z\)\}Q^\{\\beta\}\(z,a\)\-\\mathbb\{E\}\_\{a\\sim\\beta\(\\cdot\|z\)\}Q^\{\\beta\}\(z,a\)\\right\]\+\\Delta\_\{\\mathrm\{occ\}\},\(31\)where

Δocc\\displaystyle\\Delta\_\{\\mathrm\{occ\}\}=11−γ​\(𝔼z∼dπ−𝔼z∼dβ\)​\[𝔼a∼π\(⋅\|z\)​Qβ​\(z,a\)−𝔼a∼β\(⋅\|z\)​Qβ​\(z,a\)\]\.\\displaystyle=\\frac\{1\}\{1\-\\gamma\}\\left\(\\mathbb\{E\}\_\{z\\sim d^\{\\pi\}\}\-\\mathbb\{E\}\_\{z\\sim d^\{\\beta\}\}\\right\)\\left\[\\mathbb\{E\}\_\{a\\sim\\pi\(\\cdot\|z\)\}Q^\{\\beta\}\(z,a\)\-\\mathbb\{E\}\_\{a\\sim\\beta\(\\cdot\|z\)\}Q^\{\\beta\}\(z,a\)\\right\]\.\(32\)Under the standard boundedness assumption‖Qβ‖∞≤Qmax\\\|Q^\{\\beta\}\\\|\_\{\\infty\}\\leq Q\_\{\\max\}, the distribution\-shift term can be bounded by a trust\-region style argument\[[51](https://arxiv.org/html/2605.26282#bib.bib69)\]:

\|Δocc\|≤C1​DKLmax​\(π∥β\),\|\\Delta\_\{\\mathrm\{occ\}\}\|\\leq C\_\{1\}D\_\{\\mathrm\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\),\(33\)where

DKLmax\(π∥β\)≔supzDKL\(π\(⋅\|z\)∥β\(⋅\|z\)\)\.D\_\{\\mathrm\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\)\\coloneqq\\sup\_\{z\}D\_\{\\mathrm\{KL\}\}\\left\(\\pi\(\\cdot\|z\)\\\|\\beta\(\\cdot\|z\)\\right\)\.\(34\)This bound follows from the fact that the discrepancy between the occupancy distributionsdπd^\{\\pi\}anddβd^\{\\beta\}is controlled by the maximum policy divergence at each latent state, together with Pinsker’s inequality\.

Therefore,

J​\(π\)−J​\(β\)\\displaystyle J\(\\pi\)\-J\(\\beta\)≥11−γ​𝔼z∼dβ​\[𝔼a∼π\(⋅\|z\)​Qβ​\(z,a\)−𝔼a∼β\(⋅\|z\)​Qβ​\(z,a\)\]−C1​DKLmax​\(π∥β\)\.\\displaystyle\\geq\\frac\{1\}\{1\-\\gamma\}\\mathbb\{E\}\_\{z\\sim d^\{\\beta\}\}\\left\[\\mathbb\{E\}\_\{a\\sim\\pi\(\\cdot\|z\)\}Q^\{\\beta\}\(z,a\)\-\\mathbb\{E\}\_\{a\\sim\\beta\(\\cdot\|z\)\}Q^\{\\beta\}\(z,a\)\\right\]\-C\_\{1\}D\_\{\\mathrm\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\)\.\(35\)
Next, we relate the true value functionQβQ^\{\\beta\}to its approximationQ^\\hat\{Q\}\. For any policyμ\\mu, we have

\|𝔼z∼dβ,a∼μ\(⋅\|z\)​\[Qβ​\(z,a\)−Q^​\(z,a\)\]\|≤‖Qβ−Q^‖∞\.\\left\|\\mathbb\{E\}\_\{z\\sim d^\{\\beta\},a\\sim\\mu\(\\cdot\|z\)\}\\left\[Q^\{\\beta\}\(z,a\)\-\\hat\{Q\}\(z,a\)\\right\]\\right\|\\leq\\\|Q^\{\\beta\}\-\\hat\{Q\}\\\|\_\{\\infty\}\.\(36\)Applying this inequality to bothμ=π\\mu=\\piandμ=β\\mu=\\betagives

𝔼z∼dβ​\[𝔼a∼π\(⋅\|z\)​Qβ​\(z,a\)−𝔼a∼β\(⋅\|z\)​Qβ​\(z,a\)\]\\displaystyle\\mathbb\{E\}\_\{z\\sim d^\{\\beta\}\}\\left\[\\mathbb\{E\}\_\{a\\sim\\pi\(\\cdot\|z\)\}Q^\{\\beta\}\(z,a\)\-\\mathbb\{E\}\_\{a\\sim\\beta\(\\cdot\|z\)\}Q^\{\\beta\}\(z,a\)\\right\]≥𝔼z∼dβ​\[𝔼a∼π\(⋅\|z\)​Q^​\(z,a\)−𝔼a∼β\(⋅\|z\)​Q^​\(z,a\)\]−2​‖Qβ−Q^‖∞\.\\displaystyle\\geq\\mathbb\{E\}\_\{z\\sim d^\{\\beta\}\}\\left\[\\mathbb\{E\}\_\{a\\sim\\pi\(\\cdot\|z\)\}\\hat\{Q\}\(z,a\)\-\\mathbb\{E\}\_\{a\\sim\\beta\(\\cdot\|z\)\}\\hat\{Q\}\(z,a\)\\right\]\-2\\\|Q^\{\\beta\}\-\\hat\{Q\}\\\|\_\{\\infty\}\.\(37\)By the definition of the surrogate objective, the first term on the right\-hand side is exactly

J^β​\(π\)−J^β​\(β\)\.\\hat\{J\}\_\{\\beta\}\(\\pi\)\-\\hat\{J\}\_\{\\beta\}\(\\beta\)\.\(38\)Thus,

J​\(π\)−J​\(β\)\\displaystyle J\(\\pi\)\-J\(\\beta\)≥11−γ​\(J^β​\(π\)−J^β​\(β\)\)−C1​DKLmax​\(π∥β\)−C2​‖Qβ−Q^‖∞,\\displaystyle\\geq\\frac\{1\}\{1\-\\gamma\}\\left\(\\hat\{J\}\_\{\\beta\}\(\\pi\)\-\\hat\{J\}\_\{\\beta\}\(\\beta\)\\right\)\-C\_\{1\}D\_\{\\mathrm\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\)\-C\_\{2\}\\\|Q^\{\\beta\}\-\\hat\{Q\}\\\|\_\{\\infty\},\(39\)whereC1,C2=21−γ\>0C\_\{1\},C\_\{2\}=\\frac\{2\}\{1\-\\gamma\}\>0are constants depending onγ\\gammaand the boundedness of the value function\. This completes the proof\. ∎

##### Implication: search policy should be anchored to the non\-search behavioral policy\.

The above two results reveal a common source of failure when the search policyπ\\piis optimized without being anchored to the behavioral policyβ\\beta\. From \([4](https://arxiv.org/html/2605.26282#S3.E4)\), the Bellman update induced byπ\\pideviates from the Bellman update induced byβ\\betaby an amount controlled by the policy divergence:

‖Tπ​Q^−Tβ​Q^‖∞≤γ​‖Q^‖∞​2​DKLmax​\(π∥β\)\.\\\|T^\{\\pi\}\\hat\{Q\}\-T^\{\\beta\}\\hat\{Q\}\\\|\_\{\\infty\}\\leq\\gamma\\\|\\hat\{Q\}\\\|\_\{\\infty\}\\sqrt\{2D\_\{\\mathrm\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\)\}\.\(40\)This means that whenπ\\pimoves far away fromβ\\beta, the value functionQ^\\hat\{Q\}, which is learned from data generated byβ\\beta, is no longer updated consistently with the distribution on which it was trained\. As a result, applyingTπT^\{\\pi\}can propagate value estimates into latent regions that are poorly covered by the behavioral policy, whereQ^\\hat\{Q\}may be inaccurate\.

Meanwhile, \([5](https://arxiv.org/html/2605.26282#S3.E5)\) shows that the true policy improvement is lower bounded by the estimated improvement only up to two error terms:

J​\(π\)−J​\(β\)≥11−γ​\(J^β​\(π\)−J^β​\(β\)\)−C1​DKLmax​\(π∥β\)−C2​‖Qβ−Q^‖∞\.J\(\\pi\)\-J\(\\beta\)\\geq\\frac\{1\}\{1\-\\gamma\}\\left\(\\hat\{J\}\_\{\\beta\}\(\\pi\)\-\\hat\{J\}\_\{\\beta\}\(\\beta\)\\right\)\-C\_\{1\}D\_\{\\mathrm\{KL\}\}^\{\\max\}\(\\pi\\\|\\beta\)\-C\_\{2\}\\\|Q^\{\\beta\}\-\\hat\{Q\}\\\|\_\{\\infty\}\.\(41\)The first penalty is the TRPO\-style distribution\-shift term\[[51](https://arxiv.org/html/2605.26282#bib.bib69)\], which indicates that policy improvement is inherently local around the behavioral \(or reference\) policyβ\\beta\. The second penalty captures the value estimation gap introduced by using the learned value functionQ^\\hat\{Q\}instead of the true value functionQβQ^\{\\beta\}\. Together, these two terms imply that a large surrogate improvement underJ^β​\(π\)\\hat\{J\}\_\{\\beta\}\(\\pi\)does not necessarily translate into true improvement whenπ\\pideviates substantially fromβ\\beta\.

Therefore, the two bottlenecks reinforce each other\. An unconstrained search policy may exploit spurious high\-value regions ofQ^\\hat\{Q\}and move away from the data\-supported behavior distribution\. This increases the KL divergence betweenπ\\piandβ\\beta, enlarges the Bellman update mismatch in \([4](https://arxiv.org/html/2605.26282#S3.E4)\), and further increases the value estimation error appearing in \([5](https://arxiv.org/html/2605.26282#S3.E5)\)\. Consequently, the search procedure can enter a self\-reinforcing cycle of distribution shift and value overestimation, where inaccurate value estimates attract the policy, and the shifted policy further amplifies value errors\.

This analysis suggests that search in latent world models should not be performed as an unconstrained policy improvement step\. Instead, the search policy should be regularized or anchored to the behavioral policyβ\\beta, for example through a KL constraint or a policy\-prior term\. Such an anchoring mechanism keeps the search policy within the region where the learned value function is reliable, thereby reducing Bellman mismatch, controlling distribution shift, and enabling more stable monotonic improvement within the world model\.

### C\.3Derivation of Exact Score Function in \([10](https://arxiv.org/html/2605.26282#S3.E10)\)

###### Lemma C\.1\(Exact Score Function\)\.

Consider the forward and reverse diffusion processes defined in \([17](https://arxiv.org/html/2605.26282#A2.E17)\) and \([18](https://arxiv.org/html/2605.26282#A2.E18)\) in the latent world model over𝒵\\mathcal\{Z\}\. For a latent stateztz\_\{t\}and a noisy action sequenceat:t\+Hτa\_\{t:t\+H\}^\{\\tau\}at diffusion stepτ\\tau, define a proposal distribution

qτ​\(at:t\+H0\|at:t\+Hτ,zt\)=𝒩​\(at:t\+H0;1α¯τ​at:t\+Hτ,1−α¯τα¯τ​I\)\.q^\{\\tau\}\(a\_\{t:t\+H\}^\{0\}\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)=\\mathcal\{N\}\\left\(a\_\{t:t\+H\}^\{0\};\\frac\{1\}\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}a\_\{t:t\+H\}^\{\\tau\},\\frac\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\{\\bar\{\\alpha\}^\{\\tau\}\}I\\right\)\.\(42\)Then the exact score function with respect toat:t\+Hτa\_\{t:t\+H\}^\{\\tau\}can be written as

ϕ​\(zt,at:t\+Hτ,τ\)\\displaystyle\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)=−at:t\+Hτ1−α¯τ\+α¯τ1−α¯τ​𝔼at:t\+H0∼qτ\(⋅\|at:t\+Hτ\)​\[w​\(at:t\+H0\|τ\)​at:t\+H0\|τ\],\\displaystyle=\-\\frac\{a^\{\\tau\}\_\{t:t\+H\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\+\\frac\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\\mathbb\{E\}\_\{a\_\{t:t\+H\}^\{0\}\\sim q\_\{\\tau\}\(\\cdot\|a\_\{t:t\+H\}^\{\\tau\}\)\}\\left\[w\(a^\{0\|\\tau\}\_\{t:t\+H\}\)a\_\{t:t\+H\}^\{0\|\\tau\}\\right\],\(43\)where the normalized importance weight is

w​\(at:t\+H0\|τ\)=exp⁡\(G~​\(zt,at:t\+H0\)/κ\)𝔼a~t:t\+H0∼qτ\(⋅\|at:t\+Hτ\)​\[exp⁡\(G~​\(zt,a~t:t\+H0\)/κ\)\]\.w\(a\_\{t:t\+H\}^\{0\|\\tau\}\)=\\frac\{\\exp\\left\(\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\}\)/\\kappa\\right\)\}\{\\mathbb\{E\}\_\{\\tilde\{a\}\_\{t:t\+H\}^\{0\}\\sim q\_\{\\tau\}\(\\cdot\|a\_\{t:t\+H\}^\{\\tau\}\)\}\\left\[\\exp\\left\(\\tilde\{G\}\(z\_\{t\},\\tilde\{a\}\_\{t:t\+H\}^\{0\}\)/\\kappa\\right\)\\right\]\}\.\(44\)

###### Proof\.

We start from the definition of the score function\. For a fixed latent stateztz\_\{t\}, the marginal distribution of the noisy action sequenceat:t\+Hτa\_\{t:t\+H\}^\{\\tau\}is given by

πϕτ​\(at:t\+Hτ\|zt\)=∫p​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​𝑑at:t\+H0\.\\pi\_\{\\phi\}^\{\\tau\}\(a\_\{t:t\+H\}^\{\\tau\}\|z\_\{t\}\)=\\int p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)da\_\{t:t\+H\}^\{0\}\.\(45\)Therefore,

ϕ​\(zt,at:t\+Hτ,τ\)=∇at:t\+Hτlog⁡πϕτ​\(at:t\+Hτ\|zt\)=∇at:t\+Hτ​∫p​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​𝑑at:t\+H0∫p​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​𝑑at:t\+H0=∫∇at:t\+Hτp​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​𝑑at:t\+H0∫p​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​𝑑at:t\+H0\.\\begin\{split\}\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)&=\\nabla\_\{a\_\{t:t\+H\}^\{\\tau\}\}\\log\\pi^\{\\tau\}\_\{\\phi\}\(a\_\{t:t\+H\}^\{\\tau\}\|z\_\{t\}\)\\\\ &=\\frac\{\\nabla\_\{a\_\{t:t\+H\}^\{\\tau\}\}\\int p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)da\_\{t:t\+H\}^\{0\}\}\{\\int p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)da\_\{t:t\+H\}^\{0\}\}\\\\ &=\\frac\{\\int\\nabla\_\{a\_\{t:t\+H\}^\{\\tau\}\}p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)da\_\{t:t\+H\}^\{0\}\}\{\\int p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)da\_\{t:t\+H\}^\{0\}\}\.\\end\{split\}\(46\)
Under the DDPM forward process, the transition kernel is

p​\(at:t\+Hτ\|at:t\+H0,zt\)=𝒩​\(at:t\+Hτ;α¯τ​at:t\+H0,\(1−α¯τ\)​I\)∝exp⁡\(−12​\(at:t\+Hτ−α¯τ​at:t\+H0\)⊤​\(at:t\+Hτ−α¯τ​at:t\+H0\)1−α¯τ\)\.\\begin\{split\}p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)&=\\mathcal\{N\}\\left\(a\_\{t:t\+H\}^\{\\tau\};\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}a\_\{t:t\+H\}^\{0\},\(1\-\\bar\{\\alpha\}^\{\\tau\}\)I\\right\)\\\\ &\\propto\\exp\\left\(\-\\frac\{1\}\{2\}\\frac\{\\left\(a\_\{t:t\+H\}^\{\\tau\}\-\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}a\_\{t:t\+H\}^\{0\}\\right\)^\{\\top\}\\left\(a\_\{t:t\+H\}^\{\\tau\}\-\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}a\_\{t:t\+H\}^\{0\}\\right\)\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\\right\)\.\\end\{split\}\(47\)Hence,

∇at:t\+Hτp​\(at:t\+Hτ\|at:t\+H0,zt\)=−at:t\+Hτ−α¯τ​at:t\+H01−α¯τ​p​\(at:t\+Hτ\|at:t\+H0,zt\)\.\\nabla\_\{a\_\{t:t\+H\}^\{\\tau\}\}p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)=\-\\frac\{a\_\{t:t\+H\}^\{\\tau\}\-\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}a\_\{t:t\+H\}^\{0\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\.\(48\)Substituting this identity into \([46](https://arxiv.org/html/2605.26282#A3.E46)\), we obtain

ϕ​\(zt,at:t\+Hτ,τ\)=∫−at:t\+Hτ−α¯τ​at:t\+H01−α¯τ​p​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​d​at:t\+H0∫p​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​𝑑at:t\+H0=−at:t\+Hτ1−α¯τ\+α¯τ1−α¯τ​∫p​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​at:t\+H0​𝑑at:t\+H0∫p​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​𝑑at:t\+H0\.\\begin\{split\}\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)&=\\frac\{\\int\-\\frac\{a\_\{t:t\+H\}^\{\\tau\}\-\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}a\_\{t:t\+H\}^\{0\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)da\_\{t:t\+H\}^\{0\}\}\{\\int p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)da\_\{t:t\+H\}^\{0\}\}\\\\ &=\-\\frac\{a\_\{t:t\+H\}^\{\\tau\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\+\\frac\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\\frac\{\\int p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)a\_\{t:t\+H\}^\{0\}da\_\{t:t\+H\}^\{0\}\}\{\\int p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)da\_\{t:t\+H\}^\{0\}\}\.\\end\{split\}\(49\)Thus, computing the exact score reduces to estimating the posterior mean

𝔼​\[at:t\+H0\|at:t\+Hτ,zt\]=∫p​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​at:t\+H0​𝑑at:t\+H0∫p​\(at:t\+Hτ\|at:t\+H0,zt\)​πϕ0​\(at:t\+H0\|zt\)​𝑑at:t\+H0\.\\begin\{split\}\\mathbb\{E\}\\left\[a\_\{t:t\+H\}^\{0\}\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\\right\]&=\\frac\{\\int p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)a\_\{t:t\+H\}^\{0\}da\_\{t:t\+H\}^\{0\}\}\{\\int p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)da\_\{t:t\+H\}^\{0\}\}\.\\end\{split\}\(50\)
To express this posterior mean using proposal sampling, we introduce the Gaussian proposal distribution

qτ​\(at:t\+H0\|at:t\+Hτ,zt\)=𝒩​\(at:t\+H0;1α¯τ​at:t\+Hτ,1−α¯τα¯τ​I\)\.q\_\{\\tau\}\(a\_\{t:t\+H\}^\{0\}\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)=\\mathcal\{N\}\\left\(a\_\{t:t\+H\}^\{0\};\\frac\{1\}\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}a\_\{t:t\+H\}^\{\\tau\},\\frac\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\{\\bar\{\\alpha\}^\{\\tau\}\}I\\right\)\.\(51\)This proposal is proportional to the DDPM likelihoodp​\(at:t\+Hτ\|at:t\+H0,zt\)p\(a\_\{t:t\+H\}^\{\\tau\}\|a\_\{t:t\+H\}^\{0\},z\_\{t\}\)when viewed as a function ofat:t\+H0a\_\{t:t\+H\}^\{0\}\. Therefore, \([50](https://arxiv.org/html/2605.26282#A3.E50)\) can be rewritten as

𝔼​\[at:t\+H0\|at:t\+Hτ,zt\]=𝔼at:t\+H0\|τ∼qτ\(⋅\|at:t\+Hτ,zt\)​\[πϕ0​\(at:t\+H0\|τ\|zt\)​at:t\+H0\|τ\]𝔼at:t\+H0\|τ∼qτ\(⋅\|at:t\+Hτ,zt\)​\[πϕ0​\(at:t\+H0\|τ\|zt\)\]\.\\begin\{split\}\\mathbb\{E\}\\left\[a\_\{t:t\+H\}^\{0\}\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\\right\]&=\\frac\{\\mathbb\{E\}\_\{a\_\{t:t\+H\}^\{0\|\\tau\}\\sim q\_\{\\tau\}\(\\cdot\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)\}\\left\[\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\|\\tau\}\|z\_\{t\}\)a\_\{t:t\+H\}^\{0\|\\tau\}\\right\]\}\{\\mathbb\{E\}\_\{a\_\{t:t\+H\}^\{0\|\\tau\}\\sim q\_\{\\tau\}\(\\cdot\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)\}\\left\[\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\|\\tau\}\|z\_\{t\}\)\\right\]\}\.\\end\{split\}\(52\)
Substituting the posterior mean back into the score expression gives

ϕ​\(zt,at:t\+Hτ,τ\)=−at:t\+Hτ1−α¯τ\+α¯τ1−α¯τ​𝔼at:t\+H0∼qτ\(⋅\|at:t\+Hτ,zt\)​\[w​\(at:t\+H0\|τ\)​at:t\+H0\|τ\],\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)=\-\\frac\{a\_\{t:t\+H\}^\{\\tau\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\+\\frac\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\\mathbb\{E\}\_\{a\_\{t:t\+H\}^\{0\}\\sim q\_\{\\tau\}\(\\cdot\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)\}\\left\[w\(a\_\{t:t\+H\}^\{0\|\\tau\}\)a\_\{t:t\+H\}^\{0\|\\tau\}\\right\],\(53\)where the normalized importance weight is defined as

w​\(at:t\+H0\|τ\)=πϕ0​\(at:t\+H0\|τ\|zt\)𝔼a~t:t\+H0\|τ∼qτ\(⋅\|at:t\+Hτ,zt\)​\[πϕ0​\(a~t:t\+H0\|τ\|zt\)\]\.w\(a\_\{t:t\+H\}^\{0\|\\tau\}\)=\\frac\{\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\|\\tau\}\|z\_\{t\}\)\}\{\\mathbb\{E\}\_\{\\tilde\{a\}\_\{t:t\+H\}^\{0\|\\tau\}\\sim q\_\{\\tau\}\(\\cdot\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)\}\\left\[\\pi\_\{\\phi\}^\{0\}\(\\tilde\{a\}\_\{t:t\+H\}^\{0\|\\tau\}\|z\_\{t\}\)\\right\]\}\.\(54\)Finally, under the entropy\-regularized policy induced by the trajectory return,

πϕ0​\(at:t\+H0\|τ\|zt\)∝exp⁡\(G~​\(zt,at:t\+H0\|τ\)κ\)\.\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\|\\tau\}\|z\_\{t\}\)\\propto\\exp\\left\(\\frac\{\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\|\\tau\}\)\}\{\\kappa\}\\right\)\.\(55\)Hence, the normalized importance weight becomes

w​\(at:t\+H0\|τ\)=exp⁡\(G~​\(zt,at:t\+H0\|τ\)/κ\)𝔼a~t:t\+H0\|τ∼qτ\(⋅\|at:t\+Hτ,zt\)​\[exp⁡\(G~​\(zt,a~t:t\+H0\|τ\)/κ\)\]\.w\(a\_\{t:t\+H\}^\{0\|\\tau\}\)=\\frac\{\\exp\\left\(\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\|\\tau\}\)/\\kappa\\right\)\}\{\\mathbb\{E\}\_\{\\tilde\{a\}\_\{t:t\+H\}^\{0\|\\tau\}\\sim q\_\{\\tau\}\(\\cdot\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)\}\\left\[\\exp\\left\(\\tilde\{G\}\(z\_\{t\},\\tilde\{a\}\_\{t:t\+H\}^\{0\|\\tau\}\)/\\kappa\\right\)\\right\]\}\.\(56\)This completes the proof\. ∎

##### What is the role of the world model in the computation of the score function?

The world model provides the latent dynamics used to evaluate each candidate clean action sequence sampled during the score computation\. In the derivation above, the proposal distributionqτ​\(at:t\+H0\|τ∣at:t\+Hτ,zt\)q\_\{\\tau\}\(a\_\{t:t\+H\}^\{0\|\\tau\}\\mid a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)is obtained from the DDPM forward kernel and is used to sample clean action sequencesat:t\+H0\|τa\_\{t:t\+H\}^\{0\|\\tau\}that are compatible with the noisy action sequenceat:t\+Hτa\_\{t:t\+H\}^\{\\tau\}\. For each sampled action sequence, the world model recursively rolls out the latent dynamics

zh\+10\|τ=ℱ​\(zh0\|τ,ah0\|τ\),h=t,…,t\+H−1,z\_\{h\+1\}^\{0\|\\tau\}=\\mathcal\{F\}\(z\_\{h\}^\{0\|\\tau\},a\_\{h\}^\{0\|\\tau\}\),\\qquad h=t,\\ldots,t\+H\-1,\(57\)starting from the current latent statezt0\|τ=ztz\_\{t\}^\{0\|\\tau\}=z\_\{t\}\. This rollout produces a latent trajectory

\(zt,at0\|τ,zt\+10\|τ,at\+10\|τ,…,zt\+H0\|τ,at\+H0\|τ\),\\left\(z\_\{t\},a\_\{t\}^\{0\|\\tau\},z\_\{t\+1\}^\{0\|\\tau\},a\_\{t\+1\}^\{0\|\\tau\},\\ldots,z\_\{t\+H\}^\{0\|\\tau\},a\_\{t\+H\}^\{0\|\\tau\}\\right\),\(58\)which is then used to compute the trajectory\-level return

G~​\(zt,at:t\+H0\|τ\)≔∑h=tt\+H−1γh−t​ℛ~​\(zh0\|τ,ah0\|τ\)\+γH​Q^​\(zt\+H0\|τ,at\+H0\|τ\)−η​E​\(zt\+H0\|τ,at\+H0\|τ\)\.\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\|\\tau\}\)\\coloneqq\\sum\_\{h=t\}^\{t\+H\-1\}\\gamma^\{h\-t\}\\tilde\{\\mathcal\{R\}\}\(z\_\{h\}^\{0\|\\tau\},a\_\{h\}^\{0\|\\tau\}\)\+\\gamma^\{H\}\\hat\{Q\}\(z\_\{t\+H\}^\{0\|\\tau\},a\_\{t\+H\}^\{0\|\\tau\}\)\-\\eta E\(z\_\{t\+H\}^\{0\|\\tau\},a\_\{t\+H\}^\{0\|\\tau\}\)\.\(59\)The returnG~​\(zt,at:t\+H0\|τ\)\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\|\\tau\}\)determines the normalized importance weight

w​\(at:t\+H0\|τ\)=exp⁡\(G~​\(zt,at:t\+H0\|τ\)/κ\)𝔼a~t:t\+H0\|τ∼qτ\(⋅\|at:t\+Hτ,zt\)​\[exp⁡\(G~​\(zt,a~t:t\+H0\|τ\)/κ\)\]\.w\(a\_\{t:t\+H\}^\{0\|\\tau\}\)=\\frac\{\\exp\\left\(\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\|\\tau\}\)/\\kappa\\right\)\}\{\\mathbb\{E\}\_\{\\tilde\{a\}\_\{t:t\+H\}^\{0\|\\tau\}\\sim q\_\{\\tau\}\(\\cdot\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)\}\\left\[\\exp\\left\(\\tilde\{G\}\(z\_\{t\},\\tilde\{a\}\_\{t:t\+H\}^\{0\|\\tau\}\)/\\kappa\\right\)\\right\]\}\.\(60\)Thus, the world model does not directly define the Gaussian proposal used for denoising\. Instead, it provides the model\-based evaluation signal that reweights the sampled clean action sequences according to their predicted trajectory returns\. The exact score is then computed as a return\-weighted posterior mean:

ϕ​\(zt,at:t\+Hτ,τ\)=−at:t\+Hτ1−α¯τ\+α¯τ1−α¯τ​𝔼at:t\+H0\|τ∼qτ\(⋅\|at:t\+Hτ,zt\)​\[w​\(at:t\+H0\|τ\)​at:t\+H0\|τ\]\.\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)=\-\\frac\{a^\{\\tau\}\_\{t:t\+H\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\+\\frac\{\\sqrt\{\\bar\{\\alpha\}^\{\\tau\}\}\}\{1\-\\bar\{\\alpha\}^\{\\tau\}\}\\mathbb\{E\}\_\{a\_\{t:t\+H\}^\{0\|\\tau\}\\sim q\_\{\\tau\}\(\\cdot\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)\}\\left\[w\(a^\{0\|\\tau\}\_\{t:t\+H\}\)a\_\{t:t\+H\}^\{0\|\\tau\}\\right\]\.\(61\)In this sense, the world model converts each sampled clean action sequence into a trajectory\-level value estimate, and the score function shifts the noisy action sequence toward clean action sequences that achieve higher predicted return under the latent dynamics\.

A more straightforward abstraction is that the reverse transition kernel defines an iterative policy refinement process\. Starting from a simple Gaussian prior over action sequences, the reverse diffusion process gradually transforms this prior into the optimized action\-sequence policy:

πϕN​\(at:t\+HN\|zt\)⏟Gaussian prior→πϕN−1​\(at:t\+HN−1\|at:t\+HN,zt\)πϕN−1​\(at:t\+HN−1\|zt\)​⋯→πϕ0​\(at:t\+H0\|at:t\+H1,zt\)πϕ0​\(at:t\+H0\|zt\)⏟clean policy\.\\underbrace\{\\pi\_\{\\phi\}^\{N\}\(a\_\{t:t\+H\}^\{N\}\|z\_\{t\}\)\}\_\{\\text\{Gaussian prior\}\}\\xrightarrow\{\\pi\_\{\\phi\}^\{N\-1\}\(a\_\{t:t\+H\}^\{N\-1\}\|a\_\{t:t\+H\}^\{N\},z\_\{t\}\)\}\\pi\_\{\\phi\}^\{N\-1\}\(a\_\{t:t\+H\}^\{N\-1\}\|z\_\{t\}\)\\cdots\\xrightarrow\{\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|a\_\{t:t\+H\}^\{1\},z\_\{t\}\)\}\\underbrace\{\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)\}\_\{\\text\{clean policy\}\}\.\(62\)Equivalently, the final refined policy is obtained by marginalizing over all intermediate denoising steps:

πϕ0​\(at:t\+H0\|zt\)=∫at:t\+H1:NπϕN​\(at:t\+HN\|zt\)​∏τ=N1πϕτ−1​\(at:t\+Hτ−1\|at:t\+Hτ,zt\)​d​at:t\+H1:N\.\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)=\\int\_\{a\_\{t:t\+H\}^\{1:N\}\}\\pi^\{N\}\_\{\\phi\}\(a\_\{t:t\+H\}^\{N\}\|z\_\{t\}\)\\prod\_\{\\tau=N\}^\{1\}\\pi\_\{\\phi\}^\{\\tau\-1\}\(a\_\{t:t\+H\}^\{\\tau\-1\}\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)da\_\{t:t\+H\}^\{1:N\}\.\(63\)At each reverse step, the transition kernelπϕτ−1​\(at:t\+Hτ−1\|at:t\+Hτ,zt\)\\pi\_\{\\phi\}^\{\\tau\-1\}\(a\_\{t:t\+H\}^\{\\tau\-1\}\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\)is guided by the score functionϕ​\(zt,at:t\+Hτ,τ\)\\phi\(z\_\{t\},a\_\{t:t\+H\}^\{\\tau\},\\tau\)\. This score is computed by sampling candidate clean action sequences from the DDPM posterior proposalqτ\(⋅\|at:t\+Hτ,zt\)q\_\{\\tau\}\(\\cdot\|a\_\{t:t\+H\}^\{\\tau\},z\_\{t\}\), rolling them out through the world model, and reweighting them according to their predicted returns\.Consequently, the world model redefines the target for score matching by introducing model\-based evaluations\. This shift directly propagates to the stepwise transition kernel, thereby rendering the entire reverse diffusion process equivalent to policy optimization\.

From this perspective, reverse diffusion can be viewed as model\-based policy optimization in action\-sequence space\. The Gaussian prior provides exploration over possible action sequences, while the world\-model\-evaluated score progressively reshapes this prior into a policy that concentrates on action sequences with higher predicted return\. Therefore, the policy refinement in \([2](https://arxiv.org/html/2605.26282#S2.E2)\) is not merely a generative denoising procedure; it is also a KL\-regularized policy optimization process where the world model supplies the trajectory\-level objective used to guide each denoising step \(stated in Theorem[3\.2](https://arxiv.org/html/2605.26282#S3.SS2)\)\.

### C\.4Proof of Theorem[3\.2](https://arxiv.org/html/2605.26282#S3.SS2)

###### Proof\.

We prove Theorem[3\.2](https://arxiv.org/html/2605.26282#S3.SS2)in two steps: \(1\) demonstrating that the optimal distribution of the KL\-constrained policy optimization problem formulated in \([8](https://arxiv.org/html/2605.26282#S3.E8)\) follows a Gibbs distribution, and \(2\) showing that score matching within the world model is equivalent to sampling from this optimal policy\.

Step 1\. Deriving the Optimal Gibbs Policy\.

For a fixed latent stateztz\_\{t\}, the latent trajectory is deterministically generated by the learned world model

zh\+1=ℱ​\(zh,ah\),h=t,…,t\+H−1\.z\_\{h\+1\}=\\mathcal\{F\}\(z\_\{h\},a\_\{h\}\),\\qquad h=t,\\ldots,t\+H\-1\.\(64\)Therefore, onceztz\_\{t\}and the action sequenceat:t\+Ha\_\{t:t\+H\}are given, the trajectory returnG​\(zt,at:t\+H\)G\(z\_\{t\},a\_\{t:t\+H\}\)is completely determined by the rollout of the world model\. Equivalently, the world\-model constraint restricts feasible trajectories to those satisfyingzh\+1=ℱ​\(zh,ah\)z\_\{h\+1\}=\\mathcal\{F\}\(z\_\{h\},a\_\{h\}\)\.

We first consider the Lagrangian formulation of the KL\-constrained policy optimization problem \([8](https://arxiv.org/html/2605.26282#S3.E8)\)\. Letη\>0\\eta\>0be the KKT multiplier associated with the KL divergence constraint\. The regularized objective is:

maxπϕ\(⋅\|zt\)\{𝔼at:t\+H∼πϕ\(⋅\|zt\)\[G\(zt,at:t\+H\)\]−ηDKL\(πϕ\(⋅\|zt\)∥β\(⋅\|zt\)\)\}\.\\max\_\{\\pi\_\{\\phi\}\(\\cdot\|z\_\{t\}\)\}\\left\\\{\\mathbb\{E\}\_\{a\_\{t:t\+H\}\\sim\\pi\_\{\\phi\}\(\\cdot\|z\_\{t\}\)\}\\left\[G\(z\_\{t\},a\_\{t:t\+H\}\)\\right\]\-\\eta D\_\{\\mathrm\{KL\}\}\\left\(\\pi\_\{\\phi\}\(\\cdot\|z\_\{t\}\)\\\|\\beta\(\\cdot\|z\_\{t\}\)\\right\)\\right\\\}\.\(65\)
Expanding the KL term, the objective becomes

𝒥​\(πϕ\)\\displaystyle\\mathcal\{J\}\(\\pi\_\{\\phi\}\)=∫πϕ​\(at:t\+H\|zt\)​G​\(zt,at:t\+H\)​𝑑at:t\+H\\displaystyle=\\int\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)G\(z\_\{t\},a\_\{t:t\+H\}\)da\_\{t:t\+H\}−η​∫πϕ​\(at:t\+H\|zt\)​log⁡πϕ​\(at:t\+H\|zt\)β​\(at:t\+H\|zt\)​d​at:t\+H\.\\displaystyle\\quad\-\\eta\\int\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)\\log\\frac\{\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)\}\{\\beta\(a\_\{t:t\+H\}\|z\_\{t\}\)\}da\_\{t:t\+H\}\.\(66\)
Note that the behavior priorβ\\betaover the trajectory incorporates the deterministic latent dynamics constraint of world models:

β​\(at:t\+H\|zt\)=∏h=tt\+Hβ​\(ah\|zh\)⋅∏h=tt\+H−1𝟙\{zh\+1=ℱ​\(zh,ah\)\}\.\\beta\(a\_\{t:t\+H\}\|z\_\{t\}\)=\\prod\_\{h=t\}^\{t\+H\}\\beta\(a\_\{h\}\|z\_\{h\}\)\\cdot\\prod\_\{h=t\}^\{t\+H\-1\}\\mathbb\{1\}\_\{\\\{z\_\{h\+1\}=\\mathcal\{F\}\(z\_\{h\},a\_\{h\}\)\\\}\}\.\(67\)
To enforce the probability normalization constraint∫πϕ​\(at:t\+H\|zt\)​𝑑at:t\+H=1\\int\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)da\_\{t:t\+H\}=1, we introduce a second Lagrange multiplierλ\\lambdaand define the full Lagrangian:

ℒ​\(πϕ,λ\)\\displaystyle\\mathcal\{L\}\(\\pi\_\{\\phi\},\\lambda\)=∫πϕ​\(at:t\+H\|zt\)​G​\(zt,at:t\+H\)​𝑑at:t\+H\\displaystyle=\\int\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)G\(z\_\{t\},a\_\{t:t\+H\}\)da\_\{t:t\+H\}−η​∫πϕ​\(at:t\+H\|zt\)​log⁡πϕ​\(at:t\+H\|zt\)β​\(at:t\+H\|zt\)​d​at:t\+H\\displaystyle\\quad\-\\eta\\int\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)\\log\\frac\{\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)\}\{\\beta\(a\_\{t:t\+H\}\|z\_\{t\}\)\}da\_\{t:t\+H\}\+λ​\(∫πϕ​\(at:t\+H\|zt\)​𝑑at:t\+H−1\)\.\\displaystyle\\quad\+\\lambda\\left\(\\int\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)da\_\{t:t\+H\}\-1\\right\)\.\(68\)
Taking the functional derivative with respect toπϕ​\(at:t\+H\|zt\)\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)and setting it to zero gives

0=G​\(zt,at:t\+H\)−η​\(log⁡πϕ​\(at:t\+H\|zt\)β​\(at:t\+H\|zt\)\+1\)\+λ\.0=G\(z\_\{t\},a\_\{t:t\+H\}\)\-\\eta\\left\(\\log\\frac\{\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)\}\{\\beta\(a\_\{t:t\+H\}\|z\_\{t\}\)\}\+1\\right\)\+\\lambda\.\(69\)
Rearranging for the density ratio yields

log⁡πϕ​\(at:t\+H\|zt\)β​\(at:t\+H\|zt\)=G​\(zt,at:t\+H\)η\+λ−ηη\.\\log\\frac\{\\pi\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)\}\{\\beta\(a\_\{t:t\+H\}\|z\_\{t\}\)\}=\\frac\{G\(z\_\{t\},a\_\{t:t\+H\}\)\}\{\\eta\}\+\\frac\{\\lambda\-\\eta\}\{\\eta\}\.\(70\)
Exponentiating both sides, we obtain the optimal policy:

πϕ∗​\(at:t\+H\|zt\)=1Z​\(zt\)​β​\(at:t\+H\|zt\)​exp⁡\(G​\(zt,at:t\+H\)η\),\\pi\_\{\\phi\}^\{\*\}\(a\_\{t:t\+H\}\|z\_\{t\}\)=\\frac\{1\}\{Z\(z\_\{t\}\)\}\\beta\(a\_\{t:t\+H\}\|z\_\{t\}\)\\exp\\left\(\\frac\{G\(z\_\{t\},a\_\{t:t\+H\}\)\}\{\\eta\}\\right\),\(71\)whereZ​\(zt\)Z\(z\_\{t\}\)is the partition function\. Substituting the factorized form ofβ\\beta, the optimal distribution can be expressed as:

πϕ∗​\(at:t\+H\|zt\)∝∏h=tt\+Hβ​\(ah\|zh\)⋅exp⁡\(G​\(zt,at:t\+H\)η\)⋅∏h=tt\+H−1𝟙\{zh\+1=ℱ​\(zh,ah\)\}\.\\pi\_\{\\phi\}^\{\*\}\(a\_\{t:t\+H\}\|z\_\{t\}\)\\propto\\prod\_\{h=t\}^\{t\+H\}\\beta\(a\_\{h\}\|z\_\{h\}\)\\cdot\\exp\\left\(\\frac\{G\(z\_\{t\},a\_\{t:t\+H\}\)\}\{\\eta\}\\right\)\\cdot\\prod\_\{h=t\}^\{t\+H\-1\}\\mathbb\{1\}\_\{\\\{z\_\{h\+1\}=\\mathcal\{F\}\(z\_\{h\},a\_\{h\}\)\\\}\}\.\(72\)
Step 2\. Generation of the Gibbs Distribution via Score Matching\.

By Lemma[C\.1](https://arxiv.org/html/2605.26282#A3.Thmtheorem1), an expression of the form \([10](https://arxiv.org/html/2605.26282#S3.E10)\) is analytically identical to the exact score function∇at:t\+Hτlog⁡πϕτ​\(at:t\+Hτ\|zt\)\\nabla\_\{a\_\{t:t\+H\}^\{\\tau\}\}\\log\\pi\_\{\\phi\}^\{\\tau\}\(a\_\{t:t\+H\}^\{\\tau\}\|z\_\{t\}\)of a forward\-diffused distribution, whose underlying base distributionπϕ0​\(at:t\+H0\|zt\)\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)is proportional to the unnormalized numerator of the weightw​\(at:t\+H0\)w\(a\_\{t:t\+H\}^\{0\}\)in \([56](https://arxiv.org/html/2605.26282#A3.E56)\)\. Therefore, according to the properties of variance\-preserving stochastic differential equations\[[54](https://arxiv.org/html/2605.26282#bib.bib64)\], the score matching condition implicitly defines the following base distribution:

πϕ0​\(at:t\+H0\|zt\)∝exp⁡\(G~​\(zt,at:t\+H0\)η\)\.\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)\\propto\\exp\\left\(\\frac\{\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\}\)\}\{\\eta\}\\right\)\.\(73\)
We now relate this implicit base distribution back to the policy derived in Step 1\. Recall that the energy\-regularized reward is defined as:

ℛ~​\(zh,ah\)=ℛ​\(zh,ah\)−ηγh−t​E​\(zh,ah\)\.\\tilde\{\\mathcal\{R\}\}\(z\_\{h\},a\_\{h\}\)=\\mathcal\{R\}\(z\_\{h\},a\_\{h\}\)\-\\frac\{\\eta\}\{\\gamma^\{h\-t\}\}E\(z\_\{h\},a\_\{h\}\)\.\(74\)Based on the energy minimization objective \([9](https://arxiv.org/html/2605.26282#S3.E9)\), the learned energy approximates the behavior prior up to a state\-dependent constant:−E​\(zh,ah\)≈log⁡β​\(ah\|zh\)\+C​\(zh\)\-E\(z\_\{h\},a\_\{h\}\)\\approx\\log\\beta\(a\_\{h\}\|z\_\{h\}\)\+C\(z\_\{h\}\)\. In continuous state spaces with continuous dynamics, the state variations over a short rollout horizonHHare relatively small\. Thus, we can reasonably approximateC​\(zh\)≈C​\(zt\)C\(z\_\{h\}\)\\approx C\(z\_\{t\}\), which acts as a fixed constant with respect to the future action sequenceat:t\+Ha\_\{t:t\+H\}\. Summing over the trajectory and accounting for the discount factor, the energy\-regularized returnG~\\tilde\{G\}decomposes into the original returnGGand the behavior prior term:

G~​\(zt,at:t\+H0\)η≈G​\(zt,at:t\+H0\)η\+∑h=tt\+Hlog⁡β​\(ah\|zh\)\+\(H\+1\)​C​\(zt\)\.\\frac\{\\tilde\{G\}\(z\_\{t\},a\_\{t:t\+H\}^\{0\}\)\}\{\\eta\}\\approx\\frac\{G\(z\_\{t\},a\_\{t:t\+H\}^\{0\}\)\}\{\\eta\}\+\\sum\_\{h=t\}^\{t\+H\}\\log\\beta\(a\_\{h\}\|z\_\{h\}\)\+\(H\+1\)C\(z\_\{t\}\)\.\(75\)
Substituting this decomposition into \([73](https://arxiv.org/html/2605.26282#A3.E73)\), the constant term\(H\+1\)​C​\(zt\)\(H\+1\)C\(z\_\{t\}\)is absorbed into the proportionality constant, and the base distribution can be explicitly written as:

πϕ0​\(at:t\+H0\|zt\)∝\(∏h=tt\+Hβ​\(ah\|zh\)\)​exp⁡\(G​\(zt,at:t\+H0\)η\)​∏h=tt\+H−1𝟙\{zh\+1=ℱ​\(zh,ah0\)\}\.\\pi\_\{\\phi\}^\{0\}\(a\_\{t:t\+H\}^\{0\}\|z\_\{t\}\)\\propto\\left\(\\prod\_\{h=t\}^\{t\+H\}\\beta\(a\_\{h\}\|z\_\{h\}\)\\right\)\\exp\\left\(\\frac\{G\(z\_\{t\},a\_\{t:t\+H\}^\{0\}\)\}\{\\eta\}\\right\)\\prod\_\{h=t\}^\{t\+H\-1\}\\mathbb\{1\}\_\{\\\{z\_\{h\+1\}=\\mathcal\{F\}\(z\_\{h\},a\_\{h\}^\{0\}\)\\\}\}\.\(76\)This distribution is exactly identical to the optimal Gibbs policyπϕ∗​\(at:t\+H\|zt\)\\pi^\{\*\}\_\{\\phi\}\(a\_\{t:t\+H\}\|z\_\{t\}\)derived in Step 1\.

Therefore, if the score matching condition is satisfied, the reverse diffusion process uses the exact score of the noisy marginal induced by the Gibbs policy\. Consequently, the reverse diffusion sampler generates clean action sequences exactly from the optimal KL\-regularized policy distribution\.

Conversely, suppose the KL\-constrained policy optimization problem attains its optimum\. Then, by the variational argument above, the optimal policy has the Gibbs form in \([71](https://arxiv.org/html/2605.26282#A3.E71)\)\. The diffusion marginal induced by this clean\-action distribution has a unique exact score, and by Lemma[C\.1](https://arxiv.org/html/2605.26282#A3.Thmtheorem1), this score is precisely the expression in \([10](https://arxiv.org/html/2605.26282#S3.E10)\)\. Thus, the score matching condition is satisfied\.

Combining both directions, the KL\-constrained policy optimization problem attains its Gibbs\-form optimum if and only if the score matching condition in \([10](https://arxiv.org/html/2605.26282#S3.E10)\) is satisfied\. Finally, to provide flexibility in controlling the entropy and return trade\-off during practical sampling,we relax the strict Lagrange multiplierη\\etato a general temperature parameterκ\>0\\kappa\>0, yielding the target Gibbs distribution presented in Theorem[3\.2](https://arxiv.org/html/2605.26282#S3.SS2)\. This completes the proof\. ∎

## Appendix DPractical Implementation and Hyperparameters

1

Input:Encoder

ℰ\\mathcal\{E\}, latent dynamics

ℱ\\mathcal\{F\}, task embedding

ℰe​n​v\\mathcal\{E\}\_\{env\}, reward function

ℛ\\mathcal\{R\}, value function

Q^\\hat\{Q\}, implicit energy function

EE, score network

ϕ^\\hat\{\\phi\}, diffusion steps

NN, Monte Carlo samples

TT
Initialize:Replay buffer

𝒟←∅\\mathcal\{D\}\\leftarrow\\emptyset
2

3//Warmup: Train the initial world model

4for*i=1,2,…,Nwarmupi=1,2,\\dots,N\_\{\\mathrm\{warmup\}\}*do

5Sample random action

ah∼Uniform​\(𝒜\)a\_\{h\}\\sim\\mathrm\{Uniform\}\(\\mathcal\{A\}\)
6Execute action:

\(sh,□\)←env\.step​\(ah\)\(s\_\{h\},\\Box\)\\leftarrow\\mathrm\{env\.step\}\(a\_\{h\}\)
7Store transition

\(sh,ah,rh,sh\+1\)\(s\_\{h\},a\_\{h\},r\_\{h\},s\_\{h\+1\}\)into

𝒟\\mathcal\{D\}
8

9Update

ℰ,ℱ,ℛ,ℰe​n​v,Q^,E\\mathcal\{E\},\\mathcal\{F\},\\mathcal\{R\},\\mathcal\{E\}\_\{env\},\\hat\{Q\},Eby minimizing the model loss in \([14](https://arxiv.org/html/2605.26282#S3.E14)\) using

𝒟\\mathcal\{D\}
10

11for*iterationk=1,2,…,Kk=1,2,\\dots,K*do

12//Step 1: Sampling data from the environment via the score network

13Encode current observation:

zh←ℰ​\(sh\)z\_\{h\}\\leftarrow\\mathcal\{E\}\(s\_\{h\}\)
14Sample action

aha\_\{h\}iteratively via the diffusion process in \([12](https://arxiv.org/html/2605.26282#S3.E12)\)

15Execute action:

\(sh,□\)←env\.step​\(ah\)\(s\_\{h\},\\Box\)\\leftarrow\\mathrm\{env\.step\}\(a\_\{h\}\)
16Store transition

\(sh,ah,rh,sh\+1\)\(s\_\{h\},a\_\{h\},r\_\{h\},s\_\{h\+1\}\)into

𝒟\\mathcal\{D\}
17

18//Step 2: World model update

19Sample a batch of transitions

ℬ∼𝒟\\mathcal\{B\}\\sim\\mathcal\{D\}
20Update

ℰ,ℱ,ℛ,ℰe​n​v,Q^,E\\mathcal\{E\},\\mathcal\{F\},\\mathcal\{R\},\\mathcal\{E\}\_\{env\},\\hat\{Q\},Eby minimizing the joint loss in \([14](https://arxiv.org/html/2605.26282#S3.E14)\) on batch

ℬ\\mathcal\{B\}
21

22//Step 3: Diffusion policy optimization

23Estimate the score function via \([15](https://arxiv.org/html/2605.26282#S3.E15)\) using imagined trajectories in the world model

24Update score network

ϕ^\\hat\{\\phi\}by minimizing the supervised

L2L^\{2\}loss in \([16](https://arxiv.org/html/2605.26282#S3.E16)\)

25

Algorithm 1Model\-Based Diffusion Policy Optimization \(MBDPO\)Note that Algorithm[1](https://arxiv.org/html/2605.26282#algorithm1)outlines theonlinelearning paradigm\. In theofflinepretraining setting, the agent does not interact with the environment to collect new data\. Instead, the replay buffer𝒟\\mathcal\{D\}is initialized with a static, pre\-collected dataset, and Step 1 \(sampling data via the score network and environment interaction\) is entirely omitted\. Consequently, the world model and diffusion policy are optimized, relying solely on the offline dataset\.

Furthermore, for theoffline\-to\-online fine\-tuningparadigm, the initial warmup phase is completely bypassed\. Instead, the algorithm is initialized with the parameters of the world model and the score network that have been pre\-trained on the offline dataset\. The agent then proceeds directly to the main iterative loop, executingStep 1,Step 2, andStep 3to continuously interact with the environment and fine\-tune the score function and world model\.

Table 4:MBDPO hyperparameters for the default 6M model\.We use the same hyperparameters for all online\-from\-scratch tasks\.HyperparameterValueSearch via diffusion samplingHorizon \(HH\)3Diffusion steps5\-20 stepsMonte Carlo sample number512Temperatureκ\\kappa0\.5Energy regularized factorη\\eta0\.1Replay bufferCapacity1,000,000SamplingUniformArchitecture \(6M\)Encoder dim256MLP dim512Latent state dim512Task embedding dim96Task embedding norm1ActivationLayerNorm \+ MishQ^\\hat\{Q\}\-function dropout rate1%Ensemble number ofQ^\\hat\{Q\}\-functions5OptimizationBatch size256 \(online\) / 1024 \(offline\)Policy prior loss norm\.Moving \(5%, 95%\) percentilesOptimizerAdamLearning rate3×10−43\\times 10^\{\-4\}Encoder learning rate1×10−41\\times 10^\{\-4\}Gradient clip norm20Table 5:Model configurations\.We list the specifications for each model configuration \(size\) of our multi\-task offline experiments\.Encoder dimis the dimensionality of fully connected layers in the encoderℰ\\mathcal\{E\},MLP dimis the dimensionality of layers in all other components,Latent state dimis the dimensionality of the latent representationzz,\# encoder layersis the number of layers in the encoderℰ\\mathcal\{E\},\# Q\-functionsis the number of learnedQQ\-functions, andTask embedding dimis the dimensionality ofeefrom Equation \(2\)\. TD\-targets are always computed by randomly subsampling twoQ^\\hat\{Q\}\-functions, regardless of the number ofQ^\\hat\{Q\}\-functions in the ensemble\. We did not experiment with other model configurations\. \*The default \(base\) configuration used in our single\-task RL experiments has 6M parameters\.1\.7M6M\*21M54M340MEncoder dim256256102417924096MLP dim384512102417924096Latent state dim1285127687681376\# encoder layers22345\#Q^\\hat\{Q\}\-functions25558Task embedding dim9696969696