PyTorch实战：用DQN算法教AI玩平衡杆游戏（附完整代码）

PyTorch实战：用DQN算法构建AI平衡杆智能体（附完整代码与可视化分析）

1. 项目概述与核心价值

CartPole平衡杆问题是强化学习领域的经典测试环境，它模拟了一个底部装有铰接杆的小车，智能体需要通过左右移动小车来保持杆子竖直。这个看似简单的任务实际上包含了状态感知、动作决策和长期奖励优化等强化学习核心问题。

深度Q网络（DQN）将传统的Q-Learning与深度学习相结合，通过神经网络近似Q值函数，解决了传统方法在连续状态空间下的维度灾难问题。PyTorch的动态计算图和自动微分机制，使得我们可以更直观地构建和调试DQN模型。本项目将实现以下技术亮点：

• 双网络架构：分离的评估网络和目标网络设计，显著提升训练稳定性
• 经验回放机制：打破样本间的时间相关性，提高数据利用率
• ε-greedy策略：平衡探索与利用，避免陷入局部最优
• 实时可视化：训练过程动态展示，直观理解算法收敛过程

import gym
import torch
import random
from collections import deque
import matplotlib.pyplot as plt

class DQNAgent:
    def __init__(self, state_size, action_size):
        self.state_size = state_size
        self.action_size = action_size
        self.memory = deque(maxlen=2000)  # 经验回放缓冲区
        self.gamma = 0.95    # 折扣因子
        self.epsilon = 1.0   # 探索率
        self.epsilon_min = 0.01
        self.epsilon_decay = 0.995
        self.learning_rate = 0.001
        self.model = self._build_model()  # 评估网络
        self.target_model = self._build_model()  # 目标网络

2. 环境解析与状态空间设计

CartPole-v1环境提供了4个维度的状态观测值：

状态参数	物理意义	数值范围
cart_pos	小车位置	[-4.8, 4.8]
cart_vel	小车速度	[-∞, ∞]
pole_ang	杆子角度	[-0.418, 0.418]弧度
pole_vel	杆子角速度	[-∞, ∞]

奖励函数设计为每步存活奖励+1，当出现以下情况时回合终止：

• 杆子倾斜角度超过±12°
• 小车移动超出轨道边界（±2.4单位）
• 步数超过500步（v1版本）

env = gym.make('CartPole-v1')
print("观测空间:", env.observation_space)
print("动作空间:", env.action_space)

# 示例状态观测
state = env.reset()
print("初始状态示例:", state)

3. DQN网络架构与关键技术

3.1 神经网络设计

我们采用三层全连接网络作为Q函数逼近器：

1. 输入层：4个神经元，对应环境状态维度
2. 隐藏层：24个神经元，ReLU激活
3. 输出层：2个神经元，对应左右动作的Q值

import torch.nn as nn

class QNetwork(nn.Module):
    def __init__(self, state_size, action_size):
        super(QNetwork, self).__init__()
        self.fc1 = nn.Linear(state_size, 24)
        self.fc2 = nn.Linear(24, 24)
        self.fc3 = nn.Linear(24, action_size)
    
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)

3.2 经验回放实现

经验回放缓冲区存储转移样本(s, a, r, s', done)，训练时随机采样小批量数据：

class ReplayBuffer:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)
    
    def push(self, state, action, reward, next_state, done):
        self.buffer.append((state, action, reward, next_state, done))
    
    def sample(self, batch_size):
        return random.sample(self.buffer, batch_size)
    
    def __len__(self):
        return len(self.buffer)

3.3 训练流程关键步骤

1. 动作选择：ε-greedy策略平衡探索与利用
2. 目标计算：使用目标网络计算y = r + γ*maxQ(s',a')
3. 损失函数：Huber损失减少异常值影响
4. 网络更新：定期同步目标网络参数

def train(self, batch_size):
    if len(self.memory) < batch_size:
        return
    minibatch = self.memory.sample(batch_size)
    
    states = torch.FloatTensor([t[0] for t in minibatch])
    actions = torch.LongTensor([t[1] for t in minibatch])
    rewards = torch.FloatTensor([t[2] for t in minibatch])
    next_states = torch.FloatTensor([t[3] for t in minibatch])
    dones = torch.FloatTensor([t[4] for t in minibatch])
    
    current_q = self.model(states).gather(1, actions.unsqueeze(1))
    next_q = self.target_model(next_states).max(1)[0].detach()
    target = rewards + (1 - dones) * self.gamma * next_q
    
    loss = F.smooth_l1_loss(current_q.squeeze(), target)
    self.optimizer.zero_grad()
    loss.backward()
    self.optimizer.step()

4. 完整训练流程与超参数调优

4.1 训练参数配置

参数	推荐值	作用
batch_size	64	每次训练的样本量
gamma	0.99	未来奖励折扣因子
epsilon_start	1.0	初始探索率
epsilon_min	0.01	最小探索率
epsilon_decay	0.995	探索率衰减系数
target_update	100	目标网络更新间隔
lr	1e-3	学习率

4.2 训练监控与可视化

实时绘制两个关键指标：

1. 回合持续时间（越长越好）
2. 最近100轮平均得分（平滑曲线）

def plot_training(episode_durations):
    plt.figure(figsize=(10,5))
    plt.subplot(121)
    plt.plot(episode_durations)
    plt.title('Episode Durations')
    
    plt.subplot(122)
    means = [np.mean(episode_durations[max(0,i-100):i+1]) 
             for i in range(len(episode_durations))]
    plt.plot(means)
    plt.title('Last 100 Means')
    plt.tight_layout()
    plt.show()

4.3 性能优化技巧

1. 奖励塑形：对杆子角度和小车位置设计更精细的奖励函数
2. 网络结构调整：增加层宽或使用卷积层处理图像输入
3. 优先经验回放：根据TD误差优先级采样重要经验
4. 双DQN：解耦动作选择和价值评估，减少过估计

# 改进的奖励函数示例
def custom_reward(state):
    x, x_dot, theta, theta_dot = state
    r1 = (env.x_threshold - abs(x)) / env.x_threshold  # 位置奖励
    r2 = (env.theta_threshold_radians - abs(theta)) / env.theta_threshold_radians  # 角度奖励
    return r1 + r2

5. 实战演示与结果分析

经过500轮训练后，智能体通常能达到以下性能：

• 平均存活时间 > 450步（v1版本最大500步）
• 收敛时探索率ε ≈ 0.01
• 训练损失稳定在0.1以下

测试时关闭探索，观察智能体控制效果：

def test_agent(env, agent, episodes=5):
    for ep in range(episodes):
        state = env.reset()
        done = False
        total_reward = 0
        
        while not done:
            env.render()
            action = agent.act(state, epsilon=0)  # 关闭探索
            state, reward, done, _ = env.step(action)
            total_reward += reward
        
        print(f"Test Episode {ep+1}, Total Reward: {total_reward}")
    env.close()

典型训练曲线展示：

• 初期：随机策略，存活时间短（<50步）
• 中期：快速学习阶段，性能指数级提升
• 后期：稳定在接近最大步数，探索率降低

6. 进阶改进方向

对于希望进一步提升性能的开发者，可以考虑：

1. Dueling DQN：分离状态价值和优势函数

   class DuelingQNetwork(nn.Module):
       def __init__(self, state_size, action_size):
           super().__init__()
           self.fc1 = nn.Linear(state_size, 64)
           self.value_stream = nn.Sequential(
               nn.Linear(64, 64),
               nn.ReLU(),
               nn.Linear(64, 1)
           )
           self.advantage_stream = nn.Sequential(
               nn.Linear(64, 64),
               nn.ReLU(),
               nn.Linear(64, action_size)
           )
       
       def forward(self, x):
           x = torch.relu(self.fc1(x))
           values = self.value_stream(x)
           advantages = self.advantage_stream(x)
           return values + (advantages - advantages.mean())
   

2. 多步学习：使用n步回报平衡偏差和方差

3. 噪声网络：参数空间噪声替代ε-greedy

4. 分布式DQN：学习价值分布而不仅是期望值

实际部署时发现，当环境参数发生变化（如杆子质量改变）时，基础DQN的适应性有限。这时可以考虑结合元学习或在线微调技术提升模型的泛化能力。