策略梯度

# 前言

在学习大模型的训练过程中发现在微调之后需要用到强化学习,常见的方法如近端策略优化(PPO)。在查询资料后发现在没有系统性学习强化学习的基础知识的前提下想直接学习PPO会非常困难。于是我决定先放下大模型训练的学习,转而学习强化学习。

在查询了相关教程后,发现了两本比较适合入门的书籍,都有配套的在线教程。分别是Easy RL 强化学习教程 (opens new window)和动手学强化学习 (opens new window)。目前我只学习了在线教程,前者循序渐进通俗易懂,适合入门打基础,后者用来配合学习源码和复现。

# 关键词

  • 动作:演员采取的动作,一般记作aa。
  • 状态:演员和环境交互后的状态,一般记作ss。
  • 策略:策略决定了演员后续采取的动作,一般记作π\pi。
  • 回报:一个回合得到的所有奖励总和。一般记作RR。
  • 轨迹:一个回合里演员采取的所有动作aa和状态ss的集合被称为轨迹,即τ={s1,a1,s2,a2,...,st,at}\tau=\{s_1,a_1,s_2,a_2,...,s_t,a_t\}。
  • 奖励函数:表示在某个状态下采取某个动作后可得到的奖励分数,是一个函数。一般记作rr。把所有的rr加起来,就是轨迹的奖励,记作R(τ)R(\tau)。
  • 期望奖励:R¯=∑τR(τ)pθ(τ)=Eτpθ[R(τ)]\bar{R}=\sum_{\tau}R(\tau)p_{\theta}(\tau)=E_{\tau}~p_{\theta}[R(\tau)],即所有轨迹出现的概率和其奖励的乘积的累加值。
  • 蒙特卡洛采样:基于随机抽样的方法,用于估计值函数或者策略函数的期望回报。
  • REINFORCE:是一种使用蒙特卡洛采样来估计策略函数的梯度,使最大化长期奖励的期望的算法,是蒙特卡洛策略梯度算法中最简单的特例。

# 策略梯度公式推导

强化学习的目的是最大化期望奖励,可以使用梯度上升算法,需要计算期望奖励\=R的梯度。

∇Rθ¯=∑τR(τ)∇pθ(τ)\nabla\bar{R_{\theta}}=\displaystyle\sum_{\tau}R(\tau)\nabla{p_\theta}(\tau)

利用对数函数的求导公式,f(x)=logxf(x)=\log{x}的导数为1x\frac{1}{x}。

∇pθ(τ)pθ(τ)=∇logpθ(τ)\frac{\nabla{p_\theta}(\tau)}{p_{\theta}(\tau)}=\nabla\log{p_\theta}(\tau)

∇Rθ¯=∑τR(τ)∇pθ(τ)=∑τR(τ)pθ(τ)∇pθ(τ)pθ(τ)=Eτpθ(τ)[R(τ)∇logpθ(τ)]\begin{aligned} \nabla\bar{R_{\theta}} &=\displaystyle\sum_{\tau}R(\tau)\nabla{p_\theta}(\tau)\\ &=\displaystyle\sum_\tau{R(\tau)}p_\theta(\tau)\frac{\nabla{p_\theta}(\tau)}{p_{\theta}(\tau)}\\ &=E_{\tau}~p_{\theta}(\tau)[R(\tau)\nabla\log{p_\theta}(\tau)] \end{aligned}

期望Eτpθ(τ)[R(τ)∇logpθ(τ)]E_{\tau}~p_{\theta}(\tau)[R(\tau)\nabla\log{p_\theta}(\tau)]实际无法计算,因为无法采样所有轨迹,所以用采样的方式采样N个轨迹可以近似求解。

Eτpθ(τ)[R(τ)∇logpθ(τ)]≈1N∑n=1NR(τ)∇logpθ(τ)\begin{aligned} E_{\tau}~p_{\theta}(\tau)[R(\tau)\nabla\log{p_\theta}(\tau)]&\approx\frac{1}{N}\displaystyle\sum_{n=1}^NR(\tau)\nabla\log{p_\theta}(\tau) \end{aligned}

又因为某个轨迹τ\tau的概率为

pθ(τ)=p(s1)pθ(a1∣s1)p(s2∣s1,a1)pθ(a2∣s2)p(s3∣s2,a2)...=p(s1)∏t=1Tpθ(at∣st)p(st+1∣st,at)\begin{aligned} p_\theta(\tau) &=p(s_1)p_\theta(a_1|s_1)p(s_2|s_1,a_1)p_\theta(a_2|s_2)p(s_3|s_2,a_2)...\\ &=p(s_1)\displaystyle\prod_{t=1}^Tp_\theta(a_t|s_t)p(s_{t+1}|s_t,a_t) \end{aligned}

所以∇logpθ(τ)\nabla\log{p_\theta(\tau)}可变为

∇logpθ(τ)=∇logp(s1)∏t=1Tpθ(at∣st)p(st+1∣st,at)=∇logs1+∇∑t=1Tlogpθ(at∣st)+∇∑t=1Tlogp(st+1∣st,at)=∑t=1T∇logpθ(at∣st)\begin{aligned} \nabla\log{p_\theta(\tau)} &=\nabla\log{p(s_1)\displaystyle\prod_{t=1}^Tp_\theta(a_t|s_t)p(s_{t+1}|s_t,a_t)} \\ &=\nabla\log{s_1}+\nabla\displaystyle\sum_{t=1}^T\log{p_\theta(a_t|s_t)}+\nabla\displaystyle\sum_{t=1}^T\log{p(s_{t+1}|s_t,a_t)} \\ &=\displaystyle\sum_{t=1}^T\nabla\log{p_\theta(a_t|s_t)} \end{aligned}

于是公式(2)可变为

Eτpθ(τ)[R(τ)∇logpθ(τ)]≈1N∑n=1NR(τ)∇logpθ(τ)=1N∑n=1N∑t=1TnR(τn)∇logpθ(atn∣stn)\begin{aligned} E_{\tau}~p_{\theta}(\tau)[R(\tau)\nabla\log{p_\theta}(\tau)]&\approx\frac{1}{N}\displaystyle\sum_{n=1}^NR(\tau)\nabla\log{p_\theta}(\tau)\\ &=\frac{1}{N}\displaystyle\sum_{n=1}^N\sum_{t=1}^{T_n}R(\tau^n)\nabla\log{p_\theta(a_t^n|s_t^n)} \end{aligned}

梯度上升是一个迭代优化的过程,迭代时更新参数θ\theta,同时需要调整的还有学习率η\eta。每次更新参数时,参数向量沿着梯度方向移动一小步,增加目标函数的值。

θ←θ+η∇Rθ¯\theta\gets\theta+\eta\nabla\bar{R_\theta}

以上是最基本的策略梯度公式。

# 策略梯度公式优化

# 添加基线

公式(5)中的轨迹奖励存在一个问题,在部分场景中奖励总是正的,比如足球赛进球得分,失球并不扣分,这会导致每个动作的概率都会得到提升。所以我们可以把奖励减去一个基线bb,这样使奖励低于基线的动作的概率会降低,超过的概率会提高。

∇Rθ¯≈1N∑n=1N∑t=1Tn(R(τn)−b)∇logpθ(atn∣stn)\nabla\bar{R_\theta}\approx\frac{1}{N}\displaystyle\sum_{n=1}^N\sum_{t=1}^{T_n}(R(\tau^n)-b)\nabla\log{p_\theta(a_t^n|s_t^n)}

关于bb的取值,可以取奖励的平均值,在训练过程中,记录下奖励并计算平均值即bb。

# 调整分数

轨迹奖励只计算当前动作执行之后的奖励,而不是像之前计算所有奖励。因为在执行当前动作之前发生的事情和执行当前动作是无关的,所以之前的奖励不能算是当前动作的贡献。同时对未来的奖励打折扣,因为当前动作执行后会影响未来的结果,当时随着时间往后拖,该动作的影响变小。在R前面乘以一个折扣因为γ\gamma。

∇Rθ¯≈1N∑n=1N∑t=1Tn(∑t′=tTnγt′−trt′n−b)∇logpθ(atn∣stn)\nabla\bar{R_\theta}\approx\frac{1}{N}\displaystyle\sum_{n=1}^N\sum_{t=1}^{T_n}(\displaystyle\sum_{t'=t}^{T_n}\gamma^{t'-t}r_{t'}^n-b)\nabla\log{p_\theta(a_t^n|s_t^n)}

# 优势函数

R-b这一项被称为优势函数,用Aθ(st,at)A^\theta(s_t,a_t)表示。它的含义是在某个状态下执行某个动作,相比其他动作,它有多好。这里的优势其实是相对优势,因为减去了基线b。优势函数通常可以由网络估计出来,这个网络称为评论员(critic)。

# 代码实现

reinforce (opens new window)