Learning When to Stop: Prefix-Optimal Dynamic Diffusion Policies for Continuous Control
Rohit Kumar Salla, Manoj Saravanan, Simon Stepputtis
cs.LG, cs.RO
Submitted: 2026-08-05
Project page: https://rohitsalla.github.io/POGP
License: http://creativecommons.org/licenses/by/4.0/
The gist: Diffusion policies are a powerful policy class for continuous control, but their iterative denoising process creates a substantial computational bottleneck.
Terminology
Abstract
Diffusion policies are a powerful policy class for continuous control, but their iterative denoising process creates a substantial computational bottleneck. Reducing this cost requires adapting the number of denoising steps to the difficulty of each action while preserving task performance. We introduce Prefix-Optimal Generative Policies (POGP), a framework that learns a prefix value function at every intermediate denoising step through a Bellman-style recursion over the denoising chain. The prefix value function serves two purposes: it provides an auxiliary training objective that encourages intermediate outputs to become high-quality actions, and it enables a test-time stopping rule that terminates denoising when additional steps are unlikely to produce meaningful improvement. Across four MuJoCo environments and comparisons with 12 baselines, POGP reduces the required number of denoising iterations by approximately 2.7-fold while retaining near-full task performance. Compared with state-of-the-art dynamic diffusion baselines, prefix training also improves final task performance by approximately 3.5%. These results indicate that supervising intermediate denoising steps is useful not only for adaptive early stopping, but also as an auxiliary objective that improves the learned policy.
Sources
- Deep Reinforcement Learning at the Edge of the Statistical Precipice
- PonderNet: Learning to Ponder
- Diffusion Policy: Visuomotor Policy Learning via Action Diffusion
- Adaptive Computation Time for Recurrent Neural Networks
- Planning with Diffusion for Flexible Behavior Synthesis
- Distributional Soft Actor-Critic with Diffusion Policy
- Efficient Online Reinforcement Learning for Diffusion Policy
- Robot Skill Adaptation via Soft Actor-Critic Gaussian Mixture Models
- Diffusion Policy Policy Optimization
- Progressive Distillation for Fast Sampling of Diffusion Models
- Proximal Policy Optimization Algorithms
- Diffusion Actor-Critic with Entropy Regulator
- D3P: Dynamic Denoising Diffusion Policy via Reinforcement Learning
Related papers
- Polynomial-Augmented Neural Networks (PANNs) with Weak Orthogonality Constraints for Enhanced Function and PDE Approximation
- AIRL-S: Unifying Reinforcement Learning and Search-Based Test-Time Scaling via Adversarial Inverse Reinforcement Learning
- Transformers as Bayesian In-Context Experimenters: Smoothness-Adaptive Efficient ATE Estimation
- Convergence issues in Relational Concept Analysis based on AOC-posets
- Beliefs Beyond Posteriors: Local-Consistency Optimisation for Bayesian Neural Networks
- Understanding Diffusion Models via Ratio-Based Function Approximation with SignReLU Networks