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

BayesFP: Posterior Estimation for Flow-Based Policies via Feynman-Kac Sampling
Jun 19, 2026Sreevardhan Sirigiri, Weiming Zhi, Fabio RamosRobot Policy Learning Trajectory Generation
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
DEFlow: Deep Equilibrium Flow Maximum-Entropy Reinforcement Learning
Oct 12, 2026Ruslan Ageev, Huaze Tang, Runqing Miao +2Reinforcement Learning Robot Policy Learning
Tsinghua University · Peng Cheng Laboratory
ScoRe-Flow: Complete Distributional Control via Score-Based Reinforcement Learning for Flow Matching
Apr 13, 2026Xiaotian Qiu, Lukai Chen, Jinhao Li +3RL Control Robotic RL
Zhejiang University, Hangzhou, China · Shanghai Innovation Institute · Shanghai Jiao Tong University, Shanghai, China +1

Jun 19, 2026J/K move · Enter open · S save

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

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

DEFlow: Deep Equilibrium Flow Maximum-Entropy Reinforcement Learning

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