reinflow

ReinFlow：用 PPO 微调 Flow Matching 策略

论文出处：reinflow
由chatgpt整理。

ReinFlow 是一种用在线强化学习微调 Flow Matching 控制策略的方法，主要面向机器人连续控制。

它通常不从零开始训练，而是先用模仿学习得到一个 Flow Matching Policy，再通过环境奖励继续优化，使策略可以突破示范数据的性能上限。该方法发表于 NeurIPS 2025。

1. 背景

Flow Matching Policy 通常从一个高斯噪声动作开始，通过速度网络逐步将其变成实际动作：

$$a^{k+1} = a^k+v_\theta(t_k,a^k,o)\Delta t$$

其中：

• $o$ 是机器人观测；
• $a^k$ 是第 $k$ 个生成步骤的中间动作；
• $v_\theta$ 是 Flow Matching 学到的速度场；
• 最终的 $a^K$ 是机器人执行的动作或动作序列。

这个过程本质上是确定性的 ODE 离散积分。给定初始噪声后，中间轨迹就确定了。

2. 为什么直接用 PPO

PPO 等策略梯度算法需要计算动作的概率：

$$\log \pi_\theta(a\mid o)$$

以及新旧策略之间的概率比：

$$r= \frac{\pi_\theta(a\mid o)} {\pi_{\theta_{\text{old}}}(a\mid o)}$$

这里就是一个比较有意思的点，其实理论上是可以直接计算动作的概率的，论文4.1“Likelihood Computation over a Short Denoising Trajectory“中

$$\ln p_1\bigl(\psi_1(x)\bigr) = \ln p_0\bigl(\psi_0(x)\bigr) - \int_0^1 \nabla \cdot v\bigl(t,\psi_t(x)\bigr)\,\mathrm{d}t, \qquad x \sim p_0(\cdot)$$

因为ODE 解的唯一性， 在固定时间 $t$ 下不会把两个不同初始点映射到同一个点。起点处概率可知，可以通过起点到终点的局部体积压缩/膨胀比来得到终点处动作的概率。这个过程可以用Monte-Carlo近似为

$$\ln \widehat{p}_1(x_1) = \ln p_0(x_0) - \sum_{k=0}^{K-1} \operatorname{tr} \left[ Z^\top \partial_x v_\theta(t_i, X_{t_i}) Z \right] \Delta t_i$$

但这个过程计算复杂度高，且只使用一个或少量生成步骤时离散误差会更加明显，同时确定性生成过程本身也没有直接的在线探索机制。

3. ReinFlow 的核心方法

ReinFlow 在每个 Flow 生成步骤中加入一个可学习的高斯噪声：

$$a^{k+1} \sim \mathcal N \left( a^k+v_\theta(t_k,a^k,o)\Delta t,\; \sigma_\phi^2(t_k,a^k,o) \right)$$

这里增加了一个噪声网络：

$$\sigma_\phi(t_k,a^k,o)$$

它根据观测、中间动作和当前 Flow 时间，输出每个动作维度的噪声标准差。

加入噪声之后，原来的确定性 Flow 轨迹变成了一条离散马尔可夫链：

$$a^0\rightarrow a^1\rightarrow\cdots\rightarrow a^K$$

由于每一步都是已知均值和方差的高斯分布，因此可以直接计算整条生成路径的联合概率：

$$\log \bar\pi(a^{0:K}\mid o) = \log p(a^0) + \sum_{k=0}^{K-1} \log p(a^{k+1}\mid a^k,o)$$

这里精确计算的是离散生成路径的联合概率，而不是直接计算最终动作 $a^K$ 的边缘概率。ReinFlow 通过对应的策略梯度定理，使用路径联合概率来更新策略，从而不需要显式求最终动作的边缘密度。

提问：为什么不再考虑空间膨胀和收缩？

“does not directly describe the marginal probability of the final action”只计算已经采样到的整条路径的联合概率而不是直接计算最终动作的概率，但是可以使用这个来优化策略。

4. 训练流程

ReinFlow 的训练大致分为以下过程。

预训练 Flow Policy

使用示范数据进行行为克隆或 Flow Matching 训练，得到速度网络：

$$v_\theta(t,a,o)$$

添加噪声网络和价值网络

增加：

• 噪声网络 $\sigma_\phi$，控制动作生成过程中的随机性；
• Critic 网络 $V_\psi(o)$，估计状态价值。

与环境交互

对于每个观测 $o$：

1. 采样初始动作噪声 $a^0$；
2. 执行 $K$ 个带噪声的 Flow 更新；
3. 得到最终动作 $a^K$；
4. 在环境中执行动作并收集奖励；
5. 保存整条内部生成轨迹 $a^{0:K}$。

计算优势函数

使用 Critic 和 GAE 等方法计算优势：

$$\hat A_t$$

优势为正表示该动作结果高于当前价值估计，优势为负表示低于当前估计。

使用 PPO 更新

ReinFlow 使用整条 Flow 路径的联合概率计算新旧策略概率比：

$$r_t= \frac{ \bar\pi_{\theta,\phi}(a_t^{0:K}\mid o_t) }{ \bar\pi_{\theta_{\text{old}},\phi_{\text{old}}} (a_t^{0:K}\mid o_t) }$$

然后使用 PPO clipped objective，同时更新：

• Flow 速度网络 $\theta$；
• 噪声网络 $\phi$；
• Critic 网络 $\psi$。

论文也给出了使用 SAC 等其他策略优化方法的可能形式，但实验主要采用 PPO。

5. 噪声网络的作用

噪声网络主要有两个作用。

第一是让每个 Flow 转移具有明确的高斯概率密度，从而可以计算策略梯度需要的 likelihood。

第二是提供探索。不同状态、不同生成阶段可以使用不同噪声：

• 不确定或需要探索时输出较大方差；
• 策略逐渐稳定后输出较小方差。

噪声大小通常受到上下界限制，避免动作扰动超过机器人关节或控制系统允许的范围。

6. 训练和部署的区别

训练阶段：

$$\text{Flow速度网络}+\text{噪声网络}$$

动作生成过程是随机的，用于探索和计算概率。

训练完成后，论文中的做法是丢弃噪声网络，只保留经过强化学习更新的 Flow 速度网络：

$$a^{k+1} = a^k+v_\theta(t_k,a^k,o)\Delta t$$

因此部署时仍然可以使用确定性 Flow Policy，不需要保留额外的噪声网络。

7. 与普通 PPO 的关系

普通高斯 PPO 直接输出动作分布：

$$a\sim\mathcal N(\mu_\theta(o),\sigma_\theta(o))$$

ReinFlow 不直接输出最终动作的高斯分布，而是让动作通过多步 Flow 转移生成：

$$a^0\rightarrow a^1\rightarrow\cdots\rightarrow a^K$$

因此其动作分布可以比单个对角高斯更复杂，并能够表示多峰动作分布。PPO 的主体结构没有明显改变，主要变化是策略概率由整条 Flow 生成路径的概率计算。

8. 方法特点与局限

ReinFlow 可以用于 Rectified Flow、Shortcut Model，以及原则上由 ODE 定义的其他 Flow Policy。它支持少量甚至单步动作生成，并在论文中的运动控制和机械臂操作基准上用于在线微调。

当前实现仍有以下限制：

• 主要采用 on-policy PPO，需要持续与环境交互；
• 对噪声上下界和噪声强度较敏感；
• 需要预训练好的 Flow Matching Policy；
• 论文实验主要基于模拟控制任务，真实机器人在线强化学习仍需进一步验证；
• 实验网络规模相对有限，扩展到大型 Flow VLA 模型的效果尚未得到充分验证。

概括来说，ReinFlow 的主要结构是：

$$\boxed{ \text{确定性 Flow} + \text{可学习高斯噪声} \rightarrow \text{可计算路径概率} \rightarrow \text{PPO 在线微调} }$$

9.补充 ODE 解的唯一性。

设流映射为

$$\psi_t(x_0)=x(t),$$

轨迹满足

$$\frac{\mathrm d x(t)}{\mathrm dt}=v(t,x(t)),\qquad x(0)=x_0.$$

假设两个不同初始点 $x_a\neq x_b$ 在时刻 $t^\ast$ 被映射到同一个点：

$$\psi_{t^\ast}(x_a)=\psi_{t^\ast}(x_b)=y.$$

那么从时刻 $t^\ast$ 的状态 $y$ 向后求解 ODE：

$$\frac{\mathrm d x}{\mathrm dt}=v(t,x),\qquad x(t^\ast)=y.$$

若 $v(t,x)$ 对 $x$ 满足 Lipschitz 条件，ODE 的解是唯一的。因此，从 $y$ 反向积分只能得到唯一一条轨迹，于是必须有

$$x_a=x_b,$$

这与 $x_a\neq x_b$ 矛盾。

所以固定时间的流映射

$$\psi_t:x_0\mapsto x(t)$$

是单射。

10.连续时间变量替换公式

设神经 ODE 为

$$\frac{\mathrm d X_t}{\mathrm dt} = v_\theta(t,X_t), \qquad X_0=x_0,\quad X_1=x_1.$$

概率密度 $p_t(x)$ 满足连续性方程：

$$\frac{\partial p_t(x)}{\partial t} + \nabla_x\cdot\left(p_t(x)v_\theta(t,x)\right) =0.$$

展开散度：

$$\frac{\partial p_t}{\partial t} + v_\theta^\top\nabla_x p_t + p_t\nabla_x\cdot v_\theta =0.$$

另一方面，沿轨迹 $X_t$ 对 $\log p_t(X_t)$ 求全导数：

$$\frac{\mathrm d}{\mathrm dt}\log p_t(X_t) = \frac{\partial}{\partial t}\log p_t(X_t) + \nabla_x\log p_t(X_t)^\top \frac{\mathrm dX_t}{\mathrm dt}.$$

代入

$$\frac{\mathrm dX_t}{\mathrm dt} = v_\theta(t,X_t),$$

得到

$$\frac{\mathrm d}{\mathrm dt}\log p_t(X_t) = \frac{1}{p_t} \left( \frac{\partial p_t}{\partial t} + v_\theta^\top\nabla_xp_t \right).$$

由连续性方程，

$$\frac{\partial \rho}{\partial t} + \nabla\cdot(\rho v)=0$$

$$\frac{\partial p_t}{\partial t} + v_\theta^\top\nabla_xp_t = -p_t\nabla_x\cdot v_\theta.$$

因此

$$\boxed{ \frac{\mathrm d}{\mathrm dt}\log p_t(X_t) = -\nabla_x\cdot v_\theta(t,X_t) }$$

这就是瞬时变量替换公式。

散度写成 Jacobian 的迹，向量场

$$v_\theta(t,x) = \begin{bmatrix} v_1(t,x)\\ \vdots\\ v_d(t,x) \end{bmatrix}$$

的散度为

$$\nabla_x\cdot v_\theta = \sum_{j=1}^d \frac{\partial v_j}{\partial x_j}.$$

而 Jacobian 为

$$\partial_xv_\theta = \frac{\partial v_\theta}{\partial x} = \begin{bmatrix} \frac{\partial v_1}{\partial x_1} & \cdots & \frac{\partial v_1}{\partial x_d} \\ \vdots & \ddots & \vdots \\ \frac{\partial v_d}{\partial x_1} & \cdots & \frac{\partial v_d}{\partial x_d} \end{bmatrix}.$$

所以

$$\nabla_x\cdot v_\theta = \operatorname{tr} \left( \partial_xv_\theta \right).$$

于是

$$\frac{\mathrm d}{\mathrm dt}\log p_t(X_t) = - \operatorname{tr} \left[ \partial_xv_\theta(t,X_t) \right].$$

从 $t=0$ 积分到 $t=1$：

$$\boxed{ \log p_1(x_1) = \log p_0(x_0) - \int_0^1 \operatorname{tr} \left[ \partial_xv_\theta(t,X_t) \right] \mathrm dt }$$

其中

$$x_1=X_1=\psi_1(x_0).$$

将时间区间划分为

$$0=t_0<t_1<\cdots<t_K=1, \qquad \Delta t_k=t_{k+1}-t_k.$$

使用左端点黎曼和：

$$\int_0^1 \operatorname{tr} \left[ \partial_xv_\theta(t,X_t) \right]\mathrm dt \approx \sum_{k=0}^{K-1} \operatorname{tr} \left[ \partial_xv_\theta(t_k,X_{t_k}) \right] \Delta t_k.$$

因此

$$\log p_1(x_1) \approx \log p_0(x_0) - \sum_{k=0}^{K-1} \operatorname{tr} \left[ \partial_xv_\theta(t_k,X_{t_k}) \right] \Delta t_k.$$

用随机向量估计 Jacobian 的迹。直接构造 $d\times d$ Jacobian 的开销很大，因此使用 Hutchinson 迹估计。

令随机向量 $z$ 满足

$$\mathbb E[z]=0, \qquad \mathbb E[zz^\top]=I,$$

例如

$$z\sim\mathcal N(0,I)$$

或每个元素独立取 $\pm1$ 的 Rademacher 分布。

对任意矩阵 $A$，

$$\mathbb E_z[z^\top A z] = \operatorname{tr}(A).$$

证明为

$$\begin{aligned} \mathbb E[z^\top A z] &= \mathbb E\left[\operatorname{tr}(z^\top A z)\right]\\ &= \mathbb E\left[\operatorname{tr}(Azz^\top)\right]\\ &= \operatorname{tr}\left(A\mathbb E[zz^\top]\right)\\ &= \operatorname{tr}(A). \end{aligned}$$

令

$$A=\partial_xv_\theta(t_k,X_{t_k}),$$

便有

$$\operatorname{tr} \left[ \partial_xv_\theta(t_k,X_{t_k}) \right] \approx z^\top \partial_xv_\theta(t_k,X_{t_k}) z.$$

若将一个或多个随机向量写成矩阵 $Z$，并满足

$$\mathbb E[ZZ^\top]=I,$$

则

$$\operatorname{tr} \left[ \partial_xv_\theta \right] \approx \operatorname{tr} \left[ Z^\top \partial_xv_\theta Z \right].$$

代回时间离散公式：

$$\boxed{ \ln \widehat p_1(x_1) = \ln p_0(x_0) - \sum_{k=0}^{K-1} \operatorname{tr} \left[ Z^\top \partial_xv_\theta(t_k,X_{t_k}) Z \right] \Delta t_k }$$

这就是图中的公式。

多个随机向量时的归一化，若

$$Z=[z_1,\ldots,z_m]$$

且每个 $z_j$ 都满足

$$\mathbb E[z_jz_j^\top]=I,$$

那么通常应写成

$$\operatorname{tr}(A) \approx \frac1m \operatorname{tr}(Z^\top A Z) = \frac1m\sum_{j=1}^m z_j^\top A z_j.$$

图中没有 $1/m$，通常意味着已经把 $Z$ 定义为

$$Z=\frac{1}{\sqrt m}[z_1,\ldots,z_m],$$

或者 $Z$ 实际上只有一列。

最后，符号 $\widehat p_1$ 上的帽子表示它不是完全精确的密度，因为包含两种近似：

$$\text{时间积分离散误差} + \text{随机迹估计误差}.$$

为什么能够加速计算

表面上看：

$$\epsilon^\top J\epsilon$$

仍然包含 Jacobian $J$。真正节省计算的地方在于，它可以通过一次自动微分得到。

定义标量：

$$s(x)=\epsilon^\top v(x).$$

对 $x$ 求梯度：

$$\nabla_xs(x) = \nabla_x\left(\epsilon^\top v(x)\right) = J^\top\epsilon.$$

这是一次普通反向传播能够直接计算的 vector-Jacobian product。

再和 $\epsilon$ 点乘：

$$\epsilon^\top\nabla_xs = \epsilon^\top J^\top\epsilon.$$

由于它是标量：

$$\epsilon^\top J^\top\epsilon = \epsilon^\top J\epsilon.$$

因此完整计算过程是：

$$\boxed{ \begin{aligned} s&=\epsilon^\top v(x),\\ g&=\nabla_xs=J^\top\epsilon,\\ \widehat{\operatorname{tr}(J)} &=\epsilon^\top g. \end{aligned} }$$

只需要：

• 一次速度网络前向传播；
• 一次反向传播；
• 两次向量点积。

不需要生成 $D\times D$ 的矩阵。