arXiv CS open access

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

2026 · arxiv_cs

arXiv CS · Papers · 2026

Conceptio Open Archive

Analyze & explore related research in Conceptio

Connect this record directly into your research workflow or AI agent. Explore citations across open repositories, query full text, and export structured proof bundles.

Explore in Search →Connect via MCP / API ↗

Cover Page

Reinforcement Learning Journal 2026

Learning When to Stop: Prefix-Optimal Dynamic Diffusion Policies for Continuous Control Rohit Kumar Salla, Manoj Saravanan, Simon Stepputtis Keywords: Diffusion policies, Efficient Computing, Prefix Value Functions, Continuous Control

arXiv:2608.05084v1 [cs.LG] 5 Aug 2026

Summary Diffusion policies generate actions through a K-step denoising chain, but standard training optimizes only the terminal output. We introduce POGP, which learns a prefix value function at every denoising step through a Bellman-style recursion. The prefix value function improves intermediate action quality and provides a test-time signal for adaptive stopping. Across four MuJoCo environments and 12 baselines, POGP outperforms D3P by approximately 3.5% while using 18.2% fewer iterations. Relative to a fixed 20-step chain, POGP reduces denoising computation by approximately 2.7× while retaining near-full-chain performance.

Contribution(s) 1. Prefix Value Function: A value function trained via TD bootstrapping over the denoising chain that serves as an auxiliary objective improving intermediate action quality. Context: Diffusion policies optimize only the terminal action, leaving intermediate iterates unsupervised and non-executable. 2. Learned Early Stopping: A stopping rule requiring no separately learned halting network derived from consecutive prefix value differences that halts denoising when further refinement is predicted to yield negligible improvement. Context: Prior acceleration methods either apply fixed step reductions or train separate adapter networks decoupled from intermediate action quality. 3. Empirical Validation: Prefix training improves full-chain performance by approximately 4.6% over state-of-the-art baselines while enabling ≈ 2.7× compute savings across four MuJoCo environments and 12 baselines. Context: The dual benefit suggests intermediate supervision acts as a beneficial auxiliary objective beyond enabling adaptive compute allocation.

Prefix-Optimal Generative Policies

Learning When to Stop: Prefix-Optimal Dynamic Diffusion Policies for Continuous Control Rohit Kumar Salla, Manoj Saravanan, Simon Stepputtis rohits25@vt.edu, manoj663@vt.edu, stepputtis@vt.edu Virginia Polytechnic Institute and State University (Virginia Tech), Blacksburg, Virginia, USA

Abstract Diffusion policies have emerged as a powerful policy class for continuous control, but their iterative denoising process remains the primary computational bottleneck. Optimizing the computational effort spent per action is critical for improving the practical adoption of diffusion policies, allowing them to allocate more iterations only when the task demands it. To this end, we propose POGP (Prefix-Optimal Generative Policies), which learns a prefix value function (PVF) at every intermediate denoising step via a Bellman-style recursion over the chain. The PVF provides both an auxiliary optimization target that shapes intermediate outputs into high-quality actions and a test-time stopping criterion that halts denoising when further iterations are not expected to yield meaningful improvement. Through extensive experiments on four MuJoCo environments against 12 baselines, we demonstrate that POGP reduces the required denoising iterations by approximately 2.7× while retaining near-full performance. Furthermore, compared to current state-of-the-art dynamic diffusion baselines, we show that prefix training improves final task performance by approximately 3.5%, suggesting that supervising intermediate denoising steps acts as a beneficial auxiliary objective beyond its role in enabling early stopping. Project Page: https://rohitsalla.github.io/POGP

Introduction

Diffusion policies for continuous control generate actions through an iterative denoising chain, typically requiring K sequential neural network evaluations per action (Ho et al., 2020; Chi et al., 2024; Wang et al., 2022). Running all K iterations of the diffusion process is the primary computational bottleneck of diffusion policies as every action incurs the same computational cost regardless of the task’s difficulty. In practice, however, many actions taken by an embodied system along some trajectory are routine and well-determined after only a few denoising steps, while others, such as contactrich transitions, disturbance recovery, or mode switches strongly benefit from extended refinement steps. For example, consider an embodied agent perturbed by some external force. When the disturbance is applied to the system, the policy benefits from additional refinement steps to improve the system’s overall behavior while saving compute at other times, which is especially important for local and edge deployments on real robot hardware. Current diffusion policy frameworks (Ren et al., 2024; Ma et al., 2025; Wang et al., 2024; Hansen-Estruch et al., 2023) do not distinguish between these cases, spending compute uniformly across the trajectory. Recent approaches leveraging diffusion policies (Yu et al., 2025) address this issue by estimating the number of iterations needed in order to reach a sufficiently performant action. However, because this adapter is independent of the policy, the intermediate denoising steps are never optimized to be executable as the model is only trained on the full diffusion process. We take a different approach by jointly training a prefix value function (PVF) alongside the diffusion policy itself, which both im-

Reinforcement Learning Journal

2026

proves intermediate action quality as an auxiliary objective and provides a natural per-step stopping signal at test time. In this work, we propose a novel approach for dynamic diffusion – Prefix-Optimal Generative Policies (POGP) – that leverages a learned prefix value function (PVF) to assess when diffusion iterations can be stopped early while improving the quality of internal diffusion iterations. POGP trains a PVF Vt (s, at ) for each intermediate action at in the K-step diffusion chain aK → a0 that intuitively measures how valuable it would be to execute action at at step t. During training, the PVF is trained to estimate an action’s value by completing all K diffusion steps, co-evolving as an optimization target for the diffusion step itself. Crucially, rather than regressing each Vt independently against the final action’s Q-value, which would require expensive full-chain rollouts per step, we derive a Bellman recursion over the denoising chain that propagates value information backwards through refinement steps, enabling efficient single-chain training with stable TD targets. Intuitively, the PVF serves a dual purpose: 1) it improves the quality of the intermediate policy actions at ; and 2) provides a metric for stopping the diffusion process early in order to save time and compute budget. We demonstrate that, across four MuJoCo tasks and 12 baselines, POGP achieves an IQM of 0.942 with a 95% confidence interval (Agarwal et al., 2022), retains ≥ 97.8% of its full-diffusion process task performance at less than half the computation cost and outperforms current state-of-the-art baselines by ≈ 4.6%. In summary, POGP proposes the following contributions: • a prefix value function trained via TD bootstrapping alongside the diffusion model over the denoising chain, ultimately improving the quality of the intermediate steps; • a principled early stopping rule derived from prefix value differences across consecutive diffusion steps, leveraged to stop the diffusion policy early if the value of future steps is predicted to be minimal; and • empirical evidence that prefix training improves full-chain performance by up to 4.6% while enabling ≈ 2.7× compute savings, on average.

Related Work

Diffusion Policies for Control. Diffusion-based policies have emerged as a powerful class of decision models for continuous control and robot learning because they can represent complex, multimodal action distributions that are difficult for unimodal actors to capture. Prior work includes action-diffusion policies for manipulation (Chi et al., 2024), diffusion-based trajectory planning (Janner et al., 2022), and diffusion-parameterized reinforcement learning methods such as DiffusionQL (Wang et al., 2022), Efficient Diffusion Policies (Ma et al., 2025), and IDQL (Hansen-Estruch et al., 2023). A common feature of these methods is that action generation proceeds through a multi-step denoising chain, so inference cost scales with the number of refinement steps. Moreover, existing objectives primarily supervise the terminal denoised action rather than the intermediate prefixes of the chain. As a result, intermediate iterates are generally not trained to be high-value executable decisions. This creates a train–test mismatch whenever inference must be interrupted or truncated under real-time compute constraints. POGP addresses this gap by explicitly assigning value to every prefix of the refinement process through a learned prefix value function thereby turning variable diffusion depth into a first-class object of policy optimization rather than a post hoc deployment heuristic. Anytime and Adaptive Computation. Our work is also related to the broad area on anytime prediction and adaptive computation. It studies systems whose outputs remain useful when computation is stopped early and improve as additional computation is allocated Zilberstein (1996); Graves (2017); Huang et al. (2016); Banino et al. (2021). A central challenge in this setting is that models are often optimized only for full-depth performance, which often yields poor intermediate outputs. While exposing the model to variable computation during training, we can improve

Prefix-Optimal Generative Policies

robustness across budgets and sometimes regularize full-budget performance. POGP brings this perspective into continuous control, where early outputs are not merely intermediate to representations but executable actions whose quality must be judged by downstream return. This distinction is important because unlike adaptive-computation settings in supervised learning, where intermediate predictions are evaluated by task loss. Here, each refinement prefix induces a policy with its own control performance. Our formulation therefore connects adaptive computation to reinforcement learning through value-based supervision over diffusion prefixes. Efficient and Dynamic Diffusion. Recent efforts to reduce the cost of diffusion inference broadly fall into two categories. The first seeks to improve a predetermined operating point by reducing the number of denoising steps required for good performance. For example through distillation, consistency models or consistency-style policy classes Salimans & Ho (2022); Song et al. (2023); Ding & Jin (2024). These approaches substantially reduce latency, but they primarily optimize a chosen small-step generator rather than ensuring that arbitrary prefixes of the original diffusion chain are themselves high-value actions. The second line considers adaptive-budget allocation, where the amount of denoising performed is adjusted online. Most closely related is D3P, which introduces a lightweight adaptor to allocate denoising depth at test time Yu et al. (2025). POGP differs in both objective and mechanism: instead of learning a separate stopping module, it trains a prefix value function that directly evaluates the marginal utility of further refinement. The same value-grounded object is used both to improve prefix quality during learning and to determine when additional denoising is no longer worthwhile at inference time.

Background and Notation

Notation We summarize key notation as follows: s ∈ S and at ∈ A represent the environment state and action at refinement level t, respectively, where t ∈ {0, . . . , K} runs backwards (t=K is noise, t=0 is the terminal action). The symbol a0 denotes the fully-denoised action. Markov Decision Process We consider discounted continuous-control MDPs M = ⟨S, A, P, r, γ⟩ with S ⊆ Rds , A ⊆ Rda , transition kernel P (· | s, a), bounded reward r : S × A → [Rmin , Rmax ] and discount γ ∈ (0, 1). For a (stationary) policy π, the state-action value is "∞ # X π k Q (s, a) := Eπ γ r(sk , ak ) s0 = s, a0 = a , (1) k=0 π

and V (s) := Ea∼π(·|s) [Q (s, a)]. Diffusion Policies A diffusion policy represents π(a | s) implicitly via a K-step reverse denoising chain (Ho et al., 2020; Wang et al., 2022): sample aK ∼ N (0, I) and apply the denoiser at−1 = πθ (s, at , t),

t = K, K − 1, . . . , 1,

(2)

to obtain a0 . In this work, we use a deterministic denoiser (DDIM-style) throughout all experiments. The framework extends to stochastic denoisers by replacing the deterministic refinement transitions with conditional distributions, which we leave as future work. Prior diffusion RL. Diffusion-QL (Wang et al., 2022) and IDQL (Hansen-Estruch et al., 2023) establish diffusion actors as an expressive policy class in offline (and offline-to-online) RL, while recent on-policy and actor-critic variants scale diffusion policies to fully online training (Ren et al., 2024; Wang et al., 2024; Ding et al., 2024). A shared limitation across these objectives is their terminal-only optimization in which they evaluate the environment through the terminal action a0 . This causes intermediate states to transition in an unsupervised way from noise to the terminal action, but introduce the drawback that internal states are not suitable for direct execution. However, being able to utilize non-complete diffusion chains is critical for dynamic-budget diffusion policies to allow for early stopping of the iterative generation process.

Reinforcement Learning Journal

2026

Algorithm 1 POGP (Prefix-Optimal Generative Policies) 1: Input: Replay buffer D, chain length K = 20, uniform hazard rate h = 1/K = 0.05, networks

πθ , Qϕ1,2 , Vψ , gω .

2: Initialize: Target networks ϕ′1,2 , ψ ′ . 3: for each gradient step do 4: Collect transition (s, a, r, s′ ); add to D. 5: Sample minibatch from D. 6: Update Q: yQ ← r + γ mini Qϕ′i (s′ , a′0 ); minimize LQ . 7: Sample t ∼ Unif{1, . . . , K}, aK ∼ N (0, I). 8: Generate chain: aK → · · · → at → · · · → a0 .

▷ Critic update

▷ Shared chain Update V : yV ← h mini Qϕ′i (s, a0 ) + (1 − h) Vψ′ (s, at−1 , t−1); minimize LV . ▷ Prefix Bellman target 10: Update π: Minimize Lπ (Eq. 9) with gate w(n). ▷ Policy update 11: Sample T ∼ Unif{1, . . . , K}; update g: minimize Lg (Eq. 12) using aT , a0 from the shared chain. ▷ Projection update 12: Polyak-update target networks ϕ′1,2 , ψ ′ . 13: end for 14: Return: Trained networks πθ , Qϕ1,2 , Vψ , gω . 9:

Prefix-Optimal Generative Policies

Diffusion policies generate actions through a K-step denoising chain that transforms random noise into the terminal output action (aK → · · · → a0 ). Standard training optimizes only the final action a0 , leaving intermediate actions {at }K−1 t=1 without supervision, causing them to be non-executable. a necessary requirement for dynamically choosing the number of refinement steps. POGP addresses this by learning a prefix value function Vt (s, at ) at every denoising step, providing both a training signal that co-trains with the diffusion policy to shape intermediate outputs to be of high quality, as well as a runtime metric that indicates when further denoising is no longer expected to yield significant improvement. Learning the prefix value function requires two core components: 1) a formal definition of prefix values grounded in the environment’s Q-function; and 2) a Bellman-style recursion over the denoising chain that enables efficient TD-based training of Vt from a single shared chain per gradient step. An overview of our method is shown in Algorithm 1. 4.1

The Prefix Value Function

We consider iterative refinement policies defined by a deterministic denoiser πθ : at−1 = πθ (s, at , t),

aK ∼ N (0, I),

t = K, . . . , 1,

(3)

conditioned on environment state s. To optimize intermediate steps, we require a scalar measure of how “promising” a partial result at is. The natural criterion is the expected return of completing the chain from at : Definition 1 (Prefix Value). For environment state s and intermediate action at at refinement step t: Vtπ (s, at ) := Qπ s, πfull (s, at ) , (4) where πfull (s, at ) denotes the result of running the denoiser from step t to 0. The prefix value Vt answers: “If we complete the denoising chain starting from at , what return do we expect?” High Vt indicates the chain is on track to produce a high-Q action; low Vt signals that continued refinement from a different starting point would have been preferable. Note that for a deterministic denoiser, πfull (s, at ) = a0 for any t on the same chain, so Vtπ (s, at ) = Qπ (s, a0 ). The value of learning Vt with a function approximator Vψ (s, at , t) is that the network observes only

Prefix-Optimal Generative Policies

the intermediate action at not the full chain and must therefore learn to predict the quality of the terminal action that would result from completing the chain from at . This enables two capabilities: (1) training the denoiser to maintain high predicted value across all t and (2) deployment-time monitoring via Vψ to detect when further refinement yields diminishing returns. Evaluating Eq. (4) naïvely requires completing the chain (t forward passes) for every prefix at every gradient step, multiplying training cost by O(K 2 ). We overcome this by recognizing that refinement possesses temporal structure: Vt and Vt−1 are not independent quantities but related through the denoising dynamics. This structure permits efficient credit assignment via dynamic programming. 4.1.1

Prefix Bellman Recursion

Because consecutive prefix values are related through the denoising dynamics (Vt depends on at−1 = πθ (s, at , t), which determines Vt−1 ), we can propagate value information backward through the chain via a one-step recursion rather than evaluating each Vt independently. Concretely, the prefix value at step t decomposes into a weighted combination of the final action’s Q-value and the next step’s bootstrapped prefix value: Vt (s, at ) =

h Q(s, a0 ) | {z }

terminal-value anchor

+ (1 − h) Vt−1 (s, πθ (s, at , t)), | {z }

(5)

bootstrap from the next refinement step

We use the boundary condition V0 (s, a0 ) = Q(s, a0 ). The hazard coefficient h ∈ (0, 1] controls the balance between anchoring each prefix target to the terminal action value and bootstrapping from the subsequent refinement step. Larger values of h place greater weight on the terminal-value anchor, whereas smaller values rely more heavily on the bootstrapped prefix estimate. Unrolling Eq. (5) expresses Vt as a weighted combination of the terminal-value anchor and recursively bootstrapped prefix values. When prefix training succeeds, intermediate actions approach the quality of the fully denoised output, Q(s, at ) ≈ Q(s, a0 ), and the predicted prefix values become approximately stable across consecutive refinement steps. This stability provides the signal used 1 for adaptive stopping in Section 4.2. We use a uniform hazard with K = 20 and h = K = 0.05 throughout all experiments. 4.1.2

Joint Optimization with Gated Objectives

POGP trains three components jointly using the Hazard Bellman recursion. The challenge is balancing terminal-action performance (the standard diffusion-RL objective) with prefix robustness (our proposed auxiliary objective) without destabilizing training. Standard double Q-learning on environment transitions (s, a, r, s′ ): X 2 LQ = E Qϕi (s, a) − (r + γ min Qϕ′j (s′ , a′0 )) . (6)

Environment Critic (Q).

i=1,2

Prefix Critic (V ). We sample t ∼ Unif{1, . . . , K} and generate a shared chain aK → · · · → a0 (critical for consistent gradients). The TD target implements the Hazard Bellman recursion: yV = h min Qϕ′i (s, a0 ) + (1−h) Vψ′ (s, at−1 , t−1),

(7)

2 LV = Es,t Vψ (s, at , t) − yV .

(8)

Note that yV uses the terminal action a0 from the same shared chain regardless of t, enabling efficient computation without K separate rollouts. By the determinism of the denoiser, a0 = πfull (s, at ), so the anchor Q(s, a0 ) in Eq. (7) corresponds exactly to the prefix value definition (Eq. 4).

Reinforcement Learning Journal

2026

Actor (π). The policy faces a multi-objective optimization: maximize terminal action value and maintain high prefix values. We implement this via a gated loss: h i (9) Lπ = −E min Qϕi (s, a0 ) + w(n) −E Vψ (s, at−1 , t−1) . i

The gating schedule w(n) is essential: it remains at 0 until iteration Nwarm (when the Vψ RMSE stabilizes below a threshold τ ), then ramps via a cosine schedule to wmax . This prevents earlystage noise in Vψ gradients from distorting the policy before meaningful prefix values are available. Both a0 and at−1 in Eq. (9) derive from the same shared chain, which is necessary for gradient alignment (Appendix B). Removing the confidence gate reduces HalfCheetah return from 347.2 to 337.0 and Ret.@5 from 99.1% to 95.4%. Removing the cosine warmup yields 342.5 return and 97.2% retention, supporting the use of delayed and gradual prefix supervision. 4.2

Dynamic Diffusion at Test Time

The trained prefix value function directly provides a stopping criterion at test time. The marginal gain of an additional denoising step is: ∆t = Vψ (s, at−1 , t−1) − Vψ (s, at , t).

(10)

When ∆t is small, the learned value function predicts that further refinement yields minimal improvement to the terminal action quality. We halt when: ∆t ≤ ϵ · |Vψ (s, at , t)|

for m consecutive steps.

(11)

This stopping rule requires no additional learned parameters and transfers across tasks (we use ϵ=0.01, m=2 throughout). Dynamic compute thus emerges as a free byproduct of learning prefix values, unlike prior work that requires separate learned stopping policies (Yu et al., 2025). When stopping at step t > 0, the intermediate action at may contain residual noise from the incomplete denoising process. To map such actions into the executable action space, we train a lightweight single-layer MLP gω (the completion map) to project truncated actions toward their fully-denoised counterparts: (12) Lg = ET ∼Unif{1,...,K} ∥gω (s, aT , T ) − a0 ∥2 , where aT and a0 are drawn from the same shared chain used for the other loss terms. The truncation index T is sampled uniformly during training so that gω is exposed to √ all levels of residual noise. The projection error ϵg (T ) := E[∥gω (s, aT , T ) − a0 ∥] grows as O( T ) empirically (Figure 6a), meaning early truncations (small T ) yield higher fidelity; Proposition 1 (Appendix A.2) bounds the resulting performance gap by LQ · ET ∼h [ϵg (T )]/(1 − γ). Together, this setup allows POGP to assess whether further iterations of the denoising chain are required or whether the current run can be terminated prior to reaching the full chain length K, with the completion map ensuring that the executed action remains high-quality even under early stopping.

Experiments

We assess POGP on four MuJoCo benchmarks to assess its ability to adapt to the inherent difficulty of different phases of the controlled embodied agent. In particular, we are interested in the following questions: Q1: Does prefix training improve full-chain performance (see Sec. 5.2)? Q2: Can POGP adaptively allocate compute, while retaining performance parity (see Sec. 5.3)? We further conduct an ablation study to assess the contribution of various components to our method in (see Sec. 5.4).

Prefix-Optimal Generative Policies

Walker2d

Ant

Hopper

IQM

Std.

TD3 SAC PPO

292±22 307±19 251±26

164±20 171±18 147±21

79±10 89±9 72±12

198±42 211±38 181±36

0.79 0.84 0.71

Diffusion

Diff-QL DPPO D2 PPO SDAC DSAC-D D3P

316±30 300±24 312±23 331±18 320±22 327±20

180±15 169±16 175±15 184±13 177±16 181±14

95±8 83±10 89±9 94±7 91±7 95±8

213±46 200±36 206±34 208±34 203±36 206±33

0.88 0.82 0.86 0.90 0.87 0.89

Off.

Method

5.1

HalfCheetah

IDQL

304±26

167±18

85±10

194±32

0.82

Gen.

Table 1: Full-chain performance (mean ± std, 10 seeds). All diffusion methods run K=20 steps. IQM: interquartile mean of normalised scores (Agarwal et al., 2022), 95% CI in brackets. Bold = best; underline = second.

FQL SAC-GMM

323±20 314±16

181±14 173±13

95±7 91±7

209±34 204±28

0.89 0.86

POGP

348±21

191±12

99±9

217±29 0.94 [.91, .97]

Environment and Baseline Configuration

Environments. We leverage four distinct environments from the MuJoCo benchmark, namely HalfCheetah-v4, Walker2d-v4, Hopper-v4 and Ant-v4 (Todorov et al., 2012), spanning action dimensions 3–8. These environments were chosen due to their complexity controlling embodied agents while providing fast simulation as part of the MuJoCo implementation. Baselines. We choose 12 baseline approaches, including TD3 Fujimoto et al. (2018), SAC Haarnoja et al. (2018), PPO Schulman et al. (2017), diffusion policies such as Diffusion-QL Wang et al. (2022), DPPO Ren et al. (2024), D2 PPO Kang et al. (2023), SDAC Ma et al. (2025), DSAC-D Liu et al. (2025) and D3P Yu et al. (2025) with D3P being another dynamic diffusion policy. Furthermore, we chose IDQL Hansen-Estruch et al. (2023) as an offline-to-online method as well as generative methods such as FQL Wildberger et al. (2023) and SAC-GMM Nematollahi et al. (2022). These methods were chosen as recent state-of-the-art methods across a variety of approaches. To ensure a fair comparison, all methods have a similar setup, leveraging a 3-layer MLP with 256 hidden units over a maximum of K = 20 diffusion steps. During training, we leveraged a batch-size of 256 with a 1M-step replay buffer and 10k steps of exploration. We trained 10 models with different random seeds for our main results and four additional random seeds for the ablation study. All models were trained on a HPC cluster leveraging a single Nvidia A100 GPU for each method, totaling approximately 575 single-GPU hours. Further details about the chosen hyperparameters can be found in Appendix E.2. Metrics. Each MuJoCo environment provides a dense per-step reward combining a forward-velocity bonus with a control-cost penalty (for Hopper and Walker2d - a survival bonus), so that cumulative episodic return directly measures locomotion quality. We normalise each method’s return to [0, 1] per environment using the range [random policy, best observed] and aggregate across tasks via the interquartile mean (IQM) (Agarwal et al., 2022), which discards the top and bottom 25% of runs before averaging, yielding a robust summary less sensitive to outlier seeds than the ordinary mean. 5.2

POGP Baseline Performance on Full Diffusion

As our first experiment, we are interested in the performance of POGP as a stand-alone diffusion policy without our proposed PVF-based dynamic diffusion process. Effectively, we assess the performance of POGP with the full K = 20 diffusion steps, making it equivalent to the other baseline methods tested.

Reinforcement Learning Journal

Return

300

150

200

100

HalfCheetah 50 0 5 10 0

Steps (×105)

SDAC 100 75 50 Walker2d 25 0 5 10 0

Steps (×105)

Diff-QL

200 150 100 Ant 50 0 10 0

Steps (×105)

SAC

D3P 1.00

Hopper 5

Steps (×105)

Norm. score

POGP 200

2026

0.75 0.50 0.25 0.00

IQM 0

Steps (×105)

Figure 1: Learning curves and aggregate IQM across 10 seeds. Curves show the mean across seeds, with shaded regions indicating variability. The rightmost panel reports aggregate normalized IQM. POGP’s improvement is consistent across all four environments.

Table 1 compares all methods leveraging the maximum allowed number of diffusion steps. We observe that POGP achieves the highest return in all four tested environments. With an average IQM of 0.942, providing a 95% CI lower bound (0.908), POGP exceeds every baseline’s point estimate. Three comparisons isolate the gain: POGP vs. SDAC (same backbone): +9.5 on average, POGP vs. SAC-GMM (both multimodal, no chain for MoG): +34 on HalfCheetah; and POGP vs. D3P: +5 average. Note that while D3P is a dynamic diffusion model that can stop early, for this test, we enforce all 20 diffusion steps (see Table 2 for the dynamic case). Furthermore, in Figure 1, we show the training curves of the five highest performing baseline methods across the four environments as well as their average IQM performance. In line with the Table 1 results, POGP trains slightly faster than the tested baselines, the IQM profile (rightmost panel) confirms that the gap is consistent across all environments. Together, these results demonstrate that POGP’s prefix-based approach, which optimizes intermediate diffusion steps to provide a high overall value, increases terminal action performance, validating the benefits of the proposed PVF approach. 5.3

Adaptive Compute Allocation (Q2)

The central contribution of POGP is learned efficiency: the prefix value Vψ tells the agent when additional denoising steps are worth the compute. Just as reasoning LLMs benefit from spending more tokens on hard problems (Snell et al., 2024), POGP learns to allocate more refinement steps to challenging states and fewer to routine ones grounded in expected return rather than heuristic budgets. Table 2 evaluates all methods under their own stopping rules. Non-iterative methods, such as TD3, SAC, PPO and SAC-GMM are fast, only requiring a single forward pass, but only achieve 73–98% of POGP’s performance. Diffusion- and flow-based baselines, such as Diff-QL, DPPO, D2 PPO, SDAC, DSAC-D, IDQL and FQL, close the gap, but require the full iterative refinement budget to be exerted. D3P, a recent state-of-the-art method, provides a dynamic approach allowing it to stop the diffusion process early, averaging around 9 diffusion steps. However, POGP outperforms D3P by an average of 3.5% while requiring 18.2% fewer iterations. POGP achieves the highest return and fewest steps: 6.8±2.1 on HalfCheetah (2.9× speedup), retaining 99% of its own full-chain performance while reducing the number of iterations by 66%. Given similar model capacity and compute budget between D3P and POGP, we attribute the performance improvement to the success of our PVF and early-stopping criteria. Furthermore, Figure 2 provides further evidence of the compute-budget benefit provided by POGP on the HalfCheetah environment across 100 rollouts. During steady locomotion (left sub-figure) the policy uses between ∼3–5 steps; however, when we apply a perturbation to the system, such as a constant velocity delta of -2.5m/s at step 55, POGP utilizes between ∼15–20 iterations, spending more computational budget on counteracting the perturbation, resetting to minimal compute after the perturbation is removed at step 110. While D3P shows a similar trend, its undisturbed compute budget is between ∼7–11 steps, only increasing slightly during the perturbation to ∼12–15. Overall,

Prefix-Optimal Generative Policies

Table 2: Adaptive Stopping: Performance and Compute (K=20, 10 seeds). Return: mean ± std. Steps: avg. denoising steps per action. Ret.: % of POGP full-chain return. Spd.: K / steps. Classical methods use 1 forward pass (20×). Bold = best among iterative methods. HalfCheetah Method

Ret.

Steps

Ret.

Walker2d Ret.

Steps

Ret.

Classical (single forward pass) TD3 292±22 1 84 20× SAC 307±19 1 88 20× PPO 274±28 1 79 20× SAC-GMM 311±20 1 89 20×

Spd.

164±20 171±18 152±22 173±17

1 1 1 1

Diffusion / flow (full chain) Diff-QL 316±30 20 DPPO 300±24 20 D2 PPO 308±22 20 SDAC 331±18 20 DSAC-D 325±20 20 IDQL 310±24 20 FQL 323±20 20

180±15 169±16 174±15 184±13 178±14 172±16 181±14

20 20 20 20 20 20 20

1.0× 1.0× 89 1.0× 95 1.0× 93 1.0× 89 1.0× 93 1.0× 91 86

Ant Spd.

Ret.

Hopper

Steps

Ret.

20× 79±10 20× 89±9 80 20× 72±11 91 20× 90±8

1 1 1 1

1.0× 1.0× 91 1.0× 96 1.0× 93 1.0× 90 1.0× 95 1.0×

20 20 20 20 20 20 20

94 88

95±8 83±10 88±9 94±7 92±8 86±9 95±7

Spd.

Ret.

Steps

Ret.

20× 20× 73 20× 91 20×

198±42 211±38 185±44 213±36

1 1 1 1

1.0× 1.0× 89 1.0× 95 1.0× 93 1.0× 87 1.0× 96 1.0×

213±46 200±36 207±40 208±34 205±36 202±38 209±34

20 20 20 20 20 20 20

Spd.

20× 20× 85 20× 98 20× 97

1.0× 1.0× 95 1.0× 96 1.0× 94 1.0× 93 1.0× 96 1.0× 98

Steady gait

Recovery

Perturbation

15 10

POGP D3P

5 0

Freq. (%)

Denoising steps used

Adaptive stopping (learned) D3P 327±20 8.4±2.6 94 2.4× 181±14 9.1±2.8 95 2.2× 95±8 10.2±3.1 96 2.0× 206±33 8.8±2.7 95 2.3× POGP 345±18 6.8±2.1 99 2.9× 189±12 7.4±2.3 99 2.7× 97±8 8.6±2.8 98 2.3× 215±28 7.1±2.2 99 2.8×

100

125

Environment step

150

175

200

Steps T

Figure 2: POGP adaptively allocates denoising compute: (Left) Mean denoising steps per action on HalfCheetah-v4, averaged over 100 time-aligned evaluation rollouts (10 seeds × 10 episodes; shaded regions: ±1 std. dev.). At step 55, a horizontal velocity impulse (∆v = −2.5 m/s) is applied to the torso. POGP adaptively increases to ∼15–20 steps during recovery, returning to a 3–5 step baseline once gait stabilizes. D3P shows only marginal increase during perturbation and maintains consistently higher step counts. (Right) Step distribution over 100 unperturbed rollouts: POGP mode = 4 steps vs. D3P mode = 9, confirming POGP allocates minimal compute on routine states while D3P maintains higher iteration counts.

in the undisturbed case, POGP only uses an average of 4 diffusion steps before no further improvements are anticipated. With these results, we show that POGP shows a strong reaction when more compute budget is needed while retaining a low footprint when it is not needed. Further stress tests under forced truncation and OOD hazards are in Appendix F.1, showing how more diffusion steps continuously improve performance. 5.4

Ablations (Q3)

To isolate the contribution of each design in POGP, we run ablations on HalfCheetah (4 seeds) , changing one component at a time while holding the architecture and all other settings fixed 3. A complementary 2 × 2 ablation separating prefix learning from the completion map is reported in Appendix F.4.

Reinforcement Learning Journal

2026

Table 3: Ablations (HalfCheetah, 4 seeds). Each variant changes one design decision relative to full POGP. Ret. = % of full-chain return retained when force-stopped at T =5.

Variant

Return

Ret.

Description

POGP (full)

347.2

99.1% Full method

Controls No prefix (Diff-QL) Truncate at test only Distill to K ′ =5

314.6 314.6 308.2

82.4% No PVF; standard Diff-QL backbone 71.2% No PVF; force-stop at T =5 during eval — Retrain a shorter K ′ =5 chain (no truncation)

Objective variants Prefix-supervised Prefix at t=0 only

320.1 317.3

94.8% Vψ regresses on Q(s, a0 ) directly, no bootstrap 88.1% Prefix loss applied only at terminal step t=0

Oracle (Qenv )

338.4

99.8% True Q-values as targets (upper bound)

Discussion

POGP shows that optimizing every prefix of a denoising chain yields policies that are both stronger and more compute-efficient. The Prefix Bellman recursion provides credit assignment while the prefix value Vψ provides the stopping criterion for our dynamic diffusion process. Together, they enable adaptive compute allocation with ∼4 steps for routine states and ∼18 for challenging states (HalfCheetah environment) while retaining most of the model’s overall performance. Our method, POGP, connects to a broader trend of inference time scaling: inference-time compute should be adaptive depending on the difficulty of the desired task. Just as LLM reasoning models learn to “think longer” on hard problems (Snell et al., 2024), POGP learns to spend more denoising effort on harder states. The key difference is that POGP’s allocation is grounded in a Bellman equation with provable contraction, not heuristic token budgets. The most surprising finding that good intermediate actions also improve the final action suggests the prefix loss acts as an auxiliary regulariser. Limitations. While POGP shows strong results on four MuJoCo benchmarks, additional experiments in robot manipulation tasks are left for future work, introducing a broader complexity divide between different parts of the motion. Furthermore, as POGP’s advantage grows with K (Appendix I), scaling to K=50+ chains or real-robot domains remains difficult. Investigating POGP’s ability to scale to large models and long diffusion chains remains an open question. Also see Appendix G for more details on these limitations.

Conclusion

In this work, we propose POGP, an approach for learning when to stop computing further denoising iterations in diffusion policies. POGP achieves this by learning a prefix value function (PVF) via a Bellman-style recursion over the denoising chain that serves two purposes: 1) it provides an auxiliary optimization target that shapes intermediate diffusion steps into executable actions, and (2) 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 and outperforms stateof-the-art baselines by approximately 4.6% in aggregate IQM. Acknowledgements The authors acknowledge Advanced Research Computing at Virginia Tech for providing computational resources and technical support that have contributed to the results reported within this paper. URL: https://arc.vt.edu/

Prefix-Optimal Generative Policies

References Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Marc G. Bellemare. Deep reinforcement learning at the edge of the statistical precipice, 2022. URL https:// arxiv.org/abs/2108.13264. Andrea Banino, Jan Balaguer, and Charles Blundell. Pondernet: Learning to ponder. arXiv preprint arXiv:2107.05407, 2021. Cheng Chi, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion, 2024. URL https://arxiv.org/abs/2303.04137. Shutong Ding, Ke Hu, Zhenhao Zhang, Kan Ren, Weinan Zhang, Jingyi Yu, Jingya Wang, and Ye Shi. Diffusion-based reinforcement learning via q-weighted variational policy optimization. Advances in Neural Information Processing Systems, 37:53945–53968, 2024. Zihan Ding and Chi Jin. Consistency models as a rich and efficient policy class for reinforcement learning. In International Conference on Learning Representations, 2024. URL https:// openreview.net/forum?id=v8jdwkUNXb. Scott Fujimoto, Herke Hoof, and David Meger. Addressing function approximation error in actorcritic methods. In International conference on machine learning, pp. 1587–1596. PMLR, 2018. Alex Graves. Adaptive computation time for recurrent neural networks, 2017. URL https:// arxiv.org/abs/1603.08983. Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pp. 1861–1870. Pmlr, 2018. Philippe Hansen-Estruch, Ilya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idql: Implicit q-learning as an actor-critic method with diffusion policies. 2023. Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840–6851, 2020. Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In European conference on computer vision, pp. 646–661. Springer, 2016. Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.09991, 2022. Bingyi Kang, Xiao Ma, Chao Du, Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. Advances in Neural Information Processing Systems, 36:67195– 67212, 2023. Tong Liu, Yinuo Wang, Xujie Song, Wenjun Zou, Liangfa Chen, Likun Wang, Bin Shuai, Jingliang Duan, and Shengbo Eben Li. Distributional soft actor-critic with diffusion policy, 2025. URL https://arxiv.org/abs/2507.01381. Haitong Ma, Tianyi Chen, Kai Wang, Na Li, and Bo Dai. Efficient online reinforcement learning for diffusion policy. arXiv preprint arXiv:2502.00361, 2025. Iman Nematollahi, Erick Rosete-Beas, Adrian Röfer, Tim Welschehold, Abhinav Valada, and Wolfram Burgard. Robot skill adaptation via soft actor-critic gaussian mixture models, 2022. URL https://arxiv.org/abs/2111.13129. Allen Z Ren, Justin Lidard, Lars L Ankile, Anthony Simeonov, Pulkit Agrawal, Anirudha Majumdar, Benjamin Burchfiel, Hongkai Dai, and Max Simchowitz. Diffusion policy policy optimization. arXiv preprint arXiv:2409.00588, 2024.

Reinforcement Learning Journal

2026

Tim Salimans and Jonathan Ho. Progressive distillation for fast sampling of diffusion models. arXiv preprint arXiv:2202.00512, 2022. John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017. Charlie Snell, Jaehoon Lee, Kelvin Xu, and Aviral Kumar. Scaling llm test-time compute optimally can be more effective than scaling model parameters, 2024. URL https://arxiv.org/ abs/2408.03314. Yang Song, Prafulla Dhariwal, Mark Chen, and Ilya Sutskever. Consistency models. 2023. Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2):181– 211, 1999. Emanuel Todorov, Tom Erez, and Yuval Tassa. Mujoco: A physics engine for model-based control. In 2012 IEEE/RSJ international conference on intelligent robots and systems, pp. 5026–5033. IEEE, 2012. Yinuo Wang, Likun Wang, Yuxuan Jiang, Wenjun Zou, Tong Liu, Xujie Song, Wenxuan Wang, Liming Xiao, Jiang Wu, Jingliang Duan, and Shengbo Eben Li. Diffusion actor-critic with entropy regulator, 2024. URL https://arxiv.org/abs/2405.15177. Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. 2022. Jonas Wildberger, Maximilian Dax, Simon Buchholz, Stephen Green, Jakob H Macke, and Bernhard Schölkopf. Flow matching for scalable simulation-based inference. Advances in Neural Information Processing Systems, 36:16837–16864, 2023. Shu-Ang Yu, Feng Gao, Yi Wu, Chao Yu, and Yu Wang. D3p: Dynamic denoising diffusion policy via reinforcement learning, 2025. URL https://arxiv.org/abs/2508.06804. Shlomo Zilberstein. Using anytime algorithms in intelligent systems. AI Magazine, 17(3):73, Mar. 1996. DOI: 10.1609/aimag.v17i3.1232. URL https://ojs.aaai.org/aimagazine/ index.php/aimagazine/article/view/1232.

Prefix-Optimal Generative Policies

Supplementary Materials The following content was not necessarily subject to peer review.

Theoretical Proofs

A.1

Proof of Theorem 1: Contraction and Unique Fixed Point of Th

Theorem 1. Let h(t) ≥ h > 0 for all t ∈ {1, . . . , K}. The Hazard Bellman operator (Th V )t (s, at ) := h(t) Q s, πfull (s, at ) + (1 − h(t)) Vt−1 s, πθ (s, at , t) ,

(13)

with boundary V0 (s, a0 ) = Q(s, a0 ), is a contraction in the ℓ∞ norm with rate (1 − h) and therefore has a unique fixed point. Proof. Let V, V ′ be two prefix value functions and define δt := ∥Vt − Vt′ ∥∞ := sups,at |Vt (s, at ) − Vt′ (s, at )|. At the boundary t = 0: δ0 = 0 since V0 = Q(s, a0 ) is independent of the prefix value function being iterated. For t ≥ 1: ′ |(Th V )t (s, at ) − (Th V ′ )t (s, at )| = (1 − h(t)) Vt−1 (s, at−1 ) − Vt−1 (s, at−1 ) (14) ≤ (1 − h(t)) δt−1

(15)

≤ (1 − h) δt−1 .

(16)

Taking the supremum over (s, at ): δt ≤ (1 − h) δt−1 ≤ (1 − h)t δ0 = 0.

(17) QK

Since the denoising chain is finite (length K), the entire operator on the product space t=0 Vt (equipped with the max norm maxt δt ) contracts by (1 − h) per “sweep” from t = 0 to t = K. By the Banach fixed-point theorem, there exists a unique fixed point V ∗ such that Th V ∗ = V ∗ . Remark 1. The contraction rate (1 − h) tightens as h → 1 (uniform stopping) and loosens as h → 0 (rare stopping). This motivates using hazard values bounded away from zero; empirically we use h = 0.05 uniformly. A.2

Proof of Proposition 1: Completion Map Bias Bound

Proposition 1. Let gω satisfy E[∥gω (s, aT , T ) − a0 ∥] ≤ ϵg (T ) and let Qπ be LQ -Lipschitz in its action argument. Then the performance gap under truncation is bounded: |Jδ1 (π) − Jh (π)| ≤

LQ ET ∼h [ϵg (T )]. 1−γ

(18)

Proof. For each environment step k, the policy under hazard h executes ãk = gω (sk , aTk , Tk ) instead of a0 . The reward difference at each step satisfies: |r(s, a0 ) − r(s, ã)| ≤ Lr ∥a0 − ã∥ = Lr ∥a0 − gω (s, aT , T )∥.

(19) (20)

Taking expectations and summing over the infinite horizon: |Jδ1 (π) − Jh (π)| ≤ E

∞ hX

i γ k |r(sk , a0 ) − r(sk , ãk )|

(21)

k=0

Lr ET ∼h ∥a0 − gω (s, aT , T )∥ 1−γ Lr ≤ ET ∼h [ϵg (T )]. 1−γ

≤

(22) (23)

Reinforcement Learning Journal

2026

Since Qπ is LQ -Lipschitz and Qπ (s, a) = r(s, a) + γV π (s′ ), we have Lr ≤ LQ (1 − γ), giving the stated bound with LQ in place of Lr /(1 − γ). Corollary 1. If ϵg (T ) is non-decreasing in T (i.e., more truncation yields larger projection error), then performance under any early-stopping hazard √− γ) below √ h is at most LQ · ET ∼h [ϵg (T )]/(1 full-chain performance. Empirically, ϵg (T ) ≈ c T (Figure 6a), so ET ∼h [ϵg (T )] ≤ c K, giving a finite, K-dependent bound. A.3

Proposition 2: Hazard Shape Controls the Fixed-Point Landscape

Proposition 2 (Restated). Let {Vt∗ } denote the unique fixed point of Th . Unrolling the recursion gives: VK∗ (s, aK ) =

K X

wt Q s, πfull (s, at ) ,

(24)

t=1

QK PK where wt = h(t) j=t+1 (1 − h(j)) and t=1 wt = 1. If Q(s, πfull (s, at )) = Q∗ for all t (uniform action quality across prefixes), then VK∗ = Q∗ for any hazard h. Proof. By induction. At t = 0: V0∗ = Q(s, a0 ) = w0 Q(s, πfull (s, a0 )) with w0 = 1. Assume the representation holds for all t′ < t. At level t: ∗ Vt∗ (s, at ) = h(t)Q(s, a0 ) + (1 − h(t))Vt−1 (s, at−1 )

= h(t)Q(s, a0 ) + (1 − h(t))

t−1 X

wt′ Q(s, πfull (s, at′ )),

(25) (26)

t′ =1

where at−1 = πθ (s, at , t). Setting wt = h(t) and absorbing (1 − h(t)) into the weights of the PK recursed sum yields the stated formula after normalization. The weights wt satisfy t=1 wt = 1 by the telescoping identity for geometric-like products. When Q(s, πfull (s, at )) = Q∗ for all t: PK VK∗ = Q∗ t=1 wt = Q∗ . Remark 2 (Uniform hazard and flat landscapes). Under uniform h(t) = h for all t, the weights wt = h(1 − h)K−t form a geometric distribution over refinement levels. This assigns strictly positive mass to every prefix, which is why uniform-hazard training produces fixed points that are robust to arbitrary test-time hazard distributions,the optimizer cannot concentrate entirely on any single truncation depth.

Gradient Variance Analysis: Shared vs. Independent Chains

We show formally why the shared-chain mechanism (stability fix S2) is necessary. Consider the actor loss gradient with respect to θ: ∇θ Lπ = −∇θ E[Q(s, a0 )] − w(n) ∇θ E[Vψ (s, at−1 , t − 1)].

(27)

Independent chains. If a0 and at−1 are generated from independent noise draws ξ0 , ξ1 ∼ N (0, I), the two gradient terms are estimated from different trajectories of the stochastic computation graph. Let gQ = ∇θ Q(s, a0 (ξ0 )) and gV = ∇θ Vψ (s, at−1 (ξ1 )). The total gradient variance is: Var[gQ + w · gV ] = Var[gQ ] + w2 Var[gV ] + 2w Cov[gQ , gV ].

(28)

Since ξ0 ⊥ ξ1 , the covariance term is Cov[gQ , gV ] = 0. The two gradient estimates therefore point in potentially opposing directions in parameter space.

Prefix-Optimal Generative Policies

Shared chains. With a shared noise draw ξ, both a0 (ξ) and at−1 (ξ) derive from the same trajectory. By the chain rule, gQ and gV share the portion of the computational graph from aK to at , inducing positive covariance: Cov[gQ , gV ] = E[⟨gQ − ḡQ , gV − ḡV ⟩] > 0,

(29)

since both gradients push θ toward parameters that produce better denoising from step t onward. This positive covariance reduces total variance and aligns the optimization direction between the two objectives. Empirical confirmation. Panel (d) of Figure 5 shows that gradient norms for ∇θ V stabilize and converge toward ∇θ Q after the confidence gate opens. With independent chains, the two gradient norms diverge and remain unstable throughout training (not shown; available from code). The shared-chain cosine similarity between gQ and gV averages 0.63±0.09 at convergence; independent chains average 0.07 ± 0.21.

Full Ablation Results

Table 4 extends the surgical ablations to all four environments and reports both full-chain return and truncation retention at T =5. Table 4: Full ablation results across all environments (4 seeds, mean ± std). Ret.@5 = fraction of full-chain return retained at T =5. Each row changes exactly one design decision relative to full POGP, holding all five stability fixes constant. Full-chain Return

Ret.@5 (%)

Variant

POGP (full)

347.2 190.2 98.4 215.8 99.1 98.6 98.2 98.8

Objective variants No prefix (Diff-QL backbone) Truncate at test only Distill to K ′ =5 Prefix-supervised (no Th ) Th + frozen actor Prefix at t=0 only

314.6 314.6 308.2 320.1 316.8 317.3

W2d

Hop

HC W2d Ant Hop

96.0 96.0 90.7 94.1 93.8 93.5

210.1 210.1 204.4 211.9 210.5 210.0

82.4 71.2 — 94.8 96.2 88.1

80.8 68.9 — 93.7 95.8 87.5

82.1 72.3 — 93.4 95.4 87.2

81.5 70.4 — 94.1 96.0 88.3

Stability fix ablations (one fix removed at a time) − Double-Q (S1) 326.8 184.5 95.2 − Shared chains (S2) 307.5 172.4 87.6 − Confidence gate (S3) 337.0 186.1 96.8 − Cosine warmup (S4) 342.5 188.2 97.1 − Fixed anchor weight (S5) 337.1 185.8 96.5

213.1 198.6 214.2 215.0 213.7

91.3 84.7 95.4 97.2 95.6

90.8 83.1 95.0 96.9 95.2

90.5 83.4 94.8 96.7 95.0

91.0 84.2 95.3 97.1 95.5

Oracle (Qenv targets)

178.9 178.9 171.3 183.4 181.2 180.7

Ant

338.4 191.8 99.1 218.2 99.8 99.7 99.6 99.8

Interpretation. The “distill to K ′ =5” row has no retention column because distillation changes the inference protocol: the model is retrained to execute in K ′ steps, so “truncation” at T =5 would mean zero refinement, an undefined quantity. The oracle row uses true Qenv values as Hazard TD targets (removing critic approximation error) and serves as an upper bound on what prefix learning alone can achieve with a perfect critic. The two most important ablations are shared chains (S2) and no prefix (baseline). Removing shared chains drops return by ∼40 points and retention by ∼15 percentage points, larger than any other single change. The gap between prefix-supervised and full POGP (4.3 retention points) isolates the contribution of the Hazard Bellman bootstrap specifically, above and beyond any generic intermediate supervision.

Reinforcement Learning Journal

Hazard Theory and Multi-Hazard Training

D.1

When Does Uniform Training Generalize?

2026

Proposition 2 shows that VK∗ is a convex combination of Q-values across all refinement levels, with weights determined by the hazard. A natural question is: for what class of test hazards htest does training under a single hazard htrain generalize? Proposition 3 (Generalization via Support Coverage). If the training hazard htrain has htrain (t) > 0 for all t ∈ {1, . . . , K}, then the fixed point V ∗ trained under htrain achieves bounded suboptimality under any test hazard htest with htest (t) > 0. Specifically, ET ∼htest [V ∗ (s, aT , T )] ≥ ET ∼htest [Q(s, πfull (s, aT ))] − ϵ,

(30)

where ϵ depends on the ratio htest (t)/htrain (t) and the approximation error of V ∗ . Intuitively, uniform training hazards assign non-trivial weight to every prefix level, preventing the optimizer from “ignoring” early truncation depths. Under early-biased test hazards (where small T is likely), the trained policy still performs well because prefixes at all levels were optimized. This theoretical prediction matches the empirical OOD robustness in Table 15. D.2

Multi-Hazard Training

A natural extension is to train under a distribution of hazards H rather than a single h: JH (π) = Eh∼H Jh (π) .

(31)

The Hazard Bellman operator generalizes straightforwardly: replace h(t) in Eq. (5) with h̄(t) = Eh∼H [h(t)]. The contraction rate becomes (1 − mint h̄(t)), which is positive as long as the average hazard is bounded away from zero at each level. An appealing special case is a learned hazard hϕ (t; s) that adapts to the environment state. If the Qlandscape is flat (all prefixes produce similar actions), the learned hazard should assign high stopping probability early. If the landscape is steep, the hazard should concentrate near t=0 (full refinement). Table 15 already demonstrates that test-time hazard generalization is strong under uniform training; a learned hϕ could extend this to environments where the Q-landscape varies significantly across states. We leave this to future work.

Implementation Details

E.1

Network Architecture

All networks use the same backbone for fair comparison: • Denoiser πθ : MLP, 256 hidden units, ELU activations, sinusoidal timestep embedding of dimension 64 concatenated to the input. Input: (s, at , t); output: at−1 ∈ Rda with tanh output activation scaled to [−1, 1]da . • Environment critic Qϕ1,2 : Identical architecture to standard SAC/TD3 critics. MLP, 256 units, ReLU. Input: (s, a); two independent heads for double-Q. • Prefix critic Vψ : MLP, 256 units, ELU. Input: (s, at , t) where t is encoded identically to the denoiser timestep embedding. Output: scalar. • Completion map gω : MLP (single linear layer), 256 units. Input: (s, aT , T ); output: projected action ã ∈ Rda with tanh scaling. 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. We verify that increasing gω to 2 layers does not change results meaningfully (< 0.5% retention difference).

Prefix-Optimal Generative Policies

E.2

Training Protocol Table 5: Hyperparameters shared across all methods and environments. Parameter

Value

Optimizer Learning rate (Q, V , g) Learning rate (π) Batch size Replay buffer size Random exploration steps Denoising chain length K Polyak averaging ρ Discount γ Gradient clipping Target network update Evaluation frequency Evaluation episodes Total training steps

Adam 3 × 10−4 3 × 10−4 256 106 10,000 10 (main), 20 (robustness) 0.995 0.99 1.0 (global norm) Every 1 gradient step Every 10,000 steps 10 106

Table 6: POGP-specific hyperparameters (tuned on HalfCheetah seed 42 applied to all environments and seeds).

E.3

Parameter

Value

Description

Hazard h Gate threshold τ Prefix weight wmax Warmup fraction Dynamic stopping ϵ Dynamic stopping m

1/K (uniform) 5.0 0.25 0.1 0.01 2

Stopping probability per denoising step V-critic RMSE must fall below this to open gate Maximum weight on prefix term in actor loss Fraction of training steps for cosine ramp Relative gain threshold Consecutive steps below threshold to halt

Baseline Implementations

All baselines are reimplemented from scratch in the shared codebase to ensure identical data handling, replay buffer management and evaluation protocol. Where baselines have published official code, we verified our implementations match their reported numbers within 2% on a common reference task (HalfCheetah at 1M steps). • TD3: Twin-delayed DDPG (Fujimoto et al., 2018). We use the standard hyperparameters and policy smoothing with σ=0.2, clip c=0.5. • SAC: Soft actor-critic (Haarnoja et al., 2018) with automatic entropy tuning. • PPO: Proximal policy optimization (Schulman et al., 2017) with 64 parallel environments, clip ϵ=0.2, 10 epochs per update. • Diffusion-QL: Diffusion policy with Q-weighted behavioral cloning loss (Wang et al., 2022). Online variant with no offline pretraining. • DPPO / D2 PPO: Diffusion policy optimized with PPO-style updates (Ren et al., 2024). D2 PPO adds a second denoising stage. • SDAC: Score-matching diffusion actor-critic (Ma et al., 2025). Uses a score-matching auxiliary loss to avoid backpropagating through the full chain. • DSAC-D: Liu et al. (2025) Diffusion extension of soft actor-critic with direct score-based policy updates. • D3P: Yu et al. (2025) Diffusion policy with a learned depth-selection prior and dynamic early exit.

Reinforcement Learning Journal

2026

• IDQL: Offline implicit diffusion Q-learning (Hansen-Estruch et al., 2023) with 100K offline pretraining steps from a random-policy replay buffer, followed by 900K online fine-tuning. • FQL: Flow Q-learning (Wildberger et al., 2023); the flow-based policy is trained with a reflow objective to enable one-step action generation. • SAC-GMM: SAC Nematollahi et al. (2022) with a 5-component Gaussian mixture policy. Serves as a capacity-matched multimodal baseline without iterative refinement. E.4

Computational Resources

All experiments were run on a single NVIDIA A100 (80GB) GPU. Total compute: Table 7: Compute summary. Component Main results Truncation results Ablations OOD hazard Sensitivity / K sweep

Runs

GPU-hours

12 × 4 × 10 = 480 9 × 4 × 10 = 360 8 × 4 × 4 = 128 3 × 4 × 4 = 48 —

≈ 250 ≈ 200 ≈ 70 ≈ 25 ≈ 30 ≈ 575

Total

Wall-clock time per 1M-step run is approximately 1.0–1.2 hours for POGP (slightly higher than baselines due to shared chain generation) and 0.7–0.9 hours for non-diffusion baselines.

Additional Results

F.1

Full Truncation Results: All Methods, All Environments

Table 8 extends truncation results to include all 12 methods and reports retention at five stopping times. Table 8: Extended truncation stress test (K=20, 10 seeds). Classical baselines use an actionrefresh proxy; diffusion methods execute the projected prefix. HalfCheetah Method

Walker2d

Ant

Hopper

T =2

T =5

T =10

T =15 T =2

T =5

T =10

T =15 T =2

T =5

T =10

T =15 T =2

T =5

T =10

T =15

TD3 SAC PPO

96 97 95

95 96 94

95 93 96 94 94 91

92 93 90

92 91 93 91 90 89

90 90 88

90 93 90 94 88 91

92 93 90

92 93 91

Diff-QL DPPO D2 PPO SDAC DSAC-D D3P IDQL FQL SAC-GMM

93 89 91 95 93 99 91 93 96

88 84 86 92 90 98 87 89 95

84 79 82 90 88 98 84 86 95

80 75 78 88 86 97 82 83 95

90 87 89 93 91 98 89 91 94

85 82 84 91 89 98 85 88 93

82 78 80 90 88 97 83 86 93

80 76 78 89 87 97 81 85 93

90 87 89 93 91 98 88 91 92

86 83 85 91 89 97 84 88 91

83 80 82 90 88 97 82 86 91

81 78 80 89 87 97 81 85 91

91 89 90 94 92 99 90 92 94

86 83 85 92 90 98 86 89 93

83 79 81 91 89 98 83 87 93

80 76 78 90 88 97 81 85 93

POGP

100

99 99

98 99

97 100

Prefix-Optimal Generative Policies

F.2

Prefix Value Convergence Trajectories

Figure 3 shows how the prefix value function Vψ (s, at , t) evolves during training for three representative refinement levels (t = 1, 5, 10) across environments. The convergence ordering (t=1 stabilizes fastest, t=10 slowest) matches the Bellman recursion: levels closer to the boundary condition V0 = Q converge first and the signal propagates backward through the chain. Walker2d

HalfCheetah

Hopper

Ant

200

100 200

300

Vψ value

150

75 150

200 100

100

0 0

80 4

Steps (×10 )

100

0 0

100

0 0

80 4

Steps (×10 )

100

Steps (×104 )

Figure 3: Prefix value convergence trajectories. Vψ at three refinement levels: t=1 (solid), t=5 (dashed) and t=10 (dotted) across 100K training steps. The boundary-anchored level (t=1) converges fastest, consistent with backward induction through the Hazard Bellman recursion.

F.3

Completion Map Error: Per-Environment Breakdown

The main text reports completion error averaged across environments. Table 9 disaggregates by environment and action dimension. As expected, higher-dimensional action spaces (da =8 for Ant) incur higher absolute error, but relative retention remains strong because the Lipschitz constant LQ scales sub-linearly with da in the environments tested. Table 9: Completion map error ϵg (T ) and retention at T =5 per environment.

F.4

Environment

ϵg (5)

LQ (est.)

Ret.@5

HalfCheetah Walker2d Hopper Ant

6 6 3 8

0.152 0.161 0.118 0.195

9.8 10.4 8.7 11.2

99.1% 98.6% 98.8% 98.2%

Separating Prefix Learning and the Completion Map

To separate the effects of prefix learning and the completion map gω , we evaluate a 2 × 2 ablation on HalfCheetah using five random seeds. We independently remove prefix supervision and the completion map while keeping the remaining architecture and training protocol unchanged. Prefix learning is the primary source of truncation robustness and adaptive stopping. At seven denoising steps, full POGP retains 98.8% of its full-chain return, whereas the variant without prefix learning retains only 87.7%, even when the completion map is present. The completion map provides a smaller but complementary benefit, particularly under aggressive truncation: at two steps, it increases retention from 94.8% to 99.0% when prefix learning is enabled. These results indicate that prefix supervision improves the denoising trajectory itself, while gω primarily corrects residual action-space noise at early stopping points.

Reinforcement Learning Journal

2026

Table 10: Prefix-learning and completion-map ablation (HalfCheetah, 5 seeds). Adaptive Return is obtained using each variant’s stopping rule. Ret.@T denotes the percentage of the corresponding 20-step full-chain return retained when execution is force-stopped at denoising step T . Variant

Return Steps

POGP (full) POGP without gω No prefix, with gω No prefix or gω

344.7 339.6 321.8 314.9

6.8 6.8 19.8 19.7

Failure Mode Analysis

G.1

Seed 100: Late Gate Opening

Ret.@2

Ret.@7

Ret.@15

Ret.@20

99.0±0.96 94.8±1.35 77.7±2.05 72.9±2.82

98.8±0.72 96.3±0.65 87.7±1.53 83.6±1.66

99.0±0.63 98.7±0.53 96.2±1.05 95.0±0.90

99.5±0.29 99.5±0.37 99.2±0.32 99.3±0.46

As reported in the main text, seed 100 (HalfCheetah) underperforms by 12 points relative to the seed median (306 vs. 352 average). The cause is traced to a late confidence gate opening: the V-critic RMSE remained above the threshold τ =5 until step 8,200, compared to the typical 4,000–5,500 range. Figure 4 shows that the delayed gate opening correlates with an unusually slow initial critic learning phase: the environment Q-function RMSE was ∼30% higher than median at step 5K for this seed. Without reliable Q targets, the prefix critic Vψ cannot converge, so the gate correctly delays. However, because the full-chain policy has already specialized by step 8K, the prefix regularizer has less room to influence the learned denoising trajectory. Mitigation strategies. A more robust gate criterion could monitor Vψ gradient alignment with Q gradients (their cosine similarity), rather than absolute RMSE. Pilot experiments with this criterion reduced the gate-opening variance from ±2.1K steps to ±0.8K steps (2 seeds; results not included in main comparisons). We leave a thorough evaluation to future work. (b) Gate Opening Time

(a) Q RMSE (HalfCheetah) Steps to gate open (×103 )

RMSE

40 seed 100 20 median

τ =5 0 0

Steps (×103 )

10 8.2 8 late threshold 6

5.1

4.8

4.2 4 2 0 42

100

101

Seed

Figure 4: Seed 100 failure diagnosis. (a) Seed 100’s Q-critic RMSE remains elevated compared to the median seed, slowing V-critic convergence. (b) Seed 100’s gate opens at 8.2K steps (vs. 4–5K typical), leaving less training time for prefix regularization to influence the policy.

G.2

High-Dimensional Action Spaces

The per-environment attribution analysis shows smaller gains for Ant (da =8) vs. HalfCheetah (da =6). We hypothesize two contributing factors: 1. Noisier Q gradients. Higher-dimensional action spaces yield larger critic estimation variance, reducing the reliability of Hazard Bellman targets (Q-gradient reliability 0.58 vs. 0.82). 2. Larger action space volume. The completion map gω must project from R8 partial actions rather than R6 , increasing ϵg (T ) (Table 9).

Prefix-Optimal Generative Policies

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.

Sensitivity Analysis

We evaluate POGP’s sensitivity to its key hyperparameters by sweeping one parameter at a time while holding all others at their default values (Table 6). All experiments use HalfCheetah with 4 seeds. H.1

Hazard Value h

Table 11 reports full-chain return and retention at T =5 for uniform hazard values h ∈ {0.05, 0.1, 0.2, 0.5}. Table 11: Sensitivity to uniform hazard value h (HalfCheetah, 4 seeds). Hazard h

Return

Ret.@5 (%)

0.05 0.10 (default) 0.20 0.50

341.8 ± 4.2 347.2 ± 3.8 344.1 ± 5.1 335.6 ± 6.3

97.8 ± 0.4 99.1 ± 0.3 98.5 ± 0.5 96.2 ± 0.8

Performance is robust across h ∈ [0.05, 0.2], degrading only at h = 0.05 where the high stopping probability biases the prefix value toward early (noisy) truncations. The default h = 1/K = 0.1 balances full-chain quality and prefix robustness. H.2

Gate Threshold τ Table 12: Sensitivity to gate threshold τ (HalfCheetah, 4 seeds). Threshold τ

Return

Ret.@5 (%)

Gate opens (K steps)

2.0 5.0 (default) 10.0 ∞ (no gate)

340.5 ± 5.8 347.2 ± 3.8 345.8 ± 4.4 337.0 ± 7.2

98.4 ± 0.5 99.1 ± 0.3 98.7 ± 0.4 95.4 ± 1.1

7.1 ± 1.8 4.8 ± 1.2 2.3 ± 0.6 0

Setting τ too low delays the gate excessively (reducing the effective prefix training window), while removing the gate entirely (τ = ∞) allows noisy early V-critic estimates to destabilize actor training. The default τ = 5 provides a good balance. H.3

Prefix Weight wmax Table 13: Sensitivity to maximum prefix weight wmax (HalfCheetah, 4 seeds). Weight wmax

Return

Ret.@5 (%)

0.05 0.10 0.25 (default) 0.50 1.00

343.1 ± 4.0 345.9 ± 3.9 347.2 ± 3.8 342.8 ± 5.5 331.4 ± 8.1

96.8 ± 0.6 98.2 ± 0.4 99.1 ± 0.3 98.9 ± 0.5 98.6 ± 0.7

Very small wmax under-regularizes (low retention), while wmax = 1.0 over-regularizes and harms full-chain return. The sweet spot lies in [0.1, 0.5], with wmax = 0.25 offering the best trade-off.

Reinforcement Learning Journal

2026

Chain Length Sensitivity

A key question is how POGP’s advantage scales with the denoising chain length K. Longer chains increase the cost of full inference and amplify the benefit of early stopping, but also make prefix training harder (more levels to optimize). We evaluate K ∈ {5, 10, 15, 20} across all four environments. I.1

Return and Retention vs. K

Table 14: Chain length sensitivity (4 seeds per cell). Full-chain return and retention at T = ⌈K/2⌉ (half-chain truncation). POGP’s retention advantage over the strongest diffusion baseline (SDAC) grows with K. K=5 Method

K=15

K=20

Ret.@⌈K/2⌉

Ret.

Ret.@⌈K/2⌉

Ret.

Ret.@⌈K/2⌉

Ret.

Ret.@⌈K/2⌉

HalfCheetah SDAC 310.2 POGP 338.5 ∆ +28.3

96.1% 98.4% +2.3

318.4 347.2 +28.8

90.2% 99.1% +8.9

322.1 351.6 +29.5

85.8% 98.8% +13.0

324.8 354.2 +29.4

82.4% 98.5% +16.1

Walker2d SDAC 175.8 POGP 184.3

95.2% 97.9%

182.6 190.2

89.8% 98.6%

185.4 193.8

84.5% 98.2%

187.2 196.1

81.1% 97.8%

Ant SDAC POGP

92.4 95.6

94.8% 97.2%

95.1 98.4

89.4% 98.2%

96.8 100.1

83.9% 97.6%

97.5 101.3

80.6% 97.1%

Hopper SDAC POGP

208.1 211.5

95.8% 98.1%

212.4 215.8

90.5% 98.8%

214.8 218.4

85.2% 98.4%

216.1 220.1

82.0% 98.0%

I.2

Ret.

K=20

Discussion

Three trends emerge from Table 14: 1. POGP’s retention advantage grows with K. At K=5, the retention gap over SDAC is modest (∼2–3 percentage points) because short chains leave little room for truncation degradation. By K=20, the gap widens to ∼15–17 points: SDAC’s retention drops steadily as the chain lengthens, while POGP maintains >97% retention at half-chain truncation across all environments. 2. Full-chain return improves monotonically with K for both methods. Longer chains provide more expressive policies. POGP consistently outperforms SDAC in absolute return at every K, with a roughly constant gap of ∼28–30 points on HalfCheetah. The prefix training objective does not harm full-chain quality. 3. POGP’s retention is approximately constant across K. While SDAC’s retention at half-chain drops from ∼96% (K=5) to ∼82% (K=20), POGP’s retention stays in the narrow band of 97– 99% regardless of chain length. This confirms that the Hazard Bellman operator successfully propagates prefix value information across arbitrarily many refinement levels, as predicted by the contraction analysis (Theorem 1). These results suggest that POGP becomes increasingly valuable as diffusion-policy practitioners adopt longer chains for more complex tasks. Whether this scaling extends beyond K=20 to the very long chains used in image generation (K=50–100) remains an open question; we anticipate

Prefix-Optimal Generative Policies

that architectural modifications to the prefix critic (e.g., attention over refinement levels) may be needed to maintain stable backward induction over such depths.

Ablation and Dynamics Analysis

Figure 5 presents comprehensive ablation and training dynamics analysis. Panel (a) identifies the shared chain mechanism (S2) as the most critical stability component, with its removal causing a 39.7-point return drop. Panels (b–d) show training dynamics relative to the confidence gate opening at 40K steps. (a) Stability Ablation

(d) Grad Norms

−39.7

Gate (S3)

−10.2

Gate Opens

−10.4

Fraction

Anchor (S5)

τ =5 −4.7

Cosine (S4)

−40

−20

∆ Return

100

150

200

Steps (×103 ) POGP (Ours)

∥∇θ ∥

∇V −20.4

Dbl-Q (S1)

RMSE

Chain (S2)

(b) V-critic RMSE

100

150

Steps (×103 )

200

∇Q

100

Diff-QL

Figure 5: Ablation and Dynamics Analysis. (a) Sensitivity analysis identifying the shared chain (S2) as the critical stability component. (b–d) Training dynamics relative to the confidence gate (shaded). Once the RMSE threshold τ =5 is met at 40k steps, the gate opens, resulting in stabilized gradient norms and a significant reduction in prefix violation rates.

Completion Map and Calibration Analysis

Two supporting analyses round out the experimental picture. The projection network gω maps truncated actions to executable√ones. Figure 6 (left two panels) confirms the theoretical prediction: projection error grows as O( T ) and performance retention decays linearly with this error, bounded by LQ · ϵg (empirical LQ ≈ 10). Retention stays above 95% even at the largest truncation. For the adaptive stopping rule to work, Vψ must rank actions reliably. Figure 6 (right panel) reports Spearman ρ between Vψ (s, aT , T ) and realised return. Rank correlation is 0.72 at T =10 (one denoising step) and 0.94 at T =1 (full chain), enough for the stopping rule to make useful decisions.

Hazard Robustness

A well-trained policy should generalise beyond the training hazard. Table 15 evaluates four testtime hazard distributions never seen during training (POGP trains with uniform h=1/K). Even under early-biased truncation where interruption concentrates at minimal refinement, the hardest case, POGP retains 97.4%. Baselines drop below 81%.

Extended Related Work

M.1

Diffusion Models for Offline RL and Planning

150

Steps (×103 )

The use of diffusion models in RL predates online training formulations. Janner et al. (2022) use diffusion to model entire trajectory distributions, enabling planning by sampling and reranking trajectories. Wang et al. (2022) introduce Diffusion-QL, the first to use a diffusion actor with Q-weighted

200

Reinforcement Learning Journal

2026

Retention Decay

Completion Error √ O( T )

≈ LQ ϵg

Spearman ρ

0.2

(c)

(b) Retention (%)

Error ϵg (T )

(a) 0.4

Value Calibration

Reliable Range

95% threshold

0.2

0.4

Completion error ϵg

Stopping time T

Hopper

HalfCheetah

Stopping time T Walker2d

Ant

Figure 6: Completion map analysis √ and prefix value calibration. (a) Projection error ϵg (T ) grows sublinearly, matching the O( T ) theoretical trend. (b) Policy retention decays linearly with completion error, governed by the Lipschitz constant LQ . (c) Spearman correlation remains high (ρ > 0.7) even at significant truncation, validating the use of Vψ for early stopping. Table 15: Hazard robustness (HalfCheetah, 4 seeds). Return normalised to full-chain. All test hazards are unseen during training. h=δ1

Unif.

Geom.

Diff-QL 100% SDAC 100% POGP 100%

84.2% 86.1% 98.8%

86.8% 91.5% 78.3% 88.4% 93.2% 80.1% 99.1% 99.5% 97.4%

Late

Early

fine-tuning in offline settings. These offline methods implicitly assume inference always completes; POGP’s prefix training extends their objectives to the online, variable-compute setting. M.2

Relationship to Progressive Training

Progressive distillation (Salimans & Ho, 2022) and consistency models (Song et al., 2023) train models to match the output of many-step samplers with fewer steps. While both improve efficiency, they optimize for a fixed reduced K rather than arbitrary truncation depths. POGP is complementary: one could apply progressive distillation to reduce K from 20 to 10, then apply POGP to make the K=20 chain truncation-robust. M.3

Relationship to Options and Temporal Abstraction

The options framework (Sutton et al., 1999) 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 where an option’s internal computation is itself a prefix-optimal process could extend POGP to hierarchical settings. M.4

Relationship to Adaptive Computation Time

ACT (Graves, 2017) and PonderNet (Banino et al., 2021) learn when to halt iterative computation in supervised settings. POGP provides the RL analogue: the stopping rule emerges naturally from the learned prefix value Vψ , without a separate halting network. Unlike ACT, POGP’s stopping criterion is grounded in expected return rather than a cost penalizing computation.

Prefix-Optimal Generative Policies

Broader Impact and Limitations

Positive impact. POGP reduces the compute required for diffusion-policy inference in robotics and continuous control: 2–4× speedups with <5% quality loss translate directly into lower energy consumption and faster control loops. In safety-critical settings (e.g., prosthetics, industrial robots), the anytime guarantee that interruption yields a usable action is a meaningful safety property. Limitations. 1. Evaluation scope. Results are limited to MuJoCo locomotion. Whether prefix training generalizes to manipulation, navigation, or real-robot settings is unknown. 2. Stochastic denoisers. The main text uses deterministic DDIM-style denoisers. Stochastic denoisers require modifying the refinement MDP to include stochastic transitions; the Hazard Bellman operator extends, but contraction rate analysis would need to account for transition noise. 3. Long chains. The advantage of POGP over SDAC grows with K (Appendix I), but experiments are limited to K ≤ 20. Whether this continues to K=50 or K=100 (as used in image generation) is unexplored and may require architectural changes to the prefix critic. 4. Seed sensitivity. One in four evaluated seeds (seed 100) failed to benefit from prefix training due to late gate opening. A more robust gating criterion is needed for reliable deployment.

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

Analyze & explore related research in Conceptio

Related documents