Title: Analytic Energy-Guided Policy Optimization for Offline Reinforcement Learning
Venue: NeurIPS 2025
Accepted: yes
Keywords: offline reinforcement learning

decision making

generative models
Abstract: Conditional decision generation with diffusion models has shown powerful competitiveness in reinforcement learning (RL). Recent studies reveal the relation between energy-function-guidance diffusion models and constrained RL problems. The main challenge lies in estimating the intermediate energy, which is intractable due to the log-expectation formulation during the generation process. To address this issue, we propose the Analytic Energy-guided Policy Optimization (AEPO). Specifically, we first provide a theoretical analysis and the closed-form solution of the intermediate guidance when the diffusion model obeys the conditional Gaussian transformation. Then, we analyze the posterior Gaussian distribution in the log-expectation formulation and obtain the target estimation of the log-expectation under mild assumptions. Finally, we train an intermediate energy neural network to approach the target estimation of log-expectation formulation. We apply our method in 30+ offline RL tasks to demonstrate the effectiveness of our method. Extensive experiments illustrate that our method surpasses numerous representative baselines in D4RL offline reinforcement learning benchmarks.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces Analytic Energy-Guided Policy Optimization (AEPO), a novel method for offline reinforcement learning that leverages diffusion policy with analytically derived energy guidance. Inspired by the connection between guided diffusion sampling and constrained RL, the authors aim to improve action generation by analytically estimating the exact intermediate energy term, which is usually approximated or manually designed in prior work. By Taylor expansion of the Q value function and moment generating function of the Gaussian distribution, AEPO give an explicit estimation of the intermediate energy function during the diffusion sampling.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths:\n\n-\tThis paper derives a closed-form approximation of the intermediate energy used in guided diffusion policy sampling by Taylor expansion and moment-generating functions.\n-\tAEPO achieves better performance for many D4RL tasks compared with SOTA diffusion-based and non-diffusion-based offline RL methods.\n-\tThis paper is well written and easy to follow, and it also provides detail proof about the theoretical results and experimental study.\n\n### Weaknesses:\n\n-\tThe closed-form approximation of the intermediate energy in AEPO heavily relies on the Taylor expansion of the Q-value function. However, the validity of a Taylor expansion requires more than differentiability—it also assumes that the first- and second-order terms dominate, and that higher-order terms are negligible. The paper lacks sufficient discussion or analysis on whether these conditions hold in practice. If the Taylor approximation deviates significantly from the true Q-value, then the analytically derived intermediate energy remains just an estimate. In that case, it may suffer from similar biases as previous approaches that relied on inexact intermediate energy approximations.\n\n-\tWhile AEPO introduces a theoretically grounded closed-form approximation of the intermediate energy, its empirical advantage over other SOTA diffusion policy methods is relatively modest. Notably, Appendix B.3 shows that the proposed Guidance Rescaling technique for the diffusion score function yields significant performance gains. This suggests that the reported improvements may be largely attributed to this auxiliary component, rather than the core contribution of the closed-form intermediate energy itself. As a result, the experimental comparisons may not be entirely fair or isolating the effect of the main innovation. Furthermore, the paper lacks ablation studies on the impact of different Q-function learning strategies, which would be important for understanding the robustness and general applicability of AEPO.\n\n-\tThe paper omits several relevant and competitive baselines, particularly offline RL methods that incorporate uncertainty based pessimistic Q-value estimation, such as UWAC[1] and QDQ[2].\n\n[1]  Wu, Y., Zhai, S., Srivastava, N., Susskind, J., Zhang, J., Salakhutdinov, R., & Goh, H. (2021). Uncertainty weighted actor-critic for offline reinforcement learning. arXiv preprint arXiv:2105.08140.\n\n[2] Zhang, J., Fang, L., Shi, K., Wang, W., & Jing, B. (2024). Q-Distribution guided Q-learning for offline reinforcement learning: Uncertainty penalized Q-value via consistency model. Advances in Neural Information Processing Systems, 37, 54421-54462.\n\nQUESTIONS:\n1. The first two terms in Equation (15) assume that $\\bar{a}$ and $a_0$ are independent. However, Appendix J.2 shows that the value of $\\bar{a}$ depends on $a_0$. Does this affect the accuracy of Equation (15)?\n\n2.  AEPO seems sensitive to several hyperparameters, such as the inverse temperature, expectile weight, and guidance degree. Does this sensitivity affect its usefulness across different tasks? How should these hyperparameters be chosen based on the characteristics of a task?\n\n3. Does AEPO offer any computational efficiency improvements compared to other state-of-the-art methods?\n\nLIMITATIONS:\nConsider improving how the paper presents its motivation, especially the connection between guided diffusion sampling and constrained RL. Although Equation (1) and Equation (2) are similar, the paper should provide a clearer intuition for why the Q-value function can be used as an energy function.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an posterior estimation method for the exact energy guidance function if Diffusion RL. There are three key estimation employed in this paper. 1) Consider Q function as a linear function within a certain area. 2) Consider diffusion posterior \\mu(\\x_0| \\x_t) as a Gaussian distribution. With these two estimations, the method can effective compute exact energy function during training, and apply energy guidance in Offline RL tasks such as D4RL.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n1. The paper poses a neat and in-depth analysis of existing energy guidance method. The proposed method is well theoretical-grounded. The posterior estimation method contain good mathematical insight. \n2. The experiments are extensive and supportive. This work compares with a series of related baselines to prove the effectiveness of the proposed method.\n\nWeaknesses:\n1) I think there's still room for improvement on the paper writing. E.g., What is the core motivation for proposing AEPO, what problem of QGPO(CEP) is AEPO trying to solve? This one should be highlighted in Introduction and Method Section.  Also, what is the key difference between DPS and AEPO (I think they both some kind of posterior estimation). Finally I would suggest keeping only \\pi, \\mu notation or p, q notation in introduction. Keeping both is not really necessary.\n2) Regarding (1).  I think usually the advantage of Analytic xxx Algorithm is that you can compute some quantity so that you do not need to estimate them using a network. However, AEPO seems also need to train an energy guidance model (Eq. 25), which somehow weakens its advantage.\n3) Line 177:   \"Reminding that the posterior distribution µ0|t(a0|at, s) is also a Gaussian distribution\". This is wrong.  $µ0|t(a0|at, s)$ is not Gaussian, $µ0|t(am|at, s)$ can be considered as Gaussian only if m and t are close.\n\nQUESTIONS:\n1. In Eq 3, Do you mean s.t. KL(\\pi|\\mu) <=  \\eta?   It seems  <= xxx is missing.\n2. Why cannot you analytically compute the gradient of energy function from Eq. 17.\n3. Besides performance, what is the advantage of AEPO compared with previous work , especially QGPO and DPS.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 4, "DATE": "2025-06-30"}

comments: SUMMARY:
This paper proposes Analytic Energy-guided Policy Optimization (AEPO), a method designed to approximate the intermediate energy used in conditional decision generation with diffusion models. It first analyzes the limitations of existing methods that provide only inexact intermediate energy estimates, and then introduces a novel approximation technique based on Taylor expansions. Subsequently, the model approximates the Gaussian mean and covariance of the posterior distribution.
AEPO is evaluated on the D4RL offline reinforcement learning benchmark and shows promising performance compared to 30 baseline approaches.

STRENGTHS AND WEAKNESSES:
# Strengths
- The paper presents comprehensive mathematical derivations, including key theorems, proofs, and approximations. I reviewed most of the proofs and found them technically solid and convincing, demonstrating a high level of quality.
- The experimental section is thorough, with comparisons against more than 30 baseline approaches on the D4RL benchmark, highlighting the competitiveness of the proposed method.
- A reasonable ablation study and analysis are provided to support the understanding of individual components.

# Weakness
- No major weaknesses are identified at this stage, acknowledging the reviewer's limited expertise in energy-based methods and offline RL. However, I remain open to further discussion with other reviewers regarding potential limitations of the paper during the rebuttal phase.

QUESTIONS:
Except for D4RL, are there any other benchmarks that can be used to evaluate the current approach? Is the OG-Bench (ICLR25), a goal-conditioned offline RL benchmark, a reasonable choice for the current approach?

LIMITATIONS:
yes, a limitation section has been provided in the appendix.
TITLE: Official Review
CLARITY: 4
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-02

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors study the problem of offline RL with diffusion-model-based algorithms. It falls in the category of learning a diffusion model for the *behaviour policy*, using an energy function based on the Q-function to guide the generation. The paper identifies that the theoretically optimal guidance signal for this task, termed the *intermediate energy*, is defined by an intractable log-expectation formulation. Consequently, prior methods have relied on theoretically inexact heuristics or potentially unstable training procedures to approximate this guidance, creating a gap between theory and practice.\n\nThe proposed algorithm, Analytic Energy-guided Policy Optimization (AEPO), is novel with tractable analytic approximation for this intermediate energy. The authors achieve this by linearizing the Q-function with a first-order Taylor expansion and then leveraging the closed-form moment-generating function of a Gaussian distribution to resolve the expectation. Experiments have been conducted to validate the idea.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n* This paper studies a highly significant and fundamental problem in diffusion-based reinforcement learning: the intractability of the theoretically optimal guidance signal.\n* They propose a principled and effective solution. The core contribution—deriving an analytic approximation for the intermediate energy—is highly original and technically sound.\n* The use of a Taylor expansion combined with the moment-generating function of a Gaussian is a clever and novel approach that elegantly transforms an intractable problem into a tractable one. It is a distinct and well-reasoned improvement over previous heuristics (like MSE/DPS) and more complex training schemes (like contrastive learning), and its derivation is clear and well-supported.\n*  The paper is very clearly written and well-organized.\n\n### Weaknesses\n\nThe paper presents several areas for improvement.\n\nFirst, the first-order Taylor expansion of the Q-function in Equation (14) assumes local linearity. The authors should clarify whether the necessary conditions of Q-function smoothness and the locality of $\\bar{a}$ around $a_0$ are met.\n\nSecond, the method's reliance on Gaussian mechanics restricts its applicability to continuous action spaces. This significant limitation is not adequately discussed in the main paper. While it may perform well in continuous spaces, its effectiveness in discrete spaces remains unclear.\n\nThird, regarding experiments, while the results are generally strong, further analysis and experiments on computational efficiency are needed. Given the potentially long training times, more profiling would help identify areas for improvement.\n\nFinally, the evaluation lacks diversity, focusing on a specific experimental setting rather than including real-world robotic tasks or domains with different physical properties. This narrow scope limits the assessment of the method's broader generalization and robustness, thereby reducing the stated impact of the results.\n\nQUESTIONS:\n1. How is the locality and smoothness of the Q-function (Equation 14) validated, both theoretically and experimentally?\n2. Can the proposed AEPO be naturally extended to a discrete action space?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 4, "ORIGINALITY": 4, "DATE": "2025-07-08"}