ReinFlow: Fine-tuning Flow Matching Policy with Online Reinforcement Learning
TL;DR 精炼摘要
本文提出了ReinFlow,一个有效的在线强化学习框架,用于微调连续机器人控制的流匹配策略。通过引入可学习噪声并将流过程转换为离散马尔可夫过程,ReinFlow增强了探索能力并确保训练稳定性。实验表明,在腿部运动和视觉操作任务中,ReinFlow显著提升了奖励和成功率,同时减少计算时间。
摘要
We propose ReinFlow, a simple yet effective online reinforcement learning (RL) framework that fine-tunes a family of flow matching policies for continuous robotic control. Derived from rigorous RL theory, ReinFlow injects learnable noise into a flow policy's deterministic path, converting the flow into a discrete-time Markov Process for exact and straightforward likelihood computation. This conversion facilitates exploration and ensures training stability, enabling ReinFlow to fine-tune diverse flow model variants, including Rectified Flow [35] and Shortcut Models [19], particularly at very few or even one denoising step. We benchmark ReinFlow in representative locomotion and manipulation tasks, including long-horizon planning with visual input and sparse reward. The episode reward of Rectified Flow policies obtained an average net growth of 135.36% after fine-tuning in challenging legged locomotion tasks while saving denoising steps and 82.63% of wall time compared to state-of-the-art diffusion RL fine-tuning method DPPO [43]. The success rate of the Shortcut Model policies in state and visual manipulation tasks achieved an average net increase of 40.34% after fine-tuning with ReinFlow at four or even one denoising step, whose performance is comparable to fine-tuned DDIM policies while saving computation time for an average of 23.20%. Project webpage: https://reinflow.github.io/
论文精读
中文精读
1. 论文基本信息
1.1. 标题
ReinFlow: Fine-tuning Flow Matching Policy with Online Reinforcement Learning (ReinFlow:利用在线强化学习微调流匹配策略)
1.2. 作者
- Tonghe Zhang (卡内基梅隆大学,机器人研究所)
- Chao Yu (清华大学,电子工程系,通讯作者)
- Sichang Su (得克萨斯大学奥斯汀分校,航空航天工程系)
- Yu Wang (清华大学,电子工程系)
1.3. 发表期刊/会议
- 发表状态: 预印本 (Published at UTC: 2025-05-28)
- 领域影响力: 虽然目前是以预印本形式发布在 arXiv 上,但从其内容质量、实验规模(对比了 ICLR/NeurIPS 级别的基线)以及作者机构来看,这是一篇针对顶级机器人或机器学习会议(如 CoRL, NeurIPS, ICRA)的高质量投稿。
1.4. 发表年份
2025年
1.5. 摘要
该论文提出了 ReinFlow,这是一个简单而有效的在线强化学习(RL)框架,用于微调一类用于连续机器人控制的流匹配(Flow Matching)策略。基于严格的强化学习理论,ReinFlow 通过在流策略的确定性路径中注入可学习的噪声,将流过程转换为离散时间的马尔可夫过程(Markov Process)。这一转换使得似然(likelihood)计算变得精确且直接,从而促进了探索并确保了训练的稳定性。ReinFlow 能够微调多种流模型变体,包括 Rectified Flow 和 Shortcut Models,特别是能够在极少数甚至单步去噪步骤下进行微调。实验表明,ReinFlow 在具有挑战性的腿部运动和视觉操作任务中,相比于最先进的扩散强化学习微调方法(如 DPPO),不仅显著提升了奖励和成功率,还大幅减少了计算时间和推理步数。
1.6. 原文链接
-
arXiv 链接: https://arxiv.org/abs/2505.22094
-
项目主页: https://reinflow.github.io/
2. 整体概括
2.1. 研究背景与动机
- 核心问题: 在机器人学习领域,流匹配(Flow Matching) 模型因其精确建模、快速推理和易于实现的特点,已成为替代扩散策略(Diffusion Policies)生成机器人动作的热门选择。然而,这些模型通常依赖于模仿学习(Imitation Learning),其性能受限于演示数据的质量(通常混合了次优数据)。
- 现有挑战:
- 数据质量瓶颈: 仅靠模仿学习,机器人很难超越其训练数据的水平,且增加数据量往往会遇到收益递减的问题。
- 探索机制缺失: 模仿学习缺乏内在的探索机制,难以应对需要超越专家演示的复杂任务。
- 在线微调困难: 虽然在线强化学习(RL)可以解决上述问题,但将 RL 应用于流匹配策略面临技术难题。流策略通常是确定性的常微分方程(ODE),难以计算策略梯度的对数概率(log probability),特别是在极少推理步数下,离散化误差会很大,导致训练不稳定。
- 切入点: 论文提出通过“噪声注入”将确定性的流模型转化为随机的离散马尔可夫过程,从而解决概率计算和探索难题。
2.2. 核心贡献
- 算法设计 (ReinFlow): 提出了第一个能稳定微调流匹配策略的在线 RL 算法。核心创新是引入一个噪声注入网络(Noise Injection Network),将流转换为具有高斯转移概率的离散时间马尔可夫过程,实现了精确且易于计算的似然函数。
- 高效推理支持: 该方法支持在极少步数(甚至 1 步)下微调 Rectified Flow 和 Shortcut Models,在保持高性能的同时极大提升了推理速度。
- 显著的性能提升:
-
在腿部运动任务中,奖励平均净增长 135.36%。
-
在操作任务中,成功率平均净增长 40.34%。
-
相比最先进的扩散 RL 方法(DPPO),墙钟时间(Wall Time)减少了 62.82%。
-
3. 预备知识与相关工作
3.1. 基础概念
为了理解本文,初学者需要掌握以下概念:
- 流匹配 (Flow Matching): 一种生成模型技术。想象我们要把一堆无序的数据(比如噪声分布 )变成有序的数据(比如机器人动作分布 )。流匹配定义了一个随时间 变化的向量场(速度场 ),数据点沿着这个速度场移动,最终从噪声变成目标数据。这就像随波逐流,只要知道水流的速度,就能算出船的位置。
- Rectified Flow: 一种特殊的流匹配,它强制流的路径尽可能走直线。直线路径意味着推理时可以少走几步,速度更快。
- 常微分方程 (ODE): 描述流匹配过程的数学工具。位置的变化率(速度)是时间和当前位置的函数:。
- 在线强化学习 (Online RL): 智能体(Agent)在环境中实时交互,收集数据,并利用这些新数据不断更新策略。与“离线 RL”(只用历史数据)相对,在线 RL 更强调探索(Exploration),即尝试未知的动作以发现更好的策略。
- 策略梯度 (Policy Gradient): 强化学习的一类算法。其核心思想是计算动作的概率对奖励的梯度,从而调整参数使得高奖励动作的概率变大。计算的关键在于获得动作的对数概率 。
3.2. 前人工作与技术演进
- 扩散策略 (Diffusion Policies): 如 DDPM, DDIM。这是当前机器人动作生成的主流方法。已有工作尝试用 RL 微调扩散模型(如 DPPO [42], IDQL [23])。
- 局限性: 推理慢(需要多步去噪),且计算对数概率极其复杂或不准确。
- 流匹配策略 (Flow Matching Policies): 如 Flow Q-Learning (FQL) [38]。FQL 尝试离线训练流策略。
-
差异化: FQL 主要关注离线设置,缺乏在线探索机制。本文 ReinFlow 是首个针对流策略的在线微调方法,通过噪声注入解决了探索和概率计算问题。
-
4. 方法论
4.1. 方法原理
ReinFlow 的核心思想是:将确定性的流生成过程“随机化”。 原生的流匹配在推理时是确定性的:给定一个初始噪声和状态,解 ODE 得到的动作是固定的。这使得计算策略梯度所需的概率密度变得非常困难(通常涉及复杂的 ODE 积分和迹估计)。 ReinFlow 引入了一个辅助的噪声注入网络,在求解 ODE 的每一步都加入可学习的高斯噪声。这样,整个生成过程就变成了一个离散时间的马尔可夫链,每一步的转移概率都是简单的高斯分布,从而可以精确、快速地计算整个轨迹的概率。
下图(原文 Figure 7)展示了 ReinFlow 的微调流程:
该图像是示意图,展示了ReinFlow框架的整体结构。图中包含多个模块,起始于传感器采集数据,信号经过编码器处理后提取视觉运动特征。然后,特征被送入速度头和噪声注入网络,其中噪声注入网络的输出为可学习的噪声,最终生成机器人的动作。图形明确展示了数据流向和各部分的关系,直观表达了该框架的工作机制。
4.2. 核心方法详解 (逐层深入)
步骤 1: 定义流匹配与动作生成
在机器人控制中,流匹配策略将标准正态分布的噪声 转化为动作 。
对于 Rectified Flow,其对应的常微分方程(ODE)为:
其中 是由神经网络参数化的速度场。
为了生成动作,我们通常将时间 [0, 1] 离散化为 步(),并数值积分求解。
步骤 2: 噪声注入与马尔可夫化
为了能够进行强化学习微调,ReinFlow 在每一步积分时注入噪声。 设 为第 步的去噪动作。过程如下:
- 初始化: (从标准正态分布采样初始噪声)。
- 逐步去噪 (关键公式):
- 解释:
-
: 下一步的动作状态。
-
: 这是标准的欧拉积分步骤,沿着速度场方向移动一小步 。这是确定性部分。
-
: 这是注入的噪声项。
-
: 是一个专门的噪声注入网络,它根据当前时间、动作和观测 输出噪声的标准差。
-
: 采样的标准高斯噪声。
这一步将原本确定的 变成了一个服从高斯分布的随机变量:
-
- 解释:
步骤 3: 精确似然计算
由于上述过程是一个马尔可夫链,整个动作生成路径 的联合对数概率可以精确写出: 其中 ,。 相比于传统方法需要估算 ODE 的迹(trace estimator),这种计算方式是精确且解析的,即使只有 1 步去噪 () 也完全成立。
步骤 4: 策略梯度优化
有了精确的 ,我们就可以利用策略梯度定理进行优化。 ReinFlow 使用 PPO (Proximal Policy Optimization) 算法的变体。根据定理 4.1,对于参数化为马尔可夫过程的策略,其梯度为:
- 解释:
-
: 优势函数 (Advantage Function),衡量当前动作比平均水平好多少。
-
: 整个去噪轨迹的对数概率之和。
-
这意味着我们通过调整速度网络 和噪声网络 ,使得那些能获得高优势(高奖励)的去噪轨迹出现的概率变大。
具体的损失函数(结合 PPO 的 Clip 机制和正则化项)为: 其中 是新旧策略的概率比, 是正则化项(详见下文)。
-
步骤 5: 正则化与探索
为了防止策略在微调过程中偏离预训练知识太远,或为了增强探索,ReinFlow 引入了两种正则化:
- Wasserstein-2 () 正则化: 约束微调后的策略与预训练策略之间的 距离。 这实际上是让新策略生成的动作不要离旧策略太远。
- 熵正则化 (Entropy Regularization): 鼓励策略保持较高的随机性(探索)。 其中 是微分熵。最大化熵可以防止策略过早收敛到次优解。
5. 实验设置
5.1. 数据集
实验涵盖了机器人学习中的两类代表性任务:
-
OpenAI Gym Locomotion (状态输入):
-
任务: Hopper-v2, Walker2d-v2, Ant-v0, Humanoid-v3(难度递增)。
-
特点: 输入为低维状态向量,奖励是密集的(每一步都有数值奖励)。
-
预训练数据: D4RL 数据集(包含中等或专家级的演示)。
-
图示: 如下图(原文 Figure 8)所示:
该图像是一个示意图,展示了四种不同的机器人模型:Hopper、Walker2d、Ant和Humanoid。每个模型在不同的姿态下展示,突出了它们在连续机器控制中的应用。
-
-
Manipulation Tasks (操作任务):
-
Franka Kitchen: 状态输入,长程规划任务。
-
Robomimic: 视觉输入(像素),稀疏奖励(只有完成任务才给 +1,否则 0)。包括:PickPlaceCan (Can), NutAssemblySquare (Square), TwoArmTransport (Transport)。
-
特点: 极具挑战性,特别是稀疏奖励和高维视觉输入。
-
图示: 如下图(原文 Figure 9)所示:
该图像是四个示意图,展示了不同的机器人控制任务,包括(a) Franka厨房,(b) 取放罐,(c) 螺母组装方块,以及(d) 双臂运输。这些任务展示了ReinFlow框架在多种操作环境中的应用。
-
5.2. 评估指标
-
回合奖励 (Episode Reward):
- 概念定义: 一个完整回合中智能体获得的所有奖励的累加值。用于衡量 Locomotion 任务的性能。
- 公式:
- 符号解释: 为回合最大步数, 为第 步获得的奖励。
-
成功率 (Success Rate):
- 概念定义: 在多次试验中,智能体成功完成指定任务(如把罐子放到位置)的百分比。用于衡量 Manipulation 任务。
- 公式:
- 符号解释: 为成功的次数, 为总尝试次数。
5.3. 对比基线
-
DPPO (Diffusion Policy Policy Optimization) [42]: 目前最先进的针对扩散策略的在线 RL 微调方法。
-
FQL (Flow Q-Learning) [38]: 针对流匹配策略的离线 RL 方法,也可用于微调。
-
Behavior Cloning (BC): 仅使用监督学习预训练的基线,未进行 RL 微调。
6. 实验结果与分析
6.1. 核心结果分析
6.1.1. 墙钟时间效率 (Wall-time Efficiency)
ReinFlow 的一大优势是训练速度极快。 下图(原文 Figure 1)展示了在 OpenAI Gym 任务中,ReinFlow(红色和蓝色线)相比 DPPO(橙色线)在相同时间内能达到更高的奖励,且收敛极快。
该图像是包含四个子图的图表,展示了不同策略在Hopper-v2、Walker-v2、Ant-v2和Humanoid-v3任务中的平均回报与时间的关系。ReinFlow-S和ReinFlow-R在大多数任务中表现出显著的提升,相较于DPPO和FQL方法,表现更加优势。
6.1.2. 状态输入操作任务 (Franka Kitchen)
在 Franka Kitchen 任务中,ReinFlow-S (Shortcut Model) 在任务完成率上显著优于基线。 下图(原文 Figure 2)展示了结果:
该图像是三个任务完成率的曲线图,分别对应于Kitchen-complete、Kitchen-mixed和Kitchen-partial的任务类型。图中展示了ReinFlow-S、DPPO和FQL三种方法随样本数量变化的任务完成率,ReinFlow-S在各个任务中均表现优异。
- 分析: 在
Kitchen-mixed和Kitchen-partial这种数据质量不一的困难任务中,ReinFlow 展现了强大的微调能力,而 FQL 甚至难以超越 BC 基线。
6.1.3. 视觉操作任务 (Robomimic)
在极具挑战性的视觉输入、稀疏奖励任务中,ReinFlow 依然表现出色。 以下是原文 Table 4 的部分数据转录(关注成功率提升):
| 任务 (环境) | 算法 | 预训练成功率 | 微调后成功率 | 净增长 |
|---|---|---|---|---|
| Can (image) | ReinFlow-R | 59.00% | 98.67% | 39.67% |
| ReinFlow-S | 57.83% | 98.50% | 40.67% | |
| Square (image) | ReinFlow-R | 25.00% | 74.83% | 49.83% |
| ReinFlow-S | 34.50% | 74.67% | 40.17% | |
| Transport (image) | ReinFlow-S | 30.17% | 88.67% | 58.50% |
- 分析: 即使在极难的 Transport 任务(双臂协作)中,ReinFlow 也将成功率从 30% 提升到了 88% 以上。这验证了其在稀疏奖励下强大的探索能力。
6.2. 数据呈现 (训练曲线)
下图(原文 Figure 10)展示了样本效率(Sample Efficiency)的对比。
该图像是图表,展示了不同策略在多个任务(Hopper-v2、Walker2d-v2、Ant-v2 和 Humanoid-v3)上的 episode reward 随样本数量的变化。图中比较了 ReinFlow 系列(ReinFlow-S 和 ReinFlow-R)、DPPO 和 FOL 的表现。
- 分析: 在简单的 Hopper 任务中,FQL(绿色)样本效率较高,但在复杂的 Ant 和 Humanoid 任务中,ReinFlow(红色/蓝色)在达到更高奖励上限的同时,样本效率也不输 DPPO。
6.3. 消融实验与参数分析
6.3.1. 去噪步数的影响
ReinFlow 的一个关键特性是支持极少步数的推理。 下图(原文 Figure 4)展示了不同去噪步数的影响:
该图像是多个图表,展示了ReinFlow策略在不同环境中的表现。图(a)展示了Hopper-v2任务中的平均回报与推理步骤和预训练回合的关系,图(b)显示了Shortcut策略在Square环境中的成功率随预训练回合增加的变化,图(c)展示了ReinFlow策略在Square环境中的成功率随样本数量增加的趋势。
- 发现: 即使在 (单步推理,图 4a 左侧)的情况下,ReinFlow 也能实现显著的性能提升。这对于机器人实时控制至关重要。
6.3.2. 噪声注入策略
噪声网络该输入什么信息?
下图(原文 Figure 5)对比了仅输入状态()与输入状态+时间(s, t)的区别:
该图像是包含两个图表的结果展示,左侧图展示了噪声输入在Ant任务中的效果,右侧图展示了噪声条件在Kitchen-partial任务中的影响。两幅图分别呈现了不同输入配置下的平均集数奖励与任务完成率随样本变化的趋势,以表明ReinFlow策略的有效性。
- 结论: 同时以状态和时间为条件(State & Time Conditioning)能获得更高的成功率。这让噪声网络能根据去噪过程的不同阶段动态调整噪声大小。
6.3.3. 噪声水平与正则化
下图(原文 Figure 6)分析了噪声大小(Noise Scale)和正则化类型的影响:
该图像是图表,展示了噪声对探索的影响(a)和 (蓝色)以及熵(红色)正则化的效果(b)。图中可见,随着样本数量的增加,使用不同噪声标准差的策略的平均奖励呈现不同的增长趋势,而在正则化参数上, 和 的变化也显著影响了训练效果。
-
分析 (图 6a): 噪声太小(蓝色线)导致探索不足,奖励低;噪声适中(绿色/橙色)效果最好。
-
分析 (图 6b): 熵正则化(Entropy Reg,红色线)通常比 正则化(蓝色线)更能促进探索,获得更高的最终奖励。
7. 总结与思考
7.1. 结论总结
ReinFlow 是针对流匹配策略(Flow Matching Policies)的首个在线强化学习微调框架。通过创新性地引入可学习的噪声注入,它巧妙地解决了流模型在 RL 微调中面临的概率计算难和探索难的问题。
- 理论上: 它将连续流转化为离散马尔可夫过程,使得似然计算精确且高效。
- 实践上: 它在多项机器人控制基准测试中,以更少的推理步数和更短的训练时间,取得了超越现有 SOTA(如 DPPO)的性能。
7.2. 局限性与未来工作
- 噪声敏感性: 作者指出,算法性能对噪声幅度的超参数比较敏感,目前需要手动调整。未来希望实现自动调整。
- 模型规模: 目前实验主要在相对较小的网络上进行。未来值得探索将其扩展到大规模视觉-语言-动作(VLA)模型上。
- 样本效率: 目前的实现为了追求墙钟时间效率,牺牲了一定的样本效率。在真实机器人上训练(数据昂贵)时,可能需要进一步优化样本利用率。
7.3. 个人启发与批判
- 启发: 这篇论文提供了一个非常优雅的思路——“随机化确定性模型以利用 RL”。这种将确定性 ODE 视为随机马尔可夫过程特例(零噪声极限),并在训练时显式注入噪声来恢复概率属性的方法,不仅适用于流匹配,可能也对其他确定性生成模型的微调有借鉴意义。
- 批判: 虽然 ReinFlow 在仿真环境中表现优异,但其核心依赖于“噪声注入”来探索。在真实物理机器人上,高斯噪声可能会导致剧烈或危险的动作抖动(Jitter)。如何设计更平滑、符合物理约束的探索噪声(例如基于动作序列相关的噪声),可能是实机部署前需要解决的问题。此外,文章强调了单步推理的优势,但在单步下,流模型退化为 GAN 或 VAE 的某种形式,此时 ReinFlow 与传统的单步 RL 方法的本质区别和优势界限值得更深层的理论探讨。
相似论文推荐
基于向量语义检索推荐的相关论文。