PrT-Flow: Probability-Tractable Flow Matching for Reinforcement Learning

Congzhong Sun; Ning Tan

cs.ROOct 13, 2026

PrT-Flow: Probability-Tractable Flow Matching for Reinforcement Learning

Authors: Congzhong Sun, Ning Tan

Organizations: School of Computer Science and Engineering, Sun Yat-sen University

Abstract

Diffusion models and Flow Matching can represent flexible distributions over robot actions, but many reinforcement learning algorithms also require the probability of each action. We introduce PrT-Flow (Probability-Tractable Flow Matching), which computes probability density while retaining multimodal, multi-step action generation. Its Jacobian-Tractable Transformer (JTT) makes each sampling step invertible and enables efficient computation of its Jacobian determinant. We train these finite steps with Conditional Euler-step Training (CET), which supervises their displacement over one step during sampling rather than their instantaneous velocity. We prove that, under certain assumptions, CET has the same minimizers as fitting the Euler-solver updates of Flow Matching. PrT-Flow achieves 96% success on Robomimic Lift after 300 epochs of imitation learning, matching the diffusion policy reference, and 96% success on OGBench Cube-Triple after 600k online steps. These results demonstrate that invertible networks with tractable probability densities can learn effective robot policies and support SAC-style actor-critic training. Our code is available at https://github.com/Sunlight-zero/PrT-Flow .

Figures & tables

Figure 1: The overview of PrT-Flow. (a) Conditioned on an observation, the Jacobian-Tractable Transformer (JTT) transforms standard Gaussian noise through

K

invertible steps of size

h=1/K

. Each step returns both the updated sample and the Jacobian determinant

D

, without constructing a dense Jacobian matrix. (b) The deformed grid illustrates how the probability density of a sample evolves under these steps. At each step, JTT computes the absolute Jacobian determinant

D

. The local volume is scaled by

D

and the corresponding density is divided by

D

, allowing us to track the probability density of a sample. (c) A final tanh transformation maps the generated sample to a bounded action chunk

a

with its probability density

p(a)

. SAC-style actor–critic training then uses sampled actions, their log-densities, and learned Q-values to update the FM policy.

Algorithm 1 Forward sampling and action probability density

1: observation condition

c

, step count

K

, policy parameters

\theta

2: normalized action chunk

a

and

\log\pi_{\theta}(a\mid c)

h\leftarrow 1/K

x_{0}\sim\mathcal{N}(0,I_{d})

\ell\leftarrow\log\phi(x_{0})

6: for

k=0,\ldots,K-1

do 7:

t_{k}\leftarrow kh

8: Apply the composed JTT update to obtain

x_{k+1}=F_{\theta,h}(x_{k},t_{k},c)

and its log-determinant

\ell_{k}

\ell\leftarrow\ell-\ell_{k}

10: end for 11:

a\leftarrow\tanh x_{K}

12:

\ell\leftarrow\ell-\sum_{j=1}^{d}\log(1-a_{j}^{2})

13: return

a,\ \log\pi_{\theta}(a\mid c)\leftarrow\ell

Table 2

Figure 2: The architecture of Jacobian-Tractable Transformer (JTT). (A) A complete update applies

M

map heads in sequence and returns the next noisy sample

x_{k+1}

and the absolute Jacobian determinant

D

. (B) A causal Transformer with

L

decoder blocks uses shifted scalar tokens to predict monotone spline parameters. Cross-attention provides time and observation information. The projected outputs are converted to bounded parameter rates and multiplied by

h

to obtain offsets from the identity settings. These offsets determine the spline parameters. Between the two orthogonal rotations, the splines transform the unshifted scalar values. (C) For a fixed coordinate prefix, time, and observation,

y_{i}

determines where its spline is evaluated, but not the spline’s shape. The positive analytic derivatives form the diagonal of the lower-triangular spline Jacobian, and their product gives the head determinant without evaluating off-diagonal entries. The orthogonal rotations preserve its magnitude. At step

k

, summing the head log-determinants at their respective inputs gives

\ell_{k}=\log D

Algorithm 2 Conditional Euler-step Training

1: demonstration windows after normalization and inverse-tanh transformation

\mathcal{D}_{u}=\{(c,u_{D})\}

, initialized

\theta

, step count

K

, batch size

B

, update budget

U

2: trained policy parameters

\theta

h\leftarrow 1/K

4: for

r=1,\ldots,U

do 5: Sample a minibatch

\{(c_{i},u_{D,i})\}_{i=1}^{B}

from

\mathcal{D}_{u}

6: for

i=1,\ldots,B

do 7: Draw

z_{i}\sim\mathcal{N}(0,I_{d})

and

k_{i}\sim\mathrm{Unif}\{0,\ldots,K-1\}

t_{i}\leftarrow k_{i}h

;

X_{i}\leftarrow(1-t_{i})z_{i}+t_{i}u_{D,i}

;

V_{i}\leftarrow u_{D,i}-z_{i}

q_{i}\leftarrow\bigl(F_{\theta,h}(X_{i},t_{i},c_{i})-X_{i}\bigr)/h

10: end for 11:

\mathcal{L}\leftarrow\frac{1}{Bd}\sum_{i=1}^{B}\|q_{i}-V_{i}\|_{2}^{2}

12: Take one optimizer step on

\theta

using

\nabla_{\theta}\mathcal{L}

13: end for 14: return

\theta

Table 4

Policy	Evaluation epoch	Success rate
DDPM reference	300	96%
PrT-Flow (JTT, CET)	300	96%

Table 1: Lift imitation learning after 300 epochs. Both policies are evaluated on the same set of 50 episodes.

Policy	Dev. SR	Dev. return	Final SR	Final return
Flow-G, $\sigma=0.01$	88%	$-613.80$	92%	$-475.48$
Flow-G, $\sigma=0.05$	56%	$-1561.86$	50%	$-1635.42$
PrT-Flow, CET	94%	$-578.28$	96%	$-530.80$

Table 2: Cube-Triple Task 4 after 30k paired offline iterations and 600k online primitive environment steps. Development results use the same set of 50 evaluation episodes for all policies. Final results come from separate 50-episode evaluations after training.

Figure 3: Learning curves on Cube-Triple for PrT-Flow and Flow-G (

\sigma=0.01

and

0.05

). Panels A and B show success rate and mean return on the shared set of 50 development episodes. Lines connect evaluations without smoothing. At 600k online steps, success rates are 94% for PrT-Flow, 88% for Flow-G (

\sigma=0.01

), and 56% for Flow-G (

\sigma=0.05

). Values in the legend correspond to this endpoint, rather than each run’s best checkpoint.

Design choice	Configuration	Successful episodes
Supervised quantity	CET finite-step quotient	49/50
	Zero-step generator	33/50
Coordinate frame	Learned orthogonal frames	49/50
	Fixed initial random orthogonal frames	33/50

Table 3: Design ablations on Lift with

K=16

after 300 training epochs. Each block reports a separate pair of runs with matched training settings, with 50 evaluation episodes per model. Within each pair, only the design choice listed in the first column changes.

Supervised quantity / time sampling	Endpoint coordinates	Dev. SR	Dev. return	Final SR	Final return
Generator velocity / continuous	Action	88%	$-806.02$	94%	$-619.72$
CET quotient / grid	Generative space	94%	$-578.28$	96%	$-530.80$

Table 4: PrT-Flow with two offline training configurations, evaluated after 30k paired offline iterations and 600k online steps. Both use the same development and final evaluation sets, with 50 episodes in each set.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Lift	Cube-Triple
Batch size	256	256
Optimizer	AdamW	Adam
Learning rate	$10^{-4}$	$3\times 10^{-4}$
Weight decay	$10^{-6}$	—
AdamW $(\beta_{1},\beta_{2})$ / $\epsilon$	$(0.9,0.999)$ / $10^{-8}$	—
Cosine warm-up (optimizer steps)	100	—

Appendix

Table 5: Optimization settings for the main PrT-Flow experiments. Cube-Triple uses Adam for the actor, critics, and entropy temperature.

Comparison	FP32	FP64
$f^{-1}(f(x))$ vs. $x$	$9.54\times 10^{-7}$	$1.33\times 10^{-15}$
$f(f^{-1}(y))$ vs. $y$	$7.15\times 10^{-7}$	$1.78\times 10^{-15}$
Analytic vs. autodifferentiated log-derivative	$8.20\times 10^{-7}$	$1.55\times 10^{-15}$
Forward vs. inverse log probability	$4.53\times 10^{-6}$	$6.22\times 10^{-15}$

Appendix

Table 6: Numerical consistency of scalar RQS maps. Here

f

denotes a map with fixed parameters and step size. Each entry is the maximum absolute discrepancy across the five seeds and two step sizes; the last two rows also take the maximum over both directions.

Comparison	Metric	Discrepancy
Analytic vs. autodifferentiated local log-determinant	Maximum absolute difference	$\leq 9.2\times 10^{-12}$
FP32 vs. FP64 forward log probability	Maximum absolute difference	$1.20\times 10^{-4}$
FP32 vs. FP64 parameter gradient of mean log probability	Relative $L_{2}$ difference	$9.77\times 10^{-7}$

Appendix

Table 7: Numerical consistency of the main Cube-Triple CET policy at 600k online steps. FP32/FP64 comparisons use 32 real states with fixed noise; local Jacobian checks use 12 updates across three four-step sampling trajectories.

Explore similar work

Jun 19, 2026

BayesFP: Posterior Estimation for Flow-Based Policies via Feynman-Kac Sampling

Robots must generate trajectories that remain faithful to learned expert behavior while satisfying safety constraints and task-specific objectives specified only at inference time. We formulate constrained trajectory generation for pretrained diffusion and flow-matching policies as Bayesian posterior sampling, with the learned demonstration distribution as a prior and an inference-time, cost-derived likelihood tilting it toward feasible, optimal trajectories. To sample from this posterior without any retraining of the base policy, we leverage the Feynman--Kac corrector framework, originally formulated for diffusion models, and extend it to deterministic flow-matching policies. The result is a unified, inference-time, retraining-free sampler for diffusion and flow policies. We validate the approach on pretrained Diffusion Policy, GR00T-N1.6, and

π_{0.5}

checkpoints across simulated and real-world manipulation tasks, including planning around non-convex obstacles introduced at inference time, and show improvements over the base

π_{0.5}

on zero-shot tasks.

Sreevardhan Sirigiri, Weiming Zhi, Fabio Ramos

School of Computer Science, The University of Sydney, Australia · Australian Centre For Robotics, The University of Sydney, Australia · College of Connected Computing, Vanderbilt University, TN, USA. +1

Robot Policy Learning Trajectory Generation Diffusion Policy Flow-Based Generative Modeling Diffusion Posterior Sampling

Ask PDF

Oct 12, 2026

DEFlow: Deep Equilibrium Flow Maximum-Entropy Reinforcement Learning

Expressive flow policies can represent multiple high-value actions, but sampling them with few integration steps while retaining tractable action densities remains challenging. We introduce DEFlow, a deep equilibrium flow policy that replaces explicit Euler sampling with implicit Euler integration and solves all intermediate actions jointly as an equilibrium system. We combine implicit differentiation with a discrete change-of-variables formula for the final-action density, enabling off-policy training within soft actor-critic without backpropagating through solver iterations. DEFlow achieves an average relative improvement of approximately 10% in mean successful coverage over the strongest baseline across four Multi-Goal docking variants. It also achieves the highest mean time-averaged return on four of five manipulation tasks, while remaining broadly competitive across five MuJoCo tasks. These results support implicit flow policies for tasks requiring both diversity and precision, at the cost of additional nonlinear and linear solves. Primary

Ruslan Ageev, Huaze Tang, Runqing Miao +2

Tsinghua University · Peng Cheng Laboratory

Reinforcement Learning Robot Policy Learning Maximum Entropy RL

Ask PDF

Apr 13, 2026

ScoRe-Flow: Complete Distributional Control via Score-Based Reinforcement Learning for Flow Matching

Flow Matching (FM) policies have emerged as an efficient backbone for robotic control, offering fast and expressive action generation that underpins recent large-scale embodied AI systems. However, FM policies trained via imitation learning inherit the limitations of demonstration data; surpassing suboptimal behaviors requires reinforcement learning (RL) fine-tuning. Recent methods convert deterministic flows into stochastic differential equations (SDEs) with learnable noise injection, enabling exploration and tractable likelihoods, but such noise-only control can compromise training efficiency when demonstrations already provide strong priors. We observe that modulating the drift via the score function, i.e., the gradient of log-density, steers exploration toward high-probability regions, improving stability. The score admits a closed-form expression from the velocity field, requiring no auxiliary networks. Based on this, we propose ScoRe-Flow, a score-based RL fine-tuning method that combines drift modulation with learned variance prediction to achieve decoupled control over the mean and variance of stochastic transitions. Experiments demonstrate that ScoRe-Flow achieves 2.4x faster convergence than flow-based SOTA on D4RL locomotion tasks and up to 5.4% higher success rates on Robomimic and Franka Kitchen manipulation tasks.

Xiaotian Qiu, Lukai Chen, Jinhao Li +3

Zhejiang University, Hangzhou, China · Shanghai Innovation Institute · Shanghai Jiao Tong University, Shanghai, China +1

RL Control Robotic RL Reinforcement Learning RL for Generative Models

Ask PDF