""" ============================================================================= 第 6 讲 · 代码 5：用 PyTorch 实现 REINFORCE（带基线）训练 CartPole ============================================================================= 对应讲义：lectures/lecture06_policy_gradient.md 第 8 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 把前面几个文件的思想放到神经网络上： - 策略网络 π_θ(a|s)：输入状态，输出两个动作的概率（softmax） - 价值网络 V_φ(s)：作为基线，输入状态，输出一个数 - 每次收集若干个完整回合，计算"未来回报" G_t，用 (G_t - V_φ(s_t)) 作为权重更新策略 和第 5 讲的 DQN 对比一下： DQN ：学 Q，策略是 argmax Q；异策略；可以用经验回放 REINFORCE：直接学策略；同策略（每批数据用完就扔）；不需要 argmax，天然支持随机策略 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- 需要：pip install torch gymnasium python 05_reinforce_cartpole.py 带基线的 REINFORCE（约 1 分钟） python 05_reinforce_cartpole.py --no-baseline 不用基线 python 05_reinforce_cartpole.py --full-return 每一步都用整条轨迹的回报 G_0（不利用因果性） python 05_reinforce_cartpole.py --no-normalize 不做优势标准化 python 05_reinforce_cartpole.py --seed 3 换一个随机种子 最"原始"的 REINFORCE： python 05_reinforce_cartpole.py --no-baseline --no-normalize --full-return """ import argparse import time import numpy as np import torch import torch.nn as nn import gymnasium as gym CONFIG = { "env_id": "CartPole-v1", "iterations": 200, # 更新多少次 "episodes_per_batch": 8, # 每次更新前收集多少个完整回合 "gamma": 0.99, "lr_policy": 3e-3, "lr_value": 1e-2, "hidden": 64, "normalize_advantage": True, # 把优势（G_t - b）标准化为均值 0、标准差 1 "eval_every": 20, } # ============================================================================= # 第 1 部分：网络 # ============================================================================= class PolicyNet(nn.Module): """ 策略网络：状态 → 每个动作的"偏好"（logits）。概率 = softmax(logits)。 【为什么输出 logits 而不是概率？】 PyTorch 的 Categorical 分布可以直接接收 logits，内部会用数值稳定的方式计算 softmax 和 log 概率（01 文件讲过 softmax 的溢出问题）。自己先算概率再取 log， 在概率很接近 0 时会出现 log(0) = -inf。 """ def __init__(self, obs_dim, n_actions, hidden): super().__init__() self.net = nn.Sequential( nn.Linear(obs_dim, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, n_actions), ) def distribution(self, obs): """返回动作分布对象。dist.sample() 采样动作，dist.log_prob(a) 计算 log π(a|s)。""" return torch.distributions.Categorical(logits=self.net(obs)) class ValueNet(nn.Module): """价值网络（基线）：状态 → V(s) 的估计。""" def __init__(self, obs_dim, hidden): super().__init__() self.net = nn.Sequential( nn.Linear(obs_dim, hidden), nn.Tanh(), nn.Linear(hidden, hidden), nn.Tanh(), nn.Linear(hidden, 1), ) def forward(self, obs): return self.net(obs).squeeze(-1) # (batch, 1) → (batch,) # ============================================================================= # 第 2 部分：收集数据 # ============================================================================= def collect_batch(env, policy, n_episodes, gamma, full_return): """ 用当前策略收集 n_episodes 个完整回合。 【返回】 (obs, actions, returns, episode_returns)： - obs, actions, returns：把所有回合的所有时间步拼在一起的张量 returns[t] 是这一步的"权重"： full_return=False：未来回报 G_t = r_t + γ r_{t+1} + ...（利用因果性，讲义 6 节） full_return=True ：整条轨迹的回报 G_0（每一步都一样） - episode_returns：每个回合的总奖励（不打折），用于观察学习进度 【关于截断】CartPole 到 500 步会被截断。这里的回报只算到截断为止， 严格来说会低估"本可以继续"的回合的回报；简单的 REINFORCE 通常这样处理。 """ all_obs, all_actions, all_returns, episode_returns = [], [], [], [] for _ in range(n_episodes): obs, _ = env.reset() ep_obs, ep_actions, ep_rewards = [], [], [] done = False while not done: obs_t = torch.as_tensor(obs, dtype=torch.float32) with torch.no_grad(): # 采样动作时不需要梯度 action = policy.distribution(obs_t).sample().item() next_obs, reward, terminated, truncated, _ = env.step(action) ep_obs.append(obs) ep_actions.append(action) ep_rewards.append(reward) obs = next_obs done = terminated or truncated # 从后往前算未来回报 G, returns = 0.0, np.zeros(len(ep_rewards)) for t in range(len(ep_rewards) - 1, -1, -1): G = ep_rewards[t] + gamma * G returns[t] = G if full_return: returns[:] = returns[0] # 每一步都用 G_0 all_obs.extend(ep_obs) all_actions.extend(ep_actions) all_returns.extend(returns) episode_returns.append(sum(ep_rewards)) return (torch.as_tensor(np.array(all_obs), dtype=torch.float32), torch.as_tensor(all_actions), torch.as_tensor(np.array(all_returns), dtype=torch.float32), episode_returns) def evaluate(policy, env_id, n_episodes=10, seed=12345): """ 用**最可能的动作**（而不是采样）评估策略。 【注意】策略梯度学到的是随机策略，训练时就是按它采样执行的；这里取最可能的动作 只是一种常见的评估方式，两者通常差不多。 """ env = gym.make(env_id) totals = [] for i in range(n_episodes): obs, _ = env.reset(seed=seed + i) done, total = False, 0.0 while not done: with torch.no_grad(): logits = policy.net(torch.as_tensor(obs, dtype=torch.float32)) obs, r, terminated, truncated, _ = env.step(int(logits.argmax().item())) total += r done = terminated or truncated totals.append(total) env.close() return float(np.mean(totals)) # ============================================================================= # 第 3 部分：训练 # ============================================================================= def train(cfg, seed=0, use_baseline=True, full_return=False, normalize=None, verbose=True): """ REINFORCE 主循环（讲义 8.1 节的伪代码）： 重复 iterations 次： 1. 用当前策略收集 episodes_per_batch 个回合 2. 计算每一步的回报 G_t 3. 优势 A_t = G_t - V_φ(s_t)（不用基线时 A_t = G_t）；可选：标准化 4. 策略损失 L(θ) = -mean( log π_θ(a_t|s_t) · A_t )，梯度下降一步 （最小化 -J 就是最大化 J；这个"损失"的梯度恰好是策略梯度的估计） 5. 价值损失 L(φ) = mean( (V_φ(s_t) - G_t)² )，梯度下降一步 【几个关键细节】 - A_t 必须 detach（不对价值网络求导）：基线只是一个"参照点"， 策略的梯度不应该流进价值网络。 - 优势标准化（减均值、除标准差）是一个常用技巧：让每批数据的"打分"尺度一致， 学习率更好调。它引入了一点点偏差，但通常利大于弊。 注意："减均值"这一步本身就相当于用"这批数据的平均回报"作为一个常数基线！ 所以即使 --no-baseline，只要开着标准化，其实仍然有一个（较粗糙的）基线。 想看"真正没有基线"的效果，要同时加上 --no-normalize。 - 这里没有乘 γ^t 因子（04 文件的估计量 (d)）——这是几乎所有实际代码的做法。 【返回】[(迭代次数, 最近一批的平均回报, 评估回报或 None), ...] """ torch.manual_seed(seed) env = gym.make(cfg["env_id"]) env.reset(seed=seed) env.action_space.seed(seed) obs_dim, n_actions = env.observation_space.shape[0], env.action_space.n policy = PolicyNet(obs_dim, n_actions, cfg["hidden"]) value = ValueNet(obs_dim, cfg["hidden"]) opt_pi = torch.optim.Adam(policy.parameters(), lr=cfg["lr_policy"]) opt_v = torch.optim.Adam(value.parameters(), lr=cfg["lr_value"]) log, start = [], time.time() if verbose: print(f" {'迭代':>4s} {'累计回合':>8s} {'本批平均回报':>12s} {'评估回报':>8s} {'用时':>6s}") for it in range(1, cfg["iterations"] + 1): obs, actions, returns, ep_returns = collect_batch(env, policy, cfg["episodes_per_batch"], cfg["gamma"], full_return) # ---- 优势 ---- if use_baseline: with torch.no_grad(): adv = returns - value(obs) else: adv = returns.clone() do_norm = cfg["normalize_advantage"] if normalize is None else normalize if do_norm: adv = (adv - adv.mean()) / (adv.std() + 1e-8) # 加 1e-8 防止除以 0 # ---- 策略更新 ---- log_probs = policy.distribution(obs).log_prob(actions) loss_pi = -(log_probs * adv).mean() opt_pi.zero_grad() loss_pi.backward() opt_pi.step() # ---- 价值（基线）更新 ---- if use_baseline: loss_v = ((value(obs) - returns) ** 2).mean() opt_v.zero_grad() loss_v.backward() opt_v.step() eval_ret = evaluate(policy, cfg["env_id"]) if it % cfg["eval_every"] == 0 else None log.append((it, float(np.mean(ep_returns)), eval_ret)) if verbose and eval_ret is not None: print(f" {it:4d} {it * cfg['episodes_per_batch']:8d} {np.mean(ep_returns):12.1f} " f"{eval_ret:8.1f} {time.time() - start:5.0f}s") env.close() return log def main(): parser = argparse.ArgumentParser(description="REINFORCE on CartPole") parser.add_argument("--no-baseline", action="store_true", help="不使用价值基线") parser.add_argument("--full-return", action="store_true", help="每一步都用整条轨迹的回报 G_0") parser.add_argument("--no-normalize", action="store_true", help="不做优势标准化") parser.add_argument("--seed", type=int, default=0) args = parser.parse_args() torch.set_num_threads(1) name = "REINFORCE" + ("（无价值基线）" if args.no_baseline else "（有价值基线）") if args.full_return: name += "（整条轨迹回报）" if args.no_normalize: name += "（不标准化）" print("=" * 78) print(f"{name}，CartPole-v1，种子 {args.seed}，" f"{CONFIG['iterations']} 次迭代 × 每次 {CONFIG['episodes_per_batch']} 个回合") print("=" * 78) log = train(CONFIG, seed=args.seed, use_baseline=not args.no_baseline, full_return=args.full_return, normalize=not args.no_normalize) last = np.mean([r for _, r, _ in log[-20:]]) print(f"\n 最后 20 次迭代的平均训练回报：{last:.1f}（满分 500）") if __name__ == "__main__": main()