深度Q网络案例实践

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import numpy as np
import gym
import random
import matplotlib.pyplot as plt
from collections import deque
import warnings
warnings.filterwarnings('ignore')

# 设置随机种子
SEED = 42
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)

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

class DQN(nn.Module):
    def __init__(self, state_size, action_size, hidden_size=128):
        """
        DQN网络架构
        
        Args:
            state_size (int): 状态空间维度
            action_size (int): 动作空间维度  
            hidden_size (int): 隐藏层大小
        """
        super(DQN, self).__init__()
        
        self.fc1 = nn.Linear(state_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, hidden_size)
        self.fc4 = nn.Linear(hidden_size, action_size)
        
        self.dropout = nn.Dropout(0.2)
        
        # 权重初始化
        self.apply(self._init_weights)
    
    def _init_weights(self, module):
        """Xavier初始化"""
        if isinstance(module, nn.Linear):
            torch.nn.init.xavier_uniform_(module.weight)
            module.bias.data.fill_(0.01)
    
    def forward(self, state):
        """前向传播"""
        x = F.relu(self.fc1(state))
        x = self.dropout(x)
        x = F.relu(self.fc2(x))
        x = self.dropout(x)
        x = F.relu(self.fc3(x))
        x = self.fc4(x)
        return x

class ReplayBuffer:
    def __init__(self, capacity):
        """
        经验回放缓冲区
        
        Args:
            capacity (int): 缓冲区最大容量
        """
        self.buffer = deque(maxlen=capacity)
    
    def push(self, state, action, reward, next_state, done):
        """添加经验"""
        experience = (state, action, reward, next_state, done)
        self.buffer.append(experience)
    
    def sample(self, batch_size):
        """随机采样批次数据"""
        batch = random.sample(self.buffer, batch_size)
        
        states = torch.FloatTensor([e[0] for e in batch]).to(device)
        actions = torch.LongTensor([e[1] for e in batch]).to(device)
        rewards = torch.FloatTensor([e[2] for e in batch]).to(device)
        next_states = torch.FloatTensor([e[3] for e in batch]).to(device)
        dones = torch.BoolTensor([e[4] for e in batch]).to(device)
        
        return states, actions, rewards, next_states, dones
    
    def __len__(self):
        return len(self.buffer)

class DQNAgent:
    def __init__(self, state_size, action_size, lr=1e-3, gamma=0.99, 
                 epsilon=1.0, epsilon_decay=0.995, epsilon_min=0.01,
                 buffer_size=10000, batch_size=64, target_update=10):
        """
        DQN智能体
        
        Args:
            state_size (int): 状态空间维度
            action_size (int): 动作空间维度
            lr (float): 学习率
            gamma (float): 折扣因子
            epsilon (float): 初始探索率
            epsilon_decay (float): 探索率衰减
            epsilon_min (float): 最小探索率
            buffer_size (int): 经验回放缓冲区大小
            batch_size (int): 批次大小
            target_update (int): 目标网络更新频率
        """
        self.state_size = state_size
        self.action_size = action_size
        self.lr = lr
        self.gamma = gamma
        self.epsilon = epsilon
        self.epsilon_decay = epsilon_decay
        self.epsilon_min = epsilon_min
        self.batch_size = batch_size
        self.target_update = target_update
        
        # 神经网络
        self.q_network = DQN(state_size, action_size).to(device)
        self.target_network = DQN(state_size, action_size).to(device)
        self.optimizer = optim.Adam(self.q_network.parameters(), lr=lr)
        
        # 经验回放
        self.memory = ReplayBuffer(buffer_size)
        
        # 计数器
        self.update_count = 0
        
    def act(self, state, training=True):
        """
        选择动作（ε-贪婪策略）
        
        Args:
            state: 当前状态
            training (bool): 是否在训练模式
        
        Returns:
            action: 选择的动作
        """
        if training and random.random() < self.epsilon:
            return random.choice(range(self.action_size))
        
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
        q_values = self.q_network(state_tensor)
        return q_values.argmax().item()
    
    def remember(self, state, action, reward, next_state, done):
        """存储经验"""
        self.memory.push(state, action, reward, next_state, done)
    
    def replay(self):
        """经验回放学习"""
        if len(self.memory) < self.batch_size:
            return None
        
        # 采样批次数据
        states, actions, rewards, next_states, dones = self.memory.sample(self.batch_size)
        
        # 当前Q值
        current_q_values = self.q_network(states).gather(1, actions.unsqueeze(1))
        
        # 下一步Q值（使用目标网络）
        next_q_values = self.target_network(next_states).max(1)[0].detach()
        target_q_values = rewards + (self.gamma * next_q_values * ~dones)
        
        # 计算损失
        loss = F.mse_loss(current_q_values.squeeze(), target_q_values)
        
        # 反向传播
        self.optimizer.zero_grad()
        loss.backward()
        
        # 梯度裁剪
        torch.nn.utils.clip_grad_norm_(self.q_network.parameters(), 1.0)
        
        self.optimizer.step()
        
        # 更新探索率
        if self.epsilon > self.epsilon_min:
            self.epsilon *= self.epsilon_decay
        
        # 定期更新目标网络
        self.update_count += 1
        if self.update_count % self.target_update == 0:
            self.update_target_network()
        
        return loss.item()
    
    def update_target_network(self):
        """更新目标网络"""
        self.target_network.load_state_dict(self.q_network.state_dict())
    
    def save(self, filepath):
        """保存模型"""
        torch.save({
            'q_network_state_dict': self.q_network.state_dict(),
            'target_network_state_dict': self.target_network.state_dict(),
            'optimizer_state_dict': self.optimizer.state_dict(),
            'epsilon': self.epsilon
        }, filepath)
    
    def load(self, filepath):
        """加载模型"""
        checkpoint = torch.load(filepath)
        self.q_network.load_state_dict(checkpoint['q_network_state_dict'])
        self.target_network.load_state_dict(checkpoint['target_network_state_dict'])
        self.optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        self.epsilon = checkpoint['epsilon']

def train_dqn(env, agent, n_episodes=2000, max_steps=500, 
              solve_score=195, solve_episodes=100):
    """
    训练DQN智能体
    
    Args:
        env: 环境
        agent: DQN智能体
        n_episodes (int): 训练轮数
        max_steps (int): 每轮最大步数
        solve_score (int): 解决问题的平均分数
        solve_episodes (int): 连续解决问题的轮数
    
    Returns:
        scores: 每轮得分列表
        losses: 损失值列表
    """
    scores = []
    losses = []
    recent_scores = deque(maxlen=solve_episodes)
    
    for episode in range(n_episodes):
        state = env.reset()
        total_reward = 0
        episode_losses = []
        
        for step in range(max_steps):
            # 选择动作
            action = agent.act(state)
            
            # 执行动作
            next_state, reward, done, _ = env.step(action)
            
            # 存储经验
            agent.remember(state, action, reward, next_state, done)
            
            # 学习
            loss = agent.replay()
            if loss is not None:
                episode_losses.append(loss)
            
            state = next_state
            total_reward += reward
            
            if done:
                break
        
        scores.append(total_reward)
        recent_scores.append(total_reward)
        
        # 记录平均损失
        if episode_losses:
            losses.append(np.mean(episode_losses))
        else:
            losses.append(0)
        
        # 打印进度
        if episode % 100 == 0:
            avg_score = np.mean(recent_scores)
            print(f"Episode {episode}, Score: {total_reward:.1f}, "
                  f"Avg Score: {avg_score:.1f}, Epsilon: {agent.epsilon:.3f}")
        
        # 检查是否解决问题
        if len(recent_scores) >= solve_episodes:
            avg_score = np.mean(recent_scores)
            if avg_score >= solve_score:
                print(f"\n环境在第 {episode} 轮解决！")
                print(f"最近 {solve_episodes} 轮平均分数: {avg_score:.1f}")
                break
    
    return scores, losses

def evaluate_agent(env, agent, n_episodes=10, render=False):
    """
    评估训练好的智能体
    
    Args:
        env: 环境
        agent: 训练好的智能体
        n_episodes (int): 评估轮数
        render (bool): 是否渲染环境
    
    Returns:
        eval_scores: 评估得分列表
    """
    eval_scores = []
    
    for episode in range(n_episodes):
        state = env.reset()
        total_reward = 0
        
        while True:
            if render:
                env.render()
            
            # 贪婪策略（不探索）
            action = agent.act(state, training=False)
            state, reward, done, _ = env.step(action)
            total_reward += reward
            
            if done:
                break
        
        eval_scores.append(total_reward)
        print(f"评估轮次 {episode + 1}: 得分 = {total_reward}")
    
    return eval_scores

def plot_training_results(scores, losses, window=100):
    """
    绘制训练结果
    
    Args:
        scores: 得分列表
        losses: 损失列表
        window (int): 移动平均窗口大小
    """
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(15, 5))
    
    # 绘制得分曲线
    ax1.plot(scores, alpha=0.3, color='blue', label='原始得分')
    
    # 计算移动平均
    moving_avg = []
    for i in range(len(scores)):
        start_idx = max(0, i - window + 1)
        moving_avg.append(np.mean(scores[start_idx:i+1]))
    
    ax1.plot(moving_avg, color='red', linewidth=2, label=f'{window}轮移动平均')
    ax1.set_xlabel('训练轮次')
    ax1.set_ylabel('得分')
    ax1.set_title('DQN训练过程 - 得分变化')
    ax1.legend()
    ax1.grid(True, alpha=0.3)
    
    # 绘制损失曲线
    ax2.plot(losses, color='orange', alpha=0.7)
    ax2.set_xlabel('训练轮次')
    ax2.set_ylabel('损失值')
    ax2.set_title('DQN训练过程 - 损失变化')
    ax2.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.show()

def plot_q_values_heatmap(agent, env, resolution=20):
    """
    可视化Q值热力图（针对2D状态空间的简化版本）
    
    Args:
        agent: 训练好的智能体
        env: 环境
        resolution (int): 网格分辨率
    """
    # 获取状态范围
    low = env.observation_space.low[:2]  # 只考虑位置和速度
    high = env.observation_space.high[:2]
    
    # 创建网格
    x = np.linspace(low[0], high[0], resolution)
    y = np.linspace(low[1], high[1], resolution)
    X, Y = np.meshgrid(x, y)
    
    # 计算Q值
    q_values = np.zeros((resolution, resolution, agent.action_size))
    
    for i in range(resolution):
        for j in range(resolution):
            # 简化状态（假设角度和角速度为0）
            state = np.array([X[i, j], Y[i, j], 0.0, 0.0])
            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
            
            with torch.no_grad():
                q_vals = agent.q_network(state_tensor).cpu().numpy()[0]
                q_values[i, j] = q_vals
    
    # 绘制热力图
    fig, axes = plt.subplots(1, agent.action_size, figsize=(12, 5))
    actions = ['向左推', '向右推']
    
    for a in range(agent.action_size):
        im = axes[a].imshow(q_values[:, :, a], extent=[low[0], high[0], low[1], high[1]], 
                           origin='lower', cmap='RdYlBu')
        axes[a].set_title(f'Q值: {actions[a]}')
        axes[a].set_xlabel('位置')
        axes[a].set_ylabel('速度')
        plt.colorbar(im, ax=axes[a])
    
    plt.tight_layout()
    plt.show()

def main():
    """主函数"""
    # 创建环境
    env = gym.make('CartPole-v1')
    
    # 获取状态和动作空间大小
    state_size = env.observation_space.shape[0]
    action_size = env.action_space.n
    
    print(f"状态空间维度: {state_size}")
    print(f"动作空间大小: {action_size}")
    
    # 创建智能体
    agent = DQNAgent(
        state_size=state_size,
        action_size=action_size,
        lr=1e-3,
        gamma=0.99,
        epsilon=1.0,
        epsilon_decay=0.995,
        epsilon_min=0.01,
        buffer_size=10000,
        batch_size=64,
        target_update=10
    )
    
    print("开始训练...")
    scores, losses = train_dqn(env, agent, n_episodes=2000)
    
    # 保存模型
    agent.save('dqn_cartpole.pth')
    print("模型已保存到 dqn_cartpole.pth")
    
    # 绘制训练结果
    plot_training_results(scores, losses)
    
    # 评估智能体
    print("\n开始评估...")
    eval_scores = evaluate_agent(env, agent, n_episodes=10)
    print(f"评估平均得分: {np.mean(eval_scores):.2f} ± {np.std(eval_scores):.2f}")
    
    # 可视化Q值
    plot_q_values_heatmap(agent, env)
    
    env.close()

if __name__ == "__main__":
    main()

class ImprovedDQN(nn.Module):
    def __init__(self, state_size, action_size, hidden_size=128):
        super(ImprovedDQN, self).__init__()
        
        # Dueling架构
        self.feature_layer = nn.Sequential(
            nn.Linear(state_size, hidden_size),
            nn.ReLU(),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU()
        )
        
        # 状态价值流
        self.value_stream = nn.Linear(hidden_size, 1)
        
        # 优势流
        self.advantage_stream = nn.Linear(hidden_size, action_size)
    
    def forward(self, state):
        features = self.feature_layer(state)
        
        value = self.value_stream(features)
        advantage = self.advantage_stream(features)
        
        # Dueling aggregation
        q_values = value + (advantage - advantage.mean(dim=1, keepdim=True))
        
        return q_values