Atari Pong DQN实现报错求助及基础代码需求

问题分析 #

你遇到的两个随机触发的ValueError，根源是状态数据维度不统一和未正确处理LazyFrames对象：

• 错误1：采样批量状态时，FrameStack返回的LazyFrames未转换为标准numpy数组，导致生成的数组形状不均匀（32个元素为不同结构的序列）。
• 错误2：训练循环中存在未定义的transitions变量，触发维度不匹配的异常。

关键错误修复 #

1. 修复经验回放缓冲区的采样逻辑 #

LazyFrames无法直接转换为numpy数组，需显式转换并调整维度（从(84,84,4)转为PyTorch卷积要求的(4,84,84)）：

def sample(self, batch_size):
    indices = np.random.choice(len(self.experienceBuffer), batch_size, replace=False)
    states, actions, rewards, dones, next_states = zip(*[self.experienceBuffer[idx] for idx in indices])
    
    # 转换LazyFrames为numpy数组并调整维度
    states = np.array([np.array(s).transpose(2,0,1) for s in states], dtype=np.float32)
    next_states = np.array([np.array(s).transpose(2,0,1) for s in next_states], dtype=np.float32)
    return (
        states,
        np.array(actions, dtype=np.int64),
        np.array(rewards, dtype=np.float32),
        np.array(dones, dtype=np.uint8),
        next_states
    )

2. 删除冗余代码 #

训练循环中以下代码未定义transitions变量，属于冗余内容，直接删除：

obses = np.asarray([t[0] for t in transitions])
actions = np.asarray([t[1] for t in transitions])
rews = np.asarray([t[2] for t in transitions])
dones = np.asarray([t[3] for t in transitions])
new_obses = np.asarray([t[4] for t in transitions])

obses_t = torch.as_tensor(obses, dtype=torch.float32)
actions_t = torch.as_tensor(actions, dtype=torch.int64).unsqueeze(-1)
rews_t = torch.as_tensor(rews, dtype=torch.float32).unsqueeze(-1)
dones_t = torch.as_tensor(dones, dtype=torch.float32).unsqueeze(-1)
new_obses_t = torch.as_tensor(new_obses, dtype=torch.float32)"""

3. 修正Agent类的环境重置逻辑 #

Agent._reset中使用全局env改为实例的self.env，避免全局变量依赖：

def _reset(self):
    self.state = self.env.reset()  # 替换原env.reset()
    self.total_reward = 0.0

4. 调整DQN网络的输入维度适配 #

PyTorch卷积层要求通道维度在前，需确保输入张量维度为(batch, channels, height, width)，已在采样阶段处理，简化forward函数：

def forward(self, x):
    # x已为(batch, 4, 84, 84)，直接归一化后输入卷积
    out = self.conv_net(x / 255.0)
    return self.fc_out(out)

5. 修正卷积输出尺寸计算 #

_n_conv_out中输入张量需匹配实际输入维度：

def _n_conv_out(self):
    # 模拟输入维度：(batch, channels, height, width)
    inputs = torch.randn(1, self.input_channels, 84, 84)
    with torch.no_grad():
        out = self.conv_net(inputs)
    return np.prod(out.shape[1:])  # 忽略batch维度

完整原生实现代码 #

import random
import torch
from torch import nn
import numpy as np
from torch.utils.tensorboard import SummaryWriter
import gym
from stable_baselines3.common.atari_wrappers import FireResetEnv, MaxAndSkipEnv
import time
import itertools
import collections

params = {
    'env_name':         "PongNoFrameskip-v4",
    'stop_reward':      18.0,
    'run_name':         'pong',
    'replay_size':      100000,
    'replay_initial':   10000,
    'target_net_sync':  1000,
    'epsilon_frames':   10**5,
    'epsilon_start':    1.0,
    'epsilon_final':    0.02,
    'learning_rate':    0.0001,
    'gamma':            0.99,
    'batch_size':       32
}

# 环境包装（仅使用基础gym和stable-baselines包装器）
env = gym.make(params['env_name'])
env = MaxAndSkipEnv(env, skip=4)
if "FIRE" in env.unwrapped.get_action_meanings():
    env = FireResetEnv(env)
env = gym.wrappers.ResizeObservation(env, (84, 84))
env = gym.wrappers.GrayScaleObservation(env)
env = gym.wrappers.FrameStack(env, 4)

Experience = collections.namedtuple('Experience', field_names=['state', 'action', 'reward', 'done', 'new_state'])

class DQNetwork(nn.Module):
    def __init__(self, env):
        super().__init__()
        self.input_channels = env.observation_space.shape[2]  # FrameStack后通道数在最后
        self.output_size = env.action_space.n

        self.conv_net = nn.Sequential(
            nn.Conv2d(self.input_channels, 32, 8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, 4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, 3, stride=1),
            nn.ReLU(),
            nn.Flatten()
        )

        conv_out_size = self._n_conv_out()

        self.fc_out = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, self.output_size)
        )

    def _n_conv_out(self):
        # 模拟输入维度：(batch, channels, height, width)
        inputs = torch.randn(1, self.input_channels, 84, 84)
        with torch.no_grad():
            out = self.conv_net(inputs)
        return np.prod(out.shape[1:])

    def forward(self, x):
        # x维度：(batch, 4, 84, 84)
        out = self.conv_net(x / 255.0)
        return self.fc_out(out)
  
def calc_loss(batch, net, tgt_net, device="cpu", GAMMA=params['gamma']):
    states, actions, rewards, dones, next_states = batch

    states_v = torch.from_numpy(states).to(device)
    next_states_v = torch.from_numpy(next_states).to(device)
    actions_v = torch.from_numpy(actions).to(device)
    rewards_v = torch.from_numpy(rewards).to(device)
    done_mask = torch.BoolTensor(dones).to(device)

    state_action_values = net(states_v).gather(1, actions_v.unsqueeze(-1)).squeeze(-1)
    next_state_values = tgt_net(next_states_v).max(1)[0]
    next_state_values[done_mask] = 0.0
    next_state_values = next_state_values.detach()

    expected_state_action_values = next_state_values * GAMMA + rewards_v
    return nn.MSELoss()(state_action_values, expected_state_action_values)

class Agent:
    def __init__(self, env, exp_buffer):
        self.env = env
        self.exp_buffer = exp_buffer
        self._reset()

    def _reset(self):
        self.state = self.env.reset()
        self.total_reward = 0.0

    def play_step(self, net, epsilon=0.0, device="cpu"):
        done_reward = None

        if np.random.random() < epsilon:
            action = self.env.action_space.sample()
        else:
            # 将LazyFrames转为numpy并调整维度
            state_np = np.array(self.state).transpose(2,0,1)[np.newaxis, ...]
            state_v = torch.tensor(state_np, dtype=torch.float32).to(device)
            q_vals_v = net(state_v)
            _, act_v = torch.max(q_vals_v, dim=1)
            action = int(act_v.item())

        new_state, reward, ter, trunc, _ = self.env.step(action)
        is_done = ter or trunc
        self.total_reward += reward

        exp = Experience(self.state, action, reward, is_done, new_state)
        self.exp_buffer.append(exp)
        self.state = new_state
        if is_done:
            done_reward = self.total_reward
            self._reset()
        return done_reward

class ExperienceBuffer:
    def __init__(self, capacity):
        self.experienceBuffer = collections.deque(maxlen=capacity)
    
    def __len__(self):
        return len(self.experienceBuffer)
    
    def append(self, exp):
        self.experienceBuffer.append(exp)
    
    def sample(self, batch_size):
        indices = np.random.choice(len(self.experienceBuffer), batch_size, replace=False)
        states, actions, rewards, dones, next_states = zip(*[self.experienceBuffer[idx] for idx in indices])
        
        # 转换LazyFrames为numpy数组并调整维度为(4,84,84)
        states = np.array([np.array(s).transpose(2,0,1) for s in states], dtype=np.float32)
        next_states = np.array([np.array(s).transpose(2,0,1) for s in next_states], dtype=np.float32)
        
        return (
            states,
            np.array(actions, dtype=np.int64),
            np.array(rewards, dtype=np.float32),
            np.array(dones, dtype=np.uint8),
            next_states
        )

# 训练流程
writer = SummaryWriter(f"runs/DqnPong-v1")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    
dqn_network = DQNetwork(env).to(device)
optimizer = torch.optim.Adam(dqn_network.parameters(), lr=params['learning_rate'])
tgt_net = DQNetwork(env).to(device)
tgt_net.load_state_dict(dqn_network.state_dict())

replay_buffer = ExperienceBuffer(capacity=params['replay_size'])
total_rewards = []

ts = time.time()
agent = Agent(env, replay_buffer)
best_mean_reward = None
ts_frame = 0

for step in itertools.count():
    epsilon = np.interp(step, [0, params['epsilon_frames']], [params['epsilon_start'], params['epsilon_final']])
    reward = agent.play_step(dqn_network, epsilon=epsilon, device=device)

    if reward is not None:
        total_rewards.append(reward)
        speed = (step - ts_frame) / (time.time() - ts)
        ts_frame = step
        ts = time.time()
        mean_reward = np.mean(total_rewards[-100:])
        print("%d: done %d games, mean reward %.3f, eps %.2f, speed %.2f f/s" % (
            step, len(total_rewards), mean_reward, epsilon, speed
        ))
        writer.add_scalar("epsilon", epsilon, step)
        writer.add_scalar("speed", speed, step)
        writer.add_scalar("reward_100", mean_reward, step)
        writer.add_scalar("reward", reward, step)
        
        if best_mean_reward is None or best_mean_reward < mean_reward:
            torch.save(dqn_network.state_dict(), params['env_name'] + "-best.pt")
            if best_mean_reward is not None:
                print("Best mean reward updated %.3f -> %.3f, model saved" % (best_mean_reward, mean_reward))
            best_mean_reward = mean_reward
        
        if mean_reward > params['stop_reward']:
            print(f"Solved in {step} frames!")
            break

    if len(replay_buffer) < params['replay_initial']:
        continue
    
    if step % params['target_net_sync'] == 0:
        tgt_net.load_state_dict(dqn_network.state_dict())

    optimizer.zero_grad()
    batch = replay_buffer.sample(params['batch_size'])
    loss_t = calc_loss(batch, dqn_network, tgt_net, device=device)
    loss_t.backward()
    optimizer.step()

writer.close()

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