原文发表在知乎，辛苦移步：π RL(piRL)算法支持用强化学习方法训练π 0/π 0.5(pi0/pi0.5)

最近看到清华大学发了一篇文章，解决了在强化学习方法下难以去训练pi0/pi0.5这种用flow matching生成动作的VLA模型的问题，效果看起来还不错。关于piRL的介绍可以参考：《清华大学最新！πRL：用在线强化学习让机器人 “边学边做” 的通用方案》。piRL笔者最近几天详细的研究了下，笔记如下。

先说一下核心点吧，对强化学习，特别是PPO算法，不熟悉的同学，可以先去看一下强化学习方面的知识。我们知道，PPO算法中对策略进行强化学习时，一个关键的核心点就是策略输出的动作需要有概率值，然后通过这个概率值去监督策略（模型）的训练，例如这个动作的的优势更大一些，那么就奖励这个动作，然后这个动作的概率就会更大一些。对于很多简单的任务，策略输出的动作很容易加上概率，只要改一下策略（模型），让其输出动作的mean和std，然后采样一下，就可以计算出动作的概率了。但对于用flow matching技术去生成动作的模型来说，稍微有些难度了，因为flow matching通过多步(例如10步)去噪的过程去生成一个精细的动作，那么如何在flow matching的基础之上让其可以输出动作的概率，就是需要解决的问题，piRL就是解决的这个问题。

笔者学习的代码是官方代码库中用PPO算法去训练pi0/pi0.5的代码。所以涉及到PPO算法，pi0/pi0.5，还有flow matching等技术点。PPO算法是强化学习里面比较主流且效果不错的算法，只要模型能输出动作的概率，都可以使用。pi0/pi0.5是目前市面上比较主流的VLA模型，它基于VLM(例如google的paligemma)和动作专家结合的范式，其中动作专家用的是flow matching技术。这些技术点里面，PPO算法/flow matching可参考本文附录中的章节，pi0/pi0.5模型可参考专栏：《π (pi)系列模型与算法》。

常规情况下，模型如何输出动作的概率：
可用如下伪代码表示：

模型输出动作均值与方差

x_t_mean, x_t_std = get_action_mean_std(…)

在均值与方差确定的分布(例如高斯分布)下面，采样一个具体的动作

x_t = x_t_mean + self.sample_noise(x_t.shape, device) * x_t_std

通过分布的概率密度函数确定上面动作的概率值

log_prob = self.get_logprob_norm(x_t, x_t_mean, x_t_std)
上面是一个经典的获取动作概率的方法，可以看到对模型的head头做一些简单的改造，让其可以输出动作的均值与方差即可。

使用flow matching时，怎么办？

如上图所示，类似于扩散过程，flow matching也是通过多步的去噪过程，将一个完全的噪声转移成一个精细的图片（在pi0/pi0.5中就是动作）。它是一个较为复杂的过程，在这个过程中如何拿到精细的图片的同时，也能拿到生成此图片的概率值呢？

在讲论文的方法之前，让我们先看一下pi0/pi0.5中具体是如何进行flow matching的过程的：

训练：

也用伪代码表示，其中x_t就是上图中每一个步骤的输入，对于图片去噪来说，就是混合了部分噪声的一张图片，对于action同理就是混合了部分噪声的动作。timestamp就是去噪的第几步，在pi0/pi0.5配置的是10步。timestamp就是从0~1，间隔是0.1。在训练时，不会完全执行10步的去噪，只会随机抽一个步骤去训练。因为每个步骤都是一样的，所以随机抽一个就可以，减少计算量。

noise: 完全的噪声，actions:完全无噪声的动作, timestamp->(0,1)

x_t = timestamp * noise + (1 - timestamp) * actions

通过transformer等技术获取速度场v_t

v_t = get_v_t(x_t, timestamp, cond)

速度场v_t的groud truth是：完全的噪声 - 完全无噪声的动作

u_t = noise - actions

均方差loss

loss = F.mse_loss(u_t, v_t, reduction=“none”)
推理：

也用伪代码表示，可以看到跟训练时不一样，推理时必须走完10步（好像是废话）。

timestamp = 1

最开始时完全的噪声

x_t = noise
while (timestamp >= 0):
v_t = get_v_t(x_t, timestamp, cond)
x_t -= 0.1 * v_t

x_t最终就是精细的动作

return x_t
咱们不讲原理和数学公式，从代码中可以看出来，在flow matching中貌似模型是推理出了一个平均去噪的速度，然后用这个速度逐渐去噪。

从上面的过程可以看出来，中间没有任何概率相关的东西，输出的v_t是一个确定的值。在此，piRL的思想看起来呼之欲出了：把v_t参照传统的思想一样，输出成分布的均值与方差，然后在分布里面采样v_t不就可以拿到概率了吗？很接近了，不过在论文中是每个去噪的步骤输出的是x_t的均值与方差，至于是v_t还是x_t，笔者觉得都差不多。论文中提供了两种概率化的方法：

flow_noise：

对每一去噪步，如下所示，模型推理时输出的仍然是速度v_t，通过简单转换可以转换成x_t，而x_t_std就是一个简单的网络头推理出来的。

x_t_mean应该是t-1时刻的mean

x_t_mean = x_t - 0.1 * v_t
x_t_std = self.noise_head(suffix_out)
有了每一去噪步的mean&std，总共10步，在数学上很容易将最终输出的x_0的概率算出来。也就是将去噪序列（从初始噪声到最终动作）的联合概率作为动作对数似然。跟上面论文中的原理图也是能对应上的。

flow_sde：

flow_sde更复杂一些，用数学语言来说的话，“将确定性 ODE 去噪转化为随机 SDE”，ODE主是常微分方程的缩写。SDE是随机微分方程的缩写。上面讲的pi0/pi0.5中的方法相当于ODE的方法，是一个确定性的过程。而SDE在数据建模的时候，就在公式中增加了随机性，相当于自带了随机性。如上图所示。上面公式中除了v_t（通过模型输出）外，其它相关参数都是确定的，不像flow_noise一样，还需要一个单独的网络输出方差，也就是flow_sde建模完成后，每一步的方差就是确定的，可直接写在配置文件中。在强化学习采样或训练的时候，也不必将10步去噪流程走完，随机选一步来计算均值和方差即可，其它的9步跟上面讲的pi0/pi0.5中模型学习训练的时候方法一样，也就是没有随机性。

项目代码实现与上图中的公式有所不同，结果是一样的，只是公式的表达方式不同。关于ode&sde可参考附录中的两篇文章或自行检索学习。

基于flow matching的动作输出可以计算概率后，剩下的就是PPO算法本身的逻辑了，不再赘述，可自行检索或参考下面附录。

附录：

stable_baselines3库中ppo算法的使用 ：
学习ppo算法笔者建议在一个简单的代码环境中去学习，直接在上面论文的代码库中去学习ppo容易被无关的东西干扰，因为论文中的代码库所基于的框架RLinf也挺复杂。而stable_baseline3就是一个比较简单易用，易学习的强化学习库。

一些环境安装方面的指引可参考文档：《一小时实践入门 stable-baselines3》。跑了一下CartPole任务。

CartPole任务
CartPole任务介绍

每存活 1 个时间步，立即获得 +1 的奖励。因此，一个回合的 总回报就是坚持步数。算法目标：最大化 Σ reward，即让杆子尽可能晚倒下。

下面的value net与actor net就是ppo算法中的critic/actor两部分，value net评估当前state下的值(value)函数值。actor net就是策略，输出当前state下的动作及其概率。在CartPole任务中，动作是离散的，即向左或向右，策略输出这两个动作的概率值。在一些任务中，动作是连续值，策略输出动作的均值与方差（上面的flow_noise与flow_sde就是这样），然后在均值与方差确定的正态分布中采样一个动作，同时计算此动作在分布中的概率。

value net：
输入：

就是上面的4个状态变量

模型结构：

(value_net): Sequential(
(0): Linear(in_features=4, out_features=64, bias=True)
(1): Tanh()
(2): Linear(in_features=64, out_features=64, bias=True)
(3): Tanh()
)
Linear(in_features=64, out_features=1, bias=True)
输出：

输出就是强化学习中value函数的值，例如tensor([[9.0188]])

actor net：
输入：就是上面的4个状态变量

模型结构：

(policy_net): Sequential(
(0): Linear(in_features=4, out_features=64, bias=True)
(1): Tanh()
(2): Linear(in_features=64, out_features=64, bias=True)
(3): Tanh()
)
Linear(in_features=64, out_features=2, bias=True)
输出：

上面模型结构输出的二维向量经过softmax后变成向左移动和向右移动的概率（加和等于1，例如tensor([[0.4771, 0.5229]]）。因为强化学习每次采样的动作需要有一定的随机性，所以动作就用“多分类版本的伯努利试验”的思想依概率选择一个，最终输出选择的动作与其概率，举例如下。注意，虽然0.5229比0.4771大，但不总是会选择这个概率大的。

动作：actions
tensor([1], device=‘cuda:0’) # 代表向右移动
概率：log_prob
tensor([-0.6484], device=‘cuda:0’) # torch.exp([-0.6484])=0.5229
roll out：
一次roll out就是策略与环境之间产生一次交互，产生state,action,reward,next_state等强化学习所需要的信息。在一次交互中记录了以下几个信息，其中action相关信息与value相关信息分别就是上面actor net与value net输出的。

state action reward done value函数值 action的概率

rollout_buffer.add(self._last_obs, actions, rewards, self._last_episode_starts, values, log_probs)
在此项目中一次性会收集2048个roll out，原因在下章节会讲解。

在flow_noise与flow_sde算法中，比cartpole任务会多在buffer中存储一些信息，如下：

chains(存储去噪过程中的x_t序列)，主要用于后面的重要性采样里面。在采样时，每个扩散步的结果是下一个扩散步的输入。但在训练时，因为要计算采样时的动作在新的策略下的概率，所以每个扩散步的输入是采样时记录的输入（即chains里面存储x_t）。核心点是：新旧策略的输入是一样的情况下，算出的动作概率才能对比。
denoise_inds（存储采样时选择的扩散步），如果denoise_inds = [7, 7, 7, 7, 7, 7, 7, 7, 7, 7]的话，在强化学习训练时，不会把10个扩散步都计算了，只会计算跟采样时选择的那个扩散步，即第7步。
GAE：
GAE全称广义优势估计，关于其定义与计算过程大家可自行检索。不同于像sac/td3等算法中对reward的简单使用，GAE核心有思想上有两点改进：

1，比较优势。采取当前action后获取的收益，应该是一个相对的概念。例如某个状态下有3种可能的动作，如果执行后可获得的收益分别是1,2,3，平均值是2，比较优势就是相对平均收益的diff。分别是-1,0,1，也就是第3种动作的收益是1，而不是直接获取到的3。大家可以看出来比较优势更加能代码当前action相对于其它可能action的真正收益。

2，考虑未来。当前action后获取到的收益，不仅要考虑当前直接可以获取到到的reward，也应该考虑预期的未来的总体收益。例如某个动作虽然即时收益很高，但可能是饮鸩止渴，不符合“可持续发展”。这样的动作的收益也不应该很高。

PPO算法的核心之一就是GAE。从上面的分析可以看出，GAE是一种更合理的收益计算方法。

从GAE的计算过程可以发现，计算之前需要拿到episode内连续的轨迹。所以上面的roll out一次性会收集2048个step，这2048个会根据实际的情况分隔成多段。例如step=100的时候cartpole的杆子倒下了，那么这个episode就done了。GAE会在这100个连续的step基础之上来计算。

重要性采样：
重要性采样也是ppo算法的核心点之一，关于重要性采样的必要性，kimi总结得不错：

具体到ppo这个场景中来说：

问题：

强化学习中，价值函数
，其中
是从状态 s 出发的轨迹分布。直接采样
需与环境大量交互，且很多轨迹回报极低（效率差）。
就是当前训练时的最新策略所产生。

解决方案：

用重要性采样复用 “其他策略（如探索策略，也就是roll out与环境交互时的那时的策略）产生的轨迹”—— 从探索策略的轨迹分布
采样，通过权重
修正，无需重新与环境交互，大幅降低训练成本。搜索策略就是采样时那时的策略。因为采样时的轨迹是现成的，相关的信息，例如优势的计算都是ready的。复用搜索策略产生的轨迹还有一点好处就是可以重复多次使用，提升数据使用效率。

另外，还需要限制一下新旧策略的diff范围，保证每次更新不会偏离旧策略太远，从而稳定训练。常见做法是把w限制在区间
：

具体到代码层面：

ratio就是上面两个概率的比值

ratio = th.exp(log_prob - rollout_data.old_log_prob)
…

限制ratio的上下限

policy_loss_1 = advantages * ratio
policy_loss_2 = advantages * th.clamp(ratio, 1 - clip_range, 1 + clip_range)
policy_loss = -th.min(policy_loss_1, policy_loss_2).mean()
loss设计：
critic loss：

self.returns = self.advantages + self.values
…

rollout_data.returns就是上面的self.returns

values_pred就是预估的value值

value_loss = F.mse_loss(rollout_data.returns, values_pred)
policy loss：

参考重要性采样中的代码截取。

原理解读：

以下总结来源于豆包：

flow matching：
可参考这篇文章：《Flow Matching生成模型：从理论基础到Pytorch代码实现-阿里云开发者社区》

例如对于一张图片来说：一张 256×256 的 RGB 图像可以展平成一个 256×256×3 = 196608 维的向量 x。 于是：

p₀(x)：在训练开始时，图像向量 x 服从的分布（通常是简单的高斯噪声）。
p₁(x)：我们希望模型最终输出的图像向量 x 服从的分布（真实图像的分布，比如人脸、猫、棋盘格等）。
Flow Matching 做了什么？

它学习一个向量场 v(x, t)，把 p₀（简单分布）里的点 x₀（噪声）沿着一条光滑的轨迹 x(t) 送到 p₁（目标分布）里的点 x₁（真实图像）。 这条轨迹满足：

训练完成后，从 p₀里采样一个噪声 → 用向量场“流动” → 得到一张服从 p₁ 的图像。

环境安装：
本来是要把论文中的代码框架(也是一个比较健全的框架，名为RLinf，专门用于强化学习)跑起来再看下代码的执行过程，按官方的文档来操作，安装环境用docker镜像挺简单，问题是这个代码跑起来需要较多的资源，把所有能调小的配置参数都调小了，但笔者本地内存32G+GPU(24G显存)仍然跑不起来，会内存OOM。云端autodl这种环境不支持docker镜像，所以暂时放弃跑起来的想法。再说笔者的目标也不是为了研究论文中的代码框架，所以就直接硬读代码的方式去看piRL的实现细节。

其它参考文章：

关于ode/sde/diffusion的关系，有几篇文章讲得还可以，可参考：

《Diffusion Model与SDE、ODE之间的联系》

《扩散模型 Diffusion Models - 原理篇》