World Models via Policy-Guided Trajectory Diffusion

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, and interleave predicting the next state with sampling the next action from the policy. Thus, the 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 via 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 capabilities of our approach and demonstrate that it obtains competitive prediction errors to state-of-the-art autoregressive baselines. PolyGRAD also enables performant policies to be trained via on-policy RL in imagination. We believe that PolyGRAD introduces a promising paradigm for world modelling with many possible extensions to explore in future work.