4. Prioritized Experience Replay实现

任务目标

在DQN中实现优先经验回放（PER），以TD误差值作为采样优先级，在CartPole-v1环境中与均匀采样基线进行对比，验证PER对训练收敛速度的影响。

核心实现

以下代码实现了一个完整的对比实验：NaiveDQN使用均匀回放，PERDQN使用SumTree数据结构存储经验并按TD误差比例采样。两者共享网络结构与超参数，仅在回放机制上不同，确保对比公平。

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from collections import deque
import random

# ---------- SumTree（PER核心数据结构） ----------
class SumTree:
    def __init__(self, capacity):
        self.capacity = capacity
        self.tree = np.zeros(2 * capacity - 1)
        self.data = np.zeros(capacity, dtype=object)
        self.write_idx = 0
        self.size = 0

    def _propagate(self, idx, delta):
        parent = (idx - 1) // 2
        self.tree[parent] += delta
        if parent != 0:
            self._propagate(parent, delta)

    def update(self, idx, priority):
        delta = priority - self.tree[idx]
        self.tree[idx] = priority
        self._propagate(idx, delta)

    def add(self, priority, data):
        idx = self.write_idx + self.capacity - 1
        self.data[self.write_idx] = data
        self.update(idx, priority)
        self.write_idx = (self.write_idx + 1) % self.capacity
        self.size = min(self.size + 1, self.capacity)

    def get(self, s):
        idx = 0
        while True:
            left = 2 * idx + 1
            if left >= len(self.tree):
                break
            if s <= self.tree[left]:
                idx = left
            else:
                s -= self.tree[left]
                idx = left + 1
        data_idx = idx - self.capacity + 1
        return idx, self.tree[idx], self.data[data_idx]

    def total(self):
        return self.tree[0]

# ---------- 网络 ----------
class QNet(nn.Module):
    def __init__(self, obs_dim, act_dim):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(obs_dim, 64), nn.ReLU(),
            nn.Linear(64, 64), nn.ReLU(),
            nn.Linear(64, act_dim)
        )
    def forward(self, x):
        return self.fc(x)

# ---------- 经验回放基类 ----------
class ReplayBuffer:
    def __init__(self, capacity):
        self.capacity = capacity
        self.buffer = deque(maxlen=capacity)
    def push(self, transition):
        self.buffer.append(transition)
    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        return map(np.array, zip(*batch))
    def __len__(self):
        return len(self.buffer)

# ---------- PER回放 ----------
class PERBuffer:
    def __init__(self, capacity, alpha=0.6, beta=0.4, beta_increment=0.001, eps=1e-6):
        self.tree = SumTree(capacity)
        self.alpha = alpha
        self.beta = beta
        self.beta_increment = beta_increment
        self.eps = eps
        self.max_priority = 1.0

    def push(self, transition):
        self.tree.add(self.max_priority ** self.alpha, transition)

    def sample(self, batch_size):
        batch = []
        idxs = []
        segment = self.tree.total() / batch_size
        for i in range(batch_size):
            s = random.uniform(segment * i, segment * (i + 1))
            idx, p, data = self.tree.get(s)
            idxs.append(idx)
            batch.append(data)
        self.beta = min(1.0, self.beta + self.beta_increment)
        probs = np.array([self.tree.tree[idx] for idx in idxs]) / self.tree.total()
        weights = (self.tree.size * probs) ** (-self.beta)
        weights /= weights.max()
        return map(np.array, zip(*batch)), idxs, weights

    def update_priorities(self, idxs, td_errors):
        for idx, td in zip(idxs, td_errors):
            priority = (abs(td) + self.eps) ** self.alpha
            self.tree.update(idx, priority)
            self.max_priority = max(self.max_priority, priority)

    def __len__(self):
        return self.tree.size

# ---------- DQN训练器（支持两种回放） ----------
class DQN:
    def __init__(self, env, replay, use_per=False, gamma=0.99, lr=1e-3, batch_size=64):
        self.env = env
        self.replay = replay
        self.use_per = use_per
        self.gamma = gamma
        self.batch_size = batch_size
        obs_dim = env.observation_space.shape[0]
        act_dim = env.action_space.n
        self.q_net = QNet(obs_dim, act_dim)
        self.target_net = QNet(obs_dim, act_dim)
        self.target_net.load_state_dict(self.q_net.state_dict())
        self.optimizer = optim.Adam(self.q_net.parameters(), lr=lr)
        self.loss_fn = nn.MSELoss()

    def act(self, obs, epsilon):
        if random.random() < epsilon:
            return self.env.action_space.sample()
        with torch.no_grad():
            q = self.q_net(torch.FloatTensor(obs).unsqueeze(0))
        return q.argmax().item()

    def train_step(self):
        if len(self.replay) < self.batch_size:
            return None
        if self.use_per:
            (obs, act, rew, nxt, done), idxs, weights = self.replay.sample(self.batch_size)
            weights = torch.FloatTensor(weights).unsqueeze(1)
        else:
            obs, act, rew, nxt, done = self.replay.sample(self.batch_size)
            weights = torch.ones(self.batch_size, 1)
        obs = torch.FloatTensor(obs)
        act = torch.LongTensor(act).unsqueeze(1)
        rew = torch.FloatTensor(rew).unsqueeze(1)
        nxt = torch.FloatTensor(nxt)
        done = torch.FloatTensor(done).unsqueeze(1)

        q = self.q_net(obs).gather(1, act)
        with torch.no_grad():
            q_next = self.target_net(nxt).max(1, keepdim=True)[0]
            target = rew + self.gamma * q_next * (1 - done)
        td_error = (q - target).detach().numpy().flatten()
        loss = (self.loss_fn(q, target) * weights).mean()

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()

        if self.use_per:
            self.replay.update_priorities(idxs, td_error)
        return loss.item()

    def sync_target(self):
        self.target_net.load_state_dict(self.q_net.state_dict())

# ---------- 训练循环 ----------
def train(use_per, episodes=400, capacity=10000):
    env = gym.make('CartPole-v1')
    replay = PERBuffer(capacity) if use_per else ReplayBuffer(capacity)
    agent = DQN(env, replay, use_per=use_per)
    rewards = []
    epsilon = 1.0
    epsilon_min = 0.01
    epsilon_decay = 0.995

    for ep in range(episodes):
        obs, _ = env.reset()
        total_reward = 0
        done = False
        while not done:
            action = agent.act(obs, epsilon)
            nxt, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated
            replay.push((obs, action, reward, nxt, done))
            agent.train_step()
            obs = nxt
            total_reward += reward
        epsilon = max(epsilon_min, epsilon * epsilon_decay)
        if ep % 20 == 0:
            agent.sync_target()
        rewards.append(total_reward)
    env.close()
    return rewards

# ---------- 运行对比 ----------
if __name__ == '__main__':
    seed = 42
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)

    print('训练均匀采样DQN...')
    uniform_rewards = train(use_per=False)
    print('训练PER-DQN...')
    per_rewards = train(use_per=True)

    # 输出每50回合的平均回报（用于对比）
    for i in range(0, 400, 50):
        print(f'Episodes {i}-{i+50}: Uniform={np.mean(uniform_rewards[i:i+50]):.1f}, PER={np.mean(per_rewards[i:i+50]):.1f}')

运行与验证

1. 依赖安装：pip install gymnasium torch numpy
2. 运行方式：直接执行上述脚本，将依次训练两个智能体各400回合。
3. 判定标准：观察输出的分段平均回报。若PER-DQN在前期（如0-150回合）的平均回报高于均匀采样DQN，或更早达到200分的稳定水平，则说明PER在该任务上加速了收敛。CartPole-v1的回合最大回报为200。

预期输出（确定性逻辑推论）：由于PER优先回放高TD误差样本，通常在前100-200回合内平均回报增长更快。但受随机种子和超参数影响，两者最终都能收敛到接近200的平均回报。若PER的平均回报曲线整体不低于均匀采样，且前期上升更陡，即可判定PER生效。

关键说明

• PER的alpha控制优先级影响程度（0为均匀采样），beta控制重要性采样修正强度，随训练线性增至1。
• TD误差在训练步骤中计算，用于更新SumTree中对应样本的优先级；新样本以当前最大优先级插入，保证至少被采样一次。
• 目标网络每20回合同步一次，与PER机制解耦，确保对比仅反映回放策略差异。