跑通一段强化学习代码最难的地方往往不是算法本身而是环境装不上、版本对不上、训练半天 reward 纹丝不动。这篇文章我写的就是用 tianshou OpenAI Gym现在社区普遍基于维护分支 gymnasium在雅达利游戏环境上从零训练 DQN 模型的完整过程配上可以直接复制的训练代码。适合刚接触强化学习不久、想跑通一个标准图像输入任务的同学也适合那些已经在别的框架上写过 DQN、想看看 tianshou 这套工程化组件怎么用的人。我会把环境包装、版本兼容、训练器配置、训练中怎么判断异常这些细节全部摊开讲而不是只丢代码。1. 为什么选 tianshou 来跑雅达利库之间的取舍1.1 tianshou 相比手写训练循环和 SB3 的差异早几年我自己第一次实践 DQN 在 Atari 上的效果时是纯 PyTorch 手写的自己实现经验回放、自己写 target network 同步、自己处理多环境采样、自己写 epsilon 衰减……代码量轻松破一千行而且 debug 相当痛苦。等换到 tianshou 之后最大的感受是它把强化学习训练里的“基础设施”全部抽象成了可替换组件。tianshou 的核心设计是 policy、collector、trainer、replay buffer 互相独立。你可以在完全不改数据收集逻辑的前提下把 DQNPolicy 换成 RainbowPolicy 或 PPO只要接口对得上。这种组件化设计在调实验时极其舒服。相比之下 Stable-Baselines3 封装得很高开箱即用、API 简洁但一旦你要改一个内部细节经常需要绕开高层接口去动源码手写训练循环则灵活度最高但代价是重复造轮子和引入隐藏 bug 的风险。我做了一张简单的对比表帮你快速判断自己适合哪个方案方案学习价值开发速度定制灵活度实验管理纯 PyTorch 手写最高低最高低Stable-Baselines3中高中中tianshou高高高高这里说的实验管理指的是日志记录、模型保存、评估流程这些容易被忽视但实际每天都要用的东西。tianshou 内置 Logger 抽象和 Trainer 的 checkpoint 流程比大多数手写方案规范得多。1.2 这套组合到底解决了什么问题Gymnasiums 提供了统一的env.step(action)接口让同一套算法代码可以在 CartPole、Atari、Mujoco 之间无缝切换。雅达利环境则是强化学习领域最有代表性的图像输入基准输入是原始像素帧模型必须自己从图像中学习状态特征这对验证一个算法的真正能力很有说服力。你在一张 84x84 灰度图上练出来的技巧迁移到真实视觉任务时多少是有帮助的。tianshou 在中间层补上了数据收集、经验回放、策略更新调度、多环境并行这些工程细节。换句话说你自己只负责三件事定义网络结构、选择策略类、配置训练器参数。其余繁琐的事tianshou 帮你包办了。这也是我建议你在有一定 RL 基础之后直接用 tianshou 跑实验的原因——它能让你把时间花在“算法本身的改进”上而不是“今天 buffer 越界了没”。2. 环境安装与版本兼容一半的坑都在这里2.1 依赖版本怎么选才不会装完就报错tianshou、gymnasium、torch 这三者之间的版本关系一度是我见过的最容易出问题的地方。老教程里一堆代码还在用gym这个名字而新生态已经拆分成了gymnasium。tianshou 在 1.0 之后明确支持 gymnasium我建议你直接用新命名空间不要再碰老 gym。我的环境参考如下组件推荐版本说明Python3.10 / 3.11太老或太新都可能遇到轮子缺失PyTorch 2.0按你的 CUDA 版本装对应 wheelgymnasium 0.29用import gymnasium as gymale-py0.8 左右gymnasium[atari] 会自动带tianshou1.0API 比 0.4.x 清爽很多安装命令大致是这样conda create -n atari python3.10 conda activate atari # PyTorch 请按实际 CUDA 版本从官网拿命令 pip install torch # atari 相关会自动安装 ale-py、autorum 等 pip install gymnasium[atari] # tianshou 本体 pip install tianshou一个很关键的提醒不要在一个环境里同时装 gym 和 gymnasium否则import gym的指向可能忽老忽新Classic control 环境不报错Atari 环境就莫名奇妙炸了。我以前碰过好几次这种“玄学问题”最后发现是gym残留包在捣乱。2.2 ROM 缺失环境能 import一创建就报错这是 Atari 玩家必踩的一道坎。Atari 游戏本身是 ROM 文件gymnasium 和 ale-py 只负责加载 ROM 的逻辑不会把游戏数据包塞进安装包里。第一次创建环境时几乎都会遇到类似FileNotFoundError: No such file or directory: Pong.bin的报错。解决办法是用 AutoROM 下载 ROM 并接受许可协议python -m AutoROM --accept-license如果这条命令在你的环境里找不到入口可以试试AutoROM --accept-license等它跑完会自动把 ROM 文件放到 ale-py 寻找的目录里。之后重新执行环境创建脚本ALE/PongNoFrameskip-v5就能正常加载了。有个小经验如果你所在网络环境访问 ROM 源不稳定也可以手动把对应 ROM 文件放进ale-py/roms目录文件名要求比较严格最好从官方列表对照着放。这一块我不展开最省心的办法还是先跑一次 AutoROM大概率能成功。2.3 快速验证环境是否装好写任何训练代码之前先跑下面这段确认环境创建和接口返回都正常import gymnasium as gym env gym.make(ALE/PongNoFrameskip-v5, repeat_action_probability0) obs, info env.reset() print(obs.shape, obs.dtype) # 随机执行两个动作看看是否正常 for _ in range(2): action env.action_space.sample() obs, reward, terminated, truncated, info env.step(action) print(reward, terminated, truncated) env.close()如果能看到(210, 160, 3) uint8这类输出说明环境载体已经通了。接下来才进入预处理包装阶段。这里还要区分一下 v4 和 v5 语法。v5 是 gymnasium 维护的新版本构造环境时可以直接传frameskip、repeat_action_probability参数行为方式更贴近原版 ALE。老 v4 时代很多配置要靠 wrapper 手工补代码风格也是gym.make(PongNoFrameskip-v4)这种。本文统一用 v5因为 tianshou 1.x 的官方示例就是朝向 gymnasium 这套生态设计的。3. Atari 环境包装管线预处理是训练的生死线3.1 为什么不能拿原始画面直接训练如果你直接把 210x160 的 RGB 三通道图像丢给 DQN训练能跑但效率和性能都会很差。原始画面有大量无关信息比如屏幕背景、记分板这些对决策的贡献很低却成倍增加了网络的参数量和计算量。更关键的问题是 Atari 游戏的帧闪烁flickering。很多对象不是每帧都绘制而是隔帧出现的单帧画面的信息是不完整的。原始 DQN 论文里最经典的方案是使用“结束帧池化”在一个动作持续的连续两帧里取逐像素最大值把两帧的闪烁内容合并到一张图上。这个逻辑写起来不复杂但没必要自己造轮子。此外单帧静止画面无法体现运动趋势比如球往哪个方向飞、速度多快所以需要堆叠最近 N 帧作为观测。综合下来标准预处理也就三件事转灰度、裁剪缩放到 84x84、堆叠最近 4 帧。灰度可以直接把三通道压缩成一通道缩放让网络输入规模大幅下降堆叠 4 帧给模型提供时间维度上的运动信息。3.2 推荐包装组合AtariPreprocessing FrameStackgymnasium 自带的AtariPreprocessing已经把这些逻辑都打包了。注意这里有一个很多人会搞混的点处理帧跳转frame skipping应当在 AtariPreprocessing 里做所以构造底层环境时最好选NoFrameskip 变体。我的环境工厂函数如下import gymnasium as gym from gymnasium.wrappers import AtariPreprocessing, FrameStack def make_env(env_idALE/PongNoFrameskip-v5): # 底层环境不跳帧 env gym.make(env_id, repeat_action_probability0) # 跳帧、灰度、缩放、最大池化合并闪烁帧都在这个 wrapper 里完成 env AtariPreprocessing( env, noop_max30, # reset 后先随机执行 0~30 次无操作打破初始状态确定性 frame_skip4, # 每个动作持续 4 帧 screen_size84, # 缩放至 84x84 terminal_on_life_lossTrue, # 丢一条命就算一个 episode 结束 grayscale_obsTrue, # 灰度化形状从 3 通道变 1 通道 grayscale_newaxisTrue, # 把灰度图像放到第一维形成 (1, 84, 84) scale_obsTrue, # 像素值从 0~255 缩放到 0.0~1.0 ) # 堆叠最近 4 帧输出形状 (4, 84, 84) env FrameStack(env, num_stack4) return env重点解释几个参数。noop_max30的根源是原始 DQN 论文里的无操作预热重置环境后随机执行 0 到 30 次 NOOP让同一个游戏的开局不止一种画面削弱确定性对训练的影响。terminal_on_life_lossTrue表示只要游戏里的一条命丢了就把这一局当作终止状态处理。这么做的好处是能提供更密集的终止信号让模型更早学会回避死亡坏处是它会改变环境的最优价值函数语义因为“一条命丢掉就停止”不等价于“整个游戏结束”。对 Pong 这种一次性掉分场景来说体验很好所以我一直默认开启。3.3 为什么推荐 NoFrameskip 环境如果你贪图方便直接gym.make(ALE/Pong-v5, frameskip4)再套一个AtariPreprocessing(frame_skip4)那就会发生双重帧跳转底层 ALE 已经把一个动作维持 4 帧上层 wrapper 又把每帧再重复执行 4 次实际效果是一个动作被拉长到 16 帧。这会让环境时序严重变慢模型看到的“球速”也和真实游戏对不上训练出来的策略往往很迟钝。最干净的做法就是选 NoFrameskip 变体让底层环境每步只前进 1 帧帧跳转完全交给上层的 AtariPreprocessing。这也是 DeepMind 原始 DQN 论文里的标准操作。现代 gymnasium 用ALE/PongNoFrameskip-v5或者PongNoFrameskip-v5都能拿到这个环境。tianshou 官方 Atari 示例的精神也是这个只是不同时期示例代码写了几十种风格消化起来容易晕。FrameStack 返回的LazyFrames对象在存储到 replay buffer 时不会立刻复制内存等到实际转 numpy 数组才展开对动辄几十万条 transition 的 buffer 来说内存友好不少。这也是我坚持用它的一个原因。4. 完整可运行代码DQN 在 Pong 上从零训练4.1 网络和策略定义背后的逻辑Atari 是图像输入网络主体自然是一组卷积层加全连接层。tianshou 的Net类比较聪明它会自动根据state_shape的维度决定要不要塞卷积层。我们把state_shape传(4, 84, 84)即可后续的 CNN 结构由 tianshou 内部自动推断注释里记得标明通道在前。DQN 策略的几个关键参数discount_factor0.99标准的未来奖励折现。estimation_step3n 步回报。值越大reward 传播越快但偏差也越大3 是很多 Atari 任务的稳妥选择。target_update_freq1000目标网络同步频率。太频繁会让目标 Q 值跟着当前网络一起漂移训练不稳定太慢又会拖慢学习。1000 是经验值实际可以按 gradient step 调。is_doubleTrue、is_duelingTrue本质上是把 Double DQN 和 Dueling DQN 直接作为开关打开。这两项对 Atari 的表现提升已经被大量实验验证过默认开就对了。elu 还是 relutianshou 默认激活函数选的是 ReLU配上scale_obsTrue后的 0~1 输入基本够用我没有刻意改。4.2 Collector 与 VectorReplayBuffer 怎么配合Collector 是 tianshou 的数据收集引擎。训练时它负责从 vector env 里采样动作、执行 step、把 transition 存入 buffer测试时它不做探索只看策略表现。多环境并行我习惯用 8 个训练环境、2 个测试环境。并行不是为了测速度好看而是保证经验回放里的数据多样性。如果只用单环境连续很长一段时间的 transition 都来自同一个轨迹片段样本相关性太高DQN 学起来容易震荡。VectorReplayBuffer 的分片设计就是匹配这种多环境的每个环境一个 buffer 分片写入毫无竞争采样时又能从全部分片里均匀抽取。如果你的机器是 WindowsShmemVectorEnv需要用共享内存可能会出现一些平台相关的权限问题。这时候直接退回到DummyVectorEnv即可虽然每个环境串行但数据收集管线整体仍然是异步的训练照样能进行。4.3 完整训练脚本 train_atari.py下面是完整的训练脚本直接保存运行即可。代码基于 tianshou 1.x 的 API 编写如果你用的是 0.4.xTrainer 的部分参数名和调用方式会有些区别建议以官方示例为准。 tianshou Gymnasium Atari DQN 训练脚本 目标环境: ALE/PongNoFrameskip-v5 import datetime import itertools import os import gymnasium as gym import torch from gymnasium.wrappers import AtariPreprocessing, FrameStack from tianshou.data import Collector, VectorReplayBuffer from tianshou.env import DummyVectorEnv, ShmemVectorEnv from tianshou.policy import DQNPolicy from tianshou.trainer import OffpolicyTrainer from tianshou.utils import TensorboardLogger from tianshou.utils.net.common import Net from tianshou.utils.net.discrete import DQN ENV_ID ALE/PongNoFrameskip-v5 SEED 42 NUM_TRAIN_ENVS 8 NUM_TEST_ENVS 2 # 给每个子环境一个递增 seed _seed_counter itertools.count() def make_env(): env gym.make(ENV_ID, repeat_action_probability0) # 首次 reset 时设置 seed保证不同子环境初始轨迹不同 env.reset(seedSEED next(_seed_counter)) env AtariPreprocessing( env, noop_max30, frame_skip4, screen_size84, terminal_on_life_lossTrue, grayscale_obsTrue, grayscale_newaxisTrue, scale_obsTrue, ) env FrameStack(env, num_stack4) return env def main(): device cuda if torch.cuda.is_available() else cpu print(fUsing device: {device}) # 先拿一个环境确认空间信息 probe_env make_env() state_shape probe_env.observation_space.shape action_shape probe_env.action_space.n probe_env.close() # 创建 vector envWindows 上如果 Shmem 报错就改 DummyVectorEnv train_envs ShmemVectorEnv([make_env for _ in range(NUM_TRAIN_ENVS)]) test_envs ShmemVectorEnv([make_env for _ in range(NUM_TEST_ENVS)]) # 网络Net 会自动为图像输入构建卷积层DQN 输出每个动作的 Q 值 net Net(state_shapestate_shape, action_shapeaction_shape, hidden_sizes[512], devicedevice) model DQN(net, action_shapeaction_shape, devicedevice) # DQN 策略默认开启 Double Q 和 Dueling policy DQNPolicy( modelmodel, optimtorch.optim.Adam(model.parameters(), lr1e-4), action_spaceprobe_env.action_space, discount_factor0.99, estimation_step3, target_update_freq1000, is_doubleTrue, is_duelingTrue, ) # 经验回放200k 条 transition对 Pong 来说够用 buffer VectorReplayBuffer(total_size200_000, buffer_numNUM_TRAIN_ENVS) # 训练 collector 带探索噪声测试 collector 纯确定性如果策略有 eps 的话 train_collector Collector(policy, train_envs, buffer, exploration_noiseTrue) test_collector Collector(policy, test_envs, exploration_noiseFalse) # 日志目录 log_path os.path.join( log, ENV_ID.split(/)[-1], datetime.datetime.now().strftime(%Y%m%d-%H%M%S), ) os.makedirs(log_path, exist_okTrue) logger TensorboardLogger(log_path) def save_best_fn(policy): torch.save(policy.state_dict(), os.path.join(log_path, best.pth)) def save_checkpoint_fn(epoch, env_step, gradient_step): torch.save( { model: policy.state_dict(), optim: policy.optim.state_dict(), epoch: epoch, env_step: env_step, gradient_step: gradient_step, }, os.path.join(log_path, checkpoint.pth), ) # 正式训练 result OffpolicyTrainer( policypolicy, train_collectortrain_collector, test_collectortest_collector, max_epoch100, step_per_epoch100_000, step_per_collect10, update_per_step0.1, episode_per_test10, batch_size256, save_best_fnsave_best_fn, save_checkpoint_fnsave_checkpoint_fn, loggerlogger, ).run() print(fbest_reward {result.best_reward:.2f}) if __name__ __main__: main()我解释几个容易让你在心里打问号的参数。step_per_collect10表示每收集 10 步环境转换就从中抽取 batch 更新网络。update_per_step0.1表示每收集 1 步平均执行 0.1 次梯度更新换算过来就是每 10 步样本做 1 次更新。这两个参数共同控制训练频率经验多余重、探索步数不足之间要取得平衡。你不需要背公式只需要记住想更新更频繁就把update_per_step调大。save_best_fn会在测试集上打到当前最好平均 reward 时保存策略权重save_checkpoint_fn则是每个 epoch 结束时保存完整状态方便中断恢复。这两个函数看起来简单实际上没有它们跑几小时后一旦进程崩了就是一堆白工。4.4 训练后的评估脚本 eval_atari.py训练结束之后我习惯单独写一个评估脚本加载 best 权重跑固定数量的 episode统计平均 rewardimport gymnasium as gym import torch from gymnasium.wrappers import AtariPreprocessing, FrameStack from tianshou.data import Collector from tianshou.env import DummyVectorEnv from tianshou.policy import DQNPolicy from tianshou.utils.net.common import Net from tianshou.utils.net.discrete import DQN ENV_ID ALE/PongNoFrameskip-v5 MODEL_PATH log/PongNoFrameskip-v5/xxx/best.pth def make_env(): env gym.make(ENV_ID, repeat_action_probability0) env AtariPreprocessing( env, noop_max30, frame_skip4, screen_size84, terminal_on_life_lossTrue, grayscale_obsTrue, grayscale_newaxisTrue, scale_obsTrue, ) env FrameStack(env, num_stack4) return env def main(): device cuda if torch.cuda.is_available() else cpu probe_env make_env() state_shape probe_env.observation_space.shape action_shape probe_env.action_space.n probe_env.close() net Net(state_shapestate_shape, action_shapeaction_shape, hidden_sizes[512], devicedevice) model DQN(net, action_shapeaction_shape, devicedevice) policy DQNPolicy( modelmodel, optimtorch.optim.Adam(model.parameters(), lr1e-4), action_spaceprobe_env.action_space, discount_factor0.99, estimation_step3, target_update_freq1000, is_doubleTrue, is_duelingTrue, ) # 加载权重并切换为评估模式 policy.load_state_dict(torch.load(MODEL_PATH, map_locationdevice)) policy.eval() test_envs DummyVectorEnv([make_env for _ in range(1)]) test_collector Collector(policy, test_envs, exploration_noiseFalse) result test_collector.collect(n_episode10) print(average reward:, result[episode_reward].mean()) if __name__ __main__: main()注意policy.eval()这一步不能省。tianshou 的 DQNPolicy 内部有eps_test和eps_train训练时 epsilon 会随时间衰减评估时如果不显式切到 eval 状态策略可能仍保留探索噪声导致测试分数偏低。5. 训练过程中那些“看上去正常其实有问题”的观察5.1 训练初期 reward 为负别急着关训练我几乎每次跑 Pong 都会看到前几十万步 reward 稳定在负值。Pong 的回合制度是 21 分打完训练早期模型连球都碰不到每局被电脑累计斩获 20 分左右平均回报自然低。很多第一次跑强化学习的人看到这个负回报就以为代码写错了。其实判断标准应该是“有没有上升趋势”而不是“绝对值够不够高”。我建议你打开 TensorBoard 看test/reward这类的指标曲线不要盯着零散的 Console 输出判断。一条虽然全在负值区域但持续爬升的曲线比一条在正负之间剧烈抖动的曲线更值得等待。5.2 Loss 下降和 reward 上升并没有必然联系DQN 的 TD loss 表示 Q 值预测和目标之间的误差。你可能会看到 loss 稳定下降但 reward 迟迟不涨或者 reward 涨了loss 反而变大。这都是正常的。原因是 Q 值本身没有一个绝对固定的尺度它随策略变化而变化。策略越好每条 state-action 的价值可能整体抬高模型需要重新拟合的目标也随之移动。所以只看 loss 绝对值没有意义。我实际训练中更关注三个指标平均回报曲线的趋势、平均 episode 长度、动作分布的熵。如果模型开始出现明显的偏好行为比如 Pong 里持续朝同一个方向移动说明它已经学到了一些东西后面就是慢慢打磨。5.3 三个高频异常的排查链路现象一loss 变 NaN。最常见的两个原因是学习率太大导致梯度爆炸或者 buffer 里混入了非法值。Pong 的 reward 只有 1/-1 之类出现非法值的概率很低所以优先怀疑梯度爆炸。解法是先把 lr 从 1e-4 降到 3e-5 试跑如果还炸再考虑在优化器上做梯度裁剪。tianshou 的 trainer 没有暴露特别顺手的 clip 参数我通常直接在策略外面包一个 callback 或者在 optimizer 上挂 hook但初学阶段调小 lr 是最快的路径。现象二跑了几百万步 reward 毫无变化。先用随机策略在同一环境下跑 10 个 episode看一下随机基准是多少。Pong 的随机策略一般也是被压着打但你能看到环境本身有正有负的 reward 信号。如果连随机策略都是清一色 0 或者一些怪异值那多半是环境配置问题比如动作语义理解错了。Pong 的动作集合里包含“右移”、“左移”、“右转”、“左转”这些动作DQN 得自己学会用哪个不需要你管但如果你的 action_space 拿错了比如把 NOOP 之外的执行方向搞混模型永远学不会接球。现象三训练慢到无法忍受。DQN 在 Atari 上的标配训练量是 10M 到 40M 环境步如果只用 1 个 CPU 环境跑一个晚上可能才几十万步什么都看不出来。建议至少上一张普通游戏显卡PyTorch 的 CNN 推理在 GPU 上收益巨大。如果没有 GPU就减少并行环境数、把 screen_size 从 84 降到 64 甚至 42先验证代码链路再放到大算力机器上跑正式实验。另一个很容易忽略的坑并行环境过多导致内存爆炸。VectorReplayBuffer 存的是原始图像帧虽然 FrameStack 用了 LazyFrames 做延迟展开但 20 万条 transition 每条约 4x84x84 字节整体量级并不小。如果你把 buffer 开到 100 万就要做好内存规划。Pong 这种简单任务根本不需要那么大规模。6. 踩坑清单与我的调参心得6.1 一张表总结高频问题现象可能原因解决方向创建 Atari 环境报 ROM 缺失ROM 未下载执行python -m AutoROM --accept-licenseimport gymnasium后环境创建报版本错环境中残留老gym卸载 gym重建虚拟环境模型动作像“慢动作”底层 frameskip 和 wrapper 双重跳帧改用 NoFrameskip 变体环境训练初期 reward 很负但不 NaN正常Pong 早期被打穿继续跑看趋势而非绝对值loss NaNlr 过大降低 lr观察是否恢复几百万步 reward 不涨学习率不合理或动作语义误解先跑随机策略看 reward 分布再确认 action_space测试分数低于训练中间观测评估时没有切policy.eval()加载权重后显式调用 evalWindows 下多环境卡住Shmem 权限问题改用 DummyVectorEnv6.2 我跑 Atari 实验时养成的几个习惯第一正式跑长训练之前先跑一个微型实验。把max_epoch改成 1step_per_epoch改成 2000确认整条 pipeline 从数据收集到模型保存都能走通然后直接用checkpoint.pth验证一下断点恢复。这比直接开着 100 epoch 跑一夜最后发现某个参数拼错导致白跑要省心得多。第二训练日志一定要记全。至少要有 tradeoff 指标、每多少步记录一次、checkpoint 的路径清晰可查。tianshou 的TensorboardLogger和WandBLogger都支持得很好我强烈建议选一个用。我自己的经验是如果只用print一晚上跑完你很难定位到“第几个 epoch 开始 reward 突然掉下去”这种关键转折点。第三关于超参Pong 这个环境其实相对宽容。lr1e-4、batch_size256、estimation_step3、target_update_freq1000是一组很稳的起点。如果你想更快看到效果可以把step_per_collect降到 5提高更新频率但别贪心更新太频繁会让数据相关性上升训练更容易震荡。第四最后说一个经常被忽略的小技巧多开几个独立的随机种子分别保存模型。不同 seed 下的策略分化很严重尤其是 DQN 这类 off-policy 算法同一套超参在不同 seed 下的最终表现可能差出一大截。最后汇报结果时取多 seed 的中位数而非最好值才是更诚实也更接近复现水平的评估方式。这一点在你后续写论文或跑对比实验时尤为关键。这个项目扩展的空间也很大把DQNPolicy换成RainbowPolicy把普通 buffer 换成 Prioritized ReplayBuffer甚至从 DQN 迁移到 PPO 都只需要改很有限的代码。等你熟悉了这套组合之后Atari 上再跑新算法就是改配置和调参的事而不是重写训练循环的事。