用PyTorch和Gymnasium玩转倒立摆：DQN实战避坑指南（附完整代码）

用PyTorch和Gymnasium玩转倒立摆DQN实战避坑指南附完整代码倒立摆控制是强化学习入门的经典案例它能直观展示智能体如何通过试错学习平衡策略。本文将带您用PyTorch和Gymnasium实现深度Q网络(DQN)解决方案避开环境配置、训练调试中的常见陷阱并提供可复用的完整代码框架。1. 环境搭建与核心工具链1.1 工具选型与版本管理现代强化学习开发需要特别注意工具链的版本兼容性。推荐使用以下组合# 版本锁定示例requirements.txt gymnasium0.29.1 torch2.2.1 matplotlib3.8.2关键差异说明Gymnasium作为OpenAI Gym的升级维护版API更稳定且持续更新。若遇到CartPole-v1环境报错检查是否误装了旧版gym。硬件加速配置技巧device torch.device( cuda if torch.cuda.is_available() else mps if torch.backends.mps.is_available() else cpu ) print(fUsing {device} acceleration)1.2 环境观测空间解析倒立摆环境提供4维状态观测小车位置-4.8到4.8小车速度无界杆角度±0.418弧度杆角速度无界env gym.make(CartPole-v1) state, _ env.reset() print(fObservation space: {env.observation_space}) print(fInitial state: {state})2. DQN实现关键组件2.1 神经网络架构设计采用三层全连接网络处理状态输入class DQN(nn.Module): def __init__(self, n_observations, n_actions): super().__init__() self.net nn.Sequential( nn.Linear(n_observations, 128), nn.ReLU(), nn.Linear(128, 128), nn.ReLU(), nn.Linear(128, n_actions) ) def forward(self, x): return self.net(x)参数初始化技巧添加nn.init.kaiming_normal_(layer.weight)可加速收敛。2.2 经验回放机制实现经验回放池避免样本相关性Transition namedtuple(Transition, (state, action, next_state, reward)) class ReplayMemory: def __init__(self, capacity): self.memory deque([], maxlencapacity) def push(self, *args): self.memory.append(Transition(*args)) def sample(self, batch_size): return random.sample(self.memory, batch_size)提示经验池容量建议设置为10000-50000batch size通常取128或2563. 训练流程优化策略3.1 ε-贪婪策略的动态调整智能探索与利用的平衡def select_action(state): global steps_done eps_threshold EPS_END (EPS_START - EPS_END) * \ math.exp(-1. * steps_done / EPS_DECAY) steps_done 1 if random.random() eps_threshold: with torch.no_grad(): return policy_net(state).argmax(dim1, keepdimTrue) else: return torch.tensor( [[env.action_space.sample()]], devicedevice, dtypetorch.long )参数建议EPS_START0.9EPS_END0.05EPS_DECAY10003.2 目标网络更新技巧双网络架构稳定训练def update_target_net(): target_net_state_dict target_net.state_dict() policy_net_state_dict policy_net.state_dict() for key in policy_net_state_dict: target_net_state_dict[key] policy_net_state_dict[key]*TAU \ target_net_state_dict[key]*(1-TAU) target_net.load_state_dict(target_net_state_dict)注意更新率TAU通常取0.005过大可能导致震荡4. 训练监控与可视化4.1 实时训练曲线绘制def plot_durations(show_resultFalse): plt.figure(1) durations_t torch.tensor(episode_durations, dtypetorch.float) if show_result: plt.title(Result) else: plt.clf() plt.title(Training...) plt.xlabel(Episode) plt.ylabel(Duration) plt.plot(durations_t.numpy()) if len(durations_t) 100: means durations_t.unfold(0, 100, 1).mean(1).view(-1) means torch.cat((torch.zeros(99), means)) plt.plot(means.numpy()) plt.pause(0.001)4.2 关键指标日志记录建议监控以下指标平均回合时长100episode滑动平均ε值变化曲线Q值变化幅度损失函数下降趋势5. 典型问题排查指南5.1 训练不收敛的解决方案现象可能原因解决方法奖励始终很低学习率过大调低LR至1e-4~1e-5波动剧烈batch size太小增大至128或256Q值爆炸没有梯度裁剪添加clip_grad_norm_5.2 性能优化技巧GPU利用率提升# 数据预处理时添加pin_memory dataloader DataLoader(dataset, batch_size32, pin_memoryTrue, num_workers2)向量化环境使用gymnasium.vector可并行多个环境envs gym.vector.make(CartPole-v1, num_envs4) states, _ envs.reset()6. 完整代码架构# 完整导入部分 import gymnasium as gym import torch import torch.nn as nn import torch.optim as optim from collections import deque, namedtuple import random import math import matplotlib.pyplot as plt # 超参数配置部分 BATCH_SIZE 128 GAMMA 0.99 EPS_START 0.9 EPS_END 0.05 EPS_DECAY 1000 TAU 0.005 LR 1e-4 # 网络定义、经验回放等组件... # 训练主循环... if __name__ __main__: train_dqn()实际测试发现在RTX 3090上训练600回合约需15分钟最终平均持续时间可达195步以上环境最大步长为200。建议先进行50回合的快速验证确认基本流程正常后再开展完整训练。