从零开始训练一个能打雅达利游戏的强化学习模型,整个过程比想象中有意思,也比想象中更容易卡住。我这次选的环境是Atari 2600上的经典游戏Breakout,算法框架用Tianshou,环境侧用OpenAI Gym生态的后续项目Gymnasium。目标很明确:模型从原始游戏画面开始,自己学会控制挡板、反弹小球、打掉砖块,而不是靠人工写规则。文章会给出可直接复现的完整代码,并重点说明那些“看起来能运行,但结果就是不对”的坑在哪里。
这个项目适合两类人:一类是已经跑过CartPole、但不知道下一步该怎么进阶的RL新手;另一类是知道DQN理论、但第一次想在高维图像输入上看到模型真实学习的开发者。整体代码量不大,但背后的数据流如果不理顺,跑出来的大概率只是随机撞墙的智能体。
1. 为什么选雅达利,以及这个项目要解决什么问题
1.1 雅达利在强化学习里的地位
雅达利2600游戏算是深度强化学习的“高考题”。2013年,DQN第一次做到只靠CNN从像素里学出策略,并且在多个Atari游戏上超过人类玩家,这是深度强化学习真正站稳脚跟的标志。此后的很多经典论文,比如Double DQN、Dueling DQN、Rainbow、分布式RL,都拿Atari环境做基准测试,原因很简单:它比CartPole这类玩具环境复杂得多,又比自动驾驶这类真实场景可控得多。
Atari环境有这么几个特性,决定了它特别适合做入门到进阶的桥梁。第一,输入是高维图像,但分辨率不高,84x84的灰度图就够用,普通消费级显卡完全扛得住。第二,动作空间是离散的,不需要处理连续控制里那些麻烦的探索问题。第三,游戏种类多,Pong简单,Breakout中等,Montezuma's Revenge这种难到能劝退一大半人的也有,你可以根据自己的算力选难度。
我再强调一次为什么选Breakout而不是Pong。Pong收敛非常快,快到你可能还没体会到训练调参的乐趣就已经结束了;Breakout则是“有明显学习曲线又不会等太久”的中间档,每个回合的分数量化清晰,模型打得怎么样一眼就能看出来。用它做教学演示,视觉反馈最好。
1.2 从原始像素到预处理的输入管线
很多人第一次接触Atari环境,直接拿原始画面丢给网络,结果训练半天loss不降、reward不涨。问题往往不在算法,而在预处理。实际项目里的标准做法,是把原始260x160的彩色游戏帧做四个处理:灰度化、缩放到84x84、四帧堆叠、外加帧跳过。
灰度化和缩放是为了降低网络输入的维度,本质上是把“彩色大图”压缩成“黑白小图”,模型需要的颜色信息在打砖块这种游戏里几乎用不到。四帧堆叠很好理解:单帧图像看不出速度,叠上最近四帧,网络就有了“球往哪个方向飞”的线索,这和你看视频需要连续画面而不是截图是一个道理。帧跳过则是每4个游戏帧只让模型做一次决策,减少决策频率,训练速度能提不少。
还有一个很容易被忽略的细节:开局时让环境随机执行最多30次无操作,这叫Noop Reset,目的是让每一局的起始画面不完全一样,避免模型把“开局第一帧”当成固定特征,从而过拟合死记硬背。这些预处理在Gymnasium里都有现成封装,不用自己写,但要把原理搞清楚,后面环境报错和训练不收敛的时候才知道去哪里找原因。
1.3 项目目标与运行预期
这个项目的基线目标,是在预设的帧预算内观察到得分曲线明显上升。具体来说,用200个epoch、每个epoch采2万个环境步去训练,在Breakout上大约相当于400万帧原始游戏画面。一台带中等GPU的机器跑下来,大概两三个小时。跑完以后,模型的平均回报应该从0附近涨到2以上,注意这个“2”不是真实游戏分数,而是奖励被裁剪后的数值,后面会说。
如果你只想知道流程是否跑通,建议先换Pong,环境名改成“Pong”,动作数不同但代码逻辑一模一样。Pong在大概50万帧内就能看到快速收敛,适合第一次跑通闭环。Breakout则更适合观察更长周期的学习曲线。
2. 环境与框架选型:Gymnasium、Tianshou 和各自的坑
2.1 OpenAI Gym 后的时代:Gymnasium 版本与 API 差异
这里必须先说一个很多老教程害人不浅的地方:环境库已经从OpenAI Gym升级成了Gymnasium。名字上像是同一个项目,但API有破坏性改动。最典型的是环境重置这个函数,旧版写法是obs = env.reset(),新版必须写成obs, info = env.reset(),少解包一个返回值就会直接报错。
更麻烦的是环境名。旧版Atari环境叫Breakout-v0、Breakout-v4,新版建议用ALE/Breakout-v5这种命名。如果你照着一两年前的博客、视频去跑,大概率会卡在环境名不识别或者reset参数不匹配上面。
所以我的建议是:不要再用已经停止维护的旧gym,直接上Gymnasium。安装命令如下,Python版本建议3.10,PyTorch版本选2.x就好。
conda create -n rl python=3.10 -y conda activate rl pip install "gymnasium[atari]"==0.29.1 pip install tianshou==1.0.0 pip install torch --index-url https://download.pytorch.org/whl/cu118如果你机器上没有NVIDIA显卡,把最后一行换成pip install torch就行,CPU版也能跑,只是训练慢一些。这里把gymnasium[atari]这个方括号里的atari额外项写上,它会顺便装上Atari的ROM包和相关依赖。否则你后面创建环境的时候可能报“ROM not found”,那时候再补齐就多绕一段路。
2.2 Tianshou的模块化设计好在哪
选Tianshou而不选其它强化学习库,主要理由是它的模块切分特别清楚。Policy负责根据观测给动作、算梯度;Collector负责和环境交互、采集数据;ReplayBuffer负责存历史transition;Trainer负责把上面这些串起来跑循环。每一块都可以独立替换。
和Stable-Baselines3比,Tianshou更偏研究向,自定义网络、自定义采样逻辑都更容易,而且底层是PyTorch,改起来顺手。和RLlib比,Tianshou轻量很多,不用搭Ray集群,单机调试时心智负担小。缺点也有,就是版本迭代快,API有调整,很多网上教程代码跑不起来。这篇文里我会尽量给出更稳定的写法,并且说明哪些参数在个别版本里可能改名。
具体到DQN,Tianshou把Double DQN、优先经验回放、Dueling这些常见改版都做成了参数,大部分时候不需要自己实现,这非常适合快速验证想法。我第一次在Tianshou里从DQN切到PPO,只改了policy那一段,其余Collector、Buffer、Trainer全部复用,这个体验在别的框架里一般没这么顺。
2.3 安装与环境自检
安装完之后,先做一个最小化自检,确认Atari环境真的能跑。下面这段代码先创建环境,打印观测空间和动作空间:
import gymnasium as gym env = gym.make("ALE/Breakout-v5", frameskip=1) obs, info = env.reset() print("观测空间:", env.observation_space) print("动作空间:", env.action_space) print("初始观测形状:", obs.shape) env.close()正常情况下,动作空间会显示Discrete(4),意思是Breakout有4个可用动作,分别是无操作、开球、右移、左移。这里有一个坑,我们后面会说怎么处理。
如果这一步就报错,九成是ROM或版本问题。先说版本:gymnasium和ale-py版本不匹配时,可能会在env.reset()报奇怪解包错误,把gymnasium升到0.29以上基本能解决。再说ROM:如果报错信息里出现“ROM not found”或者“Invalid ROM”,执行下面这条命令手动安装ROM:
python -m autorom.install这条命令会从官方源下载全部Atari游戏ROM。下载过程中可能提示接受许可协议,确认一下就好。这一步做完,重跑自检代码,环境一般就通了。
3. 策略网络与DQN算法设计
3.1 CNN网络结构与维度计算
Breakout的输入是84x84的灰度图,一次堆叠4帧,所以网络输入形状是4x84x84。这个任务用三层卷积加两层全连接就能打得很不错,结构其实非常经典,常被叫做Nature CNN。
我用的网络结构是这样的:
| 层 | 参数 | 输出维度 |
|---|---|---|
| Conv1 | 32个8x8卷积核,步长4 | 32x20x20 |
| ReLU | - | 32x20x20 |
| Conv2 | 64个4x4卷积核,步长2 | 64x9x9 |
| ReLU | - | 64x9x9 |
| Conv3 | 64个3x3卷积核,步长1 | 64x7x7 |
| ReLU | - | 64x7x7 |
| Flatten | - | 3136 |
| FC1 | 512单元 | 512 |
| FC2 | 动作数 | 4 |
第三个卷积的输出为什么是7x7?输入84x84,经过8x8步长4的卷积,尺寸变成(84-8)/4+1=20;再过4x4步长2,变成(20-4)/2+1=9;最后3x3步长1,变成(9-3)/1+1=7。全连接层的输入就是64x7x7的展平结果,也就是3136维。
一个特别容易被忽略的地方是维度顺序。Gymnasium的FrameStack默认把堆叠帧放在最后一维,也就是说观测形状是84x84x4,而PyTorch的卷积层需要4x84x84这种通道在前排列。很多新手直接把环境给的obs扔给网络,在Conv2d那里报维度错误,就是没做permute。我会在代码的forward里显式做转换。
3.2 DQN的关键机制在Tianshou里怎么落地
DQN能稳定工作,靠三个机制:经验回放、目标网络、以及探索策略。经验回放把历史transition存进buffer,训练时随机抽一批,打破样本之间的时序相关性,这一点解决的是“连续帧高度相似,导致梯度方向不稳定”的问题。目标网络则是一个延迟更新的副本,用来计算TD目标中的下一状态Q值,避免“用自己更新自己”带来的训练震荡。
在Tianshou里,经验回放和训练调度由VectorReplayBuffer和Collector负责,你不用手写采样循环;目标网络的更新频率则直接由DQNPolicy的target_update_freq参数控制。这个参数的含义是每隔多少次梯度更新,就把目标网络参数同步一次。一般来说,1000左右是比较常用且稳妥的选择,太大了目标值长期不动,太小了又退化成普通Q学习。
探索策略方面,DQN用epsilon-greedy:大部分时候选最大Q值的动作,但有epsilon的概率随机选动作。训练开始阶段epsilon大,强迫模型多探索世界;训练后期epsilon小,让模型更多利用已经学到的策略。文章后面的代码里我会做一个阶段退火:前5万个环境步用epsilon=1.0,相当于完全随机探索,之后降到0.1。
3.3 超参数选择背后的权衡
几个关键超参数,我按“改它会影响什么”的角度一起说清楚。
学习率1e-4是Atari领域久经验证的选择。图像输入的问题在于梯度规模天然偏大,学习率太高网络很容易发散,表现为某个epoch后reward突然跌穿底部;调低到1e-4或5e-5之后反而更快。
折扣因子0.99是标准配置,几乎不用改。Buffer大小我设为20万条transition,在普通机器上这是一个内存和效果都不错的折中。缓冲区太小,采样的多样性就低,模型容易反复看到最近的片段;太大则内存开销高,而且并不一定有更大的收益。
一个容易被误解的参数是step_per_collect。它表示每采集多少环境步就做一轮梯度更新。我代码里设200,也就是每个epoch大约100次更新。如果你发现训练速度很慢、采样和更新比例失衡,可以先调整这个值。注意Tianshou里的“step”指的是环境交互步,不是原始游戏帧,在帧跳过设成4的情况下,一个环境步实际对应4个游戏帧。
4. 完整代码:从环境包装到训练闭环
4.1 环境构造与网络定义
下面的代码是整个项目的地基。make_env负责把一个原始的ALE环境包装成适合RL训练的格式。训练用的环境会裁剪奖励和提前终止,评估用的环境则展示真实得分,这个区别后面会专门讲。
import numpy as np import torch import gymnasium as gym from gymnasium.wrappers import AtariPreprocessing, FrameStack, TransformReward from tianshou.data import Collector, VectorReplayBuffer from tianshou.env import DummyVectorEnv from tianshou.policy import DQNPolicy from tianshou.trainer import OffpolicyTrainer def make_env( game="Breakout", seed=0, clip_reward=True, terminal_on_life_loss=True, ): # 注意这里先设 frameskip=1,真正的帧跳过由 AtariPreprocessing 控制, # 避免两边重复执行 skip,导致动作重复次数翻倍。 env = gym.make(f"ALE/{game}-v5", frameskip=1, repeat_action_probability=0.0) env = AtariPreprocessing( env, noop_max=30, frame_skip=4, screen_size=84, grayscale_obs=True, terminal_on_life_loss=terminal_on_life_loss, ) env = FrameStack(env, num_stack=4) if clip_reward: env = TransformReward(env, lambda r: np.clip(r, -1.0, 1.0)) env.reset(seed=seed) return env class AtariNet(torch.nn.Module): def __init__(self, action_num, device="cpu"): super().__init__() self.device = device self.cnn = torch.nn.Sequential( torch.nn.Conv2d(4, 32, kernel_size=8, stride=4), torch.nn.ReLU(), torch.nn.Conv2d(32, 64, kernel_size=4, stride=2), torch.nn.ReLU(), torch.nn.Conv2d(64, 64, kernel_size=3, stride=1), torch.nn.ReLU(), ) self.fc = torch.nn.Sequential( torch.nn.Flatten(), torch.nn.Linear(64 * 7 * 7, 512), torch.nn.ReLU(), torch.nn.Linear(512, action_num), ) def forward(self, obs, state=None, info=None): if not isinstance(obs, torch.Tensor): obs = torch.as_tensor(np.asarray(obs), dtype=torch.float32) obs = obs.to(self.device) / 255.0 # 兼容 [4,84,84] 单样本和 [B,84,84,4] 批量两种情况 if obs.dim() == 3: obs = obs.unsqueeze(0) if obs.shape[-1] == 4: obs = obs.permute(0, 3, 1, 2) return self.fc(self.cnn(obs)), state环境构造器里有两个设计点值得注意。
第一个是frameskip=1。Gymnasium创建环境时如果直接传frameskip=4,环境底层已经做了帧跳过;而AtariPreprocessing的frame_skip=4也会做一次帧跳过。两个同时开,相当于每4帧才观察一次,但一跳跳了16帧,动作响应变得非常迟钝,训练效果直线下降。正确做法是创建环境时用frameskip=1,把跳过逻辑完全交给我们的wrapper。
第二个是repeat_action_probability=0.0。ALE环境默认有一定概率重复上一步动作,这是为了模拟真实机器的随机性,但对复现实验不友好,改成0.0以后环境行为更可控。如果你想更贴近真实游戏手感,保留默认0.25也行。
4.2 训练主程序与阶段退火
网络定义好之后,训练主程序的核心就三件事:建并行环境、建Policy、建Collector并交给Trainer。完整代码如下:
def main(): seed = 42 device = "cuda" if torch.cuda.is_available() else "cpu" train_envs = DummyVectorEnv( [lambda i=i: make_env("Breakout", seed=seed + i * 100) for i in range(8)] ) test_envs = DummyVectorEnv( [lambda i=i: make_env("Breakout", seed=seed + i * 100) for i in range(4)] ) env = make_env("Breakout", seed=seed) action_num = env.action_space.n env.close() model = AtariNet(action_num, device=device).to(device) optim = torch.optim.Adam(model.parameters(), lr=1e-4) policy = DQNPolicy( model, optim, discount_factor=0.99, target_update_freq=1000, is_double=True, ).to(device) buffer = VectorReplayBuffer(total_size=200000, buffer_num=len(train_envs)) train_collector = Collector(policy, train_envs, buffer, exploration_noise=True) test_collector = Collector(policy, test_envs, exploration_noise=False) def train_fn(epoch, env_step): if env_step < 50000: policy.set_eps(1.0) else: policy.set_eps(0.1) def test_fn(epoch, env_step): policy.set_eps(0.01) trainer = OffpolicyTrainer( policy=policy, train_collector=train_collector, test_collector=test_collector, max_epoch=200, step_per_epoch=20000, step_per_collect=200, update_per_step=1, episode_per_test=20, batch_size=32, train_fn=train_fn, test_fn=test_fn, ) result = trainer.run() torch.save(policy.state_dict(), "dqn_breakout.pt") print("best_reward:", result["best_reward"]) if __name__ == "__main__": main()这里的DummyVectorEnv有8个并行环境的接口,但别误会,它并不是多进程并行,环境仍然是同步串行执行的。它的作用是让底层数据流统一成批量接口,代码上更方便,也为以后换成SubprocVectorEnv留了口子。想真正多进程加速,把DummyVectorEnv换成SubprocVectorEnv即可,但随之而来的是状态同步和调试成本,建议先跑通主流程再去折腾。
train_fn里按环境步数做了一个探索率退火。前5万步用epsilon=1.0,意思是完全随机动作,让buffer里先积累足够多样化的数据;之后降到0.1,进入“以利用为主、保留少量探索”的阶段。这个阶段的探索率如果设成0,模型遇到没见过的情况就会不知所措,所以哪怕训练后期我也保留了10%的随机动作。
运行后,控制台会输出类似“epoch 10 reward: 0.30”之类的日志。如果你看到几个epoch后reward还在0附近徘徊,别急着下“模型不行”的结论,Breakout的奖励比较稀疏,刚开场只能打掉一两个砖块,clip之后可能只有0.1。继续观察几十个epoch,曲线有抬升趋势就是正常的。
4.3 评估脚本与权重保存
训练用的环境为了稳定收敛,做了奖励裁剪,而且开了terminal_on_life_loss,也就是只要丢一条命就判定本局结束。这样做训练效率高,但评估真实水平时不能用同一套环境,否则看到的分数只是“第一条命能打多少”。正确做法是关掉生命结束判定、不做奖励裁剪,跑完整局拿真实得分。
def evaluate(model, episodes=10, game="Breakout"): model.eval() env = make_env(game, seed=0, clip_reward=False, terminal_on_life_loss=False) scores = [] for _ in range(episodes): obs, _ = env.reset() total = 0.0 done = False while not done: with torch.no_grad(): q, _ = model(obs) action = int(q.argmax(dim=-1).item()) obs, reward, terminated, truncated, _ = env.step(action) total += reward done = terminated or truncated scores.append(total) env.close() return scores if __name__ == "__main__": model = AtariNet(4, device="cpu") model.load_state_dict(torch.load("dqn_breakout.pt", map_location="cpu")) print("真实得分:", evaluate(model, episodes=5))这里再提醒一次,加载权重时网络结构、动作数必须和训练时保持一致。如果你重命名了网络文件或者改了动作数,load_state_dict会直接报错,这个错误信息比弄错超参数友好得多。
5. 训练效果、常见问题与调参避坑
5.1 训练过程中的真实预期与日志解读
先说一个最容易劝退人的现象:训练一开始reward可能长时间在0附近。Breakout这款游戏,开局如果没有打中砖块,整个回合连一分都没有,clip之后奖励就是0,所以前几十个epoch的均值在0到0.1之间浮动,完全正常。
真正要关注的是实验中期以后的趋势。Tianshou的Trainer会打印每个epoch的测试平均reward,这里有个细节:episode_per_test=20表示测试阶段跑20个episode来平均。单局波动很大的话,可以加大这个数,比如50,测试结果会平滑很多,代价是测试时间变长。
我实际跑下来的经验是,Breakout在30个epoch左右能隐约看到reward超过0.3,80个epoch以后有明显抬升,200个epoch时测试均值通常在2到5之间。有些调参良好的配置能到10以上,但会需要更长训练时间和更精细的探索率调整。如果你跑到100个epoch还是全程0,先怀疑环境bug,不要怀疑算法原理。
5.2 高频问题速查表
下面这几条是我在这类项目里被问过最多、也最容易踩的情况,整理成速查表。
| 症状 | 可能原因 | 解决办法 |
|---|---|---|
env.reset()报“too many values to unpack” | 误用旧版gym教程代码 | 改成obs, info = env.reset() |
| 环境创建时报“ROM not found” | Atari ROM缺失 | 安装gymnasium[atari]或执行python -m autorom.install |
| 网络报 Conv2d 输入维度错误 | FrameStack把帧放在最后一维 | 在forward里加permute(0, 3, 1, 2) |
| 训练很久reward一直为0 | 探索不足、奖励裁剪后数值太小 | 确认前5万步epsilon是1.0,观察更久 |
| 测试得分很低,但训练得分还行 | 开启了terminal_on_life_loss | 评估环境时改成False |
| 第N个epoch突然崩溃,reward跌穿 | 学习率过大或buffer太小 | 学习率降到5e-5,buffer提到50万再试 |
里面我想单独展开说两点。
第一点是terminal_on_life_loss的双面性。训练环境设置成True,是让模型每丢一条命就重新开始,相当于把“活得更久”这个信息也作为训练信号,能显著加快学习。但代价是训练目标与最终目标不完全一致,所以最后评估真实实力时,一定要关掉它再跑完整局。用训练环境的分数当作最终成绩,会严重低估模型。
第二点是关于随机性的。Atari环境即使固定seed,每局的随机性依然很大,单次评估的得分波动可能非常明显。判断模型好坏,最少一次跑5到10个episode取平均。更进一步,把评估的随机seed也固定下来,这样不同模型之间的对比才公平。
5.3 再往深走一步:从DQN到Rainbow和并行环境
跑通这套DQN以后,最自然的进阶路线是往Rainbow方向走。Rainbow本质上就是DQN的多个改进叠加,包括Double DQN、优先经验回放、Dueling网络、多步回报、NoisyNet和分布式Q值。Tianshou对这些支持得比较好,优先经验回放只需要换一个ReplayBuffer,Duelling网络需要自己定义输出Q值和V值的网络结构,其它部分大多还是参数。
如果你主要想要更快的训练吞吐,下一步就是并行环境。把DummyVectorEnv换成SubprocVectorEnv,环境会跑在多个子进程里,并行度直接拉满,训练速度提升非常明显。代价是你需要留意环境的pickle序列化、坑更多、调试更麻烦。我的建议是先确保单个环境流程完全稳定,再上并行。
还有一个很实用的方向:把Breakout换成别的雅达利游戏验证泛化能力。因为代码里make_env第一个参数就是游戏名,换游戏只需要改动这一处。但注意不同游戏的动作数不一样,Pong是6个动作,Space Invaders是6个,动作数变了网络最后一层的输出维度也要跟着变。训练前把action_num打印出来确认一下,能避免很多低级报错。
最后再分享一点个人体会。第一次完整跑通这套流程时,我最强烈的感受是:强化学习项目的调试难度不在写代码,而在“你很难判断模型不工作到底是代码问题、超参数问题,还是单纯需要更长时间”。Tianshou这种框架把标准流程封装得很好,但你不能因此跳过因果链的理解。每次训练完,我都会把best_reward记录成日志,连同超参配置一起存下来,方便下一次对照。强化学习的宝藏在实验记录里,而不在直觉里。