""" ============================================================================= 第 6 讲 · 代码 5：Policy gradient —— REINFORCE、reward-to-go 与 baseline ============================================================================= 对应讲义：lectures/lecture06_rl_control.md 第 5 节 ----------------------------------------------------------------------------- 【这个文件要回答的问题 What this file answers】 ----------------------------------------------------------------------------- 前面的方法都先学 V（或者模型），再从 V 推出策略。Policy gradient 直接对策略参数 θ 做梯度上升： ∇η(θ) = E[ Σ_t ∇log π_θ(a_t | s_t) · (回报) ] （REINFORCE / likelihood ratio） 演示 1：得分函数（score function）∇log π 就是第 1 讲 GLM 的梯度：softmax 策略 → x(e_a − π)， 高斯策略 → (a − μ)x/σ²。它的期望为 0 —— 这是 baseline 不引入偏差的原因。 演示 2：同一个 θ 上，四种估计量都无偏，但方差差很多：全回报 vs reward-to-go vs 减去 baseline。 给每步奖励加一个常数 c，不改变最优策略，却让不减 baseline 的估计量方差暴涨。 演示 3：学习曲线：方差直接变成学习速度和稳定性的差别（20 次独立运行，用精确的 η(θ) 评估）。 问题：一维走廊，5 个状态，从 0 出发，状态 4 每步奖励 +1，动作 左 / 右（10% 的概率走反），固定时域 H = 10。 策略：表格型 softmax，θ 形状 (5, 2)。状态少，所以 η(θ) 和 ∇η(θ) 都能用动态规划精确算出来（rl_tools.exact_return_and_grad）， 可以直接检查采样估计是否无偏、方差多大。 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python3 05_policy_gradient.py # 约 4 秒 需要 numpy。依赖本目录的 rl_tools.py（用到第 1 讲的 softmax 和 one_hot）。 """ import numpy as np from rl_tools import (exact_return_and_grad, make_corridor, print_header, reinforce_grad, sample_trajectories, softmax, softmax_policy, softmax_score) H = 10 # 固定时域：每条轨迹 10 步，没有终止状态 N_STATES = 5 def make_task(c=0.0): """ 走廊任务。c 是加到**每一步**奖励上的常数：状态 4 的奖励 1 + c，其他状态 c。 【思路】时域固定、没有终止状态时，每条轨迹都正好拿 H 次奖励，所以加常数 c 让每条轨迹的回报都多 c·H —— 对所有策略都一样，最优策略和 ∇η 完全不变（演示 2 验证）。但 REINFORCE 的估计量会被它影响。 【边界】有终止状态时（回合长度依赖策略），加常数**会**改变最优策略（代码 01 演示 3 的网格世界就是例子）。 """ mdp = make_corridor(N_STATES, {}, start=0, step_reward=c, slip=0.1) mdp.R[N_STATES - 1] = 1.0 + c return mdp def rewards_after(rewards): """ Reward-to-go：G_t = Σ_{j > t} R(s_j)，即动作 a_t **之后**拿到的奖励（a_t 影响不了它之前和当时的奖励）。 【参数】rewards：(M, H) 【返回】(M, H) 【numpy 小知识】[:, ::-1] 把每行倒过来；倒过来做 cumsum 再倒回去 = "从 t 到结尾的和"（suffix sum）。 再减去 rewards 本身，就是"严格在 t 之后"的和。 """ suffix = np.cumsum(rewards[:, ::-1], axis=1)[:, ::-1] return suffix - rewards def finite_horizon_optimum(mdp): """最优策略的 η*：有限时域价值迭代 V_t = R + max_a P V_(t+1)，从 V_H = 0 往回算 H 步。""" V = np.zeros(mdp.S) for _ in range(H): V = mdp.R + (mdp.P @ V).max(axis=1) return float(V[mdp.start]) # ============================================================================= # 演示 1：得分函数 = GLM 的梯度，期望为 0 # ============================================================================= def demo_score_function(): print_header("演示 1：得分函数 ∇log π —— 就是第 1 讲 GLM 的梯度 (target − prediction) × input") theta = np.array([[0.5, -0.5]]) # 一个状态、两个动作 pi = softmax_policy(theta, 0) print(f"(a) 表格型 softmax 策略，某个状态的 logits = [0.5, −0.5] → π = {np.round(pi, 4)}") print(f"\n{'action a':>9}{'pi(a)':>9}{'score = e_a − pi':>22}") expected = np.zeros(2) for a in range(2): sc = softmax_score(theta, 0, a)[0] expected += pi[a] * sc print(f"{a:>9}{pi[a]:>9.4f}{str(np.round(sc, 4)):>22}") print(f" E_a[score] = Σ_a π(a)·score(a) = {np.round(expected, 12) + 0.0}") rng = np.random.default_rng(51) mu_w, sigma = 0.8, 0.5 # 高斯策略 a ~ N(w·x, σ²) x = rng.uniform(-1, 1, size=200000) a = mu_w * x + sigma * rng.standard_normal(x.size) score = (a - mu_w * x) * x / sigma ** 2 # ∂/∂w log N(a; w x, σ²) print(f"\n(b) 高斯策略 a ~ N(w·x, σ²)，w = {mu_w}，σ = {sigma}：score = (a − w·x)·x / σ²") print(f" 200000 个样本上 score 的平均 = {score.mean():+.4f}（标准误 {score.std() / np.sqrt(x.size):.4f}）") print(""" 观察 Observations： · softmax 策略的得分函数是 x(s)·(e_a − π)ᵀ：和第 1 讲 softmax 回归单个样本的梯度 x·(y − p)ᵀ 一模一样， 只是"标签"y 换成了"自己刚做的动作"a。高斯策略的得分函数 (a − w·x)·x/σ² 是线性回归的梯度形状（误差 × 输入）。 所以 REINFORCE 就是"以自己的动作为标签、按回报加权的 GLM 回归"（rl_tools.reinforce_grad 就是这么写的）。 · 得分函数的期望为 0（(a) 精确为 0，(b) 在抽样误差内为 0）：E_a[∇log π(a)] = ∇Σ_a π(a) = ∇1 = 0。 所以从回报里减去任何**不依赖于 a_t** 的 baseline b(s_t)，梯度的期望不变 —— 演示 2 用数字验证。""") # ============================================================================= # 演示 2：四种估计量的偏差与方差 # ============================================================================= def exact_baselines(theta, mdp): """ 两种"理想的"baseline（用动态规划精确计算，只用于演示 2 的方差分析）： b_t = E[G_t]：只依赖时间 t 的常数（不同轨迹相同） b_t(s) = E[G_t | s_t = s] = V_t(s) − R(s)：依赖当前状态（"这个状态本来就值这么多"） 【返回】(b_time, b_state)：形状 (H,) 和 (H, S) """ pi = softmax(theta) P_pi = np.einsum("sa,sat->st", pi, mdp.P) d = np.zeros((H, mdp.S)) d[0, mdp.start] = 1.0 for t in range(1, H): d[t] = d[t - 1] @ P_pi V = np.zeros((H + 1, mdp.S)) for t in range(H - 1, -1, -1): V[t] = mdp.R + P_pi @ V[t + 1] b_state = V[:H] - mdp.R[None, :] # E[Σ_{j>t} R | s_t = s] b_time = (d * b_state).sum(axis=1) # 对 s_t 的分布求平均 return b_time, b_state def per_episode_grads(theta, mdp, rng, N): """ 采样 N 条轨迹，返回每条轨迹单独的梯度估计（四种估计量），形状 dict: name → (N, S·A)。 【思路】reinforce_grad 对一批轨迹求平均；这里想要"每条轨迹一个估计"来算方差， 所以对每种权重分别算 per-step 贡献再按轨迹求和。用 one-hot 写成矩阵运算，避免 Python 循环。 """ S, A = theta.shape st, ac, rw, acted = sample_trajectories(theta, mdp, H, N, rng) G0 = np.repeat(rw.sum(axis=1, keepdims=True), H, axis=1) Gt = rewards_after(rw) b_time, b_state = exact_baselines(theta, mdp) weights = {"full return G_0": G0, "reward-to-go G_t": Gt, "G_t - b_t": Gt - b_time[None, :], "G_t - V_t(s_t)": Gt - b_state[np.arange(H)[None, :], st]} Y = np.eye(A)[ac] # (N, H, A) one-hot 动作 P = softmax(theta[st.ravel()]).reshape(N, H, A) # (N, H, A) 每一步的 π(·|s_t) Xs = np.eye(S)[st] # (N, H, S) one-hot 状态 out = {} for name, w in weights.items(): contrib = (w * acted)[..., None] * (Y - P) # (N, H, A) out[name] = np.einsum("nhs,nha->nsa", Xs, contrib).reshape(N, S * A) # 每条轨迹：Σ_t x(s_t)(w·(e_a − π))ᵀ return out def demo_variance(): print_header("演示 2：同一个 θ 上，四种 REINFORCE 估计量 —— 都无偏，方差差很多") N = 20000 print(f"θ = 0（均匀随机策略）。每种设置采样 N = {N} 条轨迹，每条轨迹给一个梯度估计 ĝ。") print("bias = ‖mean(ĝ) − ∇η‖（∇η 由动态规划精确算出），noise = √(tr Cov / N) 是 N 个样本平均的预期抽样误差。") print("cos(10) = 每 10 条轨迹平均成一个估计，和真实梯度的夹角余弦的平均值（越接近 1 越好）。\n") print(f"{'c':>4} {'estimator':<18}{'bias':>8}{'noise':>8}{'tr Cov':>10}{'cos(10)':>9}") theta = np.zeros((N_STATES, 2)) res = {} for c in [0.0, 1.0]: mdp = make_task(c) eta, g = exact_return_and_grad(theta, mdp, H) grads = per_episode_grads(theta, mdp, np.random.default_rng(52), N) for name, G in grads.items(): bias = np.linalg.norm(G.mean(axis=0) - g.ravel()) trcov = G.var(axis=0).sum() batches = G.reshape(N // 10, 10, -1).mean(axis=1) cos = (batches @ g.ravel()) / (np.linalg.norm(batches, axis=1) * np.linalg.norm(g) + 1e-300) res[(c, name)] = (bias, trcov, cos.mean()) print(f"{c:>4.0f} {name:<18}{bias:>8.3f}{np.sqrt(trcov / N):>8.3f}{trcov:>10.2f}{cos.mean():>9.3f}") print(f" (η = {eta:.3f}, ‖∇η‖ = {np.linalg.norm(g):.3f})") f0, r0, t0, s0 = (res[(0.0, k)][1] for k in grads) f1, r1, t1, s1 = (res[(1.0, k)][1] for k in grads) print(f""" 观察 Observations： · 四种估计量都无偏：bias 和 noise（预期的抽样误差）是同一个量级，没有系统性的偏离。∇η 在 c = 0 和 c = 1 时完全一样 （加常数不改变最优策略，也不改变梯度）。 · 方差（tr Cov）却差很多。c = 0 时：全回报 {f0:.2f} → reward-to-go {r0:.2f}（动作影响不了它之前的奖励， 把那部分去掉就是去掉纯噪声）→ 减时间 baseline {t0:.2f} → 减状态 baseline {s0:.2f}。 · 给每步奖励加 c = 1：全回报的方差从 {f0:.1f} 涨到 {f1:.0f}（{f1 / f0:.0f} 倍），10 条轨迹的估计和真实梯度几乎不相关 （cos = {res[(1.0, 'full return G_0')][2]:.2f}）。原因：每个 ∇log π 都乘上了一个很大的正数 ≈ cH， "所有做过的动作都被大幅强化"，真正有用的信号（谁比谁好）淹没在里面。减去 baseline 后的两行和 c = 0 时**完全相同** —— baseline 把这个常数精确地减掉了。 · 所以 baseline 的作用是把"回报"变成"比平均好多少"（advantage）：只有比平均好的动作被强化，比平均差的被削弱。 Actor-critic 用一个学出来的 V(s) 当 baseline，就是"G_t − V_t(s_t)"这一行（CS234 第 7 讲）。""") # ============================================================================= # 演示 3：学习曲线 # ============================================================================= def run_reinforce(mdp, method, rng, iters=200, M=10, lr=0.1, checkpoints=(0, 10, 25, 50, 100, 200)): """ REINFORCE 训练：每轮采样 M 条轨迹，θ ← θ + lr · ĝ。在 checkpoints 上记录**精确的** η(θ)。 【参数】method："full"（全回报）、"rtg"（reward-to-go）、"rtg+loo"（reward-to-go 减去 leave-one-out baseline）、 "exact"（用精确梯度做梯度上升：没有噪声的参照） 【思路】leave-one-out baseline：第 i 条轨迹在时刻 t 的 baseline = 其他 M − 1 条轨迹在时刻 t 的 G_t 的平均。 它不依赖第 i 条轨迹自己的动作，所以不引入偏差；不需要学一个 V 函数，只要同一批里有多条轨迹。 （同样的想法在 LLM 的 RL 微调里叫 RLOO；GRPO 用组内平均当 baseline，见 CS234 第 16 讲。） 【返回】(curve, theta)：curve 是每个 checkpoint 的 η(θ) 的列表；theta 是训练结束时的参数 """ theta = np.zeros((mdp.S, mdp.A)) curve = [] for k in range(iters + 1): if k in checkpoints: curve.append(exact_return_and_grad(theta, mdp, H)[0]) if k == iters: break if method == "exact": g = exact_return_and_grad(theta, mdp, H)[1] else: st, ac, rw, acted = sample_trajectories(theta, mdp, H, M, rng) if method == "full": w = np.repeat(rw.sum(axis=1, keepdims=True), H, axis=1) else: w = rewards_after(rw) if method == "rtg+loo": w = w - (w.sum(axis=0, keepdims=True) - w) / (M - 1) g = reinforce_grad(theta, st, ac, acted, w) theta = theta + lr * g return curve, theta def demo_learning_curves(): print_header("演示 3：学习曲线 —— batch = 10 条轨迹，学习率 0.1，20 次独立运行") checkpoints = (0, 10, 25, 50, 100, 200) eta_star = finite_horizon_optimum(make_task(0.0)) print(f"表中是 η(θ_k) − cH（减去常数后不同 c 可以直接比较），20 次运行的平均；最后一列是 20 次里最差的那次。") print(f"最优策略（一直向右）的 η* = {eta_star:.3f}；初始的均匀随机策略是第 0 列。\n") header = f"{'c':>3} {'method':<20}" + "".join(f"{'k=' + str(k):>8}" for k in checkpoints) + f"{'worst@200':>11}" print(header) methods = [("exact gradient", "exact"), ("full return", "full"), ("reward-to-go", "rtg"), ("reward-to-go + LOO", "rtg+loo")] res, thetas = {}, {} for c in [0.0, 1.0]: mdp = make_task(c) for name, method in methods: runs = 1 if method == "exact" else 20 outs = [run_reinforce(mdp, method, np.random.default_rng((53, i)), checkpoints=checkpoints) for i in range(runs)] curves = np.array([o[0] for o in outs]) - c * H res[(c, method)] = curves thetas[(c, method)] = [o[1] for o in outs] print(f"{c:>3.0f} {name:<20}" + "".join(f"{v:>8.3f}" for v in curves.mean(axis=0)) + f"{curves[:, -1].min():>11.3f}") e0 = res[(0.0, "exact")][0] full1, rtg1, loo1 = (res[(1.0, m)] for m in ["full", "rtg", "rtg+loo"]) full0, loo0 = res[(0.0, "full")], res[(0.0, "rtg+loo")] same = np.array_equal(res[(0.0, "rtg+loo")], res[(1.0, "rtg+loo")]) worst = int(np.argmin(full1[:, -1])) th_w = thetas[(1.0, "full")][worst] g_w = np.linalg.norm(exact_return_and_grad(th_w, make_task(1.0), H)[1]) g_0 = np.linalg.norm(exact_return_and_grad(np.zeros_like(th_w), make_task(1.0), H)[1]) p_right = softmax(th_w)[:, 1] print(f"\nc = 1、full return 最差的那次运行（第 {worst} 次）训练结束时：π(右 | s) = {np.round(p_right, 3)}（s = 0 … 4），" f"‖∇η‖ = {g_w:.3f}（θ = 0 时是 {g_0:.3f}）") print(f""" 观察 Observations： · c = 0 时三种估计量的学习曲线几乎重合（200 轮后 {full0[:, -1].mean():.3f} vs {loo0[:, -1].mean():.3f}），都跟着精确梯度 （{e0[-1]:.3f}）走。演示 2 里方差只差约 2 倍，batch = 10、学习率 0.1 时这点噪声被很多步平均掉了 —— baseline 不是在任何问题上都有显著效果。 · c = 1 时差别就出来了：全回报平均只到 {full1[:, -1].mean():.3f}，最差的一次只有 {full1[:, -1].min():.3f} —— 高方差的更新 把 softmax 推进了一个很差的、几乎确定性的策略：在起点向右的概率只剩 {p_right[0]:.3f}（一直往墙上撞）， 梯度也缩小到 {g_w:.3f}，几乎出不来了（policy gradient 的过早收敛 / premature convergence）。 reward-to-go 平均 {rtg1[:, -1].mean():.3f}、最差 {rtg1[:, -1].min():.3f}；加 LOO baseline 平均 {loo1[:, -1].mean():.3f}、 最差 {loo1[:, -1].min():.3f}，""" + ("和 c = 0 时**逐位相同**（baseline 把常数精确地减掉了）。" if same else "和 c = 0 时基本一样。") + f""" · 奖励的"零点"在 RL 里本来是任意的（每步 +1 存活奖励、每步 −1 时间惩罚都很常见）， 一个好的估计量不应该被它左右 —— baseline 就是为此存在的。 · 精确梯度在 200 轮后也只到 {e0[-1]:.3f}（η* = {eta_star:.3f}）：softmax 越接近确定性，梯度越小（π(1 − π) 因子）， 普通梯度上升在"最后一段"很慢。natural policy gradient / TRPO / PPO 改的就是这个几何问题（CS234 第 7 讲）。""") def main(): demo_score_function() demo_variance() demo_learning_curves() if __name__ == "__main__": main()