PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Policy Improvement with Proposal-Conditioned Refinement Flows

Junhyun Ha, Juho Lee, Byoungwoo Park

arXiv:2609.36812v2cs.LGcs.AIcs.RO

TL;DR

KL regularization can limit improvement of expressive offline policies when high-value actions have low behavior density, and simple proposal editors can limit refinement flexibility. PReFlow selects behavior proposals with a critic and refines them with a conditional flow under a joint regularized objective. It remains competitive offline and achieves the highest aggregate score among compared methods after online fine-tuning.

  • Problem

    KL-regularized policy improvement can restrict high-value actions with low behavior density, while Gaussian or deterministic editors have limited flexibility to represent multiple refinements of one proposal.

    Source: §1 · p. 1

  • Method

    PReFlow selects behavior proposals with the critic and refines them using a conditional flow, jointly regularizing proposal selection and refinement under a Gaussian-smoothed behavior prior.

    Source: §1 · pp. 1–2

  • Results

    PReFlow is competitive in offline RL and has the highest aggregate post-fine-tuning score among compared methods; after 500K online environment steps, its aggregate score is 91.

    Source: §1 · p. 2, §5.1 · p. 8, §6 · p. 10

  • Takeaways & Limitations

    PReFlow's results suggest expressive policies help not only model multimodal behavior but also refine each proposal.

    Source: §6 · p. 10

    Proposal selection uses Qψ as an approximation to refinement value Vδ, although proposals with similar critic values can differ through the critic's local gradient and curvature.

    Source: §6 · p. 10

Abstract

from arXiv · show

Diffusion and flow policies can model complex behaviors in offline reinforcement learning (RL). However, penalizing their KL divergence from the behavior policy can discourage actions having high critic values with low behavior density. Directly refining behavior proposals may be an alternative, yet Gaussian or deterministic editors limit expressiveness to represent multiple separated modes for the same proposal. In this work, we introduce Proposal-Conditioned Refinement Flows (PReFlow), a policy extraction method combining critic-based proposal selection with a conditional refinement flow. To optimize proposal selection and refinement together, we formulate a KL-regularized objective whose optimum induces a Gibbs policy over final actions under a Gaussian-smoothed behavior prior. The refinement flow can represent multiple high value modes, while a proposal-centered Gaussian reference regulates large action changes. This Gaussian reference further enables us to make use of simulation-free, closed form adjoint matching targets from sampled endpoints and critic gradients, yielding a single velocity regression loss without a backward adjoint solve. On 50 OGBench tasks, PReFlow achieves competitive offline performance and the highest aggregate score among the compared methods after online fine-tuning, reaching 91\% after 500K environment steps.

1. Introduction

PReFlow addresses limits in policy improvement and candidate-based extraction by selecting behavior proposals and refining them with a conditional flow. Its joint objective and Gaussian reference support expressive local refinements and simulation-free training; it is competitive offline and leads the compared methods in aggregate offline-to-online performance.

1. Introduction

  • KL regularization can restrict improvement when high-value actions have low behavior density, while candidate-based extraction cannot output actions beyond its sampled proposals.Residual editors based on Gaussian or deterministic refinements also have limited flexibility to represent multiple distinct refinements for one proposal.Source: §1 · p. 1
  • PReFlow selects a behavior-flow proposal with the critic, then uses a conditional flow to add a refinement to that proposal.A proposal-centered Gaussian reference discourages large action changes without restricting the refinement distribution to a Gaussian family, allowing distinct high-value refinements.Source: §1 · p. 1
  • A joint KL-regularized objective for selection and refinement induces a Gibbs policy under a Gaussian-smoothed behavior prior.The Gaussian reference enables closed-form adjoint-matching targets from forward-sampled endpoints and critic gradients, trained with one velocity-regression loss and no backward adjoint solve.Source: §1 · pp. 1–2
  • PReFlow achieves the highest aggregate offline-to-online performance among compared methods while remaining competitive in offline RL.Controlled ablations report gains from expressive flow refinement over Gaussian editing, improvements with a single proposal, and further benefits from proposal selection.Source: §1 · p. 2

2. Preliminaries

The preliminaries frame offline RL as return maximization from behavior-policy data, then describe behavior-regularized policy extraction and flow matching. They also outline QAM's adjoint-matching targets and the sequential backward computation they require.

2. Preliminaries

  • Offline RL seeks a policy maximizing expected discounted return, with offline data collected by a behavior policy and offline-to-online RL adding fine-tuning through limited online interactions.The setting is an MDP with states, actions, transitions, discount, rewards, and an initial-state distribution.Source: §2 · p. 2
  • KL-regularized policy extraction balances critic value against deviation from the behavior policy, and optimizing over action distributions yields a Gibbs policy favoring higher-value actions.The coefficient τ(s) controls critic value relative to the KL penalty.Source: §2 · p. 2
  • Flow matching learns a behavior velocity field by transforming Gaussian noise into behavior actions through an ODE, using interpolated noise-action pairs and their velocity as regression targets.QAM uses a learned flow as the ODE velocity to approximate the Gibbs target policy.Source: §2 · p. 3
  • QAM's adjoint matching samples SDE trajectories and propagates endpoint critic gradients backward along them to construct flow-training targets.Because the adjoint depends on the reference-flow Jacobian, target construction requires sequential backward vector–Jacobian products across the sampled trajectory.Source: §2 · p. 3

3. Policy Improvement with Conditional Refinement Flows

PReFlow selects behavior-policy proposals and conditionally refines them with a flow, using a proposal-centered Gaussian reference to regulate action changes while allowing multimodal refinements. A joint KL-regularized objective connects this two-stage policy to a Gibbs policy over final actions, and its Gaussian reference supports simulation-free flow training.

3.1. Towards Expressive Conditional Refinement

  • The conditional refinement flow can represent multiple high-value modes for one proposal, while a proposal-centered Gaussian reference discourages large action changes without restricting the refinement to a Gaussian family.In the toy comparison, PReFlow captures all three modes, while QAM undercovers two modes farther from the proposal and Gaussian refinement concentrates on one.Source: §3.1 · p. 4, Figure 2 · p. 4

3.3. Simulation-free Training of Flow Refinement

  • A Gaussian reference enables simulation-free regression targets from sampled endpoints and critic gradients, without a backward adjoint solve or residual-reconstruction loss.The targets use Gaussian bridge samples and endpoint critic gradients.Source: §3.1 · p. 4, §3.3 · pp. 6–7

3.2. A Regularized Objective for Conditional Policy Improvement

  • For a fixed proposal, the optimal refinement balances critic value against a KL penalty on refinement magnitude and can remain multimodal when the critic favors distinct refinements.The reference for the refined action is a Gaussian centered at the proposal.Source: §3.2 · p. 5
  • Proposal selection is jointly regularized toward the behavior policy, with the optimal conditional refinement value guiding which proposals are favored.This value reflects the proposal's value after regularized refinement.Source: §3.2 · p. 5
  • Jointly optimizing proposal selection and conditional refinement yields a Gibbs policy over final actions under a Gaussian-smoothed behavior prior.Nearby proposals can contribute reference density to actions with low behavior density.Source: §3.2 · p. 6
  • For sufficiently smooth critics and small δ, proposal values can be approximated by critic evaluations, and the practical method greedily selects the highest-valued sampled candidate.Finite-temperature selection weights are proportional to exp(τQψ(s, a_i)); the method instead uses greedy selection.Source: §3.2 · p. 6

3.3. Simulation-free Training of Flow Refinement

  • The refinement flow models normalized additive changes by conditioning on the state and fixed proposal, then integrating its velocity from Gaussian noise to a sampled endpoint.The endpoint gives the refinement Δa = δY1.Source: §3.3 · p. 6
  • Training samples endpoints from the current refinement flow, forms Gaussian bridge states, and regresses the velocity target using a single objective.The target is constructed from the endpoint and bridge sample; sampled endpoints and targets are held fixed during the update.Source: §3.3 · pp. 6–7

4. Related Work

Related work develops critic-guided improvement for diffusion and flow policies through several training and sampling strategies. PReFlow instead combines proposal selection with locally refined actions trained by adjoint matching under proposal-centered KL regularization.

4. Related Work

5. Experiments

The experiments evaluate PReFlow across 50 OGBench tasks against policy-extraction baselines, using offline training followed by online fine-tuning. They examine aggregate performance, refinement design and cost, proposal selection, and sensitivity to refinement scale and inverse temperature.

5.1. Offline and Offline-to-Online Performance

  • 91 after 500K online steps was PReFlow’s highest aggregate score; its offline score was 53, competitive with RQL.The aggregate comparison covers offline and online success rates.Source: Table 1 · p. 8, §5.1 · p. 8

5.2. Expressiveness and Efficiency of Conditional Refinement Flow

  • With one proposal, flow refinement achieved higher mean scores than Gaussian refinement across all four domains, and both improved over the base policy.The comparison used δ=0.1 and evaluated al, hm, scene, and c2.Source: §5.2 · p. 9, Table 2 · p. 9
  • PReFlow had the lowest mean wall-clock time per update among the compared methods across five flow-step counts.The closed-form regression target requires no backward simulation.Source: §5.2 · p. 9, Figure 4 · p. 9

5.3. Effect of Proposal Selection

  • Proposal selection complemented learned refinement, with offline-performance and final-online-success gains already evident at K=2.For K>1, the critic selects the highest-value behavior proposal before refinement.Source: §5.3 · p. 9, Figure 5 · p. 10

5.4. Sensitivity to Refinement Scale and Inverse Temperature

  • Offline performance varied with both δ and τ, which control effective refinement strength under a local linear approximation.The conditional target’s refinement mean is τδ^2∇_aQ(s,a), with covariance δ^2I.Source: §5.4 · p. 9, Figure 6 · p. 10

6. Conclusion

PReFlow selects behavior proposals with a critic and refines them using a conditional flow. The paper reports competitive offline performance and the highest aggregate success rate after online fine-tuning, while noting a limitation in how proposals are valued.

6. Conclusion

  • PReFlow combines critic-based proposal selection with conditional flow refinement, using a proposal-centered Gaussian reference and closed-form adjoint matching targets.The targets use sampled refinement endpoints and critic gradients, enabling simulation-free training without a backward adjoint solve.Source: §6 · p. 10
  • Proposals with similar critic values can differ in refinement value because of local critic gradients and curvature.Efficiently estimating refinement value while capturing that dependence is identified as future work.Source: §6 · p. 10

A. Proofs and Technical Details

The appendix derives the proposal-conditioned optimum, its Gaussian-reference adjoint targets, and a practical approximation for proposal selection. It also details implementation choices, related methods, and conditions on the idealized training characterization.

A.1. Optimal proposal-conditioned policy improvement

  • The optimal refinement policy has value V_δ(s, a) = τ^-1 log Z(s, a), and optimizing refinement yields the optimal proposal-selection policy.The final-action mixture under these policies gives the unclipped action distribution.Source: §A.1 · p. 14

A.2. Gaussian reference and adjoint target

  • PReFlow’s Gaussian reference provides closed-form bridge distributions and adjoints, enabling endpoint-gradient velocity targets without an additional simulation or backward adjoint solve.Endpoints are sampled by forward rollouts; intermediate states come from the analytic bridge.Source: §A.2 · pp. 14–15

A.3. Proposal-selection approximation

  • For smooth critics, V_δ(s, a) = Q_ψ(s, a) + O(δ^2), motivating critic-based selection weights when refinement scales are small.The implementation uses greedy critic argmax, which does not generally sample the optimal proposal marginal.Source: §A.3 · p. 15

B.1. Policy extraction with diffusion and flow policies

  • PReFlow evaluates the critic at the final refined action and uses its gradient for a velocity target, avoiding intermediate-value estimation and value backups across refinement steps.Unlike QPILOTS-M, refinement rollout requires no critic gradients at inference time.Source: §B.1 · p. 16

B.2. Behavior proposals and conditional action refinement

  • PReFlow’s conditional flow can represent multiple refinement modes per proposal; its Gaussian KL regularizer discourages large edits without a hard displacement radius or behavior-score estimate.The joint optimum weights proposals by regularized refinement value and induces a Gibbs action marginal under the Gaussian-smoothed behavior prior; implementation uses critic argmax.Source: §B.2 · p. 16

B.3. Adjoint matching and tractable reference dynamics

  • The Gaussian terminal law matches the refinement prior, so the density correction vanishes and the terminal cost depends on the critic at the refined action.Analytic bridges and adjoints provide training targets without further trajectory simulation or a backward solve, although endpoint generation still uses forward rollouts.Source: §B.3 · pp. 16–17

C.1. Temperature Adaptation

  • Temperature adaptation uses endpoint dispersion as a proxy for conditional KL and updates τ toward a target dispersion c.For the linearized target, c ≥ 1 corresponds to conditional KL of (c−1)/2 per action dimension.Source: §C.1 · p. 17

C.2. Implementation details

  • Training draws refinement proposals directly from the behavior policy, uses K-candidate selection for generated actions, and clips refined actions to environment bounds.The implementation also samples and refines one proposal for each TD target.Source: §C.2 · p. 17

C.3. Training procedure

Algorithm 3 jointly updates the critic, behavior flow, and refinement flow, with optional automatic temperature adaptation and a target-network update.

C.3. Training procedure

  • Each joint update trains the critic, behavior flow, and refinement flow, then updates the target critic by exponential averaging.The refinement update uses Algorithm 2, and temperature is updated from its endpoints when automatic adaptation is enabled.Source: §C.3 · p. 18
  • The critic update generates a next action using Algorithm 1 with K = 1, clips it to action bounds, and updates ten critics.The target uses the pessimistic target value and the expectile extension of the squared TD-error loss; κ = 0.5 recovers squared-error regression up to a constant factor.Source: §C.3 · pp. 17–18
  • Behavior-flow training uses flow matching on replay actions, while refinement-flow training uses proposals with K = 1 when automatic temperature adaptation is enabled.The algorithm outputs updated behavior-flow, refinement-flow, critic, target-critic, and temperature parameters.Source: §C.3 · p. 18

D.1. Benchmark and evaluation protocol

The evaluation follows the OGBench single-task protocol across ten domains and 50 tasks, with offline training followed by online fine-tuning. The appendix specifies domain settings and how comparison baselines were sourced or reproduced.

D.1. Benchmark and evaluation protocol

  • The benchmark covers 50 tasks across ten OGBench domains, with 1M offline gradient steps followed by 500K online environment steps and scores averaged across tasks.PReFlow and reproduced TRQAM, Q-Flow, and RQL use eight seeds per task for the main benchmark.Source: §D.1 · p. 18

D.2. PReFlow hyperparameters

  • Per-domain settings were selected by evaluating candidate temperatures, dispersion targets, refinement scales, and proposal counts, then fixing the best configuration across each domain’s five tasks.Configurations were selected using the sum of offline and online success rates; the sensitivity sweep instead used one task per domain and no online fine-tuning.Source: §D.2 · p. 19, Table 3 · p. 19

D.3. Baselines

  • The comparison combines released baseline results with eight-seed reproductions of TRQAM, Q-Flow, and RQL, and released eight-seed results for QPILOTS variants.For Q-Flow, the study reports the 1M-step evaluation rather than averaging the original paper’s 800K, 900K, and 1M evaluations.Source: §D.3 · pp. 19–20

D.4. Toy example: conditional multimodality

The toy comparison tests whether a conditional flow can fit a multimodal refinement target, while OGBench experiments compare refinement classes and proposal-selection rules.

D.4. Toy example: conditional multimodality

  • The toy experiment compares flow refinement, full-covariance Gaussian refinement, and QAM using the same conditional objective, critic, proposal, reference, scale, and temperature.All methods train for 100K steps; the flow and Gaussian refiners use different batch sizes.Source: §D.4 · pp. 20–21
  • Flow refinement ends below 0.1 KL to the toy target and has the lowest mean KL at every evaluated training step.Figure 7 plots KL divergence to the target with means over three seeds and one-standard-deviation shading on logarithmic axes.Source: Figure 7 · p. 21, §D.4 · p. 21

D.5. Refinement policy class on OGBench

  • The OGBench refinement-class study compares a learned diagonal Gaussian, a learned flow, and no refinement, with all variants using one proposal.It evaluates offline success on four domains, with five tasks per domain and three seeds per task.Source: §D.5 · p. 21

D.6. Proposal selection

  • The proposal-selection study compares uniform sampling, critic-weighted sampling, and greedy critic selection before applying the learned refinement flow.Each setting is evaluated over three seeds covering all five tasks in a domain.Source: §D.6 · p. 21

D.7. Timing methodology

The appendix describes timing measurements and compares SPAR, proposal-selection rules, refinement, and hyperparameter choices. It also specifies how full task-level results are reported.

D.7. Timing methodology

  • Figure 4 reports mean full-training-update times on one NVIDIA RTX A6000 GPU, averaged over flow-step counts of 5, 10, 20, 50, and 100.Source: §D.7 · p. 21

E.2. Computational cost

  • Table 5 reports c2 training-update costs in milliseconds, with median times over 30 calls after warm-up and the lowest time highlighted per column.The protocol uses 10 Euler steps, batch size 256, and device synchronization; the full step includes all training updates.Source: Table 5 · p. 22, §D.7 · p. 22, §E.2 · p. 22

E.1. Comparison with SPAR

  • SPAR improves on its anchor in al and scene but remains below PReFlow in all four evaluated domains.The tested SPAR configuration uses a Gaussian anchor, hard filter, λ_g=1, and λ_u=0.5.Source: §E.1 · p. 22

E.3. Effect of Proposal Selection

  • Critic-guided selection improves over uniform selection on c2 and scene, with limited improvement on hm; argmax has the highest mean in four of six domain-K settings.The authors use argmax selection in their benchmark experiments.Source: §E.3 · p. 23, Table 6 · p. 23

E.4. Proposal selection without refinement

  • Without refinement, K=1 scores near zero, and selecting more proposals remains far below refinement with a single proposal.The comparison averages offline success rates over al, hm, scene, and c2.Source: §E.4 · p. 23, Figure 8 · p. 23

E.5. Domain-specific hyperparameter choices

  • Increasing critic expectile κ from 0.5 to 0.7 improves offline performance by 4.9 and 5.9 percentage points on c2 and c4, respectively.The change does not consistently improve other domains; the authors hypothesize that increased critic overestimation limits its benefit there.Source: §E.5 · p. 23
  • Increasing refinement scale δ from 0.1 offline to 0.2 online improves online performance on ag, but yields no significant gains on other domains.Source: §E.5 · p. 24

E.6. Full task-level results

  • Tables 7 and 8 report per-task offline and post-fine-tuning success rates, with aggregate rows averaging five tasks per domain.Task cells report means and 95% confidence intervals; Figure 10 supplies task-level learning curves.Source: §E.6 · p. 24

E.7. Full training curves

The training-curve figure covers all 50 OGBench tasks and includes PReFlow alongside baseline curves. Online fine-tuning follows 1M offline gradient steps and lasts 500K environment steps.

E.7. Full training curves

Related papers

Loading 2609.36812v2…