PyTorch中DQN结合PER的重要性权重采样实现与收敛问题排查

核心错误点 #

• Target计算错误：计算TD目标时错误地索引了rewards[0]和(1-terminal)[0]，导致整个batch的所有样本都复用了第一个样本的奖励和终止标志，完全破坏了目标值的正确性，直接去掉索引即可。
• 重要性采样权重使用错误：IS权重是用来对每个样本的损失做加权，而不是直接乘以Q值。原写法smooth_l1_loss(weights*Q_s, targets)逻辑完全错误，需要先求每个样本的损失再乘权重后平均。
• Beta退火逻辑错误：当前beta更新代码取max(self.beta, self.step_count/self.epochs)，初始beta为0.5，训练前期step_count/epochs远小于0.5，beta会一直卡在初始值不增长，无法完成退火消除偏差的作用，应该改为逐步递增到1的逻辑。
• 优先级更新潜在报错：如果模型运行在GPU上，直接对张量调用.numpy()会报错，需要先移动到CPU再转换。
• 超参不合理：CartPole-v0环境下gamma设为0.999过高，会导致值估计波动大，建议改为0.99。

修复后的learn函数核心代码 #

def learn(self):
    # 修正beta退火逻辑，从初始值逐步升到1，可替换为符合你全局步计数的递增逻辑
    self.beta = min(1.0, self.step_count/self.epochs)
    
    # Sample random transitions
    samples = self.memory.sample(self.batch_size, self.beta)
    indexes = np.array(samples['indexes'])
    weights = samples['weights']

    states = self.memory.data['state'][indexes]
    actions = self.memory.data['action'][indexes]
    rewards = self.memory.data['reward'][indexes]
    terminal = self.memory.data['done'][indexes]
    next_states = self.memory.data['next_state'][indexes]

    states = torch.tensor(states, dtype=torch.float32).to(self.Q_net.device)
    actions = torch.tensor(actions, dtype=torch.int64).unsqueeze(-1).to(self.Q_net.device)
    rewards = torch.tensor(rewards, dtype=torch.float32).unsqueeze(-1).to(self.Q_net.device)
    terminal = torch.tensor(terminal, dtype=torch.float32).unsqueeze(-1).to(self.Q_net.device)
    next_states = torch.tensor(next_states, dtype=torch.float32).to(self.Q_net.device)

    # 修正target计算，去掉错误的[0]索引
    targets = self.T_net.forward(next_states)
    targets = rewards + self.gamma * torch.max(targets, dim=1, keepdim=True)[0] * (1-terminal)

    Q_s = self.Q_net.forward(states)
    Q_s = torch.gather(input=Q_s, dim=1, index=actions)

    # 修正优先级更新，先移到CPU再转numpy
    priorities = torch.abs(targets - Q_s).detach().cpu().numpy() + self.min_sampling_prob
    self.memory.update_priorities(indexes, priorities.squeeze())

    weights = torch.tensor(weights, dtype=torch.float32).to(self.Q_net.device).unsqueeze(-1)

    # 修正损失计算，先求每个样本的损失再乘IS权重，再平均
    elementwise_loss = nn.functional.smooth_l1_loss(Q_s, targets, reduction='none')
    loss = torch.mean(elementwise_loss * weights)

    self.optimizer.zero_grad()
    loss.backward()
    self.optimizer.step()
    self.step_count += 1

测试脚本修改建议 #

将Agent初始化参数中的gamma从0.999改为0.99即可。

内容的提问来源于stack exchange，提问作者AF Grobler