从REINFORCE到PPO
从REINFORCE到PPO
梯度估计
在RL中,考虑最简单的情况,直接取当前状态的期望reward之和作为loss,那么最终目标就是尝试最大化这个loss。形式化地说,记
- \(\tau: (S_{0}, A_{0}, R_{0}, S_{1}, A_{1}, R_{1},\dots)\)代表model的行为序列,其中\(S\)代表当前所处的状态, \(A\)代表当前状态下采取的动作, \(R\)表示动作\(A\)之后得到的奖励
- \(\pi(a|s:\theta)\)代表在状态\(s\)下模型采用动作\(a\)的概率。策略\(\pi\)关于\(\theta\)是可微的
- \(Q_{\pi}(s, a)\)代表在当前状态为\(s\), 采取动作为\(a\)时的期望收益
- \(P((s', r)|(a,s))\)代表当前状态为\(s\),采取动作为\(a\)时, 转移到状态\(s'\),获得reward为\(r\)的概率。需要注意的是,这是一个与model,或者说策略\(\pi\)无关的概率分布,它只与交互的环境有关
那么在当前状态为\(s\)时loss可以表达为 \[ loss(s) = V_{\pi}(s) = \sum_{a} \pi(a|s:\theta)Q_{\pi}(s,a) \] 它既是loss,也代表了处于状态\(s\)时的期望收益 \(Q_{\pi}(s,a)\)包含了reward部分以及model的决策部分,前者关于参数\(\theta\)是不可导的,后者与\(\theta\)有关。如果直接将整个\(Q_{\pi}(s,a)\)看成常数的话,loss求梯度倒是很方便了, 但是显然是不够合理的。我们需要对其进行进一步的拆分。
我们可以继续展开为 \[ V_{\pi}(s) = \sum_{a}\pi(a|s:\theta) \sum_{s', r} P(s',r|a, s)(V_{\pi}(s')+ r) \] 我们成功把reward \(r\) 拆出来了,它与\(\theta\)无关,我们求个导试试 得到 \[ \nabla V_{\pi}(s) = \sum_{a}\nabla\pi(a|s:\theta)Q_{\pi}(s, a) + \pi(a|s:\theta)\sum_{s'}P(s'|a, s)\nabla V_{\pi}(s') \] 纵观整个\(\nabla V_{\pi}(s)\),唯一有问题的还剩下\(\nabla V_{\pi}(s')\),它里面还是有r,我们无法显式地得到其\(\nabla\)的封闭形式。没办法,只能展开 \[ \begin{flalign} \nabla V_{\pi}(s) &= \sum_{a}(\nabla\pi(a|s:\theta)Q_{\pi}(s, a) + \pi(a|s:\theta)\sum_{s'}P(s'|a, s))\cdot\\ &[\sum_{a'}\nabla\pi(a'|s':\theta)Q_{\pi}(s', a') + \pi(a'|s':\theta)\sum_{s''}P(s''|a', s')] \cdot \dots \end{flalign} \] 方便起见,以下记\(\pi(a'|s':\theta)\)为\(\pi(a'|s')\) 对于最终表达式,每一项\(\sum_{a}\nabla \pi(a_{k}|s_{k})Q_{\pi}(s_{k},a_{k})\)的系数为 \[ \sum_{a_{0}} \pi(a_{0}|s_{0})\sum_{s_{1}}P(s_{1}|a_{0},s_{0})\sum_{a_{1}}\pi(a_{1}|s_{1})\sum_{s_{2}}P(s_{2}|a_{1},s_{1})\dots\sum_{a_{k-1}}\pi(a_{k-1}|s_{k-1})\sum_{s_{k}}P(s_{k}|a_{k-1},s_{k-1}) \] 可以简化系数为 \[ \sum_{s_{0},s_{1},\dots s_{k-1}}P(s_{1}|s_{0})P(s_{2}|s_{1})\dots P(s_{k}|s_{k-1}) \] 它实际代表了从状态\(s_{0}\)经过k次决策之后到达状态\(s_{k}\)的概率,可以简记为 \[ P(s_{0}\rightarrow x,k, \pi) \] 这里\(x\)就代表状态,即\(s_{k}\)
回到\(\nabla V_{\pi}(s)\),现在其可以整理为 \[ \nabla V_{\pi}(s) = \sum_{x\sim \pi}\sum_{k=0}^{\infty} P(s\rightarrow x,k, \pi) \sum_{a}\nabla \pi(a|x)Q(x,a) \] \[ \nabla V_{\pi}(s) = \sum_{x\sim \pi} \eta(x) \sum_{a}\nabla \pi(a|x)Q(x,a) \] 实际消除了r,得到了一个类封闭形式。但是\(\eta(x)\)实际无法计算
归一化 \[ \mu(x) = \frac{\eta(x)}{\sum_{x_{i}\sim \pi} \eta(x_{i})} \] 代表了状态\(x\)在策略\(\pi\)下所有行为序列中出现的概率 在同一个策略\(\pi\)下\(\sum_{x_{i}\sim \pi}\eta(x_{i})\)是一个定值,系数统一除以一个定值,实际上计算得到的梯度与原梯度只是相差一个常数,这个常数可以由步长参数进行调整
最终, 在初始状态为\(s\)的情况下,得到对loss梯度的估计为 \[ \nabla loss(s) = \nabla V_{\pi}(s) \propto \sum_{x\sim \pi} \mu(x) \sum_{a}\nabla \pi(a|x)Q(x,a) = E_{x\sim \pi}[\sum_{a}\nabla \pi(a|x)Q(x,a)] \]
与一开始的式子进行对比,发现它与直接将\(Q\)看成常数的效果非常相似,只是外面多了个期望
REINFORCE
在\(\nabla loss\)中对a也进行估计
\[ \begin{flalign} \nabla loss(x) &= E_{x\sim \pi}[\sum_{a}\nabla \pi(a|x)Q(x,a)]\\ &= E_{x\sim \pi}[\sum_{a}\pi(a|x)\frac{\nabla \pi(a|x)}{\pi(a|x)} Q(x,a)]\\ &= E_{x\sim \pi}[\sum_{a}\pi(a|x)\nabla \log\pi(a|x) \ Q(x,a)]\\ &= E_{x, a\sim \pi}[\nabla \log\pi(a|x) \ Q(x,a)]\\ \end{flalign} \] 采样估计
\[ \nabla loss = \nabla \log\pi(a_{t}|s_{t}) \ Q(s_{t},a_{t}) \]
在实际操作中,\(Q(s_{t},a_{t})\)的估计直接由之后每一个阶段的奖励加权求和得到, 一般也记为\(G_{t}\),换句话说,\(G_{t}\)是\(Q(s_{t},a_{t})\)的一个估计, 具体估计形式可以见下面的算法伪代码
由于这里对\(Q(s_{t}, a_{t})\)的估计需要该步骤之后的奖励信息,因此这个算法的参数更新必须在整个决策轨迹都生成之后才能进行,不支持在线更新
重新整理一下整个过程中蕴含的一些直觉: 首先对rl背景进行建模,得到\(\nabla loss\)的一般表达式: \[ \nabla V_{\pi}(s) = \sum_{a}\nabla\pi(a|s:\theta)Q_{\pi}(s, a) + \pi(a|s:\theta)\sum_{s'}P(s'|a, s)\nabla V_{\pi}(s') \] 然后通过无限展开对上式的最后一项\(\nabla V_{\pi}(s')\)进行估计,最后可以将对奖励的求梯度操作全部转移到决策轨迹中每一个状态上对模型的求梯度操作。然后再利用采样操作来近似这个值即可,得到结果 \[ \nabla loss(x) = E_{(a_{t}, s_{t})\sim \pi}[\nabla \log\pi(a_{t}|s_{t}) \ G(t)]\\ \]
减小估计的方差
REINFORCE with baseline
REINFORCE这种方法存在方差过大的问题:更新不稳定,难以收敛 方差来源于对\(Q(s_{t}, a_{t})\)的估计\(G(t)\),而不是前者\(\nabla \log\pi(a_{t}|s_{t})\),这是精确计算的gradient
\[ \begin{flalign} Var[\nabla \log\pi(a_{t}|s_{t}) \ G(t)] &= E[(\nabla \log\pi(a_{t}|s_{t}) \ G(t))^2] - E[\nabla \log\pi(a_{t}|s_{t}) \ G(t)]^2\\ \end{flalign} \] 后一项是对实际梯度的无偏估计,对于方差没有影响,主要是前一项。对于\(Var[G(t)]\)来说,轨迹过长, Var会不断累加 \[ Var[G(t)] = Var[a_{t+1}R_{t+1}+ \dots + a_{T}R_{T}] = \sum_{R_{i}}Var[a_{i}R_{i}] \]
考虑引入一个baseline,记为\(b(s_{t})\) \[ \nabla loss' = \nabla \log\pi(a_{t}|s_{t}) \ ( G(t) - b(s_{t})) \] 这一引入不会改变估计的无偏性,这是因为 \[ \sum_{a}b(s_{t})\nabla_{\theta} \log \pi(a_{t}|s_{t}) = b(s)\nabla_{\theta}\sum_{a}\log \pi(a_{t}|s_{t}) = b(s)\nabla_{\theta}1 = 0 \]
从而 \[ \begin{flalign} Var[\nabla \log\pi(a_{t}|s_{t})(G(t)-b(s_{t}))] &= E[(\nabla \log\pi(a_{t}|s_{t})(G(t)-b(s_{t})))^2] - E[\nabla \log\pi(a_{t}|s_{t})(G(t)-b(s_{t}))]^2\\ \end{flalign} \]
同理,后一项对方差没有影响,我们只需要关注前一项 \[ E[(\nabla \log\pi(a_{t}|s_{t})(G(t)-b(s_{t})))^2] \] 为了减少方差,目标是选择\(b(s)\)来最小化 \[ \mathbb{E}[(G(t)-b(s_{t}))^2] = \mathbb{E}[G(t)^2] - 2\mathbb{E}[G(t)b(s_{t})] + \mathbb{E}[b(s_{t})^2]\ \ \ \ (1) \]
注意到\(\mathbb{E}[G(t)] = V(s_{t})\),从而 \[ \frac{\partial}{\partial b(s_{t})}\mathbb{E}[(G(t)-b(s_{t}))^2] = -2V(s_{t}) + 2b(s_{t}) \] 从而梯度最小值在\(b(s_{t})\)取为\(V(s_{t})\)的时候取到,此时\(b(s_{t})\)代表的实际含义就是当前状态下的期望收益。但是因为在状态\(s_{t}\)的时候\(V(s_{t})\)实际上也是未知的,所以可以考虑类似\(G(t)\)的方式来拟合,或者使用神经网络来拟合。 至于引入baseline的有效性:原始的REINFORCE就是取\(b(s_{t})\)为0时候的情况,显然其对应的方差是大于\(v(s_{t})\)取最优解\(V(s_{t})\)时候的情况的,具体可以参考式子(1)
从而我们可以得到更新后的梯度表达式为 \[ \nabla loss(x) = E_{(a_{t}, s_{t})\sim \pi}[\nabla \log\pi(a_{t}|s_{t}) \ A(t)]\\ \] 其中\(A(t) = G(t)- b(s_{t})\),其含义可以理解为在状态\(s_{t}\)下,采用动作\(a_{t}\)时可以获得的收益,相比状态\(s_{t}\)的期望收益的差值,也称为优势。 ## actor-critic model
\(b(s)\)的选择很重要,因此考虑由另外一个mode来进行拟合,这个model拟合的内容就是在状态\(s\)下期望获得的收益,称为critic,而\(\pi\)就被称为actor 而\(( Q(s_{t},a_{t}) - b(s_{t}))\)被称为优势,它表明了当前决策相比原本决策会获得的额外收益,是更多还是更少。这个值的正负直接决定了优化的方向是鼓励还是抑制
上述梯度更新思路仍然存在一个问题,就是对样本的利用效率太低了。每次更新的时候需要等待一整轮决策生成完毕,并且只用一次。可以使用重要性采样来进行优化
\[ \large \begin{flalign} \nabla loss(x) &= E_{(a_{t}, s_{t})\sim \pi}[\nabla \log\pi(a_{t}|s_{t}) \ A(t)]\\ &= E_{(a_{t}, s_{t})\sim \pi_{\theta_{old}}}[\frac{\pi(s_{t},a_{t})}{\pi_{old}(s_{t},a_{t})}\nabla \log\pi(a_{t}|s_{t}) \ A(t)]\\ &= E_{(a_{t}, s_{t})\sim \pi_{old}}[\frac{\pi(s_{t},a_{t})}{\pi_{old}(s_{t},a_{t})}\nabla \log\pi(a_{t}|s_{t}) \ A(t)]\\ \end{flalign} \]
可以进一步优化如下 \[ \large \begin{flalign} \nabla loss(x) &= E_{(a_{t}, s_{t})\sim \pi_{old}}[\frac{\pi(s_{t},a_{t})}{\pi_{\theta_{old}}(s_{t},s_{t})}\nabla \log\pi(a_{t}|s_{t}) \ A(t)]\\ &= E_{(a_{t}, s_{t})\sim \pi_{old}}[\frac{\pi(a_{t}|s_{t})}{\pi_{old}(a_{t}|s_{t})}\frac{1}{\pi(a_{t}|s_{t})}\nabla \pi(a_{t}|s_{t}) \ A(t)]\\ &= E_{(a_{t}, s_{t})\sim \pi_{old}}[\frac{\nabla \pi(a_{t}|s_{t})}{\pi_{old}(a_{t}|s_{t})}A(t)]\ \ \ \ (2)\\ \end{flalign} \] 上述变形保持估计的无偏,但是估计的方差会随着策略\(\pi\)与\(\pi_{old}\)的比值 \[ \frac{\pi(a_{t}|s_{t})}{\pi_{old}(a_{t}|s_{t})}\ \ \ (3) \] 而正比变化。这是因为link(之后再补充)
我们需要让\((2)\)式尽可能小,也就是让更新尽可能的平缓,由此可以引入KL Loss来对两者得到的分布的距离进行限制
从而得到新的\(loss\)为 \[ \large loss = E_{(a_{t}, s_{t})\sim \pi_{old}}[\frac{\pi(a_{t}|s_{t})}{\pi_{old}(a_{t}|s_{t})}A(t) - \lambda KL(\pi_{old}(a_{t}|s_{t}),\pi(a_{t}|s_{t}))]\ \ \ (4) \] 原本用cross entropy更合理,但是\(\pi_{old}\)的分布相对当前参数策略\(\pi\)来说是固定的,所以可以直接减掉,那就用KL Loss就可以了。
PPO
但是实际计算KL Loss开销还是太大了,不现实,所以经验性地直接对\(\large r_{t} = \frac{\pi(a_{t}|s_{t})}{\pi_{old}(a_{t}|s_{t})}\)进行clip操作,将其限制在\([1-\epsilon,1+\epsilon]\)之间
得到loss为 \[ \large loss = E_{(a_{t}, s_{t})\sim \pi_{old}}[min(r_{t}A(t) , clip(r_{t}, 1-\epsilon, 1+\epsilon)A(t))]\ \ \ (4) \]
这就是最终的PPO的式子了。