Method
The challenge: policy updates need action ratios
PPO and GRPO use the action ratio \(r_\theta(o,a)=\pi_\theta(a\mid o)/\pi_{\mathrm{old}}(a\mid o)\). A diffusion or flow policy generates an action through a trajectory, so its exact action likelihood is generally intractable. DoobPO separates the update rule from the generative model: first specify the desired change in the action distribution, then learn a generative process that realizes it.
1. An update rule defines a policy tilt
Different RL updates express different preferences over actions. A KL-proximal policy mirror descent update has an analytic target ratio proportional to \(\exp(\widehat A(o,a)/\lambda)\). PPO and GRPO do not generally yield a closed-form target ratio, so DoobPO learns a lightweight positive ratio model \(r_\beta(o,a)\) with their respective RL objectives. A regularizer encourages \(\mathbb{E}_{a\sim\pi_{\mathrm{old}}(\cdot\mid o)}[r_\beta(o,a)]=1\). Both updates start from the same old policy; their learned ratios specify different target action distributions.
For PPO and GRPO, the positive ratio is learned with the chosen update objective; for PMD, it has an analytic form. Once the ratio is available, every rule follows the same distribution-level step: reweight the old policy and normalize. The negative log ratio is the endpoint energy that describes this tilt.
\[ \begin{aligned} \pi_{\mathrm{target}}(a\mid o) &=\frac{r(o,a)\,\pi_{\mathrm{old}}(a\mid o)}{Z(o)},\\ &\propto e^{-E(o,a)}\pi_{\mathrm{old}}(a\mid o),\\ E(o,a)&=-\log r(o,a). \end{aligned} \]
\(Z(o)=\mathbb E_{a\sim\pi_{\mathrm{old}}(\cdot\mid o)}[r(o,a)]\) normalizes the tilt; the update rule sets its action-dependent energy.
2. Doob turns the endpoint tilt into a generative path
For diffusion policies, Doob's h-transform propagates the endpoint preference through the pretrained denoising process. The resulting process has the ratio-tilted action distribution as its endpoint (Lemma 3.1). The diffusion drift correction changes intermediate distributions during generation, reaching the chosen action law. For flow policies, terminal reweighting defines the corresponding tilted marginals, learned through weighted velocity matching.
The action-endpoint ratio induces an intermediate weight \(h_t\): its expected value conditioned on the current state of the reference diffusion. The spatial log gradient of this weight supplies the added drift shown in the diagram. The next step trains that drift without sampling the ideal corrected path.
\[ \begin{aligned} h_t(x,o)&=\mathbb{E}\!\left[r(o,X_0)\mid X_t=x,o\right],\\ \Delta b_t(x,o)&=-\sigma_t^2\nabla_x\log h_t(x,o). \end{aligned} \]
\(\Delta b_t\) is the extra diffusion drift. \(X_0\) is the action endpoint and \(X_T\) the noise endpoint in the paper's diffusion time convention.
3. Conditional drift matching makes it trainable
The ideal tilted path is unavailable for direct training. DoobPO instead weights conditional score or velocity regression by the endpoint ratio, using samples from the old policy:
\[ \mathcal L_{\mathrm{CDM}}(\theta) =\mathbb E_{\mathrm{old}} \!\left[r(o,a)\,\|\psi_\theta-\psi_{\mathrm{cond}}\|_2^2\right]. \]
\(\psi_\theta=\psi_\theta(x_t,t,o)\); \(\psi_{\mathrm{cond}}=\psi_{\mathrm{cond}}(x_t,t,o,a)\) is a conditional score target for diffusion or a velocity target for flow.
At each fixed observation, the conditional objective's gradient is a positive multiple of the ideal drift-matching gradient for both diffusion and flow (Theorems 3.1 and 3.2). The implementation uses the Gaussian diffusion or affine flow regression targets described in the paper.
Why it works: a minimum-energy diffusion update
Among controlled reverse diffusions with the specified noise and action marginals, the Doob-transformed process minimizes quadratic control energy:
\[ \begin{aligned} \inf_u\quad &\frac{1}{2}\, \mathbb E_{Q^u}\!\int_0^T \|\sigma_t u_t(Y_t,o)\|_2^2\,\mathrm dt,\\ &Y_T\sim p_T^h,\quad Y_0\sim p_0^h. \end{aligned} \]
Proposition 3.1 is stated for each observation and assumes a realizable model class and an attained minimum. This optimal-control guarantee concerns diffusion; the flow result establishes gradient equivalence.
One framework, several update rules
The generative policy can be diffusion or flow, while the update rule determines the tilt. Analytic ratios lead directly to weighted matching; PPO and GRPO use learned ratios. This view includes existing methods as special cases and yields DoobPPO, DoobGRPO, and DoobPMD.
| Method | Generative policy | Policy update | Ratio |
|---|---|---|---|
| DPMD | Diffusion | Policy mirror descent | Analytic |
| SDAC | Diffusion | Soft actor-critic | Analytic |
| DoobPMD (ours) | Flow | Policy mirror descent | Analytic |
| DoobPPO (ours) | Diffusion / Flow | PPO | Learned |
| DoobGRPO (ours) | Diffusion / Flow | GRPO | Learned |
