DiPOD: Diffusion Policy Optimization without Drifting Apart

TL;DR

DiPOD diagnoses why existing diffusion policy-gradient methods are unstable as a “double-drift” phenomenon: optimizing the standard variational ELBO surrogate lets the surrogate detach from the true log-likelihood, which in turn makes the proxy policy gradient misalign with the true policy gradient of expected return. The fix is structurally simple — interleave a self-distillation step with each policy-improving update, which amounts to adding an on-policy ELBO regularizer to the diffusion policy-gradient loss. The recipe is reported to stabilize training and reach higher rewards than prior methods across both diffusion language model post-training and continuous-control diffusion policies, making it one of the first filed papers to treat dLLM-RL and robotics diffusion-policy RL as instances of the same underlying problem.

Key claims

Existing diffusion policy-gradient methods are unstable because of a “double-drift” phenomenon: optimizing a variational surrogate (the ELBO) lets the surrogate separate from the true log-likelihood, which then misaligns the resulting proxy policy gradient from the true policy gradient of expected return [Abstract, §1].
Maintaining a tight ELBO bound throughout training restores alignment between the surrogate and true policy gradients [Abstract, §3].
A simple practical algorithm — augmenting each diffusion policy-gradient update with an on-policy ELBO regularizer (equivalently, interleaving self-distillation with policy-improving updates) — achieves the tight-bound condition [Abstract, §3].
The recipe applies uniformly to both diffusion language model post-training and continuous-control diffusion policies, suggesting the diagnosis is a property of the diffusion-policy class rather than a domain-specific issue [Abstract, §4].
DiPOD substantially stabilizes training and reaches higher rewards than prior diffusion policy-gradient methods across both settings [Abstract, §4].

Method

DiPOD targets a structural mismatch in policy-gradient methods built on top of diffusion policies. When the log-likelihood log p_θ(x) is replaced with the variational ELBO L_ELBO(x; θ) ≤ log p_θ(x) for tractability — which is unavoidable in dLLMs and continuous-control diffusion policies alike — the gap log p_θ(x) − L_ELBO(x; θ) is itself θ-dependent, so the gradient of the surrogate is not the gradient of the true objective. Under policy improvement the surrogate gap can grow without bound (the “drifting apart” in the title), producing a misaligned proxy gradient that drives unstable updates.

The algorithm is a one-line modification of any diffusion policy-gradient update: at each step, in addition to the reward-maximization gradient, add an on-policy ELBO regularizer evaluated on rollouts from the current policy. Mechanically this is equivalent to interleaving a self-distillation step (tighten the ELBO of p_θ against itself on freshly sampled trajectories) with each policy-improving step. The two terms compete: the policy-gradient term pulls the policy toward high-reward regions while the self-distillation term keeps the ELBO surrogate tight to the true log-likelihood at exactly those regions, so the proxy gradient stays aligned with the true policy gradient throughout training.

Results

DiPOD is reported to substantially stabilize training and reach higher rewards than previous diffusion policy-gradient methods across two regimes: diffusion language model post-training and continuous-control diffusion policies [Abstract]. The paper is short (8 pages of body in the v2 preprint) and headline numbers per benchmark/task are not in the abstract; this page will be refreshed when those are extracted.

Why it’s interesting

DiPOD bridges two lineages on the wiki that have so far been studied separately. Stabilizing Reinforcement Learning for Diffusion Language Models (StableDRL) diagnoses dLLM-RL instability at the ratio-estimation layer — GRPO’s importance ratio is ELBO-estimated and develops heavy-tailed noise, producing the self-reinforcing spike loop that StableDRL fixes with unconditional clipping and convex-hull normalization. DiPOD diagnoses the same instability one layer up, at the objective-surrogate layer — the ELBO itself drifts away from log p_θ under policy improvement, so even before noisy ratio estimation enters the picture the gradient is misaligned. The two diagnoses are complementary rather than competing: StableDRL constrains how the noisy ratio is used in each update; DiPOD constrains how the ELBO surrogate evolves across updates. A natural ablation no filed paper has run is to stack them.

The unification with continuous-control diffusion policies is the second reason this matters. The wiki’s existing flow/diffusion-RL recipes — DanceGRPO (DanceGRPO: Unleashing GRPO on Visual Generation), DDRL (Data-regularized Reinforcement Learning for Diffusion Models at Scale), FPO (Flow Matching Policy Gradients), QGF (Test-Time Gradient Guidance of Flow Policies in Reinforcement Learning) — all attack pieces of the “policy gradient through a diffusion sampler is unstable” problem, but each is anchored to a specific domain (visual generation, Cosmos-scale T2V, continuous-control). DiPOD’s “self-distillation as on-policy ELBO regularizer” framing claims the same primitive works for dLLM math benchmarks and continuous-control diffusion policies, which would make it the first filed paper to argue the dLLM-RL and robotics diffusion-policy-RL communities are solving the same numerical problem. Worth a closer read once a v3 or follow-up reports headline numbers separable from the StableDRL / DDRL baselines.