PaperlayerView exact revision on arXiv

Source-linked AI summary

Doob-Based Generative Policy Optimization

Maojiang Su, Qijie Zhu, Yu Xiao, Shuyang Yu, Minshuo Chen, Zhaoran Wang, Han Liu

arXiv:2610.28990v1cs.LG

TL;DR

Diffusion and flow policies have generally intractable action likelihoods, limiting direct use of policy-optimization methods such as PPO and GRPO. DoobPO separates policy improvement from generative-policy learning and distills a ratio-tilted target through conditional drift matching. It reports strong MuJoCo results and characterizes the diffusion modification as a minimum-energy realization of the policy tilt.

  • Problem

    Generative-policy action likelihoods are generally intractable, limiting direct use of widely adopted policy-optimization methods such as PPO and GRPO.

    Source: §1 · p. 2

  • Method

    DoobPO represents policy improvement with a target ratio, then distills the ratio-tilted target into diffusion or flow policies through conditional drift matching.

    Source: §1 · p. 2

  • Results

    DoobPO achieves the strongest overall performance among diffusion-policy RL baselines, and DoobGRPO performs best on 9 out of 10 flow-policy tasks.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The Doob-transformed diffusion realizes the terminal policy tilt through a minimum-energy modification of the pretrained process.

    Source: §1 · p. 3

    The empirical evaluation is limited to MuJoCo continuous-control benchmarks and Robomimic manipulation tasks, while the theoretical results characterize population objectives under stated assumptions.

    Source: §5 · p. 15

Abstract

from arXiv · show

Diffusion and flow models provide expressive policy classes for robotic and sequential decision-making systems. However, their action likelihoods are generally intractable, limiting the direct use of widely adopted policy optimization methods such as PPO and GRPO. We propose Doob Policy Optimization (DoobPO), a framework that leverages Doob's $h$-transform for likelihood-free reinforcement learning fine-tuning of generative policies. DoobPO separates policy improvement from generative policy learning. It first learns a lightweight ratio model that represents the desired policy update, and then distills the resulting ratio-tilted target policy into the generative policy via conditional drift matching. This formulation applies uniformly to both diffusion and flow policies and supports a broad family of policy improvement rules, including PPO and GRPO. We prove that the practical conditional drift matching objective is gradient-equivalent to the ideal Doob drift matching objective. We further show that the induced diffusion process solves an endpoint-constrained stochastic optimal control problem. This establishes that it achieves the minimum-energy modification required to realize the desired policy tilt. Across MuJoCo continuous-control benchmarks, DoobPO achieves strong performance against competitive diffusion- and flow-policy baselines, showing a principled and practical route to RL fine-tuning for generative policies.

1 Introduction

DoobPO addresses the intractable action likelihoods that limit direct reinforcement-learning fine-tuning of diffusion and flow policies. It separates policy improvement from generative policy learning, with strong reported results across MuJoCo continuous-control benchmarks.

Motivation

  • Generative-policy action likelihoods are generally intractable because actions require simulating a denoising or flow trajectory, limiting direct use of methods such as PPO [Schulman et al., 2017] and GRPO [Shao et al., 2024].Prior approaches can involve biased estimates, limited data efficiency, computational overhead, or restrictions to particular policy classes or update rules.Source: §1 · p. 2

Method

  • DoobPO first represents a desired policy improvement as a target policy ratio, then distills the ratio-tilted target into a generative policy through conditional drift matching.The framework supports both diffusion and flow policies and policy updates including PPO [Schulman et al., 2017] and GRPO [Shao et al., 2024].Source: §1 · p. 2

Theory

  • For diffusion and flow policies, conditional drift matching has a gradient proportional to ideal Doob drift matching, and the Doob-transformed diffusion solves an endpoint-constrained control problem.The resulting modification is the minimum-energy change to the pretrained diffusion process that realizes the terminal tilt.Source: §1 · p. 3

Experiments

  • DoobPO achieves the strongest overall performance among diffusion-policy RL baselines, while DoobGRPO performs best on 9 out of 10 flow-policy tasks.The reported flow-policy gains are large on challenging locomotion environments.Source: §1 · p. 3

2 Preliminaries

This section introduces the reinforcement-learning setup and the generative policy classes considered. It explains why likelihood-based trust-region methods are difficult to apply to diffusion and flow policies.

Reinforcement Learning

  • An MDP is specified by state and action spaces, a transition kernel, reward function, and discount factor; a policy and transition kernel induce trajectory distributions.The section denotes these components by (S, A, P, R, γ).Source: §2 · p. 3

Policy Optimization

  • Policy-gradient methods ascend an estimate of the objective gradient, often using an advantage estimate in place of cumulative future reward to reduce variance.Trust-region methods constrain policy changes between successive iterates; PPO uses a clipped surrogate that limits updates away from the behavior policy.Source: §2 · pp. 3–4

Diffusion Models and Score Matching

  • Diffusion policies use the action distribution at the end of a reverse-time denoising process, whereas flow policies use the terminal distribution of a learned-velocity ODE.Diffusion generates actions by simulating an SDE; flow generates them by simulating an ODE.Source: §2 · pp. 4–5

Policy Optimization

  • Evaluating diffusion- and flow-policy likelihoods is expensive, making the probability ratios required by trust-region methods such as PPO and GRPO intractable.This limits direct application of those methods to improving these generative policies.Source: §2 · p. 5

3 Likelihood-Free Generative Policy Optimization

DoobPO separates policy improvement from generative policy learning: it represents an update with a target policy ratio, then distills the ratio-tilted target into a generative policy. Its diffusion and flow formulations use conditional drift matching, and the framework covers multiple policy-update methods.

3.1 Diffusion Policy Optimization with DoobPO

  • For diffusion policies, DoobPO learns a lightweight ratio model and distills the resulting policy tilt through conditional drift matching, whose gradient direction matches the ideal objective.The ideal drift-matching objective requires intractable tilted-distribution samples and drift access; the conditional objective avoids these requirements.Source: §3.1 · pp. 6–8

3.2 Flow Policy Optimization with DoobPO

  • For flow policies, DoobPO reuses ratio learning and distills the terminally tilted policy into a velocity model through conditional drift matching.The approximation is described as reliable when the pretrained velocity model is well trained under the specified affine path and the trust-region update keeps the behavior policy close to the base behavior.Source: §3.2 · p. 9

3.3 Unified Framework and Instantiations

  • DoobPO combines a generative policy class with a policy-update method, covering closed-form updates and learned-ratio instantiations such as DoobPPO and DoobGRPO.When an update has a closed-form target ratio, ratio learning is unnecessary and the method reduces to a reweighted drift-matching step.Source: §3.3 · p. 10

3.4 Optimality from the Perspective of Stochastic Optimal Control

  • Conditional drift matching solves an endpoint-constrained stochastic optimal control problem by learning the minimum-energy drift modification that realizes the tilted policy.The result concerns controlled reverse diffusions with prescribed endpoint marginals; its stated minimizer result assumes a realizable model class and an attained minimum.Source: §3.4 · pp. 10–11

4 Experiments

The experiments evaluate DoobPO on MuJoCo continuous-control benchmarks with diffusion and flow policies, and on Robomimic manipulation tasks.

4.1 Experimental Setup

  • 10 MuJoCo tasks span locomotion, manipulation, and stabilization; diffusion and flow methods train from scratch for 4M and 5M steps, respectively.Figure 1 tracks mean undiscounted return every 200k steps over 3 seeds, with shaded regions showing one standard deviation.Source: §4.1 · pp. 11–12, Figure 1 · p. 12

4.2 Diffusion Policy Comparison

  • 8 of 10 MuJoCo tasks (including ties): DoobPO has the highest mean returns among diffusion-policy RL methods, with DoobPPO and DoobGRPO leading on different tasks.DoobPPO leads on HalfCheetah, Humanoid, and Swimmer; DoobGRPO leads on Reacher and Pusher, while DACER leads on Ant and Hopper. Against DIPO, DoobPPO improves by 22.6% on HalfCheetah, 20.9% on Humanoid, 175.2% on Ant, and 12.2% on Walker2d; against SDAC, gains are 8.2% on Humanoid, 8.8% on Hopper, and 8.9% on Walker2d.Source: Figure 1 · p. 12, §4.2 · pp. 12–13, Table 2 · p. 13

4.4 Robomimic Manipulation

  • 99.46% versus 99.23% on Can and 99.78% versus 99.59% on Square: DoobPPO matches or slightly improves DPPO on near-saturated tasks.The evaluation uses sparse binary success rewards and extends DoobPO to Robomimic manipulation settings.Source: §4.4 · p. 14

4.5 Implementation Details and Ablation Studies

  • Implementation details include computational resources, training hyperparameters, and reproducibility information; ablations examine ratio-network size and sampling-step count.Source: §4.5 · p. 14

4.3 Flow Policy Comparison

  • 9 of 10 MuJoCo tasks (including ties): DoobGRPO is strongest under the 5M-step budget, improving over ReinFlow by 373.8% on HalfCheetah.DoobGRPO also improves over ReinFlow by 71.7% on Hopper, 213.7% on Swimmer, and 738.0% on Walker2d.Source: §4.3 · p. 13, Figure 2 · p. 14, Table 3 · p. 15

5 Discussion and Conclusion

DoobPO separates policy improvement from generative policy learning, with theory connecting conditional drift matching to an endpoint-constrained control problem. The authors report strong overall MuJoCo results while noting limits in evaluation scope and theoretical assumptions.

5 Discussion and Conclusion

  • DoobPO’s conditional drift matching has a gradient proportional to the ideal objective, and under stated assumptions its diffusion minimizes the energy needed to realize tilted endpoint marginals.The authors interpret this as a control-theoretic account of transferring policy updates into generative dynamics.Source: §5 · p. 15
  • The empirical study is limited to continuous generation, MuJoCo continuous control, and Robomimic manipulation, while theory characterizes population objectives under stated assumptions.The authors identify quantifying ratio-estimation and approximate drift-matching errors as future work.Source: §5 · p. 15

Impact Statement

LLM Usage Disclosure

  • Large language models aided writing and literature retrieval, while the authors state that technical content, proofs, experiments, and results are original contributions.Source: LLM Usage Disclosure · p. 17

A Notation

  • The notation distinguishes diffusion and flow endpoint conventions and denotes the learned ratio model as rβ(o, a), with terminal tilt proportional to that ratio.The stated realizability condition requires a score-model parameter setting to equal the Doob drift for all listed states, times, and observations.Source: A Notation · p. 17

B Related Works

  • For diffusion policies, earlier work reports strong performance in visuomotor control and robotic manipulation, as well as imitation and conditioned skill learning [Reuss et al., 2023, Chen et al., 2023].Source: B Related Works · p. 18
  • DoobPO generalizes analytic-ratio ideas from DPMD and SDAC to flow policies and broader updates, including PPO [Schulman et al., 2017] and GRPO [Shao et al., 2024].Flow-policy likelihoods remain intractable because evaluation requires solving an ODE and computing its Jacobian log-determinant; prior workarounds use stepwise likelihoods or surrogates.Source: B Related Works · pp. 18–19

C Proof of Lemma 3.1

The proof fixes a state and applies reverse-time Doob-transform and bridge-marginal identities to characterize the tilted denoising process.

C Proof of Lemma 3.1

D Proof of Theorem 3.1

The proof expands the conditional drift-matching loss and relates its gradient to the ideal drift-matching gradient through ratio-weighted expectations and a normalizing constant.

D Proof of Theorem 3.1

  • The proof establishes that the conditional drift-matching gradient equals the ideal objective’s gradient multiplied by Zβ(o).Source: D Proof of Theorem 3.1 · p. 21
  • The proof fixes the conditioning state and omits it initially, then extends the argument pointwise across states.Source: D Proof of Theorem 3.1 · p. 21
  • The normalizing constant is the expected learned ratio under the reference endpoint distribution, and the loss expansion uses ratio-weighted conditional expectations.Source: D Proof of Theorem 3.1 · pp. 21–22

E Proof of Theorem 3.2

The flow proof derives gradient equivalence by relating expectations under the ratio-tilted flow path to the reference path and comparing the resulting expressions.

E Proof of Theorem 3.2

F Proof of Proposition 3.1

The proof identifies conditional drift matching with the optimal control for an endpoint-constrained stochastic optimal control problem. The conclusion holds when the model class is realizable and the objective minimum is attained.

F Proof of Proposition 3.1

  • Under realizability and attainment of the minimum, any drift-matching minimizer satisfies the stated minimizer condition.The proposition concerns the conditional objective and the endpoint distributions defined in (3.3).Source: F Proof of Proposition 3.1 · p. 26
  • The proof connects the reverse-time denoising process to forward-time stochastic optimal control, then identifies the corresponding optimizer.The cited SOC formulation is written in forward time, while the denoising process is written in reverse time.Source: F Proof of Proposition 3.1 · pp. 26–28
  • Any minimizer of the conditional drift matching objective induces the optimal control for the stochastic optimal control problem.The proof states the result conditionally on the state o.Source: F Proof of Proposition 3.1 · pp. 26, 28

G.1 Policy Evaluation

The section describes PPO- and GRPO-style evaluation signals for ratio learning, an unbiased regularizer estimate, and computational costs that remain comparable to strong diffusion-policy baselines.

G.1 Policy Evaluation

  • PPO evaluation estimates advantages with generalized advantage estimation and uses them as the first-stage ratio network’s policy-improvement signal.The value function is trained by regression against empirical returns or bootstrap targets.Source: §G.1 · pp. 28–29
  • For DoobGRPO, returns from sampled action groups are normalized within each state to produce the relative-advantage signal for ratio learning.The group size is G, and each sample’s return is denoted Ri.Source: §G.1 · p. 29

G.2 Unbiased Ratio Regularizer

  • Double sampling replaces the squared ratio-centering moment with the product of two conditionally independent estimates, yielding an unbiased estimate.For each sampled state, two independent action batches are drawn from the old policy.Source: §G.2 · p. 29

G.3 Ratio Regularizer in the Off-Policy Setting

  • In replay-buffer training, centering uses actions sampled from the current actor because stored transitions have no single well-defined behavior policy.The authors report that a small coefficient, λ = 0.01, stabilizes ratio learning during training.Source: §G.3 · p. 29

G.4 Computation and Memory Cost

  • DoobPPO uses 1578 MB GPU memory and 25.28 minutes, while DoobGRPO uses 1795 MB and 28.12 minutes; both remain comparable to several strong diffusion-policy baselines.The measurements cover 200K iterations and 1 million environment interactions.Source: Table 4 · p. 30, §G.4 · p. 30

H Ablation Studies

The ablations favor lightweight ratio networks for both diffusion and flow policies, while the preferred number of generation steps differs in robustness across the two settings. The algorithms alternate ratio learning with drift matching to update the generative policy.

H.1 Diffusion Policy.

  • On HalfCheetah-v4, the 3×256 MLP achieves the best DoobPPO performance, and T = 20 diffusion steps outperform T = 10 and T = 30.The deeper MLP and larger ResNet underperform the smaller architecture; T = 30 substantially degrades performance.Source: §H.1 · pp. 30–31, Figure 3 · p. 32

H.2 Flow Policy.

  • For DoobGRPO flow policies, the 3×256 MLP reaches a peak return of 6433, while the deeper MLP reaches 23 and the ResNet reaches −94.T = 20 has the best peak performance at 6433; T = 10 (6120) and T = 30 (6282) remain competitive.Source: §H.2 · p. 31, Figure 4 · p. 33

I.1 Diffusion Policy Optimization with DoobPO

  • The diffusion algorithm collects rollouts, estimates advantages, learns the ratio network, then updates the drift model using conditional drift matching.For the drift-matching update, it samples a time and generates a noised action from each replay-buffer state-action pair.Source: §I.1 · p. 34

I.2 Flow Policy Optimization with DoobPO

  • The flow algorithm likewise alternates rollout collection, advantage estimation, ratio learning, and a drift-matching update to the velocity model.Its intermediate actions are generated from replay-buffer state-action pairs before minimizing the conditional drift-matching objective.Source: §I.2 · p. 35

Related papers

Loading 2610.28990v1…