使用REINFORCE算法强化梯度策略

一、整体概述

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

二、依赖库 gym：OpenAI 开发的强化学习环境库，用于创建和管理各种强化学习任务的环境，这里使用其 CartPole-v1 环境。torch：PyTorch 深度学习框架，用于构建神经网络模型、进行张量运算和自动求导。torch.nn：PyTorch 中用于定义神经网络层和模型结构的模块。torch.optim：PyTorch 中的优化器模块，用于更新神经网络的参数。numpy：用于进行数值计算和数组操作。torch.distributions.Categorical：PyTorch 中用于处理分类分布的模块，用于从策略网络输出的动作概率分布中采样动作。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。 训练流程： 数据收集阶段： 每个回合开始时，重置环境状态。在回合中，不断与环境交互，直到回合结束。对于每个时间步，将状态转换为张量，通过策略网络得到动作的 logits，使用 Categorical 分布采样动作，并记录动作的对数概率。执行动作，获取下一个状态、奖励和回合是否结束的信息。将状态、动作、奖励和对数概率存储在 episode_data 字典中。 计算蒙特卡洛回报： 从最后一个时间步开始，反向计算每个时间步的累计折扣奖励 G。将计算得到的回报存储在 returns 列表中，并转换为张量。 可选基线处理： 如果 baseline 为 True，对回报进行标准化处理，以降低方差。 计算策略梯度损失： 对于每个时间步的对数概率和回报，计算策略梯度损失。将所有时间步的损失相加得到总损失。 参数更新： 清零优化器的梯度。进行反向传播计算梯度。使用优化器更新策略网络的参数。 学习率衰减： 如果 lr_decay 不为 None，每 100 个回合衰减一次学习率。 记录训练进度： 记录每个回合的总奖励。每 50 个回合输出一次平均奖励和当前学习率。如果平均奖励达到环境的奖励阈值，输出解决信息并提前结束训练。 3.3 主程序部分

收起

python

if __name__ == "__main__": ...

功能：创建环境，初始化策略网络和优化器，进行训练，保存模型，并可视化训练进度。步骤： 创建环境：使用 gym.make 创建 CartPole-v1 环境，并获取状态维度和动作维度。初始化网络和优化器： 创建 REINFORCEPolicy 策略网络实例。使用 Adam 优化器，设置较高的初始学习率（lr = 1e-2）。 训练模型：调用 reinforce_train 函数进行训练，设置训练回合数为 800。保存模型：使用 torch.save 保存训练好的策略网络的参数。可视化训练进度： 使用 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 = False while 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 = 0 for r in reversed(episode_data['rewards']): G = r + gamma * G returns.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}!") break return rewards_history if __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()