arXiv：cs.AI· Congzhong Sun, Ning Tan·· 6 小时前
PrT-Flow：面向强化学习的概率可解流匹配
PrT-Flow: Probability-Tractable Flow Matching for Reinforcement Learning
AI 导读
PrT-Flow 通过 Jacobian-Tractable Transformer（JTT）让每个采样步可逆，在保留多模态、多步动作生成的同时计算概率密度，并用 Conditional Euler-step Training（CET）监督单步位移。
正文
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 .
| Subjects: | Robotics (cs.RO); Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2610.17667 [cs.RO] |
| (or arXiv:2610.17667v1 [cs.RO] for this version) | |
| https://doi.org/10.48550/arXiv.2610.17667 arXiv-issued DOI via DataCite (pending registration) |
Submission history
From: Congzhong Sun [view email]
[v1]
Tue, 13 Oct 2026 03:51:30 UTC (378 KB)
来源：arXiv：cs.AI · arxiv.org