Environment and Baseline Configuration
Environments
We leverage four distinct environments from the MuJoCo benchmark: HalfCheetah-v4, Walker2d-v4, Hopper-v4, and Ant-v4, spanning action dimensions from 3 to 8. These environments were chosen because they involve complex control of embodied agents while providing fast simulation through the MuJoCo implementation.
POGP Baseline Performance on Full Diffusion
As our first experiment, we are interested in the performance of POGP as a standalone diffusion policy without the proposed PVF-based dynamic diffusion process. We assess the performance of POGP with the full number of diffusion steps, making it equivalent to the other baseline methods tested.
Conclusion
POGP is an approach for learning when to stop computing further denoising iterations in diffusion policies. POGP achieves this by learning a prefix value function (PVF) through a Bellman-style recursion over the denoising chain. The PVF serves two purposes:
- It provides an auxiliary optimization target that shapes intermediate diffusion steps into executable actions.
- It offers a test-time stopping criterion that halts denoising when further iterations are not expected to yield meaningful improvement.
Across four MuJoCo tasks and 12 baselines, POGP reduces the diffusion compute budget by approximately 66% while retaining near-full performance. It outperforms state-of-the-art baselines by approximately 4.6% in aggregate interquartile mean (IQM).
The authors acknowledge Advanced Research Computing at Virginia Tech for providing computational resources and technical support that contributed to the reported results.
Multi-Hazard Training
A natural extension is to train under a distribution of hazards rather than a single hazard.
Network Architecture
All networks use the same backbone for fair comparison.
Environment Critic
The environment critic uses an architecture identical to standard SAC and TD3 critics: an MLP with 256 units and ReLU activations. Its input is the environment state and action representation. It uses two independent heads for double-Q value estimation.
Prefix Critic
The prefix critic is an MLP with 256 units and ELU activations. Its input is the environment state, partially denoised action, and diffusion timestep, where the diffusion timestep is encoded identically to the denoiser timestep embedding. The output is a scalar.
The completion map is deliberately shallow: its purpose is projection, not re-refinement. A deeper map would add parameters that could implicitly learn to refine, confounding the ablation. Increasing the completion map depth to two layers does not meaningfully change results, with a small retention difference.
Baseline Implementations
All baselines are reimplemented from scratch in the shared codebase to ensure identical data handling, replay buffer management, and evaluation protocols. Where baselines have published official code, we verified that our implementations match their reported numbers within 2% on a common reference task: HalfCheetah at 1 million steps.
- TD3: Twin-delayed DDPG. We use the standard hyperparameters and policy smoothing with the policy-smoothing noise scale, clipping the smoothing noise.
- SAC: Soft actor-critic with automatic entropy tuning.
- PPO: Proximal policy optimization with 64 parallel environments, clipping the policy-ratio objective, and 10 epochs per update.
- Diffusion-QL: Diffusion policy with a timestep-weighted behavioral cloning loss. This is an online variant with no offline pretraining.
- DPPO / D²PPO: Diffusion policy optimized with PPO-style updates. D²PPO adds a second denoising stage.
- SDAC: Score-matching diffusion actor-critic. It uses a score-matching auxiliary loss to avoid backpropagating through the full chain.
- DSAC-D: Diffusion extension of soft actor-critic with direct score-based policy updates.
- D3P: Diffusion policy with a learned depth-selection prior and dynamic early exit.
- IDQL: Offline implicit diffusion Q-learning with 100,000 offline pretraining steps from a random-policy replay buffer, followed by 900,000 online fine-tuning steps.
- FQL: Flow Q-learning. The flow-based policy is trained with a reflow objective to enable one-step action generation.
- SAC-GMM: SAC with a five-component Gaussian mixture policy. It serves as a capacity-matched multimodal baseline without iterative refinement.
Computational Resources
All experiments were run on a single NVIDIA A100 GPU with 80 GB of memory.
The total compute was as follows:
- Wall-clock time per 1-million-step run was approximately 1.0–1.2 hours for POGP, slightly higher than for baselines because of shared chain generation.
- Wall-clock time per 1-million-step run was approximately 0.7–0.9 hours for non-diffusion baselines.
High-Dimensional Action Spaces
The per-environment attribution analysis shows smaller gains for Ant (the Ant action dimension) than for HalfCheetah (the HalfCheetah action dimension). We hypothesize two contributing factors.
Noisier Gradients
Higher-dimensional action spaces yield greater critic-estimation variance, reducing the reliability of Hazard Bellman targets. Q-gradient reliability was 0.58 for Ant compared with 0.82 for HalfCheetah.
A deeper completion map or an ensemble V-critic might address both issues, but we leave this to future work to avoid conflating the architecture with the objective.
Hazard Value
Table 11 reports full-chain return and retention at the specified diffusion depth for uniform hazard values across the evaluated hazard range.
Discussion
Three trends emerge from Table 14:
- Full-chain return improves monotonically with the number of diffusion steps for both methods.
- Longer chains provide more expressive policies.
- POGP consistently outperforms SDAC in absolute return at every evaluated chain length, with a roughly constant gap of 28–30 points on HalfCheetah.
The prefix training objective does not harm full-chain quality.
Diffusion Models for Offline Reinforcement Learning and Planning
The use of diffusion models in reinforcement learning predates online training formulations. Diffusion models have been used to model entire trajectory distributions, enabling planning by sampling and reranking trajectories. Diffusion-QL introduced the first use of a diffusion actor with a timestep-weighted fine-tuning objective in offline settings.
These offline methods implicitly assume that inference always completes. POGP’s prefix training extends their objectives to the online, variable-compute setting.
Relationship to Options and Temporal Abstraction
The options framework decomposes environment-time policies into temporally extended behaviors with initiation conditions and termination functions. POGP’s refinement MDP shares the termination-over-time structure but operates in refinement time rather than environment time.
A formal unification in which an option’s internal computation is itself a prefix-optimal process could extend POGP to hierarchical settings.
Relationship to Adaptive Computation Time
Adaptive Computation Time (ACT) and PonderNet learn when to halt iterative computation in supervised settings. POGP provides the reinforcement-learning analogue: the stopping rule emerges naturally from the learned prefix value function, without a separate halting network.
Unlike ACT, POGP’s stopping criterion is grounded in expected return rather than a cost that penalizes computation.
Frequently Asked Questions
What is POGP? POGP is an approach for learning when to stop denoising iterations in diffusion policies for continuous control.
What is the purpose of the prefix value function? It shapes intermediate diffusion steps into executable actions and provides a test-time criterion for stopping denoising.
How much diffusion computation does POGP reduce? Across four MuJoCo tasks and 12 baselines, POGP reduces the diffusion compute budget by approximately 66% while retaining near-full performance.
Which environments were used for evaluation? The evaluation used HalfCheetah-v4, Walker2d-v4, Hopper-v4, and Ant-v4 from the MuJoCo benchmark.