PrT-Flow: Probability-Tractable Flow Matching for Reinforcement Learning
Congzhong Sun, Ning Tan
cs.RO
Oct 13, 2026 · v1
cs.AI
TL;DR
The paper's two propositions on Flow Matching density tracking and the equivalence of CET minimizers are formally verified in Lean 4.
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 .
Problem
Diffusion and Flow Matching robot policies represent multimodal action distributions, but they do not readily provide action probabilities. RL algorithms such as SAC and PPO need those probabilities. A finite-step Euler sampler is not generally invertible, and computing its density is costly.
Approach
PrT-Flow parameterizes each Euler sampling step with a Jacobian-Tractable Transformer (JTT). JTT combines causal monotone splines with orthogonal coordinate transforms, so each step is invertible with an analytic Jacobian determinant. Log-densities therefore accumulate during the same forward pass that generates an action. Conditional Euler-step Training (CET) supervises finite one-step displacements rather than instantaneous velocities. Propositions 4.1 and 4.2 on density tracking and on CET's minimizers are formally verified in Lean 4.
Results
CET is proved to have the same minimizers as fitting the Euler-solver updates of Flow Matching, under stated assumptions. PrT-Flow reaches 96% success on Robomimic Lift with imitation learning, matching the Diffusion Policy reference. It also reaches 96% success on OGBench Cube-Triple after offline-to-online SAC-style training.
