本文深入浅出地介绍了Policy Gradient强化学习算法，重点关注其核心损失函数，并详细阐述了在CartPole问题中的具体实现。文章首先解释了Policy Gradient的原理，然后通过代码示例展示了如何构建策略网络、运行游戏、收集数据并更新策略，最终成功解决了CartPole问题。文章还提供了完整的代码实现，方便读者理解和实践。

💡Policy Gradient 是一种直接优化策略的强化学习算法，其核心在于通过计算策略梯度来更新策略。

🔑损失函数是Policy Gradient的关键，它基于动作的对数概率和动作价值的乘积来计算期望。

🕹️在CartPole问题中，首先需要构建一个策略网络，该网络输出动作的概率分布。

🔄通过运行游戏，收集每个动作的对数概率和奖励，然后计算动作价值，并利用这些数据更新策略网络。

📈更新策略时，需要将奖励转换为动作价值，并进行归一化处理，以稳定训练过程。

一、前言

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

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

二、损失函数

Policy Gradient 的损失函数如下：

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

2.1 对数概率

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

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

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

# 根据当前状态，让策略返回动作的概率分布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 的价值：

同理可以计算动作 a2、a3 的价值：

这部分计算代码如下：

# 在完成一次游戏后，会得到n个rewardsrewards = [...]discounted_rewards = []R = 0for 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 网络，代码如下：

import torchfrom torch import nnclass 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，因此我们需要收集这两个东西并返回。完整代码如下：

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 的损失函数即可。完整代码如下：

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 整合

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

import numpy as npimport torchfrom torch import nnimport torch.nn.functional as Fimport gymnasium as gymdevice = torch.device("cuda" if torch.cuda.is_available() else "cpu")if torch.backends.mps.is_available():    device = torch.device("mps")# 超参数GAMMA = 0.99LEARNING_RATE = 0.01EPISODES = 3000RENDER_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中的问题。