cs.LGOct 3, 2026

Score-Calibrated Flow for Sampling from Unnormalized Densities with Applications to Generative Online Reinforcement Learning

Authors: Zeyang Li, Yunan Wang, Risheek Garrepalli, Mohammad Ghavamzadeh, Navid Azizan

Organizations: Massachusetts Institute of Technology · Qualcomm AI Research

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Diffusion and flow models provide expressive policy classes for online reinforcement learning (RL), enabling multimodal behaviors and improved performance. However, training these policies remains challenging: the critic specifies the desired policy as an unnormalized Boltzmann density but does not provide direct samples from it. Many existing methods rely on importance sampling to construct training signals, which can suffer from high variance, increasing computational cost and destabilizing training. We propose Score-Calibrated Flow (SCF), a simple and efficient algorithm for training generative models to sample from unnormalized densities without importance sampling or backpropagation through the sampling trajectory. We learn the desired flow by enforcing self-consistency, bypassing target posterior mean estimation. By jointly exploiting the prescribed target score and the structure of flow matching, we establish these self-consistency requirements as score-calibrated optimality conditions, first for the terminal density and then for the trainable velocity field. We prove that their unique solutions are, respectively, the target density and the ideal flow model that conditional flow matching (CFM) would recover if target samples were available. We formulate the velocity condition as a fixed-point equation and exploit its conditional-expectation structure to construct a stop-gradient objective for enforcing it. The resulting training procedure retains the scalable sample-interpolate-regress structure of CFM despite the absence of target samples, using endpoints generated by the current flow. For online RL, the critic gradient supplies the target score at the generated actions, yielding a direct approach to actor training. Experiments on RL benchmarks demonstrate that SCF matches or improves upon state-of-the-art generative-policy baselines, while substantially reducing training time.

Figures & tables

Figure 1: Visualization of samples from the five trained models on the eight-Gaussian target. SCF achieves the lowest discrepancy from ground-truth samples while being the cheapest to train and to sample from. The detailed descriptions are provided in Appendix C.1 .
Figure 2: Training curves on six tasks from HumanoidBench. The solid lines correspond to the mean across five seeds, and the shaded regions represent minimum-to-maximum range.
Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

 DeepMind Control SuiteHumanoidBench
Parallel environments115050
Environment steps per training step115050
Replay buffer size2.5×1052.5\times 10^{5}1×1061\times 10^{6}
Batch size256256256256
Hidden layers2×2562\times 2563×2563\times 256
Normalizationnonebatch renormalization
Appendix
Table 1: Shared settings of the RL experiments.
Figure 3: Training curves on eight tasks from the DeepMind Control Suite. The solid lines correspond to the mean across five seeds, and the shaded regions represent minimum-to-maximum range.
TaskSACSCF (ours)QSMRFMQFlexSDACDIMEQVPOMaxEntDPDQS
walker-run0.470.871.082.671.942.433.244.406.117.51
h1-run1.021.121.212.082.993.153.334.586.2811.66
Appendix
Table 2: Milliseconds per training iteration on walker-run (DeepMind Control Suite) and h1-run (HumanoidBench), measured on one NVIDIA GeForce RTX 5090 with batch size 256256 and averaged over 10001000 steps.
Figure 4: Training curves of SCF on the h1-run task under different design choices. Left: experiments on NFEs for sampling. Right: experiments on the coefficients of the regression target. The solid lines correspond to the mean across five seeds, and the shaded regions represent minimum-to-maximum range.
More in this topic

Sampling from Unnormalized Distributions

See all 17 papers

Explore similar work

CardsList
May 22, 2026

Score-Based One-step MeanFlow Policy Optimization

Diffusion and flow matching have emerged as expressive policy classes in reinforcement learning, but their reliance on multi-step denoising imposes substantial computational overhead at inference time, which is particularly problematic in online RL. MeanFlow offers a promising alternative by learning an average velocity field that maps noise to data in a single network evaluation. However, MeanFlow typically requires samples from the target distribution to construct its target velocity field, which are unavailable in online RL. We propose Score-Based One-step MeanFlow Policy Optimization (SOM), an actor-critic algorithm that resolves this by constructing the target velocity field directly from the Q-function via score estimation and a probability flow ODE, thereby concentrating probability mass on high-value modes. In the fully online RL setting, SOM achieves state-of-the-art performance on locomotion tasks with a single generation step, while substantially reducing both training and inference time compared to prior diffusion- and flow-matching-based policies.

Kyungyoon Kim, Donghyeon Ki, Hee-Jun Ahn +1

Korea University, Decision Making Lab · Gauss Labs Inc.

Flow MatchingMeanFlowPolicy OptimizationMean FlowRL for Diffusion ModelsFlow-Based Policy Learning
AskPDF
Jun 29, 2026

FlowAWR: Online Adaptive Flow Reinforcement via Advantage-Weighted Rectification

Aligning generative flow models on continuous spaces via online reinforcement learning is constrained by intractable trajectory likelihoods. Existing density-approximated policy gradient methods rely on stochastic SDE samplers to construct tractable transition kernels, which introduce training-inference inconsistencies and necessitates Classifier-Free Guidance (CFG). While implicit frameworks such as DiffusionNFT directly optimize forward-process velocity fields, its heuristic fixed-magnitude corrections prevent optimization strength from relative intra-group quality. We propose Flow Advantage-Weighted Rectification (FlowAWR), a paradigm that recasts continuous generative policy optimization as supervised regression toward a theoretically optimal velocity field. Starting from the optimal policy of a KL-constrained reward maximization, FlowAWR derives the optimal velocity field that admits a magnitude-aware, advantage-weighted rectification form, yielding SDE-free optimization and CFG-free generation. In comparative evaluations on SD3.5-Medium, FlowAWR achieves improved alignment performance alongside a 2×\times to 5×\times convergence acceleration over DiffusionNFT (e.g., reaching a 24.12 PickScore in 1.2k steps, versus 23.82 in 2.0k steps for DiffusionNFT and 23.50 in >>4k steps for FlowGRPO). Under multi-reward constraints, FlowAWR sustains generation quality, satisfying structural rules while maintaining stable out-of-domain performance. Code at FlowAWR.

Zheming Fu, Ruizhe He, Wei Shang +5

Beihang University · Joy Future Academy · Zhongguancun Academy +2

Flow MatchingGenerative ModelingFlow-Based Generative ModelingRL for Generative Models
AskPDF
Oct 12, 2026

A Survey of Flow Matching in Reinforcement Learning

Flow Matching (FM) has recently emerged as a principled and efficient generative modeling framework for reinforcement learning (RL), enabling expressive, multimodal policy parameterizations via deterministic probability transport. Compared to diffusion-based policies that rely on stochastic denoising chains, FM uses sampling based on ordinary differential equations (ODEs), with learned velocity fields, which can substantially reduce inference latency and simplify the incorporation of RL objectives. As research in flow-based RL rapidly accelerates across offline continuous control, online fine-tuning, and foundation model alignment, the literature has become highly fragmented. In this survey, we provide a comprehensive taxonomy of flow-matching approaches in reinforcement learning. We organize the literature along two axes: the target distribution being modeled (e.g., action policies, value critics, transition dynamics) and the mechanism of RL signal integration (e.g., energy-weighted regression, flow-based policy gradients, and group relative policy optimization). Furthermore, we survey emerging frontiers such as discrete and non-Euclidean action spaces, provide a systematic comparative analysis against Gaussian and diffusion baselines, and outline critical open problems. Ultimately, this survey serves as a foundational roadmap for the next generation of generative reinforcement learning and alignment.

Nabuat Zaman Nahim, Fairoz Nower Khan, Peizhong Ju

Department of Computer Science, University of Kentucky, Lexington, Kentucky 40506, USA

Reinforcement LearningFlow-Based Generative ModelingRL for Generative Models
AskPDF