Policy Gradient 极简教程

一、前言

Policy Gradient 是一种非常强大的强化学习算法，正如其名，Policy Gradient 通过直接对 Policy 本身计算梯度来实现学习的目的。

Policy Gradient 算法非常简洁，除了 Loss 的计算，其余和常规深度学习训练差别不大。今天我们聚焦 Loss 的实现（而非数学推导），介绍 Policy Gradient，并解决 CartPole 问题。

二、损失函数

Policy Gradient 的损失函数如下：
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> L ( θ ) = − E [ l o g π ( a ∣ s ; θ ) × A ( s , a ) ] L(θ) = -E[log π(a|s; θ) × A(s,a)] </math>L(θ)=−E[logπ(a∣s;θ)×A(s,a)]

可以看到这个损失就是计算一个期望。期望内部分为两个部分，log π(a|s; θ)表示策略π(·;θ)，在给定状态 s 时，执行动作 a 的概率的对数。而 A(s, a)表示在给定状态 s 的情况下，执行动作 a 的价值。

2.1 对数概率

首先要知道，策略π(·;θ)是一个输出概率分布的神经网络，假设动作空间为 3，则策略会输出 3 个概率值，且概率值和为一。

在我们从策略获取动作时，会伴随一个概率值p，对这个概率执行对数操作就是对数概率了。

下面用伪代码理解这个过程：

python 复制代码

# 根据当前状态，让策略返回动作的概率分布
action_probs = policy(state)
# 获取动作概率分布
dist = torch.distributions.Categorical(action_probs)
# 采样一个动作
action = dist.sample()
# 获取采样出来动作的对数概率
log_prob = log(action_probs[action])

2.2 动作价值

在一个游戏中，我们游玩的整个流程如下：

游戏开始，初始化为状态 s0
在状态 s0 下执行动作 a1，转换到状态 s1，并获得奖励 r1
在状态 s1 下执行动作 a2，转换到状态 s2，并获得奖励 r2
在状态 s2 下执行动作 a3，转换到状态 s3，并获得奖励 r3，并结束游戏

现在来明确一下，动作 a1 和哪些奖励有关？在执行 a1 后，我们陆续获得了 r1、r2、r3 三个奖励，其中 r1 是直接相关，而 r2、r3 是间接相关，并且 r2 与 a1 的相关性更高。

由此我们可以算出动作a1 的价值：
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> A ( s 0 , a 1 ) = r 1 + γ ( r 2 + γ r 3 ) A(s0, a1) = r_1 + \gamma (r_2 + \gamma r_3) </math>A(s0,a1)=r1+γ(r2+γr3)

同理可以计算动作 a2、a3 的价值：
<math xmlns="http://www.w3.org/1998/Math/MathML" display="block"> A ( s 1 , a 2 ) = r 2 + γ r 3 A ( s 2 , a 3 ) = r 3 A(s1, a2) = r_2 + \gamma r_3 \\ A(s2, a3) = r_3 </math>A(s1,a2)=r2+γr3A(s2,a3)=r3

这部分计算代码如下：

python 复制代码

# 在完成一次游戏后，会得到n个rewards
rewards = [...]
discounted_rewards = []
R = 0
for r in reversed(rewards):
    R = r + GAMMA * R
    discounted_rewards.insert(0, R)

有了这两部分的理解，后面的内容就简单了。

三、Policy Gradient的实现

下面我们使用Policy Gradient解决 CartPole 问题。

整个流程非常简单，就是初始化环境和 Policy，然后利用现有的 Policy 玩一遍游戏，并收集 log_prob和 rewards，然后再利用已有数据更新 Policy，然后一直循环即可。

3.1 创建Policy网络

CartPole 是一个相对简单的问题，因此我们使用简单的 MLP 网络，代码如下：

python 复制代码

import torch
from torch import nn


class PolicyNetwork(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(PolicyNetwork, self).__init__()
        self.fc1 = nn.Linear(input_dim, 128)
        self.fc2 = nn.Linear(128, output_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return F.softmax(self.fc2(x), dim=-1)

这里的 input_dim 和 output_dim 稍后从环境中获取。

3.2 运行游戏

运行游戏需要环境和 Policy 两个东西，因此我们定义一个接收两个参数的函数 run_episode，另外更新 Policy 需要 log_probs 和 rewards，因此我们需要收集这两个东西并返回。完整代码如下：

ini 复制代码

def run_episode(env, policy):
    state = env.reset()[0]
    done = False
    log_probs = []
    rewards = []

    while not done:
        # 选择 action
        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)
        action_probs = policy(state_tensor)
        dist = torch.distributions.Categorical(action_probs)
        action = dist.sample()

        # 执行 action
        next_state, reward, done, _, _ = env.step(action.item())

        # 计算 log prob
        log_prob = dist.log_prob(action)
        log_probs.append(log_prob)
        rewards.append(reward)

        state = next_state
    return log_probs, rewards

这里需要注意我们必须使用 Policy 来采样动作，否则收集的数据无法用于更新 Policy。

3.3 更新Policy

更新 Policy 则是将 rewards 转换成动作价值，然后代入 Policy Gradient 的损失函数即可。完整代码如下：

python 复制代码

def update_policy(optimizer, log_probs, rewards):
    discounted_rewards = []
    R = 0
    for r in reversed(rewards):
        R = r + GAMMA * R
        discounted_rewards.insert(0, R)

    # 归一化
    discounted_rewards = torch.FloatTensor(discounted_rewards).to(device)
    discounted_rewards = (discounted_rewards - discounted_rewards.mean()) / (discounted_rewards.std() + 1e-7)

    policy_loss = []
    for log_prob, reward in zip(log_probs, discounted_rewards):
        policy_loss.append(-log_prob * reward)

    optimizer.zero_grad()
    loss = torch.stack(policy_loss).sum()
    loss.backward()
    optimizer.step()

在这里我们计算动作价值后，对动作价值进行归一化，以稳定训练。

-log_prob * reward是计算损失的核心部分。常规情况下我们是执行梯度下降，而在 Policy Gradient 中，我们的损失函数说明了回报的期望，而我们的目的是希望回报最大化，因此这里使用负号改成梯度上升。

3.4 整合

最后，我们将整个流程整合：

ini 复制代码

import numpy as np
import torch
from torch import nn
import torch.nn.functional as F
import gymnasium as gym

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
if torch.backends.mps.is_available():
    device = torch.device("mps")
# 超参数
GAMMA = 0.99
LEARNING_RATE = 0.01
EPISODES = 3000
RENDER_EVERY = 100


...


def train():
    # env = gym.make('CartPole-v1', render_mode='human')
    env = gym.make('CartPole-v1')
    input_dim = env.observation_space.shape[0]
    output_dim = env.action_space.n
    
    # 初始化 Policy
    policy = PolicyNetwork(input_dim, output_dim).to(device)
    optimizer = torch.optim.Adam(policy.parameters(), lr=LEARNING_RATE)

    scores = []

    for episode in range(EPISODES):
        # 玩一轮游戏
        log_probs, rewards = run_episode(env, policy)
        # 更新Policy
        update_policy(optimizer, log_probs, rewards)
        
        score = sum(rewards)
        scores.append(score)

        print(f"Episode {episode + 1}, Score: {score:.2f}")
        
        # 当分数达到某个阈值时，认为已经解决了该问题
        if len(scores) > 50 and np.mean(scores[-50:]) > 200:
            print("Environment solved!")
            break

    env.close()
    torch.save(policy.cpu().state_dict(), 'policy.pth')


if __name__ == "__main__":
    train()

五、总结

Policy Gradient 是一种非常简洁且强大的强化学习算法，使用 Policy Gradient 可以很快解决 CartPole 问题。另外在只需修改几行代码的情况下，Policy Gradient 还可以解决其他许多问题，感兴趣的读者可以尝试用上述代码解决更多gymnasium中的问题。