| 第382行： |
第382行： |
| | | | |
| | # 训练循环 | | # 训练循环 |
| − | env = gym.make('Pong-v4', frameskip=4)
| + | ... |
| − | action_dim = env.action_space.n
| |
| − | agent = DQNAgent(action_dim)
| |
| − | | |
| − | episodes = 500000
| |
| − | total_rewards = []
| |
| − | | |
| − | for episode in range(episodes):
| |
| − | state, info = env.reset()
| |
| − | state = preprocess(state)
| |
| − | | |
| − | # 初始化 4 帧堆叠
| |
| − | state_stack = deque([state] * 4, maxlen=4)
| |
| − | | |
| − | # 构造当前状态张量：(1,4,84,84)
| |
| − | state_tensor = torch.tensor(np.array(state_stack), dtype=torch.float32).unsqueeze(0).to(device)
| |
| − | | |
| − | total_reward = 0
| |
| − | done = False
| |
| − | | |
| − | while not done:
| |
| − | # 根据策略选择动作
| |
| − | action = agent.select_action(state_tensor)
| |
| − | next_state, reward, done, truncated, info = env.step(action)
| |
| − | done = done or truncated
| |
| − | | |
| − | # 预处理后压入堆叠
| |
| − | next_state = preprocess(next_state)
| |
| − | state_stack.append(next_state)
| |
| − | next_state_tensor = torch.tensor(np.array(state_stack), dtype=torch.float32).unsqueeze(0).to(device)
| |
| − | | |
| − | # 存储到经验回放
| |
| − | agent.memory.push(state_tensor.cpu().numpy(), action, reward, next_state_tensor.cpu().numpy(), done)
| |
| − | | |
| − | # 训练一步
| |
| − | agent.train()
| |
| − | | |
| − | # 更新状态
| |
| − | state_tensor = next_state_tensor
| |
| − | total_reward += reward
| |
| | | | |
| | # 定期更新目标网络 | | # 定期更新目标网络 |
| 第429行： |
第390行： |
| | total_rewards.append(total_reward) | | total_rewards.append(total_reward) |
| | print(f"Episode {episode}: Total Reward = {total_reward}") | | print(f"Episode {episode}: Total Reward = {total_reward}") |
| − | | + | ... |
| − | env.close()
| |
| | </syntaxhighlight> | | </syntaxhighlight> |
| | | | |