arXiv cs.AI· Congzhong Sun, Ning Tan·· 3 天前AI 评分31
PrT-Flow：面向强化学习的概率可计算流匹配
PrT-Flow: Probability-Tractable Flow Matching for Reinforcement Learning
AI 导读
PrT-Flow 是一种概率可计算的流匹配方法，在保留多模态、多步动作生成的同时计算概率密度。其 Jacobian-Tractable Transformer（JTT）使每个采样步骤可逆，配合 Conditional Euler-step Training（CET）训练，并证明 CET 与拟合 Flow Matching 的 Euler 求解器更新具有相同极小值点。
正文 · 原文
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