使用 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()