PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning a Diffusion Model Policy from Rewards via Q-Score Matching

Michael Psenka, Alejandro Escontrela, Pieter Abbeel, Yi Ma

arXiv:2312.11752v5cs.LGcs.AI

TL;DR

Continuous-action RL needs expressive policies that remain easy to sample, while existing diffusion-policy approaches have not fully used their score structure for actor-critic learning. This paper develops Q-score matching to align policy scores with Q-function action gradients, and reports performance similar to or better than SAC and TD3 on six control tasks. Its analysis focuses on the score while assuming a fixed noise model.

  • Problem

    Continuous-action RL poses a policy-parameterization challenge, and prior diffusion-policy work has not fully exploited score structure for actor-critic learning.

    Source: §1 · p. 1

  • Method

    Q-score matching iteratively aligns a diffusion policy’s score with the action gradient of its learned Q-function for off-policy RL.

    Source: §1 · p. 1, §4 · p. 4

  • Results

    QSM achieves similar or better performance than SAC and TD3 across six Deepmind control-suite tasks.

    Source: §5.1 · p. 7

  • Takeaways & Limitations

    QSM provides a geometric policy-update approach that differentiates through the denoising model rather than the entire diffusion evaluation.

    Source: §6 · p. 8

    The work focuses on optimizing the score Ψ and assumes a fixed noise model Σa.

    Source: §6 · p. 8

Abstract

from arXiv · show

Diffusion models have become a popular choice for representing actor policies in behavior cloning and offline reinforcement learning. This is due to their natural ability to optimize an expressive class of distributions over a continuous space. However, previous works fail to exploit the score-based structure of diffusion models, and instead utilize a simple behavior cloning term to train the actor, limiting their ability in the actor-critic setting. In this paper, we present a theoretical framework linking the structure of diffusion model policies to a learned Q-function, by linking the structure between the score of the policy to the action gradient of the Q-function. We focus on off-policy reinforcement learning and propose a new policy update method from this theory, which we denote Q-score matching. Notably, this algorithm only needs to differentiate through the denoising model rather than the entire diffusion model evaluation, and converged policies through Q-score matching are implicitly multi-modal and explorative in continuous domains. We conduct experiments in simulated environments to demonstrate the viability of our proposed method and compare to popular baselines. Source code is available from the project website: https://michaelpsenka.io/qsm.

1. Introduction

Diffusion policies offer expressivity and easy sampling for continuous-action RL, but prior work has not fully exploited their score structure for off-policy actor learning. The paper introduces Q-score matching, motivated by a link between policy scores and action gradients of Q-functions.

1. Introduction

  • Diffusion policies suit continuous-action RL because they combine expressive distributions with easy sampling, while Gaussian policies trade expressiveness against sampling ease.Diffusion sampling does not require computing normalization constants.Source: §1 · p. 1
  • Q-score matching (QSM) iteratively matches a policy’s score to the action gradient of its Q-function, addressing the lack of optimal-policy samples for direct score matching.The paper presents this as a geometric view of policy optimization: pushing the policy-score vector field toward the Q-gradient field.Source: §1 · p. 1
  • The paper applies QSM to off-policy RL, where the absence of a fixed sampling distribution makes diffusion-policy training unclear, and reports it as viable.This setting is presented as important but previously unexplored for diffusion-model policies.Source: §1 · p. 1

1.1.1. DIFFUSION MODELS IN REINFORCEMENT LEARNING

  • Earlier diffusion-policy work emphasizes behavior cloning [Janner et al., 2022; Reuss et al., 2023] and offline Q-learning [Wang et al., 2022; Suh et al., 2023; Kang et al., 2023; Hansen-Estruch et al., 2023; Lu et al., 2023].The fixed datasets in those settings align more closely with diffusion models’ distribution-learning use.Source: §1.1.1 · p. 2
  • Prior diffusion policy-gradient formulas can depend on internal action samples, while action-gradient methods require an explicit policy distribution unavailable for diffusion models.Diffusion-QL instead backpropagates through the full diffusion evaluation, which the paper describes as computationally challenging.Source: §1.1.1 · p. 2

1.1.2. STOCHASTIC OPTIMAL CONTROL

  • The work relates to stochastic optimal control but assumes a model of expected discounted rewards through the Q-function rather than a state-dynamics model.The authors identify a deeper connection to stochastic optimal control as future work.Source: §1.1.2 · p. 2

2. Problem formulation

The paper formulates continuous-time state and action dynamics so policy optimization can be analyzed through an action-score vector field. Time-discretizing the action process at higher fidelity yields a diffusion-like generative policy, subject to discretization error.

2. Problem formulation

  • The formulation uses continuous-time dynamics for both state and action, an uncommon RL choice that supports the link to diffusion models.The authors state that continuous time simplifies theoretical statements.Source: §2 · p. 2, §2.2 · p. 3

2.2. Definitions

  • State and action evolve with continuous dynamics and separate uncertainty covariances, while the action vector field Ψ is the main policy-optimization parameter.The covariances correspond to uncertainty in the respective state and action dynamics.Source: §2.2 · p. 3

2.1. Notation

  • The policy parameter Ψ is an action-vector-field “score”; in some settings it is analogous to, but not identical with, the distribution score ∇a log π(a|s).The notation uses dot derivatives in deterministic settings and SDE notation in stochastic settings.Source: §2.2 · p. 3, §2.1 · p. 3

2.2. Definitions

  • The objective maximizes expected discounted path-integral reward over trajectories generated by the stochastic dynamics and initial state-action pairs.The reward is defined on states, and γ is the fixed discount factor.Source: §2.2 · p. 3

2.3. Time discretization

  • Actions are modeled as a smooth flow over time, then discretized at K times the fidelity of state dynamics to obtain a score-based generative model.The discretized model is analogous to a depth-K diffusion model, with Ψ as denoising mean and Σa as variance.Source: §2.3 · p. 3
  • The continuous-time theory is expected to approximately apply to the discretized diffusion-like system, with error from time discretization.This is an expectation stated by the authors, not an exact equivalence.Source: §2.3 · p. 3
  • Increasing K lets actions depend more on the current state than the previous action, making the initial previous-action condition less influential.The passage states this effect as K increases.Source: §2.3 · p. 4

3. Policy gradient for diffusion policies

The standard policy-gradient expression for diffusion policies is difficult to use directly because the policy exposes incremental transitions rather than a global action probability. Rewriting it yields an expectation over internal action states, motivating an alternative score-structured update.

3. Policy gradient for diffusion policies

  • The global policy probability is unavailable for a time-discretized diffusion policy, which exposes only incremental action transitions and makes path integration quickly infeasible as K grows.The policy-gradient expectation can be rewritten, but then depends explicitly on internal action states rather than only executed actions.Source: §3 · p. 4
  • The rewritten policy-gradient formula uses internal action states and does not exploit diffusion-model structure, so its estimates are expected to be less sample efficient.The formula also holds for models with internal actions generally.Source: §3 · p. 4
  • The paper therefore pursues alternative diffusion-policy updates that use score-based structure.The authors motivate this direction after identifying the internal-action dependence of the rewritten formula.Source: §3 · p. 4

4. Policy optimization via matching the score to the Q-function

The section develops Q-score matching as an alternative diffusion-policy update: align the policy score with the learned Q-function’s action gradient. It establishes this alignment in deterministic and stochastic settings and illustrates its entropy effect in gridworld.

4. Policy optimization via matching the score to the Q-function

  • Q-score matching addresses unavailable optimal-policy samples by iteratively matching the policy score Ψ to ∇aQΨ, with gradients through only the denoising model.Source: §4 · p. 4

4.1. Non-stochastic setting

  • In deterministic dynamics, the optimal score aligns positively with the action gradient wherever it is nonzero; aligning a mismatch strictly increases Q toward Q*.Source: §4.1 · p. 5

4.2. Stochastic setting

  • Under stochastic dynamics, the same positive alignment holds on the support generated by P, π, and the dynamics; iterative matching is theoretically justified to increase Q globally.Source: §4.2 · p. 5

4.3. Pedagogical reduction in gridworld

  • In gridworld, the update matches soft policy iteration using standard Q, and lower α yields higher entropy in the converged policy.Source: §4.3 · p. 6
  • Across eight continuous-control tasks, QSM matches and sometimes outperforms TD3 and SAC, particularly in samples needed to reach high rewards, while matching Gaussian and tanh-model sample efficiency.Source: Figure 3 · p. 7

5. Experiments

The experiments assess QSM’s ability to learn with limited interaction, produce multimodal policies, and compare with popular baselines. On continuous-control tasks, QSM performs similarly to or better than SAC and TD3.

5. Experiments

  • The practical implementation uses N-step returns and a DDPG-style critic; each update targets the score model with ∇aQ from noised replay-buffer actions.Updates occur every environment step with batches of 256; the critic and score model are 2-layer MLPs with 256 units per layer, and small Gaussian noise is added to sampled actions.Source: §5 · p. 7

5.1. Continuous Control Evaluation

  • QSM achieves similar or better performance than SAC and TD3 across six DeepMind Control Suite tasks.Source: §5.1 · p. 7
  • QSM produces multimodal action samples; quadruped samples compare 1,000 actions per model at one time step, with training method as the only difference.The compared models share the same architecture and sampling procedure; QSM’s extreme-action mass corresponds to optimal-trajectory initial actions that later converge after trajectory commitment.Source: Figure 5 · p. 8, §5.1 · p. 8
  • Policy-gradient updates to diffusion policies perform notably worse than QSM on the selected continuous-control tasks.Source: Figure 6 · p. 8

6. Conclusion

The conclusion presents QSM as a way to train diffusion policies by matching their learned score to the Q-function’s action gradient. It reports favorable results against popular RL algorithms and identifies open directions for the method and its theory.

6. Conclusion

  • QSM iteratively matches the policy’s parameterized score to the action gradient of its Q-function, with favorable results against popular RL algorithms.Source: §6 · p. 8
  • The implementation focuses on the score while assuming a fixed noise model Σa; optimizing that model and answering broader theoretical questions remain future work.Source: §6 · p. 8
  • Convergence-rate analysis is a proposed direction that could highlight further improvements to QSM.Source: §6 · p. 9

A. Additional experiments

The appendix adds experiments to contextualize the main results and illustrate possible extensions to more complex tasks.

A. Additional experiments

B. Proofs

The proofs are intended to be read sequentially, with the first part of Theorem 4.1 carrying most of the analysis.

B. Proofs

  • The proofs build on one another and should be read in sequence.Source: B. Proofs · p. 12
  • The “aligning at the origin” portion of Theorem 4.1 is the longest, densest proof segment and contains most of the Q-function analysis.The authors describe the remaining theory as following from that analysis.Source: B. Proofs · p. 12

B.1. Proof of policy gradient formula in (9)

The proof establishes an equality between policy-gradient expectations for policies represented through incremental steps. It uses the identity for conditional marginal probabilities to rewrite the gradient of the integrated policy probability.

B.1. Proof of policy gradient formula in (9)

  • The stated proof introduces a policy-gradient formula for diffusion models through an incremental sequence of steps.The setup assumes access to the conditional marginal probabilities for those steps.Source: §B.1 · pp. 12–13
  • The theorem states that the two policy-gradient expectations are equal for the incremental policy parameterization.The parameterization uses conditional marginal probabilities satisfying the stated identity.Source: §B.1 · pp. 12–13
  • The proof substitutes the identity into one side, then rewrites the gradient of a log-integral to obtain the matching expression.Source: §B.1 · p. 13

B.2. Proof of Theorem 4.1

The proof argues that a score field not aligned with the action gradient of its Q-function can be locally modified to strictly increase Q. It builds the argument from a decomposition of Q and treats trajectories with and without re-entry into a neighborhood.

B.2. Proof of Theorem 4.1

  • The local decomposition separates Q into a state-only integral and a boundary term involving the same Q-function.The state-only integral can be made arbitrarily small.Source: §B.2 · p. 14
  • A score field misaligned with ∇aQ can be modified toward that gradient to strictly increase Q, including when the misalignment occurs at the origin.The proof constructs a modified field Ψ′ and compares QΨ′(0, 0) with QΨ(0, 0).Source: §B.2 · pp. 14, 15, 18
  • The proof handles trajectories that do not re-enter the neighborhood separately from those that re-enter, including finite or infinite re-entry sequences.For re-entry cases, it repeatedly decomposes Q until the remaining terms yield a strict increase.Source: §B.2 · pp. 15–18
  • The extension applies the strict-increase argument to a fixed starting condition and to expectations over distributions of starting conditions.Local perturbations affect a measure-positive region around a non-aligned point without decreasing Q at other state-action pairs.Source: §B.2 · p. 18

B.3. Proof of Theorem 4.3

The stochastic proof extends the local score-perturbation argument by conditioning on trajectory events and using the support of the joint distribution. It concludes that such perturbations strictly increase Q for fixed and distributed starting conditions.

B.3. Proof of Theorem 4.3

  • The proof defines support through positive probability of the state-action trajectory entering every neighborhood of a point at some finite time.Source: §B.3 · p. 19
  • Local score perturbations strictly increase Q both at a fixed starting condition and across a distribution of starting conditions.The stochastic proof conditions on how often trajectories re-cross the neighborhood around the origin.Source: §B.3 · p. 20
  • The theoretical optimization uses bounded smooth vector fields with a Lipschitz condition, but the chosen space is not closed.Tightening the space, for example with a fixed Lipschitz constant, is left for future work.Source: §B.3 · p. 20
Loading 2312.11752v5…