一、整体概述

此代码利用 REINFORCE 算法（一种基于策略梯度的强化学习算法）来解决 OpenAI Gym 中的 CartPole-v1 环境问题。CartPole-v1 环境的任务是控制一个小车，使连接在小车上的杆子保持平衡。代码通过构建一个神经网络作为策略网络，在与环境的交互中不断学习，以找到能获得最大累计奖励的策略。

二、依赖库

1. gym：OpenAI 开发的强化学习环境库，用于创建和管理各种强化学习任务的环境，这里使用其 CartPole-v1 环境。
2. torch：PyTorch 深度学习框架，用于构建神经网络模型、进行张量运算和自动求导。
3. torch.nn：PyTorch 中用于定义神经网络层和模型结构的模块。
4. torch.optim：PyTorch 中的优化器模块，用于更新神经网络的参数。
5. numpy：用于进行数值计算和数组操作。
6. torch.distributions.Categorical：PyTorch 中用于处理分类分布的模块，用于从策略网络输出的动作概率分布中采样动作。
7. matplotlib.pyplot：用于绘制训练过程中的奖励曲线，可视化训练进度。

三、代码详细解释

3.1 REINFORCE 专用策略网络类 REINFORCEPolicy

收起

python

class REINFORCEPolicy(nn.Module):def __init__(self, input_dim, output_dim):super().__init__()self.net = nn.Sequential(nn.Linear(input_dim, 64),nn.ReLU(),nn.Linear(64, output_dim))def forward(self, x):return self.net(x)

• 功能：定义一个简单的前馈神经网络作为策略网络，用于根据环境状态输出动作的概率分布。
• 参数：
  • input_dim：输入状态的维度，即环境状态的特征数量。
  • output_dim：输出的维度，即环境中可用动作的数量。
• 结构：
  • 包含两个全连接层（nn.Linear），中间使用 ReLU 激活函数（nn.ReLU）引入非线性。
  • 第一个全连接层将输入维度映射到 64 维，第二个全连接层将 64 维映射到输出维度。
• 前向传播方法 forward：接收输入状态 x，并通过定义的网络层计算输出。

3.2 REINFORCE 训练函数 reinforce_train

收起

python

def reinforce_train(env, policy_net, optimizer, num_episodes=1000, gamma=0.99, lr_decay=0.995, baseline=True):...

• 功能：使用 REINFORCE 算法训练策略网络。
• 参数：
  • env：OpenAI Gym 环境对象。
  • policy_net：策略网络模型。
  • optimizer：用于更新策略网络参数的优化器。
  • num_episodes：训练的总回合数，默认为 1000。
  • gamma：折扣因子，用于计算未来奖励的折扣，默认为 0.99。
  • lr_decay：学习率衰减因子，默认为 0.995。
  • baseline：布尔值，指示是否使用基线来降低方差，默认为 True。
• 训练流程：
  1. 数据收集阶段：
     • 每个回合开始时，重置环境状态。
     • 在回合中，不断与环境交互，直到回合结束。
     • 对于每个时间步，将状态转换为张量，通过策略网络得到动作的 logits，使用 Categorical 分布采样动作，并记录动作的对数概率。
     • 执行动作，获取下一个状态、奖励和回合是否结束的信息。
     • 将状态、动作、奖励和对数概率存储在 episode_data 字典中。
  2. 计算蒙特卡洛回报：
     • 从最后一个时间步开始，反向计算每个时间步的累计折扣奖励 G。
     • 将计算得到的回报存储在 returns 列表中，并转换为张量。
  3. 可选基线处理：
     • 如果 baseline 为 True，对回报进行标准化处理，以降低方差。
  4. 计算策略梯度损失：
     • 对于每个时间步的对数概率和回报，计算策略梯度损失。
     • 将所有时间步的损失相加得到总损失。
  5. 参数更新：
     • 清零优化器的梯度。
     • 进行反向传播计算梯度。
     • 使用优化器更新策略网络的参数。
  6. 学习率衰减：
     • 如果 lr_decay 不为 None，每 100 个回合衰减一次学习率。
  7. 记录训练进度：
     • 记录每个回合的总奖励。
     • 每 50 个回合输出一次平均奖励和当前学习率。
     • 如果平均奖励达到环境的奖励阈值，输出解决信息并提前结束训练。

3.3 主程序部分

收起

python

if __name__ == "__main__":...

• 功能：创建环境，初始化策略网络和优化器，进行训练，保存模型，并可视化训练进度。
• 步骤：
  1. 创建环境：使用 gym.make 创建 CartPole-v1 环境，并获取状态维度和动作维度。
  2. 初始化网络和优化器：
     • 创建 REINFORCEPolicy 策略网络实例。
     • 使用 Adam 优化器，设置较高的初始学习率（lr = 1e-2）。
  3. 训练模型：调用 reinforce_train 函数进行训练，设置训练回合数为 800。
  4. 保存模型：使用 torch.save 保存训练好的策略网络的参数。
  5. 可视化训练进度：
     • 使用 matplotlib.pyplot 绘制每个回合的总奖励曲线。
     • 设置 x 轴标签为 “回合数”，y 轴标签为 “总奖励”，标题为 “REINFORCE 训练进度”。
     • 显示绘制的图形。

四、注意事项

• 代码使用了新版 Gym API，确保你的 Gym 库版本支持 env.reset() 和 env.step() 的返回值格式。
• 可以根据实际情况调整超参数，如 num_episodes、gamma、lr_decay 和初始学习率，以获得更好的训练效果。
• 训练可能需要一定的时间，尤其是在计算资源有限的情况下，可以适当减少 num_episodes 来加快训练速度。

完整代码

import gym
import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from torch.distributions import Categorical
import matplotlib.pyplot as plt# REINFORCE专用策略网络
class REINFORCEPolicy(nn.Module):def __init__(self, input_dim, output_dim):super().__init__()self.net = nn.Sequential(nn.Linear(input_dim, 64),nn.ReLU(),nn.Linear(64, output_dim))def forward(self, x):return self.net(x)def reinforce_train(env, policy_net, optimizer, num_episodes=1000, gamma=0.99, lr_decay=0.995, baseline=True):rewards_history = []lr = optimizer.param_groups[0]['lr']for ep in range(num_episodes):# 数据收集阶段state, _ = env.reset()episode_data = {'states': [], 'actions': [], 'rewards': [], 'log_probs': []}done = Falsewhile not done:state_tensor = torch.FloatTensor(state)logits = policy_net(state_tensor)policy = Categorical(logits=logits)action = policy.sample()log_prob = policy.log_prob(action)next_state, reward, terminated, truncated, _ = env.step(action.item())done = terminated or truncated# 存储轨迹数据episode_data['states'].append(state_tensor)episode_data['actions'].append(action)episode_data['rewards'].append(reward)episode_data['log_probs'].append(log_prob)state = next_state# 计算蒙特卡洛回报returns = []G = 0for r in reversed(episode_data['rewards']):G = r + gamma * Greturns.insert(0, G)returns = torch.tensor(returns)# 可选基线（降低方差）if baseline:returns = (returns - returns.mean()) / (returns.std() + 1e-8)# 计算策略梯度损失policy_loss = []for log_prob, G in zip(episode_data['log_probs'], returns):policy_loss.append(-log_prob * G)total_loss = torch.stack(policy_loss).sum()  # 使用 torch.stack() 代替 torch.cat()# 参数更新optimizer.zero_grad()total_loss.backward()optimizer.step()# 学习率衰减if lr_decay:new_lr = lr * (0.99 ** (ep//100))optimizer.param_groups[0]['lr'] = new_lr# 记录训练进度total_reward = sum(episode_data['rewards'])rewards_history.append(total_reward)# 进度输出if (ep+1) % 50 == 0:avg_reward = np.mean(rewards_history[-50:])print(f"Episode {ep+1} | Avg Reward: {avg_reward:.1f} | LR: {optimizer.param_groups[0]['lr']:.2e}")if avg_reward >= env.spec.reward_threshold:print(f"Solved at episode {ep+1}!")breakreturn rewards_historyif __name__ == "__main__":env = gym.make('CartPole-v1')state_dim = env.observation_space.shape[0]action_dim = env.action_space.n# 初始化REINFORCE专用网络policy_net = REINFORCEPolicy(state_dim, action_dim)optimizer = optim.Adam(policy_net.parameters(), lr=1e-2)  # 更高初始学习率# 训练rewards = reinforce_train(env, policy_net, optimizer, num_episodes=800)# 保存与测试（同原代码）torch.save(policy_net.state_dict(), 'reinforce_cartpole.pth')plt.plot(rewards)plt.xlabel('Episode')plt.ylabel('Total Reward')plt.title('REINFORCE Training Progress')plt.show()