> ## Documentation Index
> Fetch the complete documentation index at: https://agent-rl.tacetai.com/llms.txt
> Use this file to discover all available pages before exploring further.
# 6. 策略梯度
[第 2 章](/tutorial/zh/chapters/02-rl-fundamentals#2-5-最优策略，以及为什么我们不去计算它)中，我们已知环境的精确模型，因此可以直接对目标 $J(\theta)$ 求导。第 5 章面对的情形更接近实际训练：TinyAgent 与环境交互，最后拿到隐藏测试给出的分数；测试过程和离散采样都不可微。怎样仅凭采样轨迹和这些分数更新模型？本章推导得分函数（score function）估计量，也就是 REINFORCE，并把它用于 token 序列。
推导很短，但每一行都有工程上的后果。对轨迹概率求导时，环境转移项不依赖模型参数，因此无需知道环境的可微模型。工具 token 也不会成为策略动作对应的得分项，这解释了第 5 章为何需要动作掩码（action mask）。反向传播使用的是*代理损失*（surrogate loss）：它的梯度对应我们要估计的策略梯度，但损失数值本身不代表策略表现。基线（baseline）使估计量保持无偏，并且可以降低其方差。逐 token 损失怎样归约成标量，决定了不同长度的轨迹（trajectory）各占多大权重。在 gym 上尝试这些之前，我们还需要一个至少偶尔能成功的策略（policy）。我们构建这样一个策略，即本书的**基座智能体**（base agent），然后对 TinyAgent 进行第一次强化学习训练。
本章保存的代码通过 `arl.` 调用前面的章节，因此这些导入属于一个保存块。
```python theme={null}
#@save
import copy
import torch
import agentrl as arl
```
```python theme={null}
import torch.nn.functional as F
```
## 6.1 得分函数估计量
目标是从策略采样的轨迹的期望回报（return）：
$$
J(\theta) = \mathbb{E}_{\tau \sim p_\theta}\big[R(\tau)\big] = \sum_\tau p_\theta(\tau)\, R(\tau).
$$
回报 $R(\tau)$ 由环境计算，这里由测试计算，它对 $\theta$ 没有导数。对 $\theta$ 的依赖完全位于分布 $p_\theta$ 中。在求和号内求导，并利用 $\nabla p = p\, \nabla \log p$（[附录 A.5](/tutorial/zh/chapters/A-math#a-5-对数导数技巧)）：
$$
\nabla_\theta J(\theta) = \sum_\tau \nabla_\theta p_\theta(\tau)\, R(\tau) = \mathbb{E}_{\tau \sim p_\theta}\big[R(\tau)\, \nabla_\theta \log p_\theta(\tau)\big].
$$
右侧的期望可以从样本中估计。剩下的是整条轨迹的得分 $\nabla_\theta \log p_\theta(\tau)$。[2.2 节](/tutorial/zh/chapters/02-rl-fundamentals#2-2-策略、轨迹与回报)把轨迹概率分解为智能体能控制的部分和不能控制的部分：
$$
\log p_\theta(\tau) = \log \rho_0(s_0) + \sum_{t=0}^{T-1} \log \pi_\theta(a_t \mid s_t) + \sum_{t=0}^{T-1} \log P(s_{t+1} \mid s_t, a_t).
$$
初始分布 $\rho_0$ 和转移 $P$ 都不依赖于 $\theta$，因此它们的梯度为零，于是
$$
\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim p_\theta}\Big[R(\tau) \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t \mid s_t)\Big].
$$
这就是式 (6.1)，即针对轨迹的 REINFORCE 估计量（Williams, 1992）。推导中用到了三个假设，每一个都是真实系统可能打破它的地方。回报不能直接依赖于 $\theta$，如果学习得到的奖励模型（reward model）与策略共享权重，这一点就不成立。轨迹必须从 $\pi_\theta$ 本身采样，一旦在更新后复用数据，或者推理引擎计算出的概率略有不同，这一点就不成立（[第 8 章](/tutorial/zh/chapters/08-ppo)、[第 13 章](/tutorial/zh/chapters/13-rollout-systems)）。对轨迹的求和必须有限或性质良好，轮数预算保证了这一点。
估计量只需要用 $P$ 来*采样*。它从不计算 $P$ 的值，也从不对它求导。我们在第 2 章的修复 MDP 上检验这一点，那里可以得到精确梯度。下面的采样器并行运行许多回合，并为每个回合记录每个状态-动作对被访问的次数以及总回报。
```python theme={null}
def sample_repair(P, R, policy, n, generator, max_steps=50):
"""Run n repair-MDP episodes in parallel -> visit counts [n, S, A] and returns [n]."""
S, A = R.shape
state = torch.zeros(n, dtype=torch.long) # every episode starts 'unread'
visits = torch.zeros(n, S, A, dtype=P.dtype)
returns = torch.zeros(n, dtype=P.dtype)
alive = torch.ones(n, dtype=torch.bool)
for _ in range(max_steps):
action = torch.multinomial(policy[state], 1, generator=generator).squeeze(-1) # [n]
visits[torch.arange(n), state, action] += alive.to(P.dtype)
returns += alive * R[state, action]
probs = P[state, action].clone() # [n, S]
probs[~alive, -1] = 1.0 # finished episodes stay in 'done'
state = torch.multinomial(probs, 1, generator=generator).squeeze(-1)
alive &= state != S - 1
return visits, returns
```
策略是对形状为 `[S, A]` 的 logits 表 $\theta$ 做 softmax，随机初始化，使它既不均匀也不确定。第 2 章的 `policy_evaluation` 精确给出 $J(\theta)$，自动微分给出精确梯度。蒙特卡洛一侧则把每个回合的 $\sum_t \log \pi_\theta(a_t \mid s_t)$ 计算为 `visits * log_softmax(theta)`，乘以回报，取平均，再求导。
```python theme={null}
P, R = arl.make_repair_mdp()
torch.manual_seed(0)
theta = (0.5 * torch.randn(7, 4, dtype=torch.float64)).requires_grad_() # policy logits [S, A]
J = arl.policy_evaluation(P, R, torch.softmax(theta, -1), gamma=1.0)[0]
exact = torch.autograd.grad(J, theta)[0] # [S, A]
n = 20_000
gen = torch.Generator().manual_seed(0)
visits, returns = sample_repair(P, R, torch.softmax(theta, -1).detach(), n, gen)
log_pi = torch.log_softmax(theta, -1) # [S, A]
surrogate = (returns[:, None, None] * visits * log_pi).sum() / n # mean of R(tau) * log p(tau)
estimate = torch.autograd.grad(surrogate, theta)[0]
cos = F.cosine_similarity(estimate.flatten(), exact.flatten(), dim=0)
print(f'J exact {J.item():.4f} | mean return {returns.mean().item():.4f}')
print(f'relative error {((estimate - exact).norm() / exact.norm()).item():.3f} | cosine {cos.item():.4f}')
```
Output:
```text theme={null}
J exact 0.1323 | mean return 0.1264
relative error 0.068 | cosine 0.9988
```
该策略在约 13% 的回合中通过（$J = 0.132$）。估计值与精确梯度的余弦相似度为 0.999，6.8% 的相对误差是 20,000 个回合的蒙特卡洛噪声。看看代码对什么求了导：`returns`、`visits` 和 `log_pi`。只有最后一个依赖于 $\theta$，而 $P$ 根本没有出现在被求导的表达式中。编码环境的处境与 $P$ 相同：我们可以运行它，但不能对它求导。
### 代理损失
我们求导的表达式 $\frac{1}{N}\sum_n R(\tau_n)\, \log p_\theta(\tau_n)$（回报保持固定）是一个**代理目标**（surrogate objective）。它在当前参数处的梯度等于式 (6.1) 的估计量。它的*值*不是 $J$，也不估计任何有用的量。回报作为常数进入，因此把每个回报平移同一个量，会任意改变代理目标的值，而它的期望梯度仍为 $\nabla J$，因为 $\mathbb{E}[\nabla \log p_\theta(\tau)] = 0$。
```python theme={null}
for shift in [0.0, -J.item(), 5.0]:
surrogate = ((returns + shift)[:, None, None] * visits * torch.log_softmax(theta, -1)).sum() / n
grad = torch.autograd.grad(surrogate, theta)[0]
cos = F.cosine_similarity(grad.flatten(), exact.flatten(), dim=0).item()
print(f'return shift {shift:+.3f}: surrogate value {surrogate.item():9.4f} | cosine to exact {cos:.4f}')
```
Output:
```text theme={null}
return shift +0.000: surrogate value -0.8342 | cosine to exact 0.9988
return shift -0.132: surrogate value -0.4056 | cosine to exact 0.9992
return shift +5.000: surrogate value -17.0281 | cosine to exact 0.7833
```
三个代理目标的值分别为 $-0.83$、$-0.41$ 和 $-17.0$，而它们的梯度估计的都是同一个向量。这些值都与 $J = 0.132$ 无关。平移 $-J$ 使梯度估计略微*更接近*精确梯度（余弦 0.9992 对 0.9988），而平移 $+5$ 使它偏离得多（0.78）。平移回报就是一个基线，6.4 节会解释为什么一种平移有帮助而另一种有害。由此立即得到一条实践规则：永远不要把策略梯度的损失曲线当作进展的度量。应该记录奖励、通过率和梯度范数。
## 6.2 可供学习的基座策略
式 (6.1) 用每条轨迹的回报为其得分加权。如果每条采样轨迹的回报都相同，估计量就是该常数乘以 $\sum_n \nabla\log p_\theta(\tau_n)$，这个量的期望为零，而加上基线后它恰好为零。第 5 章测量了未训练的 TinyAgent：384 次调用中 384 次格式错误，每个回合奖励都是 0。强化学习无法从那里起步。
实际的大模型 RL 通常也不会从这种随机策略起步。预训练模型已经能写出流畅的文本，经过指令微调的模型大多数时候已经能输出格式正确的工具调用。强化学习随后改进的是它做出*哪些*调用。我们用**行为克隆**（behavior cloning）给 TinyAgent 同样的起点：在示范数据上做监督训练，交叉熵损失只作用于策略 token。示范来自第 4 章的 `random` 和 `blind` 脚本化策略在训练任务上的运行。它们格式正确但毫无技巧：`random` 在九个合法调用中均匀选择，`blind` 不读代码就猜一个修复操作。模仿它们的模型学会了调用格式和回合的节奏，但学不会如何选择正确的操作。我们把得到的结果称为**基座智能体**。它是本书中预训练基座模型的对应物，后面每一章都从它开始做强化学习。[第 15 章](/tutorial/zh/chapters/15-sft-and-rejection-sampling)会系统研究监督式热启动（warm start），包括它对探索（exploration）的影响。
行为克隆损失是在策略 token 上的平均，我们经常用到它，值得保存为一个函数。
```python theme={null}
#@save
def masked_mean(x, mask, dim=None):
"""Mean of x where mask is 1: x, mask [B, T] -> scalar (dim=None) or reduced over dim."""
if dim is None:
return (x * mask).sum() / mask.sum().clamp(min=1e-8)
return (x * mask).sum(dim) / mask.sum(dim).clamp(min=1e-8)
```
`make_base_agent` 从每个脚本化策略为每个训练任务收集 `episodes_per_task` 条示范，用 `scripted_rollout`（5.6 节）编码，然后在随机小批次（minibatch）上执行 `steps` 次 Adam 更新。所有内容都设定了随机种子，因此相同的调用返回相同的权重。
```python theme={null}
#@save
def make_base_agent(tok, seed=0, kinds=('random', 'blind'), episodes_per_task=8, steps=200,
batch_size=32, lr=3e-3):
"""TinyAgent behavior-cloned on well-formed but unskilled scripted play on the train tasks."""
torch.manual_seed(seed)
gen = torch.Generator().manual_seed(seed)
demos = []
for k, kind in enumerate(kinds):
policy = arl.scripted_policy(kind, seed=seed + k)
demos += [arl.scripted_rollout(policy, arl.BugFixEnv(), task, tok)
for task in arl.make_tasks('train') for _ in range(episodes_per_task)]
data = arl.collate(demos, tok.pad_id) # ids [N, T], action_mask [N, T-1]
model = arl.TinyAgent(tok.vocab_size)
opt = torch.optim.Adam(model.parameters(), lr=lr)
for _ in range(steps):
idx = torch.randint(len(demos), (batch_size,), generator=gen)
logp = arl.sequence_log_probs(model, data['ids'][idx]) # [b, T-1]
loss = -masked_mean(logp, data['action_mask'][idx]) # cross-entropy on policy tokens
opt.zero_grad()
loss.backward()
opt.step()
return model
```
为了衡量一个策略，我们需要它的 rollout 的汇总统计。`rollout_stats` 报告平均奖励、通过率（pass rate，已提交且所有隐藏测试都通过）、格式正确且合法的调用比例、提交率，以及每个回合的平均轮数和策略 token 数。后面的章节记录同样的数字。
```python theme={null}
#@save
def rollout_stats(trajs):
"""Summary of a list of Trajectory -> dict of reward, pass_rate, valid_calls, submit_rate, turns, action_tokens."""
n, turns = len(trajs), sum(t.info['n_turns'] for t in trajs)
return dict(reward=sum(t.reward for t in trajs) / n,
pass_rate=sum(bool(t.info['submitted'] and t.info['hidden_pass']) for t in trajs) / n,
valid_calls=1 - sum(t.info['n_invalid'] for t in trajs) / turns,
submit_rate=sum(bool(t.info['submitted']) for t in trajs) / n,
turns=turns / n,
action_tokens=sum(float(t.action_mask.sum()) for t in trajs) / n)
```
```python theme={null}
def sample(model, tasks, n_per_task, seed):
"""n_per_task episodes of each task, sampled at temperature 1 in lockstep."""
batch_tasks = [t for t in tasks for _ in range(n_per_task)]
envs = [arl.BugFixEnv() for _ in batch_tasks]
return arl.rollout_batch(model, envs, batch_tasks, tok, generator=torch.Generator().manual_seed(seed))
def show(name, stats):
print(f'{name:<22}' + ' '.join(f'{k} {v:.3f}' for k, v in stats.items()))
def short(action):
"""' edit fix_init ' -> 'fix_init'; ' read ' -> 'read'."""
call = arl.parse_call(action)
return (call[1] or call[0]) if call else 'malformed'
```
我们完全按照第 5 章的方式构建分词器，然后构建基座智能体，并在两个划分上与未训练的 TinyAgent 比较，每个任务 16 个回合。
```python theme={null}
tok = arl.build_tokenizer(arl.make_tasks('all'))
train_tasks, heldout_tasks = arl.make_tasks('train'), arl.make_tasks('heldout')
base = make_base_agent(tok)
torch.manual_seed(0)
show('untrained, train', rollout_stats(sample(arl.TinyAgent(tok.vocab_size), train_tasks, 16, seed=0)))
base_train = rollout_stats(sample(base, train_tasks, 16, seed=0))
show('base agent, train', base_train)
show('base agent, held-out', rollout_stats(sample(base, heldout_tasks, 16, seed=0)))
```
Output:
```text theme={null}
untrained, train reward 0.000 pass_rate 0.000 valid_calls 0.000 submit_rate 0.000 turns 6.000 action_tokens 35.681
base agent, train reward 0.119 pass_rate 0.119 valid_calls 0.979 submit_rate 0.781 turns 3.134 action_tokens 11.312
base agent, held-out reward 0.138 pass_rate 0.138 valid_calls 0.988 submit_rate 0.825 turns 3.013 action_tokens 10.706
```
未训练模型从不通过，它的调用没有一个合法，并且总是耗尽六轮预算。基座智能体在约 98% 的轮次中调用合法，约五个回合中有四个会提交，在训练回合中通过率为 12%，在留出（held-out）回合中为 14%。这正是它的示范数据所决定的位置：介于 `random`（4.6 节中训练集上为 0.10）和 `blind`（0.21）之间，远低于 oracle 的 1.00。它在留出任务上的数字与训练任务上的数字在采样噪声范围内一致，因为它所模仿的东西都不依赖于正在修复的是哪个函数。本章其他样本中它的通过率在 0.12 到 0.19 之间；对几百个回合而言，这个范围就是采样误差。撰写本章时，在检查主机上 `make_base_agent` 约用时五秒。构建过程不检查计时。
一个采样回合展示了强化学习要面对的行为。
```python theme={null}
for traj in sample(base, train_tasks[:1], 3, seed=1):
print(f'reward {traj.reward:.0f} | stop {traj.info["stop"]:<9} | ' + ' -> '.join(map(short, traj.info['actions'])))
```
Output:
```text theme={null}
reward 0 | stop submit | submit
reward 0 | stop submit | add_zero_guard -> submit
reward 0 | stop budget | add_zero_guard -> test -> swap_minmax -> fix_range_end -> test -> fix_range_end
```
这些回合的调用格式都合法，但决策还缺少依据。在 `clamp` 上，一个回合提交了未修改的文件，一个猜了 `add_zero_guard` 后提交，还有一个在编辑和测试之间徘徊直到预算耗尽。基座智能体确实读取文件时（只在少数回合中如此），它接下来的操作与读到的内容无关。成功主要来自碰巧猜中修复操作。然而，每一次成功都是一条回报与其他轨迹不同的轨迹，因此策略梯度不再为零。
## 6.3 token 上的估计量
在 [2.8 节](/tutorial/zh/chapters/02-rl-fundamentals#2-8-把-llm-智能体建模为-mdp-的三种方式)的 token 层面，一条轨迹是一个 token 序列 $y_1, \dots, y_T$，其前缀为 $c_i = y_{ 0 # [B, T-1] positions with a logit gradient
print('positions with a direct (logit) gradient:', int(direct.sum()), '| policy tokens of successful episodes:',
int((batch['action_mask'] * batch['rewards'][:, None]).sum()))
print('direct gradient outside the action mask:', int((direct & (batch['action_mask'] == 0)).sum()))
base.zero_grad()
```
Output:
```text theme={null}
positions with a direct (logit) gradient: 70 | policy tokens of successful episodes: 70
direct gradient outside the action mask: 0
```
恰好 70 个位置的 logits 收到梯度，它们恰好就是成功回合的策略 token。失败回合的 $R = 0$，完全没有贡献，任何提示、模板或工具位置都没有收到直接梯度。间接梯度在输入一侧可以看到：工具位置上的嵌入梯度不为零。
```python theme={null}
owner = arl.pad_rows([t.owner for t in trajs], batch['ids'].shape[1], -1) # [B, T]
emb_grad = embedded['h'].grad.norm(dim=-1) # [B, T]
for name, code in [('prompt', 1), ('policy', 2), ('tool', 3)]:
rows = owner == code
print(f'{name:<7} positions: {int(rows.sum()):5d} | mean |d loss / d embedding| = {emb_grad[rows].mean():.2e}')
```
Output:
```text theme={null}
prompt positions: 480 | mean |d loss / d embedding| = 1.87e-04
policy positions: 954 | mean |d loss / d embedding| = 1.49e-04
tool positions: 1149 | mean |d loss / d embedding| = 3.18e-05
```
提示和工具位置通过它们提供给之后策略 token 的表示获得梯度。在工具位置上，梯度大小约为策略位置上的五分之一，但不为零。策略正是这样学会*使用*观测的，例如在读取文件后选择操作，而从不被训练去*产生*观测。掩码去掉了直接项，注意力保留了间接项。
## 6.4 基线
6.1 节的平移实验提示我们：从回报中减去某个量，可以在不引入偏差的情况下降低估计量的噪声。一般的表述如下。令 $b(s_t)$ 为时刻 $t$ 状态的任意函数，或者是任何在 $a_t$ 被采样之前就已确定的量的函数。那么
$$
\mathbb{E}_{a_t \sim \pi_\theta(\cdot \mid s_t)}\big[b(s_t)\, \nabla_\theta \log \pi_\theta(a_t \mid s_t)\big] = b(s_t) \sum_a \pi_\theta(a \mid s_t)\, \frac{\nabla_\theta \pi_\theta(a \mid s_t)}{\pi_\theta(a \mid s_t)} = b(s_t)\, \nabla_\theta \sum_a \pi_\theta(a \mid s_t) = b(s_t)\, \nabla_\theta 1 = 0.
$$
对 $s_t$ 取外层期望（塔性质，tower property），对每个 $t$ 都得到 $\mathbb{E}_\tau[b(s_t) \nabla \log \pi_\theta(a_t \mid s_t)] = 0$。因此，对任何这样的 $b$，
$$
\nabla_\theta J(\theta) = \mathbb{E}_\tau\Big[\sum_t \big(R(\tau) - b(s_t)\big)\, \nabla_\theta \log \pi_\theta(a_t \mid s_t)\Big]
$$
这就是式 (6.3)。证明只用到 $b$ 在抽取 $a_t$ 时已经确定这一点。基线可以依赖于任务、提示和更早的 token。它不能依赖于 $a_t$、任何在 $a_t$ 之后采样的东西，也不能依赖于该轨迹自身的回报。用*包含*该轨迹自身奖励的批次均值计算的基线会轻微违反最后一条，[第 7 章](/tutorial/zh/chapters/07-advantages-and-credit)会量化其后果。
基线不能改变均值，因此它只能改变方差。对常数基线 $b$ 和逐轨迹得分 $g = \nabla \log p_\theta(\tau)$，总方差 $\mathbb{E}\lVert (R - b) g\rVert^2 - \lVert \nabla J \rVert^2$ 是 $b$ 的二次函数，在下式处取最小值：
$$
b^* = \frac{\mathbb{E}\big[R\, \lVert g \rVert^2\big]}{\mathbb{E}\big[\lVert g \rVert^2\big]}.
$$
这是一个加权平均奖励，权重正比于得分范数的平方。在实践中，普通的平均奖励与它很接近，而 $V(s_0)$，即从任务初始状态出发的期望回报，是自然的依赖于状态的选择（[附录 A.3](/tutorial/zh/chapters/A-math#a-3-控制变量) 把基线作为控制变量（control variate）来讨论）。
我们先在精确梯度已知的地方确认无偏性。我们采样 400 个独立批次，每批 64 个修复 MDP 回合，对四种基线分别在各批次上平均估计量：无基线、精确的 $J$、批次自身的平均回报，以及一个故意选得很差的常数 1。
```python theme={null}
batches, size = 400, 64
visits, returns = sample_repair(P, R, torch.softmax(theta, -1).detach(), batches * size,
torch.Generator().manual_seed(1))
score = (visits - visits.sum(-1, keepdim=True) * torch.softmax(theta, -1).detach()) # d log p(tau) / d theta [N, S, A]
score, returns = score.view(batches, size, 7, 4), returns.view(batches, size)
for name, b in [('none', 0.0), ('exact J', J.item()), ('batch mean', returns.mean(1, keepdim=True)),
('constant 1', 1.0)]:
per_batch = ((returns - b)[..., None, None] * score).mean(1) # [batches, S, A]
error = ((per_batch.mean(0) - exact).norm() / exact.norm()).item()
noise = ((per_batch.var(0).sum() / batches).sqrt() / exact.norm()).item() # Monte Carlo scale
print(f'{name:<11} relative error of the mean {error:.3f} (noise level {noise:.3f})',
f'| per-batch variance {per_batch.var(0).sum():.4f}')
```
Output:
```text theme={null}
none relative error of the mean 0.033 (noise level 0.034) | per-batch variance 0.0071
exact J relative error of the mean 0.028 (noise level 0.030) | per-batch variance 0.0057
batch mean relative error of the mean 0.038 (noise level 0.030) | per-batch variance 0.0056
constant 1 relative error of the mean 0.080 (noise level 0.056) | per-batch variance 0.0191
```
对 softmax 表而言，得分有闭式形式：每次访问 $(s, a)$ 都给第 $s$ 行加上 $\mathbf{e}_a - \pi_\theta(\cdot \mid s)$，前几行代码计算的正是这个。噪声水平是 400 个批次平均值的标准误差相对于梯度大小的比值。对每种基线，平均值与精确梯度的差距都在约一个噪声水平之内（常数 1 为 1.4 个），这正是无偏估计量应有的表现。方差则各不相同：减去精确的 $J$ 使每批方差降低 20%（从 0.0071 降到 0.0057），减去 1 则使其变为 2.7 倍。批次均值在原理上是无偏性的一个例外：它包含每条轨迹自身的回报，这使期望梯度缩小为原来的 $1 - 1/64$ 倍，1.6% 的误差远低于这里的噪声水平。
现在在真实系统上做同样的测量。对 240 个基座智能体回合（每个训练任务 12 个），我们在 TinyAgent 的全部参数上计算每个回合的得分 $\nabla_\theta \sum_i m_i \log \pi_\theta(y_i \mid c_i)$，每个回合一次反向传播。
```python theme={null}
def episode_gradient(model, traj, weights):
"""Flattened gradient of sum_i weights_i * log pi(y_i | c_i) for one trajectory -> [n_params]."""
logp = arl.sequence_log_probs(model, traj.ids[None])[0] # [T-1]
grads = torch.autograd.grad((weights * logp).sum(), list(model.parameters()))
return torch.cat([g.flatten() for g in grads])
def total_variance(per_episode):
"""Trace of the covariance of per-episode gradient estimates [N, n_params]."""
return per_episode.var(0).sum().item()
probe = sample(base, train_tasks, 12, seed=3)
scores = torch.stack([episode_gradient(base, t, t.action_mask) for t in probe]) # [N, n_params]
rewards = torch.tensor([t.reward for t in probe]) # [N]
print('scores:', tuple(scores.shape), '| mean reward:', round(rewards.mean().item(), 3))
```
Output:
```text theme={null}
scores: (240, 157244) | mean reward: 0.188
```
这些得分构成一个 240 乘 157,244 的矩阵，每行一个回合，而这个探测样本的平均奖励恰好是 0.188。把每一行乘以 $R_n - b$，就得到该回合对估计量的贡献，其协方差的迹就是基线应当最小化的量。我们评估几种基线。逐任务基线使用在一个**单独**批次的回合上测得的每个任务的平均奖励，因此它在探测回合被抽取之前就已确定，符合证明的要求。
```python theme={null}
squared_norms = scores.pow(2).sum(-1)
b_star = ((rewards * squared_norms).sum() / squared_norms.sum()).item()
other = sample(base, train_tasks, 12, seed=4)
task_mean = {t.name: sum(o.reward for o in other if o.info['task'] == t.name) / 12 for t in train_tasks}
per_task = torch.tensor([task_mean[t.info['task']] for t in probe])
for name, b in [('none', 0.0), ('mean reward', rewards.mean().item()), ('b* (optimal constant)', b_star),
('per-task mean', per_task), ('constant 0.5', 0.5), ('constant 1.0', 1.0)]:
shown = f'{b:.3f}' if isinstance(b, float) else 'varies'
print(f'{name:<22} b = {shown:<7} total variance {total_variance((rewards - b)[:, None] * scores):8.1f}')
```
Output:
```text theme={null}
none b = 0.000 total variance 44.7
mean reward b = 0.188 total variance 40.9
b* (optimal constant) b = 0.123 total variance 39.3
per-task mean b = varies total variance 40.8
constant 0.5 b = 0.500 total variance 91.9
constant 1.0 b = 1.000 total variance 322.9
```
图 6.1 画出了完整的曲线。最优常数 $b^* = 0.123$ 给出 39.3 的总方差，平均奖励给出 40.9。由于平均奖励很小，完全不用基线也只差 14%（44.7）。抛物线在右侧急剧上升：基线为 1 时方差扩大七倍以上，因为失败占大多数；基线为 1 时，每条失败轨迹都会以权重 $-1$ 贡献完整的得分向量。逐任务基线（40.8）并不比全局均值更好，尽管按任务估计基线有机会利用难度差异。每个任务的均值只由 12 个独立回合估计，因此该估计的噪声被加到每个优势（advantage）上，而在低于 0.2 的通过率下，这抵消了跟踪每个任务难度所带来的收益。[第 7 章](/tutorial/zh/chapters/07-advantages-and-credit)会用组基线重新讨论这一权衡。
## 6.5 因果性与 reward-to-go
基线论证还有第二个推论。假设奖励在回合进行中到达，$R(\tau) = \sum_{t'} r_{t'}$。对 $t' < t$，奖励 $r_{t'}$ 在 $a_t$ 被采样之前就已确定，因此它可以作为 $a_t$ 的基线，且 $\mathbb{E}[r_{t'} \nabla \log \pi_\theta(a_t \mid s_t)] = 0$。去掉这些项得到
$$
\nabla_\theta J(\theta) = \mathbb{E}_\tau\Big[\sum_t \nabla_\theta \log \pi_\theta(a_t \mid s_t) \sum_{t' \ge t} r_{t'}\Big],
$$
即式 (6.4)。内层求和就是 **reward-to-go**（剩余回报）$G_t$。因果性是说：一个动作的学习权重只需包含从它发生时起的奖励。在 token 层面，$t$ 遍历 token 位置，$G_i = \sum_{j \ge i} r_j$。
对于放在最后一个策略 token 上的纯结果奖励（outcome reward），每个策略 token 上都有 $G_i = R(\tau)$，因果性不改变任何东西。当奖励包含中间项时它才重要：每轮的成本、对格式错误调用的格式惩罚，或者每轮的塑形奖励（[第 11 章](/tutorial/zh/chapters/11-rewards-and-verifiers)、[第 12 章](/tutorial/zh/chapters/12-multi-turn-and-long-horizons)）。我们给探测回合加上每次调用 0.1 的成本。成本放在每次调用的最后一个 token 上，隐藏测试奖励放在最后一次调用的最后一个 token 上。
```python theme={null}
def step_cost_rewards(traj, cost=0.1):
"""Per-position rewards [T-1]: -cost at the last token of every call, plus the outcome at the final one."""
m = traj.action_mask
call_end = m * (1 - torch.cat([m[1:], torch.zeros(1)])) # 1 at the last token of each call
rewards = -cost * call_end
rewards[call_end.nonzero().max()] += traj.reward
return rewards
traj = next(t for t in probe if t.reward > 0 and t.info['n_turns'] >= 3)
r = step_cost_rewards(traj) # [T-1]
to_go = r.flip(-1).cumsum(-1).flip(-1) * traj.action_mask # G_i = sum_{j >= i} r_j
m = traj.action_mask
starts = (m * (1 - torch.cat([torch.zeros(1), m[:-1]]))).nonzero().flatten() # first token of each call
print(f'return of the episode: {r.sum().item():+.2f} (the weight of every token without causality)')
for action, j in zip(traj.info['actions'], starts.tolist()):
print(f' {short(action):<15} reward-to-go at its first token {to_go[j].item():+.2f}')
```
Output:
```text theme={null}
return of the episode: +0.60 (the weight of every token without causality)
add_zero_guard reward-to-go at its first token +0.60
swap_minmax reward-to-go at its first token +0.70
fix_range_end reward-to-go at its first token +0.80
submit reward-to-go at its first token +0.90
```
使用总回报时，回合中的每个 token 获得相同的权重。使用 reward-to-go 时，较早的调用要为它自己以及之后每一次调用的成本负责，而最后一次调用只为它自己的成本负责。我们在 240 个探测回合上测量两种估计量的方差。总回报估计量复用已有的得分；reward-to-go 每个回合需要多一次反向传播，因为它的权重逐 token 不同。
```python theme={null}
total = torch.stack([step_cost_rewards(t).sum() for t in probe]) # [N]
var_total = total_variance(total[:, None] * scores)
to_go_grads = torch.stack([episode_gradient(base, t, step_cost_rewards(t).flip(-1).cumsum(-1).flip(-1)
* t.action_mask) for t in probe])
var_to_go = total_variance(to_go_grads)
print(f'total return: variance {var_total:.1f} | reward-to-go: variance {var_to_go:.1f}',
f'| reduction {1 - var_to_go / var_total:.0%}')
```
Output:
```text theme={null}
total return: variance 104.9 | reward-to-go: variance 60.2 | reduction 43%
```
从每个 token 的权重中去掉过去的成本降低了方差，降低的幅度随中间奖励的数量和大小而增大。[第 7 章](/tutorial/zh/chapters/07-advantages-and-credit)保存了一个通用的、带折扣并跳过工具 token 的 `returns_to_go`，并用优势的估计替换剩下的 reward-to-go。
## 6.6 损失聚合
式 (6.2) 是对 token 求和、对轨迹取平均。实现中很少恰好这样计算。它们在一个填充后的 `[B, T]` 张量上计算逐 token 项 $\ell_{n,i} = -\hat A_{n,i} \log \pi_\theta(y_{n,i} \mid c_{n,i})$，再把它归约为一个标量，而所选的归约方式给每个 token 设定了一个权重 $w_{n,i}$：$\nabla \mathcal{L} = \sum_{n,i} w_{n,i}\, \nabla \ell_{n,i}$。令 $L_n = \sum_i m_{n,i}$ 为轨迹 $n$ 的策略 token 数。常用的四种归约是
| 模式 | 损失 | 一个策略 token 的权重 $w_{n,i}$ | 轨迹 $n$ 的总权重 |
| - | - | - | - |
| `token_mean` | $\sum_{n,i} m\ell / \sum_{n} L_n$ | $1 / \sum_k L_k$ | $L_n / \sum_k L_k$ |
| `seq_mean_token_sum` | $\frac{1}{B}\sum_n \sum_i m\ell$ | $1 / B$ | $L_n / B$ |
| `seq_mean_token_mean` | $\frac{1}{B}\sum_n \frac{1}{L_n}\sum_i m\ell$ | $1 / (B L_n)$ | $1 / B$ |
| `constant` | $\sum_{n,i} m\ell / C$ | $1 / C$ | $L_n / C$ |
只有 `seq_mean_token_sum` 是推导所得的式 (6.2)。其他模式改变归一化方式；其中有的只是整体缩放，有的还会改变不同长度轨迹之间的相对权重。`token_mean` 除以批次中的策略 token 数，因此当批次中的轨迹变长时，有效学习率（learning rate）会缩小。`seq_mean_token_mean` 给每条轨迹相同的总权重而不论其长度，因此长轨迹中单个 token 的权重比短轨迹中更小。`constant` 除以一个固定的数 $C$，例如批大小乘以策略 token 的最大数目，这使逐 token 权重在不同批次间保持固定。哪些公开方案使用哪种模式，以及 Dr. GRPO 对按序列归一化的批评（Liu et al., 2025），是[第 9 章](/tutorial/zh/chapters/09-grpo-and-variants)的主题。这里我们构建这些归约并测量它们的作用。
```python theme={null}
#@save
def aggregate_loss(per_token, mask, mode='token_mean', norm_const=None):
"""Reduce per-token terms [B, T] under mask [B, T] to a scalar; mode sets each token's weight."""
if mode == 'token_mean': # each policy token in the batch weighs 1 / sum(mask)
return masked_mean(per_token, mask)
if mode == 'seq_mean_token_sum': # sum within a trajectory, mean over trajectories
return (per_token * mask).sum(-1).mean()
if mode == 'seq_mean_token_mean': # mean within a trajectory, then mean over trajectories
return masked_mean(per_token, mask, dim=-1).mean()
if mode == 'constant': # sum over the batch divided by a fixed constant
if norm_const is None:
raise ValueError("mode='constant' needs norm_const, e.g. batch_size * max_policy_tokens")
return (per_token * mask).sum() / norm_const
raise ValueError(f'unknown aggregation mode {mode!r}')
def reinforce_loss(logp, advantages, mask, mode='token_mean', norm_const=None):
"""Surrogate -A * log pi on policy tokens: logp, mask [B, T]; advantages [B] or [B, T], detached."""
if advantages.dim() == 1:
advantages = advantages[:, None] # one advantage per trajectory -> [B, 1], broadcast over T
return aggregate_loss(-advantages.detach() * logp, mask, mode, norm_const)
```
`reinforce_loss` 对优势做了 detach。在推导中它们是常数，6.8 节展示了不这样做时出现的偏差。表中的权重就是损失相对于每个逐 token 项的梯度，因此可以用自动微分直接读出。取三条轨迹，分别有 4、8 和 24 个策略 token：
```python theme={null}
lengths = torch.tensor([4, 8, 24])
mask = (torch.arange(24)[None] < lengths[:, None]).float() # [3, 24]
per_token = torch.zeros(3, 24, requires_grad=True)
for mode in ['token_mean', 'seq_mean_token_sum', 'seq_mean_token_mean', 'constant']:
loss = aggregate_loss(per_token, mask, mode, norm_const=3 * 24)
weight = torch.autograd.grad(loss, per_token)[0] # dL / d(per-token term)
print(f'{mode:<20} per-token weight', [f'{w:.4f}' for w in weight[:, 0].tolist()],
'| trajectory total', [f'{w:.3f}' for w in weight.sum(-1).tolist()])
```
Output:
```text theme={null}
token_mean per-token weight ['0.0278', '0.0278', '0.0278'] | trajectory total ['0.111', '0.222', '0.667']
seq_mean_token_sum per-token weight ['0.3333', '0.3333', '0.3333'] | trajectory total ['1.333', '2.667', '8.000']
seq_mean_token_mean per-token weight ['0.0833', '0.0417', '0.0139'] | trajectory total ['0.333', '0.333', '0.333']
constant per-token weight ['0.0139', '0.0139', '0.0139'] | trajectory total ['0.056', '0.111', '0.333']
```
输出就是填入了数字的那张表。在 `token_mean`、`seq_mean_token_sum` 和 `constant` 下，三条轨迹的逐 token 权重相同，因此 24 个 token 的轨迹的总权重是 4 个 token 轨迹的六倍。这三种模式只差一个整体缩放，且只有 `token_mean` 的缩放依赖于批次（这里是 $1/36$）。在 `seq_mean_token_mean` 下，每条轨迹的总权重为 $1/3$，短轨迹中一个 token 的权重是长轨迹中一个 token 的六倍。
在智能体任务中，长度不是随机的。在基座智能体的探测批次中，成功回合比失败回合短，因为侥幸猜中会以较早的提交结束，而失败则会四处徘徊。使用 $R - \bar R$ 作为优势时，成功被向上推，失败被向下推，而聚合方式决定推的力度。
```python theme={null}
probe_batch = arl.collate(probe, tok.pad_id)
mask, adv = probe_batch['action_mask'], probe_batch['rewards'] - probe_batch['rewards'].mean()
length = mask.sum(-1)
success = probe_batch['rewards'] > 0
long = ~success & (length > length[~success].median()) # the longer half of failures
short_fail = ~success & ~long
print(f'policy tokens: successes {length[success].mean():.1f} | short failures {length[short_fail].mean():.1f}',
f'| long failures {length[long].mean():.1f}')
logp = torch.zeros_like(mask, requires_grad=True)
for mode in ['token_mean', 'seq_mean_token_sum', 'seq_mean_token_mean']:
weight = torch.autograd.grad(reinforce_loss(logp, adv, mask, mode), logp)[0] # = -A * w per token [B, T-1]
push_up, push_down = -weight[success].sum().item(), weight[~success].sum().item()
per_token = weight.sum(-1) / length # push on each token of a trajectory
print(f'{mode:<20} push up {push_up:.3f} | push down {push_down:.3f} | per-token push, long / short failures',
f'{per_token[long].mean() / per_token[short_fail].mean():.2f}')
```
Output:
```text theme={null}
policy tokens: successes 9.4 | short failures 6.4 | long failures 19.6
token_mean push up 0.132 | push down 0.157 | per-token push, long / short failures 1.00
seq_mean_token_sum push up 1.435 | push down 1.712 | per-token push, long / short failures 1.00
seq_mean_token_mean push up 0.152 | push down 0.152 | per-token push, long / short failures 0.33
```
优势是中心化的，$\sum_n \hat A_n = 0$，但在两种与长度成正比的模式下，推力并不平衡。它们的总推力正比于 $\sum_n \hat A_n L_n$，而由于失败平均比成功长，向下的推力比向上的推力多 19%（0.157 对 0.132）。每个 token 被推的力度相同，因此一个长失败（平均 19.6 个 token）每个 token 受到的推力与一个短失败（6.4 个 token）相同。`seq_mean_token_mean` 精确地平衡了总量，因为每条轨迹的总权重为 $1/B$，而优势之和为零。然而按 token 计，每个长失败被向下推的力度只有每个短失败的三分之一。如果策略能让一个失败回合变长，例如在放弃之前多加几次调用，`seq_mean_token_mean` 就会降低这些调用中每一次的逐 token 惩罚。这就是 Dr. GRPO 归因于按序列归一化的长度增长背后的机制（[第 9 章](/tutorial/zh/chapters/09-grpo-and-variants)）。
图 6.2 把两种效应分开。左图是轨迹*自身*的长度：三种与长度成正比的模式落在一条过原点的直线上，而 `seq_mean_token_mean` 是平的，因此批次中最短和最长轨迹的总权重相同。右图是*批次中其余部分*的长度：只有 `token_mean` 会在其他轨迹变长时改变一条轨迹的逐 token 权重，当其他轨迹的平均长度从 10 增长到 40 时，该权重降到四分之一。因此，随着训练中回合变长，`token_mean` 会降低每个 token 的有效学习率，而其他三种模式不会。本书默认使用 `token_mean`，因为大多数公开方案使用它，而且依赖于批次的缩放容易监控。[附录 B 问题 1](/tutorial/zh/chapters/B-design-review#问题-1：为变长工具轨迹选择损失聚合方式) 会详细论证这一选择。
## 6.7 用 REINFORCE 训练 TinyAgent
第一次强化学习运行所需的一切都已就绪。一次更新从当前策略采样一批回合，计算优势，并在代理损失上走一步梯度。`policy_gradient_step` 在一个整理好的批次上执行更新。它裁剪（clip）全局梯度范数，这是一项标准的保护措施，我们同时记录其值；它返回损失以及裁剪之前的梯度范数。第 8 章和第 9 章会替换损失，但保持这个形状。
```python theme={null}
#@save
def policy_gradient_step(model, optimizer, batch, advantages, mode='token_mean', norm_const=None,
max_grad_norm=1.0):
"""One REINFORCE update on a collated batch; advantages [B] or [B, T-1] -> dict(loss, grad_norm)."""
logp = arl.sequence_log_probs(model, batch['ids']) # [B, T-1], with grad
loss = reinforce_loss(logp, advantages, batch['action_mask'], mode, norm_const)
optimizer.zero_grad()
loss.backward()
grad_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
optimizer.step()
return dict(loss=loss.item(), grad_norm=float(grad_norm))
```
训练循环每一步从随机的训练任务中采样 64 个回合。优势是奖励减去基线 $b$，$b$ 是过去各批次均值的指数移动平均，初始化为 6.2 节测得的基座智能体训练通过率。基线在见到批次之前计算，因此满足式 (6.3) 的条件。每十步我们用每个留出任务 8 个回合测量留出通过率。模型是基座智能体的一个副本，基座智能体本身保持不变，留作后续比较。
```python theme={null}
policy = copy.deepcopy(base)
opt = torch.optim.Adam(policy.parameters(), lr=2e-3)
gen = torch.Generator().manual_seed(0)
baseline, history, evals = base_train['pass_rate'], [], []
for step in range(80):
if step % 10 == 0:
evals.append((step, rollout_stats(sample(policy, heldout_tasks, 8, seed=100 + step))['pass_rate']))
tasks = [train_tasks[i] for i in torch.randint(len(train_tasks), (64,), generator=gen).tolist()]
trajs = arl.rollout_batch(policy, [arl.BugFixEnv() for _ in tasks], tasks, tok, generator=gen)
batch = arl.collate(trajs, tok.pad_id)
stats = policy_gradient_step(policy, opt, batch, batch['rewards'] - baseline) # advantage [B]
baseline = 0.9 * baseline + 0.1 * batch['rewards'].mean().item()
history.append(dict(step=step, reward=batch['rewards'].mean().item(), **stats))
if step % 10 == 0:
print(f'step {step:2d} | batch reward {history[-1]["reward"]:.3f} | held-out pass {evals[-1][1]:.3f}',
f'| loss {stats["loss"]:+.4f} | grad norm {stats["grad_norm"]:.2f}')
evals.append((80, rollout_stats(sample(policy, heldout_tasks, 8, seed=180))['pass_rate']))
print(f'step 80 | held-out pass {evals[-1][1]:.3f}')
```
Output:
```text theme={null}
step 0 | batch reward 0.156 | held-out pass 0.263 | loss +0.0067 | grad norm 0.10
step 10 | batch reward 0.328 | held-out pass 0.200 | loss +0.0325 | grad norm 0.10
step 20 | batch reward 0.469 | held-out pass 0.375 | loss +0.0194 | grad norm 0.08
step 30 | batch reward 0.531 | held-out pass 0.375 | loss -0.0815 | grad norm 0.13
step 40 | batch reward 0.719 | held-out pass 0.525 | loss -0.0045 | grad norm 0.07
step 50 | batch reward 0.781 | held-out pass 0.537 | loss +0.0173 | grad norm 0.07
step 60 | batch reward 0.828 | held-out pass 0.475 | loss +0.0092 | grad norm 0.09
step 70 | batch reward 0.844 | held-out pass 0.537 | loss -0.0123 | grad norm 0.06
step 80 | held-out pass 0.450
```
批次奖励在第 70 步时从 0.16 攀升到 0.84，图 6.3 展示了完整曲线。损失一列印证了 6.1 节：它正负皆有，始终在零附近百分之几的范围内，对进展毫无说明，而奖励一列则能说明。梯度范数保持在 0.06 到 0.13 之间，远低于裁剪阈值 1，因此本次运行中裁剪从未生效。留出通过率每个点只用 80 个回合测量，因此在相邻两次评估之间会变动好几个百分点。这次运行在检查主机上约用时七秒。在从留出一列中解读任何东西之前，我们用每个任务 16 个回合测量两个划分，并看看策略实际在做什么。
```python theme={null}
for name, tasks in [('train', train_tasks), ('held-out', heldout_tasks)]:
show(f'base, {name}', rollout_stats(sample(base, tasks, 16, seed=7)))
show(f'REINFORCE, {name}', rollout_stats(sample(policy, tasks, 16, seed=7)))
```
Output:
```text theme={null}
base, train reward 0.144 pass_rate 0.144 valid_calls 0.973 submit_rate 0.750 turns 3.369 action_tokens 12.109
REINFORCE, train reward 0.775 pass_rate 0.775 valid_calls 0.951 submit_rate 0.975 turns 3.416 action_tokens 12.934
base, held-out reward 0.125 pass_rate 0.125 valid_calls 0.981 submit_rate 0.800 turns 2.931 action_tokens 10.506
REINFORCE, held-out reward 0.456 pass_rate 0.456 valid_calls 0.974 submit_rate 0.944 turns 4.481 action_tokens 17.144
```
```python theme={null}
from collections import Counter
for name, tasks in [('train', train_tasks), ('held-out', heldout_tasks)]:
trajs = sample(policy, tasks, 16, seed=8)
reads = sum(' read ' in t.info['actions'] for t in trajs) / len(trajs)
edits = sum(len(t.info['ops']) for t in trajs) / len(trajs)
fix_op = {t.name: t.fix_op for t in tasks}
first_right = sum(bool(t.info['ops']) and t.info['ops'][0] == fix_op[t.info['task']] for t in trajs) / len(trajs)
print(f'{name}: read {reads:.2f} | edits per episode {edits:.2f} | first edit is the right op {first_right:.2f}')
for calls, count in Counter(' -> '.join(map(short, t.info['actions'])) for t in trajs).most_common(3):
print(f' {count:3d} x {calls}')
```
Output:
```text theme={null}
train: read 0.00 | edits per episode 2.18 | first edit is the right op 0.50
49 x swap_minmax -> submit
45 x fix_comparison -> add_zero_guard -> fix_comparison -> submit
44 x fix_init -> fix_comparison -> submit
held-out: read 0.00 | edits per episode 3.48 | first edit is the right op 0.23
16 x swap_minmax -> swap_minmax -> swap_minmax -> swap_minmax -> swap_minmax -> submit
13 x swap_minmax -> swap_minmax -> fix_comparison -> swap_minmax -> swap_minmax -> submit
13 x fix_comparison -> submit
```
训练通过率从 0.14 上升到 0.78，留出通过率从 0.13 上升到 0.46。行为统计解释了这两个数字。策略从不读取，也从不在编辑前测试，因此它选择第一次编辑时掌握的唯一任务相关信息，就是问题描述中的函数名。它的第一次编辑在 50% 的训练回合和 23% 的留出回合中是正确的操作，而在五个修复操作中均匀猜测的正确率为 20%。在训练任务上，它部分记住了哪个操作对应哪个函数名。在留出任务上，函数名是它在训练中从未见过的单词，它的第一选择只是猜测。
其余的收益来自第二种策略：在提交之前应用多个操作，训练任务上每回合 2.2 个，留出任务上 3.5 个。错误的操作常常使源代码保持不变（`no change`）或做出无害的改动，因此像 `fix_comparison -> add_zero_guard -> fix_comparison` 这样的序列同时修复了 `compare` 和 `zero` 两种缺陷类型。最常见的留出回合将 `swap_minmax` 应用了五次，按奇偶性，`min` 和 `max` 最终仍处于交换后的状态。在这个奖励下，叠加编辑是一种合法的修复，但没有哪个工程师会接受这种策略，而它之所以能迁移到留出任务，只是因为编辑菜单很小。读取，这种本可以干净迁移的行为，从未出现。基座智能体只在少数回合中读取，而且读完后选择的操作与读到的内容无关，因此在它自己的样本中，读取与奖励不相关，而 REINFORCE 只能强化它采样到的东西。
> **证据。** 本书测得，基于一个 158k 参数模型、20 个训练任务、单个随机种子。这次运行演示的是估计量的机制，并不是关于大模型会学到哪些策略的论断。定性的教训是估计量的一般性质，而非这个模型的性质：强化学习会放大在策略自身样本中已经得到奖励的行为，而记住训练任务或利用一个很小的动作菜单，都是低成本的得分方式。[第 4 章](/tutorial/zh/chapters/04-environments-and-tools#4-6-脚本化基线)为“学到的不只是格式”设定了 0.2 的留出门槛，而这次运行表明，越过这条门槛本身并不能证明智能体学会了读取。第 [7](/tutorial/zh/chapters/07-advantages-and-credit)、[9](/tutorial/zh/chapters/09-grpo-and-variants) 和 [12](/tutorial/zh/chapters/12-multi-turn-and-long-horizons) 章会回到这次运行的两种失败模式：为无用步骤分配学习信号，以及缺乏信息收集。
## 6.8 实践中的设计问题
> **问题：稀疏二值奖励下的高方差。**
>
> **现象：** 每步奖励曲线噪声很大，不同随机种子的学习速度差异很大，简单任务饱和的同时困难任务上的进展停滞。
>
> **原因：** 在二值奖励和低通过率下，大多数回合以一个小的负权重贡献得分，少数回合贡献大的正权重。估计量的方差相对于其均值很大（图 6.1），而一个所有采样回合都失败的任务没有任何贡献。
>
> **可选方案：** (a) 每步更多回合，rollout 计算量线性增长，换来 $1/\sqrt{B}$ 的噪声降低。(b) 更好的基线，尤其是由组内样本得到的逐任务基线（[第 7 章](/tutorial/zh/chapters/07-advantages-and-credit)）。(c) 更稠密的奖励，但有被钻空子的风险（[第 11 章](/tutorial/zh/chapters/11-rewards-and-verifiers)）。(d) 让任务保持在接近 50% 通过率的课程（[第 14 章](/tutorial/zh/chapters/14-scaling-and-efficiency)）。
>
> **建议做法：** 在调整任何其他东西之前，先直接测量梯度噪声。下面的代码块把 240 个探测回合分成互不相交的批次，计算同样大小的独立批次所得梯度估计之间的余弦相似度。
```python theme={null}
centered = (rewards - rewards.mean())[:, None] * scores # [N, n_params]
gen = torch.Generator().manual_seed(0)
for size in [8, 32, 120]:
cos = []
for _ in range(20): # 20 random disjoint pairs
order = torch.randperm(len(centered), generator=gen)
first, second = centered[order[:size]].mean(0), centered[order[size:2 * size]].mean(0)
cos.append(F.cosine_similarity(first, second, dim=0).item())
print(f'batches of {size:3d} episodes: mean cosine between independent estimates {sum(cos) / len(cos):.3f}')
```
Output:
```text theme={null}
batches of 8 episodes: mean cosine between independent estimates -0.092
batches of 32 episodes: mean cosine between independent estimates 0.165
batches of 120 episodes: mean cosine between independent estimates 0.418
```
来自两个独立的 8 回合批次的梯度估计互不相关（平均余弦 $-0.09$ 在 20 对样本的噪声范围内为零），32 回合的批次弱相关（0.17），即使是 120 回合的批次也只部分一致（0.42）。像 6.7 节那样的 64 回合批次，在任何单独一步中大部分都是噪声。学习之所以仍然有效，是因为噪声在多步之间被平均掉，而信号不会。这也是学习率与批大小相互权衡的原因。
> **问题：聚合引入的长度偏差。**
>
> **现象：** 平均回合长度在训练中增长，而奖励没有相应上升，增长往往集中在失败回合中。
>
> **原因：** `seq_mean_token_mean` 把每条轨迹的推力除以其自身长度（6.6 节），因此更长的失败在每个 token 上受到的惩罚更少。`token_mean` 把步长与批次的总长度耦合在一起。
>
> **可选方案：** (a) `seq_mean_token_sum` 或 `constant`，它们保持逐 token 权重固定，但需要针对更大的梯度重新调整学习率。(b) 广泛使用的 `token_mean`，同时记录批次的策略 token 数。(c) 显式的长度成本，这会改变任务。
>
> **建议做法：** `token_mean`，同时分别记录成功和失败回合的平均长度。如果失败回合变长，就切换到 `constant`。
> **问题：detach 了错误的张量。**
>
> **现象：** 训练能运行，损失也在下降，但学习更慢或朝奇怪的方向漂移，梯度与参考实现不一致。
>
> **原因：** 某个在代理损失中应当是常数的量，例如由模型自身输出计算的基线或优势，被留在了计算图上。于是它的梯度被加到估计量上，而推导中没有任何依据支持这一项。
>
> **可选方案：** (a) 在损失函数内部对优势做 detach，如 `reinforce_loss` 所做。(b) 在 `torch.no_grad()` 下计算优势。
>
> **建议做法：** 两者都做，再加一个在固定批次上把梯度与参考实现比较的单元测试。下面的代码块在一个五动作老虎机（bandit）上复现这种偏差，基线直接取自策略自身的概率。
```python theme={null}
logits = torch.tensor([0.5, -0.2, 0.1, 0.0, 0.3], requires_grad=True) # a 5-action policy
reward_of = torch.tensor([1.0, 0.0, 0.0, 1.0, 0.0]) # pass / fail per action
probs = torch.softmax(logits, -1)
exact = torch.autograd.grad((probs * reward_of).sum(), logits, retain_graph=True)[0]
actions = torch.multinomial(probs.detach(), 200_000, replacement=True, generator=torch.Generator().manual_seed(0))
logp = torch.log_softmax(logits, -1)[actions] # [N]
baseline = probs[0] # a function of theta
for name, b in [('detached', baseline.detach()), ('attached', baseline)]:
loss = -((reward_of[actions] - b) * logp).mean()
estimate = -torch.autograd.grad(loss, logits, retain_graph=True)[0]
print(f'{name:<9} relative error to exact gradient {((estimate - exact).norm() / exact.norm()).item():.3f}')
```
Output:
```text theme={null}
detached relative error to exact gradient 0.001
attached relative error to exact gradient 1.538
```
基线被 detach 时，估计值在采样噪声范围内与精确梯度一致。不 detach 时，梯度会多出一项 $-\nabla b \cdot \overline{\log \pi}$，其期望为 $H(\pi)\, \nabla b$，误差与梯度本身同一量级。当价值头与策略共享主干、且其预测未经 `detach()` 就进入策略损失时，会出现同样的 bug（[第 7 章](/tutorial/zh/chapters/07-advantages-and-credit)）。
> **问题：优势全为零的批次。**
>
> **现象：** 在某些步中每个回合的奖励都相同，记录的损失和梯度恰好为零，但策略仍在变化，有时还会退化。
>
> **原因：** 所有奖励相等且使用均值基线时，每个优势都为零，因此梯度为零。Adam 仍会按照之前各步的一阶矩估计移动参数。权重衰减也会移动参数。优势为零的一步并不是空操作。在 `token_mean` 下，优势为零的轨迹还会往归一化分母中加入 token，从而稀释其他每条轨迹的权重。
>
> **可选方案：** (a) 当批次没有信号时跳过优化器步骤。(b) 在组成批次之前过滤掉优势为零的组并重新采样，即动态采样（[第 9 章](/tutorial/zh/chapters/09-grpo-and-variants)）。(c) 接受这种漂移。
>
> **建议做法：** 使用组时选 (b)，否则选 (a)。下面的代码块先走一个普通的步，再在全零优势上走一步。
```python theme={null}
model = copy.deepcopy(base)
opt = torch.optim.Adam(model.parameters(), lr=2e-3)
batch = probe_batch
policy_gradient_step(model, opt, batch, batch['rewards'] - batch['rewards'].mean())
before = torch.cat([p.detach().flatten() for p in model.parameters()])
stats = policy_gradient_step(model, opt, batch, torch.zeros(len(batch['rewards'])))
after = torch.cat([p.detach().flatten() for p in model.parameters()])
print(f'zero-advantage step: loss {stats["loss"]:.1f} | grad norm {stats["grad_norm"]:.1f}',
f'| parameter change {(after - before).norm():.4f}')
```
Output:
```text theme={null}
zero-advantage step: loss 0.0 | grad norm 0.0 | parameter change 0.4903
```
损失和梯度范数恰好为零，而参数仍然移动了一个非零的量，即上一步的动量。
## 小结
* REINFORCE 估计量 $\mathbb{E}[R(\tau) \sum_t \nabla \log \pi_\theta(a_t \mid s_t)]$ 只需要用环境来采样。环境动态从得分中消失；在 token 层面，所有提示、模板和工具 token 也同样消失，这就是动作掩码。工具 token 仍通过注意力影响梯度。
* 代理损失的梯度就是估计量，而它的值没有意义。应记录奖励和通过率，而不是损失。
* 任何在动作被采样之前就已确定的基线都使估计量保持无偏。在基座智能体上，平均奖励接近最优常数，基线为 1 会使方差扩大七倍以上，而由 12 个回合估计的逐任务基线并不比全局均值更好。reward-to-go 从每个 token 的权重中去掉过去的奖励，在每次调用有成本的设定下使方差降低了 43%。
* 损失聚合设定每个 token 的权重。三种常用模式按长度成比例地为轨迹加权，`seq_mean_token_mean` 为它们等权加权，而 `token_mean` 把步长与批次的长度耦合在一起。
* 基座智能体，即在 `random` 和 `blind` 的运行数据上做行为克隆的 TinyAgent，98% 的时候写出合法调用，回合通过率在 12% 到 19% 之间。八十步 REINFORCE 把它在训练任务上的通过率从 0.14 提高到 0.78，在留出任务上从 0.13 提高到 0.46，靠的是部分记住函数名和叠加编辑，而从未靠读取。
## 练习
1. 用来自 `policy_evaluation` 的依赖于状态的基线 $b(s) = V^\pi(s)$，在每一步应用，重复 6.4 节的无偏性检验。与精确常数 $J$ 相比，每批方差下降了多少？然后改用 $b(s_t, a_t) = Q^\pi(s_t, a_t)$，并说明此时均值是错的。证明中的哪一行不再成立？
2. 证明最优常数基线为 $b^* = \mathbb{E}[R \lVert g\rVert^2] / \mathbb{E}[\lVert g\rVert^2]$，并解释它为什么不同于平均奖励。在 TinyAgent 中，哪些类型的回合的 $\lVert g \rVert^2$ 较大？
3. 改变一个假设：分别用 `seq_mean_token_mean` 和 `seq_mean_token_sum` 做 REINFORCE 训练（为后者重新调整学习率，并说明为什么需要重新调整）。记录失败回合中策略 token 的平均数目。失败回合的长度是否如 6.6 节预测的那样变化？
4. 只用 `blind` 的运行数据构建基座智能体，再只用 `random` 的运行数据构建一次。测量通过率、读取率，然后测量 80 步 REINFORCE 后的结果。哪种初始分布带来更多的留出提升？为什么？
5. 让读取对 REINFORCE 变得有用：用一个 50% 的时候会出错的 oracle（它先读取，再以 0.5 的概率选择正确的操作）的示范来训练基座智能体。REINFORCE 现在会学会读取吗？测量留出通过率。
6. token 级估计量把每个 token 视为一个动作。推导一个把整次调用视为一个动作的轮级估计量。证明它的得分是该调用中各 token 得分之和，因此两个估计量完全相同。这两种建模方式在哪里不再等价？
7. 设计评审：一位同事提议用代理损失的滑动平均绝对值对其做归一化来稳定训练，“让损失保持在 1 左右”。这对梯度有什么影响？在什么条件下它等价于一个学习率调度？什么时候它会改变更新的方向？
This documentation is built and hosted on [Mintlify](https://mintlify.com), a developer documentation platform.