Prior-Guided Diffusion Planning for Offline Reinforcement Learning

TL;DR

This work tackles distributional shift in offline RL by improving diffusion-planner planning via Prior Guidance (PG). PG replaces the standard Gaussian prior with a learnable prior $p_ψ(\mathbf{x}_T|\mathbf{s})$ and uses a latent value function to steer diffusion trajectories toward high-value outcomes without backpropagating through the denoising process. By formulating behavior-regularized planning in the latent-prior space and alternating training between the prior and the latent value function, PG achieves efficient, single-trajectory planning with tractable regularization. Empirically, PG delivers state-of-the-art performance on long-horizon offline RL benchmarks (D4RL), while offering substantial reductions in inference cost compared with MCSS and robust ablations that validate the design choices. Overall, PG advances diffusion-based offline RL by enabling expressive priors and practical, scalable planning for complex tasks.

Abstract

Diffusion models have recently gained prominence in offline reinforcement learning due to their ability to effectively learn high-performing, generalizable policies from static datasets. Diffusion-based planners facilitate long-horizon decision-making by generating high-quality trajectories through iterative denoising, guided by return-maximizing objectives. However, existing guided sampling strategies such as Classifier Guidance, Classifier-Free Guidance, and Monte Carlo Sample Selection either produce suboptimal multi-modal actions, struggle with distributional drift, or incur prohibitive inference-time costs. To address these challenges, we propose Prior Guidance (PG), a novel guided sampling framework that replaces the standard Gaussian prior of a behavior-cloned diffusion model with a learnable distribution, optimized via a behavior-regularized objective. PG directly generates high-value trajectories without costly reward optimization of the diffusion model itself, and eliminates the need to sample multiple candidates at inference for sample selection. We present an efficient training strategy that applies behavior regularization in latent space, and empirically demonstrate that PG outperforms state-of-the-art diffusion policies and planners across diverse long-horizon offline RL benchmarks.Our code is available at https://github.com/ku-dmlab/PG.

Figures (9)

• Figure 1: Comparison between the behavior-cloned diffusion model and Prior Guidance. (Left) The behavior diffusion model samples latent noise $\mathbf{x}_T$ from a standard Gaussian and generates trajectories $\mathbf{x}_0$ via a learned denoising process $\tilde{\pi}_{\beta}(\mathbf{x}_0|\mathbf{x}_T, \mathbf{s})$, approximating the dataset distribution. (Right) Prior Guidance instead samples $\mathbf{x}_T$ from a learned prior $p_\psi(\mathbf{x}_T | \mathbf{s})$ concentrated in high-value regions, leading to improved trajectory generation through the same denoising process.
• Figure 2: Visualization of trajectory generation using three different guided sampling methods on maze2d-large-v1. Each method produces 20 planned trajectories. For MCSS, each trajectory is selected from a set of 10 candidates ($N=10$); the unselected candidates are visualized in gray.
• Figure 3: Predicted state sequences in maze2d-large-v1 based on MCSS varying the number of candidate trajectories $N$. When $N$ is small ($N=1$), MCSS often fails to sample trajectories with high value estimates, as the likelihood of capturing an optimal trajectory within a limited set is low. Conversely, when $N$ is excessively large ($N=500$ or $5000$), the method tends to select OOD trajectories with inflated value estimates.
• Figure 4: Training and inference time comparison. (Left) Training time until convergence for DQL (planner), PG with and without backpropagation through the denoising process across different environments. DQL (planner) denotes a baseline where the diffusion model is directly trained via Eq. \ref{['eq:behavior_regularized_planner']}; due to the intractability of computing the model's density, Eq. \ref{['eq:diffusion_loss']} is used as a surrogate. (Right) Inference time for MCSS and PG. The inference time is measured on the antmaze-medium-play-v2, with MCSS evaluated under different numbers of sampled trajectories.
• Figure 5: Average normalized score comparison between MCSS with varying numbers of samples $N$ and PG across four D4RL tasks. PG consistently matches or outperforms MCSS while requiring significantly fewer samples. Full results are provided in the Appendix \ref{['appendix:dv']}.

Theorems & Definitions (3)

• Proposition 1
• Proposition 1
• proof