PaperlayerView exact revision on arXiv

Source-linked AI summary

Schr\"odinger--F\"ollmer Actor--Critic: Diffusion Policy Improvement with Finite-Sample Analysis

Yuling Jiao, Lican Kang, Jerry Zhijian Yang, Jincheng Ying

arXiv:2609.33239v1stat.MLcs.LG

TL;DR

Diffusion policies can represent multimodal actions, but an advantage-weighted update specifies a target distribution without specifying how to sample from it. SFAC uses a Doob h-transform drift correction, estimated with paired SNIS and fitted by supervised regression, to implement the update. Its reference-anchored implementation improves final-window returns over offline initialization on six evaluated continuous-control tasks.

  • Problem

    An advantage-weighted update defines a target distribution for a diffusion policy but does not specify the diffusion dynamics needed to sample from it.

    Source: §1 · p. 3

  • Method

    SFAC uses a minimax Bellman critic and a Doob h-transform drift correction, estimated with paired SNIS and fitted by supervised regression without critic action gradients.

    Source: §1 · pp. 3–4

  • Results

    On six D4RL continuous-control tasks, the reference-anchored implementation achieved higher final-window normalized returns than the corresponding offline initialization.

    Source: §1 · p. 4, §7 · p. 44

  • Takeaways & Limitations

    In the small-update regime, the ideal KL-regularized tilt follows the natural policy-gradient direction, and the Doob correction represents the same first-order change in diffusion-drift space.

    Source: §1.1 · p. 5

    The analysis does not separately quantify the effects of replay-buffer management.

    Source: §3.1 · p. 18

Abstract

from arXiv · show

Diffusion policies represent multimodal action distributions, but an advantage-weighted update does not specify how to sample from the resulting target distribution. We propose Schr\"odinger--F\"ollmer Actor--Critic (SFAC), an offline-to-online reinforcement learning (RL) method for Kullback--Leibler (KL)-regularized policy improvement through conditional diffusion. A minimax Bellman critic estimates the advantage function, which defines an exponentially tilted target policy. A Doob $h$-transform expresses this update as a correction to the reference diffusion drift. We derive a posterior-mean representation of the correction and estimate it using paired self-normalized importance sampling (SNIS). Supervised regression on these drift targets updates the neural actor without critic action gradients. In the small-update regime, the KL-regularized update follows the natural policy-gradient direction, and the Doob correction represents the same local change in the space of diffusion drifts. Under suitable conditions, we derive finite-sample bounds that separate the effects of critic estimation, neural drift regression, finite-sample SNIS, diffusion discretization, and inherited actor error on expected average policy suboptimality. Synthetic experiments assess the accuracy of approximation to prescribed advantage-tilted targets and sensitivity to sampling budgets. On six offline-to-online continuous-control tasks, a reference-anchored implementation achieves higher final-window returns than those of its corresponding offline initialization.

1 Introduction

The introduction frames a gap between defining an advantage-tilted target for a diffusion policy and providing dynamics that sample from it. It presents SFAC’s drift-correction approach, its finite-sample analysis, and its evaluation on continuous-control tasks.

1 Introduction

  • Diffusion actors represent multimodal conditional action distributions, but an advantage-weighted target does not specify diffusion dynamics for sampling from that target.In offline-to-online RL, early online value estimates remain influenced by offline data coverage, while restrictive policy updates can limit adaptation.Source: §1 · pp. 2–3

1.1 Contributions

  • SFAC realizes the advantage tilt through a Doob h-transform drift correction, estimated as the difference between tilted and reference posterior means using paired SNIS.The estimated correction supplies supervised regression targets for the next actor, without critic action gradients; locally, the update follows the natural policy-gradient direction.Source: §1 · pp. 3–4, §1.1 · p. 5
  • On six D4RL continuous-control tasks, the reference-anchored implementation achieved higher final-window normalized returns than the corresponding offline initialization.Synthetic experiments assess approximation to prescribed advantage-tilted targets and sensitivity to endpoint-sample and diffusion-step budgets.Source: §1 · p. 4, §1.1 · p. 5
  • Under suitable conditions, the analysis bounds expected average policy suboptimality by combining critic estimation, actor regression, SNIS, discretization, and inherited actor-error contributions.The bounds account for dependent transition data and separate statistical and computational budgets.Source: §1 · p. 4, §1.1 · p. 5

1.2 Related Work

  • SFAC differs from generative-policy approaches that guide synthetic transitions by updating conditional action distributions through posterior-mean drift targets and supervised actor refitting.The introduction distinguishes this drift-target construction from KL regularization or actor refitting alone.Source: §1.2 · pp. 6–7

2 Preliminaries

The preliminaries set up the discounted MDP and function classes, then present the Schrödinger–Föllmer diffusion and Doob transform used to represent conditional tilted distributions.

2.1 Notation

  • The paper’s notation specifies vector and function norms, probability laws, the standard Gaussian measure, and asymptotic shorthand.It distinguishes the supremum norm from the measure-essential supremum norm.Source: §2.1 · p. 9

2.2 Markov Decision Process

  • A discounted MDP defines state and action spaces, transition and reward kernels, and policy value and advantage quantities used to describe performance.The paper also models temporal dependence with β-mixing.Source: §2.2 · pp. 9–11

2.3 Deep Neural Networks and Hölder Classes

  • The function classes comprise bounded ReLU networks for approximation and Hölder classes for population functions.Network classes constrain depth, width, sparsity, and output bounds.Source: §2.3 · pp. 12–13

2.4 Schrödinger–Föllmer Process and Doob h-Transforms

  • A Schrödinger–Föllmer process transports the origin to a target distribution over finite time, while a Doob h-transform changes its terminal law by exponential tilting.For conditional generation, the construction applies statewise; its drift also has a conditional-mean representation that supports sample-based estimation.Source: §2.4 · pp. 14–15

3 Method

SFAC estimates the current policy’s advantage with a minimax Bellman critic, then realizes the exponentially tilted target through a conditional diffusion drift correction. Paired SNIS estimates that correction, and regression fits the next actor.

3.1 Advantage Function Learning

  • The minimax Bellman critic estimates the current policy’s action-value function, and subtracting a Monte Carlo state-value estimate yields the advantage used in the plug-in target.Under the stated realizability and completeness conditions, the true action-value function is a population minimizer of the critic objective.Source: §3 · p. 16, §3.1 · pp. 17–18

3.2 Conditional Schrödinger–Föllmer Policy Improvement

  • SFAC represents the current policy as a conditional Schrödinger–Föllmer process and uses a Doob h-transform to realize the exponentially tilted target through a drift correction.The target drift differs from the reference drift by the difference between advantage-tilted and reference posterior means; paired SNIS estimates these means from shared terminal-action samples.Source: §3.2 · pp. 19–21
  • The actor is updated by regressing a neural drift toward the current learned drift plus the Monte Carlo correction, then sampling the fitted policy with Euler–Maruyama.The implementation pretrains its initial drift network on offline state–action data.Source: §3 · p. 17, §3.2 · pp. 21–22
  • As λ^-1 approaches zero, the ideal exponential tilt follows the natural policy-gradient direction to first order, and the Doob correction represents the same local change in diffusion-drift space.The analysis describes the population policy sequence as monotonically improving under the cited policy-improvement result.Source: §3 · p. 16, §3.2 · p. 23

4 Theoretical Analysis

The analysis bounds critic estimation, diffusion-actor approximation, and their propagation into policy performance. Its guarantees are conditional on coverage, regularity, and sampling assumptions, and account for inherited actor error.

4 Theoretical Analysis

  • The analysis combines advantage-estimation and conditional-diffusion sampling errors to bound expected average policy suboptimality.The section organizes these results into critic estimation, actor approximation, and overall policy-performance bounds.Source: §4 · p. 23

4.1 Convergence Rate of Advantage Estimation

  • The advantage estimator is mean-square consistent uniformly over iterations, with its rate balancing neural approximation and statistical error under dependent observations.Exponential mixing affects the stated polynomial rate through logarithmic factors, while the discounted Bellman equation contributes a (1 − γ)^-2 factor.Source: §4.1 · pp. 24–25

4.2 Convergence Rate of the Amortized Conditional SF Actor

  • The actor-error bound separates neural regression and approximation, Euler discretization, finite-m SNIS, and inherited reference-drift error.Increasing training sample size alone is insufficient for convergence: m and T must be chosen jointly, and inherited error must also vanish.Source: §4.2 · p. 27

4.3 Overall Policy-Performance Guarantee

  • The overall guarantee propagates critic, drift-regression, SNIS, diffusion-discretization, and inherited-drift errors, yielding an approximate policy-iteration bound rather than a stationary-point guarantee.Direct comparison with gradient-based actor–critic sample complexities is not immediate because the analysis separates budgets K, n, m, and T.Source: §4.3 · p. 30

5 Proof Sketches

The proof sketches build the policy-performance guarantee from dependent-data critic bounds and diffusion-actor error decompositions. They then propagate those errors through policy iteration.

5.1 Proof Sketch of Theorem 4.1

  • The advantage-error proof controls action-value and state-value errors separately, then combines them with the initial advantage-error inequality and Theorem 4.1 parameter choices.The action-value bound is obtained from Bellman-residual excess risk.Source: §5.1 · p. 30, §5.1.2 · p. 33

5.1.1 Action-Value Function Estimation Error

  • Action-value error is bounded through Bellman-residual excess risk, split into statistical and neural approximation components.The statistical component uses empirical-process bounds for dependent data; approximation uses neural-network bounds.Source: §5.1.1 · p. 31

5.1.2 State-Value Function Estimation Error

  • The state-value error combines Monte Carlo integration fluctuation with inherited action-value error; the fluctuation has conditional mean zero given the fitted critic and state.The action-value bound controls the inherited component.Source: §5.1.2 · p. 33

5.2 Proof Sketch of Theorem 4.2

  • The proof decomposes actor error into drift estimation and time discretization, then combines their bounds to obtain Wasserstein and KL error bounds.Drift estimation includes SNIS, actor statistical and approximation errors, and inherited reference discrepancy.Source: §5.2 · pp. 34–36

5.3 Proof Sketch of Theorem 4.3

  • The performance-gap proof bounds critic error with Theorem 4.1 and controls policy-approximation terms using Wasserstein and KL bounds under Assumption 4.7.These bounds are used to propagate errors through policy iteration.Source: §5.3 · pp. 36–37

6 Numerical Experiments

The experiments evaluate SFAC’s offline-to-online performance on D4RL tasks and its sampling accuracy and sensitivity to computational budgets. The conclusion summarizes the method and finite-sample analysis, while the appendix lists supporting derivations and proofs.

6.1 Experimental Setup

  • Evaluation covers six D4RL tasks across medium and medium-replay variants of HalfCheetah, Hopper, and Walker2d, reporting normalized returns.Main experiments use 500K online steps, three seeds, and the average of the final five evaluations as the primary online metric; published baselines are contextual rather than budget-matched.Source: §6.1 · pp. 38–39

6.2 Offline-to-Online Performance

  • SFAC’s final-window score exceeded its offline initialization on all six tasks, with the average increasing from 34.09 to 83.81.The online score averages the last five evaluations; published baselines are contextual because protocols and interaction budgets differ.Source: Table 2 · p. 39, §6.2 · p. 39

6.3 Sampling Accuracy and Sensitivity

  • Synthetic tests use single-state problems with two- and four-component Gaussian-mixture reference policies to isolate sampling error from critic learning and environment noise.The four-component mixture has separated modes.Source: §6.3 · p. 40
  • At (T, m) = (8, 32), SFAC samples recovered the multimodal support and relative mode-mass shifts of both prescribed targets.Figure 2 compares reference, tilted-target, and sampled distributions for two- and four-component mixtures.Source: §6.3 · p. 41, Figure 2 · p. 41
  • Increasing m from 8 to 16 gave the largest Monte Carlo improvement, while discretization error continued falling through T = 32 for the four-component target.After T = 16, error was essentially flat for the two-component target; curves plateaued after m = 16, and (T, m) = (8, 32) was selected as an accuracy–computation trade-off.Source: §6.3 · p. 42, Figure 3 · p. 43

7 Conclusion

  • The conclusion presents conditional Schrödinger–Föllmer dynamics for KL-regularized diffusion-policy improvement and bounds separating critic, regression, SNIS, discretization, and inherited actor errors.The stated bound concerns expected average policy suboptimality under suitable conditions.Source: §7 · p. 44

Appendix

  • The appendix collects derivations and proofs, including the diffusion natural-policy-gradient interpretation, conditional drift estimator, and policy-performance analysis.It also includes auxiliary results and the offline pretraining algorithm.Source: Appendix · p. 44

A Connection to Natural Policy Gradient

The finite KL-regularized policy update has the natural policy-gradient direction as its first-order limit. The Doob correction represents that local change in diffusion-drift space.

A Connection to Natural Policy Gradient

  • The finite KL-regularized update has the natural policy-gradient direction as its first-order limit.The result interprets the Doob correction as an approximation to the policy gradient.Source: A Connection to Natural Policy Gradient · p. 45
  • The first-order Doob drift correction represents the natural policy-gradient direction in the Schrödinger–Föllmer process.The passage identifies this as the diffusion-based representation of the same local update.Source: A Connection to Natural Policy Gradient · p. 45

B.1 Derivation of the Reference SDE Drift Term in (9)

  • The reference diffusion drift is expressed through a heat-semigroup representation and a Brownian-bridge representation.The derivation takes a logarithmic gradient of the density representation.Source: §B.1 · pp. 49–50

B.2 Derivation of the Doob h-transform SDE Drift Term in (11)

  • The Doob h-transform changes the reference drift by adding the logarithmic gradient of h.The displayed drift identity writes the transformed drift as the reference drift plus ∇y log hk(y, t; s).Source: §B.2 · p. 50

B.3 Derivation of the Drift Estimator (12)

The drift estimator subtracts reference and advantage-tilted posterior means, estimated with paired self-normalized importance sampling. The accompanying proofs control minimax, empirical-process, and approximation terms.

B.3 Derivation of the Drift Estimator (12)

  • The Doob h-transform replaces the reference posterior mean with its advantage-tilted counterpart, and paired SNIS estimates both using the same endpoint samples.Subtracting the paired estimates and dividing by 1 − t gives the drift estimator.Source: §B.3 · p. 52

C.1 Proof of Lemma 5.2

  • The Lemma 5.2 proof uses empirical minimax optimality to make one term nonpositive and bounds the other displayed terms.The passage identifies the third term as nonpositive and says the second, fourth, first, and fifth terms are bounded.Source: §C.1 · pp. 52–53

C.2 Proof of Lemma 5.3

  • The Lemma 5.3 proof uses independent conditional draws that preserve stationarity, alongside rescaling and product-cover arguments for the network classes.The proof also accounts for incomplete blocks and truncated tails.Source: §C.2 · pp. 53–54

C.3 Proof of Lemma 5.4

  • The Lemma 5.4 proof relates population risk to the squared Bellman residual, then uses Hölder-function approximation by ReLU networks on a truncated domain.The proof invokes the Bellman identity and Jensen’s inequality, and splits approximation error over the truncated domain and its complement.Source: §C.3 · pp. 55–56

C.4 Proof of Theorem 4.1

The proof combines component bounds to establish the theorem’s policy-suboptimality bound. Its supporting lemmas treat path discretization and finite-sample importance-sampling error.

C.4 Proof of Theorem 4.1

  • The theorem proof combines the preceding lemmas, and with NV ≍ n the mixing remainder and Monte Carlo error are of smaller order.The proof introduces Lemma 5.1 to relate Bellman residual to action-value error before proving Theorem 4.1.Source: §C.4 · pp. 57–58

D.1 Proof of Lemma 5.5

  • On the final interval [1 − 1/T, 1], the drift energy and synchronous squared error are each bounded by CdA/T.Novikov’s condition also holds for the comparison between exact and learned paths.Source: §D.1 · p. 59

D.2 Proof of Lemma 5.6

  • The Lemma 5.6 proof conditions on the fitted critic and actor and applies a self-normalized importance-sampling error decomposition.The lemma’s stated conditions include Assumptions 4.1 and 4.2.Source: §D.2 · p. 61

D.3 Proof of Lemma 5.7

The proofs control actor-regression excess risk, drift approximation, and the transfer from grid-based drift bounds to policy-level KL bounds.

D.3 Proof of Lemma 5.7

  • The Lemma 5.7 proof compares an empirical risk minimizer with network candidates and controls the resulting excess risk through empirical-process terms.The analysis conditions on the fitted quantities and uses the actor-regression sample.Source: §D.3 · pp. 63–64

D.4 Proof of Lemma 5.8

  • For Lemma 5.8, Hölder regularity supports network approximation, and clipping does not increase the approximation error.The resulting bound is uniform over fitted quantities, so the proof takes expectations to obtain (24).Source: §D.4 · pp. 66–67

D.5 Proof of Lemma 5.9

  • The Lemma 5.9 proof replaces continuous drift energy with its Euler-grid counterpart, then transfers the state expectation using Assumption 4.7.A uniformly bounded likelihood ratio makes the two KL directions equivalent up to a constant.Source: §D.5 · p. 67

D.6 Proof of Lemma 5.10

The proof balances diffusion-grid and SNIS sampling choices to control the drift-regression errors. With that balance, the displayed error terms and inherited reference error have matching order.

D.6 Proof of Lemma 5.10

  • The final bound follows by combining the approximation bound with the preceding decomposition and transferred bounds.The transfers apply to the learned and exact target paths up to a constant.Source: §D.6 · p. 68
  • With T ≍ n^α and m ≍ n^(ακρ), the displayed error terms are O(n^-α).The passage defines κρ = max{1, 4/(2 + δρ)}.Source: §D.6 · p. 69
  • The proof’s final assertion holds when the inherited reference error has the same order as the displayed rate.Source: §D.6 · p. 69

D.7 Proof of Theorem 4.2

The proof controls actor approximation and sampling errors to bound the learned policy’s distance from the target. It then combines these controls with the overall convergence argument.

D.7 Proof of Theorem 4.2

  • The offline drift-regression result uses a ReLU empirical risk minimizer under the stated assumptions and dimension-dependent network choices.The proof reduces squared prediction error to conditional-mean regression and controls the design and label second moments by C(1 + log T).Source: §D.7 · pp. 70–71
  • The actor proof couples the learned Euler scheme to the exact target diffusion and uses Wasserstein and path-law KL bounds.The KL argument invokes Girsanov’s theorem; terminal projection transfers the path-law bound to the action law.Source: §D.7 · pp. 72–73

E Proof of the Overall Convergence Rate

E.1 Proof of Theorem 4.3

The proof bounds one-step policy improvement while accounting for critic and actor errors. Under the balanced choices, the actor term decays no more slowly than the critic term, while inherited initialization error is smaller order.

E.1 Proof of Theorem 4.3

  • The proof compares the practical policy with the plug-in target and bounds the resulting error terms using the stated norms.It inserts the practical policy into the performance differences and applies the performance-difference lemma.Source: §E.1 · p. 74
  • The proof sums the one-step bound across iterations and uses total variation control to relate the plug-in target and practical policy.Source: §E.1 · p. 75
  • Under the balanced choices, the actor term decays no more slowly than the critic term, while inherited initialization error is smaller order.Source: §E.1 · p. 76

F Auxiliary Results

The auxiliary results provide tools for neural-network approximation, concentration, and change of measure. The offline pretraining algorithm fits a reference Schrödinger–Föllmer drift from Brownian-bridge data.

F Auxiliary Results

  • The neural-network results are stated on compact domains; larger boxes require rescaling, and unbounded designs require tail control.Source: F Auxiliary Results · p. 76
  • The auxiliary results include a change-of-measure statement under which the process is Brownian motion under the tilted measure.Source: F Auxiliary Results · pp. 77–78

G Offline Pretraining Algorithm

Related papers

Loading 2609.33239v1…