理解后训练所需的强化学习

监督微调像看示范：题目给定，正确回答也摆在旁边。强化学习换了一种反馈方式：你先采取行动，再观察结果和奖励。

假设任务是计算 $(2+3)\times4$ 。系统没有逐字提供解答，只允许先选一种运算合并 2 和 3，再选一种运算处理 4。最后得到 20，记 1 分；否则记 0 分。

哪一步该得到奖励？失败了，是第一步选错，还是第二步选错？如果只保留那 1 分、0 分，参数又怎样更新？我们用这个只有四条路径的小环境，把几个核心概念连起来。

1. 状态、动作与策略：当前在哪，接下来做什么

在第一步，系统面对 2 和 3，可以选加法或乘法：

选加法 → 中间结果 5
选乘法 → 中间结果 6

第二步可以把中间结果乘 4，也可以加 4。于是所有路径都能画出来：

两步算式环境，第一步加法或乘法得到5或6，第二步乘4或加4，四个终点中只有20获得奖励1。

图 1：本环境刻意很小。第一步选加法还不够，第二步也要选对，整条路径才得到奖励。

我们给图里的角色起几个名字：

名称	在这个例子里是什么
状态 $s$	当前步骤和已有中间结果
动作 $a$	选择加法或乘法
策略 $\pi_\theta(a\mid s)$	在该状态选择各动作的概率
环境	根据动作计算中间结果，推进到下一状态
奖励 $r$	对一步行动或终局结果给出的分数
回合	从开始到终止的一次完整尝试

策略可以很复杂，也可以只是一张表。这里有三个非终止状态：开始、得到 5、得到 6。每个状态都有两个动作分数，把分数做 Softmax 就得到概率。

“状态”不是随便挑几个数字就够了。我们希望它包含判断后续所需的信息。语言模型生成时，已经看到的提示与已经写出的前缀可以作为状态；下一个词元是动作，追加词元后进入新的前缀状态。

2. 奖励与回报：眼前一分，还是最终一分

奖励是某一步收到的反馈；回报把从当前开始的后续奖励合起来。

常见的折扣回报写作：

$G_t=r_t+\gamma r_{t+1}+\gamma^2r_{t+2}+\cdots$

$\gamma$ 是折扣因子。若 $\gamma=0.9$ ，三步奖励为 $[0,0,1]$ ，那么三步开始时的回报分别是：

$[0.81,\ 0.9,\ 1]$

折扣让较远的奖励在当前回报中权重更小。是否需要折扣，要看任务。有限长度的文本回答常关心最终质量；本章的两步任务取 $\gamma=1$ ，两步都使用终局的 1 或 0。

给每写一个词就加一分，会让模型偏向写长；只在最终答案正确时加一分，则让反馈变得稀疏。奖励怎样定义，会改变策略实际追求什么。这件事比“选哪种优化器”更早决定训练方向。

3. 价值函数：从这里出发，平均能拿多少分

如果当前策略在两步都各有一半概率选对，成功概率是：

$\frac12\times\frac12=\frac14$

开始状态的平均回报为 0.25。到达中间结果 5 后，只要最后一步乘 4 就能得分，因此这里的平均回报为 0.5；到达 6 后，两个动作都得不到 20，平均回报为 0。

状态价值函数把这个“从这里按当前策略继续，预期拿多少分”记为：

$V^\pi(s)=\mathbb E_\pi[G_t\mid s_t=s]$

$\mathbb E$ 表示平均或期望，下标 $\pi$ 提醒我们：价值取决于后续怎样行动。换一个策略，状态价值也会改变。

如果先指定一个动作，再按策略继续，得到的是动作价值：

$Q^\pi(s,a) =\mathbb E_\pi[G_t\mid s_t=s,a_t=a]$

比较该动作与状态平均水平，可以定义优势：

$A^\pi(s,a)=Q^\pi(s,a)-V^\pi(s)$

在刚才均匀随机的策略下，开始时选加法的动作价值是 0.5，减去平均值 0.25，优势为 0.25；选乘法的优势则为 $-0.25$ 。它们表达的是“比在这里通常能得到的结果更好还是更差”，不是这个动作脱离上下文后的永久好坏。

实际问题很难枚举所有未来，价值函数通常也要学习。我们的程序用三个可训练数估计三个状态的价值，再用采样到的回报调整它们。预测值与真实期望接近，但不保证完全相同。

4. 没有标准动作，策略怎样求梯度

先把问题缩成只有一步的情形。每个动作有奖励 $r(a)$ ，平均奖励是：

$J(\theta)=\sum_a\pi_\theta(a)r(a)$

我们想让 $J$ 增大。利用 $\nabla p=p\nabla\log p$ ，可以把梯度改写成：

$\begin{aligned} \nabla_\theta J &=\sum_a r(a)\nabla_\theta\pi_\theta(a)\\ &=\sum_a\pi_\theta(a)r(a) \nabla_\theta\log\pi_\theta(a) \end{aligned}$

最后一行的形式很有用：按策略采样一个动作，就能用它的奖励乘上对数概率的梯度，得到一个有噪声的梯度估计。反复采样后，估计趋向所需的平均方向。

这就是策略梯度方法的基本想法。Williams 的 REINFORCE 研究是这条路线的经典来源。

多步尝试，把沿路的动作连起来

本例是 $\gamma=1$ 的有限回合。每次尝试会有两个被选中的动作，策略损失可以写成：

$\mathcal L_{\mathrm{actor}} =-\sum_t \hat A_t \log\pi_\theta(a_t\mid s_t)$

$\hat A_t$ 是采样得到的优势估计。前面的负号把“最大化奖励方向”变成优化器熟悉的“最小化损失”。实际代码再对一批回合取平均。

若 $\hat A_t>0$ ，训练倾向于提高所选动作的概率；若 $\hat A_t<0$ ，倾向于压低它。这里的反馈不是一份标准动作标签，而是行动以后获得的相对表现。

5. 基线：一次成功，比平时好多少

直接拿回报 $G_t$ 乘梯度也能训练，不过波动可能较大。常用的做法是减掉一个只依赖状态的基线：

$\hat A_t=G_t-V_\phi(s_t)$

$V_\phi$ 是当前学到的价值估计，参数记为 $\phi$ ，与策略参数 $\theta$ 区分。

为什么减掉一个数不会凭空改变目标？对固定状态，基线不依赖这次抽到的动作；所有动作概率之和为 1，其梯度为 0，所以它在策略梯度期望中的贡献为零。合适的基线可以减少估计波动。

程序中，价值用作基线时要停止梯度：

advantages = returns - values.detach()[states]
actor_loss = -(advantages * chosen_log_probs).sum(1).mean()

否则，优化器可能通过改变基线本身来减少这个表达式，而不是只按我们期望的策略梯度更新策略。价值函数有自己的任务：让它预测的回报接近观察到的回报，通常可以使用平方误差。

信用分配：成功属于哪一步

我们的奖励只在最后出现。成功时，正确的两步都得到正向信号；失败时，却未必每一步都错。

例如，先算出 5，最后错选加 4，整条路径得 0 分。第一步其实有用，但这一次尝试本身没有告诉训练器“第一步没问题”。学习依靠许多次尝试：从 5 出发有时能成功，从 6 出发始终失败，两类路径的统计逐渐区分开来。

把长期结果归因到之前的行动，称为信用分配。回合越长、奖励越稀疏，区分哪些行动真正有帮助就越困难。价值估计、中间反馈和更好的探索可以提供帮助，也可能引入新的误差。

6. 动手：从随机选择到高概率成功

policy_gradient.py 用三个状态、两个动作训练一张策略表，同时学习状态价值。它不训练语言模型，也不声称学会了通用算术。

每次更新采样 32 个两步回合，共更新 200 次。初始动作分数全为零，因此每个动作概率为一半。运行：

uv run --python 3.12 --with-requirements requirements.txt python src/examples/policy_gradient.py

任务：先合并 2 和 3，再与 4 运算；最终结果 20 得 1 分。
训练前成功概率=0.250；训练后=0.973
第一步选择加法的概率=0.986
得到 5 后选择乘法的概率=0.987
学得的状态价值：0.962, 0.981, 0.000
当前策略的精确状态价值：0.973, 0.987, 0.000
奖励 [0, 0, 1]、gamma=0.9 的回报：0.81, 0.90, 1.00
每 20 次更新的成功概率：0.250, 0.573, 0.788, 0.873, 0.914, 0.935, 0.951, 0.960, 0.965, 0.969, 0.973

这里的成功概率不是从另一次小样本测试中估出的频率。环境只有四条路径，能够根据学得的策略直接算出精确期望。两条正确动作概率相乘，就得到约 0.973。

策略训练过程中精确成功概率逐步提高，起点为四条路径中一条正确对应的0.25。

图 2：固定采样种子的一次训练记录。当前环境可枚举，所以曲线表示策略的精确成功概率；不是大模型推理基准成绩。

价值估计没有与精确值完全重合。这很正常：它来自有限采样，而且策略一直在变化，预测目标也跟着移动。若把这张价值表当作绝不会出错的“裁判”，就误解了它的角色。

展开完整代码：采样、回报、策略梯度与价值学习

"""第 28 章：两步算式环境中的策略梯度与状态价值，不训练语言模型。"""

import torch
from torch import nn


def discounted_returns(rewards, gamma):
    result, future = [], 0.0
    for reward in reversed(rewards):
        future = reward+gamma*future
        result.append(future)
    return list(reversed(result))


def expected_reward(logits):
    probabilities = logits.softmax(-1)
    # 只有先 2+3、再乘 4 的路径得到 1 分。
    return probabilities[0, 0]*probabilities[1, 0]


def train(steps=200):
    torch.set_num_threads(1)
    generator = torch.Generator().manual_seed(2026)
    logits = nn.Parameter(torch.zeros(3, 2, dtype=torch.float64))
    values = nn.Parameter(torch.zeros(3, dtype=torch.float64))
    optimizer = torch.optim.SGD([{"params": [logits], "lr": .2},
                                {"params": [values], "lr": .1}])
    history = [(0, expected_reward(logits).item())]
    for step in range(1, steps+1):
        probabilities = logits.softmax(-1)
        first = torch.multinomial(probabilities[0].detach(), 32, replacement=True, generator=generator)
        second_states = 1+first  # 状态 1 的中间结果为 5，状态 2 的中间结果为 6。
        second = torch.multinomial(probabilities[second_states].detach(), 1, generator=generator).squeeze(1)
        states = torch.stack([torch.zeros_like(first), second_states], dim=1)
        actions = torch.stack([first, second], dim=1)
        final_results = torch.where(second == 0, (5+first)*4, (5+first)+4)
        rewards = (final_results == 20).double()
        # gamma=1；两步都使用同一个终局回报。价值作为基线时停止梯度。
        returns = rewards[:, None].expand(-1, 2)
        advantages = returns-values.detach()[states]
        chosen_log_probs = logits.log_softmax(-1)[states, actions]
        actor_loss = -(advantages*chosen_log_probs).sum(1).mean()
        value_loss = .5*((values[states]-returns)**2).mean()
        optimizer.zero_grad(set_to_none=True)
        (actor_loss+value_loss).backward()
        optimizer.step()
        if step % 20 == 0:
            history.append((step, expected_reward(logits).item()))
    return logits.detach(), values.detach(), history


def main():
    logits, values, history = train()
    probabilities = logits.softmax(-1)
    print("任务：先合并 2 和 3，再与 4 运算；最终结果 20 得 1 分。")
    print(f"训练前成功概率={history[0][1]:.3f}；训练后={history[-1][1]:.3f}")
    print(f"第一步选择加法的概率={probabilities[0, 0]:.3f}")
    print(f"得到 5 后选择乘法的概率={probabilities[1, 0]:.3f}")
    print("学得的状态价值：" + ", ".join(f"{value:.3f}" for value in values))
    exact = [expected_reward(logits), probabilities[1, 0], torch.tensor(0.0)]
    print("当前策略的精确状态价值：" + ", ".join(f"{value:.3f}" for value in exact))
    print("奖励 [0, 0, 1]、gamma=0.9 的回报：" +
          ", ".join(f"{value:.2f}" for value in discounted_returns([0, 0, 1], .9)))
    print("每 20 次更新的成功概率：" + ", ".join(f"{value:.3f}" for _, value in history))


if __name__ == "__main__":
    main()

7. 回到语言模型

生成回答时，可以把提示和已生成的前缀视为状态，词元视为动作，整个回答视为一条轨迹。奖励可能来自人的偏好、一个预测偏好的模型，或者程序化的答案检查。

这种对应关系让策略梯度能够用于语言模型，但困难也随之增大：词表可能有很多候选，一段回答可能很长，奖励只覆盖部分质量要求。某条回答得了高分，不表示其中每句话都正确；某次低分也不表示所有中间步骤都没有价值。

下一章会把“这两份回答，我更喜欢哪份”变成训练信号，并区分奖励模型、参考策略和当前策略各自做什么。

8. 停下来试一试

在两步都均匀随机的策略下，为什么开始状态价值是 0.25，得到 5 后却是 0.5？
得到 6 后的价值为什么是 0？这与数字 6 本身有没有永久关系？
奖励 $[0,1,2]$ 、 $\gamma=0.5$ 时，三个回报分别是多少？
为什么策略损失前面有负号？
减去一个只依赖状态的基线，为什么可以不改变策略梯度的期望？
为什么计算 actor 损失时要对价值基线停止梯度？
一次失败里包含一个正确步骤，为什么不能直接把所有动作都标为“错误示范”？
把终局奖励改成“结果越大越好”。策略会学到什么，是否仍解决原题？

参考答案与提示

开始时需要连续两次选对；到达 5 后只剩最后一次选择。
当前允许的后续动作只能得到 24 或 10，均不满足目标。换一个任务或后续动作集合，价值就可能不同。
从最后往前算，分别为 $1$ 、 $2$ 、 $2$ ；第一项是 $0+0.5\times1+0.25\times2=1$ 。
常用优化器最小化损失，而我们要提高正优势动作的对数概率。
对固定状态，基线可提出动作求和；概率之和的梯度为零。它不提供偏向某个动作的额外期望项。
基线在该梯度估计中被当作已给定的数，价值本身用单独的回报预测损失学习。
终局反馈没有逐步标注错误，多次探索与价值估计才逐渐区分有用步骤。
可能选择先乘得 6、再乘得 24。奖励增长了，却偏离了计算 $(2+3)\times4$ 的目标。

9. 继续阅读

Williams，Simple Statistical Gradient-Following Algorithms for Connectionist Reinforcement Learning：理解采样回报与对数概率梯度的联系。
Schulman 等，Proximal Policy Optimization Algorithms：下一章要用到的策略更新方法。

Keyboard shortcuts

从机器学习到大模型：原理、实践与前沿