PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Congzhong Sun, Ning Tan

arXiv:2610.17667v1cs.ROcs.AI

TL;DR

Robot policies need to represent multiple valid actions, but reinforcement-learning algorithms also require probabilities for sampled actions, which finite Flow Matching sampling does not directly provide. PrT-Flow uses invertible finite steps with tractable determinants and trains their displacements with CET. It reports high success in imitation learning and SAC-style online control.

  • Problem

    Robot policies must represent multiple valid motions, while reinforcement-learning algorithms require action probabilities and finite Flow Matching sampling can differ from its continuous-flow distribution.

    Source: §1 · pp. 1–2

  • Method

    PrT-Flow parameterizes finite sampling steps as invertible networks with tractable Jacobian determinants and trains their displacements using CET.

    Source: §1 · pp. 2–3

  • Results

    PrT-Flow achieves high task success in imitation learning and SAC-style online control, including 96% success on Robomimic Lift.

    Source: §1 · p. 3

  • Takeaways & Limitations

    The reported results show that invertible, density-evaluable networks can support effective robot policies in imitation learning and SAC-style control.

    Source: §1 · p. 3

    Density propagation assumes each fixed-step Euler map is a piecewise continuously differentiable bijection with a nonzero spatial Jacobian determinant almost everywhere.

    Source: §4.1 · p. 5, §4.2 · p. 7

Abstract

from arXiv · show

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 .

1 Introduction

Robot policies must represent multiple valid actions, while reinforcement-learning algorithms also need the probability of each sampled action. PrT-Flow makes finite Flow Matching steps invertible with tractable Jacobian determinants, trains their displacements with CET, and reports high success in imitation learning and online control.

Motivation

  • A policy must represent a distribution over action chunks because averaging several valid motions may produce an invalid motion, and replanning follows execution of a specified number of actions.Many reinforcement-learning algorithms also require action probabilities: SAC uses sampled-action log probability and PPO uses a probability ratio.Source: §1 · pp. 1–2
  • Flow Matching's finite numerical sampler can have an endpoint distribution different from the continuous flow, motivating density computation for the actual sampled action chunk.PrT-Flow's invertible steps return updated samples and determinants; a final tanh maps the sample to a bounded action chunk with its density.Source: §1 · p. 2, Figure 1 · p. 2

Contributions

  • PrT-Flow and JTT compute action probabilities during multi-step generation using tractable per-step Jacobian determinants.The method also trains finite-step displacements with CET, which is introduced as equivalent in minimizers to marginal Flow Matching over the same family of finite maps.Source: §1 · p. 3

Experiments

  • 96% success on Robomimic Lift and high task success on OGBench Cube-Triple show effective learning in imitation and SAC-style online control.The introduction reports high success across both settings but gives no Cube-Triple percentage in these passages.Source: §1 · p. 3

2 Related Work

Prior work uses diffusion and Flow Matching for flexible robot action generation, while several reinforcement-learning approaches avoid directly evaluating the final action probability. PrT-Flow instead computes the final density of its multi-step generated action.

2.1 Diffusion and Flow Matching Policies

  • Diffusion Policy and Flow Matching retain multi-step generation for multimodal robot actions; PointFlowMatch applies conditional Flow Matching to manipulation.PrT-Flow likewise retains multi-step generation while providing action probabilities needed for reinforcement learning.Source: §2.1 · p. 3

2.2 Reinforcement Learning for Generative Policies

  • Several generative-policy RL methods avoid evaluating the exact final action probability, whereas PrT-Flow provides that probability directly for policy optimization.The cited approaches include transition-level PPO ratios, a distilled one-step actor, proxy ratios, and transport updates.Source: §2.2 · p. 3

2.3 Density Evaluation for Policy Optimization

  • PrT-Flow computes the final action density from deterministic invertible Euler steps, without an auxiliary solver-history variable.Unlike continuous normalizing flows and maximum-entropy Flow Matching RL, it accumulates log-determinants of its actual sampling maps; CET supervises finite maps against Euler targets.Source: §2.3 · pp. 3–4

3 Preliminaries

The preliminaries define observation-conditioned action chunks, conditional Flow Matching paths, and finite-step sampling. They also introduce SAC, whose policy uses sampled actions and their log probabilities in learning.

3.1 Robot Policies and Action Chunks

  • An action-chunk policy predicts consecutive actions from an observation window, executes the first specified actions, and then replans from updated observations.The predicted chunk is flattened into a single vector, and its conditional density is denoted πθ(a | c).Source: §3.1 · p. 4

3.2 Conditional Flow Matching

  • Conditional Flow Matching draws Gaussian noise and a time, then trains a network to fit the velocity along a straight path to a demonstration endpoint.With an MSE loss, the network predicts the average target velocity for a given state, time, and observation.Source: §3.2 · p. 4

3.3 Sampling by Numerical Integration

  • Sampling numerically solves the learned velocity-field ODE with K equal Euler steps of size h = 1/K, and the generated-action distribution is determined by those finite updates.The endpoint is mapped to the action chunk after the final step.Source: §3.3 · p. 5

3.4 Soft Actor-Critic

  • SAC learns a policy and action-value function from replay-buffer transitions, and its policy must provide both sampled actions and their log probabilities.The log probability also enters the critic's soft Bellman target.Source: §3.4 · p. 5

4 PrT-Flow

PrT-Flow represents each sampling step as an invertible finite map, allowing action-chunk densities to be tracked through the sampler. JTT makes the maps’ determinants tractable, while CET trains their finite-step displacements along Flow Matching paths.

4.1 Motivation

  • Directly parameterizing each Euler-step as an invertible network with an analytic determinant addresses the extra determinant calculations and finite-step invertibility issues of arbitrary velocity networks.Under the stated bijectivity and nonzero-determinant conditions, density corrections follow from change of variables and accumulate across steps.Source: §4.1 · pp. 5–6

4.2 Design Requirements for the Euler-step Network

  • The Euler-step network is designed to be identity at zero step size, globally bijective with nonzero spatial determinant, tractable to differentiate, and numerically stable across repeated updates.The numerical goal is to limit directional compression or stretching that could make updates and density gradients sensitive to perturbations.Source: §4.2 · pp. 6–7

4.3 Network Architecture

  • A JTT feedforward computation composes M map heads, each combining a learned orthogonal frame, a Transformer conditioner, and monotone scalar splines.The conditioner predicts spline parameters from preceding scalar coordinates, time, and observations; the splines perform the invertible state transformation.Source: §4.3 · p. 7

4.3.1 Causal Scalar Transformations

  • Prefix-conditioned monotone rational-quadratic splines make each scalar transformation bijective, while their positive analytic derivatives yield the diagonal entries used for the lower-triangular Jacobian determinant.Right-shifted tokens and causal attention enforce dependence on preceding coordinates, and one Transformer pass predicts parameters for all coordinates in parallel.Source: §4.3.1 · pp. 7–8

4.3.2 Step-Size Parameterization and Orthogonal Heads

  • Scaling spline-parameter offsets by h makes every head the identity at h = 0; composing heads in different orthogonal frames changes coordinate dependencies without adding determinant corrections.The same update parameters are used at each sampling step, with time supplied to each conditioner; the composed update’s log-determinant comes from scalar derivatives.Source: §4.3.2 · pp. 8–9

4.4.1 Conditional Euler-step Training

  • CET supervises the map’s finite displacement at the positive sampling step size, rather than only its zero-step velocity, using one map evaluation and ordinary backpropagation.Training samples noise and a demonstration endpoint, interpolates them to a grid state, and targets the endpoint-pair velocity; deployment composes all K updates.Source: §4.4.1 · p. 9

A One JTT step

A JTT update composes map heads to transform a noisy action chunk and track its density change. The section also describes CET supervision and SAC-style actor–critic updates that use the sampled chunk’s log probability.

A One JTT step

  • One JTT update composes M map heads and returns the next noisy sample with its determinant, obtained by multiplying head determinants whose scalar derivatives suffice because orthogonal rotations preserve volume.The figure shows the complete update, an individual conditioned spline head, and its determinant calculation.Source: Figure 2 · p. 10, A One JTT step · p. 10
  • CET trains a finite update from endpoint-pair velocities at noisy intermediate states, and under the stated assumptions its conditional supervision shares minimizers with marginal Flow Matching.Each example uses a sampled noise–demonstration pair and one local map evaluation; deployment applies all K updates sequentially.Source: A One JTT step · pp. 10–11

4.4.2 SAC-Style Policy Learning

  • Offline-to-online SAC-style learning performs actor and critic updates on the same batch, using CET offline and the SAC-style actor objective online.Online interaction stores sampled-chunk transitions in replay, and later updates train on replay windows.Source: §4.4.2 · p. 11
  • The SAC-style actor objective combines entropy, critic value, and an action-matching penalty against replay actions; actor and critic targets use the complete chunk’s log probability.The action and its log probability are computed together, allowing policy gradients to backpropagate through sampled actions.Source: §4.4.2 · pp. 11–12

5 Evaluation

The evaluation tests PrT-Flow on imitation learning, online control, and design choices, then discusses how finite-step probabilities support policy learning and where broader evaluation is still needed.

5.1 Experimental Setup

  • The evaluation tests imitation learning, offline-to-online actor–critic learning, and the effects of finite-step supervision and learned coordinate frames across Lift and Cube-Triple.Lift uses 27,582 training windows and 50-episode evaluations; Cube-Triple compares PrT-Flow with Flow-G at two noise scales under shared offline and online training budgets.Source: §5 · p. 12, §5.1 · pp. 12–13

5.2 Main Results

  • 96% success on Lift matches the DDPM baseline after 300 training epochs.Both policies were evaluated on the same 50 episodes.Source: §5.2 · p. 13, Table 1 · p. 13
  • 48 of 50 final Cube-Triple episodes succeed with PrT-Flow, versus 46 for Flow-G (σ = 0.01) and 25 for Flow-G (σ = 0.05), after 600k online steps.On the shared development set, PrT-Flow succeeds in 47 of 50 episodes, compared with 44 and 28, respectively; it also improves more rapidly during online training.Source: §5.2 · p. 13, Table 2 · p. 13, Figure 3 · p. 14

5.3 Ablations and Offline Training Configurations

  • 49 of 50 Lift episodes succeed with CET, compared with 33 under zero-step generator supervision in a matched comparison.The policies share their architecture, initialization, time grid, and demonstration endpoints; the objective is the changed design choice.Source: §5.3 · p. 14, Table 3 · p. 14
  • 49 of 50 Lift episodes succeed with learned coordinate frames, compared with 33 when frames remain fixed.The matched policies use CET and start from the same initial weights and random orthogonal frames.Source: §5.3 · pp. 14–15, Table 3 · p. 14
  • CET has higher Cube-Triple development success at 40 of 42 checkpoints, with a final rate of 96% versus 94% for velocity supervision.This comparison changes the complete offline configurations, including estimator, time sampling, and target coordinates, while retaining paired offline actor–critic training and SAC-style online learning.Source: §5.3 · p. 15, Table 4 · p. 15

5.4 Discussion

  • PrT-Flow computes final action-chunk densities for its finite sampling procedure through deterministic invertible maps; its Cube-Triple comparison shows that Flow-G’s sampling-noise scale can substantially affect learning trajectories and fixed-budget performance.PrT-Flow does not require an additional sampling-noise scale to compute action probabilities.Source: §5.4 · p. 15
  • The same policy can be extended to likelihood ratios, KL regularization, importance weighting, and replay-action rescoring; longer horizons, executed-action densities, visual conditioning, and real-robot evaluations remain areas for establishing its broader scope.The discussion presents these as possible extensions and future evaluations.Source: §5.4 · pp. 15–16

6 Conclusion

PrT-Flow makes action-chunk probabilities computable during finite-step, multimodal Flow Matching generation. Experiments show high success in imitation learning and SAC-style online control, while a proposition establishes a condition under which CET and marginal Flow Matching regression share minimizers.

6 Conclusion

  • PrT-Flow computes action-chunk densities while preserving multimodal, multi-step Flow Matching generation by making finite sampling steps invertible with tractable Jacobian determinants.JTT combines causal monotone splines and orthogonal coordinate transformations; CET supervises finite displacements at the deployment step size.Source: §6 · p. 16
  • At a fixed step size and within the same finite map family, CET and marginal Flow Matching regression have the same minimizing parameters.This is the scope of Proposition 4.2.Source: §6 · p. 16
  • 48 of 50 evaluation episodes succeeded on Robomimic Lift, and 48 of 50 final episodes succeeded on OGBench Cube-Triple under offline-to-online SAC-style learning.The Cube-Triple policy served as the action generator; the results provide evidence that probability-tracking constraints do not prevent effective robot learning.Source: §6 · p. 16
  • The finite-step model generates actions from noise and returns explicit probability densities, a property proposed as a basis for studying more policy-optimization methods.The conclusion describes this broader use as future work and says such policies may connect Flow Matching expressivity with likelihood-based reinforcement learning.Source: §6 · p. 16

A Experimental Details

The experimental details specify preprocessing choices and training schedules, while Table 5 is identified as the optimization-settings reference for the main PrT-Flow experiments.

A Experimental Details

  • Table 5 summarizes optimization settings for the main PrT-Flow experiments, with Adam used for Cube-Triple’s actor, critics, and entropy temperature.The supplied caption does not give the remaining table entries.Source: Table 5 · p. 20
  • Lift observations are min–max normalized from training statistics to a nominal [−1, 1] range without clipping, and 10% of demonstrations are reserved for validation.Cube-Triple offline actions are clipped with margin 10−5; remaining action coordinates at ±1 are moved to ±(1 − 10−6), then transformed using atanh for CET.Source: A Experimental Details · p. 20
  • The main Lift result uses the epoch-300 checkpoint from a 3,000-epoch cosine schedule, while matched ablations use a full 300-epoch schedule.For incomplete replay windows, invalid positions are masked, but the loss denominator includes all batch–time–coordinate entries.Source: A Experimental Details · p. 20

A.1 Flow-G Baseline

The Flow-G configurations vary sampling noise while retaining the released implementation’s fixed-key perturbations and log-probability calculation.

A.1 Flow-G Baseline

  • The two Flow-G configurations differ only in sampling noise, using σ = 0.01 or 0.05; both use transition noise 0.05 in SAC actor training.Both retain noise reuse and log-probability calculation from the authors’ released offline-to-online implementation.Source: §A.1 · p. 20
  • The released implementation’s summed Gaussian transition terms, base log probability, and tanh correction serve as an entropy surrogate, not the exact log probability of the sampled path or action.The passage distinguishes this calculation from SAC-Flow’s joint path density and action density obtained by integrating out intermediate states.Source: §A.1 · p. 20
  • All three policies share the development evaluation panel and undergo separate 50-episode final evaluations.Source: §A.1 · p. 20

B Numerical Verification

Numerical verification checks scalar rational-quadratic-spline maps and a trained Cube-Triple policy, comparing analytic quantities with autodifferentiated or higher-precision calculations.

B Numerical Verification

  • Table 6 reports maximum absolute discrepancies for scalar RQS derivative, inverse-composition, and log-probability checks across five seeds and two step sizes.The checks use independent input and output sets, fixed parameters, CPU FP32 and FP64 evaluation, and include identity-tail points; evaluated derivatives were positive and zero-step and identity-tail checks passed.Source: B Numerical Verification · pp. 20–21, Table 6 · p. 21
  • The main Cube-Triple CET policy at 600k online steps generates 25-dimensional action chunks through four finite updates, and its forward log probabilities are compared in FP32 and FP64.The comparison uses the same parameters, 32 real dataset states, and fixed Gaussian noise; the relative discrepancy of the mean-log-probability parameter gradient is also reported.Source: B Numerical Verification · p. 21
  • Table 7 reports close numerical agreement between FP32 and FP64 forward probabilities and between analytic and autodifferentiated local log-determinants.The local checks cover 12 updates across three four-step sampling trajectories.Source: B Numerical Verification · p. 21, Table 7 · p. 21

C Mathematical Guarantees

The appendix proves the inversion and Jacobian identities used to compute policy densities, alongside the minimization equivalence underlying CET. Its causal scalar transformations are invertible and have tractable Jacobians under the stated construction.

C Mathematical Guarantees

  • The appendix formally verifies Propositions 4.1 and 4.2 in Lean 4.These propositions cover the inversion and Jacobian identities and the minimization equivalence underlying CET.Source: C Mathematical Guarantees · p. 21

C.1 Causal Inversion and Scalar RQS Derivatives

  • At h = 0, equal bin widths and heights, unit slopes, and coincident knots make each scalar map and its orthogonal compositions the identity.The construction uses cumulative widths and heights to define knots on [−B, B], with the standard RQS formula inside bins.Source: §C.1 · p. 21
  • Causal inversion recovers each input coordinate sequentially because each conditioned scalar map has a unique inverse, proving bijectivity.The Jacobian is lower triangular before rotation; orthogonal conjugation preserves its determinant, and composition adds log-determinants across heads and steps.Source: §C.1 · p. 22
  • Positive RQS derivatives make each bin strictly increasing, while identity tails preserve bijectivity on the real line.The endpoint slopes may differ from the tails’ unit slope, so the proof uses almost-everywhere derivatives rather than assuming global C1 regularity.Source: §C.1 · p. 22

C.2 CET Minimizers and Comparison with Generator Training

Under the stated assumptions, CET and fitting the Flow Matching Euler-solver updates have the same minimizers. This equivalence does not extend to replacing the finite-step quotient with zero-step generator training.

C.2 CET Minimizers and Comparison with Generator Training

  • Subtracting a θ-independent loss constant and scaling by positive h2 preserve all minimizers for the corresponding objectives.The proof preserves pairwise loss comparisons, including when the minimizer sets are empty.Source: §C.2 · p. 22
  • If the finite map is identity at h = 0 and differentiable there, its finite-step quotient converges to its local velocity as h tends to zero.The expansion is Fθ,h(x, t, c) = x + hvθ(x, t, c) + o(h).Source: §C.2 · p. 22
  • For jointly differentiable heads that are identity at h = 0, composition adds their generators, but this establishes neither a semigroup law nor an exact flow.Source: §C.2 · p. 22
  • An invertible translation example has different optima for finite-step quotient and generator training: θ = −h versus θ = 0 for a zero target.Thus invertibility and identity at h = 0 alone do not justify substituting the zero-step derivative for the deployment quotient.Source: §C.2 · p. 22

Related papers

Loading 2610.17667v1…