跳至主要内容

策略梯度定理详解:上手入门

了解 RL 中的策略梯度定理及其数学推导。基于策略梯度实现算法,在 Gymnasium 中解决一个简单的 RL 环境。
· 15分钟 阅读

使用 AI 探索

强化学习(RL)中的策略梯度是一类直接通过估计“期望回报”相对于策略参数的梯度来优化智能体策略的算法。

在本教程中,我们将解释策略梯度定理及其推导,并展示如何使用 PyTorch 实现策略梯度算法。

什么是策略梯度定理?

在强化学习中,智能体的策略指的是其根据对环境的观测来决定行动的算法。RL 问题的目标是最大化智能体与环境交互所获得的奖励。产生最大奖励的策略就是最优策略。

为了最大化回报,常用的算法大体分为两类:基于策略的方法与基于价值的方法:

  • 基于策略的方法(如策略梯度算法)通过对期望回报做梯度上升直接学习最优策略,而不依赖价值函数。策略以参数化形式表示。当策略由神经网络实现时,策略参数即为网络权重。网络通过对策略参数执行梯度上升来学习最优策略。
  • 基于价值的方法(如 Q-learning)估计状态或状态-动作对的价值。 它们通过选择价值最高的动作间接得到策略。导出最优价值函数的策略即被选为最优策略。贝尔曼方程描述了最优状态价值函数与状态-动作价值函数。

根据策略梯度定理,期望回报的导数等于“回报”与“策略对数(通常表示为概率分布)导数”的乘积的期望值。

策略通常建模为一个参数化函数。当策略由神经网络建模时,策略参数指网络权重。因此,计算期望回报(累积奖励)关于策略参数的梯度,会得到用于更新策略以提升其性能的方向。我们可以用该梯度迭代更新策略参数,使期望回报不断增加。训练过程应当收敛到最大化期望回报的最优策略。

在后续部分,我们将详细解释并推导该定理。

为何使用策略梯度方法?

策略梯度方法的一个关键优势是能处理复杂的动作空间,而传统的基于价值的方法在这方面较为吃力。

处理高维动作空间

基于价值的方法(如 Q-learning)通过对所有可能动作估计价值来工作。当环境的动作空间是连续的或虽离散但规模很大时,这变得困难。

策略梯度方法对策略进行参数化,并估计累积奖励相对于策略参数的梯度。它们利用该梯度直接通过更新参数来优化策略。因此,可以高效处理高维或连续动作空间。策略梯度还是人类反馈强化学习(RLHF)方法的基础。

通过参数化策略并基于梯度调整其参数,策略梯度能够高效应对连续且高维的动作。该直接方式有助于更好的泛化与更灵活的探索,适用于机器人控制等复杂环境中的任务。

学习随机策略

给定一组观测:

  • 确定性策略明确规定智能体采取的动作。
  • 随机策略给出一组可选动作以及智能体选择每个动作的概率。

在随机策略下,同一观测在不同迭代中可能导致选择不同的动作。这促进了对动作空间的探索,并防止策略陷入局部最优。因此,在需要探索以发现通往最大回报路径的环境中,随机策略尤为有用。

在基于策略的方法中,策略输出被转换为一个概率分布,每个可能动作都被分配一个概率。智能体通过从该分布中采样选择动作,从而实现随机策略。因此,策略梯度方法将探索与利用相结合,适合具有复杂奖励结构的环境。

策略梯度定理的推导

在开始推导之前,先明确整个证明中使用的数学记号与关键概念。

数学记号与预备知识

如前文所述,策略梯度定理指出:期望回报的导数等于“回报”与“策略对数的导数”的乘积的期望。

在推导策略梯度定理之前,先引入以下记号:

  • E[X] 表示随机变量 X 的概率期望。
  • 从数学上,策略可表示为一个概率矩阵,按不同观测给出选择不同动作的概率。策略通常建模为一个参数化函数,其参数记为 θ。
    • πθ 表示由 θ 参数化的策略。实践中,这些参数是建模策略的神经网络权重。
  • 轨迹 τ 指从随机初始状态开始到当前时间步或终止状态的一系列状态序列。
  • ∇θf 表示函数 f 对参数 θ 的梯度。
  • J(πθ) 表示智能体遵循策略 πθ 所获得的期望回报。这也是梯度上升的目标函数。
  • 环境在每个时间步根据智能体的动作给出奖励。回报指从初始状态到当前时间步的累积奖励。
    • R(τ) 表示沿轨迹 τ 产生的回报。

推导步骤

我们从目标函数的展开与对数导数技巧出发,自底向上推导并证明策略梯度定理。

目标函数(式 1)

策略梯度方法中的目标函数是沿轨迹累积的回报

J,其基于参数为 θ 的策略 π 得到。该目标函数为:

上述等式中:

  • 左侧(LHS)是遵循策略 πθ.
  • 右侧(RHS)是对每一步中遵循策略 τ (所生成轨迹)的期望,即轨迹上回报 R(τ) 的期望。

目标函数的微分(式 2)

对上式两边(关于 θ)求导得:

期望的梯度(式 3)

右侧的期望可表示为以下两者乘积的积分:

  • 遵循某条轨迹 τ 的概率
  • 沿该轨迹 τ 产生的回报

因此,式 2 的右侧可改写为:

积分的梯度等于梯度的积分。所以我们可将上式中的梯度 ∇θ 移入积分号内。于是右侧变为:

因此,式 2 可改写为:

轨迹的概率(式 4)

接下来更仔细地看一下 P(τ|θ),即在给定策略参数 τ (以及策略 πθ)时,智能体遵循轨迹的概率。轨迹由一系列步骤组成。因此:

  • 获得轨迹 τ 的概率是以下部分的乘积:
    • 每个独立时间步发生的概率。

  • 在时间步 t,智能体从状态 s 转移到状态 st+1,其采取的动作为 at。其发生的概率为以下部分的乘积:
    • 策略在状态 st 预测动作 at
    • 在给定动作 at 和状态 st 的条件下,转移到状态 st+1 的概率

因此,从初始状态 s0 出发,基于策略 πθ,智能体遵循轨迹 τ 的概率为:

为便于处理,我们希望将右侧的乘积表示为求和形式,因此对上式两边取对数:

 

对数概率的导数(式 5)

现在对上式中的对数概率(关于 θ)求导。

在上述等式右侧:

  • 第一项 log ρ0(s0) 与 θ 无关,因此其导数为 0。
  • 求和中的第一项 P(st+1|st, at) 也与 θ 无关,其关于 θ 的导数也为 0。

去掉上述为零的项后,得到(式 5):

回顾式 2:

式 5 计算了式 2 右侧第一部分的对数。我们需要将某项的导数与其对数联系起来,可通过链式法则与对数导数技巧实现。

对数导数技巧

我们稍作绕行,利用微积分规则导出一个结果,用以化简前式并使其适于计算实现。

在微积分中,对数的导数可表示为:

据此变形,可将 x 的导数表示为其对数导数的形式:

这有时被称为对数导数技巧。

链式法则

根据链式法则,给定 z(y) 是 y 的函数,而 y 又是 θ 的函数,即 y(θ),则 z 关于 θ 的导数为:

此处 y(θ) 代表 P(θ),而 z(y) 代表 log(y)。因此,

应用链式法则

我们知道 d(log(y)) / dy = 1/y。将其用于上式右侧第一项。

将 y 移到左侧,并使用 记号:

y 代表 P(θ)。因此上式等价于:

应用对数导数技巧

上述结果给出了式 2(如下所示)右侧第一项:

将该结果代入式 2 右侧,得到:

我们对右侧积分号内的项重新排列如下:

推导最终结果

注意上式包含了期望的积分展开: ∫P(θ)∇logP(θ) = E[∇logP(θ)]

因此,上式右侧可表示为期望:

我们将对数概率的导数代入期望回报的表达式:

在上式中,用式 5 中的 ∇logP(θ) 的值替换,得到:

这就是策略梯度定理给出的回报函数梯度形式。

策略梯度的直观理解

策略梯度方法将策略输出转换为一个概率分布。智能体从该分布中采样以选择动作。策略梯度方法会调整策略参数,从而在每次迭代中更新该概率分布。更新后的分布更倾向于选择能带来更高奖励的动作。

策略梯度算法计算期望回报关于策略参数的梯度。通过沿该梯度方向移动策略参数,智能体在训练过程中提高选择高回报动作的概率。

本质上,带来更好结果的动作在未来被选择的可能性更大,从而逐步改进策略以最大化长期回报。

在 Python 中实现策略梯度

在讨论完策略梯度的基本原理后,我们将展示如何使用 PyTorch 与 Gymnasium 来实现它们。

环境准备

第一步需要安装 gymnasium 以及 NumPy 和 PyTorch 等依赖库。

在服务器或本地机器上安装 gymnasium 及其依赖:

$ pip install gymnasium 

在 Notebook(如 Google Colab 或 DataLab)中安装:

!pip install gymnasium

在 Python 环境中导入这些包:

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torch.distributions as distributions
import numpy as np
import gymnasium as gym  

编写一个简单的策略梯度智能体

使用 .make() 方法创建环境实例。

env = gym.make('CartPole-v1')

与其他机器学习方法类似,我们用神经网络来实现策略梯度智能体。

CartPole-v1 是一个简单环境,因此我们设计一个包含 1 个隐藏层(64 个神经元)的简单网络。输入层维度等于观测空间的维度;输出层维度等于环境动作空间的大小。于是,策略网络从观测状态映射到动作。给定观测作为输入,网络按照策略输出预测的动作。

下面的代码实现了策略网络:

class PolicyNetwork(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, dropout):
        super().__init__()

        self.layer1 = nn.Linear(input_dim, hidden_dim)
        self.layer2 = nn.Linear(hidden_dim, output_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = self.layer1(x)
        x = self.dropout(x)
        x = F.relu(x)
        x = self.layer2(x)
        return x

训练智能体

环境在每个时间步根据智能体的状态与动作给出奖励。策略梯度方法通过对累积奖励(回报)执行梯度下降(或上升)来优化,目标是最大化总回报。

要计算一个回合(episode)的回报,需对该回合所有时间步的奖励进行带折扣的累加。此外,对回报做标准化有助于训练更加平滑稳定。下面的代码展示了如何实现:

def calculate_stepwise_returns(rewards, discount_factor):
    returns = []
    R = 0

    for r in reversed(rewards):
        R = r + R * discount_factor
        returns.insert(0, R)
    returns = torch.tensor(returns)
    normalized_returns = (returns - returns.mean()) / returns.std()
    return normalized_returns

前向传播的每次迭代,我们执行以下步骤:

  • 使用 .step() 函数按当前策略运行智能体。策略在每个时间步预测所选动作的概率。
  • 根据智能体动作从环境接收奖励。
  • 直到到达终止状态,累计每步奖励和动作的对数概率。

下面的代码实现了前向传播:

def forward_pass(env, policy, discount_factor):
    log_prob_actions = []
    rewards = []
    done = False
    episode_return = 0

    policy.train()
    observation, info = env.reset()

    while not done:
        observation = torch.FloatTensor(observation).unsqueeze(0)
        action_pred = policy(observation)
        action_prob = F.softmax(action_pred, dim = -1)
        dist = distributions.Categorical(action_prob)
        action = dist.sample()
        log_prob_action = dist.log_prob(action)

        observation, reward, terminated, truncated, info = env.step(action.item())
        done = terminated or truncated

        log_prob_actions.append(log_prob_action)
        rewards.append(reward)
        episode_return += reward

    log_prob_actions = torch.cat(log_prob_actions)
    stepwise_returns = calculate_stepwise_returns(rewards, discount_factor)

    return episode_return, stepwise_returns, log_prob_actions

使用反向传播与梯度上升更新策略

在传统机器学习中:

  • 损失指预测输出与真实输出之间的差异。
  • 我们用梯度下降最小化损失。

在 RL 中:

  • 损失是用于进行梯度下降(或上升)的代理量。
  • 我们用梯度上升最大化回报(累积奖励)。
  • 期望回报值被用作梯度下降中的“损失”代理。其等于以下两者的乘积:
    • 每一步的期望回报,及
    • 每一步中被采样动作的选择概率。
  • 为使用反向传播做梯度上升,我们取损失的相反数。

下面的代码计算损失:

def calculate_loss(stepwise_returns, log_prob_actions):
    loss = -(stepwise_returns * log_prob_actions).sum()
    return loss

与标准机器学习算法类似,为更新策略,您需要对损失函数执行反向传播。下面的 update_policy() 方法会调用 calculate_loss(),随后对该损失执行反向传播以更新策略参数,即策略网络的模型权重。

def update_policy(stepwise_returns, log_prob_actions, optimizer):
    stepwise_returns = stepwise_returns.detach()
    loss = calculate_loss(stepwise_returns, log_prob_actions)
    
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    return loss.item()

训练循环

我们使用前文定义的函数来训练策略。在开始训练前,需要:

  • 一个未训练的策略,即 PolicyNetwork 类的随机初始化实例。
  • 一个使用 Adam 算法的优化器。
  • 折扣因子、学习率、dropout 比例、奖励阈值和最大训练轮数等超参数。

我们迭代训练循环,直到平均回报超过奖励阈值。每次迭代执行以下步骤:

  • 对每个回合运行一次前向传播。收集该回合的动作对数概率、每步回报以及总回报,并将回合回报累积到数组中。
  • 使用对数概率与每步回报计算损失。在损失上运行反向传播,并用优化器更新策略参数。
  • 检查 N_TRIALS 个回合的平均回报是否超过奖励阈值。

下面的代码实现了这些步骤:

def main(): 
    MAX_EPOCHS = 500
    DISCOUNT_FACTOR = 0.99
    N_TRIALS = 25
    REWARD_THRESHOLD = 475
    PRINT_INTERVAL = 10
    INPUT_DIM = env.observation_space.shape[0]
    HIDDEN_DIM = 128
    OUTPUT_DIM = env.action_space.n
    DROPOUT = 0.5

    episode_returns = []

    policy = PolicyNetwork(INPUT_DIM, HIDDEN_DIM, OUTPUT_DIM, DROPOUT)

    LEARNING_RATE = 0.01
    optimizer = optim.Adam(policy.parameters(), lr = LEARNING_RATE)

    for episode in range(1, MAX_EPOCHS+1):
        episode_return, stepwise_returns, log_prob_actions = forward_pass(env, policy, DISCOUNT_FACTOR)
        _ = update_policy(stepwise_returns, log_prob_actions, optimizer)

        episode_returns.append(episode_return)
        mean_episode_return = np.mean(episode_returns[-N_TRIALS:])

        if episode % PRINT_INTERVAL == 0:
            print(f'| Episode: {episode:3} | Mean Rewards: {mean_episode_return:5.1f} |')

        if mean_episode_return >= REWARD_THRESHOLD:
            print(f'Reached reward threshold in {episode} episodes')
            break

通过调用 main() 函数来运行训练程序:

main()

此 DataLab 工作簿包含了上述策略梯度算法的实现。您可以直接运行,或作为修改算法的起点。

策略梯度方法的优势与挑战

策略梯度方法具有多项优势,例如:

  • 处理连续动作空间:基于价值的方法(如 Q-learning)在连续动作空间中效率不高,因为需要在整个动作空间上估值。策略梯度方法可直接利用期望回报的梯度优化策略,适用于连续动作分布。因此,策略梯度适合如基于连续动作空间的机器人控制等任务。
  • 随机策略:策略梯度方法可以学习随机策略——为每个可能的动作给出被选择的概率。这使智能体能尝试多样化的动作,降低陷入局部最优的风险。它有助于在复杂环境中探索动作空间以找到最优策略。随机性有助于在探索(尝试新动作)与利用(选择已知最优动作)之间取得平衡,这对具有不确定性或稀疏奖励的环境至关重要。
  • 直接优化策略:策略梯度直接优化策略,而非借助价值函数。在连续或高维动作空间中,为每个动作近似估值的计算成本很高,因此基于策略的方法在此类环境中表现良好。

尽管优势明显,策略梯度方法也存在一些固有挑战:

  • 梯度估计方差高:策略梯度方法通过对概率分布采样来选择动作,本质上通过采样轨迹来估计期望回报。由于采样具有随机性,后续迭代中的回报估计可能具有较高方差。这会导致智能体学习效率下降,因为策略更新在不同迭代间可能大幅波动。
  • 训练不稳定:
    • 策略梯度方法对超参数较为敏感(如学习率)。学习率过高会导致策略参数更新过大,从而错过最优参数;反之过小则收敛缓慢。
    • 策略梯度需要在探索与利用间取得平衡。探索不足可能无法到达最优策略的邻域;探索过多则难以收敛,可能在动作空间附近震荡。
  • 采样效率低:策略梯度方法通过从初始到终止完整执行策略并累积每步奖励来估计回报。因此需要与环境进行大量交互以获取大量轨迹样本,对于状态或动作空间很大的环境,这既低效又代价高。

稳定性的解决方案

由于训练不稳定在策略梯度方法中较为常见,开发者采用了多种方法来稳定训练过程。下面介绍常见的稳定化方案:

使用基线函数

由于采样效率不高,训练迭代中估计回报的梯度方差可能很大,导致训练不稳定且缓慢。降低方差的常用方法是使用基线函数,如优势演员-评论家(A2C)。其思想是用一个代理(优势函数)替代目标函数中的回报估计。

优势定义为:采样轨迹的实际回报与给定初始状态的期望回报之差。该方法利用价值函数表示状态和状态-动作对的期望值。将损失表示为“实际回报 − 期望回报”(而非仅回报)可降低损失函数及其梯度的方差,使训练更稳定。

使用熵正则化

在某些环境(如稀疏奖励,仅极少状态给出奖励)中,策略会很快变得确定化,并变得贪心地利用已探索过的路径。这会抑制进一步探索,常导致收敛到局部最优和次优策略。

解决方案是鼓励探索:当策略过于确定化时给予惩罚。这可通过在目标函数中加入基于熵的项来实现。熵度量策略中的随机性:熵越大,智能体选择动作的随机性越强。该熵项等于熵系数与当前策略熵的乘积。

将熵纳入目标函数有助于在利用与探索之间取得平衡。

策略梯度的扩展

在众多策略梯度扩展中,最基础的一种是 REINFORCE 算法。它为策略梯度定理提供了直接实现,并成为更高级技术的基础。

REINFORCE 算法

REINFORCE(也称蒙特卡罗强化)是策略梯度定理的基本实现之一。它使用蒙特卡罗方法估计回报与策略梯度。遵循 REINFORCE 时,智能体直接从环境中采样整条轨迹(从初始到终止状态)的所有动作。这与 TD 学习和动态规划等方法形成对比,后者基于价值函数估计进行自举。

下面给出 REINFORCE 算法的基本步骤:

  • 用随机参数初始化策略
  • 重复多个训练回合。对每个回合:
    • 按以下方式生成整条回合的每一步:
      • 将状态输入策略函数。
      • 策略函数为每个可能动作生成概率。
      • 从该概率分布中随机采样一个动作。
    • 对回合中的每个状态,估计直至该步的回报(折扣累积奖励)。
    • 根据策略梯度定理估计目标函数的梯度,表示为每步回报与每步动作概率的乘积。
    • 应用梯度更新策略参数

对于每个策略,您可以采样单条轨迹来估计梯度(如上所示),也可以在同一策略下对多条轨迹的梯度取平均。

演员-评论家方法

演员-评论家方法将基于策略的方法(如 REINFORCE)与价值函数结合起来。

  • 演员与策略梯度方法类似。演员实现策略,在每一步基于策略选择动作,并沿期望回报的梯度更新策略。
  • 评论家实现价值函数,作为基线(见上一节)。这有助于提升训练效率与稳定性。

像 REINFORCE 这样的策略梯度方法使用原始回报沿轨迹估计梯度。由于轨迹由采样过程得到,回报与梯度可能具有较大方差。用优势函数替代原始回报可以缓解该问题。优势函数等于实际回报与期望回报(即价值函数)之差。演员-评论家是一类算法,当评论家用优势函数实现(最常见做法)时,也称为优势演员-评论家(A2C)。

近端策略优化(PPO)

在复杂环境中,仅靠 A2C 等演员-评论家方法不足以控制回报与梯度的方差。在此情形下,限制每次迭代中策略变化的幅度会有所帮助。这会强制更新后的策略(梯度上升后)位于旧策略的邻域内。

如近端策略优化(PPO)等方法对策略梯度做了两点修改:

  • 使用优势函数。通常以价值函数作为基线,在这一点上与 A2C 类似。
  • 限制每次迭代中策略参数的变化幅度。通过裁剪的代理目标函数实现:算法规定新旧策略之比必须落在某个区间内,当该比值(梯度更新后)超过预设范围时,将其裁剪到区间内。

因此,PPO 相比基础策略梯度显著提升了稳定性,尤其在复杂环境中。裁剪目标函数可防止因回报与梯度方差过大而导致的策略更新不稳定。为在探索与利用间取得平衡,还可在 PPO 中引入熵正则化:向目标函数添加熵项(策略熵乘以缩放系数)。

最新进展

策略梯度是解决 RL 问题的最早方法之一。随着 GPU 的发展,出现了多种将现代机器学习技术应用于策略梯度的新方法。

梯度提升强化学习

近年来,人们在将梯度提升等方法用于 RL 算法方面取得了进展。梯度提升通过组合多个弱模型的预测来生成一个强模型,这被称为梯度提升强化学习(GBRL)。GBRL 是一个与 XGBoost 类似的 Python 包,用于在 RL 算法中实现这些技术。

迁移强化学习

迁移学习(TL)是一种将一个模型学到的知识用于提升另一个模型性能的技术。由于从零开始训练 ML 模型的成本很高,迁移学习非常有用。TL 方法已被用于策略梯度以提升 RL 模型性能,这种方法称为迁移强化学习(TRL)。

总结

策略梯度是解决 RL 问题最基础的方法之一。

本文从第一性原理出发介绍了策略梯度,并推导了策略梯度定理。我们还演示了如何在 Gymnasium 环境中使用 PyTorch 实现一个简洁的基于梯度的算法。最后,我们讨论了该基本算法在实践中的挑战以及常见扩展。

如果您希望进一步学习强化学习与基于 PyTorch 的深度学习,欢迎查看以下课程:

Arun Nanda

作者

Arun Nanda

Arun 曾是一位初创公司创始人,喜欢构建新事物。他目前正探索人工智能的技术与数学基础。他乐于分享自己的所学,因此会写作记录。

除了 DataCamp,您还可以在 Medium、Airbyte 和 Vultr 阅读他的文章。

主题