Truncated Rectified Flow Policy for Reinforcement Learning with One-Step Sampling

📄 arXiv: 2604.09159v1 📥 PDF

作者: Xubin Zhou, Yipeng Yang, Zhan Li

分类: cs.LG

发布日期: 2026-04-10

💡 一句话要点

提出截断修正流策略TRFP，解决最大熵强化学习中策略建模的局限性问题。

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

关键词: 强化学习 最大熵 修正流 多峰策略 单步采样

📋 核心要点

  1. 传统最大熵强化学习使用高斯策略，难以建模复杂的多峰动作分布，限制了其应用。
  2. 论文提出截断修正流策略（TRFP），通过混合确定性-随机架构，实现可处理的熵正则化优化。
  3. 实验表明，TRFP在多目标环境和MuJoCo基准测试中，能有效捕获多峰行为，性能优于现有方法。

📝 摘要（中文）

最大熵强化学习(MaxEnt RL)已成为序列决策的标准框架，但其标准高斯策略参数化本质上是单峰的，限制了其对复杂多峰动作分布的建模能力。这种限制激发了人们对基于扩散和流匹配的生成策略的日益增长的兴趣，作为更具表现力的替代方案。然而，将此类策略纳入MaxEnt RL具有挑战性，原因有两个：连续时间生成策略的可能性和熵通常难以处理，并且多步采样会引入长程反向传播不稳定性和大量的推理延迟。为了应对这些挑战，我们提出了截断修正流策略(TRFP)，这是一个建立在混合确定性-随机架构上的框架。这种设计使得熵正则化优化易于处理，同时支持通过梯度截断和流线化进行稳定的训练和有效的单步采样。在玩具多目标环境和10个MuJoCo基准上的实验结果表明，TRFP有效地捕获了多峰行为，在标准采样下优于大多数基准上的强大基线，并且在单步采样下仍然具有很强的竞争力。

🔬 方法详解

问题定义：最大熵强化学习（MaxEnt RL）旨在寻找最大化累积奖励和策略熵的策略。然而，标准MaxEnt RL通常使用高斯分布来参数化策略，这限制了其表达能力，尤其是在需要建模多峰动作分布的复杂环境中。此外，使用基于扩散模型或流匹配的生成策略时，计算似然和熵变得困难，并且多步采样会导致训练不稳定和推理延迟。

核心思路：TRFP的核心思路是设计一种混合确定性-随机架构，该架构既能保持策略的表达能力，又能使熵正则化优化变得易于处理。通过截断梯度和流线化，TRFP能够实现稳定的训练和有效的单步采样，从而克服了传统生成策略在MaxEnt RL中的局限性。

技术框架：TRFP框架包含以下主要组成部分： 1. 确定性部分：使用神经网络学习一个确定性的动作映射。 2. 随机部分：使用修正流（Rectified Flow）生成随机扰动，并将其添加到确定性动作中。 3. 梯度截断：在反向传播过程中截断梯度，以提高训练稳定性。 4. 流线化：通过优化修正流的参数，使流线更加平直，从而减少采样步骤。

关键创新：TRFP的关键创新在于其混合确定性-随机架构和梯度截断与流线化技术。这种架构允许TRFP在保持策略表达能力的同时，实现高效的训练和推理。与传统的基于扩散模型或流匹配的策略相比，TRFP避免了多步采样带来的问题，并简化了熵的计算。

关键设计：TRFP的关键设计包括： 1. 修正流的选择：选择合适的修正流函数，以保证生成的随机扰动具有良好的性质。 2. 梯度截断阈值的设置：合理设置梯度截断的阈值，以平衡训练的稳定性和收敛速度。 3. 损失函数的设计：设计合适的损失函数，以鼓励策略学习到最大化累积奖励和熵的动作分布。损失函数通常包含奖励项、熵正则化项和流线化项。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TRFP在玩具多目标环境和10个MuJoCo基准测试中表现出色。在标准采样下，TRFP在大多数基准测试中优于强大的基线方法。更重要的是，在单步采样下，TRFP仍然具有很强的竞争力，这表明其具有高效的推理能力。例如，在某些MuJoCo任务中，TRFP的性能提升超过10%。

🎯 应用场景

TRFP具有广泛的应用前景，包括机器人控制、自动驾驶、游戏AI等领域。在这些领域中，智能体需要学习复杂的策略以应对各种环境和任务。TRFP能够有效地建模多峰动作分布，提高智能体的适应性和鲁棒性，从而实现更高效、更可靠的决策。

📄 摘要（原文）

Maximum entropy reinforcement learning (MaxEnt RL) has become a standard framework for sequential decision making, yet its standard Gaussian policy parameterization is inherently unimodal, limiting its ability to model complex multimodal action distributions. This limitation has motivated increasing interest in generative policies based on diffusion and flow matching as more expressive alternatives. However, incorporating such policies into MaxEnt RL is challenging for two main reasons: the likelihood and entropy of continuous-time generative policies are generally intractable, and multi-step sampling introduces both long-horizon backpropagation instability and substantial inference latency. To address these challenges, we propose Truncated Rectified Flow Policy (TRFP), a framework built on a hybrid deterministic-stochastic architecture. This design makes entropy-regularized optimization tractable while supporting stable training and effective one-step sampling through gradient truncation and flow straightening. Empirical results on a toy multigoal environment and 10 MuJoCo benchmarks show that TRFP captures multimodal behavior effectively, outperforms strong baselines on most benchmarks under standard sampling, and remains highly competitive under one-step sampling.