PyTorch强化学习实战（9）——深度Q学习

PyTorch强化学习实战（9）------深度Q学习

• • [0. 前言](#0. 前言)
  • [1. 深度Q学习](#1. 深度Q学习)
  • • [1.1 环境交互机制](#1.1 环境交互机制)
    • [1.2 SGD 优化](#1.2 SGD 优化)
    • [1.3 步骤之间的相关性](#1.3 步骤之间的相关性)
    • [1.4 马尔可夫性质](#1.4 马尔可夫性质)
  • [2. 深度Q学习训练的最终形式](#2. 深度Q学习训练的最终形式)
  • [3. 从零开始实现深度Q网络](#3. 从零开始实现深度Q网络)
  • • [3.1 Atari 环境包装器](#3.1 Atari 环境包装器)
    • [3.2 DQN 模型](#3.2 DQN 模型)
    • [3.3 模型训练](#3.3 模型训练)
    • [3.4 模型性能](#3.4 模型性能)
  • [4. 模型实战演示](#4. 模型实战演示)
  • 小结
  • 系列链接

0. 前言

我们已经讨论价值迭代方法的问题，并了解了它的变体，Q学习。在本节中，我们将通过神经网络预处理环境状态来扩展Q学习方法，这将极大提升Q学习方法的灵活性和适用性。

1. 深度Q学习

Q学习方法解决了遍历完整状态集的问题，但当可观测状态数量极大时仍会面临挑战。例如，Atari 游戏屏幕画面组合千变万化------若将原始像素作为独立状态处理，很快就会面临状态数量爆炸导致的价值函数难以估算问题。

在一些环境中，可观测状态量近乎无限。例如，在 CartPole 中，环境提供的状态由 4 个浮点数构成，虽然理论组合数有限(由比特位决定)，但其规模已达 2 ( 4 ⋅ 32 ) ≈ 3.4 × 10 38 2^{(4·32)}≈3.4×10^{38} 2(4⋅32)≈3.4×1038 量级。尽管实际取值受环境限制会减少，所以并不是所有 4 个 float32 值的位组合都是可能的，但状态空间仍然过于庞大。虽然可以创建一些区间来离散化这些值，但这往往引发更多问题；我们需要人工判定参数区间的区分粒度，但这与我们"不窥探环境内部机制"的通用实现原则相悖。

对于 Atari 来说，单个像素变化通常无关紧要，相似图像可归为同一状态。然而，我们仍然需要区分一些状态。以下图片显示了 Pong 游戏中的两种不同情况。我们通过控制一个球拍与人工智能 (Artificial Intelligence, AI) 对手进行对战(我们的球拍在右边，而对手的球拍在左边)。游戏的目标是将乒乓球打过对手的球拍，同时防止球打过我们的球拍。我们可以认为以下两种情况完全不同，在右侧显示的情况中，球离对手很近，所以我们可以保持静止，而左图球飞向我方时则需快速移动球拍防守。虽然这两种状态来自 10 70802 10^{70802} 1070802 种可能组合，但智能体必须做出差异化响应。

针对这一问题，我们可以采用一种将状态和动作映射为价值的非线性表示方法。在机器学习中，这称为"回归问题"。虽然这种表示方法的实现形式和训练方式多种多样，但使用深度神经网络 (Neural Network, NN) 是最流行的选择之一，尤其是在处理以屏幕图像呈现的观测数据时。基于这一思路，我们对Q学习算法进行以下修改：

1. 使用初始近似值初始化 Q ( s , a ) Q(s,a) Q(s,a)
2. 通过与环境的交互获取数据元组 ( s , a , r , s ′ ) (s,a,r,s′) (s,a,r,s′)
3. 计算损失：
   L = { ( Q ( s , a ) − r 2 ) 回合结束 ( Q ( s , a ) − ( r + γ m a x a ′ ∈ A Q ( s ′ , a ′ ) ) ) 2 其它 \mathcal L= \begin{cases} (Q(s,a)-r^2) & 回合结束\\ (Q(s,a)-(r+\gamma \underset {a'\in A}{max}Q(s',a')))^2& 其它 \end{cases} L={(Q(s,a)−r2)(Q(s,a)−(r+γa′∈AmaxQ(s′,a′)))2回合结束其它
4. 使用随机梯度下降 (Stochastic Gradient Descent, SGD) 算法，通过最小化损失来更新 Q ( s , a ) Q(s, a) Q(s,a)，并调整模型参数
5. 从步骤 2 重复直到收敛

这个算法看似简单，但实际效果往往不尽如人意。下面我们将探讨可能导致问题的若干因素，以及应对这些情况的潜在解决方案。

1.1 环境交互机制

首先，我们需要与环境进行交互以获取训练数据。在简单的环境中，比如 FrozenLake，随机动作尚可接受，但这并非最优策略以乒乓球游戏为例：随机移动球拍赢得一分的概率虽然不是零，但微乎其微------这意味着我们需要耗费极长时间才能捕捉到这种罕见情况。替代方案是采用Q函数近似作为动作策略(正如在价值迭代法中通过测试阶段积累经验的做法)。

当Q函数表征准确时，环境交互产生的经验数据能为智能体提供有效的训练素材。但若近似效果欠佳(例如训练初期)，智能体可能陷入某些状态的次优行动中而无法尝试新策略,这就是探索-利用困境。一方面，智能体需要探索环境以构建完整的状态转移与动作结果图谱。必须高效利用交互过程，避免重复尝试已知效果的动作。

由此可见，在训练初期Q函数近似较差时，随机动作更具优势------它能提供更均匀的环境状态信息。随着训练推进，随机动作效率下降，此时应转向依赖Q函数决策。

ε-贪婪算法 (epsilon-greedy method) 正是协调这两种行为的经典方法：通过超参数 ε ε ε 控制随机策略与Q策略之间的切换。通过调整epsilon，我们可以选择随机动作的比例，通常做法是将 ε ε ε 从 1.0 (100% 随机)逐步衰减至一个较小值(如 5% 或 2%)。这种方法既保证了初期的充分探索，又实现了后期的策略稳定。探索与利用问题还存在其他解决方案，但这仍是强化学习领域尚未得到根本解决的核心问题，也是当前研究的前沿方向之一。

1.2 SGD 优化

Q学习的核心流程借鉴了监督学习范式：试图用神经网络逼近复杂的非线性 Q ( s , a ) Q(s,a) Q(s,a) 函数，通过贝尔曼方程计算目标值后，将其转化为监督学习问题进行处理。这种思路本身可行，但 SGD 优化的基本前提是训练数据必须满足独立同分布 (independent and identically distributed, iid) 特性------即数据应从目标分布中随机采样获得。而我们的训练数据存在两个关键缺陷：

• 样本非独立性：即便积累大批量数据，由于同属一个训练回合 (episode)，样本间具有强相关性
• 分布偏移问题：训练数据分布与待学习最优策略的样本分布不一致。现有数据来自其他策略(当前策略/随机策略/ε-贪婪策略的混合)，而我们的目标是学习能获取最高奖励的最优策略

为解决这个问题，通常需要建立历史经验的大容量缓冲区，并从中抽样训练数据(而非直接使用最新经验)，该技术称为经验回放缓冲区。最简单的实现方式是固定容量的先进先出缓冲区，新数据不断覆盖最旧经验。

经验回放缓冲区能提供基本独立的训练数据，同时保证数据时效性(仍由近期策略生成)。

1.3 步骤之间的相关性

默认训练流程的另一个实际问题同样与数据非独立同分布相关，但表现形式略有不同。贝尔曼方程通过 Q ( s ′ , a ′ ) Q(s′,a′) Q(s′,a′) 来计算 Q ( s , a ) Q(s,a) Q(s,a) 的值(这种递归使用公式的过程称为自举法)。然而，状态 s s s 与 s ′ s′ s′ 之间仅一步之遥，导致二者高度相似，神经网络很难区分它们。当我们更新网络参数使 Q ( s , a ) Q(s,a) Q(s,a) 逼近目标值时，可能间接影响 Q ( s ′ , a ′ ) Q(s′,a′) Q(s′,a′) 及邻近状态的输出结果。这种机制会使训练陷入"追尾循环"的不稳定状态：更新状态 s s s 的Q值后，后续可能发现 Q ( s ′ , a ′ ) Q(s′,a′) Q(s′,a′) 表现变差，而对其的修正又可能破坏先前对 Q ( s , a ) Q(s,a) Q(s,a) 的拟合，如此恶性循环。

为了使训练更加稳定，可采用目标网络技巧：保留一个网络副本，专门用于计算贝尔曼方程中的 Q ( s ′ , a ′ ) Q(s′,a′) Q(s′,a′)。该副本仅定期与主网络同步(例如每 N 步同步一次，其中 N 通常取较大超参数值，如 1000 或 10000 次训练迭代)。

1.4 马尔可夫性质

我们的强化学习方法以马尔可夫决策过程 (Markov Decision Process, MDP)为理论基础，其核心假设是环境满足马尔可夫性质：仅凭当前观测即可做出最优决策。换言之，观测数据应能完全区分不同状态。

如以上 Pong 游戏屏幕截图中所见，单帧画面显然无法捕捉全部关键信息(仅凭一帧无法判断球与对手球拍的运动速度和方向)。这显然违反了马尔可夫性质，使单帧观测的 Pong 环境退化为部分可观测马尔可夫决策过程 (Partially Observable MDP, POMDP)。POMDP 本质上是缺失马尔可夫性质的 MDP，在实际中极为常见------例如多数卡牌游戏中，由于无法看到对手手牌，游戏观察值就是 POMDP，当前观测(己方手牌与桌面牌组)可能对应对手不同的持牌组合。

通过简单技巧使环境回归 MDP 范畴：维护历史观测序列并将其整合为状态表征。针对 Atari 游戏，通常将连续 k 帧画面堆叠作为状态观测(经典取值为 k=4)，使智能体能推断当前状态动态(如球体速度与方向)。尽管这种方法无法解决环境中的长程依赖问题，但对多数游戏已足够有效。

2. 深度Q学习训练的最终形式

研究者们已发现众多提升深度Q学习 (Deep Q-learning) 训练稳定性和效率的技巧，但仅凭ε-贪婪策略、经验回放缓冲区和目标网络这三个核心组件，就能成功在 49 款 Atari 游戏上完成深度Q网络 (Deep Q-Network, DQN) 训练，证明了该方法在复杂环境中的卓越效能。

原始论文 Playing Atari with deep reinforcement learning (没有目标网络)于 2013 年发布，使用了七个游戏进行测试。随后，在 2015 年初，修订版 Human-level control through deep reinforcement learning 发布，将测试规模扩展至 49 款游戏。这两篇论文提出的 DQN 算法流程如下：

1. 随机初始化 Q ( s , a ) Q(s,a) Q(s,a) 和 Q ^ ( s , a ) \hat Q(s,a) Q^(s,a) 网络参数，设 ε ← 1.0 ε←1.0 ε←1.0，清空经验回放缓冲区
2. 以概率 ε ε ε 选择随机动作 a a a，否则执行 a = a r g m a x a Q ^ ( s , a ) a = \underset {a}{argmax} \hat Q(s,a) a=aargmaxQ^(s,a)
3. 在模拟器中执行动作 a a a，获取奖励 r r r 和下一状态 s ′ s′ s′
4. 将转移样本 ( s , a , r , s ′ ) (s,a,r,s′) (s,a,r,s′) 存入经验回放缓冲区
5. 从缓冲区随机采样小批量转移样本
6. 对缓冲区中的每个样本计算目标值：
   y = { r 回合结束 r + γ m a x a ′ ∈ A Q ^ ( s ′ , a ′ ) 其它 \mathcal y= \begin{cases} r & 回合结束\\ r+\gamma \underset {a'\in A}{max}\hat Q(s',a')& 其它 \end{cases} y={rr+γa′∈AmaxQ^(s′,a′)回合结束其它
7. 计算损失函数： L = ( Q ( s , a ) − y ) 2 \mathcal L = (Q(s, a) − y)^2 L=(Q(s,a)−y)2
8. 通过 SGD 算法最小化损失，更新 Q ( s , a ) Q(s,a) Q(s,a) 网络参数
9. 每执行 N 步后，将 Q Q Q 网络权重同步至 Q ^ \hat Q Q^ 网络
10. 从第 10 步开始重复，直到收敛。

接下来，实现该算法算法进行 Atari 游戏。

3. 从零开始实现深度Q网络

FrozenLake 和 CartPole 等环境对计算资源要求较低------观测空间小、神经网络参数量少，训练循环中节省几毫秒并无实质意义。但 Atari 环境的情况截然不同：单次观测就包含 10 万个数据点，需要经过预处理、缩放并存入经验回放缓冲区。即便一个额外的数据拷贝操作，在最先进的 GPU 上也可能导致训练时间从秒级延长至小时级。

神经网络训练循环同样可能成为瓶颈。虽然强化学习模型的参数量远不及当今大语言模型 (Large Language Model, LLM)，但原始 DQN 就有超过 150 万个参数需要数百万次调整。简言之，性能优化至关重要------特别是当需要调试超参数并同时训练数十个模型时。
PyTorch 虽然表达力强，编写相对高效的代码比优化 TensorFlow 计算图更直观，但仍存在低效操作的风险。例如：逐样本计算 DQN 损失的实现比并行版本慢 2 倍；单个多余的数据拷贝操作可使代码运行速度降低 13 倍。

基于代码长度、逻辑结构和复用性考虑，我们将代码实现分为三个模块：：

• wrappers.py：Atari 环境包装器
• dqn_model.py：DQN 神经网络层
• dqn_pong.py：主模块(包含训练循环、损失函数计算和经验回放缓冲区)

3.1 Atari 环境包装器

在本节中，我们使用了 stable-baseline3 中的 AtariWrapper 类，并禁用了部分非必要包装器，代码实现保存为 wrappers.py，关于环境包装器的具体实现细节，参考《Atari 游戏包装器》一节。了解了包装器后，接下来，查看模型部分。

3.2 DQN 模型

模型采用三层卷积网络接两层全连接层结构，所有层之间都使用了 ReLU 非线性激活函数，代码实现保存为 dqn _model.py。模型直接输出环境中所有可用动作的Q值(未施加非线性约束，因Q值理论上可取值任意范围)。这种单次前向传播即可计算出所有动作Q值的设计，相比传统逐动作计算 Q ( s , a ) Q(s,a) Q(s,a) 的方式能显著提升运算效率。

import torch
import torch.nn as nn

class DQN(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(DQN, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU(),
            nn.Flatten(),
        )
        size = self.conv(torch.zeros(1, *input_shape)).size()[-1]
        self.fc = nn.Sequential(
            nn.Linear(size, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions)
        )

为了实现网络架构的通用性，我们将其分为卷积模块和线性模块两部分实现。卷积模块负责处理 4×84×84 维度的输入张量(对应四帧堆叠的 84×84 灰度图像)，最终将最后一个卷积层的输出展平为一维向量后馈入两个全连接层。

另一问题是在于，由于输入尺寸变化会影响卷积层输出值的总数，我们无法预先确定最终展平向量的具体维度。虽然对于固定 84×84 的输入，经过特定卷积配置后输出值恒为 3136 个(可通过硬编码解决)，但这种方式会降低代码对输入尺寸变化的适应性。更优的方案是在模型初始化时，通过向卷积模块传入模拟输入张量来动态获取该维度值------这种方法只需在创建模型时执行一次前向计算，既能保证运行效率，又能使代码具备通用性。
forward()函数接收4维输入张量，第一维度是批大小，第二维度对应堆叠帧的色彩通道，第三、四维度则为图像的高度和宽度：

    def forward(self, x: torch.ByteTensor):
        # scale on GPU
        xx = x / 255.0
        return self.fc(self.conv(xx))

在将数据输入网络之前，我们需要先进行数据缩放和类型转换处理。
Atari 图像中的每个像素都以无符号字节 (unsigned byte) 形式存储，取值范围为 0 到 255。这种表示方式在两方面具有优势：内存效率和 GPU 带宽利用率。从内存角度来看，由于经验回放池需要存储大量观测数据(通常数千帧)，使用 uint8 类型的 numpy 数组能最大限度节省内存空间。另一方面，训练过程中需要频繁将观测数据从主机内存传输至 GPU 显存，而内存与显存之间的带宽是有限资源，因此压缩数据体积能显著提升传输效率。

因此，我们始终以 dtype=uint8 格式存储原始观测数据，网络输入张量也采用 ByteTensor 类型。但 PyTorch的Conv2D 层要求输入为浮点张量，因此通过将输入张量除以 255.0，我们既完成了 [0,1] 范围的归一化，也实现了类型转换 (ByteTensor→FloatTensor)。由于此时字节张量已位于 GPU 显存中，该转换过程非常高效。最终，我们将缩放后的张量送入网络的卷积模块和线性模块进行处理。

3.3 模型训练

dqn_pong.py 模块包含经验回放缓冲区、智能体、损失函数计算以及训练循环。在代码实现之前，首先说明训练超参数的设置。
DeepMind 的论文提供了完整超参数表，用于在 49 款 Atari 游戏上训练模型。对所有游戏采用相同参数配置(但为每款游戏单独训练模型)，旨在证明该方法具有足够鲁棒性------仅凭单一模型架构和超参数组合，就能攻克不同复杂度、动作空间、奖励机制等特性的多款游戏。但本节的目标相对简单，只需解决 Pong 游戏。

与 Atari 中的其他游戏相比，Pong 十分简单直接，因此原论文中的超参数对本节的任务而言有些过于复杂。例如，为了在所有 49 个游戏中获得最佳结果，DeepMind 使用了百万级观测数据的回放缓冲区，这需要约 20GB 内存存储，且需从环境中采集大量样本填充缓冲区。

原论文所使用的ε-衰减策略对于简单的 Pong 游戏也非最优。在训练过程中，DeepMind 将 ε ε ε 从 1.0 线性衰减到 0.1，直到获得环境中的第一百万帧。但实验表明，对 Pong 而言，在前 15 万帧完成 ε ε ε 衰减并保持稳定即可。回放缓冲区也可大幅缩小，仅需 1 万条状态转移数据就足够了。

需要注意的是，这种加速优化仅针对特定环境，可能破坏其他游戏的收敛性。我们也可自由尝试不同参数配置或挑战 Atari 游戏集中的其他游戏。

(1) 首先，导入所需的模块：

import gymnasium as gym
import dqn_model
import wrappers

from dataclasses import dataclass
import argparse
import time
import numpy as np
import collections
import typing as tt

import torch
import torch.nn as nn
import torch.optim as optim

from torch.utils.tensorboard.writer import SummaryWriter
from ale_py import ALEInterface

(2) 定义超参数：

DEFAULT_ENV_NAME = "PongNoFrameskip-v4"
MEAN_REWARD_BOUND = 19

以上两个值分别设定了默认训练环境和停止训练的阈值条件(最近 100 轮的平均奖励达到该值时终止训练)。如需修改训练环境，可通过命令行参数 --env 重新指定。

继续定义以下超参数：

• 用于贝尔曼近似的 γ γ γ 值 (GAMMA)
• 从经验回放缓冲区采样的批大小 (BATCH_SIZE)
• 经验回放缓冲区的最大容量 (REPLAY_SIZE)
• 开始训练前等待填充经验回放缓冲区的帧数 (REPLAY_START_SIZE)
• Adam 优化器使用的学习率 (LEARNING_RATE)
• 将训练模型权重同步至目标模型的频率，目标模型用于获取贝尔曼近似计算中下一状态的价值 (SYNC_TARGET_FRAMES)

GAMMA = 0.99
BATCH_SIZE = 32
REPLAY_SIZE = 10000
LEARNING_RATE = 1e-4
SYNC_TARGET_FRAMES = 1000
REPLAY_START_SIZE = 10000

最后一批超参数与 ε ε ε 衰减策略相关。为实现有效探索，我们在训练初期设定 ε = 1.0 ε=1.0 ε=1.0，此时所有动作都将被随机选择。随后在前 150000 帧训练过程中， ε ε ε 会线性衰减至 0.01，这意味着仅有 1% 的步骤会采取随机动作。原始 DeepMind 论文采用了类似机制，但其衰减周期延长了近 10 倍(达到 ε = 0.01 ε=0.01 ε=0.01 需要经过百万帧训练)：

EPSILON_DECAY_LAST_FRAME = 150000
EPSILON_START = 1.0
EPSILON_FINAL = 0.01

(3) 定义类型别名及数据类 Experience，用于存储经验回放缓冲区的记录项。该数据类包含以下字段：当前状态、采取的动作、获得的奖励、终止/截断标志位，以及新状态：

State = np.ndarray
Action = int
BatchTensors = tt.Tuple[
    torch.ByteTensor,           # current state
    torch.LongTensor,           # actions
    torch.Tensor,               # rewards
    torch.BoolTensor,           # done || trunc
    torch.ByteTensor            # next state
]

@dataclass
class Experience:
    state: State
    action: Action
    reward: float
    done_trunc: bool
    new_state: State

(4) 接下来，定义经验回放缓冲区 (experience replay buffer)，其作用是存储从环境中获取的状态转移数据：

class ExperienceBuffer:
    def __init__(self, capacity: int):
        self.buffer = collections.deque(maxlen=capacity)

    def __len__(self):
        return len(self.buffer)

    def append(self, experience: Experience):
        self.buffer.append(experience)

    def sample(self, batch_size: int) -> tt.List[Experience]:
        indices = np.random.choice(len(self), batch_size, replace=False)
        return [self.buffer[idx] for idx in indices]

每次在环境中执行一步动作时，我们会将该状态转移数据存入缓冲区，但仅保留固定数量的步骤(本节中为 1 万条转移数据)。训练时，我们从回放缓冲区随机采样一批转移数据，这种方法能有效打破环境连续步骤之间的相关性。

经验回放缓冲区主要利用了 deque 类的特性来维护缓冲区中固定数量的条目。在 sample() 方法中，我们生成随机索引列表，并返回一组 Experience 实例，这些实例重新打包并转换为张量。

(5) 接下来，实现 Agent 类，负责与环境交互，并将交互结果存储到经验回放缓冲区中：

class Agent:
    def __init__(self, env: gym.Env, exp_buffer: ExperienceBuffer):
        self.env = env
        self.exp_buffer = exp_buffer
        self.state: tt.Optional[np.ndarray] = None
        self._reset()

    def _reset(self):
        self.state, _ = env.reset()
        self.total_reward = 0.0

在智能体初始化过程中，我们需要存储对环境和经验回放缓冲区的引用，同时跟踪当前观察值和累计总奖励。

(6) Agent 的主要方法是在环境中执行一步动作并将结果存储在经验回放缓冲区中。为此，首先需要选择动作：

    @torch.no_grad()
    def play_step(self, net: dqn_model.DQN, device: torch.device,
                  epsilon: float = 0.0) -> tt.Optional[float]:
        done_reward = None

        if np.random.random() < epsilon:
            action = env.action_space.sample()
        else:
            state_v = torch.as_tensor(self.state).to(device)
            state_v.unsqueeze_(0)
            q_vals_v = net(state_v)
            _, act_v = torch.max(q_vals_v, dim=1)
            action = int(act_v.item())

我们以 epsilon 概率(作为参数传入)随机选择动作；否则，就通过模型获取所有可能动作的Q值并选择最优动作。在此方法中，我们使用 PyTorch 的 no_grad() 装饰器来禁用整个方法期间的梯度追踪，因为我们并不需要这些梯度信息。

(7) 当动作选定后，我们将其传递给环境以获取下一个观测和奖励，并将数据存入经验回放缓冲区，随后处理回合结束的情况：

        new_state, reward, is_done, is_tr, _ = self.env.step(action)
        self.total_reward += reward

        exp = Experience(
            state=self.state, action=action, reward=float(reward),
            done_trunc=is_done or is_tr, new_state=new_state
        )
        self.exp_buffer.append(exp)
        self.state = new_state
        if is_done or is_tr:
            done_reward = self.total_reward
            self._reset()
        return done_reward

如果当前步骤导致回合结束，则该函数的返回值为累计总奖励；否则返回 None。

(8) 函数 batch_to_tensors 接收一批 Experience 对象，并将其重新打包为 PyTorch 张量，返回一个包含状态、动作、奖励、完成标志和新状态的元组，各元素均转换为相应类型的张量：

def batch_to_tensors(batch: tt.List[Experience], device: torch.device) -> BatchTensors:
    states, actions, rewards, dones, new_state = [], [], [], [], []
    for e in batch:
        states.append(e.state)
        actions.append(e.action)
        rewards.append(e.reward)
        dones.append(e.done_trunc)
        new_state.append(e.new_state)
    states_t = torch.as_tensor(np.asarray(states))
    actions_t = torch.LongTensor(actions)
    rewards_t = torch.FloatTensor(rewards)
    dones_t = torch.BoolTensor(dones)
    new_states_t = torch.as_tensor(np.asarray(new_state))
    return states_t.to(device), actions_t.to(device), rewards_t.to(device), \
           dones_t.to(device),  new_states_t.to(device)

在与状态数据交互时，通过 np.asarray() 函数避免内存拷贝操作，这一优化至关重要。因为 Atari 游戏的观测数据量庞大(包含 4 帧 84×84 字节的图像)，而我们每次需要处理一批包含 32 个这样的观测数据。若不进行此优化，性能大约会下降 20 倍。

(9) 训练模块的 calc_loss 函数用于计算采样批次的损失。该函数采用向量化操作并行处理整个批次的样本，最大限度地发挥了 GPU 的并行计算能力。虽然这种实现方式比直观的循环批处理更难理解，但优化效果显著，并行版本的速度比显式循环快两倍以上。需要计算的损失函数表达式如下：
L = ( Q ( s , a ) − ( r + γ m a x a ′ ∈ A Q ^ ( s ′ , a ′ ) ) ) 2 \mathcal L=(Q(s,a)-(r+\gamma \underset{a'\in A}{max}\hat Q(s',a')))^2 L=(Q(s,a)−(r+γa′∈AmaxQ^(s′,a′)))2

对于非终止时间步，使用以上方程计算；而对于终止时间步，则采用以下公式：
L = ( Q ( s , a ) − r ) 2 \mathcal L=(Q(s,a)-r)^2 L=(Q(s,a)−r)2

def calc_loss(batch: tt.List[Experience], net: dqn_model.DQN, tgt_net: dqn_model.DQN,
              device: torch.device) -> torch.Tensor:
    states_t, actions_t, rewards_t, dones_t, new_states_t = batch_to_tensors(batch, device)

在函数参数中，我们传入以下对象：当前处理的批数据、正在训练的主网络，以及目标网络(它会周期性与训练好的主网络同步)。

第一个模型(通过 net 参数传入)用于计算梯度；而目标网络( tgt_net 参数)则用于计算下一状态的价值，并且该计算过程不应影响梯度。为此，我们调用 PyTorch 张量的 detach() 方法，防止梯度流向目标网络的计算图。

在函数开始时，首先调用 batch_to_tensors 函数，将批数据解包为独立的张量变量：

    state_action_values = net(states_t).gather(
        1, actions_t.unsqueeze(-1)
    ).squeeze(-1)

接下来，将观测数据输入第一个模型，并通过张量的 gather() 操作提取已执行动作对应的特定Q值。gather() 调用的第一个参数指定了执行收集操作的维度索引(本节中为 1，对应动作维度)。

第二个参数是待选取元素的索引张量。此处需要通过 unsqueeze() 和 squeeze() 调用来分别完成两项操作：为 gather() 函数构建索引参数，以及消除我们创建的额外维度(索引张量的维度必须与处理数据的维度保持一致)。下图展示了 gather() 在典型场景下的运作机制：该示例包含 6 个批次样本和 4 个可选动作，通过可视化方式呈现了索引收集过程。

需要注意的是，对张量执行 gather() 操作的结果是可微分的，该操作会保留与最终损失值相关的所有梯度信息。

(10) 接下来，我们暂时禁用梯度计算(这能带来小幅性能提升)，将目标网络应用于下一状态的观测数据，并沿着相同的动作维度(维度1)计算最大Q值：

    with torch.no_grad():
        next_state_values = tgt_net(new_states_t).max(1)[0]

max() 函数会同时返回最大值及其索引(即同时计算 max 和 argmax)，这一特性非常实用。不过在当前场景中，我们仅需要最大值，因此取返回结果的第一个元素(最大值)：

        next_state_values[dones_t] = 0.0

如果批次中的状态转换是回合的最后一步，那么该动作的价值将不包含下一状态的折扣奖励，因为终止状态之后不存在可获取奖励的后续状态。这个细节看似微小，但在实际训练中极其关键------忽略这一点会导致模型无法收敛。

(10) 接下来，使用 detach() 方法将目标值从其计算图中分离，防止梯度流向用于计算下一状态Q值近似值的神经网络：

        next_state_values = next_state_values.detach()

这一操作至关重要。若不进行分离，损失值的反向传播将同时影响当前状态和下一状态的预测值。然而根据贝尔曼方程的要求，我们不应调整下一状态的预测值------它们仅用于计算目标Q值的参考基准。通过调用张量的 detach() 方法，我们切断了该张量与计算历史的关联，确保梯度不会流入计算图的这个分支。

(11) 最后，计算贝尔曼近似值和均方误差损失：

    expected_state_action_values = next_state_values * GAMMA + rewards_t
    return nn.MSELoss()(state_action_values, expected_state_action_values)

为全面理解损失函数的计算逻辑，完整查看 calc_loss() 函数的实现代码：

def calc_loss(batch: tt.List[Experience], net: dqn_model.DQN, tgt_net: dqn_model.DQN,
              device: torch.device) -> torch.Tensor:
    states_t, actions_t, rewards_t, dones_t, new_states_t = batch_to_tensors(batch, device)

    state_action_values = net(states_t).gather(
        1, actions_t.unsqueeze(-1)
    ).squeeze(-1)
    with torch.no_grad():
        next_state_values = tgt_net(new_states_t).max(1)[0]
        next_state_values[dones_t] = 0.0
        next_state_values = next_state_values.detach()

    expected_state_action_values = next_state_values * GAMMA + rewards_t
    return nn.MSELoss()(state_action_values, expected_state_action_values)

这就结束了我们的损失函数计算。

(12) 完成损失函数的计算流程，接下来实现训练循环。首先创建命令行参数解析器：

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--dev", default="cpu", help="Device name, default=cpu")
    parser.add_argument("--env", default=DEFAULT_ENV_NAME,
                        help="Name of the environment, default=" + DEFAULT_ENV_NAME)
    args = parser.parse_args()
    device = torch.device(args.dev)

支持指定计算设备(如 GPU / CPU )并允许训练非默认配置的环境。

(13) 初始化训练环境：

    env = wrappers.make_env(args.env)
    net = dqn_model.DQN(env.observation_space.shape, env.action_space.n).to(device)
    tgt_net = dqn_model.DQN(env.observation_space.shape, env.action_space.n).to(device)

该环境已加载所有必要的包装器 (wrapper)，包含待训练的主神经网络和结构相同的目标网络。虽然两者初始随机权重不同，但这并不影响效果------我们将每 1000 帧(约对应一局 Pong 游戏)同步一次网络参数。

(14) 随后创建指定容量的经验回放缓冲区，并将其传递给智能体：

    writer = SummaryWriter(comment="-" + args.env)
    print(net)

    buffer = ExperienceBuffer(REPLAY_SIZE)
    agent = Agent(env, buffer)
    epsilon = EPSILON_START

(15) epsilon 初始值设为 1.0，但会在每次迭代中逐步衰减：

    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE)
    total_rewards = []
    frame_idx = 0
    ts_frame = 0
    ts = time.time()
    best_m_reward = None

创建优化器、用于存储完整回合奖励的缓冲区、帧计数器，以及若干用于追踪训练速度和历史最佳平均奖励的变量。每当平均奖励突破记录时，将模型保存至文件中。

(16) 训练循环开始时，首先计算已完成的迭代次数，并按照预定策略衰减 epsilon 值：

    while True:
        frame_idx += 1
        epsilon = max(EPSILON_FINAL, EPSILON_START - frame_idx / EPSILON_DECAY_LAST_FRAME)

在指定帧数范围内( EPSILON_DECAY_LAST_FRAME=150000)，epsilon 将线性衰减，最终稳定在 EPSILON_FINAL=0.01。

(17) 指示智能体在环境中执行单步动作(使用当前网络和 epsilon 值)：

        reward = agent.play_step(net, device, epsilon)
        if reward is not None:
            total_rewards.append(reward)
            speed = (frame_idx - ts_frame) / (time.time() - ts)
            ts_frame = frame_idx
            ts = time.time()
            m_reward = np.mean(total_rewards[-100:])
            print(f"{frame_idx}: done {len(total_rewards)} games, reward {m_reward:.3f}, "
                  f"eps {epsilon:.2f}, speed {speed:.2f} f/s")
            writer.add_scalar("epsilon", epsilon, frame_idx)
            writer.add_scalar("speed", speed, frame_idx)
            writer.add_scalar("reward_100", m_reward, frame_idx)
            writer.add_scalar("reward", reward, frame_idx)

仅在遇到回合终止步骤时返回浮点奖励值。在这种情况下，我们会展示训练进度。具体来说，计算以下通过控制台和 TensorBoard 展示的指标：

• 处理速度(帧/秒)
• 已完成回合数
• 最近 100 回合平均奖励
• 当前 epsilon 值

(18) 每当最近 100 回合平均奖励突破历史记录时，会展示结果并保存模型参数：

            if best_m_reward is None or best_m_reward < m_reward:
                torch.save(net.state_dict(), args.env + "-best_%.0f.dat" % m_reward)
                if best_m_reward is not None:
                    print(f"Best reward updated {best_m_reward:.3f} -> {m_reward:.3f}")
                best_m_reward = m_reward
            if m_reward > MEAN_REWARD_BOUND:
                print("Solved in %d frames!" % frame_idx)
                break

若平均奖励超过预设阈值(例如 Pong 设定为 19.0，意味着 21 局游戏中获胜超过 19 局)，则终止训练。

(19) 检查缓冲区训练数据是否充足：

        if len(buffer) < REPLAY_START_SIZE:
            continue
        if frame_idx % SYNC_TARGET_FRAMES == 0:
            tgt_net.load_state_dict(net.state_dict())

首先需要等待足够的数据积累，在本节中为 10000 条转移样本，其次每完成 SYNC_TARGET_FRAMES (默认 1000 帧)同步一次主网络参数到目标网络。

(20) 训练循环的最后将梯度清零，从经验回放缓冲区采样数据批次、计算损失值，最后执行优化步骤以最小化损失：

        optimizer.zero_grad()
        batch = buffer.sample(BATCH_SIZE)
        loss_t = calc_loss(batch, net, tgt_net, device)
        loss_t.backward()
        optimizer.step()
    writer.close()

3.4 模型性能

在 Pong 上，模型需要约 40 万帧才能达到平均奖励 17 分(即胜率超过 80%)。从 17 分提升至 19 分仍需相近的训练帧数，因为此时学习进度会进入平台期，模型需要精细调整策略才能继续提高得分。因此完整训练通常需要约 100 万游戏帧。值得注意的是，这只是针对相对简单的 Pong 游戏。其他游戏可能需要数亿帧的训练数据，以及扩大 100 倍的经验回放缓冲区。

训练过程中奖励动态如下所示：

训练过程的控制台输出如下所示：

在前 1 万帧训练阶段，由于未执行任何模型训练(代码中最耗时的操作)，处理速度会保持极高水平。当超过 1 万帧后开始采样训练批次时，性能指标会回落至更具代表性的数值范围。值得注意的是，随着训练进行，系统性能还会因 epsilon 衰减而出现轻微下降------当 epsilon 值较高时，动作选择完全随机；而当 epsilon 趋近于 0 时，每次动作选择都需要执行神经网络推理计算Q值，这会额外增加时间开销。

经过数十局游戏训练后，DQN 模型通常会开始掌握获胜策略(此时 epsilon 约降至 0.5)，能够在 21 局比赛中赢得 1-2 局：

最终经过大量对局训练，DQN 能完全压制内置的 Pong AI 对手：

由于训练过程中存在随机性，实际训练动态可能存在差异。约 10% 的情况会出现训练完全不收敛的现象，表现为长期保持 -21 分的奖励值。这种情形在深度学习中并不罕见(源于训练随机性)，而在强化学习中更为常见(叠加了环境交互的随机性)。如果模型在前 10-20 万次迭代后仍未呈现正向提升趋势，建议重启训练。

4. 模型实战演示

训练过程仅是整体工作的一部分。我们的终极目标不仅是训练模型，更要让模型能出色地完成游戏对战。训练期间，每当最近 100 局平均奖励刷新记录时，系统会自动将模型保存为 PongNoFrameskip-v4-best_<score>.dat 文件。dqn_play.py 程序可加载该模型文件并进行单局对战演示，直观展示模型表现。

(1) 首先，导入 PyTorch 和 Gym 模块：

import gymnasium as gym
from ale_py import ALEInterface
import argparse
import numpy as np
import typing as tt

import torch

import wrappers
import dqn_model

import collections

DEFAULT_ENV_NAME = "PongNoFrameskip-v4"

(2) 脚本接收保存的模型文件名参数，并支持指定 Gymnasium 环境(需确保模型与环境匹配)：

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-m", "--model", required=True, help="Model file to load")
    parser.add_argument("-e", "--env", default=DEFAULT_ENV_NAME,
                        help="Environment name to use, default=" + DEFAULT_ENV_NAME)
    parser.add_argument("-r", "--record", required=True, help="Directory for video")
    args = parser.parse_args()

(3) 此外，需要通过 -r 参数指定一个不存在的目录名，该目录将用于保存游戏过程的录像：

    env = wrappers.make_env(args.env, render_mode="rgb_array")
    env = gym.wrappers.RecordVideo(env, video_folder=args.record)
    net = dqn_model.DQN(env.observation_space.shape, env.action_space.n)
    state = torch.load(args.model, map_location=lambda stg, _: stg, weights_only=True)
    net.load_state_dict(state)

    state, _ = env.reset()
    total_reward = 0.0
    c: tt.Dict[int, int] = collections.Counter()

创建环境并用 RecordVideo 包装器包装它，初始化模型后，从参数指定的文件加载权重。传递给 torch.load() 的 map_location 参数用于将加载的张量从 GPU 映射到 CPU。默认情况下，PyTorch 会尝试在保存时的设备上加载张量，但如果将训练好的模型从带 GPU 的设备迁移到无 GPU 的设备上，就需要重新映射设备位置。本节并未使用 GPU，因为推理过程即使不加速也足够快。

(4) 以下代码基本是训练代码中 Agent 类的 play_step() 方法的副本，只是移除了ε-贪婪动作选择策略：

    while True:
        state_v = torch.tensor(np.expand_dims(state, 0))
        q_vals = net(state_v).data.numpy()[0]
        action = int(np.argmax(q_vals))
        c[action] += 1

(5) 仅将观测数据输入智能体，并选择具有最大Q值的动作：

        state, reward, is_done, is_trunc, _ = env.step(action)
        total_reward += reward
        if is_done or is_trunc:
            break
    print("Total reward: %.2f" % total_reward)
    print("Action counts:", c)
    env.close()

(6) 将动作传递给环境，累计总奖励，并在回合结束时终止循环。回合结束后，程序会显示总得分及智能体执行动作的次数。通过以下命令运行程序：

$ python dqn_play.py -m PongNoFrameskip-v4-best_-13.dat -r ./video

如果希望探索本节内容，也可以尝试使用 Atari 系列中的其他游戏，如 Breakout、Atlantis 或 River Raid，以提升实践能力。

小结

在本节中，我们深入讨论了如何利用神经网络近似Q值，以及这种近似方法带来的额外复杂性。在深度Q网络 (Deep Q Network, DQN) 部分，还介绍了几种提高 DQN 训练稳定性和收敛性的技术，包括经验回放缓冲池、目标网络和帧堆叠机制。最终，我们将这些扩展技术整合为一个完整的 DQN 实现方案，成功解决了 Atari 游戏系列中的 Pong 环境。

系列链接

PyTorch强化学习实战（1）------强化学习（Reinforcement Learning，RL）详解
PyTorch强化学习实战（2）------强化学习环境库Gymnasium
PyTorch强化学习实战（3）------Gymnasium API扩展功能
PyTorch强化学习实战（4）------PyTorch基础
PyTorch强化学习实战（5）------PyTorch Ignite 事件驱动机制与实践
PyTorch强化学习实战（6）------交叉熵方法详解与实现
PyTorch强化学习实战（7）------表格学习与贝尔曼方程
PyTorch强化学习实战（8）------Q学习详解与实现