Constrained Flow Policy Updates: A Generalized Schrödinger Bridge View

Boyang Li; Matthew Kim; Sylvia Herbert

cs.LGSep 26, 2026

Constrained Flow Policy Updates: A Generalized Schrödinger Bridge View

Authors: Boyang Li, Matthew Kim, Sylvia Herbert

Organizations: University of California San Diego

Abstract

Online safe reinforcement learning (RL) seeks policies that maximize reward while satisfying safety constraints. Reward and safety can induce multimodal action distributions, challenging the prevailing primal-dual methods: Gaussian actors may collapse onto a single suboptimal mode, and optimization over the nonconvex Lagrangian landscape can be unstable. Diffusion and flow policies can represent such distributions, but recent work with a diffusion actor relies on estimating and matching the score of an augmented-Lagrangian target policy. Instead, we differentiate the augmented objective directly through the generation path of a flow policy, so no score needs to be estimated. Because a flow policy lacks a readily available action log-density for entropy regularization, we build on the density-free kinetic-energy regularizer of FLAC, a recent reward-only method, and propose Reparameterized Augmented-Lagrangian Flow Actor with Least Energy (RAFALE), an off-policy actor-critic method for safe RL. We formulate its update as a constrained one-ended generalized Schrödinger bridge and show that, for each source draw, this path-space problem is exactly an entropy-regularized problem in action space. At positive noise, its solution reweights the reward-only action distribution only where the estimated cost exceeds a threshold set by the Lagrange multiplier. As the noise vanishes, the optimal value converges to that of a least-energy map objective that the flow policy optimizes directly. Across seven Safety-Gymnasium tasks, RAFALE achieves competitive reward with mean final cost within budget on every task, whereas strong baselines trade one for the other; ablations support the necessity of both its augmented objective and its flow actor.

Figures & tables

Method	Setting	Actor	Constraint	Actor update
PID-Lag ( Stooke et al., 2020 )	online	Gaussian	expected cost	policy gradient
CAL ( Wu et al., 2024 )	online	Gaussian	expected cost	reparameterization
RCRL ( Yu et al., 2022 )	online	Gaussian	state-wise	reparameterization
RESPO ( Ganai et al., 2023 )	online	Gaussian	state-wise	policy gradient
QSM ( Psenka et al., 2024 )	online	diffusion	none	Q-value guidance
RSM, RFM, FLAG ( Ma et al., 2025 ; Li et al., 2026b ; Kim et al., 2026 )	online	diffusion, flow	none	weighted matching

Table 1: Positioning by setting, actor class, constraint type, and actor-update route; a related taxonomy appears in Gao et al. (2026) . RAFALE occupies the online, constrained, generative-policy, reparameterization cell.

Task	Metric	RESPO	CAL	RCRL	ALGD	RAFALE (ours)
Swimmer	Reward	$36\pm 3$	$41\pm 3$	$20\pm 15$	$\bm{52\pm 13}$	$\bm{93\pm 23}$
Swimmer	Cost	$7.5\pm 0.6$	$16.6\pm 6.7$	$4.0\pm 7.6$	$2.9\pm 3.5$	$5.4\pm 2.7$
HalfCheetah	Reward	$\bm{2344\pm 320}$	$2650\pm 111$	$1796\pm 2493$	$2817\pm 69$	$\bm{3025\pm 55}$
HalfCheetah	Cost	$9.8\pm 5.9$	$15.5\pm 14.7$	$195.2\pm 434.8$	$22.4\pm 44.7$	$3.5\pm 4.7$
Hopper	Reward	$\bm{1073\pm 308}$	$1664\pm 24$	$1037\pm 575$	$1702\pm 15$	$\bm{1655\pm 46}$
Hopper	Cost	$3.1\pm 5.1$	$12.2\pm 6.2$	$0.2\pm 0.3$	$15.7\pm 12.3$	$0.3\pm 0.4$

Table 2: Final episodic reward and cost (mean

\pm

SD, five seeds;

b=10

, or

25

on CarGoal); RESPO trains for 9M–10M environment steps, the others for 1M–3M ( Appendix E ). Red: mean cost above budget. Among entries within budget, bold blue and bold green mark the highest and second-highest reward.

Figure 1: Learning curves (mean

\pm 1

standard error over five seeds): test reward (upper rows) and test cost (lower rows) against environment steps. Dashed lines mark the budget, and cost axes are truncated. RESPO, trained for up to 10M steps, appears in Figure 2 .

Task	Variant	Reward	Cost	Peak	Excess	Feas.
Hopper	Augmented $A$ (RAFALE)	$\bm{1655\pm 46}$	$0.3\pm 0.4$	$\bm{174.6}$	$\bm{92.3\pm 15.9}$	5/5
	Hinge $C$	$1547\pm 33$	$0.8\pm 1.2$	902.9	$251.2\pm 30.3$	5/5
	Standard $U$	$1208\pm 573$	$0.0\pm 0.1$	767.8	$240.6\pm 70.2$	5/5
Walker2d	Augmented $A$ (RAFALE)	$\bm{3177\pm 26}$	$1.3\pm 1.3$	$\bm{125.9}$	$\bm{20.3\pm 3.8}$	5/5
	Hinge $C$	$3075\pm 89$	$0.0\pm 0.0$	390.1	$55.9\pm 28.4$	5/5
	Standard $U$	$2939\pm 65$	$1.2\pm 2.1$	443.8	$83.7\pm 42.9$	5/5

Table 3: Lagrangian ablation (five seeds; mean

\pm

SD where shown): final reward and cost, peak evaluated cost (Peak), integrated excess above budget (Excess), and feasible final seeds (Feas.). Bold blue: highest reward; bold black: lowest Peak and Excess, most feasible seeds.

Task	Flow (RAFALE)	Gaussian	GMM
Hopper	${\color[rgb]{0,0.2734,0.6602}\bm{1655\!\pm\!46}}\;/\;\bm{0.3\!\pm\!0.4}$ (5)	$1408\!\pm\!506\;/\;7.6\!\pm\!17.0$ (4)	$1645\!\pm\!51\;/\;1.5\!\pm\!2.4$ (5)
Walker2d	$3177\!\pm\!26\;/\;\bm{1.3\!\pm\!1.3}$ (5)	$3151\!\pm\!15\;/\;3.5\!\pm\!4.0$ (4)	${\color[rgb]{0,0.2734,0.6602}\bm{3190\!\pm\!17}}\;/\;1.4\!\pm\!1.1$ (5)
Humanoid	${\color[rgb]{0,0.2734,0.6602}\bm{6513\!\pm\!27}}\;/\;\bm{3.2\!\pm\!3.1}$ (5)	$6510\!\pm\!51\;/\;8.1\!\pm\!2.6$ (4)	$6503\!\pm\!116\;/\;8.5\!\pm\!5.1$ (3)

Table 4: Actor families with their native regularizers: final reward / cost (mean

\pm

SD, five seeds; feasible seeds in parentheses). Bold blue: highest reward; bold black: lowest cost.

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

Route	Critic signal	What is trained	Extra requirement
Q-value guidance	$\nabla_{a}Q$	score network, toward $\nabla_{a}Q/\beta$	scores at every noise level
Weighted matching	$Q$ of each candidate (also $\nabla_{a}Q$ in RFM)	weighted denoising or flow-matching loss	candidate actions and weights
Policy gradient	advantage	clipped ratio objective, on-policy	surrogate for the density ratio
Reparameterization	$\nabla_{a}Q$ at the generated action	sampler, by backpropagation	density-free regularizer

Appendix

Table 6

Setting	Value	Setting	Value
Batch size	256	Replay capacity	1M
Hidden network	$2\times 256$	Gradient-norm cap	10
Actor/critic/ $\alpha$ LR	$3{\times}10^{-4}$	Initial $\log\alpha$	$-2$
Environment warmup $t_{0}$	5k	Dual warmup $t_{\lambda}$	200k
Update cycle	16 env. steps	Updates per cycle $n_{u}$	$16\times\mathrm{UTD}$
Policy delay	2 updates	Dual cadence $d_{\lambda}$	2000 env. steps

Appendix

Table 5: Common training settings; LR denotes learning rate.

Task	$b$	Steps	UTD	Anneal start
Swimmer	10	1M	1	700k
HalfCheetah	10	1M	1	700k
Hopper	10	2M	2	1.7M
Walker2d	10	1M	1	700k
Ant	10	3M	1	2.7M
Humanoid	10	3M	2	2.4M

Appendix

Table 6: Task-level settings for the reported configurations.

Task	$\eta_{\lambda}$	$\eta_{p}$	$\lambda_{\max}$	Warm-start value $z_{\mathrm{w}}$
Swimmer	.001	0	0.5	0.3999828
HalfCheetah	.001	0	4.2	0
Hopper	.001	0	6.7	0
Walker2d	.001	0	16.5	1.654
Ant	.001	0	15.0	0
Humanoid	.001	.05	9.8	0

Appendix

Table 7: Task-specific multiplier settings used by the reported full method.

Setting	Swimmer	HalfCheetah	Hopper	Walker2d	Ant	Humanoid	CarGoal
Steps	1M	1M	2M	1M	3M	3M	1M
$c_{\mathrm{CAL}}$	10	10	100	100	100	1000	10
$k_{\mathrm{CAL}}$	0.5	0.5	0.5	0.5	0.5	0	0.5
$N_{Q_{c}}$	8	8	8	8	8	4	4
Max-over-heads target	no	no	no	no	no	yes	no
Critic layer norm	yes	yes	yes	yes	yes	no	yes

Appendix

Table 8: CAL settings per task.

N_{Q_{c}}

is the cost-critic ensemble size and

k_{\mathrm{CAL}}

the weight of its standard deviation in CAL’s cost estimate. Mask “t/t” denotes a cost-critic bootstrap mask of zero on termination or truncation; “1” a mask forced to one.

Setting	Swimmer	HalfCheetah	Hopper	Walker2d	Ant	Humanoid	CarGoal
Steps	1M	1M	2M	1M	3M	3M	1M
$\gamma_{h}$	0.99	1.0	1.0	1.0	1.0	1.0	1.0
Learning rates	decay (1M)	decay (3M)	constant	decay (3M)	decay (3M)	decay (3M)	decay (3M)
Multiplier update	every 12	every 12	every step	every 12	every 12	every 12	every 12
Multiplier cap	none	none	100	none	none	none	none
Safety-critic layer norm	yes	no	no	no	no	no	no

Appendix

Table 9: RCRL settings per task.

Task	RESPO	CAL	RCRL	ALGD	RAFALE
Swimmer	5	1	4	5	5
HalfCheetah	2	2	4	4	4
Hopper	4	1	5	1	5
Walker2d	2	4	4	2	5
Ant	5	4	4	1	5
Humanoid	5	5	4	5	5

Appendix

Table 10: Seeds, out of five, whose final mean episodic cost is at or below the budget, for the endpoints of Table 2 . RAFALE’s HalfCheetah seed 0 ends at 10.9 against a budget of 10.

Figure 2: All five methods under the convention of Figure 1 , with environment steps on a logarithmic axis so that RESPO’s 10M-step trajectories fit alongside the others; the dotted line marks the step budget of the off-policy methods.

Task	Step	RESPO reward	RESPO cost	RESPO final reward
Swimmer	1M	$17\pm 9$	$16.6\pm 9.8$	36
HalfCheetah	1M	$434\pm 231$	$1.9\pm 2.1$	2344
Hopper	2M	$776\pm 292$	$0.5\pm 0.7$	1073
Walker2d	1M	$373\pm 86$	$1.2\pm 0.6$	2925
Ant	3M	$117\pm 41$	$0.7\pm 0.3$	1950
Humanoid	3M	$916\pm 326$	$0.3\pm 0.3$	5380

Appendix

Table 11: RESPO at RAFALE’s step budget (mean

\pm

SD over five seeds), read at that step, interpolating linearly between the two bracketing checkpoints where it was not evaluated.

Figure 3: Learning curves of the three Lagrangian variants: mean

\pm 1

standard error over five seeds at the evaluation steps shared by all seeds; dashed lines mark the budget. The augmented Lagrangian

L_{A}

reduces transient cost excursions relative to the hinge control

L_{C}

and the standard Lagrangian

L

, without a systematic loss of final reward.

Figure 4: Learning curves of the flow, tanh-Gaussian, and eight-component tanh-GMM actors: mean

\pm 1

standard error over five seeds at the evaluation steps shared by all seeds, with 50 evaluation episodes per checkpoint; dashed lines mark the budget.

Figure 5: Frozen-critic multiplier response. For each fixed $\lambda$ , a flow actor is trained from the same random initialization (dotted) for the same number of updates. Larger multipliers shift the generated-action cost-critic CDF left. The dashed lines mark the budget $h$ in cost-critic units, and $\hat{\lambda}$ is the multiplier stored in the checkpoint.

Flow-Based Policy Learning

See all 37 papers

Safe Flow Q-Learning: Offline Safe Reinforcement Learning with Reachability-Based Flow Policies
Mar 16, 2026Mumuksh Tayal, Manan Tayal, Ravi PrakashFlow Matching Offline RL
Indian Institute of Science, India
Augmented Lagrangian Multiplier Network for State-wise Safety in Reinforcement Learning
May 1, 2026Jiaming Zhang, Yujie Yang, Yao Lyu +2Constrained RL Augmented Lagrangian Method
School of Vehicle and Mobility & College of AI, Tsinghua University, Beijing, 100084, China
Reachability-Aware Diffusion Policy Optimization
Oct 5, 2026Hikmet Simsir, Kutay Demiray, Ozgur S. OguzConstrained RL Diffusion Policy
Department of Computer Engineering, Bilkent University

Mar 16, 2026J/K move · Enter open · S save

Safe Flow Q-Learning: Offline Safe Reinforcement Learning with Reachability-Based Flow Policies

Mumuksh Tayal, Manan Tayal, Ravi Prakash

Indian Institute of Science, India

Flow Matching Offline RL Safe RL Flow-Based Policy Learning

Open paper AskPDF arXiv

Offline safe reinforcement learning (RL) seeks reward-maximizing policies from static datasets under strict safety constraints. Existing methods often rely on soft expected-cost objectives or iterative generative inference, which can be insufficient for safety-critical real-time control. We propose Safe Flow Q-Learning (SafeFQL), which extends FQL to safe offline RL by combining a Hamilton--Jacobi reachability-inspired safety value function with an efficient one-step flow policy. SafeFQL learns the safety value via a self-consistency Bellman recursion, trains a flow policy by behavioral cloning, and distills it into a one-step actor for reward-maximizing safe action selection without rejection sampling at deployment. Empirically, SafeFQL trades modestly higher offline training cost for substantially lower inference latency than diffusion-style safe generative baselines, which is advantageous for real-time safety-critical deployment. Across boat navigation, and Safety Gymnasium MuJoCo tasks, SafeFQL matches or exceeds prior offline safe RL performance while substantially reducing constraint violations.

Abstract

Flow-Based Policy Learning

Explore similar work

Safe Flow Q-Learning: Offline Safe Reinforcement Learning with Reachability-Based Flow Policies

Augmented Lagrangian Multiplier Network for State-wise Safety in Reinforcement Learning

Reachability-Aware Diffusion Policy Optimization