Paper deep dive
PrT-Flow: Probability-Tractable Flow Matching for Reinforcement Learning
Congzhong Sun, Ning Tan
Intelligence
Status: not_run | Model: - | Prompt: - | Confidence: 0%
Entities (0)
Relation Signals (0)
No relation signals yet.
Cypher Suggestions (0)
No Cypher suggestions yet.
Abstract
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 this https URL .
Tags
Links
- Source: https://arxiv.org/abs/2610.17667v1
- Canonical: https://arxiv.org/abs/2610.17667v1
Trouble viewing inline? Open PDF directly →
Full Text
71,694 characters extracted from source content.
Expand or collapse full text
PrT-Flow: Probability-Tractable Flow Matching for Reinforcement Learning Congzhong Sun †thanks: Email: sunczh3@mail2.sysu.edu.cn. Ning Tan †thanks: Corresponding author. Email: tann5@mail.sysu.edu.cn. Affiliation: 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. 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/Kh=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)p(a). SAC-style actor–critic training then uses sampled actions, their log-densities, and learned Q-values to update the FM policy. 1 Introduction Learning robot manipulation from demonstrations requires a policy that maps observations, such as camera images or robot states, to actions. For a given observation, several motions may accomplish the same task, yet their average may not correspond to any valid motion. Therefore, the policy should represent a distribution over actions. With action chunking, it predicts a sequence of actions from the current observation and replans from a new observation after executing a specified number of actions. Given an observation c, a stochastic policy treats an action chunk A as a random vector following the conditional distribution π(⋅∣c)π(· c). Diffusion Policy learns this distribution from robot demonstrations using a conditional diffusion model (Chi et al., 2023). Flow Matching offers a related approach to probability transport, learning to move noise toward demonstration actions along a prescribed probability path (Lipman et al., 2023; Liu et al., 2023). PointFlowMatch uses this formulation in place of diffusion to learn robot manipulation policies (Chisari et al., 2024). Both approaches use flexible generative models to capture the multiple valid behaviors that a robot may exhibit. Imitation learning (IL) and reinforcement learning (RL) provide complementary ways to learn robot policies: IL trains a policy to reproduce demonstrations, while RL uses task rewards to improve behavior through interaction. In the actor–critic framework widely used for continuous control, the actor represents the policy and the critic estimates state or action values. Many RL algorithms also require policy probabilities. Soft Actor-Critic (SAC) uses the log probability of a sampled action in its entropy objective (Haarnoja et al., 2018), while Proximal Policy Optimization (PPO) uses the probability ratio for an action under the updated and previous policies (Schulman et al., 2017). Therefore, a generative policy must provide both action samples and their probabilities to support these algorithms. Flow Matching provides a sampling procedure, but computing the probability density of its actions requires additional work. Continuous normalizing flows track density changes by integrating the divergence of the velocity field (Chen et al., 2018). However, a robot policy generates actions through finitely many numerical steps, whose endpoint distribution may differ from that of the continuous flow. We seek to compute the density of each action chunk under this finite sampling procedure as the chunk is generated. We introduce PrT-Flow (Probability-Tractable Flow Matching; Figure 1), which parameterizes each sampling step with an invertible network whose Jacobian determinant is analytically tractable. The Jacobian-Tractable Transformer (JTT) combines causal conditioning with monotone scalar transformations. Each forward pass returns an updated sample and the corresponding determinant, allowing the policy to compute the probability of denoised actions using the change-of-variables formula. We train these steps with Conditional Euler-step Training (CET), which supervises their displacement over one sampling step. Each training example supervises a single step along an interpolation path between noise and a demonstration endpoint; at deployment, the policy composes the steps to generate an action chunk. PrT-Flow thus combines the computable probabilities of normalizing flows with local regression along Flow Matching paths, supporting multimodal, multi-step action generation while providing the action log probabilities used by SAC. The causal and invertibility constraints in JTT may affect its ability to represent effective behavior. We therefore evaluate PrT-Flow in both imitation learning and offline-to-online reinforcement learning on several benchmarks. PrT-Flow achieves high task success in these settings. We summarize our contributions as follows. 1. We propose PrT-Flow and its JTT architecture, which compute the probability density of an action during multi-step generation using tractable per-step Jacobian determinants. 2. We introduce CET, a training objective whose minimizers coincide with those of marginal Flow Matching over the same family of finite maps. We also show its equivalence to regression against Euler-solver updates of the marginal velocity field. 3. Our experiments on Lift and Cube-Triple show that the invertible architecture supports high success rates in both imitation learning and SAC-style online control. Our code is available at https://github.com/Sunlight-zero/PrT-Flow. 2 Related Work 2.1 Diffusion and Flow Matching Policies As discussed in the Introduction, Diffusion Policy generates action sequences by iterative denoising (Chi et al., 2023). Flow Matching and rectified flow learn velocity fields that transport noise to data (Lipman et al., 2023; Liu et al., 2023), and PointFlowMatch applies conditional Flow Matching to robot manipulation (Chisari et al., 2024). PrT-Flow retains multi-step generation to preserve the multimodal ability of Diffusion Policy and Flow Matching policies while providing the action probabilities needed for RL. 2.2 Reinforcement Learning for Generative Policies DPPO fine-tunes diffusion policies with PPO over an augmented denoising decision process (Ren et al., 2025). Its PPO ratios use probabilities of individual denoising transitions, rather than the marginal probability of the final action chunk. FQL distills a Flow Matching behavior policy into a separate one-step actor optimized with Q-values (Park et al., 2025b). FPO constructs proxy probability ratios from conditional Flow Matching losses (McAllister et al., 2025), while RLDT improves a pretrained flow policy through transport updates (Lei et al., 2026). These approaches enable policy improvement without directly evaluating the exact probability of the final action samples produced by the multi-step generator. PrT-Flow instead provides the final action probability directly for policy optimization. 2.3 Density Evaluation for Policy Optimization Normalizing-flow policies have been used in SAC (Ward et al., 2019; Mazoure et al., 2020) and dexterous fine-tuning with action-chunked critics (Yang et al., 2026). NFPO uses stabilized Real NVP couplings to provide action probabilities for PPO ratios (Shi et al., 2026); PrT-Flow instead learns multi-step maps through local Flow Matching regression with CET. SAC-Flow uses noise-augmented sampling and a joint probability density over intermediate states as an entropy surrogate for offline-to-online SAC (Zhang et al., 2025a). PrT-Flow obtains the final action density directly by applying deterministic, invertible Euler-solver steps after sampling the initial Gaussian noise. We mainly compare the performance of our methods to SAC-Flow in Section 5. GenPO++ uses a reversible solver with paired consecutive states, whose fixed Jacobian determinants cancel in PPO-style joint probability ratios (Hu et al., 2026). PrT-Flow computes the final action probability density without an auxiliary solver-history variable. FALCON learns interval transport through average-velocity regression, uses cycle consistency to encourage invertibility, and evaluates discrete-map Jacobian determinants for Boltzmann sampling (Rehman et al., 2025). CET instead directly supervises finite maps against Euler targets at the actual sampling step size, while JTT guarantees invertibility by construction and computes determinants from analytic scalar derivatives. Related architectures include coupling and spline flows (Dinh et al., 2017; Durkan et al., 2019), causal Transformers (Zhai et al., 2025), and learned orthogonal projections (Dai and Seljak, 2021). JTT combines learned orthogonal coordinates with causal monotone splines and step-size conditioning. RegFlow regresses against precomputed endpoint transports (Rehman et al., 2026), whereas CET supervises local finite steps. Continuous normalizing flows and maximum-entropy Flow Matching RL track probability density through velocity-field divergence (Chen et al., 2018; Zhang et al., 2025b), whereas PrT-Flow accumulates the log-determinants of its actual sampling maps. 3 Preliminaries 3.1 Robot Policies and Action Chunks At environment time τ, the policy receives a window of observations (oτ−To+1,…,oτ)(o_τ-T_o+1,…,o_τ), where ToT_o is the observation horizon. We denote this observation condition by cτc_τ, or simply c when the context is clear. Each action is a vector τ∈ℝDa a_τ ^D_a. An action-chunk policy predicts TpT_p consecutive actions, executes the first TaT_a, and then replans from the latest observations (Chi et al., 2023). We flatten the predicted chunk into a single vector: a=vec(τ,…,τ+Tp−1)∈ℝd,d=TpDa.a=vec( a_τ,…, a_τ+T_p-1) ^d, 20.00003ptd=T_pD_a. We denote the conditional probability density of the predicted action chunk by πθ(a∣c) _θ(a c), where θ denotes the policy parameters. The demonstration dataset D pairs each observation condition c with a recorded action chunk aDa_D in the environment’s bounded control coordinates [−1,1][-1,1]. The generative model operates in ℝdR^d and maps its endpoint u to an action through a=tanhua= u. The corresponding demonstration endpoint is uD=atanha~Du_D=atanh a_D, where a~D a_D is the demonstration action after the boundary preprocessing described in Appendix A. We use t∈[0,1]t∈[0,1] for time within the generative process, distinct from environment time τ. 3.2 Conditional Flow Matching Given a training pair (c,uD)(c,u_D), draw Gaussian noise z∼(0,Id)z (0,I_d) and a time t∼Unif[0,1]t [0,1] independently. Let ϕφ denote the standard Gaussian probability density. The straight path from the noise sample to the demonstration endpoint has position and velocity Xt=(1−t)z+tuD,V=uD−z.X_t=(1-t)z+tu_D, 20.00003ptV=u_D-z. (1) Conditional Flow Matching fits these path velocities with a network vθ:ℝd×[0,1]×→ℝdv_θ:R^d×[0,1]×C ^d, where C is the observation condition space, by minimizing ℒCFM=(c,uD),z,t1d‖vθ(Xt,t,c)−V‖22,L_CFM=E_(c,u_D),z,t 1d\|v_θ(X_t,t,c)-V\|_2^2, (2) where the expectation over (c,uD)(c,u_D) follows the demonstration distribution. With an MSE loss, the network learns to predict the average target velocity for a given state, time, and observation. We use the straight-path formulation of Lipman et al. (2023); Liu et al. (2023); for a broader treatment of the foundations, conditional objectives, and implementation, see Flow Matching Guide and Code by Lipman et al. (2024). 3.3 Sampling by Numerical Integration To generate an action for a fixed observation, Flow Matching draws a Gaussian sample z and solves an ODE from the learned velocity field: dxtdt=vθ(xt,t,c),x0=z. dx_tdt=v_θ(x_t,t,c), 20.00003ptx_0=z. (3) With K equal steps of size h=1/Kh=1/K and grid points tk=kht_k=kh, we can approximately solve this ODE via an Euler solver: xk+1=xk+hvθ(xk,tk,c),k=0,…,K−1.x_k+1=x_k+h\,v_θ(x_k,t_k,c), 20.00003ptk=0,…,K-1. (4) Here xkx_k is the numerical state at tkt_k. After K steps, the endpoint u=xKu=x_K is mapped to the action chunk. Throughout the paper, XtX_t denotes a point constructed directly on a training path, whereas xkx_k denotes a state reached by the sampler. The distribution of generated actions is determined by this finite sequence of numerical updates. 3.4 Soft Actor-Critic Soft Actor-Critic (SAC) is an off-policy algorithm that learns a policy and an action-value function Qψ(c,a)Q_ψ(c,a) from transitions stored in a replay buffer ℬB (Haarnoja et al., 2018). Its policy objective combines value maximization with entropy regularization: ℒSAC=c∼ℬ,a∼πθ(⋅∣c)[αlogπθ(a∣c)−Qψ(c,a)],L_SAC=E_c ,\,a _θ(· c)[α _θ(a c)-Q_ψ(c,a)], (5) where ψ denotes the critic parameters and α>0α>0 sets the entropy weight. The policy must provide both an action sample and its log probability, which also enters the critic’s soft Bellman target. Section 4.4.2 gives the action-chunk objectives used in our experiments. 4 PrT-Flow PrT-Flow generates action chunks by repeatedly transforming noisy actions and tracking the change in probability density at each step. The Jacobian-Tractable Transformer (JTT) implements these Euler-solver updates, returning the updated state together with the update’s Jacobian determinant. Both computations are differentiable with respect to the network parameters, allowing SAC to optimize the policy through sampled actions and their log probabilities. Conditional Euler-step Training (CET) learns the Euler-solver updates by supervising their displacements along demonstration paths. We first explain how to compute probability densities during sampling, then describe the network architecture and training procedure. 4.1 Motivation Consider the Euler update from Section 3.3, xk+1=xk+hvθ(xk,tk,c)x_k+1=x_k+h\,v_θ(x_k,t_k,c). Viewed as a transformation of the state xkx_k, it defines Fh(x,t,c)=x+hvθ(x,t,c)F_h(x,t,c)=x+h\,v_θ(x,t,c), which we call an Euler-step. The sampler starts from Gaussian noise and applies K such transformations. If, for fixed h,t,ch,t,c, each Fh(⋅,t,c)F_h(·,t,c) is a piecewise continuously differentiable bijection with a nonzero Jacobian determinant almost everywhere, the probability density after one step follows from the change-of-variables formula: pk+1(xk+1∣c)=pk(xk∣c)|detJFh(xk,tk,c)|,xk+1=Fh(xk,tk,c).p_k+1(x_k+1 c)= p_k(x_k c)| J_F_h(x_k,t_k,c)|, 20.00003ptx_k+1=F_h(x_k,t_k,c). (6) Here pkp_k is the probability density before the transformation, and JFhJ_F_h is the Jacobian of Fh(⋅,t,c)F_h(·,t,c). Its absolute determinant measures the local change in volume: expansion lowers the density, while contraction raises it. Taking logarithms makes these corrections additive across sampling steps. For an arbitrary velocity network vθv_θ, differentiating the Euler-step with respect to x, while holding h,t,ch,t,c fixed, gives JFh(x,t,c)=I+hJvθ(x,t,c).J_F_h(x,t,c)=I+hJ_v_θ(x,t,c). (7) Each probability density correction therefore requires det(I+hJvθ) (I+hJ_v_θ), which generally entails additional Jacobian and determinant calculations. Moreover, a finite Euler step need not be globally invertible. We address both issues by parameterizing Fθ,hF_θ,h directly as an invertible network with an analytic determinant. Trained on Euler-step targets along Flow Matching paths, this transformation serves as a learned Euler-step. Figure 1 illustrates how it generates samples and tracks their densities. For step size h=1/Kh=1/K, our sampler is x0=z∼(0,Id),xk+1=Fθ,h(xk,tk,c),tk=kh,u=xK,a=tanhu.x_0=z (0,I_d), 10.00002ptx_k+1=F_θ,h(x_k,t_k,c), 10.00002ptt_k=kh, 10.00002ptu=x_K, 10.00002pta= u. (8) The final transformation, a=tanhua= u, maps the generated endpoint into the bounded action space. Writing ℓk=log|detJFθ,h(xk,tk,c)| _k= | J_F_θ,h(x_k,t_k,c)|, we obtain the sampled action chunk’s log probability density: logπθ(a∣c)=logϕ(z)−∑k=0K−1ℓk−∑j=1dlog(1−aj2). _θ(a c)= φ(z)- _k=0^K-1 _k- _j=1^d (1-a_j^2). (9) The three terms account for the Gaussian probability density, the determinants of the learned updates, and the final tanh transformation. Under the conditions in Section 4.2, Eq. (9) gives the log probability density of the sampler in Eq. (8) almost everywhere on (−1,1)d(-1,1)^d. Algorithm 4.1 computes it alongside the sample by subtracting the log-determinant at each step. Algorithm 1 Forward sampling and action probability density 1: observation condition c, step count K, policy parameters θ 2: normalized action chunk a and logπθ(a∣c) _θ(a c) 3: h←1/Kh← 1/K 4: x0∼(0,Id)x_0 (0,I_d) 5: ℓ←logϕ(x0) ← φ(x_0) 6: for k=0,…,K−1k=0,…,K-1 do 7: tk←kht_k← kh 8: Apply the composed JTT update to obtain xk+1=Fθ,h(xk,tk,c)x_k+1=F_θ,h(x_k,t_k,c) and its log-determinant ℓk _k 9: ℓ←ℓ−ℓk ← - _k 10: end for 11: a←tanhxKa← x_K 12: ℓ←ℓ−∑j=1dlog(1−aj2) ← - _j=1^d (1-a_j^2) 13: return a,logπθ(a∣c)←ℓa,\ _θ(a c)← 4.2 Design Requirements for the Euler-step Network Section 4.1 derives the action probability density from the determinant of each Euler-step. To compute these determinants while learning from Flow Matching paths, we impose four requirements on the Euler-step network. 1. Identity at zero step size and a local velocity. The step size h specifies the duration of each transport step. A zero-length step must therefore leave the state unchanged: Fθ,0(x,t,c)=xF_θ,0(x,t,c)=x. Without this condition, the transformation could produce a finite displacement even as the time interval tends to zero. Here the local velocity is induced by the directly parameterized finite map family. Wherever the derivative with respect to h exists at zero, the identity condition gives the expansion vθ(x,t,c)=∂hFθ,h(x,t,c)|h=0,Fθ,h(x,t,c)=x+hvθ(x,t,c)+o(h).v_θ(x,t,c)= . _hF_θ,h(x,t,c) |_h=0, 20.00003ptF_θ,h(x,t,c)=x+h\,v_θ(x,t,c)+o(h). (10) Thus, the displacement divided by h approaches vθv_θ as h→0h→ 0, connecting the learned finite update to a velocity field in the zero-step limit. 2. Global invertibility and spatial regularity. To apply the change-of-variables formula, we require the Euler update, for each fixed (t,c,h)(t,c,h), to be a piecewise continuously differentiable bijection of ℝdR^d with a nonzero spatial Jacobian determinant almost everywhere. Every output then has a unique preimage, allowing us to propagate its probability density through the full sequence of updates. 3. Tractable determinants during sampling. An invertible network does not necessarily provide a tractable determinant. For a typical neural network, computing the Jacobian determinant requires additional calculations beyond the forward pass. Our network must therefore compute the Jacobian determinant efficiently during the forward pass, without constructing the full Jacobian matrix, so that log|detJFθ,h| | J_F_θ,h| can be accumulated to obtain the sample’s probability density. 4. Numerically well-behaved transformations. A transformation can be analytically invertible yet sharply compress or stretch particular directions, making repeated updates and probability density gradients sensitive to numerical perturbations. The architecture and implementation should limit such directional distortions and allow stable evaluation of log-determinants and their gradients. This is a numerical design goal. We will describe how JTT addresses these requirements in the next section. 4.3 Network Architecture One JTT feedforward computation consists of M sequential map heads. Each head combines a learned orthogonal coordinate frame, a Transformer conditioner, and monotone scalar splines. A spline is a one-dimensional function made of connected curve segments; we use it to transform each scalar coordinate monotonically. Figure 2 shows the complete feedforward calculation (A), the operations within a head (B), and the determinant calculation (C). A head first rotates the state x∈ℝdx ^d into its learned frame, y=Qxy=Qx, then embeds the right-shifted scalar sequence (0,y1,…,yd−1)(0,y_1,…,y_d-1) as d tokens. Together with time and observation information, these tokens enter the Transformer, which predicts spline parameters for each coordinate. The splines transform the unshifted values yiy_i, and Q⊤Q rotates the result back to the original frame. A map head thus denotes a complete invertible transformation, distinct from the attention heads inside its Transformer. The Transformer predicts parameters, while the splines perform the state transformation. The conditioner can therefore use nonlinear representations of the state and observations, with invertibility guaranteed by the transformation it parameterizes. Causality specifies the dependency order among scalar coordinates of the generative state; all commands in an action chunk share the current observation condition c. 4.3.1 Causal Scalar Transformations For coordinate i, the Transformer predicts parameters ηθ,i(y<i,t,c) _θ,i(y_<i,t,c) using only the preceding coordinates y<i=(y1,…,yi−1)y_<i=(y_1,…,y_i-1), time, and observations. The resulting scalar transformation is Sθ,h(y)i=sθ,h(yi;ηθ,i(y<i,t,c)),i=1,…,d.S_θ,h(y)_i=s_θ,h\! (y_i; _θ,i(y_<i,t,c) ), 20.00003pti=1,…,d. (11) For example, in three dimensions, suppressing the fixed conditions (t,c,h)(t,c,h) gives w1=s1(y1,η1),w2=s2(y2,η2(y1)),w3=s3(y3,η3(y1,y2)),JS(y)=(∂y1w100∂y1w2∂y2w20∂y1w3∂y2w3∂y3w3). aligned w_1&=s_1(y_1; _1),\\ w_2&=s_2(y_2; _2(y_1)),\\ w_3&=s_3(y_3; _3(y_1,y_2)), aligned 20.00003ptJ_S(y)= pmatrix _y_1w_1&0&0\\ _y_1w_2& _y_2w_2&0\\ _y_1w_3& _y_2w_3& _y_3w_3 pmatrix. The Jacobian is lower triangular, so its determinant equals the product of the diagonal scalar derivatives, evaluated with the prefix held fixed. Since ηθ,i _θ,i does not depend on yiy_i, these entries require only the derivative of each scalar transformation with respect to its own input, without differentiating through the conditioner. Right-shifted tokens and a causal self-attention mask enforce this prefix dependence: token i can attend to tokens 11 through i, whose scalar inputs extend only to yi−1y_i-1. Cross-attention supplies time and observation information. Since all input coordinates are available during forward sampling, one Transformer pass predicts the parameters for all coordinates in parallel. We ensure a nonzero determinant by choosing scalar transformations with positive derivatives. Each coordinate uses a monotone rational-quadratic spline (RQS) (Durkan et al., 2019), which divides a fixed interval into bins and maps each input bin to an output bin with a rational-quadratic curve. The predicted bin widths, bin heights, and knot slopes control how the spline nonlinearly expands or contracts different parts of the coordinate range. Positive widths, heights, and slopes make the transformation strictly increasing within the interval; identity tails extend it to a bijection of ℝR. Every scalar transformation is therefore bijective at a fixed prefix, as is the full causal transformation. Analytic RQS derivatives give the diagonal entries needed to evaluate its determinant. 4.3.2 Step-Size Parameterization and Orthogonal Heads To obtain the identity at h=0h=0, we parameterize the spline through deviations from its identity settings, controlled by the step size. We convert the conditioner’s raw outputs b into bounded parameter rates and scale them by h: r(b)=Rtanh(b/R),ξh=hr(b)r(b)=R (b/R), 20.00003pt _h=h\,r(b) (12) Here R>0R>0 controls the rate magnitude, and ξh _h gives the parameter offsets for the current step. We then convert these offsets into valid bin widths, bin heights, and knot slopes. Thus h scales the parameter rates before this conversion; it does not multiply the final widths, heights, or slopes. Zero offsets produce an identity spline: the input and output bin boundaries coincide, and every knot slope equals one. Hence h=0h=0 ensures ξh=0 _h=0 and Sθ,0=idS_θ,0=id for any network output. For h>0h>0, the predicted offsets change the spline’s shape. This parameterization embeds the identity condition from Section 4.2 directly in the architecture. Appendix C.1 gives the RQS parameterization and its positivity and identity conditions. Within one causal head, early coordinates cannot depend on later ones. We therefore compose heads in different learned coordinate frames, so that subsequent heads act on different linear combinations of the preceding outputs. Each head takes the form y=Qx,Hθ,h(x)=Q⊤Sθ,h(Qx),Q⊤Q=I.y=Qx, 20.00003ptH_θ,h(x)=Q S_θ,h(Qx), 20.00003ptQ Q=I. (13) Time and observations are held fixed in this expression. Orthogonal frames preserve lengths, contribute no additional determinant factor, and can be inverted by taking the transpose. Composing M heads gives one sampling update: Fθ,h=H(M)θ,h∘⋯∘H(1)θ,h.F_θ,h=H^(M)_θ,h ·s H^(1)_θ,h. Each head has its own coordinate frame and conditioner. Since the orthogonal transformations introduce no probability density correction, the scalar log-derivatives suffice to track the density through the composition. Every head, and therefore the complete update, becomes the identity at h=0h=0. The same update parameters are used at all K sampling steps, with tkt_k supplied to each conditioner. The resulting JTT update is invertible, and its determinant depends only on scalar derivatives, as formalized below. Proposition 4.1 (Bijective feedforward transformations with tractable determinants). For fixed (t,c,h)(t,c,h), suppose every scalar transformation in Eq. (11) is a bijection ℝ→ℝR for every fixed prefix. Then each head in Eq. (13) is bijective. At differentiable points with positive scalar derivatives, log|detJHθ,h(x)|=∑i=1dlog∂sθ,h(yi,ηθ,i)∂yi. | J_H_θ,h(x)|= _i=1^d ∂ s_θ,h(y_i; _θ,i)∂ y_i. (14) Composing heads and sampling updates adds these log-determinants at their respective input states. The proposition is proved in Appendix C. 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 xk+1x_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, yiy_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 ℓk=logD _k= D. 4.4 Training 4.4.1 Conditional Euler-step Training Standard Flow Matching trains a velocity field along a continuous-time path, whereas our policy samples through a discrete sequence of finite maps. The identity condition Fθ,0(x,t,c)=xF_θ,0(x,t,c)=x defines a local velocity, ∂hFθ,h(x,t,c)|h=0 . _hF_θ,h(x,t,c) |_h=0, wherever the derivative exists (Section 4.2). This derivative describes only the map family’s first-order behavior as h→0h→ 0. At sampling time, the policy applies the finite map at a fixed positive step size h=1/Kh=1/K. Two map families with the same derivative at h=0h=0 can therefore produce different displacements at the sampling step size: equal zero-step velocities do not imply equal finite updates. CET therefore supervises the finite displacement at the sampling step size. Dividing this displacement by h gives a prediction of the Flow Matching path velocity that requires only one forward map evaluation. No differentiation with respect to h is needed, and parameter gradients follow by ordinary backpropagation. Appendix C.2 explains how regression with this finite quotient differs from regression with the zero-step generator. For each training example, we sample a demonstration action chunk and apply the inverse-tanh transformation described in Section 3.1 to obtain the target uDu_D. We also sample Gaussian noise z and a time tkt_k from the sampling grid. Appendix A specifies how boundary values are handled. Interpolating between z and uDu_D gives the input XtkX_t_k on a straight path with velocity uD−zu_D-z. Applying the update network at this state and dividing its displacement by h yields the predicted path velocity: Xt=(1−t)z+tuD,V=uD−z,qθ,h(x,t,c)=Fθ,h(x,t,c)−xh.X_t=(1-t)z+tu_D, 20.00003ptV=u_D-z, 20.00003ptq_θ,h(x,t,c)= F_θ,h(x,t,c)-xh. (15) The CET loss is similar to the Conditional Flow Matching loss (Eq. (2)). Using the sampling step size h=1/Kh=1/K and a uniformly drawn grid index k∈0,…,K−1k∈\0,…,K-1\, CET minimizes ℒCET,h=c,uD,z,k1d‖qθ,h(Xtk,tk,c)−V‖22;L_CET,h=E_c,u_D,z,k 1d\|q_θ,h(X_t_k,t_k,c)-V\|_2^2; (16) for each endpoint pair, Xtk+hV=Xtk+1X_t_k+hV=X_t_k+1, so this loss trains the map to advance the path sample toward its next grid position. Both XtkX_t_k and the target velocity V are constructed directly from the endpoints, requiring only one local map evaluation per training example. At deployment, the policy starts from noise and applies all K updates in sequence, using each generated state as the input to the next update. The training and sampling procedures thus use the same positive step size and time grid. Algorithm 4.4.1 describes how to train our model with demonstrations. Algorithm 2 Conditional Euler-step Training 1: demonstration windows after normalization and inverse-tanh transformation u=(c,uD)D_u=\(c,u_D)\, initialized θ, step count K, batch size B, update budget U 2: trained policy parameters θ 3: h←1/Kh← 1/K 4: for r=1,…,Ur=1,…,U do 5: Sample a minibatch (ci,uD,i)i=1B\(c_i,u_D,i)\_i=1^B from uD_u 6: for i=1,…,Bi=1,…,B do 7: Draw zi∼(0,Id)z_i (0,I_d) and ki∼Unif0,…,K−1k_i \0,…,K-1\ 8: ti←kiht_i← k_ih; Xi←(1−ti)zi+tiuD,iX_i←(1-t_i)z_i+t_iu_D,i; Vi←uD,i−ziV_i← u_D,i-z_i 9: qi←(Fθ,h(Xi,ti,ci)−Xi)/hq_i← (F_θ,h(X_i,t_i,c_i)-X_i )/h 10: end for 11: ℒ←1Bd∑i=1B‖qi−Vi‖22L← 1Bd _i=1^B\|q_i-V_i\|_2^2 12: Take one optimizer step on θ using ∇θℒ _θL 13: end for 14: return θ Different noise–demonstration pairs may pass through the same state with different velocities. Given the state, time, and observation condition, the MSE loss targets the conditional average over velocities of endpoint pairs consistent with that input. The following proposition shows that, within our finite-map family, this conditional supervision and supervision using the marginal velocity have the same minimizers. Proposition 4.2 (Minimization equivalence of CET and marginal Flow Matching). Fix h=1/K>0h=1/K>0 and a parameter domain Θ . Under the training distribution in Eq. (16), let S=(k,Xtk,c)S=(k,X_t_k,c) denote the information available to the network, and define U=[V∣S]U=E[V S] and qθ,h=qθ,h(Xtk,tk,c)q_θ,h=q_θ,h(X_t_k,t_k,c). Assume that the distribution of (S,V)(S,V) is independent of θ and, for every θ∈Θθ∈ , V,qθ,h∈L2V,q_θ,h∈ L^2 and qθ,hq_θ,h is measurable with respect to S. Define ℒFM,h(θ)=1d‖qθ,h−U‖2,C=1d‖V−U‖2.L_FM,h(θ)= 1dE\|q_θ,h-U\|^2, 20.00003ptC= 1dE\|V-U\|^2. (17) Then, for every θ∈Θθ∈ , ℒCET,h(θ)=C+ℒFM,h(θ),L_CET,h(θ)=C+L_FM,h(θ), (18) where C is independent of θ. For the marginal Euler target Th(S)=Xtk+hUT_h(S)=X_t_k+hU, also define ℒE,h(θ)=1d‖Fθ,h(Xtk,tk,c)−Th(S)‖2=h2ℒFM,h(θ).L_E,h(θ)= 1dE\|F_θ,h(X_t_k,t_k,c)-T_h(S)\|^2=h^2L_FM,h(θ). (19) Consequently, argminθ∈ΘℒCET,h=argminθ∈ΘℒFM,h=argminθ∈ΘℒE,h. *argmin_θ∈ L_CET,h= *argmin_θ∈ L_FM,h= *argmin_θ∈ L_E,h. (20) Similar to Conditional Flow Matching, CET uses velocities from directly sampled endpoint pairs without explicitly computing the marginal velocity U (Lipman et al., 2024). See Appendix C.2 for the proof. 4.4.2 SAC-Style Policy Learning We follow an offline-to-online actor–critic schedule based on SAC-Flow (Zhang et al., 2025a). Each offline iteration performs two successive updates on the same batch. The first trains the actor with CET and the critics using the SAC critic loss. The second applies a SAC-style actor–critic update to the same batch. During online learning, the policy executes sampled action chunks and stores the resulting transitions in replay. After warm-up, it trains the actor and critics on replay windows using the objectives below. CET is used only offline; online actor updates use the SAC-style objective. Let aBa_B denote the action chunk in a replay window, drawn from either the offline dataset or online interaction. The SAC-style actor objective is ℒactor=c∼ℬ,a∼πθ(⋅∣c)[αlogπθ(a∣c)−Q¯ψ(c,a)]+λ(c,aB)∼ℬ,a∼πθ(⋅∣c)1d∥a−aB∥22,L_actor=E_c ,\,a _θ(· c)[α _θ(a c)- Q_ψ(c,a)]+ _(c,a_B) ,\,a _θ(· c) 1d\|a-a_B\|_2^2, (21) where Q¯ψ Q_ψ averages the two critics Qψ1Q_ _1 and Qψ2Q_ _2, and λ weights the action-matching penalty. These objectives act on different quantities. CET uses demonstrations to supervise local transport at noisy intermediate states. The Q term evaluates the complete sampled chunk in terms of task reward, while the entropy term uses its log probability. Action matching penalizes deviations from replay actions directly in action space, constraining reward-driven updates even when no CET path is sampled. During online learning, the reference actions change as the policy collects new replay data. Algorithm 4.1 computes a and its log probability logπθ(a∣c) _θ(a c) together, allowing policy gradients to backpropagate. For a complete predicted chunk executed from environment time τ with Tp=TaT_p=T_a, define its discounted return as Rτ(Ta)=∑i=0Ta−1γirτ+iR_τ^(T_a)= _i=0^T_a-1γ^ir_τ+i, where rτ+ir_τ+i is the reward at each step and γ∈[0,1)γ∈[0,1) is the discount factor. Let mτ(Ta)m_τ^(T_a) indicate whether bootstrapping is permitted. Sampling the next chunk a′∼πθ(⋅∣cτ+Ta)a _θ(· c_τ+T_a) gives the soft target yτ=Rτ(Ta)+γTamτ(Ta)[Q¯target(cτ+Ta,a′)−αlogπθ(a′∣cτ+Ta)],a′∼πθ(⋅∣cτ+Ta).y_τ=R_τ^(T_a)+γ^T_am_τ^(T_a) [ Q_target(c_τ+T_a,a )-α _θ(a c_τ+T_a) ], 10.00002pta _θ(· c_τ+T_a). (22) Here Q¯target Q_target averages the two target critics. Each critic minimizes the squared Bellman residual between its value for the replayed chunk and yτy_τ. Both the actor objective and the critic target use the log probability of the complete executed chunk. Appendix A specifies the settings of the replay masks, optimizers, target-network updates, and initial entropy temperature. 5 Evaluation Our experiments address three questions: can JTT learn effective imitation policies under its structural constraints; can the same policy support offline-to-online actor–critic learning; and how do finite-step supervision and learned coordinate frames affect task performance? Lift evaluates imitation learning and the two design choices, while Cube-Triple evaluates online learning and compares offline training configurations. 5.1 Experimental Setup Benchmarks. For Robomimic Lift (Mandlekar et al., 2022), we use the low-dimensional Multi-Human demonstration dataset, comprising 27,582 training windows. The robot must grasp and lift a cube based on a 19-dimensional state observation, with each action chunk containing four seven-dimensional end-effector and gripper commands (d=28d=28). OGBench Cube-Triple Task 4 (Park et al., 2025a) uses the cube-triple-play-v0 dataset and requires the robot to move three cubes to a target configuration. The policy maps a 46-dimensional state observation to a chunk of five commands, each with five control coordinates. Both tasks use To=1T_o=1 and Tp=TaT_p=T_a: the robot executes the full predicted chunk before replanning. Policies and training. PrT-Flow applies four orthogonal JTT heads per update, starting from a Gaussian base distribution and using a final tanh transformation to map outputs to the normalized action bounds. On Lift, we use 16 sampling steps with CET and evaluate the exponential moving average (EMA) weights after 300 training epochs. On Cube-Triple, we use four sampling steps and train the randomly initialized policy and critics for 30,000 paired offline iterations, followed by 600,000 primitive environment steps of online learning. The main configuration uses finite-step CET at the four sampling grid times and an online action-matching weight of λ=100λ=100. Appendix A gives the hyperparameters and action preprocessing details. Baselines and metrics. The Lift reference is a denoising diffusion probabilistic model (DDPM) policy trained on the same data, with the same action horizons and evaluation episodes. On Cube-Triple, we compare against Flow-G, the gated velocity policy from the SAC-Flow family (Zhang et al., 2025a), with sampling noise standard deviations σ=0.01σ=0.01 and 0.050.05. All three runs use 30k paired offline iterations and 600k online primitive steps. Flow-G samples with a gated MLP and added noise, whereas PrT-Flow uses deterministic JTT updates; all three use the same raw demonstration corpus. Success rate is the fraction of episodes satisfying the environment’s success criterion, and mean return is the average undiscounted episode reward. Each evaluation comprises 50 episodes, with a maximum of 500 steps per episode on Lift and 1,000 on Cube-Triple. We report Lift results after 300 training epochs and evaluate Cube-Triple every 10k online steps, reporting results at 600k steps. 5.2 Main Results Imitation learning on Lift. PrT-Flow achieves a 96% success rate, matching the DDPM baseline when both policies are evaluated after 300 training epochs (Table 1). This result shows that causal, invertible JTT updates can support effective closed-loop manipulation despite their structural constraints. Table 1: Lift imitation learning after 300 epochs. Both policies are evaluated on the same set of 50 episodes. Policy Evaluation epoch Success rate DDPM reference 300 96% PrT-Flow (JTT, CET) 300 96% Offline-to-online learning on Cube-Triple. At 600k online steps, PrT-Flow succeeds in 47 of 50 development evaluation episodes, compared with 44 for Flow-G (σ=0.01σ=0.01) and 28 for Flow-G (σ=0.05σ=0.05). Their mean returns are −578.28-578.28, −613.80-613.80, and −1561.86-1561.86, respectively (Table 2). Figure 3 shows the full learning trajectories on this shared development set. In separate final evaluations after training, PrT-Flow succeeds in 48 of 50 episodes, compared with 46 for Flow-G (σ=0.01σ=0.01) and 25 for Flow-G (σ=0.05σ=0.05). Although Flow-G (σ=0.01σ=0.01) achieves a slightly higher final mean return (−475.48-475.48 versus −530.80-530.80), PrT-Flow improves more rapidly during online training (Figure 3). At half the online training budget (300k primitive environment steps), PrT-Flow achieves a development success rate of 68% and a mean return of −1158.22-1158.22, compared with 38% and −2035.18-2035.18 for Flow-G (σ=0.01σ=0.01). These learning curves indicate an advantage in online sample efficiency for PrT-Flow in this comparison. The resulting policy achieves effective online control using JTT and the action probabilities computed during its forward pass. Appendix B reports numerical checks of the trained policy’s forward log probabilities and their parameter gradients. 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. Policy Dev. SR Dev. return Final SR Final return Flow-G, σ=0.01σ=0.01 88% −613.80-613.80 92% −475.48-475.48 Flow-G, σ=0.05σ=0.05 56% −1561.86-1561.86 50% −1635.42-1635.42 PrT-Flow, CET 94% −578.28-578.28 96% −530.80-530.80 Figure 3: Learning curves on Cube-Triple for PrT-Flow and Flow-G (σ=0.01σ=0.01 and 0.050.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 (σ=0.01σ=0.01), and 56% for Flow-G (σ=0.05σ=0.05). Values in the legend correspond to this endpoint, rather than each run’s best checkpoint. 5.3 Ablations and Offline Training Configurations We first examine two design choices using pairs of Lift runs with matched training settings: the supervised quantity and the orthogonal coordinate frames. Both comparisons use 16 sampling steps and evaluate EMA policies after 300 training epochs (Table 3). These are separate runs with a 300-epoch cosine schedule; the PrT-Flow benchmark in Table 1 uses the epoch-300 checkpoint of a 3,000-epoch schedule. Table 3: Design ablations on Lift with K=16K=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. 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 Finite-step supervision. The two policies share the network architecture, initialization, training time grid, and demonstration endpoints in generative coordinates. We change only whether the training objective regresses the finite-step quotient (Fθ,h−id)/h(F_θ,h-id)/h at h=1/16h=1/16 or the zero-step generator ∂hFθ,h|h=0 _hF_θ,h|_h=0 against the conditional path velocity. CET succeeds in 49 of 50 episodes, compared with 33 for generator supervision. This result supports supervising the finite update used during sampling. Learned coordinate frames. In a separate matched pair, both policies use CET and start from the same initial weights and random orthogonal frames. One policy learns the frames jointly with the rest of JTT; the other keeps them fixed at initialization. The learned-frame policy succeeds in 49 of 50 episodes, compared with 33 for fixed frames. Thus, learning the coordinate system improves task performance in this comparison, supporting the use of trainable orthogonal transformations alongside the causal scalar maps. Offline training configurations on Cube-Triple. We also train the same JTT architecture with an alternative offline velocity objective, regressing the zero-step generator ∂hFθ,h|h=0 _hF_θ,h|_h=0 against the conditional path velocity at continuously sampled times, with endpoints in normalized action coordinates. CET instead supervises the finite displacement divided by h at h=1/4h=1/4, using the four sampling times and demonstration endpoints transformed as in Section 3.1. Both configurations retain the same paired offline actor–critic schedule and SAC-style online learning. They thus differ in how the local maps are supervised offline, while preserving invertible sampling and tractable probability density evaluation. 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. Supervised quantity / time sampling Endpoint coordinates Dev. SR Dev. return Final SR Final return Generator velocity / continuous Action 88% −806.02-806.02 94% −619.72-619.72 CET quotient / grid Generative space 94% −578.28-578.28 96% −530.80-530.80 CET achieves higher development success at 40 of the 42 checkpoints from 190k to 600k steps, with a final success rate of 96% versus 94% for velocity supervision. Although the final difference is one episode, the checkpoint comparison indicates different training trajectories. Because the estimator, time sampling, and target coordinates all change, this comparison evaluates the complete offline configurations. 5.4 Discussion The analysis shows how PrT-Flow computes the action probabilities defined by the finite-step sampling procedure used at deployment. JTT provides analytic Jacobian corrections through invertible finite maps, while CET supervises finite displacements at the deployment step size and relates this objective to marginal Flow Matching regression over the same map family. The experiments show that the invertible and causal structure needed for this computation still permits effective robot behavior: PrT-Flow achieves high imitation success on Lift and supports offline-to-online actor–critic learning on Cube-Triple with the same action generation and probability evaluation procedure. Thus, one policy can combine multi-step Flow Matching generation with an explicit probability density for the actions it executes. PrT-Flow and methods such as SAC-Flow differ in how they obtain policy probabilities. SAC-Flow introduces random perturbations into finite-step sampling to define probabilities over sampling paths. PrT-Flow applies deterministic, invertible maps after sampling the initial Gaussian noise, computing the probability density of the final action chunk directly using the change-of-variables formula. In our Flow-G experiments (Figure 3), changing only the sampling noise standard deviation from σ=0.01σ=0.01 to 0.050.05 produces markedly different learning trajectories and fixed-budget performance, indicating that the noise scale itself can substantially affect training. PrT-Flow does not need stepwise noise to compute action probabilities, so it requires no additional sampling noise scale σ for this purpose. Its high online control performance on Cube-Triple shows that explicit finite-step action densities can support effective policy learning without a stochastic path-based probability construction. This finite-step probability computation also provides a basis for extensions beyond SAC’s action-entropy objective. Inverting the transformations allows the same policy to re-evaluate a given action under updated parameters, making it naturally applicable to likelihood ratios, KL regularization, importance weighting, and replay-action rescoring. Extensions to longer prediction horizons, probability densities of executed action prefixes, visually conditioned policies, and real robots would help establish the scope of probability-tractable Flow Matching as a general generative policy representation. 6 Conclusion We introduced PrT-Flow to make the probability density of a robot’s executed action chunk directly computable while preserving multimodal multi-step Flow Matching generation. The central idea is to make the finite sampling procedure itself invertible with tractable Jacobian determinants, so that action probabilities can be computed directly from the sampling steps. JTT combines causal monotone splines and orthogonal coordinate transformations to accumulate density changes in the same forward pass that generates an action; CET supervises finite displacements at the deployment step size, aligning training with the sampling maps actually used. Proposition 4.2 establishes that, at a fixed step size and within the same finite map family, CET and marginal Flow Matching regression have the same set of minimizing parameters. The experiments show that the structural constraints needed to track probabilities do not prevent effective robot learning: PrT-Flow achieves 48 successes in 50 imitation-learning evaluation episodes on Robomimic Lift and serves as the action generator in offline-to-online SAC-style learning on OGBench Cube-Triple, also reaching 48 successes in 50 final evaluation episodes. These results provide evidence that a Flow Matching policy can retain multi-step generation while explicitly evaluating the probability density of the actions it executes. Constraining the finite sampling maps allows a single policy to provide both capabilities. PrT-Flow thus provides a broader basis for studying generative robot policies: the same finite-step model generates actions from noise and returns their explicit probability densities. Future work can exploit this property across a wider range of policy optimization methods. Finite-step generative policies with tractable probabilities may thereby offer a general policy representation that connects the expressive capacity of Flow Matching with likelihood-based reinforcement learning. AI Use Statement ChatGPT and OpenAI Codex were used to assist with Python code implementation, Lean 4 formalization, English-language editing, and LaTeX formatting during the preparation of this manuscript. References Chen et al. (2018) R. T. Q. Chen, Y. Rubanova, J. Bettencourt, and D. K. Duvenaud Neural ordinary differential equations. In Advances in Neural Information Processing Systems, Vol. 31. External Links: Link Cited by: §1, §2.3. Chi et al. (2023) C. Chi, S. Feng, Y. Du, Z. Xu, E. Cousineau, B. C. Burchfiel, and S. Song Diffusion policy: visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems, Daegu, Republic of Korea. External Links: Document, Link Cited by: §1, §2.1, §3.1. Chisari et al. (2024) E. Chisari, N. Heppert, M. Argus, T. Welschehold, T. Brox, and A. Valada Learning robotic manipulation policies from point clouds with conditional flow matching. arXiv preprint arXiv:2409.07343. External Links: Link Cited by: §1, §2.1. Dai and Seljak (2021) B. Dai and U. Seljak Sliced iterative normalizing flows. In Proceedings of the 38th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 139, p. 2352–2364. External Links: Link Cited by: §2.3. Dinh et al. (2017) L. Dinh, J. Sohl-Dickstein, and S. Bengio Density estimation using real NVP. In International Conference on Learning Representations, External Links: Link Cited by: §2.3. Durkan et al. (2019) C. Durkan, A. Bekasov, I. Murray, and G. Papamakarios Neural spline flows. In Advances in Neural Information Processing Systems, Vol. 32. External Links: Link Cited by: §C.1, §2.3, §4.3.1. Haarnoja et al. (2018) T. Haarnoja, A. Zhou, P. Abbeel, and S. Levine Soft actor-critic: off-policy maximum entropy deep reinforcement learning with a stochastic actor. In Proceedings of the 35th International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 80, p. 1861–1870. External Links: Link Cited by: §1, §3.4. Hu et al. (2026) K. Hu, S. Ding, P. Tao, J. Wang, and Y. Shi GenPO++: generative policy optimization with Jacobian-free likelihood ratios. arXiv preprint arXiv:2606.06967. External Links: Link Cited by: §2.3. Lei et al. (2026) B. Lei, K. Daniilidis, and A. Loquercio Reinforcement learning for flow-matching policies with density transport. arXiv preprint arXiv:2606.08602. External Links: Document, Link Cited by: §2.2. Lipman et al. (2023) Y. Lipman, R. T. Q. Chen, H. Ben-Hamu, M. Nickel, and M. Le Flow matching for generative modeling. In International Conference on Learning Representations, External Links: Link Cited by: §1, §2.1, §3.2. Lipman et al. (2024) Y. Lipman, M. Havasi, P. Holderrieth, N. Shaul, M. Le, B. Karrer, R. T. Q. Chen, D. Lopez-Paz, H. Ben-Hamu, and I. Gat Flow matching guide and code. arXiv preprint arXiv:2412.06264. External Links: Document, Link Cited by: §3.2, §4.4.1. Liu et al. (2023) X. Liu, C. Gong, and Q. Liu Flow straight and fast: learning to generate and transfer data with rectified flow. In International Conference on Learning Representations, External Links: Link Cited by: §1, §2.1, §3.2. Mandlekar et al. (2022) A. Mandlekar, D. Xu, J. Wong, S. Nasiriany, C. Wang, R. Kulkarni, L. Fei-Fei, S. Savarese, Y. Zhu, and R. Martín-Martín What matters in learning from offline human demonstrations for robot manipulation. In Proceedings of the 5th Conference on Robot Learning, Proceedings of Machine Learning Research, Vol. 164, p. 1678–1690. External Links: Link Cited by: §5.1. Mazoure et al. (2020) B. Mazoure, T. Doan, A. Durand, J. Pineau, and R. D. Hjelm Leveraging exploration in off-policy algorithms via normalizing flows. In Proceedings of the Conference on Robot Learning, Proceedings of Machine Learning Research, Vol. 100, p. 430–444. External Links: Link Cited by: §2.3. McAllister et al. (2025) D. McAllister, S. Ge, B. Yi, C. M. Kim, E. Weber, H. Choi, H. Feng, and A. Kanazawa Flow matching policy gradients. arXiv preprint arXiv:2507.21053. External Links: Link Cited by: §2.2. Park et al. (2025a) S. Park, K. Frans, B. Eysenbach, and S. Levine OGBench: benchmarking offline goal-conditioned RL. In International Conference on Learning Representations, External Links: Link Cited by: §5.1. Park et al. (2025b) S. Park, Q. Li, and S. Levine Flow q-learning. In Proceedings of the 42nd International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 267, p. 48104–48127. External Links: Link Cited by: §2.2. Rehman et al. (2025) D. Rehman, T. Akhound-Sadegh, A. Gazizov, Y. Bengio, and A. Tong FALCON: few-step accurate likelihoods for continuous flows. arXiv preprint arXiv:2512.09914. External Links: Link Cited by: §2.3. Rehman et al. (2026) D. Rehman, O. Davis, J. Lu, J. Tang, M. Bronstein, Y. Bengio, A. Tong, and A. J. Bose Efficient regression-based training of normalizing flows for boltzmann generators. In The Fourteenth International Conference on Learning Representations, External Links: Link Cited by: §2.3. Ren et al. (2025) A. Z. Ren, J. Lidard, L. L. Ankile, A. Simeonov, P. Agrawal, A. Majumdar, B. Burchfiel, H. Dai, and M. Simchowitz Diffusion policy policy optimization. In International Conference on Learning Representations, External Links: Link Cited by: §2.2. Schulman et al. (2017) J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347. External Links: Link Cited by: §1. Shi et al. (2026) D. Shi, Y. Tang, Z. Zhuang, and D. Wang NFPO: stabilized policy optimization of normalizing flow for robotic policy learning. arXiv preprint arXiv:2603.11470. External Links: Link Cited by: §2.3. Ward et al. (2019) P. N. Ward, A. Smofsky, and A. J. Bose Improving exploration in soft-actor-critic with normalizing flows policies. arXiv preprint arXiv:1906.02771. External Links: Link Cited by: §2.3. Yang et al. (2026) C. Yang, D. Tarasov, D. Liconti, R. Guntz, H. Zheng, and R. K. Katzschmann SERNF: sample-efficient real-world dexterous policy fine-tuning via action-chunked critics and normalizing flows. arXiv preprint arXiv:2602.09580. External Links: Link Cited by: §2.3. Zhai et al. (2025) S. Zhai, R. Zhang, P. Nakkiran, D. Berthelot, J. Gu, H. Zheng, T. Chen, M. Á. Bautista, N. Jaitly, and J. M. Susskind Normalizing flows are capable generative models. In Proceedings of the 42nd International Conference on Machine Learning, Proceedings of Machine Learning Research, Vol. 267, p. 74348–74369. External Links: Link Cited by: §2.3. Zhang et al. (2025a) Y. Zhang, S. Yu, T. Zhang, M. Guang, H. Hui, K. Long, Y. Wang, C. Yu, and W. Ding SAC Flow: sample-efficient reinforcement learning of flow-based policies via velocity-reparameterized sequential modeling. arXiv preprint arXiv:2509.25756. External Links: Link Cited by: §A.1, §2.3, §4.4.2, §5.1. Zhang et al. (2025b) Y. Zhang, Y. Hu, B. Dai, and N. Li Max-entropy reinforcement learning with flow matching and a case study on LQR. arXiv preprint arXiv:2512.23870. External Links: Document, Link Cited by: §2.3. Appendix A Experimental Details Table 5: Optimization settings for the main PrT-Flow experiments. Cube-Triple uses Adam for the actor, critics, and entropy temperature. Parameter Lift Cube-Triple Batch size 256 256 Optimizer AdamW Adam Learning rate 10−410^-4 3×10−43× 10^-4 Weight decay 10−610^-6 — AdamW (β1,β2)( _1, _2) / ϵε (0.9,0.999)(0.9,0.999) / 10−810^-8 — Cosine warm-up (optimizer steps) 100 — EMA power / maximum decay 0.750.75 / 0.9950.995 — Critic hidden layers / width — 4 / 512 Critic aggregation — Mean of twin Q-values Discount γ / target coefficient τtarget _target — 0.990.99 / 0.0050.005 Replay buffer-size option — 2×1062× 10^6 Online warm-up (environment steps) / UTD — 5,000 / 1 Initial entropy temperature / target entropy — 0.20.2 / 0 Online action-MSE weight — 100 Preprocessing. Lift observations use per-coordinate min–max normalization from training statistics to a nominal [−1,1][-1,1] range without clipping; 10%10\% of demonstrations are reserved for validation. Cube-Triple’s offline actions are clipped with margin 10−510^-5. After task preprocessing, any remaining action coordinates at ±1± 1 are replaced by ±(1−10−6)±(1-10^-6), leaving interior values unchanged, to obtain a~D a_D; CET then uses uD=atanha~Du_D=atanh a_D. Training details. The main Lift result is the epoch-300 checkpoint of a 3,000-epoch cosine schedule, whereas the matched ablations use a full 300-epoch schedule; both use the warm-up and EMA settings in Table 5. For incomplete replay windows, invalid positions are masked, but the loss is divided by all entries in the batch–time–coordinate array, including masked entries. The alternative generator-supervised configuration in Table 4 uses continuous times and the latent endpoint aDa_D, corresponding to tanh(aD) (a_D) in action space, whereas CET uses the inverse-tanh endpoint. A.1 Flow-G Baseline The two Flow-G configurations differ only in sampling noise (σ=0.01σ=0.01 or 0.050.05); both use transition noise 0.050.05 in SAC actor training. They retain the noise reuse and log-probability calculation of the authors’ released offline-to-online implementation.11 1 Official Flow-G sampling and log-probability routines, revision 3df8b9d. SAC-Flow’s Gaussian construction defines a joint path probability density, with action probability density obtained by integrating out intermediate states (Zhang et al., 2025a, Appendix A.2, Eqs. (15)–(16)). The released implementation instead reuses fixed-key perturbations across Euler steps while summing Gaussian transition log terms, the base log probability, and the tanh correction. This sum serves as its entropy surrogate rather than the exact log probability of the sampled path or action. All three policies share the development evaluation panel and undergo separate 50-episode final evaluations. Appendix B Numerical Verification Scalar-map checks. We test the policy’s scalar RQS configuration (eight bins, B=4B=4, rate bound 1616) at h=1/16h=1/16 and h=1/4h=1/4. For each of five random seeds, we draw independent sets of 4,096 inputs and outputs, including 128 points in the identity tails in each set, and independent standard-normal raw parameters. With parameters held fixed, we compare analytic and autodifferentiated log-derivatives, check both inverse compositions, and compare forward and inverse log probabilities under a standard Gaussian base. The same draws are evaluated on CPU in FP32 and FP64, giving 40,960 checks per direction and precision. Table 6 reports the maximum absolute discrepancies. All evaluated derivatives are positive; the zero-step and identity-tail checks also pass. 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 FP32 FP64 f−1(f(x))f^-1(f(x)) vs. x 9.54×10−79.54× 10^-7 1.33×10−151.33× 10^-15 f(f−1(y))f(f^-1(y)) vs. y 7.15×10−77.15× 10^-7 1.78×10−151.78× 10^-15 Analytic vs. autodifferentiated log-derivative 8.20×10−78.20× 10^-7 1.55×10−151.55× 10^-15 Forward vs. inverse log probability 4.53×10−64.53× 10^-6 6.22×10−156.22× 10^-15 Trained-policy checks. We check probability density evaluation on the main Cube-Triple CET policy at 600k online steps, which generates 25-dimensional action chunks through four finite updates. On the CPU, we compare FP32 and FP64 forward log probabilities using the same parameters, 32 real dataset states, and fixed Gaussian noise. For the concatenated parameter gradient g of the mean log probability, we report ‖g32−g64‖2/‖g64‖2\|g_32-g_64\|_2/\|g_64\|_2. For three of these sampling trajectories, we also compare each update’s analytic log-determinant with that of its FP64 autodifferentiated Jacobian, giving 12 local comparisons. Table 7 shows close numerical agreement for both the local determinants and the forward probability computations. 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. Comparison Metric Discrepancy Analytic vs. autodifferentiated local log-determinant Maximum absolute difference ≤9.2×10−12≤ 9.2× 10^-12 FP32 vs. FP64 forward log probability Maximum absolute difference 1.20×10−41.20× 10^-4 FP32 vs. FP64 parameter gradient of mean log probability Relative L2L_2 difference 9.77×10−79.77× 10^-7 Appendix C Mathematical Guarantees This appendix proves the inversion and Jacobian identities for computing the policy’s probability density and the minimization equivalence underlying CET. It also distinguishes regression at a fixed step size from zero-step generator training. Propositions 4.1 and 4.2 have both been formally verified in Lean 4. C.1 Causal Inversion and Scalar RQS Derivatives RQS parameterization. Each head learns an unconstrained matrix A and sets Q=exp(A−A⊤)∈SO(d)Q= (A-A ) (d). Using the step-scaled offsets ξh=hr(b) _h=h\,r(b) from Eq. (12), the widths and heights of n=nbinn=n_bin bins on [−B,B][-B,B] are wj w_j =εw+(2B−nεw)softmax(ξhw)j, = _w+(2B-n _w)softmax( _h^w)_j, (23) qj q_j =εh+(2B−nεh)softmax(ξh)j,j=1,…,n. = _h+(2B-n _h)softmax( _h^h)_j, 18.49988ptj=1,…,n. For j=0,…,nj=0,…,n, the knot slopes are mj=expclip[log(εd+softplus(ξh,jr+softplus−1(1−εd))),−L,L].m_j= \! \clip\! [ \! ( _d+softplus\! ( _h,j^r+softplus^-1(1- _d) ) ),-L,L ] \. (24) Here B>0B>0, 0<εw,εh<2B/n0< _w, _h<2B/n, 0<εd<10< _d<1, and L≥0L≥ 0. Cumulative widths and heights define knots with endpoints ±B± B; we use the standard RQS formula (Durkan et al., 2019) within bins and identity tails. At h=0h=0, wj=qj=2B/nw_j=q_j=2B/n and mj=1m_j=1, so the scalar map and its orthogonal compositions are the identity. Proof of Proposition 4.1. Given an output w, invert the first scalar map to recover y1y_1. Once y<iy_<i is known, the conditioner fixes the parameters of the iith scalar map, whose inverse uniquely recovers yiy_i. By induction, every w has exactly one input, proving surjectivity and injectivity. At any differentiable point, changing yjy_j for j>ij>i leaves the iith output unchanged, so the Jacobian of S is lower triangular. Its diagonal entries are the scalar derivatives with the prefix held fixed, giving detJS(y)=∏i∂yisi J_S(y)= _i _y_is_i. For H=Q⊤SQH=Q SQ, the chain rule yields JH(x)=Q⊤JS(Qx)QJ_H(x)=Q J_S(Qx)Q. Orthogonality of Q then gives |detJH(x)|=|detJS(Qx)|| J_H(x)|=| J_S(Qx)|. The scalar derivatives are positive, so taking logarithms is valid. Composition adds the log-determinants of successive heads and steps, each evaluated at its own input. The proof requires neither a triangular Jacobian after rotation nor an exact ODE flow. ∎ For a scalar bin of width w>0w>0, height q>0q>0, and endpoint slopes d0,d1>0d_0,d_1>0, let σ=q/wσ=q/w and τ=(x−x0)/w∈[0,1]τ=(x-x_0)/w∈[0,1]. The RQS derivative is S′(x)=σ2[d1τ2+2στ(1−τ)+d0(1−τ)2][σ+(d0+d1−2σ)τ(1−τ)]2>0.S (x)= σ^2[d_1τ^2+2στ(1-τ)+d_0(1-τ)^2][σ+(d_0+d_1-2σ)τ(1-τ)]^2>0. (25) Before squaring, the denominator is σ[τ2+(1−τ)2]+(d0+d1)τ(1−τ)σ[τ^2+(1-τ)^2]+(d_0+d_1)τ(1-τ), which is positive throughout the bin, as is the numerator. Each bin thus maps its input interval continuously and strictly increasingly onto the prescribed output interval and has a unique within-bin inverse. Equal widths and heights, unit slopes, and coincident input and output knots give the identity. Adjacent bins share knots and derivatives, and identity tails preserve bijectivity on the real line. The learned endpoint slopes may differ from the tails’ unit slope, so we use almost-everywhere derivatives without assuming global C1C^1 regularity. C.2 CET Minimizers and Comparison with Generator Training Proof of Proposition 4.2. Let q=qθ,h(Xtk,tk,c)q=q_θ,h(X_t_k,t_k,c) and U=[V∣S]U=E[V S]. Conditional expectation preserves square integrability, and the Cauchy–Schwarz inequality ensures that all terms below are integrable. Since q−Uq-U is S-measurable and [V−U∣S]=0E[V-U S]=0, we have ⟨q−U,V−U⟩=⟨q−U,[V−U∣S]⟩=0.E q-U,V-U =E q-U,E[V-U S] =0. Expanding the squared norm of q−V=(q−U)−(V−U)q-V=(q-U)-(V-U) and dividing by d gives Eq. (18). The additive constant is the same for every θ∈Θθ∈ . Subtracting it preserves every pairwise loss comparison and hence the full sets of minimizers, including when those sets are empty. Finally, Fθ,h(Xtk,tk,c)−Th(S)=h(qθ,h−U).F_θ,h(X_t_k,t_k,c)-T_h(S)=h (q_θ,h-U ). Taking squared norms and expectations, then dividing by d, gives Eq. (19). The factor h2h^2 is positive and independent of θ, which proves Eq. (20). ∎ If Fθ,0(x,t,c)=xF_θ,0(x,t,c)=x and h↦Fθ,h(x,t,c)h F_θ,h(x,t,c) is differentiable at zero, the definition of derivative gives Fθ,h(x,t,c)=x+hvθ(x,t,c)+o(h),qθ,h(x,t,c)⟶vθ(x,t,c).F_θ,h(x,t,c)=x+hv_θ(x,t,c)+o(h), 18.49988ptq_θ,h(x,t,c) v_θ(x,t,c). (26) For jointly differentiable heads that equal the identity at h=0h=0, the generator of their finite composition is the sum of their generators. Conjugating a head by a fixed linear map conjugates its generator in the same way. Neither statement establishes a semigroup law or implies that the finite map is the exact flow of its generator. Fixed-step objectives differ. Consider the translation Fθ,h(x)=x+hθ+h2F_θ,h(x)=x+hθ+h^2, which is invertible everywhere. Its generator is vθ=θv_θ=θ, but its finite-step quotient is qθ,h=θ+hq_θ,h=θ+h. For a zero target, the finite-step risk is minimized at θ=−hθ=-h, whereas the generator risk is minimized at θ=0θ=0. This example shows that invertibility and Fθ,0=idF_θ,0=id alone do not justify replacing the quotient used at deployment with the zero-step derivative in CET.