""" ============================================================================= 第 6 讲 · 代码 4：策略梯度定理的数值验证，以及三种估计量的方差对比 ============================================================================= 对应讲义：lectures/lecture06_policy_gradient.md 第 5、6、7 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 在一个很小的 MDP 上，我们可以**精确地**算出策略梯度 ∇J(θ)。有了标准答案，就能检验： 1. 策略梯度定理给出的公式是否正确？ 2. REINFORCE 的几种变体是否都是**无偏**的（样本均值 ≈ 精确梯度）？ 3. 它们的**方差**差多少？ 三种估计量（讲义第 6、7 节）： (a) 整条轨迹回报： G_0 · Σ_t ∇log π(a_t|s_t) (b) 利用因果性（未来回报）：Σ_t γ^t G_t ∇log π(a_t|s_t) (c) 再减去基线 V(s_t)： Σ_t γ^t (G_t - V(s_t)) ∇log π(a_t|s_t) 外加一个"实践中常见但有偏"的版本： (d) 去掉 γ^t 因子： Σ_t (G_t - V(s_t)) ∇log π(a_t|s_t) ----------------------------------------------------------------------------- 【环境：回合制的火星车】 ----------------------------------------------------------------------------- [+1 终点] ← s2 s3 s4 s5 s6 → [+10 终点] ▲ 起点 - 5 个非终止状态（代码中 0~4），从中间出发 - 动作左/右：80% 成功移动，20% 原地不动 - 从最左边走出去得 +1、从最右边走出去得 +10，回合结束 - 另外每走一步都有 -1 的时间成本 （为什么要加它？如果只有最后一步有奖励，"动作之前的奖励"恒为 0， 估计量 (a) 和 (b) 会完全相同，就看不出"利用因果性"的效果了。） - γ = 0.9 - 策略：表格型 softmax，每个状态 2 个参数 θ[s, 左]、θ[s, 右]，共 10 个参数 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- python 04_gradient_variance.py （约 2 秒） """ import numpy as np np.set_printoptions(precision=4, suppress=True) N = 5 # 非终止状态数 START = 2 GAMMA = 0.9 STEP_COST = -1.0 LEFT, RIGHT = 0, 1 # ============================================================================= # 第 1 部分：环境模型（用于精确计算）与采样（用于估计） # ============================================================================= def build_model(): """ 返回 (P, Rsa)： P[s, a, s'] ：非终止状态之间的转移概率（走出边界的概率不在其中，所以每行和可能 < 1） Rsa[s, a] ：在 s 执行 a 的期望即时奖励 """ P = np.zeros((N, 2, N)) Rsa = np.zeros((N, 2)) for s in range(N): for a in (LEFT, RIGHT): target = s - 1 if a == LEFT else s + 1 P[s, a, s] += 0.2 # 打滑，原地不动 Rsa[s, a] = STEP_COST # 每一步的时间成本 if target < 0: Rsa[s, a] += 0.8 * 1.0 # 走出左边界：+1，回合结束 elif target >= N: Rsa[s, a] += 0.8 * 10.0 # 走出右边界：+10，回合结束 else: P[s, a, target] += 0.8 return P, Rsa def policy_probs(theta): """表格型 softmax：theta 形状 (N, 2)，返回形状 (N, 2) 的 π(a|s)。""" z = theta - theta.max(axis=1, keepdims=True) e = np.exp(z) return e / e.sum(axis=1, keepdims=True) def sample_episode(theta, rng, max_steps=1000): """按策略采样一个回合，返回 [(s, a, r), ...]。""" pi = policy_probs(theta) s, episode = START, [] for _ in range(max_steps): a = RIGHT if rng.random() < pi[s, RIGHT] else LEFT if rng.random() < 0.8: s_next = s - 1 if a == LEFT else s + 1 else: s_next = s if s_next < 0: episode.append((s, a, STEP_COST + 1.0)) return episode if s_next >= N: episode.append((s, a, STEP_COST + 10.0)) return episode episode.append((s, a, STEP_COST)) s = s_next return episode # ============================================================================= # 第 2 部分：精确的 J(θ)、V、Q 和梯度 # ============================================================================= def exact_values(theta, P, Rsa): """ 精确计算策略的 V^π（长度 N）和 Q^π（形状 (N, 2)）。第 2 讲的方法： P_π = Σ_a π(a|s) P(·|s,a)，R_π = Σ_a π(a|s) R(s,a)，V = (I - γ P_π)^{-1} R_π Q(s,a) = R(s,a) + γ Σ_s' P(s'|s,a) V(s') """ pi = policy_probs(theta) P_pi = np.einsum("sa,sat->st", pi, P) R_pi = np.sum(pi * Rsa, axis=1) V = np.linalg.solve(np.eye(N) - GAMMA * P_pi, R_pi) Q = Rsa + GAMMA * P @ V return V, Q, P_pi def J(theta, P, Rsa): """目标函数：从起点出发的期望折扣回报 J(θ) = V^π(起点)。""" return exact_values(theta, P, Rsa)[0][START] def gradient_by_finite_difference(theta, P, Rsa, eps=1e-6): """标准答案 1：对 J(θ) 做中心差分（不依赖任何理论，只要 J 算对了就对）。""" grad = np.zeros_like(theta) for idx in np.ndindex(theta.shape): # 【Python 小知识】np.ndindex(形状) 依次给出所有下标元组，如 (0,0), (0,1), (1,0), ... e = np.zeros_like(theta) e[idx] = eps grad[idx] = (J(theta + e, P, Rsa) - J(theta - e, P, Rsa)) / (2 * eps) return grad def gradient_by_pg_theorem(theta, P, Rsa): """ 标准答案 2：策略梯度定理（讲义 5 节）： ∇J(θ) = Σ_s η(s) Σ_a ∇π(a|s) Q^π(s, a) 其中 η(s) = Σ_t γ^t Pr(s_t = s) 是从起点出发的**折扣访问量**，满足 η = e_start + γ P_π^T η。 对表格型 softmax，∇π(a|s) 只对 θ[s, :] 非零：∂π(a|s)/∂θ[s, b] = π(a|s) (1[a=b] - π(b|s))。 所以 ∂J/∂θ[s, b] = η(s) Σ_a π(a|s)(1[a=b] - π(b|s)) Q(s,a) = η(s) π(b|s) (Q(s,b) - V(s))。 """ pi = policy_probs(theta) V, Q, P_pi = exact_values(theta, P, Rsa) e_start = np.zeros(N) e_start[START] = 1.0 eta = np.linalg.solve(np.eye(N) - GAMMA * P_pi.T, e_start) return eta[:, None] * pi * (Q - V[:, None]) # ============================================================================= # 第 3 部分：四种单回合梯度估计 # ============================================================================= def grad_log_pi(theta, s, a): """表格型 softmax 的得分函数：只有 θ[s, :] 这一行非零，等于 e_a - π(·|s)。""" g = np.zeros_like(theta) pi_s = policy_probs(theta)[s] g[s] = -pi_s g[s, a] += 1.0 return g def estimators(theta, episode, V): """ 对一个回合，计算四种梯度估计。返回字典 {名字: 与 theta 同形状的梯度估计}。 【参数】V：精确的状态价值，用作基线（实践中需要另外学习一个价值函数来近似它）。 """ T = len(episode) rewards = np.array([r for _, _, r in episode]) returns = np.zeros(T) # G_t（从 t 开始、以 t 为起点折扣） G = 0.0 for t in range(T - 1, -1, -1): G = rewards[t] + GAMMA * G returns[t] = G scores = [grad_log_pi(theta, s, a) for s, a, _ in episode] out = { "(a) G_0·Σ∇logπ": returns[0] * sum(scores), "(b) Σγ^t G_t ∇logπ": sum(GAMMA ** t * returns[t] * scores[t] for t in range(T)), "(c) Σγ^t (G_t-V)∇logπ": sum(GAMMA ** t * (returns[t] - V[s]) * scores[t] for t, (s, _, _) in enumerate(episode)), "(d) Σ(G_t-V)∇logπ 无γ^t": sum((returns[t] - V[s]) * scores[t] for t, (s, _, _) in enumerate(episode)), } return out def main(): P, Rsa = build_model() rng = np.random.default_rng(0) theta = np.array([[0.3, -0.2], [0.1, 0.0], [0.0, 0.2], [-0.1, 0.1], [0.2, 0.0]]) pi = policy_probs(theta) V, Q, _ = exact_values(theta, P, Rsa) # --------------------------------------------------------------------- print("=" * 78) print("演示 1：当前策略与精确价值") print("=" * 78) print(f" π(右|s) = {pi[:, RIGHT]}") print(f" V^π(s) = {V}") print(f" J(θ) = V^π(起点) = {V[START]:.4f}") # --------------------------------------------------------------------- print("\n" + "=" * 78) print("演示 2：验证策略梯度定理——两种方法算出的精确梯度（只显示 ∂J/∂θ[s, 右]）") print("=" * 78) g_fd = gradient_by_finite_difference(theta, P, Rsa) g_pg = gradient_by_pg_theorem(theta, P, Rsa) print(f" 有限差分（直接对 J 求导）：{g_fd[:, RIGHT]}") print(f" 策略梯度定理公式 ：{g_pg[:, RIGHT]}") print(f" 全部 10 个分量的最大差：{np.abs(g_fd - g_pg).max():.1e}") print(" → 两者一致。策略梯度定理说：梯度只需要 π、Q^π 和状态访问量，") print(" **不需要对环境的转移概率求导**——这正是它能用于无模型学习的原因。") print(f" 另外注意：∂J/∂θ[s,左] = -∂J/∂θ[s,右]（softmax 只有两个动作时的对称性）：" f"{np.allclose(g_fd[:, LEFT], -g_fd[:, RIGHT])}") # --------------------------------------------------------------------- print("\n" + "=" * 78) print("演示 3：四种 REINFORCE 估计量——无偏吗？方差多大？（20000 个回合）") print("=" * 78) n_episodes = 20_000 samples = {} for _ in range(n_episodes): ep = sample_episode(theta, rng) for name, g in estimators(theta, ep, V).items(): samples.setdefault(name, []).append(g) # 【Python 小知识】dict.setdefault(键, 默认值)：键不存在时先放入默认值，再返回它 print(f" 精确梯度（∂J/∂θ[s,右]）：{g_fd[:, RIGHT]}\n") print(" 估计量 样本均值（∂J/∂θ[s,右]） 偏差大小 总方差") for name, gs in samples.items(): gs = np.array(gs) mean = gs.mean(axis=0) bias = np.linalg.norm(mean - g_fd) total_var = gs.var(axis=0).sum() print(f" {name:<26s}{mean[:, RIGHT]} {bias:.3f} {total_var:8.2f}") print("\n （\"偏差大小\"是样本均值与精确梯度之差的长度；20000 个样本下，无偏估计量也会有") print(" 约 sqrt(总方差 / 20000) 量级的随机误差。）") print("\n 观察：") print(" - (a)(b)(c) 的样本均值都接近精确梯度：它们都是**无偏**的。") print(" - 方差逐级下降：(a) 用整条轨迹的回报给每一步\"打分\"，连动作之前的奖励也算了进去；") print(" (b) 只用动作之后的回报（因果性：现在的动作影响不了过去的奖励），方差降到约 1/5；") print(" (c) 再减去基线 V(s_t)，只看\"比平均好多少\"，方差进一步下降。") print(" （如果把时间成本改成 -0.2，\"动作之前的奖励\"很小，(a) 和 (b) 的方差就几乎一样——") print(" 利用因果性的收益取决于\"过去的奖励\"带来了多少噪声。可以改 STEP_COST 试试。）") print(" - (d) 去掉了 γ^t 因子（很多实际代码都这么写）：样本均值明显偏离精确梯度——它是**有偏**的，") print(" 不再是 J(θ) 的真正梯度（相当于把后面时间步的\"打分\"放大了，所以方差也更大）。") print(" 实践中人们常常接受这个偏差，把它理解为\"每一步都当作新的起点来优化\"，") print(" 但要知道它并不是严格意义上的梯度。") if __name__ == "__main__": main()