POGP — Prefix-Optimal Dynamic Diffusion Policies

Abstract

Diffusion policies build an action through a K-step denoising chain, but conventional training only ever supervises the final step. POGP fixes this by giving every denoising step its own value estimate: a prefix value function trained with a Bellman-style recursion over the chain itself. During training, this auxiliary signal shapes intermediate steps into actions that are already useful. At test time, the change in value between one step and the next becomes a stopping rule, yielding adaptive early stopping without a tuned threshold or a separate controller. Across four MuJoCo control tasks and eleven baselines, POGP improves on the strongest dynamic-diffusion baseline by roughly 3.5% while cutting the number of denoising steps by about 2.7×, and prefix training alone lifts full-chain performance by close to 4.6%.

Adaptive Compute in Action

A single HalfCheetah perturbation demo clearly illustrates POGP’s adaptive behavior.

Single qualitative demo

HalfCheetah · perturbation and recovery

MuJoCo physics

HalfCheetah-v4 · external disturbance

Step 0 · steady

Animated HalfCheetah perturbation and recovery rollout

Steady locomotion low compute

Perturbation extra refinement

Recovery return to low compute

Demo timing Steady locomotion

The timing bar and force overlay make the disturbance window explicit: the perturbation is active from steps 55–110. The effort panel summarizes the paper's 100 time-aligned HalfCheetah evaluation rollouts.

Adaptive denoising effort

POGP spends compute when the state becomes difficult

Reported trend: POGP uses roughly 3–5 denoising steps during steady locomotion, increases to roughly 15–20 under perturbation, then returns to its low-compute regime. D3P maintains a higher routine budget and reacts less strongly. The curve above is a clean schematic of that reported trend; the paper contains the exact 100-rollout plot.

3–5

routine POGP steps

15–20

steps under perturbation

≈2.7×

overall compute reduction

Method

POGP attaches a value estimate to every denoising prefix, which improves intermediate actions during training and gives a learned early-stopping rule at test time.

Each denoising iterate receives a prefix-value estimate. At inference time, the chain stops when the next refinement is predicted to add no value.

Why the demo matters

The single demo is intentionally focused on the behavior that is easiest to see and directly supports the adaptive-compute claim.

During training, prefix supervision improves intermediate iterates rather than only the final denoised action.
During inference, POGP can stop early once the expected value gain from another denoising step is negligible.
During perturbation, the adaptive step budget temporarily rises, reflecting the need for extra refinement under more difficult control conditions.
After recovery, the budget falls back toward its efficient operating regime, preserving nearly the full return while using fewer denoising steps.

+3.5%

over D3P

+4.6%

prefix training lift

98–99%

full-chain return retained

18.2%

fewer iterations than D3P

Results

POGP improves aggregate full-chain performance while retaining nearly all of its return under adaptive stopping.

Aggregate comparison of full-chain performance, dynamic denoising cost, and retained return.

Method	HalfCheetah	Walker2d	Ant	Hopper	IQM
SAC	307±19	171±18	89±9	211±38	0.84
Diff-QL	316±30	180±15	95±8	213±46	0.88
SDAC	331±18	184±13	94±7	208±34	0.90
D3P	327±20	181±14	95±8	206±33	0.89
FQL	323±20	181±14	95±7	209±34	0.89
POGP	348±21	191±12	99±9	217±29	0.94

Abridged Table 1 from the paper page.

Method	HalfCheetah	Steps	Ret.	Speedup
SAC (1 pass)	307±19	1	88%	20×
Diff-QL (full)	316±30	20	91%	1.0×
D3P (dynamic)	327±20	8.4±2.6	94%	2.4×
POGP (dynamic)	345±18	6.8±2.1	99%	2.9×

Adaptive stopping preserves nearly full return while substantially reducing the denoising budget.

Learning When to Stop: Prefix-Optimal Dynamic Diffusion Policies for Continuous Control

Every denoising prefix is trained