强化学习实战：从零实现DQN玩转CartPole

简介面向强化学习初学者的乒乓球对战模拟项目提供可直接运行的代码与可视化界面免去复杂的环境配置通过直观的交互过程展示智能体如何根据状态、动作、奖励与策略不断优化决策。压缩包共包含十个文件具体包括三个Python脚本分别承担DQN算法实现、验证逻辑与环境封装六个.pth模型权重文件对应从初始到后期不同训练阶段的智能体参数另有一个.pyc缓存文件整体大小约三十五MB文件组织清晰便于逐一对照学习。目前已有四千余人学习或下载。运行核心脚本即可观察智能体从随机击球到策略逐渐稳定的完整训练过程替换不同的权重文件能快速对比中间模型与最终模型的行为差异。同时代码保留了算法与参数的修改空间可尝试将DQN替换为其他强化学习算法或调整奖励函数、探索率等超参数进一步理解各要素对训练效果的影响是一份兼具演示功能与实验价值的入门级资源。1. 为什么会写这个强化学习实例我最早接触强化学习的时候最烦的一件事就是书和教程看了不少什么马尔可夫决策过程、贝尔曼方程、策略梯度理论一套一套的但真要自己上手写个能跑的程序反而不知道从哪里开始。网上能搜到的代码要么封装太重一上来就是几千行的工程框架要么直接用现成库调接口跑完也不知道内部到底发生了什么。后来给团队做内部培训我定了一个原则代码必须能直接跑逻辑必须简单到能一行一行看懂算法必须能像换螺丝一样方便替换。这个实例就是在那段时间里整理出来的用的是强化学习入门最经典的CartPole环境加上一个结构足够清晰的DQN实现。整套代码没有复杂的依赖没有隐晦的缩写每个函数负责的事情都写在名字里适合刚接触强化学习的读者也适合想快速验证自己想法的研究者。标题里说的几个关键词其实是经过反复打磨之后才做到的直接运行pip安装依赖后复制粘贴就能看到训练曲线不需要修改任何路径和参数。直观每一步的决策过程、奖励变化、loss变化全部打印出来训练过程像看仪表盘一样透明。方法可替换算法实现与训练流程解耦换一个算法只需要替换一个类其他代码一律不动。初学者代码友好没有装饰器、没有元类、没有隐式类型转换变量命名完全展开注释给的也是“为什么这么做”而不是“这行代码做了什么”。下面我会从设计思路开始把整个代码的骨架、核心细节、实操过程以及我踩过的一些坑完整写出来。2. 整体设计与思路拆解2.1 环境选型为什么是CartPole而不是Atari很多新手上来就选Atari游戏或者MuJoCo机器人控制环境结果连环境都安装不好要么是ROM文件缺失要么是许可证问题。我选CartPole-v1有四个原因第一动作空间是离散的左移或右移不需要处理连续动作的采样和重参数化第二状态空间只有4个维度位置、速度、角度、角速度用一个两层的全连接网络就能拟合Q值第三单次训练收敛速度快CPU上几分钟就能看到明显效果不用等GPU第四失败信号清晰杆子倾斜超过15度或小车偏离中心超过2.4个单位就结束reward的语义非常容易理解适合用来演示“奖励信号如何引导智能体行为”。注意CartPole-v1比v0多了最大步数限制500步这意味着如果智能体表现足够好一个episode会因达到步数上限而结束而不是因失败结束。这个细节在判断收敛时很重要后面我会详细说。2.2 代码架构训练主循环与算法分离整个项目拆成三个文件train.py负责训练流程agent.py放算法实现config.py放所有超参数。reinforcement-learning-demo/ ├── train.py # 训练主循环环境交互日志输出 ├── agent.py # DQN算法实现可整体替换 └── config.py # 超参数配置为什么要这样拆因为强化学习的代码逻辑可以分为两部分和环境交互的流程与根据交互数据更新策略的方法。第一部分对几乎所有离散动作空间的算法都是一样的——观察状态、选择动作、执行动作、获得奖励、进入下一状态这个循环永远不会变。第二部分则完全不同DQN需要经验回放和目标网络SARSA需要在episode内更新Policy Gradient需要计算累积回报。如果这两部分耦合在一起换算法就得重写全部代码。拆开之后训练主循环只需要知道agent暴露了三个接口choose_action(state)、store_transition(...)、learn()至于内部是神经网络还是查表主循环完全不关心。这就是“方法可替换”的底层设计。3. 核心实现从零搭一个DQN3.1 超参数配置每个数字背后都有原因先看config.py的内容# config.py # 环境参数 ENV_NAME CartPole-v1 MAX_EPISODES 400 # 最大训练回合数 MAX_STEPS 500 # 每个回合最大步数 # 网络参数 HIDDEN_SIZE 128 # 隐藏层神经元数量 LEARNING_RATE 1e-3 # Adam优化器学习率 # 经验回放参数 MEMORY_CAPACITY 10000 # 经验池容量 BATCH_SIZE 64 # 每次采样训练样本数 # 探索与利用 EPSILON_START 1.0 # 初始探索率 EPSILON_END 0.01 # 最小探索率 EPSILON_DECAY 0.995 # 每回合衰减系数 # 目标网络 TARGET_UPDATE_FREQ 100 # 目标网络同步间隔步数 # 折扣因子 GAMMA 0.99 # 奖励折扣系数几个值得解释的选择HIDDEN_SIZE 128CartPole的状态只有4维128个神经元足够了。我以前试过256、512训练速度变慢但效果没有明显提升128是性价比最高的选择。EPSILON_DECAY 0.995每过一个episode探索率乘以0.995。400个episode后探索率约为0.01这正好让智能体在前期充分探索后期完全利用学到的策略。TARGET_UPDATE_FREQ 100目标网络每100步同步一次。这个值太大会导致目标Q值长期不变训练不稳定太小则失去了目标网络的意义。100是在实践中比较稳的默认值。GAMMA 0.99折扣因子接近1表示智能体比较“有远见”愿意为了长期回报而牺牲短期奖励。在CartPole这种单步奖励恒为1的任务里GAMMA越大智能体越倾向于尽可能延长episode。3.2 DQN Agent三层结构的完整实现接下来是核心文件agent.py我直接把完整代码放出来然后逐段拆解# agent.py import numpy as np import torch import torch.nn as nn import torch.optim as optim import random from collections import deque class QNetwork(nn.Module): 两层全连接Q网络 def __init__(self, state_dim, action_dim, hidden_size128): super(QNetwork, self).__init__() self.fc1 nn.Linear(state_dim, hidden_size) self.fc2 nn.Linear(hidden_size, hidden_size) self.fc3 nn.Linear(hidden_size, action_dim) self.relu nn.ReLU() def forward(self, x): x self.relu(self.fc1(x)) x self.relu(self.fc2(x)) return self.fc3(x) # 输出每个动作的Q值 class ReplayBuffer: 经验回放缓冲区存储并采样转移元组 def __init__(self, capacity): self.buffer deque(maxlencapacity) def push(self, state, action, reward, next_state, done): self.buffer.append((state, action, reward, next_state, done)) def sample(self, batch_size): batch random.sample(self.buffer, batch_size) states, actions, rewards, next_states, dones zip(*batch) return ( torch.FloatTensor(np.array(states)), torch.LongTensor(np.array(actions)).unsqueeze(1), torch.FloatTensor(np.array(rewards)).unsqueeze(1), torch.FloatTensor(np.array(next_states)), torch.FloatTensor(np.array(dones)).unsqueeze(1), ) def __len__(self): return len(self.buffer) class DQNAgent: DQN算法主体包含策略选择、经验存储和网络更新 def __init__(self, state_dim, action_dim, config): self.action_dim action_dim self.config config # 两个网络评估网络实时更新目标网络定期同步 self.eval_net QNetwork(state_dim, action_dim, config.HIDDEN_SIZE) self.target_net QNetwork(state_dim, action_dim, config.HIDDEN_SIZE) self.target_net.load_state_dict(self.eval_net.state_dict()) self.target_net.eval() self.optimizer optim.Adam(self.eval_net.parameters(), lrconfig.LEARNING_RATE) self.loss_fn nn.MSELoss() # 经验池和探索率 self.memory ReplayBuffer(config.MEMORY_CAPACITY) self.epsilon config.EPSILON_START # 记录总步数用于目标网络同步 self.total_steps 0 def choose_action(self, state): epsilon-greedy策略以epsilon概率随机探索否则选择最大Q值动作 if random.random() self.epsilon: return random.randrange(self.action_dim) state torch.FloatTensor(np.array(state)).unsqueeze(0) with torch.no_grad(): q_values self.eval_net(state) return q_values.argmax().item() def store_transition(self, state, action, reward, next_state, done): self.memory.push(state, action, reward, next_state, done) def decay_epsilon(self): 每个episode结束后衰减探索率 self.epsilon max(self.config.EPSILON_END, self.epsilon * self.config.EPSILON_DECAY) def learn(self): 从经验池采样一个批次更新评估网络 if len(self.memory) self.config.BATCH_SIZE: return # 经验不足时先不学习 states, actions, rewards, next_states, dones self.memory.sample(self.config.BATCH_SIZE) # 计算当前Q值根据实际采取的动作选取对应Q值 q_values self.eval_net(states).gather(1, actions) # 计算目标Q值r gamma * max(Q_target(s)) with torch.no_grad(): next_q_values self.target_net(next_states).max(1, keepdimTrue)[0] target_q_values rewards self.config.GAMMA * next_q_values * (1 - dones) # 计算loss并反向传播 loss self.loss_fn(q_values, target_q_values) self.optimizer.zero_grad() loss.backward() self.optimizer.step() # 定期同步目标网络 self.total_steps 1 if self.total_steps % self.config.TARGET_UPDATE_FREQ 0: self.target_net.load_state_dict(self.eval_net.state_dict()) def save(self, path): torch.save(self.eval_net.state_dict(), path) def load(self, path): self.eval_net.load_state_dict(torch.load(path)) self.target_net.load_state_dict(self.eval_net.state_dict())下面逐段解释几个容易被初学者忽略但又极其关键的地方第一处gather(1, actions)在做什么Q网络输出的是形状为(batch_size, action_dim)的张量表示每个状态下每个动作的Q值。我们需要的是“实际采取的那个动作的Q值”gather(1, actions)就是按actions中记录的索引从每个样本中挑出对应的Q值。比如actions是[[1], [0], [1]]就会从每行分别取索引为1、0、1的元素。这是PyTorch里实现DQN最标准的做法写成q_values[range(batch_size), actions]效果相同但gather表达更直观。第二处为什么target计算要套torch.no_grad()目标Q值公式里有max(Q_target(s))这个Q_target来自目标网络而目标网络的参数是定期从评估网络同步过来的。在计算目标值时我们并不希望梯度流经目标网络——因为那样会同时更新两个网络导致训练发散。no_grad()就是告诉PyTorch“这段计算只用于取值不追踪梯度”。我在很多初学代码里看到有人忘记这一点结果loss能降到负的训练完全乱掉。第三处(1 - dones)乘法的含义当done1时即下一状态是终止状态意味着没有未来奖励了目标Q值就应该等于当前奖励不是r gamma * max(Q)。所以要用(1 - dones)把终止状态的目标值“清零”掉未来部分。这个细节是DQN实现里最常见的bug来源少乘了这个掩码智能体在任务失败后还会幻想未来收益导致永远学不会。3.3 训练主循环让整个流程可视化train.py的完整代码如下# train.py import gym import numpy as np import matplotlib.pyplot as plt from collections import deque from agent import DQNAgent import config def train(): env gym.make(config.ENV_NAME) state_dim env.observation_space.shape[0] action_dim env.action_space.n agent DQNAgent(state_dim, action_dim, config) rewards_history [] recent_rewards deque(maxlen50) # 用于计算最近50回合的平均奖励 for episode in range(config.MAX_EPISODES): state, _ env.reset() total_reward 0 episode_loss [] # 记录本回合的loss用于观察学习状态 for step in range(config.MAX_STEPS): action agent.choose_action(state) next_state, reward, terminated, truncated, _ env.step(action) done terminated or truncated agent.store_transition(state, action, reward, next_state, done) agent.learn() state next_state total_reward reward if done: break agent.decay_epsilon() rewards_history.append(total_reward) recent_rewards.append(total_reward) # 每50回合打印一次详细日志 if (episode 1) % 50 0: avg_reward np.mean(recent_rewards) print(fEpisode {episode 1:3d} | fReward: {total_reward:3d} | fAvg (last 50): {avg_reward:5.1f} | fEpsilon: {agent.epsilon:.3f}) env.close() # 绘制并保存训练曲线 plt.plot(rewards_history) plt.xlabel(Episode) plt.ylabel(Total Reward) plt.title(DQN Training on CartPole-v1) plt.grid(True) plt.savefig(training_curve.png) plt.show() agent.save(dqn_model.pth) print(训练完成模型已保存为 dqn_model.pth) if __name__ __main__: train()这里有几个值得细说的设计terminated和truncated的区别。CartPole-v1从Gymnasium版本开始step()返回5个值其中terminated表示“因任务失败而结束”杆子倒了truncated表示“因达到步数上限而结束”坚持了500步。在训练时这两种情况对智能体来说意义不同前者是真的失败后者实际上是成功。我在这里把两者都当作done处理来存储经验这是最常见的做法。但如果你想做更精细的控制可以只把terminated当作done写入经验池这样智能体不会把“坚持到500步”误判为失败。训练日志的节奏。每步都打印reward会刷屏刷到没法看每回合都打印又太频繁。我选择每50回合打印一次并且用deque(maxlen50)计算滑动平均。当平均奖励稳定在450以上时基本可以判断策略已经收敛。4. 方法替换实战从DQN切换到SARSA和Policy Gradient4.1 替换的核心保持接口不变如果你理解了前文的设计思路替换算法其实很简单——只需要修改agent.py保证四个公开接口不变choose_action(state)返回动作store_transition(state, action, reward, next_state, done)存储一条经验learn()从历史经验中学习decay_epsilon()调整探索率如果算法不用探索率可以是空函数train.py完全不需要改动。下面我以SARSA为例展示替换的具体做法。SARSA与DQN最大的区别在于DQN更新时用的是max(Q(s))——即“下一状态所有动作中最大的Q值”这是离线策略off-policy因为学习的策略与采样的策略可以不同SARSA更新时用的是Q(s, a)——即“下一状态实际采取的那个动作的Q值”这是在线策略on-policy因为学习的数据必须由当前策略产生。所以在SARSA的实现里learn()需要在采样经验之外额外知道每个样本的下一步动作是什么。这就需要对store_transition做一点扩展——把next_action也存进去。好在train.py只调用了store_transition(state, action, reward, next_state, done)这5个参数并没有强制规定内部怎么处理。SARSA版本的store_transition可以这样实现def store_transition(self, state, action, reward, next_state, done): # SARSA需要额外确定next_action next_action self.choose_action(next_state) self.memory.push(state, action, reward, next_state, done, next_action)因为choose_action(state)是公开接口所以完全可以在内部调用。这样改完之后train.py依然不需要任何变化——这是接口设计的最大好处。4.2 DQN换Policy Gradient网络输出彻底改变Policy GradientREINFORCE与DQN思路完全不同。DQN输出的是每个动作的Q值而Policy Gradient输出的是每个动作的概率分布DQN用Q值直接选择动作Policy Gradient根据概率分布采样动作DQN的loss来自Q值与目标值的均方误差Policy Gradient的loss来自负对数概率乘以累积回报。这样大的差异agent.py内部几乎要重写但train.py依然可以保持原样——因为Policy Gradient同样只需要choose_action、store_transition、learn三个接口。decay_epsilon也不需要了Policy Gradient天然有探索能力概率采样本身就是随机的写成空函数或直接不调用都行。一个典型的Policy Gradient Agent核心实现如下class PolicyGradientAgent: def __init__(self, state_dim, action_dim, config): self.policy_net PolicyNetwork(state_dim, action_dim, config.HIDDEN_SIZE) self.optimizer optim.Adam(self.policy_net.parameters(), lrconfig.LEARNING_RATE) self.episode_states [] self.episode_actions [] self.episode_rewards [] def choose_action(self, state): state torch.FloatTensor(np.array(state)).unsqueeze(0) with torch.no_grad(): probs self.policy_net(state) dist torch.distributions.Categorical(probs) action dist.sample().item() return action def store_transition(self, state, action, reward, next_state, done): # 累积一个回合的数据回合结束时统一更新 self.episode_states.append(state) self.episode_actions.append(action) self.episode_rewards.append(reward) def learn(self): # 回合结束时计算折扣回报并更新策略网络 ...可以看到store_transition的5个参数只用了前3个next_state和done对蒙特卡洛型的Policy Gradient确实用不上。真正的更新逻辑在回合结束时触发这需要train.py在done之后额外调用一次agent.learn_single_episode()。如果你不想改train.py也可以在store_transition里检测到doneTrue时自动调用学习逻辑——这完全合法因为经验是在step()里一个个传进去的。提示这里的核心思想是train.py是数据生产者agent是数据消费者。只要保证数据从这个接口进去具体怎么消费是agent自己的事。这也是工程上“依赖倒置”的一种体现。4.3 算法对比什么时候该换哪个算法更新方式适用场景优缺点DQNoff-policy离散动作空间状态维度不高样本效率高但调参空间大需要目标网络和合理探索率SARSAon-policy对安全敏感的离散任务更保守训练过程更平稳但样本效率略低Policy Gradienton-policy Monte Carlo动作空间离散或连续需要随机策略实现简单天然探索但方差大、收敛慢如果后续想继续扩展还可以考虑把DQN的QNetwork从全连接换成卷积网络做图像输入把Policy Gradient的Monte Carlo回报换成GAE广义优势估计或者加一个Critic网络变成Actor-Critic架构。骨架不变换的永远只是agent.py内部这比我以前写的单体代码——换算法等于重写全文——舒服太多了。5. 常见问题与排查技巧实录我整理了几类最常见的报错和异常结果都是实际运行中踩过的坑。5.1env.reset()返回两个值解包报错如果你用的是旧版gym0.25以下env.reset()只返回一个状态数组新版gymnasium返回(state, info)两个值。如果发现ValueError: too many values to unpack或者正好反过来说明版本不匹配。最省事的办法是安装gymnasium而不是gympip install gymnasium matplotlib numpy torch然后用import gymnasium as gym。这也是上面代码能直接运行的依赖组合。5.2 训练曲线永远上不去奖励始终在10左右徘徊这种情况通常是探索率衰减太快导致智能体还没充分探索就过早进入“利用”模式。解决方向把EPSILON_DECAY从0.995调到0.999让高探索状态持续更久把EPSILON_START从1.0调成0.5如果智能体一开始随机乱跳很容易快速失败探索收益很低检查GAMMA是不是设得太低比如0.9以下会让智能体变得短视5.3 训练到一半loss变成NaNNaN基本只可能是梯度爆炸引起的。在CartPole环境下最常见的原因是学习率太高。把LEARNING_RATE从1e-3降到1e-4通常能解决。另外检查一下输入状态是否有异常大值——CartPole的观察值范围比较小一般不是这个问题。5.4 训练曲线看起来很平滑但实际表现很差这是奖励设计陷阱。CartPole里每个step都1所以模型只要“多活一步”就能获得更高奖励。如果你的评估方式是实时渲染看智能体表现会因为渲染导致的步频变化产生偏差。我建议训练完之后单独跑一轮评估模式关闭探索agent.epsilon 0.0 # 关闭随机探索 for _ in range(5): state, _ env.reset() total 0 for _ in range(500): action agent.choose_action(state) state, reward, terminated, truncated, _ env.step(action) total reward if terminated or truncated: break print(fEval total reward: {total})连续跑5次奖励都在460以上才算真正的收敛。5.5 换算法时报错采样出的batch不够如果从DQN换成SARSA后发现learn()里random.sample(self.buffer, batch_size)报错说明经验池还没来得及攒够一个batch。我的实现里在learn()开头判断了len(self.memory) BATCH_SIZE就直接返回正常不会触发但如果你把MEMORY_CAPACITY设置得小于BATCH_SIZE就会出问题。思路很简单经验池容量必须大于等于batch大小一般建议至少是batch的10倍以上。5.6 目标网络同步频率的影响我试过把TARGET_UPDATE_FREQ分别设成10、100、1000观察到的现象是设10时训练极其不稳定loss曲线剧烈震荡设1000时前期训练速度明显变慢但后期平稳。100是个兼顾稳定性和速度的折中值。如果你愿意等200到300之间的效果也完全可以接受。6. 从“能跑”到“能改”下一步应该做什么代码跑通只是第一步强化学习的美妙之处在于同一个框架能演化出无数变体。以下三个方向是我建议的后续扩展路径方向一改进探索策略。当前用epsilon-greedy可以改成逐渐衰减的softmax探索或者加上噪声网络NoisyNet在choose_action里就能实现不需要动其他任何代码。方向二换成连续动作空间。把环境从CartPole换成Pendulum-v1或LunarLanderContinuous-v2DQN就用不了了这时候需要上DDPG或SAC。但好消息是train.py的主循环几乎不用改choose_action从输出离散索引变成输出连续数值learn()从Q学习变成Actor-Critic更新骨架依然成立。方向三加入状态可视化。在learn()里把每个batch的平均Q值打印出来。你会发现训练初期平均Q值很低因为随机探索频繁失败随着训练推进稳步上升。Q值曲线的形状本身就是判断训练健康度的绝佳指标。最后再分享一个小技巧每次改动算法之后先固定随机种子跑一遍老算法再跑新算法对比曲线差异。强化学习训练本身有随机性如果不控制种子你很难判断效果变好是算法改进还是运气使然。我在代码里没写固定种子的逻辑因为那会让代码看起来更复杂但你在自己的实验里一定要加上import random import numpy as np import torch random.seed(42) np.random.seed(42) torch.manual_seed(42)这样出来的曲线才有可比性。这个项目我已经用了大半年从培训新人到验证新的算法思路都靠它如果你照着文章跑通了再往深走就会发现强化学习真正困难的地方不在于写代码而在于理解每个设计选择背后的权衡。代码只是思想的载体思路清晰了代码自然就顺了。本文还有配套的精品资源点击获取