Prior-Guided Diffusion Planning for Offline Reinforcement Learning

Dear Reviewer rn8g,

We appreciate your thoughtful and detailed review. We hope that the explanations provided below will address your questions and resolve any concerns..

W3. Theoretical justification hinges on DDIM being bijective which in parctice it may not hold (because it's an approximate ODE solver)

We acknowledge that exact bijectivity is not guaranteed in practice under standard DDIM sampling with finite discretization steps. This is a known limitation of DDIM inversion [1, 2], which we have also noted in the Limitations section. To fully ensure bijectivity, one could employ an ODE-based generative model or a flow-matching framework, both of which offer theoretical invertibility.

Nevertheless, our empirical results indicate that the bijective assumption holds approximately in practice and is sufficient for effective guidance. As shown in Figure 6, noise samples drawn from the learned prior consistently denoise into high‑value trajectories, indicating that the mapping from latent space to trajectory space is preserved well enough to enable strong performance. Consequently, as reported in Table 1, PG achieves state‑of‑the‑art results across diverse benchmarks, including long‑horizon planning tasks such as AntMaze, Kitchen, and Maze2D.

These results suggest that while exact bijectivity is not guaranteed, the approximate invertibility of DDIM is sufficient in practice for guiding diffusion models via learned priors.

W4. PG uses a relatively simple Gaussian prior which may not be sufficient in stochastic behavior setting.

We appreciate the reviewer’s insightful comment regarding the potential limitations of using a simple Gaussian prior in stochastic behavior settings. We agree that a more expressive prior model, such as a Gaussian mixture model (GMM), could be beneficial in scenarios where multiple distinct trajectories achieve similarly high returns, even under a deterministic optimal policy. To investigate this, we evaluated PG with a multi‑modal prior in our experiments (Table 2) to assess whether such an expressiveness would improve performance, and found that the improvement was not substantial. Nonetheless, PG can be easily extended to incorporate a more expressive prior, and our framework is fully compatible with such an extension if needed.

W2, Q1. Can you clarify which specific Kitchen tasks were used? Is the reported score averaged over several tasks (e.g., kitchen-mixed-v0, kitchen-partial-v0, etc.) or based on a single one?

We thank the reviewer for the helpful question and apologize for the lack of clarity. We clarify that the "Kitchen" domain refers to the FrankaKitchen environment from the D4RL benchmark suite [1]. In Table 1, we report performance on two standard Kitchen tasks from the D4RL benchmark suite: kitchen-mixed-v0 (denoted as Kitchen-M) and kitchen-partial-v0 (denoted as Kitchen-P). The “Kitchen” score, as reported in Figure 5, and Figure 7(b), refers to the average of the normalized scores across these two tasks. We will add detailed clarification on Section 4.1 in the final version.

W1, Q2. Why does PG underperform on standard MuJoCo tasks? These tasks are supposed to be easy and saturated within the RL domain. Do you attribute this to horizon length, simplicity of dynamics, or limitations of your prior formulation?

We appreciate the reviewer’s insightful comment on PG’s performance in standard MuJoCo tasks. PG is a diffusion planner-based method designed for long-horizon planning, where capturing temporally extended dependencies is essential. However, MuJoCo locomotion tasks involve short-horizon, dense-reward control focused on making agents run faster, which does not require such lookahead capabilities [2]. As a result, diffusion planner-based methods, including PG, often underperform compared to diffusion policy-based approaches in these settings, as shown in Table 1 of our paper. While the additional structure and trajectory-level modeling of planners benefit challenging long-horizon domains such as AntMaze, Kitchen, or Maze2D, they provide little advantage and can even be slightly detrimental in simpler environments where policies can already achieve near-optimal performance.

To examine whether the observed performance gap in MuJoCo tasks stems from the use of a diffusion planner rather than a limitation of the PG framework itself, we conduct an additional experiment by applying PG to a diffusion policy. Specifically, we replaced the standard Gaussian prior of a behavior-cloned diffusion policy with a learnable, behavior-regularized prior, following the same training procedure used in PG.

The above table shows that replacing the diffusion planner in PG with a diffusion policy (PG policy) consistently improved performance across MuJoCo tasks, raising the average score from 84.6 (PG planner) to 88.9 and surpassing the DQL baseline score of 87.9. This demonstrates that PG remains effective even in short-horizon environments, suggesting that the lower performance of PG in MuJoCo tasks is due to the nature of trajectory-level planning rather than a limitation of the method itself. The results for PG (policy) are averaged over 3 random seeds and reported as the mean $\pm$ standard error.

We will include these results to further clarify that the PG framework is effective across both policy- and planner-based approaches.

Q3. How is the inverse dynamics model trained and evaluated?

We apologize for the lack of explanation regarding the inverse dynamics model. In Prior Guidance (PG), the inverse dynamics model $\epsilon_w$ is trained following the setup established in Diffusion Veteran (DV), which PG builds upon. The model maps a pair of temporally spaced states $(\mathbf{s}^{(\tau)}, \mathbf{s}^{(\tau+m)})$ to the action $\mathbf{a}^{(\tau)}$ , where $\tau$ denotes the current environment step, and $m$ is the planning stride. Then $\epsilon_w$ is used to recover executable actions from denoised state trajectories $\mathbf{x}_0 = \left[ \mathbf{s}^{(\tau)}, \mathbf{s}^{(\tau + m)}, \dots, \mathbf{s}^{(\tau + (H-1)m)} \right]$ generated by the diffusion planner.

Training:

• The inverse model is trained via supervised regression from the dataset $\mathcal{D}$. The objective is to minimize the following loss:

  $$\mathbb{E}_{\mathbf{s}^{(\tau)}, \mathbf{s}^{(\tau + m)}, \mathbf{a}^{(\tau)},\epsilon,t} \left[\left\Vert \epsilon - \epsilon_w \left(\sqrt{\bar{\alpha}_t} \mathbf{a}^{(\tau)} + \sqrt{1 - \bar{\alpha}_t}\epsilon, t, s^{(\tau)},s^{(\tau+m)}\right) \right\Vert^2\right]$$

Inference:

• At test time, after PG generates a full predicted state trajectory $\mathbf{x}_0 = \left[ \mathbf{s}^{(\tau)}, \mathbf{s}^{(\tau + m)}, \dots, \mathbf{s}^{(\tau + (H-1)m)} \right]$, and the action $\mathbf{a}^{(\tau)}$ is computed by applying the inverse dynamics model $\epsilon_w$ to the first two states $\left(\mathbf{s}^{(\tau)}, \mathbf{s}^{(\tau + m)}\right)$
• This action is then executed in the environment, and the process repeats at the next timestep.

We will include this detailed description of the inverse dynamics model in the next revision.

References

[1] Fu, Justin, et al. "D4rl: Datasets for deep data-driven reinforcement learning." arXiv preprint arXiv:2004.07219 (2020).

[2] Lu, Haofei, et al. "What makes a good diffusion planner for decision making?." arXiv preprint arXiv:2503.00535 (2025).

Dear Reviewer 6FdX,

Thank you for your valuable comments. We hope that the following responses will answer your questions and address the concerns you have raised.

W1. Additional Complexity: comparing to MCSS, the proposed method requires training and tuning of a latent-space value function and a learnable prior, adding architectural and training complexity.

We acknowledge that Prior Guidance (PG) introduces additional networks, which we have also noted in the Limitations section, and that this indeed increase architectural and training complexity relative to MCSS. However, we highlight that this added complexity is both minimal and reasonable given the substantial gains in performance and efficiency.

1. Lower Training Cost

   The introduced components are lightweight and modular. The prior is parameterized by a simple GRU and the latent value function is a shallow MLP, both adding minor overhead compared to the diffusion planner backbone. Moreover, these networks are trained separately from the planner without backpropagation through the diffusion model, keeping the overall computational cost manageable (see Figure 4, left).

2. Avoidance of Costly Sampling

   MCSS requires sampling and scoring a large number of candidate trajectories at inference time, which is both computationally expensive and difficult to scale, especially in long-horizon or high-dimensional settings. PG eliminates this bottleneck by generating a single trajectory, yielding up to a 500× reduction in inference cost (see Figure 4, right).

3. Performance Validates Overhead

   Despite the small additional training components, PG consistently outperforms MCSS across most long-horizon benchmarks (Table 1, Figure 5), demonstrating that the modest increase in training complexity is outweighed by substantial gains in both efficiency and performance.

W2. in order to get a tractable regularization term, the author assumes bijectivity of the denoising scheme, which is not really satisfied.

We acknowledge that exact bijectivity is not guaranteed in practice under standard DDIM sampling with finite discretization steps. This is a known limitation of DDIM inversion [1, 2], which we have also noted in the Limitations section. To fully ensure bijectivity, one could employ an ODE-based generative model or a flow-matching framework, both of which offer theoretical invertibility.

Nevertheless, our empirical results indicate that the bijective assumption holds approximately in practice and is sufficient for effective guidance. As shown in Figure 6, noise samples drawn from the learned prior consistently denoise into high‑value trajectories, indicating that the mapping from latent space to trajectory space is preserved well enough to enable strong performance. Consequently, as reported in Table 1, PG achieves state‑of‑the‑art results across diverse benchmarks, including long‑horizon planning tasks such as AntMaze, Kitchen, and Maze2D.

These results suggest that while exact bijectivity is not guaranteed, the approximate invertibility of DDIM is sufficient in practice for guiding diffusion models via learned priors.

W3. Empirically speaking, while the proposed seem to work well, the performance boost is not too important. I think the performance is quite saturated in D4RL Benchmark. it would be important in include more challenging domains (e.g long-horizon tasks in ogbench, humanoid-large-maze ?)

We thank the reviewer for this valuable suggestion. We agree that D4RL has been extensively studied and that evaluating on more challenging domains, such as long‑horizon tasks in OGBench or humanoid‑large‑maze, would further validate the generality of our approach. However, given the limited rebuttal period, conducting such large‑scale experiments is unfortunately infeasible. We believe that PG clearly outperforms other baselines in long-horizon planning tasks within D4RL, while also reducing inference time substantially compared to the previous state-of-the-art method DV, constitutes a meaningful contribution. We thank the reviewer again for this suggestion and plan to extend our evaluation to more challenging environments in future work.

References

[1] Song, Jiaming, Chenlin Meng, and Stefano Ermon. "Denoising diffusion implicit models." arXiv preprint arXiv:2010.02502 (2020).

[2] Dhariwal, Prafulla, and Alexander Nichol. "Diffusion models beat gans on image synthesis." Advances in neural information processing systems 34 (2021): 8780-8794.

Thanks for the rebuttal. I maintain my score. My main concern is still the limited evaluation to D4RL where we see very moderate boost of performance.

Dear Reviewer J4nL,

We appreciate your valuable comments. We hope that our responses below will clarify the points raised and resolve any concerns.

W2. It is unclear how much of the performance gain comes from the incorporation of Diffusion Veteran (DV). The authors are encouraged to include an ablation study using the standard Diffuser as the base model to better isolate the contribution of PG.

We thank the reviewer for highlighting the importance of isolating the contribution of PG from that of Diffusion Veteran (DV). Due to the limited time available during the rebuttal period, we were unable to conduct an ablation using the standard Diffuser as the base model. Instead, we performed a related experiment by applying PG to a diffusion policy that uses the same diffusion model architecture as DQL, allowing us to assess whether PG provides meaningful performance gains even without incorporating DV. As detailed in our response to Q1, applying PG by substituting the standard Gaussian prior of a behavior-cloned diffusion policy with a learned, behavior-regularized prior achieved an average score of 88.9, outperforming other diffusion-based policies on MuJoCo tasks. These results demonstrate that PG provides a meaningful performance gain even without DV, indicating that its benefits are not tied to a specific base model.

W1, Q1. Why does the proposed PG method not perform well on some datasets? For example, in Table 1, PG does not outperform some diffusion-based methods such as DQL in MuJoCo tasks. Does the method only work well on long-horizon tasks?

We appreciate the reviewer’s insightful comment on PG’s performance in MuJoCo tasks. PG is a diffusion planner-based method designed for long-horizon planning, where capturing temporally extended dependencies is essential. However, MuJoCo locomotion tasks involve short-horizon, dense-reward control focused on making agents run faster, which does not require such lookahead capabilities [1]. As a result, diffusion planner-based methods, including PG, often underperform compared to diffusion policy-based approaches in these settings, as shown in Table 1. While the additional structure and trajectory-level modeling of planners benefit challenging long-horizon domains such as AntMaze, Kitchen, or Maze2D, they provide little advantage and can even be slightly detrimental in simpler environments where policies can already achieve near-optimal performance.

To examine whether the observed performance gap in MuJoCo tasks stems from the use of a diffusion planner rather than a limitation of the PG framework itself, we conduct an additional experiment by applying PG to a diffusion policy. Specifically, we replaced the standard Gaussian prior of a behavior-cloned diffusion policy with a learnable, behavior-regularized prior, following the same training procedure used in PG. The above table shows that replacing the diffusion planner in PG with a diffusion policy (PG policy) consistently improved performance across MuJoCo tasks, raising the average score from 84.6 (PG planner) to 88.9 and surpassing the DQL baseline score of 87.9. This demonstrates that PG remains effective even in short-horizon environments, suggesting that the lower performance of PG in MuJoCo tasks is due to the nature of trajectory-level planning rather than a limitation of the method itself. The results for PG (policy) are averaged over 3 random seeds and reported as the mean $\pm$ standard error.

We will include these results to further clarify that the PG framework is effective across both policy- and planner-based approaches.

Q2. How does the method achieve good performance on long-horizon tasks, consistently outperforming MCSS? It seems that prior guidance alone does not provide this capability. Does the improvement come from Diffusion Veteran (DV)?

We apologize for the lack of clarity and appreciate the reviewer for raising this important point. We would like to clarify that in Figure 5, the method labeled as MCSS with 50 samples corresponds directly to Diffusion Veteran (DV). That is, DV is implemented by sampling 50 trajectories from a behavior-cloned diffusion planner and selecting the highest-value one using a learned value function. Importantly, all methods in Figure 5, including MCSS (with different sample counts) and PG, share the same diffusion model architecture and hyperparameters, all trained using the DV procedure. The performance differences arise solely from the inference-time guidance algorithms.

While DV (i.e., MCSS-50) achieves strong performance, MCSS inherently couples the number of samples with the degree of behavior regularization. Increasing the number of samples improves the chance of finding higher-value trajectories, but it also raises the risk of selecting out-of-distribution ones, forcing a trade-off between trajectory quality and regularization strength. Because of this coupling, there is a practical ceiling on how precisely MCSS can optimize for the best high-value trajectory under a given level of regularization. Beyond a certain point, increasing samples no longer yields better results without sacrificing regularization.

In contrast, PG completely decouples behavior regularization from the search for high-value trajectories. The $f$-divergence regularization is applied explicitly and independently, allowing the latent space prior to be optimized solely for maximizing trajectory value within the regularization constraint. This enables PG to accurately reach the behavior-regularized optimum, without approximation error from finite sampling, and without the computational cost or OOD risk that MCSS faces at large sample counts.

We will revise the paper to clearly indicate that DV corresponds to MCSS with 50 samples, and that all methods share the same diffusion model. The observed improvements thus stem from our proposed prior-guided inference, rather than differences in model capacity or training setup.

Q3. It is unclear whether the method modifies the training of the base Diffusion Veteran. Is the proposed prior guidance applied only at the sampling stage?

We apologize for the confusion and we appreciate the reviewer’s careful reading. To clarify, PG does not modify the training of the underlying diffusion model used in Diffusion Veteran (DV). The diffusion planner is trained exactly as in DV, using standard behavior cloning on offline trajectories. PG is applied only during inference, where it replaces the MCSS procedure used in DV. Specifically:

• During training we introduce an additional learnable prior distribution over the noise variable $\mathbf{x}_T$ that is optimized to maximize the expected value under the value function subject to an $f$‑divergence constraint maintaining closeness to the standard Gaussian.
• This prior is trained after the diffusion model has been fixed through behavior cloning. At the sampling stage, the standard Gaussian prior is replaced with the learned prior distribution, after which the behavior‑cloned diffusion model is used for denoising.
• The diffusion model itself is not fine-tuned or altered in any way.

In summary, PG is a post‑hoc, inference‑time guidance method applied to a fixed diffusion model. All performance gains come from learned prior distribution rather than from altering the base diffusion model’s training. We will include this clarification in the next revision of the paper.

References

[1] Lu, Haofei, et al. "What makes a good diffusion planner for decision making?." arXiv preprint arXiv:2503.00535 (2025).

Dear Reviewer tREs,

We are grateful for your thoughtful feedback. We hope the following responses address your questions and help alleviate your concerns.

W1. The D4RL results are not a significant improvement compared to the strongest baselines (DV$^*$, DQL).

We appreciate the reviewer for their careful reading and insightful comment regarding PG’s performance. In Table 1 of the paper, PG shows substantial performance improvements over DQL and DV* in long-horizon tasks such as Kitchen, AntMaze, and Maze2D. However, in MuJoCo locomotion tasks such as Walker2d, Hopper, and HalfCheetah, the gains over DV* are modest, and the overall performance is lower than DQL. We attribute this not to a limitation of PG itself, but to the use of a diffusion planner, which is less effective in short-horizon, dense-reward settings where trajectory-level modeling provides little benefit and can even be slightly detrimental in simpler environments where policies can already achieve near-optimal performance.

To verify this, we conducted an additional experiment by applying PG to a diffusion policy instead of a planner. Specifically, we replaced the standard Gaussian prior of a behavior-cloned diffusion policy with a learnable, behavior-regularized prior, following the same training procedure used in PG. The above table shows that replacing the planner with a policy (PG policy) consistently improved performance across MuJoCo tasks, increasing the average score from 84.6 (PG planner) to 88.9, surpassing not only DV* (83.0) but also the DQL (87.9). This result demonstrates that PG remains effective even in short-horizon environments and the lower performance observed in MuJoCo tasks is due to the planner architecture rather than the PG framework itself. The results for PG (policy) are averaged over 3 random seeds and reported as the mean $\pm$ standard error.

We will include these results to further clarify that the PG framework is effective across both policy- and planner-based approaches.

Q1. Why is the DDIM denoising process bijective in the limit of small discretization steps? Is there a prior work that discussed this?

We thank the reviewer for requesting this clarification. The DDIM denoising process [1] becomes bijective in the limit of small discretization steps because it corresponds to solving a particular ordinary differential equation (ODE) using Euler integration. In Section 4.3 of the DDIM paper, the authors show that DDIM sampling is equivalent to numerically solving an ODE of the form.

where $\bar{\mathbf{x}}(t)=\frac{\mathbf{x}(t)}{\sqrt{\bar{\alpha}_t}}$ and $\sigma(t) = \sqrt{\frac{1-\bar{\alpha}_t}{\bar{\alpha}_t}}$. This ODE describes a deterministic transformation from noise $\bar{\mathbf{x}}(T)$ to data $\bar{\mathbf{x}}(0)$, and vice versa. Therefore, with sufficiently fine discretization, the DDIM process approximates the ODE solution arbitrarily well. Furthermore, Section 4.3 and Appendix B of [1] explicitly note that DDIM corresponds to the probability flow ODE for the variance-exploding SDE from Song et al. (2020) [2].

This equivalence explains why the DDIM sampling path becomes bijective when using small discretization steps. Several works leveraged this property for DDIM inversion, the process of mapping real data samples back into the latent noise space [1, 3, 4].

Q2. Equation 2 does not explain what the variables are, e.g $V$ and $f$. It would help to explain them immediately after the equation.

We thank the reviewer for pointing out the lack of variable definitions in Equation (2). We agree that the meanings of $V$ and $f$ should be clarified immediately following the equation to improve readability, as the original explanation in the Preliminary section is far from Equation (2) and may disrupt the reader’s flow. To address this, we will revise the main paper to include explicit definitions right after Equation (2). Specifically, we will add the following sentence in the next revision of the paper:

Here, $V(\mathbf{x}_0)$ denotes the estimated value of the denoised trajectory $\mathbf{x}_0$ and $f(\cdot)$ is a regularization function (e.g., for $f$-divergence) that penalizes deviation from the behavior distribution.

Q3. There is a recent related work on steering flow policies by learning the prior distribution [1]. They also train a noise-space value function to alias the sample-space value function. Can you discuss PG in relation to this work?

We appreciate the reviewer for bringing up this connection. A referenced concurrent work, "Steering Your Diffusion Policy with Latent Space Reinforcement Learning" (DSRL) [5], proposes steering a behavior-cloned diffusion policy by optimizing a learnable prior via reinforcement learning. While the high-level goal of guiding generative models through learned priors is shared with our method, there are both important commonalities and key differences:

Commonalities: Training Latent Prior and Value function

Both DSRL and PG aim to improve decision quality by learning the prior over the noise space ($\mathbf{x}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$) that feeds into the diffusion model’s denoising process, allowing for the sampling of high-value actions or trajectories. In both methods, the prior is trained to maximize the value function defined over $\mathbf{x}_T$, which estimates the value of the corresponding denoised sample $\mathbf{x}_0$.

Key Differences

1. Diffusion Planner and Diffusion Policy

   Unlike DSRL, which employs a diffusion policy that directly maps noise to action, PG utilizes a diffusion planner that outputs entire state trajectories and uses an inverse dynamics model to reconstruct executable actions. This makes PG particularly suitable for long-horizon and sparse-reward settings, where trajectory-level modeling is essential.

2. Theoretical Insight

   In our work, we provide a formal result (Proposition 1) demonstrating that, with a sufficiently large number of discretization steps, the objective for optimizing the diffusion model can be equivalently reformulated as an optimization over the latent prior distribution. This equivalence is enabled by the bijective nature of the DDIM sampling process, which allows for an exact mapping between the data space and the latent prior. This reformulation is not explored in DSRL.

3. Handling OOD Samples in Offline RL

   DSRL assumes that the learned diffusion model generates only in‑distribution samples and therefore does not incorporate explicit behavior constraints. However, as shown in Figures 3 and 5, diffusion models trained on offline datasets can still produce out‑of‑distribution samples that degrade performance. To mitigate this issue, PG introduces behavior regularization by imposing an $f$‑divergence penalty between the learned prior and the standard normal distribution. This explicitly penalizes deviation from the behavior distribution and induces regularization at the prior level.

We will revise the Related Works section to clearly position our contributions relative to DSRL and highlight both the shared foundations and the novel aspects of our approach.

References

[1] Song, Jiaming, Chenlin Meng, and Stefano Ermon. "Denoising diffusion implicit models." arXiv preprint arXiv:2010.02502 (2020).

[2] Song, Yang, et al. "Score-based generative modeling through stochastic differential equations." arXiv preprint arXiv:2011.13456 (2020).

[3] Dhariwal, Prafulla, and Alexander Nichol. "Diffusion models beat gans on image synthesis." Advances in neural information processing systems 34 (2021): 8780-8794.

[4] Zeng, Yan, Masanori Suganuma, and Takayuki Okatani. "Inverting the generation process of denoising diffusion implicit models: Empirical evaluation and a novel method." 2025 IEEE/CVF Winter Conference on Applications of Computer Vision (WACV). IEEE, 2025.

[5] Wagenmaker, Andrew, et al. "Steering Your Diffusion Policy with Latent Space Reinforcement Learning." arXiv preprint arXiv:2506.15799 (2025).