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

📄 arXiv: 2610.17667v1 📥 PDF

作者: Congzhong Sun, Ning Tan

分类: cs.RO, cs.AI

发布日期: 2026-10-13

🔗 代码/项目: GITHUB

💡 一句话要点

提出PrT-Flow以解决强化学习中的概率密度计算问题

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 强化学习 概率密度 机器人控制 可逆网络 流匹配 模仿学习 条件欧拉步训练

📋 核心要点

  1. 现有的强化学习算法在处理机器人动作时，往往无法有效计算每个动作的概率密度，限制了其灵活性和适用性。
  2. PrT-Flow通过引入雅可比可处理变换器（JTT）和条件欧拉步训练（CET），实现了在多模态动作生成中高效计算概率密度。
  3. 在Robomimic Lift和OGBench Cube-Triple的实验中，PrT-Flow分别达到了96%的成功率，显示出其在模仿学习中的有效性。

📝 摘要（中文）

扩散模型和流匹配能够表示灵活的机器人动作分布，但许多强化学习算法也需要每个动作的概率。我们提出PrT-Flow（概率可处理流匹配），它在保留多模态、多步动作生成的同时计算概率密度。其雅可比可处理变换器（JTT）使每个采样步骤可逆，并能够高效计算雅可比行列式。我们通过条件欧拉步训练（CET）来训练这些有限步骤，该方法在采样过程中监督其位移而非瞬时速度。我们证明在某些假设下，CET与流匹配的欧拉求解器更新具有相同的最小化解。PrT-Flow在Robomimic Lift上经过300个模仿学习周期后取得96%的成功率，匹配扩散策略参考，并在OGBench Cube-Triple上经过60万在线步骤后也取得96%的成功率。这些结果表明，具有可处理概率密度的可逆网络能够学习有效的机器人策略，并支持SAC风格的演员-评论家训练。

🔬 方法详解

问题定义：现有的强化学习方法在生成机器人动作时，往往无法同时满足灵活性和概率密度计算的需求，导致在复杂任务中的表现不佳。

核心思路：PrT-Flow通过引入雅可比可处理变换器（JTT），使得每个采样步骤可逆，从而能够高效计算其雅可比行列式，解决了动作生成中的概率计算问题。

技术框架：整体架构包括两个主要模块：雅可比可处理变换器（JTT）用于实现可逆采样，条件欧拉步训练（CET）用于监督采样过程中的位移。这一框架确保了多模态动作生成的同时，能够计算出每个动作的概率密度。

关键创新：最重要的创新在于引入了条件欧拉步训练（CET），它在采样过程中关注位移而非瞬时速度，从而与流匹配的欧拉求解器更新保持一致，提升了学习效率。

关键设计：在网络结构上，PrT-Flow采用了可逆网络设计，确保了每个步骤的可逆性；在损失函数上，CET的设计使得训练过程更加稳定，能够有效地优化动作生成的概率密度。

🖼️ 关键图片

img_0
img_1
img_2

📊 实验亮点

在Robomimic Lift上，PrT-Flow经过300个模仿学习周期后取得了96%的成功率，匹配了扩散策略的参考表现；在OGBench Cube-Triple上，经过60万在线步骤也达到了96%的成功率，显示出其在强化学习中的优越性。

🎯 应用场景

该研究具有广泛的应用潜力，尤其在机器人控制和自动化领域。通过有效学习机器人策略，PrT-Flow能够提升机器人在复杂环境中的自主决策能力，未来可能在智能制造、无人驾驶等领域发挥重要作用。

📄 摘要（原文）

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 .