PaperlayerView exact revision on arXiv

Source-linked AI summary

Policy Representation via Diffusion Probability Model for Reinforcement Learning

Long Yang, Zhixiong Huang, Fenghao Lei, Yucun Zhong, Yiming Yang, Cong Fang, Shiting Wen, Binbin Zhou, Zhouchen Lin

arXiv:2305.13122v1cs.LG

TL;DR

Unimodal policy representations can be limited in expressiveness for complicated policies, while prior diffusion-RL work considered empirical or offline settings. This paper represents diffusion policies as stochastic processes, analyzes their convergence, and proposes DIPO for model-free online RL; DIPO achieves the best score across the five evaluated MuJoCo tasks.

  • Problem

    The paper addresses how to characterize and analyze diffusion policies and how to design one for online model-free RL, where prior diffusion-RL work was empirical or focused on offline settings.

    Source: §1 · p. 4

  • Method

    The paper models diffusion policy as forward and reverse stochastic processes and proposes DIPO, which improves policies by updating dataset actions along the state-action value gradient.

    Source: §1.1 · p. 4, Convergence Analysis of Diffusion Policy · p. 5

  • Results

    DIPO achieves the best score across the five evaluated MuJoCo tasks and learns faster than SAC, PPO, and TD3 on Ant-v3.

    Source: §7.1 · p. 14

  • Takeaways & Limitations

    The paper provides a convergence guarantee for diffusion policy and presents DIPO as an implementation for online model-free RL.

    Source: §1.1 · p. 5, §5 · p. 11

    The convergence analysis assumes Lipschitz policy and score estimator conditions and a policy distribution satisfying a ν-Log-Sobolev inequality.

    Source: §4.3 · p. 10

Abstract

from arXiv · show

Popular reinforcement learning (RL) algorithms tend to produce a unimodal policy distribution, which weakens the expressiveness of complicated policy and decays the ability of exploration. The diffusion probability model is powerful to learn complicated multimodal distributions, which has shown promising and potential applications to RL. In this paper, we formally build a theoretical foundation of policy representation via the diffusion probability model and provide practical implementations of diffusion policy for online model-free RL. Concretely, we character diffusion policy as a stochastic process, which is a new approach to representing a policy. Then we present a convergence guarantee for diffusion policy, which provides a theory to understand the multimodality of diffusion policy. Furthermore, we propose the DIPO which is an implementation for model-free online RL with DIffusion POlicy. To the best of our knowledge, DIPO is the first algorithm to solve model-free online RL problems with the diffusion model. Finally, extensive empirical results show the effectiveness and superiority of DIPO on the standard continuous control Mujoco benchmark.

1 Introduction

The paper addresses how to represent diffusion policies and use them for online model-free reinforcement learning. It frames policies as stochastic processes, analyzes their convergence, and introduces DIPO, which improves policies through action gradients.

1 Introduction

  • The introduction identifies limited expressiveness of unimodal policies and prior diffusion-RL work focused on empirical or offline settings as the motivation for the paper.The authors describe DIPO as, to their knowledge, the first diffusion approach for online model-free RL.Source: §1 · p. 4

1.1 Our Main Work

  • Diffusion policy represents a policy through forward and reverse SDE-induced stochastic processes, rather than through a value function or parametric function.The forward process disturbs the input policy toward noise; the reverse process infers the policy.Source: §1.1 · pp. 4–5
  • Theorem 4.3 gives a convergence guarantee under mild conditions and relates sufficiently accurate score estimation to inferring actions from the data-generating policy.The paper connects this result to diffusion policy's ability to represent multimodal distributions.Source: §1.1 · p. 5
  • DIPO adapts diffusion policy to model-free online RL by updating dataset actions along the state-action value gradient, which the paper calls the action gradient.The authors state that ordinary policy improvement cannot be naively replaced by diffusion policy because diffusion policy fits, but does not improve, the current policy.Source: Convergence Analysis of Diffusion Policy · p. 5

1.2 Paper Organization

  • The paper organization assigns policy-representation motivation to Section 3, diffusion-policy theory to Section 4, online-RL implementation to Section 5, and experiments to Section 7.Figures 2 and 3 are captioned as the standard model-free online RL training framework and the DIPO framework, respectively.Source: §1.2 · p. 5, Figure 2 · p. 5, Figure 3 · p. 5

2 Reinforcement Learning

The section formulates reinforcement learning as a Markov decision process with continuous actions, state transitions, rewards, discounting, and an initial-state distribution. A policy specifies action probabilities given states, and the objective is to find a policy that maximizes expected return.

2 Reinforcement Learning

  • RL is formulated as an MDP with state space, continuous action space, transition probabilities, rewards, discount factor, and initial-state distribution.The passage names the MDP as M = (S, A, P(·), r, γ, d0).Source: §2 · p. 6
  • A policy π is a probability distribution over state-action pairs, with π(a|s) giving the probability of playing action a in state s.Trajectories sample actions from the policy and next states from the transition distribution.Source: §2 · p. 6
  • The RL objective is to find a policy that maximizes expected return.Source: §2 · p. 6

3 Motivation: A View from Policy Representation

The section contrasts value-based, parametric, and stochastic-process policy representations, then motivates diffusion policies as a way to model multimodal action distributions. In a four-goal example, diffusion policy captures multimodality more accurately than SAC, TD3, and PPO.

3 Motivation: A View from Policy Representation

  • The section presents diffusion models as capable of representing complex policy distributions and motivates this view through exploration and a numerical experiment.Source: §3 · p. 6, §3.1 · p. 6, §3.2 · p. 7

3.1.1 Policy Representation via Value Function

  • Value-based policy representations connect actions to a value function, including ε-greedy and energy-based forms.The cited examples include value-based methods such as Q-Learning and DQN, and energy-based methods such as SAC.Source: §3.1.1 · p. 6

3.1.2 Policy Representation via Parametric Function

  • Parametric policies represent a policy with a function such as a neural network, and the policy gradient theorem is central to learning its parameters.The passage lists TRPO, DDPG, PPO, and IMPALA as examples.Source: §3.1.2 · p. 6

3.1.3 Policy Representation via Stochastic Process

  • Diffusion policy generates actions through a stochastic process, which the paper presents as a policy representation distinct from value-based and parametric approaches.The paper notes earlier diffusion RL methods as offline approaches and describes its own focus as online model-free RL.Source: §3.1.3 · p. 7

3.2 Diffusion Model is Powerful to Policy Representation

  • Gaussian policies are unimodal and concentrate actions around their mean, while multimodal policies distribute actions across distinct regions.The paper argues that unimodality weakens expressiveness for complicated policies and decays exploration ability.Source: §3.2 · pp. 7–8, Figure 4 · p. 7
  • In the multi-goal environment, diffusion policy accurately captures a multimodal action distribution, whereas SAC, TD3, and PPO do not capture the multimodality well.The four symmetric goals require actions directed uniformly toward each goal; diffusion policy also maintains more gradual and steady action sizes than the compared methods.Source: Figure 5 · p. 7, §3.2 · p. 8

4 Diffusion Policy

Diffusion policy is defined by forward and reverse stochastic dynamics, implemented with an exponential-integrator discretization. Its convergence analysis bounds error through forward-process convergence, discretization, and score estimation, under stated regularity and distribution assumptions.

4.1 Stochastic Dynamics of Diffusion Policy

  • For a given state, the forward Ornstein–Uhlenbeck process starts from the policy distribution and transforms it toward Gaussian noise; the reverse process starts from that noise and recovers the policy distribution.The reverse dynamics use the score function, which the paper estimates with a function approximator.Source: §4.1 · pp. 8–9

4.2 Exponential Integrator Discretization for Diffusion Policy

  • The reverse process is discretized with an exponential integrator over a partition of the time interval, and the resulting iteration is presented in Algorithm 1.The discretization uses step size h and reverse length K.Source: §4.2 · p. 9, §4.1 · p. 9

4.3 Convergence Analysis of Diffusion Policy

  • The convergence bound separates error from forward-process convergence, exponential-integrator discretization, and score matching.The forward-process term reflects closeness of the input policy to Gaussian noise; discretization error has first-order convergence in step size h; score-matching error reflects the score estimator's accuracy.Source: §4.3 · pp. 10–11
  • Under the theorem's assumptions, sufficiently accurate score estimation and a suitable step size yield a diffusion policy whose KL divergence from the input policy is at most ϵ.The stated assumptions include Lipschitz conditions for the score estimator and policy, and a ν-Log-Sobolev inequality for the policy distribution.Source: §4.3 · pp. 10–11

5 DIPO: Implementation of Diffusion Policy for Model-Free Online RL

DIPO implements diffusion policy for model-free reinforcement learning and addresses the challenge that fitting a policy from data does not itself improve that policy.

5.1 Training Loss of DIPO

  • Because the score is unknown, DIPO uses denoising score matching, sampling replay state-action pairs, noise levels, and Gaussian noise and minimizing a noise-prediction loss.The objective provides a way to learn the noise-prediction parameters from samples.Source: §5.1 · pp. 11–12

5.2 Playing Action of DIPO

  • DIPO generates an action by running the reverse-time process with noise at each step and executing the final output action.The steps run from k = 0 to K − 1.Source: §5.2 · p. 12

5.3 Policy Improvement of DIPO

  • DIPO improves stored actions by ascending an action gradient based on estimated Q-values, then uses the updated state-action pairs to learn a new diffusion policy.The Q network is trained by minimizing Bellman residual error, and the action gradient is the gradient of expected Q with respect to actions.Source: §5.3 · pp. 12–13

6 Related Work

The related work reviews diffusion methods for reinforcement learning and other generative approaches to policy learning. It contrasts prior offline diffusion-RL work with DIPO’s online model-free setting.

6 Related Work

  • The related-work section focuses on research connecting diffusion models and generative models to reinforcement learning.It notes a separate comprehensive diffusion-model survey [Yang et al., 2022].Source: §6 · p. 13

6.1 Diffusion Models for Reinforcement Learning

  • Prior diffusion-RL work includes trajectory refinement and diffusion-generated actions [Chen et al., 2023a, Wang et al., 2023, Chi et al., 2023]; the reviewed methods address offline RL.The paper presents DIPO as the first diffusion approach to online model-free RL and notes that its training framework can use other function fitters.Source: §6.1 · p. 13

6.2 Generative Models for Policy Learning

  • Policy-learning generative models include VAE behavior priors and latent skills, GAN imitation, flow and GFlowNet methods, and Decision Transformer, energy-based, and autoregressive alternatives.The review places these methods in behavior cloning, imitation learning, goal-conditioned imitation, or offline RL contexts.Source: §6.2 · pp. 13–14

7 Experiments

Experiments compare DIPO with SAC, PPO, and TD3 on MuJoCo, visualize state visitation, compare diffusion with VAE and MLP, and examine reverse length K.

7.1 Comparative Evaluation and Illustration

  • DIPO achieves the best score on all five MuJoCo tasks and learns faster than SAC, PPO, and TD3 on Ant-v3 and Walker2d-v3.Its asymptotic reward is similar to baselines on the other three tasks, while its initial reward is best on all five.Source: §7.1 · pp. 14–15, Figure 6 · p. 15

7.2 State-Visiting Visualization

  • On Ant-v3, DIPO spans wider states than TD3, SAC, and PPO and concentrates densely late; PPO's aimless exploration partly explains its weaker task result.Source: Figure 7 · p. 15, §7.2 · pp. 15–16
  • On HalfCheetah-v3, DIPO covers more initial regions and earns better early reward than SAC, but both reach similar final reward; SAC starts near its final good region.Source: Figure 8 · p. 15, §7.2 · p. 16

7.3.1 Comparison to VAE and MLP

  • Replacing DIPO's diffusion policy with VAE or MLP plus action gradient yielded weaker reward performance, supporting diffusion as the more powerful distribution learner.Source: Figure 9 · p. 16, §7.3.1 · p. 17

7.3.2 Comparison with Different Reverse Lengths

8 Conlusion

The section summarizes diffusion policies as stochastic-process representations for reinforcement learning, their convergence analysis, and DIPO, while examining the k = 0 process through comparison with the true backward-process distribution over t ∈ [0, h].

8 Conlusion

  • DIPO combines a theoretical foundation for diffusion-policy representation via a stochastic process, convergence analysis, and a model-free online RL implementation, with reported effectiveness on MuJoCo tasks.The conclusion identifies these as the paper’s main contributions.Source: §8 · p. 17

B.1 Diffusion Probability Model (DPM).

  • The DPM background describes forward diffusion from data toward Gaussian noise and a corresponding reverse SDE that can sample data when its score function is known.It also states that the transition distribution for the Ornstein–Uhlenbeck process is Gaussian conditional on its initial value.Source: §B.1 · p. 27, §B.2 · p. 27

B.3 Exponential Integrator Discretization

  • The appendix derives an exponential-integrator discretization for the diffusion process, which DIPO uses to generate actions through a sampling iteration.It also presents an equivalent backward-time version of the iteration used in much diffusion-model literature.Source: §B.3 · p. 28, §C.3 · pp. 34–35

B.4 Fokker–Planck Equation

  • The appendix reviews the Fokker–Planck equation for density evolution and the Donsker–Varadhan representation of KL divergence, then introduces a time derivative of KL between the diffusion policy and the true reverse process.The derivative analysis concerns the policy from Algorithm 1 and the true reverse process defined in equation (4).Source: §B.4 · p. 29, §B.5 · p. 29, D Time Derivative of KL Divergence… · p. 36

B.6 Some Basic Results for Diffusion Policy

  • Under stated assumptions, the appendix gives basic diffusion-policy results including preservation of a time-varying log-Sobolev constant along the Ornstein–Uhlenbeck flow and Lipschitz regularity of the evolved policy’s score.The log-Sobolev constant is specified as ν_t in the appendix.Source: §B.6 · pp. 29–30

C.2 Loss Function of DIPO

  • DIPO’s score-matching objective has an empirical estimator that is unbiased, and the policy parameters are learned by gradient descent on that loss.The loss derivation presents an equivalent expectation form before describing learning from sampled experience.Source: §C.2 · p. 32, §C.2.2 · p. 34

D.1 Time Derivative of KL Divergence at Reverse Time k = 0

  • For k = 0, the diffusion-policy SDE is reduced on 0 ≤ t ≤ h and is driven by a standard Wiener process starting at w_0 = 0.The reduced SDE itself is not included in the supplied text.Source: §D.1 · p. 36
  • Proposition D.1 compares the diffusion-policy and true backward-process distributions over t ∈ [0, h] under Assumptions 4.1 and 4.2.The supplied passage states the comparison setup but does not include the proposition’s conclusion.Source: §D.1 · p. 36
  • The proof setup defines ρ_0,t as the joint conditional distribution of (â_0, â_t) given state s along the k = 0 process.The passage says this joint distribution can be expressed using conditionals and marginals, but does not provide the expression.Source: §D.1 · p. 36

D.2 Auxiliary Results For Reverse Time k = 0

These appendices derive auxiliary bounds for the reverse-time diffusion process, then extend the analysis from the initial reverse-time step to arbitrary steps. The results support the paper’s convergence analysis under stated assumptions.

D.2 Auxiliary Results For Reverse Time k = 0

  • The KL-divergence time derivative between the evolving distributions is decomposed and bounded through a sequence of auxiliary lemmas.The derivation uses identities for relative entropy and relative Fisher information.Source: §D.2 · pp. 38, 39, 41
  • Under Assumption 4.1, a short-time bound holds for the update, with the condition 0 ≤ t ≤ 1/(12L_s).The update is analyzed over a single step of the reverse-time process.Source: §D.2 · p. 42

D.3 Proof for Result at Reverse Time k = 0

  • The proof uses an L_p-Lipschitz condition on [0, T_0] rather than uniform L-smoothness over [0, T], which the authors describe as a weaker condition.The authors relate this proof choice to a similar convergence rate from Langevin-based algorithms.Source: §D.3 · pp. 43, 45

D.4 Proof for Result at Arbitrary Reverse Time k

  • Under Assumptions 4.1 and 4.2, Proposition D.9 considers the distribution of iteration (9) at times tk = hk, initialized at N(0, I), for step size 0 < h ≤ τ.Source: §D.4 · p. 45

E Proof of Theorem 4.3

This appendix proves Theorem 4.3 by bounding the KL-divergence between the discretized diffusion policy and its input policy. It also documents the multimodal task and the visualizations used to assess policy representation.

E Proof of Theorem 4.3

  • Theorem 4.3 upper-bounds the KL-divergence between the discretized diffusion policy and the input policy by terms attributed to discretization, forward-process convergence, and score-matching errors.The theorem concerns a reverse process initialized from a standard Gaussian and assumes conditions stated in the theorem.Source: E Proof of Theorem 4.3 · p. 46

G.1.1 Task

  • The multimodal task places a 2D point mass on a 7 × 7 plane and asks it to reach one of four symmetrically placed goals.The goals are (0, 5), (0, −5), (5, 0), and (−5, 0).Source: §G.1.1 · p. 51

G.1.2 Reward

  • Because the four goals are symmetric, the expected reward curve has four equal peaks, and a reasonable policy should distribute actions uniformly toward them.The reward includes a distance-based term and a reward of 10 for reaching a target.Source: §G.1.2 · p. 51

G.2 Plots Details of Visualization

  • The 2D visualizations plot action vectors from each of 49 states, while the 3D visualizations show critic state-action values, or PPO’s value function, across a 100 × 100 state grid.Vector direction and length represent action direction and magnitude, respectively.Source: §G.2.1 · p. 52, §G.2.2 · p. 53

G.2.3 Shape of Reward Curve

  • The appendix treats uniform action coverage of the four goals and a four-peaked reward shape as visual indicators of multimodal policy representation.The authors connect multimodal policy distributions with exploration and potentially better performance.Source: §G.2.3 · pp. 53–54

G.3 Results Report

In the multimodal experiment, the diffusion policy captures the reward landscape’s multimodality, whereas SAC, TD3, and PPO do not capture it as well. The appendix adds implementation details and state-visiting comparisons.

G.3 Results Report

  • Figure 11 shows the diffusion policy capturing the multimodal reward landscape, while Figures 12–14 show SAC, TD3, and PPO are less suited to capturing that multimodality.Diffusion, SAC, and TD3 all maintain the highest reward around −20, but the diffusion policy matches the environment and performance shape.Source: §G.3 · p. 54, Figure 14 · p. 55
  • PPO stays near its initial reward and does not improve performance, which the authors interpret as failure to fit multimodality.The report says PPO learns no information about multimodality.Source: §G.3 · p. 55, Figure 14 · p. 55
  • The diffusion policy maintains more gradual and steady action sizes than SAC, TD3, and PPO during learning of the multimodal reward.The authors associate this pattern with more sufficient exploration and better performance.Source: §G.3 · p. 55

H.2 Additional Tricks for Implementation of DIPO

  • The supplementary implementation details include double Q-learning, a soft update rule, and an action-gradient update for replay-buffer actions.The critic and noise-term networks use four-layer feedforward architectures with 256 hidden nodes and gradient normalization.Source: §H.2.1 · p. 57, §H.2.2 · p. 57

H.3 Details and Additional Reports for State-Visiting

  • Additional state-visiting analyses list poor exploration with poor initial reward, dense state-visiting with good final reward, and PPO as a counterexample to both patterns.These are presented as observations, not as universal rules.Source: §H.3 · p. 58

H.3.1 Implementation Details for 2D State-Visiting

  • For the state-visiting visualizations, the authors collect ten episodes per algorithm using the same ten random seeds, then reduce state data to two dimensions with t-SNE.Figures 15–19 report results for Humanoid, Walker2d, Ant, HalfCheetah, and Hopper.Source: Figure 15 · p. 57, §H.3.1 · p. 58, Figure 16 · p. 58

H.3.2 Observation 1: Poor Exploration Result in Poor Initial Reward Performance

The exploration analyses connect broader early state-visiting with better initial reward performance, while later state-visiting patterns distinguish DIPO from PPO. In the ablation, diffusion achieves the best reward performance across five tasks, and MLP with action gradient also performs well.

H.3.2 Observation 1: Poor Exploration Result in Poor Initial Reward Performance

  • For Hopper, DIPO and SAC show broader state-visiting and better initial reward performance than TD3 and PPO; similar patterns appear in other tasks.The authors report sparse early state-visiting alongside worse initial reward performance for TD3 and PPO, and broader visitation alongside better initial performance in cited Walker2d and Humanoid examples.Source: §H.3.2 · p. 59

H.4 Ablation Study on MLP and VAE

  • Diffusion achieves the best reward performance among all five tasks in the comparison with MLP and VAE.Figure 20 reports average MuJoCo Gym performance over iterations with ± standard-deviation shading.Source: Figure 20 · p. 60, §H.4 · p. 61

H.3.3 Observation 2: Good Final Reward Performance along with Dense State-Visiting

  • At final iterations, DIPO, SAC, and TD3 show denser state-visiting in Hopper, Walker2d, and Humanoid, while PPO explores aimlessly in Ant-v3 and HalfCheetah.The authors relate the denser visitation to avoiding “bad” states after sufficient training, and describe the pattern as early exploration followed by final exploitation.Source: §H.3.3 · p. 60

H.3.4 Observation 3: PPO Violates above Two Observations

  • Across five tasks, PPO’s state-visiting pattern reverses the described pattern: it shifts from narrow early visitation to wider visitation later, unlike DIPO.For Hopper-v3, PPO reaches asymptotic reward performance as DIPO despite their fundamentally different state-visiting distributions.Source: §H.3.4 · pp. 60–61

H.4 Ablation Study on MLP and VAE

  • MLP with action gradient performs well across all five tasks; on Humanoid-v3, its final reward is near PPO, SAC, DIPO, and TD3.The authors identify the simple MLP pairing as a direction for further research.Source: §H.4 · p. 61
Loading 2305.13122v1…