Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Analytic Energy-Guided Policy Optimization for Offline Reinforcement Learning

Jifeng Hu, Sili Huang, Zhejian Yang, Shengchao Hu, Li Shen, Hechang Chen, Lichao Sun, Yi Chang, Dacheng Tao

TL;DR

Offline reinforcement learning with diffusion-based policy generation faces the challenge of computing the intermediate energy in the log-expectation term. AEPO derives a closed-form intermediate energy under conditional Gaussian diffusion, and uses Taylor expansion plus a Gaussian posterior to approximate the log-expectation term $\log\mathbb{E}_{\mu_{0|t}}[e^{\beta Q(s,a_0)}]$, enabling analytic guidance. It trains a Q-function with expectile regression and an intermediate energy network, and introduces guidance rescaling to stabilize inference. Across 30+ D4RL tasks, AEPO achieves competitive or state-of-the-art performance against dozens of baselines, demonstrating the practical effectiveness of analytic energy guidance for offline diffusion-based RL.

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.

Analytic Energy-Guided Policy Optimization for Offline Reinforcement Learning

TL;DR

Offline reinforcement learning with diffusion-based policy generation faces the challenge of computing the intermediate energy in the log-expectation term. AEPO derives a closed-form intermediate energy under conditional Gaussian diffusion, and uses Taylor expansion plus a Gaussian posterior to approximate the log-expectation term , enabling analytic guidance. It trains a Q-function with expectile regression and an intermediate energy network, and introduces guidance rescaling to stabilize inference. Across 30+ D4RL tasks, AEPO achieves competitive or state-of-the-art performance against dozens of baselines, demonstrating the practical effectiveness of analytic energy guidance for offline diffusion-based RL.

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.
Paper Structure (41 sections, 2 theorems, 87 equations, 6 figures, 4 tables, 1 algorithm)

This paper contains 41 sections, 2 theorems, 87 equations, 6 figures, 4 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Preliminary
  3. Diffusion Probabilistic Models
  4. Guided Sampling
  5. Diffusion Offline RL
  6. Method
  7. Intermediate Energy
  8. Posterior Approximation
  9. Q Function Training
  10. Guidance Rescaling
  11. Experiments
  12. Environments
  13. Metrics
  14. Baselines
  15. Results
  16. ...and 26 more sections

Key Result

Theorem 3.1

Suppose $p_0(x_0)$ and $q_0(x_0)$ has the relation of Equation solution of energy-function-guided diffusion models. $p_{t|0}(x_t|x_0)$ and $q_{t|0}(x_t|x_0)$ are defined by for all $t\in(0, T]$. According to the Law of Total Probability, the marginal distribution $p_t(x_t)$ and $q_t(x_t)$ are given by $p_t(x_t)=\int p_{t|0}(x_t|x_0) p_0(x_0)dx_0$ and $q_t(x_t)=\int q_{t|0}(x_t|x_0) q_0(x_0)dx_0$.

Figures (6)

  • Figure 1: Q function training and posterior approximation ablation of AEPO on D4RL Gym-MuJoCo tasks.
  • Figure 2: Guidance rescale ablation of AEPO on Gym-MuJoCo walker2d-medium-expert task. The x-axis denotes training steps.
  • Figure 3: Guidance rescale ablation of AEPO on D4RL Gym-MuJoCo halfcheetah-medium-expert task. The y-axis and x-axis denote the normalized score and training steps, respectively.
  • Figure 4: Guidance rescale ablation of AEPO on D4RL Gym-MuJoCo hopper-medium-expert task. The y-axis and x-axis denote the normalized score and training steps, respectively.
  • Figure 5: Guidance rescale ablation of AEPO on D4RL Gym-MuJoCo walker2d-medium task. The y-axis and x-axis denote the normalized score and training steps, respectively.
  • ...and 1 more figures

Theorems & Definitions (7)

  • Theorem 3.1: Inexact and Exact Intermediate Energy
  • proof
  • Theorem 4.1: Intermediate Energy Guidance
  • proof
  • proof
  • proof
  • proof