# World Models via Policy-Guided Trajectory Diffusion

arXiv: https://arxiv.org/abs/2312.08533 · HTML: https://mlsanity.com/paper/2312.08533/

## Key facts

- Cited by 34 papers in the map (data through 14 Oct 2026).
- Placed in Latent dynamics models and model-based RL, one of 11 directions in World models.
- Builds on 51 of its 61 references that are in the map.
- First posted on arXiv 13 Dec 2023.

## Placed in

- World models / [Latent dynamics models and model-based RL](https://mlsanity.com/wm/latent-dynamics-models-and-model-based-rl/)
- Diffusion models / [Diffusion and flow policies for robotics](https://mlsanity.com/diffusion/diffusion-and-flow-policies-for-robotics/)

## Builds on (most cited first, 30 of 51)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Score-Based Generative Modeling through Stochastic Differential Equations](https://mlsanity.com/paper/2011.13456/) (2020)
- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Deep Unsupervised Learning using Nonequilibrium Thermodynamics](https://mlsanity.com/paper/1503.03585/) (2015)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Generative Modeling by Estimating Gradients of the Data Distribution](https://mlsanity.com/paper/1907.05600/) (2019)
- [Improved Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2102.09672/) (2021)
- [Diffusion Policy: Visuomotor Policy Learning via Action Diffusion](https://mlsanity.com/paper/2303.04137/) (2023)
- [Trust Region Policy Optimization](https://mlsanity.com/paper/1502.05477/) (2015)
- [Asynchronous Methods for Deep Reinforcement Learning](https://mlsanity.com/paper/1602.01783/) (2016)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [RePaint: Inpainting using Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2201.09865/) (2022)
- [Mastering Atari, Go, Chess and Shogi by Planning with a Learned Model](https://mlsanity.com/paper/1911.08265/) (2019)
- [Learning Latent Dynamics for Planning from Pixels](https://mlsanity.com/paper/1811.04551/) (2018)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Planning with Diffusion for Flexible Behavior Synthesis](https://mlsanity.com/paper/2205.09991/) (2022)
- [Mastering Atari with Discrete World Models](https://mlsanity.com/paper/2010.02193/) (2020)
- [Deep Reinforcement Learning at the Edge of the Statistical Precipice](https://mlsanity.com/paper/2108.13264/) (2021)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Deep Reinforcement Learning in a Handful of Trials using Probabilistic Dynamics Models](https://mlsanity.com/paper/1805.12114/) (2018)
- [When to Trust Your Model: Model-Based Policy Optimization](https://mlsanity.com/paper/1906.08253/) (2019)
- [MOPO: Model-based Offline Policy Optimization](https://mlsanity.com/paper/2005.13239/) (2020)
- [Diffusion Policies as an Expressive Policy Class for Offline Reinforcement Learning](https://mlsanity.com/paper/2208.06193/) (2022)
- [Is Conditional Generative Modeling all you need for Decision-Making?](https://mlsanity.com/paper/2211.15657/) (2022)
- [Model-Based Reinforcement Learning for Atari](https://mlsanity.com/paper/1903.00374/) (2019)
- [Learning Interactive Real-World Simulators](https://mlsanity.com/paper/2310.06114/) (2023)

## Built on by (most cited first, 30 of 34)

- [Diffusion Policy Policy Optimization](https://mlsanity.com/paper/2409.00588/) (2024)
- [Efficient Online Reinforcement Learning for Diffusion Policy](https://mlsanity.com/paper/2502.00361/) (2025)
- [AVID: Adapting Video Diffusion Models to World Models](https://mlsanity.com/paper/2410.12822/) (2024)
- [Diffusion Model Predictive Control](https://mlsanity.com/paper/2410.05364/) (2024)
- [The Essential Role of Causality in Foundation World Models for Embodied AI](https://mlsanity.com/paper/2402.06665/) (2024)
- [Diffusion World Model: Future Modeling Beyond Step-by-Step Rollout for Offline Reinforcement Learning](https://mlsanity.com/paper/2402.03570/) (2024)
- [Energy-Guided Diffusion Sampling for Offline-to-Online Reinforcement Learning](https://mlsanity.com/paper/2407.12448/) (2024)
- [Diffusion Spectral Representation for Reinforcement Learning](https://mlsanity.com/paper/2406.16121/) (2024)
- [Prioritized Generative Replay](https://mlsanity.com/paper/2410.18082/) (2024)
- [Generalizing Consistency Policy to Visual RL with Prioritized Proximal Experience Regularization](https://mlsanity.com/paper/2410.00051/) (2024)
- [Implicit Search via Discrete Diffusion: A Study on Chess](https://mlsanity.com/paper/2502.19805/) (2025)
- [DMWM: Dual-Mind World Model with Long-Term Imagination](https://mlsanity.com/paper/2502.07591/) (2025)
- [Multi-agent Coordination via Flow Matching](https://mlsanity.com/paper/2511.05005/) (2025)
- [Revisiting Multi-Agent World Modeling from a Diffusion-Inspired Perspective](https://mlsanity.com/paper/2505.20922/) (2025)
- [Fine-tuning Diffusion Policies with Backpropagation Through Diffusion Timesteps](https://mlsanity.com/paper/2505.10482/) (2025)
- [World-Action Models for Robot Learning and Control: A Survey](https://mlsanity.com/paper/2609.16074/) (2026)
- [Adversarial Diffusion for Robust Reinforcement Learning](https://mlsanity.com/paper/2509.23846/) (2025)
- [Multimodal Diffusion Forcing for Forceful Manipulation](https://mlsanity.com/paper/2511.04812/) (2025)
- [Edge General Intelligence Through World Models and Agentic AI: Fundamentals, Solutions, and Challenges](https://mlsanity.com/paper/2508.09561/) (2025)
- [Goal-Driven Reward by Video Diffusion Models for Reinforcement Learning](https://mlsanity.com/paper/2512.00961/) (2025)
- [Maximize Your Diffusion: A Study into Reward Maximization and Alignment for Diffusion-based Control](https://mlsanity.com/paper/2502.12198/) (2025)
- [Extendable Planning via Multiscale Diffusion](https://mlsanity.com/paper/2503.20102/) (2025)
- [STITCH-OPE: Trajectory Stitching with Guided Diffusion for Off-Policy Evaluation](https://mlsanity.com/paper/2505.20781/) (2025)
- [Consistency Trajectory Planning: High-Quality and Efficient Trajectory Optimization for Offline Model-Based Reinforcement Learning](https://mlsanity.com/paper/2507.09534/) (2025)
- [DAWM: Diffusion Action World Models for Offline Reinforcement Learning via Action-Inferred Transitions](https://mlsanity.com/paper/2509.19538/) (2025)
- [A Review of Online Diffusion Policy RL Algorithms for Scalable Robotic Control](https://mlsanity.com/paper/2601.06133/) (2026)
- [Answerability Fields: Answerable Location Estimation via Diffusion Models](https://mlsanity.com/paper/2407.18497/) (2024)
- [Horizon Imagination: Efficient On-Policy Rollout in Diffusion World Models](https://mlsanity.com/paper/2602.08032/) (2026)
- [Inference Time Policy Optimization for Offline RL with Differentiable World Models](https://mlsanity.com/paper/2603.22430/) (2026)
- [Advantage-Guided Diffusion for Model-Based Reinforcement Learning](https://mlsanity.com/paper/2604.09035/) (2026)

## Abstract (arXiv)

World models are a powerful tool for developing intelligent agents. By predicting the outcome of a sequence of actions, world models enable policies to be optimised via on-policy reinforcement learning (RL) using synthetic data, i.e. in "in imagination". Existing world models are autoregressive in that they interleave predicting the next state with sampling the next action from the policy. Prediction error inevitably compounds as the trajectory length grows. In this work, we propose a novel world modelling approach that is not autoregressive and generates entire on-policy trajectories in a single pass through a diffusion model. Our approach, Policy-Guided Trajectory Diffusion (PolyGRAD), leverages a denoising model in addition to the gradient of the action distribution of the policy to diffuse a trajectory of initially random states and actions into an on-policy synthetic trajectory. We analyse the connections between PolyGRAD, score-based generative models, and classifier-guided diffusion models. Our results demonstrate that PolyGRAD outperforms state-of-the-art baselines in terms of trajectory prediction error for short trajectories, with the exception of autoregressive diffusion. For short trajectories, PolyGRAD obtains similar errors to autoregressive diffusion, but with lower computational requirements. For long trajectories, PolyGRAD obtains comparable performance to baselines. Our experiments demonstrate that PolyGRAD enables performant policies to be trained via on-policy RL in imagination for MuJoCo continuous control domains. Thus, PolyGRAD introduces a new paradigm for accurate on-policy world modelling without autoregressive sampling.