Nagi-ovo

本文主要基于 Umar Jamil 的课程$^{[1]}$进行学习和记录。我们的目标是让 LLM 的行为与我们的期望的输出相一致，RLHF 则是最著名的技术之一。其标准流程涉及四个模型（听上去就很占显存，所以很多方法是去掉部分模型），这里只需记得一共需要四个即可：Reward、Actor、Critic 和 Reference Model，我们最后优化得到的模型是这里说的 Actor Model。

LLM to RL#

以前对 RL 的认识中，策略是告诉你在当前 State 下把你应该采取的 Action 的概率的东西，那这么说来，语言模型本身就可以看作一个 Policy：接收一个 Prompt (state)，输出下一个 token (action) 的概率，采样后得到一个新 state（token 被拼到 prompt 后），即相当于一个有 vocab_size 大小 Action Space 的 Policy，也是一个 RL Agent。

那这么说，还差一个提供 Reward 的东西（传统 RL 中一般是环境内置的奖励函数）

做一个 “Q-A-Reward” 的数据集可以实现这点，但是人类并不擅长寻找共识，但在比较优劣这点却很擅长。所以我们把方向转为：模型在 High Temperature 下 generate 多个 A，然后请领域专家（可以是人也可以是 AI Model）来选择出 Chosen / Prefer 的答案，标注出一个偏好数据集，用此训练出一个生成数值奖励的 Reward Model。

Reward Model#

这个 RM 是用一个预训练 LLM 如 Llama 来实现的。

当我们不想生成文本而是生成数值奖励时，可以投影到词汇表中的 Linear 替换为一个 one output feature (输出一个标量) 的 Linear，用来产生整个文本序列的单一评分值。

Reward Model Loss#

类似 Bradley-Terry 的参数化形式：

$Loss = -\log \sigma(r(x, y_w) - r(x, y_l))$

$y_w$ 表示 Chosen，$y_l$ 反之。因此当模型为 chosen 给出高 reward 时，

$r(x, y_w) - r(x, y_l)>0$

$\sigma(r(x, y_w) - r(x, y_{l))}\in(0.5,1)$

$-\log \sigma(r(x, y_w) - r(x, y_l))\in(0,1)$

这样损失会低，而模型为选择的答案给出低 reward 时损失会很高

HuggingFace 中 RewardTrainer 类接收一个 AutoModelForSequenceClassification 输入（即我们上面提到的模型结构）

Actor & Critic Model#

Trajectories#

如前所述，强化学习（RL）的核心目标是找到一个策略（policy, $\pi$），该策略能指导 agent 的行动，以获得最大可能的期望回报（expected return）。

数学上，我们将其表示为找到最大化目标函数 $J(\pi)$ 的最优策略 $\pi^*$：

期望回报 $J(\pi)$ 代表了智能体遵循策略 $\pi$ 时，在许多可能的生命周期或回合（episodes）中，预期能累积到的平均总回报。

它的计算方法是：考虑所有可能发生的轨迹（trajectory, $\tau$），并将每个轨迹的总回报 $R(\tau)$ 乘以该轨迹在策略 $\pi$ 下发生的概率 $P(\tau|\pi)$ 进行加权平均（或积分）。

• $E_{\tau \sim \pi}[\cdot]$ 表示当轨迹 $\tau$ 是根据策略 $\pi$ 生成时的期望值（Expected Value）。
• $R(\tau)$ 是在单个轨迹 $\tau$ 上获得的总回报（奖励）。
• $P(\tau|\pi)$ 是当智能体使用策略 $\pi$ 时，特定轨迹 $\tau$ 发生的概率。

轨迹 $\tau$ 就是 agent 经历的一系列状态和动作的序列，从初始状态开始。它是 agent 与环境交互的一种可能的 “故事” 或 “路径”。

• $s_t$：时间步 $t$ 的状态（State）。
• $a_t$：在时间步 $t$ 采取的动作（Action）（通常基于状态 $s_t$ 和策略 $\pi$）。

我们通常将环境建模为随机的（stochastic）。这意味着在同一个状态 $s_t$ 下执行相同的动作 $a_t$，并不总会导致完全相同的下一个状态 $s_{t+1}$。其中涉及到随机性。

下一个状态 $s_{t+1}$ 是从一个以当前状态 $s_t$ 和采取的动作 $a_t$ 为条件的概率分布中抽取的：

考虑到随机性的状态转移和 Agent 的策略，我们可以计算整个轨迹发生的概率。它由以下各项连乘得到：

1. Agent 在初始状态 $s_0$ 的概率：$p_0(s_0)$。
2. 对于轨迹中的每一个时间步 $t$：
   • 环境在给定 $s_t$ 和 $a_t$ 的条件下转移到状态 $s_{t+1}$ 的概率：$P(s_{t+1}|s_t, a_t)$。
   • 智能体根据其策略在状态 $s_t$ 选择动作 $a_t$ 的概率：$\pi(a_t|s_t)$。

（其中 $T$ 是轨迹的长度）。

在计算轨迹的总回报 $R(\tau)$ 时，我们几乎总是使用折扣回报（discounted rewards）。这意味着较早收到的回报比较晚收到的回报更有价值。

为什么？

• 反映了现实场景（今天到手的一美元比明天画饼的一美元更有价值）。
• 在持续性任务（没有固定终点的任务）中避免了无限回报的问题。
• 提供了数学上的便利性。

我们引入一个折扣因子（discount factor）$\gamma$，其中 $0 \le \gamma < 1$。$\gamma$ 越接近 0，Agent 越 “短视”（更关注眼前利益）；$\gamma$ 越接近 1，Agent
越 “有远见”（更关注长期回报）。

轨迹的总折扣回报计算如下：

• $r_t$ 是在时间步 $t$ 收到的即时回报（immediate reward）。
• $\gamma^t$ 是应用于时间步 $t$ 回报的折扣系数。

那在 LLM 中，轨迹是什么呢？前面提到了，模型是策略，prompt 是状态，next token 是 action，因此自回归生成中的这些 s, a 序列组成了轨迹。

Policy Gradient#

我们确定了强化学习的目标：找到一个最优策略 $π^∗$ 来最大化期望回报 $J(π)$。很好。但我们实际上如何表示并找到这个策略呢？

通常，尤其是在处理复杂问题时，我们不会去搜索所有可能的策略。我们会定义一个参数化策略（parameterized policy），记作
$π_θ$。你可以把 $θ$ 想象成一组 “旋钮” 或参数 —— 如果我们的策略是一个神经网络，那 $θ$ 可能就是网络的权重和偏置。

我们现在的目标就变成了：如何调整这些旋钮 $θ$ 来最大化我们的期望回报？

这意味着期望回报依赖于轨迹 $\tau$，而轨迹的分布又依赖于我们特定策略 $\pi_{\theta}$ 所选择的动作。改变 $\theta$，就改变了策略，改变了轨迹，也就改变了期望回报。

我们的策略 $\pi_{\theta}$ 是一个神经网络，我们会迭代地调整其参数 $\theta$ 来增加 $J(\pi_{\theta})$。这个更新规则看起来会非常熟悉（只是把梯度下降里的减号换成了加号）：

• $\theta_k$：我们在第 $k$ 次迭代时的参数。
• $\alpha$：学习率（步长大小）。
• $\nabla_{\theta} J(\pi_{\theta})|_{\theta_k}$：期望回报 $J$ 相对于参数 $\theta$ 的梯度（gradient），在当前参数 $\theta_k$ 处计算。它告诉我们在参数空间中哪个方向能最大程度地增加 $J$。

显然得到所有轨迹的成本是极高的，例如我们要采样出 max_token_length=100 的所有生成结果，因此我们可以用样本均值来近似期望：

应用到 LM Policy 上#

通过图中所示的生成流程得到这条采样轨迹中每个 state action pair 的对数概率，现在就可以反向传播来计算梯度，

然后将每个梯度乘以从 RM 中 reward 送入表达式来运行梯度上升优化：

High Variance#

PG 算法对于小问题效果较好，但是用于语言建模会有一些问题。

如何在不增加样本量的情况下减少方差呢？

1. 移除历史奖励： reward-to-go
   首先必须要承认的是，当前的 action 无法影响到过去已经获得的奖励，而过去的奖励增加了不必要的噪声，这和 RL 中的信用分配问题应该有点关系。因此去掉过去的项可以避免增加噪音，让估计的梯度与真实梯度拉近。所以与其从零开始计算轨迹的奖励，我们可以只考虑从当前时间步开始的动作的奖励

2. 引入 baseline
   RL 的研究已证实，引入一个依赖于 state 的项（比如计算轨迹奖励的函数，也可以是一个常数）可以减少方差。这里我们选择 价值函数 $V^\pi(s)$ .

Value Function#

$V^\pi(s)$ 告诉你根据当前策略进行行动，剩余轨迹的期望奖励是多少。

经典 RL 场景和 LM 场景中的价值定义例子：

实际操作中，用的是我们试图优化的那个 LM 做初始化，在其顶部再添加一个线性层来预估 value，这样 Transformer 层的参数可以同时用于语言建模（用将 token 投射到词表的层）和价值估计

前面说的 reward-to-go 在 RL 中被称为 Q 函数，即从当前状态开始采取这个动作，然后得到即时奖励，按照策略完成后续行动的预期奖励：

再通过引入 Value 函数得到 Q 与 V 的差异，这个差异被称为 Advantage 函数

这个 $A^\pi(s_{t,}a_t)$ 优势项表示这个特定动作，相对于在状态 $s$ 中可以采取的平均动作好多少

在图中红箭头指向的状态，向下移动的优势函数将高于其它动作的优势函数

$\nabla_{\theta}(J(\theta)) \approx \frac{1}{N}\sum_{i=1}^{N}\left(\sum_{t=0}^{T} \nabla_{\theta} \log \pi_{\theta}(a_{i,t} | s_{i,t})\right) A^{\pi}(s_{i,t}, a_{i,t})$

梯度乘上优势函数后，效果就变为了让策略增加具有高优势 action 的 logprob，并降低带来低平均回报 action 的 log prob。

GAE#

为了解决这个偏差 - 方差问题，可以使用 GAE（广义优势估计），本质上就是对所有的优势项的加权和，每项乘上一个衰减因子。

GAE 的目的是在策略梯度算法中，提供一个比原始回报 $R(τ)$ 或简单 TD 误差 $δ_t$​ 更好的优势函数 $A^π(s,a)$ 的估计值 $$A^t$​$，以降低梯度估计的方差，提高学习的稳定性和效率。

这个公式递归地定义了广义优势估计 $\hat{A}_t$。它不是只看一步的 TD 误差 $\delta_t$，而是综合了未来多步的 TD 误差信息。

这个递归式从轨迹（episode）的末尾（假设 $T$ 是最后一步，$\hat{A}_{T+1}=0$）向前计算：
* $\hat{A}_T = \delta_T + 0 = \delta_T$
* $\hat{A}_{T-1} = \delta_{T-1} + \gamma \lambda \hat{A}_T = \delta_{T-1} + (\gamma \lambda) \delta_T$
* $\hat{A}_{T-2} = \delta_{T-2} + \gamma \lambda \hat{A}_{T-1} = \delta_{T-2} + (\gamma \lambda) \delta_{T-1} + (\gamma \lambda)^2 \delta_T$
* ...
* 一般形式：$\hat{A}_t = \sum_{k=0}^{\infty} (\gamma \lambda)^k \delta_{t+k}$ （假设无限步长或在终止状态后的 $\delta$ 为 0）

参数 $\lambda$（$0 \le \lambda \le 1$）是 GAE 的关键，它控制着估计 $\hat{A}_t$ 的偏差（bias）和方差（variance）：

• 当 $\lambda = 0$ 时：
  • $\hat{A}_t = \delta_t$。GAE 退化为简单的单步 TD 误差。这种估计方差较低（因为它只依赖于下一步的信息），但可能偏差较高（因为它严重依赖于可能不准确的 $V^{\pi}(s_{t+1})$ 的估计）。
• 当 $\lambda = 1$ 时：
  • $\hat{A}_t = \sum_{k=0}^{\infty} (\gamma)^k \delta_{t+k}$。经过推导可以证明，这等价于 $\hat{A}_t = (\sum_{k=0}^{\infty} \gamma^k r_{t+k}) - V^{\pi}(s_t)$，也就是蒙特卡洛（Monte Carlo）回报减去基线（baseline）。这种估计偏差较低（因为它使用了从 $t$ 时刻开始的完整实际回报），但方差通常很高（因为它累积了多个时间步的随机性）。
• 当 $0 < \lambda < 1$ 时：
  • GAE 在上述两种极端情况之间进行插值。$\lambda$ 越接近 0，越偏向于低方差高偏差的 TD 估计；$\lambda$ 越接近 1，越偏向于高方差低偏差的 MC 估计。
  • 通过选择合适的 $\lambda$（例如 0.97），GAE 试图在偏差和方差之间取得一个良好的平衡，从而得到一个既相对准确（偏差可控）又相对稳定（方差较小）的优势估计。

语言模型的 Advantage#

如图，目标是提高在当前状态 “上海” token 的 logprob，降低 “巧克力” 的 logprob，因为选择” 上海 “的优势高于选择 “巧克力”（胡乱说的 token ）的优势。

重要性采样和离线学习#

在很多情况下，我们可能想要计算 $E_{x \sim p(x)}[f(x)]$，但是：

1. 我们很难或者无法直接从目标分布 $p(x)$ 中采样得到 $x$。
2. 或者从 $p(x)$ 采样效率很低，语言模型中就是这个问题，LM 采样成本太高了。

但是，我们可能可以很容易地从另一个替代的，或称为提议（Proposal ）的分布 $q(x)$ 中进行采样。

重要性采样 (Importance Sampling, IS) 是一种通过从不同的分布采样来估计目标分布期望的技术，将一个在概率分布 $p(x)$ 下计算的期望值 $E_{x \sim p(x)}[f(x)]$，转换为在另一个不同的概率分布 $q(x)$ 下计算一个相关函数的期望值 $E_{x \sim q(x)}\left[\frac{p(x)}{q(x)} f(x)\right]$。

这里的关键是引入了重要性权重（importance weight） $w(x) = \frac{p(x)}{q(x)}$。这个权重的作用是修正偏差：对于一个从 $q(x)$ 中抽到的样本 $x_i$，如果它在目标分布 $p(x)$ 中出现的概率更高（$p(x_i) > q(x_i)$），就会给它一个大于 1 的权重；反之，如果它在 $p(x)$ 中出现的概率更低（$p(x_i) < q(x_i)$），就会给它一个小于 1 的权重。这样加权平均后，就能得到对原始期望 $E_{x \sim p(x)}[f(x)]$ 的一个（通常是无偏或一致的）估计。

重要性采样允许我们：

1. 从容易采样的分布 $q(x)$ 中抽取样本 $x_1, x_2, ..., x_N$。
2. 通过计算加权平均来估计原始的期望值：

回到我们的场景。前面我们得到了 On-Policy 的策略梯度估计：

[!info]
On-Policy 的含义：即采集数据用的策略和训练时用的策略是同一个。由于计算时需要使用从当前策略 $\pi_{\theta}$ 采样生成的轨迹。这意味着每次策略更新后，旧数据就不能直接用了，样本效率低。至于 mini_batch_num > 1 时，后面用于更新的数据还算是 On-Policy 吗？严格意义上感觉不算，所以也可以理解为 Semi-On-Policy？（表达不一定严谨）。

而 On-Policy 则是强调当前策略模型能否和环境进行交互。$^{[2]}$

我们希望利用过去由旧策略 $\pi_{\theta_{OFFLINE}}$ 生成的数据（这些数据可能是大量存在的）来估计当前新策略 $\pi_{\theta_{ONLINE}}$ 的梯度。这样可以重复利用数据，提高样本效率。

回忆下重要性采样 IS 的原理： $E_{x \sim p(x)}[f(x)] = E_{x \sim q(x)}\left[\frac{p(x)}{q(x)} f(x)\right]$。
对应到我们的 PG（简化考虑单步决策）：
* 目标分布 $p(x)$ 对应新策略 $\pi_{\theta_{ONLINE}}(a|s)$。
* 采样分布 $q(x)$ 对应旧策略 $\pi_{\theta_{OFFLINE}}(a|s)$。
* 重要性权重为 $w_t = \frac{\pi_{\theta_{ONLINE}}(a_t|s_t)}{\pi_{\theta_{OFFLINE}}(a_t|s_t)}$。

将重要性权重应用到 On-Policy 梯度的每一项（每个时间步 $t$），得到标准的 Off-Policy 估计：

现在我们找到了可以在不每次从我们正在优化的策略（要训练的模型）中采样的情况下完整梯度上升优化，而是只采样一次，将轨迹保存到内存 / 数据库中，用 mini-batch 优化策略，然后用新策略初始化离线策略（采样的策略）

PPO Loss#

PPO Loss 主要由三个部分组成：策略损失（$L_{POLICY}$）、价值函数损失（$L_{VF}$）和熵奖励（$L_{ENTROPY}$）。

1. 策略损失 ($L_{POLICY}$)#

Clipped Surrogate Objective (裁剪的替代目标)

这是 PPO 的核心。你会注意到它和我们前面用重要性采样推导出来的 off-policy 策略梯度目标有点像，但有一个关键的改动。

• $\frac{\pi_{\theta}(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$: 这就是重要性采样比率，称之为 $r_t(\theta)$。它是在状态 $s_t$ 下，根据当前（在线）策略 $\pi_{\theta}$ 采取动作 $a_t$ 的概率，除以根据收集轨迹数据时使用的旧（离线）策略 $\pi_{\theta_{old}}$ 采取该动作的概率。这个比率修正了数据来自一个与我们当前试图改进的策略略有不同的策略这一事实。

• $\hat{A}_t$: 这是优势函数估计值，是使用 GAE 计算出来的，有助于平衡偏差和方差。它告诉我们，在状态 $s_t$ 下采取动作 $a_t$ 比在该状态下采取平均动作要好多少或差多少（根据当前的价值函数判断）。

• clip 函数: 这就是 PPO 的关键点所在。
  $\text{clip} \left( r_t(\theta), 1-\epsilon, 1+\epsilon \right)$
  它基本上是说：如果概率比率 $r_t(\theta)$ 偏离 1 太远（过高或过低），我们就把它 “裁剪” 掉。所以，如果 $r_t(\theta)$ 试图变成 $1.5$，而 $\epsilon$ 是 $0.2$，它就会被裁剪到 $1.2$。如果它试图变成 $0.5$，它就会被裁剪到 $0.8$。
  参数 $\epsilon$ (epsilon) 是一个小的超参数（例如 0.1 或 0.2），它定义了裁剪范围 $[1-\epsilon, 1+\epsilon]$。

• min 函数: 这个目标函数取了下面两项中的较小者：

  1. 未裁剪的目标: $r_t(\theta) \hat{A}_t$
  2. 裁剪后的目标: $\text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) \hat{A}_t$

  为什么要这样做呢？ 策略梯度的目标是增加具有正优势的动作的概率，并减少具有负优势的动作的概率。
  然而，使用重要性采样时，如果 $r_t(\theta)$ 变得非常大，可能会导致巨大的更新和不稳定性。PPO 通过裁剪这个比率来尝试保持新策略与旧策略的接近。

  • 如果 $\hat{A}_t > 0$ (好动作): 我们希望增加 $\pi_{\theta}(a_t|s_t)$。min 函数意味着如果 $r_t(\theta)$ 增长超过 $1+\epsilon$，目标函数就会被限制在 $(1+\epsilon)\hat{A}_t$。这可以防止策略在单次更新中变化过大，即使未裁剪的目标会建议一个更大的增幅。
  • 如果 $\hat{A}_t < 0$ (坏动作): 我们希望减少 $\pi_{\theta}(a_t|s_t)$。如果 $r_t(\theta)$ 缩小到 $1-\epsilon$ 以下，目标函数就会被限制在 $(1-\epsilon)\hat{A}_t$。（注意：当 $\hat{A}_t < 0$ 时，$r_t(\theta)\hat{A}_t$ 这一项在 $r_t(\theta)$ 较小时值较大（更接近零或为正），而 $\text{clip}(...) \hat{A}_t$ 也是在 $\text{clip}(...)$ 较小时值较大。这里的 min 实际上意味着当比率超出裁剪边界时，我们采取的是更悲观的更新步骤，或者说，导致 log 概率变化幅度更小的那一步。）
    更准确地说，当 $\hat{A}_t < 0$ 时，乘积 $r_t(\theta)\hat{A}_t$ 会随着 $r_t(\theta)$ 的增加而变得更负。min 操作确保了如果 $r_t(\theta)$ 偏离了 $[1-\epsilon, 1+\epsilon]$ 区间，我们不会让目标变得过于负（也就是说，我们不会过分地降低该动作的概率）。

2. 价值函数损失 ($L_{VF}$)#

这和前面的内容完全一样：

• $V_{\theta}(s)$ 是价值网络的输出（即 LLM 顶部再加一个线性层，用来预测从状态 $s$ 开始的期望累积奖励）。
• $\sum \gamma^{t'} r_{t'}$ 这一项（称之为 $G_s$ 或目标价值）是从状态 $s$ 开始，并遵循当前策略直到回合结束所观察到的实际折扣奖励总和。这是我们为 $V_{\theta}(s)$ 设定的经验目标。
• 这个损失函数就是预测值 $V_{\theta}(s)$ 和观察到的目标值 $G_s$ 之间的均方误差（MSE）。我们希望价值函数能够准确预测未来的奖励。这个价值函数对于计算优势 $\hat{A}_t$至关重要。

3. 熵奖励 ($L_{ENTROPY}$)#

• 这里的 $p(x)$（或者更准确地说是 $\pi_{\theta}(a|s)$，对于给定状态 $s$ 下所有可能的动作 $a$）代表了当前策略在给定状态下输出的动作概率分布。
• $\sum_x p(x) \log p(x)$ 这一项是这个概率分布的熵。熵衡量的是分布的随机性或不确定性。均匀分布（非常随机）具有高熵，而尖峰分布（对某个动作非常确定）具有低熵。
• 损失项是负熵。当我们在总损失 $L_{PPO}$ 中最小化这个 $L_{ENTROPY}$ 时（假设 $c_2$ 是正的），我们实际上是在最大化策略的熵。

鼓励更高的熵可以促进探索，会让策略变得更随机一些，尝试不同的动作（在 LLM 的情况下就是尝试不同的 token），而不是过快地收敛到一个可能是次优的确定性策略。这有助于 Agent 发现更好的策略。

最终形式 $L_{PPO}$#

最终的 PPO 损失是这三个部分的加权和：

• $c_1 L_{VF}$: 价值函数损失，由 $c_1$ 加权。$c_1$ 的一个常见值是 $0.5$ 左右。
• $c_2 L_{ENTROPY}$: 熵奖励（如果 $c_2 > 0$，实际上是对低熵的惩罚），由 $c_2$ 加权。$c_2$ 通常是一个小的正常数（例如 $0.01$），用于鼓励探索，同时又不会压倒主要的策略目标。

Agent 的参数（即 LLM 的权重）通过计算这个组合损失 $L_{PPO}$ 的梯度并执行梯度下降来进行更新。

Reference Model#

Reward Hacking#

RL 的一大问题就是 reward-hacking，模型可能会学会总是输出带来好奖励但对人类来说没意义的 token 或序列，比如连续说十遍 “谢谢你” 来提升礼貌分数，所以我们希望对齐的模型（RL post-training 后）的输出尽量和原本模型的输出较为相近。

因此会有另一个冻结了权重的模型（ref model），我们要优化的模型在每个轨迹的每一步中通过 reward model 生成奖励的时候，这个奖励会减去 ref model 与优化的模型 log prob 之间的 KL 散度，作为惩罚项来防止模型生成与原始模型差异过大的答案，以此防止上面所说的模型作弊现象。

Code walk through#

trl#

class AutoModelForCausalLMWithValueHead(PreTrainedModelWrapper):
    # ... (class attributes like transformers_parent_class) ...

这个类的核心目的是将一个标准的因果语言模型（Causal LM）（我们的 Actor Model，负责生成文本的策略 $π_θ$​）与一个 Value Head（即 Critic Model，负责估计状态价值 V (s)）捆绑在一起。在 PPO / Actor Critic 算法中，我们同时需要策略和价值函数，这个类就提供了一个统一的模型结构来同时输出这两者。

    def __init__(self, pretrained_model, **kwargs):
        super().__init__(pretrained_model, **kwargs) # 基础设置
        v_head_kwargs, _, _ = self._split_kwargs(kwargs) # 分离出给ValueHead的参数

        # 确保传入的是个有语言模型输出能力的模型
        if not any(hasattr(self.pretrained_model, attribute) for attribute in self.lm_head_namings):
            raise ValueError("The model does not have a language model head...")

        # 创建ValueHead实例，它将学习预测状态的价值 V(s)
        self.v_head = ValueHead(self.pretrained_model.config, **v_head_kwargs)

        # 初始化ValueHead的权重
        self._init_weights(**v_head_kwargs) # 默认随机初始化，也可以指定正态分布初始化

1. 充当 Actor: 即我们的语言模型 pretrained_model，它会根据当前 prompt（状态 s）生成回应（动作 a，即一系列 token）。
2. Critic: 评估 Actor 在某个状态 s 的 “好坏”，即输出 $V(s)$。这就是线性层 self.v_head 的任务。

    def forward(
        self,
        input_ids=None, # 输入的token IDs (状态 s)
        attention_mask=None,
        past_key_values=None, # 用于加速生成
        **kwargs,
    ):
        # 强制底层模型输出 hidden_states，ValueHead 需要它们作为输入
        kwargs["output_hidden_states"] = True
        # ... (处理 past_key_values 和 PEFT 的一些细节，PPO核心理解中可先忽略)

        # 1. Actor (基础语言模型) 进行计算
        base_model_output = self.pretrained_model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            **kwargs,
        )

        # 2. 提取 Actor 的输出 (用于策略更新) 和 Critic 的输入
        lm_logits = base_model_output.logits # Actor 的输出：预测下一个token的概率分布
        # 这是计算 PPO 中 L_POLICY 和 L_ENTROPY 的基础

        last_hidden_state = base_model_output.hidden_states[-1] # Critic 的输入：LM最后一层的隐藏状态，
        # 代表了当前状态 s 的表征

        # (可选) 语言模型本身的损失，在RL阶段通常不直接用
        loss = base_model_output.loss

        # (确保数据和模型在同一设备)
        if last_hidden_state.device != self.v_head.summary.weight.device:
            last_hidden_state = last_hidden_state.to(self.v_head.summary.weight.device)

        # 3. Critic (ValueHead) 进行计算
        # ValueHead接收状态表征，输出对该状态的价值估计 V(s)
        value = self.v_head(last_hidden_state).squeeze(-1) # 这是计算 PPO 中价值损失 L_VF 和优势 A_hat 的基础

        # (确保 logits 是 float32，为了数值稳定性)
        if lm_logits.dtype != torch.float32:
            lm_logits = lm_logits.float()

        # 返回 Actor 的 logits, LM loss (可能为None), 和 Critic 的value
        return (lm_logits, loss, value)

对于 PPO-RLHF 训练的每一步：

1. 我们把当前的一批 prompt (序列 input_ids) 输入模型。
2. self.pretrained_model (Actor) 会计算（Rollout）出 lm_logits。这些 logits 代表了在当前 prompt 下，模型认为接下来应该生成哪些词元的概率分布。PPO 的策略损失 $L_{POLICY}​$ 和熵奖励 $L_{ENTROPY​}$ 都需要基于这个概率分布$π_θ​(a_t​∣s_t​)$来计算。
3. 同时，我们从 base_model_output 中取出 last_hidden_state。这可以看作是当前 prompt (状态 s) 的一个向量表示。
4. 这个 last_hidden_state 被送入 self.v_head (Critic)，输出一个标量 value。这个 value 就是模型对当前状态 s 的价值估计 $V_θ​(s)$。PPO 的价值函数损失 $L_{VF}$​ 就是要优化这个 $V_θ​(s)$，使其尽可能接近真实的回报。并且，这个 $V_θ​(s)$也是计算优势函数 $A^t​$ 的关键组成部分，而 $A^t​$ 又会指导 $L_{POLICY​}$ 的计算。
5. 同样的 prompt + response 序列输入给 Reward 和 Reference model 做推理，得到 reward 和 log probs（计算 KL 惩罚）。

所以一次 forward 调用，我们就同时获得了更新 Actor (策略) 和 Critic (价值函数) 所需的核心信息。
训练的流程可以借下图帮助了解：

verl#

和 OpenRLHF 等都是优秀的 RLHF 框架，一个比较好的导读：【AI Infra】VeRL 框架入门 & 代码带读

Reference#

• [1] Umar Jamil: Reinforcement Learning from Human Feedback explained with math derivations and the PyTorch code. - YouTube
• [2] [AI Infra] VeRL 框架入门 & 代码带读
• [3] Chayenne Zhao: HybridFlow /veRL 原文浅析