""" ============================================================================= 第 5 讲 · 代码 4：深度 Q 网络（DQN）——用 PyTorch 训练 CartPole ============================================================================= 对应讲义：lectures/lecture05_function_approximation_dqn.md 第 7、8 节 ----------------------------------------------------------------------------- 【这个文件要解决什么问题？】 ----------------------------------------------------------------------------- 把第 4 讲的 Q-learning 和一个神经网络结合起来：Q(s, a; θ) 由神经网络输出。 直接这样做很不稳定（03 文件的"致命三角"：函数逼近 + 自举 + 异策略全都有）。 DQN（Mnih et al., 2015）用两个关键技巧让它稳定下来： 1. 经验回放（experience replay）：把经历过的转移存进一个大缓冲区， 训练时**随机**抽一小批来学。打破数据之间的相关性，并且重复利用数据。 2. 目标网络（target network）：TD 目标用一个"冻结"的旧网络来计算， 每隔一段时间才同步一次。避免"追自己的尾巴"——目标随着网络一起变。 ----------------------------------------------------------------------------- 【环境：CartPole-v1（Gymnasium）】 ----------------------------------------------------------------------------- 一根杆子立在小车上，智能体通过向左/向右推小车来保持杆子不倒。 - 状态：4 个连续数（小车位置、小车速度、杆子角度、杆子角速度） - 动作：0 = 向左推，1 = 向右推 - 每坚持一步奖励 +1；杆子倾斜超过约 12° 或小车出界则失败；最多 500 步 - 满分 500 ----------------------------------------------------------------------------- 【运行方式】 ----------------------------------------------------------------------------- 需要额外安装：pip install torch gymnasium python 04_dqn_cartpole.py 标准 DQN（约 30 秒） python 04_dqn_cartpole.py --double Double DQN python 04_dqn_cartpole.py --no-replay 去掉经验回放（消融实验） python 04_dqn_cartpole.py --no-target 去掉目标网络（消融实验） python 04_dqn_cartpole.py --seed 3 换一个随机种子 ----------------------------------------------------------------------------- 【如果你没有学过 PyTorch】 ----------------------------------------------------------------------------- 只需要知道几件事（代码里遇到时都有注释）： - torch.Tensor：类似 numpy 数组，但能自动求导、能在 GPU 上算。 - nn.Module：神经网络的基类。定义好层之后，net(x) 就是前向计算。 - loss.backward()：自动计算 loss 对所有参数的梯度。 - optimizer.step()：按梯度更新参数（这里用 Adam，一种自适应步长的梯度下降）。 - with torch.no_grad()：这段计算不需要梯度（比如计算 TD 目标时）。 """ import argparse import time import numpy as np import torch import torch.nn as nn import gymnasium as gym # 所有超参数集中放在一个字典里，方便查看和修改 CONFIG = { "env_id": "CartPole-v1", "total_steps": 80_000, # 与环境交互的总步数 "buffer_size": 10_000, # 经验回放缓冲区容量 "batch_size": 64, # 每次训练抽取的样本数 "gamma": 0.99, # 折扣因子 "lr": 5e-4, # 学习率（Adam） "eps_start": 1.0, # ε-贪心：初始探索率 "eps_end": 0.05, # ε-贪心：最终探索率 "eps_decay_steps": 10_000, # ε 在多少步内从 eps_start 线性降到 eps_end "learning_starts": 1_000, # 先随机收集这么多步数据再开始训练 "target_update": 250, # 每隔多少步把在线网络的参数复制给目标网络 "hidden": 128, # 隐藏层神经元数 "grad_clip": 10.0, # 梯度裁剪阈值（防止偶尔的大梯度把网络"打坏"） "eval_every": 5_000, # 每隔多少步评估一次（用贪心策略，不探索） "eval_episodes": 10, } # 这组超参数是在 3 个随机种子上试了十几种组合后选出来的。DQN 对超参数相当敏感： # 学习率、缓冲区大小、目标网络更新频率稍微一变，结果就可能差很多。这是深度 RL 的常态。 # ============================================================================= # 第 1 部分：Q 网络 # ============================================================================= class QNetwork(nn.Module): """ 输入状态 s（4 个数），输出每个动作的 Q 值（2 个数）：Q(s, ·; θ)。 【为什么一次输出所有动作的 Q 值，而不是输入 (s, a) 输出一个数？】 选动作和计算 max_a' Q(s', a') 时需要所有动作的 Q 值。一次前向计算全部得到， 比对每个动作各算一次快得多。这是 DQN 论文的设计。 【结构】两层隐藏层的全连接网络（MLP）：4 → 128 → 128 → 2，激活函数 ReLU。 ReLU(x) = max(0, x)，给网络引入非线性；没有它，多层线性层叠起来还是线性的。 """ def __init__(self, obs_dim, n_actions, hidden): super().__init__() # 【Python 小知识】调用父类的初始化，PyTorch 要求这样写 self.net = nn.Sequential( # Sequential：把几层按顺序串起来 nn.Linear(obs_dim, hidden), # 全连接层：y = W x + b nn.ReLU(), nn.Linear(hidden, hidden), nn.ReLU(), nn.Linear(hidden, n_actions), # 输出层不加激活：Q 值可以是任意实数 ) def forward(self, x): """前向计算。x 形状 (batch, 4)，返回形状 (batch, 2)。""" return self.net(x) # ============================================================================= # 第 2 部分：经验回放缓冲区 # ============================================================================= class ReplayBuffer: """ 固定容量的"环形"缓冲区：存满之后，新数据覆盖最旧的数据。 【为什么要经验回放？】（讲义 7.2 节） 1. 打破相关性：连续的转移高度相关（相邻两帧几乎一样）。用相关的数据做 SGD， 梯度估计偏差很大，网络会"过拟合"最近的经历。随机抽样让一个批次里的样本近似独立。 2. 数据复用：每条转移可以被抽中很多次，提高样本效率。 3. 平滑分布：训练数据来自过去很多个策略的混合，不会因为策略的一点变化而剧烈摆动。 这之所以可行，是因为 Q-learning 是**异策略**的（第 4 讲 4.3 节）： 旧策略收集的数据，同样可以用来学最优策略的 Q 值。 【实现】用 numpy 数组预先分配空间，比用 Python 列表存元组快得多、也省内存。 """ def __init__(self, capacity, obs_dim, rng): self.capacity = capacity self.rng = rng self.obs = np.zeros((capacity, obs_dim), dtype=np.float32) self.next_obs = np.zeros((capacity, obs_dim), dtype=np.float32) self.actions = np.zeros(capacity, dtype=np.int64) self.rewards = np.zeros(capacity, dtype=np.float32) self.terminated = np.zeros(capacity, dtype=np.float32) self.pos = 0 # 下一条数据写入的位置 self.size = 0 # 当前存了多少条 def add(self, obs, action, reward, next_obs, terminated): """存入一条转移。写到 pos 位置，然后 pos 前进一格（到末尾后绕回开头）。""" i = self.pos self.obs[i], self.actions[i], self.rewards[i] = obs, action, reward self.next_obs[i], self.terminated[i] = next_obs, float(terminated) self.pos = (self.pos + 1) % self.capacity # % 实现"绕回开头" self.size = min(self.size + 1, self.capacity) def sample(self, batch_size): """ 均匀随机抽取 batch_size 条，返回 PyTorch 张量组成的元组。 【边界】调用前应保证 size >= batch_size（训练开始前先收集 learning_starts 步）。 """ idx = self.rng.integers(0, self.size, size=batch_size) return (torch.as_tensor(self.obs[idx]), torch.as_tensor(self.actions[idx]), torch.as_tensor(self.rewards[idx]), torch.as_tensor(self.next_obs[idx]), torch.as_tensor(self.terminated[idx])) # ============================================================================= # 第 3 部分：DQN 的核心——损失函数 # ============================================================================= def dqn_loss(batch, q_net, target_net, gamma, double): """ 计算一个批次的 DQN 损失。 【公式】（讲义 7.4 节） TD 目标： y = r + γ (1 - terminated) · max_a' Q(s', a'; θ⁻) （标准 DQN） y = r + γ (1 - terminated) · Q(s', argmax_a' Q(s', a'; θ); θ⁻) （Double DQN） 损失： L(θ) = mean( Huber( y - Q(s, a; θ) ) ) 其中 θ 是在线网络的参数（要训练的），θ⁻ 是目标网络的参数（冻结的）。 【几个关键细节】 1. y 的计算放在 torch.no_grad() 里：**目标不参与求导**。这就是"半梯度"（01 文件）。 2. (1 - terminated)：真正终止时（杆子倒了），没有未来，目标就是 r。 但因为"到达 500 步时间上限"而结束的回合**不算**终止——那时杆子还立着， 未来本来还有奖励，应当继续用 Q(s') 自举。Gymnasium 用 terminated 和 truncated 两个标志区分这两种情况，混淆它们是 DQN 代码里非常常见的 bug。 3. Huber 损失（smooth L1）：误差小时像平方损失，误差大时像绝对值损失， 对偶尔出现的超大 TD 误差不那么敏感，训练更稳定。 4. Double DQN（第 4 讲 6.3 节的思想）：用在线网络**挑**动作，用目标网络**评估**它， 缓解最大化偏差。 【参数】 batch : ReplayBuffer.sample 返回的元组。 q_net : 在线网络。 target_net : 目标网络（不使用目标网络时，传入 q_net 本身）。 gamma : 折扣因子。 double : 是否使用 Double DQN。 【返回】 (loss, 平均 Q 值)——平均 Q 值用来监控是否出现高估。 """ obs, actions, rewards, next_obs, terminated = batch # Q(s, a)：先算出所有动作的 Q 值 (batch, 2)，再用 gather 取出实际执行的那个动作的值 # 【PyTorch 小知识】gather(1, idx) 沿第 1 维按 idx 取元素；unsqueeze(1) 把 (batch,) 变成 (batch, 1) q_sa = q_net(obs).gather(1, actions.unsqueeze(1)).squeeze(1) with torch.no_grad(): if double: best_next = q_net(next_obs).argmax(dim=1, keepdim=True) # 在线网络挑 next_q = target_net(next_obs).gather(1, best_next).squeeze(1) # 目标网络评 else: next_q = target_net(next_obs).max(dim=1).values target = rewards + gamma * (1.0 - terminated) * next_q loss = nn.functional.smooth_l1_loss(q_sa, target) return loss, q_sa.mean().item() # ============================================================================= # 第 4 部分：动作选择与评估 # ============================================================================= def epsilon_by_step(step, cfg): """线性衰减的 ε：前 eps_decay_steps 步从 eps_start 线性降到 eps_end，之后保持不变。""" frac = min(1.0, step / cfg["eps_decay_steps"]) return cfg["eps_start"] + frac * (cfg["eps_end"] - cfg["eps_start"]) def select_action(q_net, obs, epsilon, n_actions, rng): """ε-贪心选择动作。""" if rng.random() < epsilon: return int(rng.integers(n_actions)) with torch.no_grad(): q = q_net(torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0)) # 加一个 batch 维度 return int(q.argmax(dim=1).item()) def evaluate(q_net, env_id, n_episodes, seed): """ 用**贪心**策略（不探索）跑 n_episodes 个回合，返回平均回报。 【为什么评估时不探索？】训练时的回报包含了 ε 探索的"噪声"，不能反映学到的策略本身有多好。 评估要用单独的环境，以免打乱训练环境的状态。 """ env = gym.make(env_id) returns = [] for i in range(n_episodes): obs, _ = env.reset(seed=seed + i) done, total = False, 0.0 while not done: with torch.no_grad(): a = int(q_net(torch.as_tensor(obs, dtype=torch.float32).unsqueeze(0)).argmax().item()) obs, r, terminated, truncated, _ = env.step(a) total += r done = terminated or truncated returns.append(total) env.close() return float(np.mean(returns)) # ============================================================================= # 第 5 部分：训练主循环 # ============================================================================= def train(cfg, seed=0, use_replay=True, use_target=True, double=False, verbose=True): """ DQN 训练主循环（讲义 7.5 节的伪代码）。 【消融选项】 use_replay=False：缓冲区容量设为 batch_size，于是每次训练用的就是**最近的** batch_size 条连续转移——相当于没有经验回放（数据高度相关、只用一次左右）。 use_target=False：TD 目标直接用在线网络计算（目标网络 = 在线网络）。 【返回】 (evals, best_net)： - evals : 评估记录列表 [(步数, 贪心策略平均回报), ...] - best_net : 评估成绩最好时的网络副本 【为什么要保存"最好的"网络？】 DQN 的性能常常在学到高分之后又突然下降（讲义 8.3 节的"灾难性遗忘"）。 实践中的标准做法是：定期评估，保留评估成绩最好的那一版参数，最后使用它。 """ torch.manual_seed(seed) # 固定 PyTorch 的随机性（网络初始化等） rng = np.random.default_rng(seed) env = gym.make(cfg["env_id"]) obs_dim = env.observation_space.shape[0] n_actions = env.action_space.n q_net = QNetwork(obs_dim, n_actions, cfg["hidden"]) if use_target: target_net = QNetwork(obs_dim, n_actions, cfg["hidden"]) target_net.load_state_dict(q_net.state_dict()) # 初始时两者参数相同 else: target_net = q_net # 不用目标网络：目标用在线网络算 optimizer = torch.optim.Adam(q_net.parameters(), lr=cfg["lr"]) capacity = cfg["buffer_size"] if use_replay else cfg["batch_size"] buffer = ReplayBuffer(capacity, obs_dim, rng) obs, _ = env.reset(seed=seed) episode_return, recent_returns = 0.0, [] losses, q_means, evals = [], [], [] best_return, best_net = -np.inf, QNetwork(obs_dim, n_actions, cfg["hidden"]) start_time = time.time() if verbose: print(f" {'步数':>6s} {'ε':>6s} {'最近10回合训练回报':>16s} {'贪心评估回报':>12s} {'平均Q值':>8s} {'损失':>8s} {'用时':>6s}") for step in range(1, cfg["total_steps"] + 1): # ---- 1. 与环境交互一步，存入缓冲区 ---- eps = epsilon_by_step(step, cfg) action = select_action(q_net, obs, eps, n_actions, rng) next_obs, reward, terminated, truncated, _ = env.step(action) buffer.add(obs, action, reward, next_obs, terminated) # 注意：只存 terminated，不存 truncated obs = next_obs episode_return += reward if terminated or truncated: recent_returns.append(episode_return) episode_return = 0.0 obs, _ = env.reset() # ---- 2. 从缓冲区抽一批数据，做一步梯度下降 ---- if step >= cfg["learning_starts"] and buffer.size >= cfg["batch_size"]: loss, q_mean = dqn_loss(buffer.sample(cfg["batch_size"]), q_net, target_net, cfg["gamma"], double) optimizer.zero_grad() # 清空上一步的梯度（PyTorch 默认会累加梯度） loss.backward() # 反向传播：计算梯度 nn.utils.clip_grad_norm_(q_net.parameters(), cfg["grad_clip"]) optimizer.step() # 更新参数 losses.append(loss.item()) q_means.append(q_mean) # ---- 3. 定期同步目标网络 ---- if use_target and step % cfg["target_update"] == 0: target_net.load_state_dict(q_net.state_dict()) # ---- 4. 定期评估 ---- if step % cfg["eval_every"] == 0: eval_return = evaluate(q_net, cfg["env_id"], cfg["eval_episodes"], seed=10_000 + step) evals.append((step, eval_return)) if eval_return > best_return: # 保存目前为止最好的网络 best_return = eval_return best_net.load_state_dict(q_net.state_dict()) if verbose: train_ret = np.mean(recent_returns[-10:]) if recent_returns else float("nan") q_m = np.mean(q_means[-500:]) if q_means else float("nan") l_m = np.mean(losses[-500:]) if losses else float("nan") print(f" {step:6d} {eps:6.3f} {train_ret:16.1f} {eval_return:12.1f} {q_m:8.2f} {l_m:8.4f} " f"{time.time() - start_time:5.0f}s") env.close() return evals, best_net def main(): parser = argparse.ArgumentParser(description="DQN on CartPole") # 【Python 小知识】argparse：解析命令行参数。action="store_true" 表示"写了这个参数就是 True"。 parser.add_argument("--double", action="store_true", help="使用 Double DQN") parser.add_argument("--no-replay", action="store_true", help="去掉经验回放") parser.add_argument("--no-target", action="store_true", help="去掉目标网络") parser.add_argument("--seed", type=int, default=0) parser.add_argument("--steps", type=int, default=CONFIG["total_steps"]) args = parser.parse_args() torch.set_num_threads(1) # 小网络在 CPU 上单线程反而更快（避免线程调度开销） cfg = dict(CONFIG, total_steps=args.steps) name = "Double DQN" if args.double else "DQN" if args.no_replay: name += "（无经验回放）" if args.no_target: name += "（无目标网络）" print("=" * 78) print(f"{name}，CartPole-v1，种子 {args.seed}，共 {cfg['total_steps']} 步") print("=" * 78) evals, best_net = train(cfg, seed=args.seed, use_replay=not args.no_replay, use_target=not args.no_target, double=args.double) best = max(r for _, r in evals) final = np.mean([r for _, r in evals[-3:]]) print(f"\n 最好的一次评估：{best:.1f}；最后 3 次评估的平均：{final:.1f}（满分 500）") recheck = evaluate(best_net, cfg["env_id"], 20, seed=99_999) print(f" 用保存下来的\"最好的网络\"重新评估 20 个新回合：平均 {recheck:.1f}") print(" 提示：DQN 在 CartPole 上的表现波动较大——即使已经学到满分，之后也可能突然下降") print(" （\"灾难性遗忘\"）。换几个 --seed 试试，并用 --no-replay / --no-target 做消融对比。") if __name__ == "__main__": main()