深度Q网络案例实践

1
import torch
2
import torch.nn as nn
3
import torch.optim as optim
4
import torch.nn.functional as F
5
import numpy as np
6
import gym
7
import random
8
import matplotlib.pyplot as plt
9
from collections import deque
10
import warnings
11
warnings.filterwarnings('ignore')
12

13
# 设置随机种子
14
SEED = 42
15
random.seed(SEED)
16
np.random.seed(SEED)
17
torch.manual_seed(SEED)
18

19
# 检查GPU
20
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
21
print(f"使用设备: {device}")

1
class DQN(nn.Module):
2
    def __init__(self, state_size, action_size, hidden_size=128):
3
        """
4
        DQN网络架构
5

6
        Args:
7
            state_size (int): 状态空间维度
8
            action_size (int): 动作空间维度
9
            hidden_size (int): 隐藏层大小
10
        """
11
        super(DQN, self).__init__()
12

13
        self.fc1 = nn.Linear(state_size, hidden_size)
14
        self.fc2 = nn.Linear(hidden_size, hidden_size)
15
        self.fc3 = nn.Linear(hidden_size, hidden_size)
16
        self.fc4 = nn.Linear(hidden_size, action_size)
17

18
        self.dropout = nn.Dropout(0.2)
19

20
        # 权重初始化
21
        self.apply(self._init_weights)
22

23
    def _init_weights(self, module):
24
        """Xavier初始化"""
25
        if isinstance(module, nn.Linear):
26
            torch.nn.init.xavier_uniform_(module.weight)
27
            module.bias.data.fill_(0.01)
28

29
    def forward(self, state):
30
        """前向传播"""
31
        x = F.relu(self.fc1(state))
32
        x = self.dropout(x)
33
        x = F.relu(self.fc2(x))
34
        x = self.dropout(x)
35
        x = F.relu(self.fc3(x))
36
        x = self.fc4(x)
37
        return x

1
class ReplayBuffer:
2
    def __init__(self, capacity):
3
        """
4
        经验回放缓冲区
5

6
        Args:
7
            capacity (int): 缓冲区最大容量
8
        """
9
        self.buffer = deque(maxlen=capacity)
10

11
    def push(self, state, action, reward, next_state, done):
12
        """添加经验"""
13
        experience = (state, action, reward, next_state, done)
14
        self.buffer.append(experience)
15

16
    def sample(self, batch_size):
17
        """随机采样批次数据"""
18
        batch = random.sample(self.buffer, batch_size)
19

20
        states = torch.FloatTensor([e[0] for e in batch]).to(device)
21
        actions = torch.LongTensor([e[1] for e in batch]).to(device)
22
        rewards = torch.FloatTensor([e[2] for e in batch]).to(device)
23
        next_states = torch.FloatTensor([e[3] for e in batch]).to(device)
24
        dones = torch.BoolTensor([e[4] for e in batch]).to(device)
25

26
        return states, actions, rewards, next_states, dones
27

28
    def __len__(self):
29
        return len(self.buffer)

1
class DQNAgent:
2
    def __init__(self, state_size, action_size, lr=1e-3, gamma=0.99,
3
                 epsilon=1.0, epsilon_decay=0.995, epsilon_min=0.01,
4
                 buffer_size=10000, batch_size=64, target_update=10):
5
        """
6
        DQN智能体
7

8
        Args:
9
            state_size (int): 状态空间维度
10
            action_size (int): 动作空间维度
11
            lr (float): 学习率
12
            gamma (float): 折扣因子
13
            epsilon (float): 初始探索率
14
            epsilon_decay (float): 探索率衰减
15
            epsilon_min (float): 最小探索率
16
            buffer_size (int): 经验回放缓冲区大小
17
            batch_size (int): 批次大小
18
            target_update (int): 目标网络更新频率
19
        """
20
        self.state_size = state_size
21
        self.action_size = action_size
22
        self.lr = lr
23
        self.gamma = gamma
24
        self.epsilon = epsilon
25
        self.epsilon_decay = epsilon_decay
26
        self.epsilon_min = epsilon_min
27
        self.batch_size = batch_size
28
        self.target_update = target_update
29

30
        # 神经网络
31
        self.q_network = DQN(state_size, action_size).to(device)
32
        self.target_network = DQN(state_size, action_size).to(device)
33
        self.optimizer = optim.Adam(self.q_network.parameters(), lr=lr)
34

35
        # 经验回放
36
        self.memory = ReplayBuffer(buffer_size)
37

38
        # 计数器
39
        self.update_count = 0
40

41
    def act(self, state, training=True):
42
        """
43
        选择动作（ε-贪婪策略）
44

45
        Args:
46
            state: 当前状态
47
            training (bool): 是否在训练模式
48

49
        Returns:
50
            action: 选择的动作
51
        """
52
        if training and random.random() < self.epsilon:
53
            return random.choice(range(self.action_size))
54

55
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
56
        q_values = self.q_network(state_tensor)
57
        return q_values.argmax().item()
58

59
    def remember(self, state, action, reward, next_state, done):
60
        """存储经验"""
61
        self.memory.push(state, action, reward, next_state, done)
62

63
    def replay(self):
64
        """经验回放学习"""
65
        if len(self.memory) < self.batch_size:
66
            return None
67

68
        # 采样批次数据
69
        states, actions, rewards, next_states, dones = self.memory.sample(self.batch_size)
70

71
        # 当前Q值
72
        current_q_values = self.q_network(states).gather(1, actions.unsqueeze(1))
73

74
        # 下一步Q值（使用目标网络）
75
        next_q_values = self.target_network(next_states).max(1)[0].detach()
76
        target_q_values = rewards + (self.gamma * next_q_values * ~dones)
77

78
        # 计算损失
79
        loss = F.mse_loss(current_q_values.squeeze(), target_q_values)
80

81
        # 反向传播
82
        self.optimizer.zero_grad()
83
        loss.backward()
84

85
        # 梯度裁剪
86
        torch.nn.utils.clip_grad_norm_(self.q_network.parameters(), 1.0)
87

88
        self.optimizer.step()
89

90
        # 更新探索率
91
        if self.epsilon > self.epsilon_min:
92
            self.epsilon *= self.epsilon_decay
93

94
        # 定期更新目标网络
95
        self.update_count += 1
96
        if self.update_count % self.target_update == 0:
97
            self.update_target_network()
98

99
        return loss.item()
100

101
    def update_target_network(self):
102
        """更新目标网络"""
103
        self.target_network.load_state_dict(self.q_network.state_dict())
104

105
    def save(self, filepath):
106
        """保存模型"""
107
        torch.save({
108
            'q_network_state_dict': self.q_network.state_dict(),
109
            'target_network_state_dict': self.target_network.state_dict(),
110
            'optimizer_state_dict': self.optimizer.state_dict(),
111
            'epsilon': self.epsilon
112
        }, filepath)
113

114
    def load(self, filepath):
115
        """加载模型"""
116
        checkpoint = torch.load(filepath)
117
        self.q_network.load_state_dict(checkpoint['q_network_state_dict'])
118
        self.target_network.load_state_dict(checkpoint['target_network_state_dict'])
119
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
120
        self.epsilon = checkpoint['epsilon']

1
def train_dqn(env, agent, n_episodes=2000, max_steps=500,
2
              solve_score=195, solve_episodes=100):
3
    """
4
    训练DQN智能体
5

6
    Args:
7
        env: 环境
8
        agent: DQN智能体
9
        n_episodes (int): 训练轮数
10
        max_steps (int): 每轮最大步数
11
        solve_score (int): 解决问题的平均分数
12
        solve_episodes (int): 连续解决问题的轮数
13

14
    Returns:
15
        scores: 每轮得分列表
16
        losses: 损失值列表
17
    """
18
    scores = []
19
    losses = []
20
    recent_scores = deque(maxlen=solve_episodes)
21

22
    for episode in range(n_episodes):
23
        state = env.reset()
24
        total_reward = 0
25
        episode_losses = []
26

27
        for step in range(max_steps):
28
            # 选择动作
29
            action = agent.act(state)
30

31
            # 执行动作
32
            next_state, reward, done, _ = env.step(action)
33

34
            # 存储经验
35
            agent.remember(state, action, reward, next_state, done)
36

37
            # 学习
38
            loss = agent.replay()
39
            if loss is not None:
40
                episode_losses.append(loss)
41

42
            state = next_state
43
            total_reward += reward
44

45
            if done:
46
                break
47

48
        scores.append(total_reward)
49
        recent_scores.append(total_reward)
50

51
        # 记录平均损失
52
        if episode_losses:
53
            losses.append(np.mean(episode_losses))
54
        else:
55
            losses.append(0)
56

57
        # 打印进度
58
        if episode % 100 == 0:
59
            avg_score = np.mean(recent_scores)
60
            print(f"Episode {episode}, Score: {total_reward:.1f}, "
61
                  f"Avg Score: {avg_score:.1f}, Epsilon: {agent.epsilon:.3f}")
62

63
        # 检查是否解决问题
64
        if len(recent_scores) >= solve_episodes:
65
            avg_score = np.mean(recent_scores)
66
            if avg_score >= solve_score:
67
                print(f"\n环境在第 {episode} 轮解决！")
68
                print(f"最近 {solve_episodes} 轮平均分数: {avg_score:.1f}")
69
                break
70

71
    return scores, losses

1
def evaluate_agent(env, agent, n_episodes=10, render=False):
2
    """
3
    评估训练好的智能体
4

5
    Args:
6
        env: 环境
7
        agent: 训练好的智能体
8
        n_episodes (int): 评估轮数
9
        render (bool): 是否渲染环境
10

11
    Returns:
12
        eval_scores: 评估得分列表
13
    """
14
    eval_scores = []
15

16
    for episode in range(n_episodes):
17
        state = env.reset()
18
        total_reward = 0
19

20
        while True:
21
            if render:
22
                env.render()
23

24
            # 贪婪策略（不探索）
25
            action = agent.act(state, training=False)
26
            state, reward, done, _ = env.step(action)
27
            total_reward += reward
28

29
            if done:
30
                break
31

32
        eval_scores.append(total_reward)
33
        print(f"评估轮次 {episode + 1}: 得分 = {total_reward}")
34

35
    return eval_scores

1
def plot_training_results(scores, losses, window=100):
2
    """
3
    绘制训练结果
4

5
    Args:
6
        scores: 得分列表
7
        losses: 损失列表
8
        window (int): 移动平均窗口大小
9
    """
10
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
11

12
    # 绘制得分曲线
13
    ax1.plot(scores, alpha=0.3, color='blue', label='原始得分')
14

15
    # 计算移动平均
16
    moving_avg = []
17
    for i in range(len(scores)):
18
        start_idx = max(0, i - window + 1)
19
        moving_avg.append(np.mean(scores[start_idx:i+1]))
20

21
    ax1.plot(moving_avg, color='red', linewidth=2, label=f'{window}轮移动平均')
22
    ax1.set_xlabel('训练轮次')
23
    ax1.set_ylabel('得分')
24
    ax1.set_title('DQN训练过程 - 得分变化')
25
    ax1.legend()
26
    ax1.grid(True, alpha=0.3)
27

28
    # 绘制损失曲线
29
    ax2.plot(losses, color='orange', alpha=0.7)
30
    ax2.set_xlabel('训练轮次')
31
    ax2.set_ylabel('损失值')
32
    ax2.set_title('DQN训练过程 - 损失变化')
33
    ax2.grid(True, alpha=0.3)
34

35
    plt.tight_layout()
36
    plt.show()
37

38
def plot_q_values_heatmap(agent, env, resolution=20):
39
    """
40
    可视化Q值热力图（针对2D状态空间的简化版本）
41

42
    Args:
43
        agent: 训练好的智能体
44
        env: 环境
45
        resolution (int): 网格分辨率
46
    """
47
    # 获取状态范围
48
    low = env.observation_space.low[:2]  # 只考虑位置和速度
49
    high = env.observation_space.high[:2]
50

51
    # 创建网格
52
    x = np.linspace(low[0], high[0], resolution)
53
    y = np.linspace(low[1], high[1], resolution)
54
    X, Y = np.meshgrid(x, y)
55

56
    # 计算Q值
57
    q_values = np.zeros((resolution, resolution, agent.action_size))
58

59
    for i in range(resolution):
60
        for j in range(resolution):
61
            # 简化状态（假设角度和角速度为0）
62
            state = np.array([X[i, j], Y[i, j], 0.0, 0.0])
63
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
64

65
            with torch.no_grad():
66
                q_vals = agent.q_network(state_tensor).cpu().numpy()[0]
67
                q_values[i, j] = q_vals
68

69
    # 绘制热力图
70
    fig, axes = plt.subplots(1, agent.action_size, figsize=(12, 5))
71
    actions = ['向左推', '向右推']
72

73
    for a in range(agent.action_size):
74
        im = axes[a].imshow(q_values[:, :, a], extent=[low[0], high[0], low[1], high[1]],
75
                           origin='lower', cmap='RdYlBu')
76
        axes[a].set_title(f'Q值: {actions[a]}')
77
        axes[a].set_xlabel('位置')
78
        axes[a].set_ylabel('速度')
79
        plt.colorbar(im, ax=axes[a])
80

81
    plt.tight_layout()
82
    plt.show()

1
def main():
2
    """主函数"""
3
    # 创建环境
4
    env = gym.make('CartPole-v1')
5

6
    # 获取状态和动作空间大小
7
    state_size = env.observation_space.shape[0]
8
    action_size = env.action_space.n
9

10
    print(f"状态空间维度: {state_size}")
11
    print(f"动作空间大小: {action_size}")
12

13
    # 创建智能体
14
    agent = DQNAgent(
15
        state_size=state_size,
16
        action_size=action_size,
17
        lr=1e-3,
18
        gamma=0.99,
19
        epsilon=1.0,
20
        epsilon_decay=0.995,
21
        epsilon_min=0.01,
22
        buffer_size=10000,
23
        batch_size=64,
24
        target_update=10
25
    )
26

27
    print("开始训练...")
28
    scores, losses = train_dqn(env, agent, n_episodes=2000)
29

30
    # 保存模型
31
    agent.save('dqn_cartpole.pth')
32
    print("模型已保存到 dqn_cartpole.pth")
33

34
    # 绘制训练结果
35
    plot_training_results(scores, losses)
36

37
    # 评估智能体
38
    print("\n开始评估...")
39
    eval_scores = evaluate_agent(env, agent, n_episodes=10)
40
    print(f"评估平均得分: {np.mean(eval_scores):.2f} ± {np.std(eval_scores):.2f}")
41

42
    # 可视化Q值
43
    plot_q_values_heatmap(agent, env)
44

45
    env.close()
46

47
if __name__ == "__main__":
48
    main()

1
class ImprovedDQN(nn.Module):
2
    def __init__(self, state_size, action_size, hidden_size=128):
3
        super(ImprovedDQN, self).__init__()
4

5
        # Dueling架构
6
        self.feature_layer = nn.Sequential(
7
            nn.Linear(state_size, hidden_size),
8
            nn.ReLU(),
9
            nn.Linear(hidden_size, hidden_size),
10
            nn.ReLU()
11
        )
12

13
        # 状态价值流
14
        self.value_stream = nn.Linear(hidden_size, 1)
15

16
        # 优势流
17
        self.advantage_stream = nn.Linear(hidden_size, action_size)
18

19
    def forward(self, state):
20
        features = self.feature_layer(state)
21

22
        value = self.value_stream(features)
23
        advantage = self.advantage_stream(features)
24

25
        # Dueling aggregation
26
        q_values = value + (advantage - advantage.mean(dim=1, keepdim=True))
27

28
        return q_values