# 用Python手把手实现REINFORCE算法:从理论到代码的完整指南
如果你已经啃完了策略梯度那一堆让人眼花缭乱的数学公式,感觉脑子里塞满了期望、梯度和各种符号,但一打开编辑器却不知道从何下手,那么这篇文章就是为你准备的。我们跳过那些冗长的理论推导,直接进入实战环节。我会带你从零开始,用Python一步步搭建一个真正能跑起来的REINFORCE算法,解决从环境交互、策略网络设计到梯度计算和训练调试中的所有工程细节。你会发现,把理论变成代码的过程,远比想象中更有趣,也更能加深你对算法本质的理解。
## 1. 环境搭建与核心概念澄清
在动手写代码之前,我们需要先统一“战场”。REINFORCE算法属于策略梯度家族,它的核心思想非常直观:**直接优化策略本身**,而不是像Q-learning那样先去估计价值函数。策略通常用一个参数化的函数(比如神经网络)来表示,我们通过调整这个函数的参数,使得它倾向于选择能带来更高累计回报的动作。
> 注意:本文所有代码示例均基于Python 3.8+,并主要依赖`gymnasium`(OpenAI Gym的维护分支)、`torch`和`numpy`。请确保你的环境已安装这些库。
我们选择`CartPole-v1`作为演示环境。这个环境状态简单(小车位置、速度、杆角度、角速度),动作空间离散(左、右),非常适合作为强化学习的“Hello World”。但别小看它,要想让杆子不倒,算法也得足够聪明才行。
首先,让我们初始化环境并理解其基本结构:
```python
import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.distributions import Categorical
import matplotlib.pyplot as plt
# 创建环境
env = gym.make('CartPole-v1')
print(f"状态空间维度: {env.observation_space.shape}") # 输出: (4,)
print(f"动作空间: {env.action_space}") # 输出: Discrete(2)
```
这里有几个关键点需要立刻明确,它们直接影响后续的代码设计:
1. **On-policy特性**:REINFORCE是严格的**同策略**算法。这意味着用于采样轨迹的策略(行为策略)和我们要优化的策略(目标策略)是同一个。你无法使用旧的、过时的经验回放池,每一批数据用完即弃,必须用最新的策略重新采样。
2. **回合更新**:REINFORCE属于蒙特卡洛方法,它必须等待一个完整的回合(episode)结束,拿到从某个状态开始到结束的所有回报后,才能进行参数更新。这带来了两个影响:一是训练数据利用率低,二是只能处理有终止状态的环境(或能人为设定截断步数)。
3. **高方差**:使用整个回合的累计回报作为动作价值的估计,虽然是无偏的,但方差非常大。这是REINFORCE训练不稳定、收敛慢的根源,我们后续会引入技巧来缓解。
理解了这些,你就知道为什么我们常说REINFORCE是策略梯度算法中最“朴素”的一个。它直接,但也粗暴。接下来,我们就来构建它的核心——策略网络。
## 2. 策略网络的设计与实现
策略网络的任务是:输入当前环境状态,输出每个可选动作的概率分布。对于`CartPole-v1`,状态是4维向量,动作是2维(左或右)。因此,一个最简单的策略网络可以是一个多层感知机(MLP),输出层接上Softmax函数,确保所有动作概率之和为1。
下面是一个经典而实用的策略网络实现:
```python
class PolicyNetwork(nn.Module):
def __init__(self, state_dim, action_dim, hidden_dim=128):
super(PolicyNetwork, self).__init__()
self.fc1 = nn.Linear(state_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, hidden_dim)
self.fc3 = nn.Linear(hidden_dim, action_dim)
self.relu = nn.ReLU()
self.softmax = nn.Softmax(dim=-1)
def forward(self, x):
x = self.relu(self.fc1(x))
x = self.relu(self.fc2(x))
logits = self.fc3(x) # 输出“得分”,未归一化
return logits
def get_action(self, state):
"""
根据状态选择动作,并返回动作、对数概率和熵(用于后续计算)。
这是与环境交互的核心函数。
"""
# 将numpy数组转换为torch张量,并增加一个批次维度
state_tensor = torch.from_numpy(state).float().unsqueeze(0)
logits = self.forward(state_tensor)
# 使用Categorical分布,它内部会应用Softmax
action_distribution = Categorical(logits=logits)
action = action_distribution.sample() # 采样一个动作
# 计算所选动作的对数概率,这是梯度计算的关键
log_prob = action_distribution.log_prob(action)
# 计算分布的熵,可用于鼓励探索(后续会讲)
entropy = action_distribution.entropy()
return action.item(), log_prob, entropy
```
这个类有几个工程实现上的细节值得深究:
* **`logits` vs `probabilities`**:网络最后一层不直接输出概率,而是输出`logits`(原始分数)。这样做在数值上更稳定。`Categorical`类会内部处理Softmax和采样。
* **`log_prob`的重要性**:策略梯度定理最终推导出的更新公式中,梯度项包含了对数概率的梯度 `∇ log π(a|s)`。因此,我们必须记录在状态`s`下选择动作`a`的**对数概率**,而不是概率本身。
* **熵(Entropy)**:熵衡量了概率分布的随机性。熵越大,策略越“不确定”,探索性越强。在损失函数中加入熵的负值作为正则项,可以防止策略过早收敛到次优的确定性策略,这是提高REINFORCE性能的一个关键技巧。
有了策略网络,我们就可以让智能体在环境中尝试一个回合,收集一条轨迹(trajectory)。轨迹数据通常包括状态、动作、奖励、对数概率和熵(可选)。
## 3. 核心训练循环与梯度计算
这是REINFORCE算法的引擎所在。我们将把理论公式 `θ = θ + α * G_t * ∇ log π(a_t|s_t)` 翻译成可运行的PyTorch代码。其中`G_t`是从时刻`t`开始的折扣累计回报。
完整的训练循环包含以下步骤:
1. 用当前策略运行一个回合,收集数据。
2. 回合结束后,计算每个时间步的折扣回报 `G_t`。
3. 构造损失函数:`loss = -sum( G_t * log_prob_t )`。注意负号,因为我们要做梯度**上升**,而PyTorch的优化器默认执行梯度**下降**,所以加负号将最大化问题转化为最小化问题。
4. 反向传播,更新策略网络参数。
让我们看看代码如何实现:
```python
def compute_discounted_returns(rewards, gamma=0.99):
"""
计算折扣累计回报。
rewards: 一个列表,包含一个回合中每一步的即时奖励。
gamma: 折扣因子,范围(0, 1]。gamma=1表示无折扣。
返回: 一个与rewards等长的列表,其中每个元素是对应时间步的G_t。
"""
returns = []
R = 0
# 从后往前计算,因为G_t = r_t + gamma * G_{t+1}
for r in reversed(rewards):
R = r + gamma * R
returns.insert(0, R) # 在列表头部插入,保证顺序
return returns
def train_one_episode(policy_net, optimizer, gamma=0.99, entropy_coef=0.01):
"""
运行一个回合,并用该回合的数据进行一次策略更新。
返回该回合的总奖励(用于监控训练进度)。
"""
state, _ = env.reset()
log_probs = []
entropies = []
rewards = []
done = False
# 步骤1: 采样轨迹
while not done:
action, log_prob, entropy = policy_net.get_action(state)
next_state, reward, terminated, truncated, _ = env.step(action)
done = terminated or truncated
# 存储数据
log_probs.append(log_prob)
entropies.append(entropy)
rewards.append(reward)
state = next_state
total_reward = sum(rewards)
# 步骤2: 计算折扣回报
returns = compute_discounted_returns(rewards, gamma)
returns_tensor = torch.tensor(returns)
# 步骤3: 构造损失函数
# 将log_probs列表堆叠成一个张量
log_probs_tensor = torch.stack(log_probs)
# 基础策略梯度损失
policy_loss = -(log_probs_tensor * returns_tensor).sum()
# 熵正则项(负熵,鼓励探索)
entropy_loss = -torch.stack(entropies).sum() # 注意前面的负号
# 总损失
loss = policy_loss + entropy_coef * entropy_loss
# 步骤4: 反向传播与优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
return total_reward
```
这里有几个**极易出错**的坑点,我结合自己的调试经验重点强调:
* **回报的标准化**:直接使用原始折扣回报`G_t`与对数概率相乘,会因为回报值的量级(特别是回合很长时)导致梯度爆炸或更新步长不合理。一个几乎**必须**采用的技巧是对每个回合的回报进行**标准化**(减去均值,除以标准差),使其均值为0,方差为1。
```python
# 在计算损失前,添加回报标准化
returns = compute_discounted_returns(rewards, gamma)
returns = np.array(returns)
returns = (returns - returns.mean()) / (returns.std() + 1e-8) # 防止除零
returns_tensor = torch.tensor(returns, dtype=torch.float32)
```
这个简单的操作能极大提升训练的稳定性和收敛速度。
* **熵系数的选择**:`entropy_coef`是一个超参数。太大,策略会过于随机,无法有效学习;太小,又起不到鼓励探索的作用。通常从`0.01`开始尝试,并根据训练情况调整。
* **梯度裁剪**:即使标准化了回报,在训练初期策略还很差的时候,也可能产生极端大的梯度。在`optimizer.step()`之前加入梯度裁剪是良好的实践。
```python
torch.nn.utils.clip_grad_norm_(policy_net.parameters(), max_norm=1.0)
```
现在,我们已经有了一个可以工作的REINFORCE实现。但它的性能可能并不理想,训练曲线像坐过山车。接下来,我们就进入更关键的环节:调试与优化。
## 4. 训练调试、可视化与性能优化
把代码跑起来只是第一步,让模型真正学到东西并稳定收敛,才是真正的挑战。我们需要一套工具来监控和分析训练过程。
首先,实现一个简单的训练循环并记录历史数据:
```python
def train(policy_net, optimizer, num_episodes=1000, gamma=0.99, entropy_coef=0.01):
episode_rewards = []
moving_avg_rewards = []
for episode in range(num_episodes):
reward = train_one_episode(policy_net, optimizer, gamma, entropy_coef)
episode_rewards.append(reward)
# 计算最近100个回合的平均奖励,平滑曲线以便观察趋势
if len(episode_rewards) >= 100:
moving_avg = np.mean(episode_rewards[-100:])
else:
moving_avg = np.mean(episode_rewards)
moving_avg_rewards.append(moving_avg)
if (episode + 1) % 50 == 0:
print(f"Episode {episode+1}, Total Reward: {reward:.1f}, Moving Avg (last 100): {moving_avg:.1f}")
# 简单的早停条件:连续30个回合平均奖励大于环境的最大值(CartPole是500)
if len(episode_rewards) >= 30 and np.mean(episode_rewards[-30:]) > 495:
print(f"Solved at episode {episode+1}!")
break
return episode_rewards, moving_avg_rewards
```
运行训练后,最重要的就是绘制学习曲线。一张图能告诉你很多信息:
```python
def plot_training_progress(episode_rewards, moving_avg_rewards):
plt.figure(figsize=(12, 5))
plt.subplot(1, 2, 1)
plt.plot(episode_rewards, alpha=0.6, label='Episode Reward')
plt.plot(moving_avg_rewards, linewidth=2, label='Moving Avg (100 episodes)')
plt.xlabel('Episode')
plt.ylabel('Total Reward')
plt.title('REINFORCE Training Progress on CartPole-v1')
plt.legend()
plt.grid(True, alpha=0.3)
plt.subplot(1, 2, 2)
# 绘制最后100个回合的奖励分布直方图,查看稳定性
if len(episode_rewards) >= 100:
last_100 = episode_rewards[-100:]
plt.hist(last_100, bins=20, edgecolor='black', alpha=0.7)
plt.xlabel('Total Reward')
plt.ylabel('Frequency')
plt.title('Distribution of Last 100 Episodes')
plt.axvline(np.mean(last_100), color='red', linestyle='--', label=f'Mean: {np.mean(last_100):.1f}')
plt.legend()
plt.grid(True, alpha=0.3)
plt.tight_layout()
plt.show()
```
典型的REINFORCE在`CartPole-v1`上的学习曲线可能表现出以下特征,我们可以据此诊断问题:
| 曲线特征 | 可能原因 | 调试建议 |
| :--- | :--- | :--- |
| **奖励毫无增长,始终很低** | 学习率太大导致策略震荡;或太小导致学习停滞;网络结构太简单。 | 尝试降低学习率(如从1e-2调到1e-3, 1e-4);适当增加网络层宽度或深度;检查回报标准化和梯度裁剪是否已实现。 |
| **奖励快速增长后突然崩溃** | 策略因过度更新而“遗忘”了之前学到的好的行为;探索不足,陷入局部最优。 | **引入策略更新约束**,如使用更小的学习率,或采用PPO/TRPO等高级算法;**增大熵系数**,鼓励探索。 |
| **奖励波动极大,方差高** | REINFORCE的蒙特卡洛估计本身方差大。 | **引入基线(Baseline)**,这是最有效的改进。用状态价值函数V(s)作为基线,将更新项改为 `(G_t - V(s_t)) * ∇ log π`,能显著降低方差。这其实就是Actor-Critic的雏形。 |
| **收敛速度非常慢** | 原始REINFORCE数据效率低。 | 尝试**批量更新**:收集多个回合的数据(如10个回合)后,用这批数据的平均梯度更新一次,比单回合更新更稳定。 |
如果引入了基线,我们的损失函数计算就需要调整。假设我们有一个简单的状态价值网络`ValueNetwork`来估计`V(s)`:
```python
# 在损失计算部分,假设我们已经有了state_values(每个状态的价值估计)
advantages = returns_tensor - state_values.detach() # 优势函数,detach避免更新价值网络时影响策略梯度
policy_loss = -(log_probs_tensor * advantages).sum()
```
这个改进版的算法通常被称为 **REINFORCE with Baseline**,它已经非常接近最简单的Actor-Critic了。在实践中,对于`CartPole`这样的简单环境,原始的REINFORCE配合回报标准化和熵正则化,通常能在几百到几千个回合内达到满分(500)。如果达不到,请优先检查回报标准化和梯度裁剪是否实现正确。
## 5. 超越CartPole:挑战更复杂的环境与高级技巧
当你成功在`CartPole-v1`上驯服REINFORCE后,是时候挑战更复杂的领域了,比如`LunarLander-v2`(登月器)或`Acrobot-v1`(倒立双摆)。这些环境的状态和动作空间更复杂,对算法的鲁棒性要求更高。
此时,单纯的REINFORCE可能力不从心。你需要一个更强大的工具箱。下面这个表格对比了从原始REINFORCE到更先进方法的核心改进点:
| 技术/技巧 | 目的 | 实现复杂度 | 对性能的影响 |
| :--- | :--- | :--- | :--- |
| **回报标准化 (Return Normalization)** | 稳定梯度更新幅度,防止爆炸或消失。 | 低 | **巨大**。几乎是必备技巧。 |
| **熵正则化 (Entropy Regularization)** | 鼓励探索,防止策略过早收敛到次优确定性策略。 | 低 | 中等。能有效提升收敛稳定性和最终性能。 |
| **基线 (Baseline)** | 降低梯度估计的方差,加速收敛。 | 中 | **巨大**。REINFORCE with Baseline是质的飞跃。 |
| **广义优势估计 (GAE)** | 在TD(λ)框架下更平滑、偏差-方差权衡更好的优势估计。 | 中高 | 巨大。是现代策略梯度算法(如PPO)的标配。 |
| **信任域/裁剪 (PPO/TRPO)** | 约束每次策略更新的幅度,避免灾难性的大更新。 | 高 | **巨大**。解决了训练不稳定的核心痛点。PPO已成为实际应用的主流。 |
对于想深入下去的开发者,我的建议是:**以REINFORCE with Baseline为起点,逐步实现PPO**。PPO的核心——比例裁剪目标函数——理解起来并不比REINFORCE复杂太多,但效果却有天壤之别。
例如,PPO的损失函数核心部分如下,它通过裁剪来限制新旧策略的差异:
```python
# ratio = π_new(a|s) / π_old(a|s)
ratios = torch.exp(log_probs_new - log_probs_old.detach())
advantages = ... # 计算优势,例如用GAE
surr1 = ratios * advantages
surr2 = torch.clamp(ratios, 1 - clip_epsilon, 1 + clip_epsilon) * advantages
policy_loss = -torch.min(surr1, surr2).mean() # PPO的裁剪目标
```
从REINFORCE到PPO,你走过的路正是深度强化学习从理论到实践、从简单到成熟的缩影。每一次对高方差的克服,每一次对更新步长的约束,都让智能体在探索与利用的钢丝上走得更稳。
最后,别忘了实践出真知。多跑实验,多观察曲线,多调整超参数。记录下你每次改动带来的影响,这比死记硬背公式更有价值。我在调试过程中就曾发现,对于`LunarLander`,将折扣因子`gamma`从0.99调到0.995,同时把熵系数从0.01逐步衰减到0.001,能让学习过程平滑很多。这些细微的调整,正是工程实现中不可或缺的经验。