强化学习实验一CartPole车杆平衡—— 最简单的入门任务描述一根杆子竖在一辆小车上你要左右移动小车让杆子尽可能长时间保持直立。状态空间4个连续值小车位置、速度、杆子角度、角速度动作空间2个离散动作左移、右移奖励每坚持一帧得1分推荐网络模型浅层全连接网络MLPimport gymnasium as gymimport torchimport torch.nn as nnimport torch.optim as optimimport numpy as npfrom collections import dequeimport randomfrom tqdm import trange# 1. 定义 Q 网络 class DQN(nn.Module):def __init__(self, state_dim, action_dim, hidden_dim128):super(DQN, self).__init__()self.net nn.Sequential(nn.Linear(state_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, hidden_dim),nn.ReLU(),nn.Linear(hidden_dim, action_dim))def forward(self, x):return self.net(x)# 2. 经验回放缓冲区 class ReplayBuffer:def __init__(self, capacity):self.buffer deque(maxlencapacity)def push(self, state, action, reward, next_state, done):self.buffer.append((state, action, reward, next_state, done))def sample(self, batch_size):batch random.sample(self.buffer, batch_size)states, actions, rewards, next_states, dones zip(*batch)return (torch.FloatTensor(np.array(states)),torch.LongTensor(np.array(actions)).unsqueeze(1),torch.FloatTensor(np.array(rewards)).unsqueeze(1),torch.FloatTensor(np.array(next_states)),torch.FloatTensor(np.array(dones)).unsqueeze(1))def __len__(self):return len(self.buffer)# 3. DQN Agent class DQNAgent:def __init__(self, state_dim, action_dim, lr1e-3, gamma0.99,epsilon_start1.0, epsilon_end0.01, epsilon_decay0.995,buffer_capacity10000, batch_size64, target_update_freq100):self.action_dim action_dimself.gamma gammaself.epsilon epsilon_startself.epsilon_end epsilon_endself.epsilon_decay epsilon_decayself.batch_size batch_sizeself.target_update_freq target_update_freqself.step_count 0# 主网络和目标网络self.q_net DQN(state_dim, action_dim)self.target_net DQN(state_dim, action_dim)self.target_net.load_state_dict(self.q_net.state_dict())self.target_net.eval()self.optimizer optim.Adam(self.q_net.parameters(), lrlr)self.loss_fn nn.MSELoss()self.replay_buffer ReplayBuffer(buffer_capacity)def select_action(self, state, trainingTrue):ε-贪婪策略选择动作if training and random.random() self.epsilon:return random.randrange(self.action_dim)else:with torch.no_grad():state torch.FloatTensor(state).unsqueeze(0)q_values self.q_net(state)return q_values.argmax().item()def update(self):if len(self.replay_buffer) self.batch_size:return None# 从缓冲区采样states, actions, rewards, next_states, dones self.replay_buffer.sample(self.batch_size)# 计算当前 Q 值current_q self.q_net(states).gather(1, actions)# 计算目标 Q 值使用目标网络with torch.no_grad():next_q self.target_net(next_states).max(1, keepdimTrue)[0]target_q rewards self.gamma * next_q * (1 - dones)# 计算损失并更新loss self.loss_fn(current_q, target_q)self.optimizer.zero_grad()loss.backward()self.optimizer.step()# 衰减 εself.epsilon max(self.epsilon_end, self.epsilon * self.epsilon_decay)# 定期更新目标网络self.step_count 1if self.step_count % self.target_update_freq 0:self.target_net.load_state_dict(self.q_net.state_dict())return loss.item()# 4. 训练主循环 def train_dqn(env_nameCartPole-v1, num_episodes500, renderFalse):env gym.make(env_name, render_modehuman if render else None)state_dim env.observation_space.shape[0]action_dim env.action_space.nagent DQNAgent(state_dim, action_dim)episode_rewards []for episode in trange(num_episodes, descTraining):state, _ env.reset()total_reward 0done Falsetruncated Falsewhile not (done or truncated):action agent.select_action(state)next_state, reward, done, truncated, _ env.step(action)# 保存经验agent.replay_buffer.push(state, action, reward, next_state, done or truncated)# 更新网络loss agent.update()state next_statetotal_reward rewardepisode_rewards.append(total_reward)# 每 20 个 episode 打印一次进度if (episode 1) % 50 0:avg_reward np.mean(episode_rewards[-50:])print(fEpisode {episode1}, Avg Reward (last 50): {avg_reward:.2f}, Epsilon: {agent.epsilon:.3f})env.close()return episode_rewards# 5. 测试训练好的模型 def test_agent(agent, env_nameCartPole-v1, episodes5):env gym.make(env_name, render_modehuman)for episode in range(episodes):state, _ env.reset()total_reward 0done Falsetruncated Falsewhile not (done or truncated):action agent.select_action(state, trainingFalse)state, reward, done, truncated, _ env.step(action)total_reward rewardprint(fTest Episode {episode1}: Total Reward {total_reward})env.close()# 6. 运行入口 if __name__ __main__:# 训练rewards train_dqn(CartPole-v1, num_episodes300, renderFalse)# 创建 agent 并加载训练好的参数这里直接复用训练时的 agent# 如果想单独测试需要重新实例化 agent 并加载权重# 这里我们直接在训练后测试env gym.make(CartPole-v1)state_dim env.observation_space.shape[0]action_dim env.action_space.nenv.close()test_agent DQNAgent(state_dim, action_dim)# 注实际使用时需要保存和加载模型权重这里简化处理# 你可以用 torch.save(agent.q_net.state_dict(), dqn_cartpole.pth) 保存print(\n训练完成)
