ArXiv

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

Authors
Rohit Kumar Salla, Manoj Saravanan, Simon Stepputtis
Categories
cs.LG, cs.RO
arXiv
https://arxiv.org/abs/2608.05084v1
PDF
https://arxiv.org/pdf/2608.05084v1

Brief

Prefix-Optimal Generative Policies (POGP) addresses the computational cost of diffusion policies for continuous control by learning a prefix value function at each denoising step via a Bellman-style recursion. The prefix value both trains intermediate outputs to be higher-quality actions and yields a test-time stopping rule. In four MuJoCo tasks (vs 12 baselines) POGP cuts iterations ~2.7× and yields ~3.5% higher final performance versus SOTA dynamic diffusion; full text was not available (abstract-only).

Why it matters

Prefix-Optimal Generative Policies (POGP) trains a prefix value function at every intermediate denoising step using a Bellman-style recursion to both supervise intermediate outputs and enable test-time early stopping (authors: Rohit Kumar Salla, Manoj Saravanan, Simon Stepputtis; posted 2026-08-05).

Key details

  • On four MuJoCo environments versus 12 baselines, POGP reduces required denoising iterations by ~2.7× while retaining near-full task performance and improves final task performance by ~3.5% compared with state-of-the-art dynamic diffusion baselines.
Source evidence

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.