TRPO详解

使用优势函数进行策略梯度的方法,需要先计算梯度

$g = \mathbb{E}_{p(s,a,k) \sim \tau} \left[ A(s, a) \nabla_{\theta}(log \pi(a \mid s)) \right]$

其中$A(s, a)$需要以某种形式直接或是间接由critic网络给出。暂时不考虑critic的拟合与优化,$\tau$的获取方式需要某种actor策略$\pi_{traj}$。这个策略与需要更新训练的$\pi(a \mid s)$有可能是同样的(on policy),也可能有不同(off policy)。

策略梯度需要一个步长超参,然而在比较陡峭的Loss平面上,步长很难决定。

策略梯度的优化最终目标都是$J = \mathbb{E}_{s_0 \sim p(s_0)}[V_{\pi}(s_o)]$,考虑一次策略的更新$ \pi_{old}(\theta_{old}) \rightarrow \pi_{new}(\theta_{new})$,由于步长的问题,无法保证每次按照$\nabla_{\theta}J$方向更新时,J都能增加。所以尝试找到一种方法,使得$J(\pi) = J(\pi_{old}) + something$能够一直增加。


1.  问题转化

首先需要完成一个看似无意义的问题转化,转化后方便进行后续几个优化操作。

考虑$A_{\pi_{old}}(s, a)=Q(s, a) – V(s) = Q(s, a)-\sum\limits_{a’ \sim \pi_{old}} \pi_{old}(a’ \mid s) Q(s, a’)$

1步td时存在$A_{\pi_{old}}(s, a)= \mathbb{E}_{s’} \left[ R_{s \rightarrow s’}^{a \sim \pi_{old}} + \gamma V_{\pi_{old}}(s’) – V_{\pi_{old}}(s) \right]$

则当采样是从新策略$\tau \mid \pi$获得的trajectory时

$\mathbb{E}_{s_0,a_0,… \mid \pi} \left[ \sum\limits_{t=0}^{\infty} \gamma^t A_{\pi_{old}}(s_t, a_t) \right]= \mathbb{E}_{s \sim \tau_{new}} \left[ \sum\limits_{t=0}^{\infty} \gamma^t r(s_t) \right] + \mathbb{E}_{s_0 \sim \tau?} \left[ -V_{\pi_{old}}(s_0) \right]$

这里第二项中的预期,无论$\tau \mid \pi$还是$\tau \mid \pi_{old}$,起点$s_0$都是一样的,而第一项的reward需要从$\tau \mid pi$新策略的trajectory里累积。而第一项又正好符合定义$\mathbb{E}_{s_0}[V_{\pi}(s_0)]=J(\pi)$,所以

$J(\pi) – J(\pi_{old}) = \mathbb{E}_{\tau \mid \pi} \left[ \sum\limits_{t=0}^{\infty} \gamma^t A_{\pi_{old}}(s_t, a_t) \right]$

重写一下得到

$J(\pi) = J(\pi_{old}) + \mathbb{E}_{\tau \mid \pi} \left[ \sum\limits_{t=0}^{\infty} \gamma^t A_{\pi_{old}}(s_t, a_t) \right]$

要评估右边的预期值括号时,就出现了一个问题:1步td的$A_{\pi_{old}}(s, a)$使用到的状态转移来计算$V(s’)$和$V(s)$时,要依赖于旧策略$\pi_{old}$,然而在$\sum\limits_{t=0}^{\infty}$时,又需要t在新策略$\pi$生成的trajectory上推进。这让该项变得难以评估。

于是退一步,让sum的trajectory也使用$\pi_{old}$,其他地方维持不变,就定义了一个新的东西,作为$J(\pi)$的某种近似,叫做替代函数surrogate function:

$L_{\pi_{old}}(\pi) = J(\pi_{old})  + \mathbb{E}_{\tau \mid \pi_{old}} \left[ \sum\limits_{t=0}^{\infty} \gamma^t A_{\pi_{old}}(s_t, a_t) \right]$

这个函数我们可以方便地按照正常的方式采样计算,但是它和我们的最终目标$J(\pi)$仍然存在区别。这个区别要怎样评估呢?

首先,$A_{\pi_{old}}(s, a)$可正可负,当$a$因为新旧$\pi(a \mid s)$导致的差异而改变时,优势函数的最大误差不会超过
$$2 \mathop{\max}_{s, a} | A_{\pi_{old}}(s, a)|$$

假设$P(a_{\pi} \neq a_{\pi_{old}}) \le \alpha$,即两种策略得出的action不同的概率存在一个最大值$\alpha$

那么每一步的优势函数的误差,定义为左边,满足$$|\bar{A}(s)| \le 2 \alpha \mathop{\max}_{s, a} | A_{\pi_{old}}(s, a)| $$

有了一条路径上的单步误差上限,那么当计算$J(\pi) – L_{\pi_{old}}(\pi)$,涉及到两条路径时,

$\mathbb{E}_{\tau \mid \pi} […] – \mathbb{E}_{\tau \mid \pi_{old}} […]$两个预期括号内的单步总误差,最坏情况就是

$|\bar{A}(s \sim \tau_{\pi}) – \bar{A}(s \sim \tau_{\pi_{old}})| \le |\bar{A}(s \sim \tau_{\pi})| + | \bar{A}(s \sim \tau_{\pi_{old}})| =  4 \alpha \mathop{\max}_{s, a} | A_{\pi_{old}}(s, a)|$

为了方便,定义$\epsilon :=\mathop{\max}_{s, a} | A_{\pi_{old}}(s, a)| $

定义$n_t$是在某条路径上,截至到第t步之前的所有时间i,$a_i^{(\pi)} \neq a_i^{(\pi_{old})}$的次数,则有

$P(n_t=0) \ge (1-\alpha)^t$,以及$P(n_t > 0) \le 1 – (1-\alpha)^t$

对于任意一条路径,存在$\mathbb{E}_{\tau \mid pi?}[…] = P(n_t=0) \mathbb{E}_{\tau \mid pi?, n_t=0}[…] + P(n_t>0) \mathbb{E}_{\tau \mid pi?, n_t>0}[…]$

其中第一项$n_t=0$意味着新旧策略路径一样(截至到t时),括号内的优势函数也一样,所以可以完全抵消,只剩$n_t>0$的部分。

所以$|J(\pi) – L_{\pi_{old}}(\pi)| \le \sum\limits_{t}^{\infty} \gamma^t P(n_t>0)  4 \alpha \epsilon$

$ = 4 \epsilon \alpha \left( \sum\limits_{t=0}^{\infty}\gamma^t – \sum\limits_{t=0}^{\infty}(1-\alpha)^t \right)$

$ =\frac{4 \epsilon \alpha^2 \gamma } {(1-\gamma)[1 – \gamma(1-\alpha)]}$

$\le \frac{4 \epsilon \alpha^2 \gamma } {(1-\gamma)^2} $

注意到$\alpha$实际上就是将$\pi$与$\pi_{old}$在每一种操作上产生不一样的结果的概率加起来,得到单步错误的概率,再取所有步数内的最大错误概率,即$$ \alpha = \mathop{\max}_{s_t \sim \tau?} D_{TV}(\pi(\cdot \mid s_t) || \pi_{old}(\cdot \mid s_t))$$

这里对于tv距离的使用,因为对称性所以在哪种策略使用的路径上计算,都是可用的。

然后利用不等式$D_{TV}(p || q)^2 \le D_{KL}(p || q)$

可得 $J(\pi)  \ge L_{\pi_{old}}(\pi) – \frac{4 \epsilon \gamma}{(1-\gamma)^2} D_{KL}^{max_{s \sim \tau_{old}}}(\pi_{old} || \pi) := M(\pi)$

接下来试着把$\pi = \pi_{old}$的情况带入,分析

$L(\pi=\pi_{old})= J(\pi_{old}) + \sum\limits_{s \sim \tau \mid \pi_{old}} \rho(s) \sum\limits_{a} \pi(a \mid s, \theta=\theta_{old}) A_{\pi_{old}}(s, a))$

根据优势函数的定义,右边第二项里$\sum_a$的操作,会使得右侧的预期值为0,结果就是让$L(\pi(\theta_0))=J(\pi(\theta_0))$

此时KL距离也是0,导致$M(\pi=\pi_{old}) = J(\pi_{old})$

所以每次从$\pi_{old}$更新到$\pi$,只要优化$M$,就能确保$J$的增量不会少于M的增量,由此证明训练问题可以转化。


2. 模型主要思想

在考虑优化M的问题之前,可以把$M(\pi)$看作是优化$L_{\pi_{old}}(\pi)$的同时,使用惩罚项$$\mathop{\max}_{s_t \sim \tau_{old}} D_{KL}[\pi_{old}(\cdot \mid s_t) || \pi(\cdot \mid s_t)]$$以及惩罚系数$C=\frac{4 \epsilon \gamma}{(1-\gamma)^2} $

直接使用这个惩罚系数的话,因为它很大,会导致步长太小,收敛太慢,所以放宽惩罚,改用约束条件,即:

优化$L_{\pi_{old}}(\pi)$的同时,要求$D_{KL}^{max_{s \sim \tau_{old}}}(\pi_{old} || \pi) \le \delta$

$L_{\pi_{old}}(\pi) = J(\pi_{old}) + \sum\limits_{s \sim \tau \mid \pi_{old}} \rho(s) \sum\limits_{a} \pi(a \mid s, \theta) A_{\pi_{old}}(s, a))$

为了能让第二项在$\tau \mid \pi_{old}$内使用旧策略采样,与需要进行更新训练的策略(即与优势函数相乘的策略)允许存在一些差距(off policy),于是使用importance sample方法,L的第二项修改为:

$\mathbb{E}_{s, a \sim \tau \mid \pi_{old}} \left[ \frac{\pi(a \mid s, \theta)}{\pi(a \mid s,  \theta_{old}) } A_{\pi_{old}}(s, a)\right]$

求1阶导数(梯度)时,$J(\pi_{old})$作为初始$\theta_{old}$的起点,相当于已经detach()了,无论几阶导数都是0。

 

所以$g = \nabla_{\theta} L_{\pi_{old}}(\pi(\theta)) = \nabla_{\theta} \mathbb{E}_{s, a \sim \tau \mid \pi_{old}} \left[ \frac{\pi(a \mid s, \theta)}{\pi_{old}(a \mid s) } A_{\pi_{old}}(s, a)\right] $

$=\mathbb{E}_{s, a \sim \tau \mid \pi_{old}} \left[ \frac{ \nabla_{\theta} \pi(a \mid s, \theta)}{\pi_{old}(a \mid s) } A_{\pi_{old}}(s, a)\right]$

对$J(\pi)$求导:

$\nabla_{\theta} J(\pi(\theta)) = \nabla_{\theta} \sum\limits_{s \sim \tau \mid \pi} \rho_{\pi_{\theta}}(s) \sum\limits_a \pi_{\theta}(a \mid s) A_{\pi_{old}}(s, a)$

product rule的两项,对$\rho_{\pi_{\theta}}(s)$求导的一项是$ \nabla_{\theta}\rho_{\pi_{\theta}}(s) \sum\limits_{a} \pi_{\theta}(a \mid s) A_{\pi_{old}}(s, a)$,里面$\sum\limits_a$的部分会变成0。所以只剩下product rule的另一项:

$\nabla_{\theta} J(\pi(\theta))  = \sum\limits_{s \sim \tau \mid \pi} \rho_{\pi_{\theta}}(s) \sum\limits_a  \nabla_{\theta} \pi_{\theta}(a \mid s) A_{\pi_{old}}(s, a)$

当$\theta=\theta_{old}$时,

$\nabla_{\theta} J(\pi(\theta))  \mid _{\theta=\theta_{old}} = \sum\limits_{s \sim \tau \mid \pi} \rho_{\pi_{\theta}}(s) \sum\limits_a \pi_{\theta=\theta_{old}} \frac{\nabla_{\theta} \pi_{\theta}(a \mid s)} {\pi_{\theta=\theta_{old}}} A_{\pi_{old}}(s, a)$

$ = \mathbb{E}_{s, a \sim \tau \mid \pi_{old}} \left[ \frac{ \nabla_{\theta} \pi(a \mid s, \theta) \mid _{\theta=\theta_{old}}}{\pi_{old}(a \mid s) } A_{\pi_{old}}(s, a)\right]$

可以得到$\nabla_{\theta} J(\pi(\theta))  \mid _{\theta=\theta_{old}} = \nabla_{\theta} L_{\pi_{old}}(\pi(\theta))  \mid _{\theta=\theta_{old}}$

J与L在$\theta_{old}$处的梯度是一样的,又因为$M(\pi)=L_{\pi_{old}}(\pi)-CD_{KL}(\pi_{old} || \pi)$,那么在$\theta=\theta_{old}$处,$D_{KL}(\pi_{old} || \pi)$的一阶导是0,因为$D_{KL} \ge 0$,那么只要是连续的,在$\theta=\theta_{old}$时,$KL=0$且一阶导也必定是0。

另一方面对约束项求导时,max操作不方便处理,于是改成mean操作,约束项变成

$D = \frac{1}{T} \sum\limits_{t \sim \tau_{old}} D_{KL}\left[ \pi_{old}(\cdot \mid s_t) || \pi(\cdot \mid s_t) \right] \le \delta$

接下来只需要找到合适的lagrange multiplier $\beta$,优化$L – \beta D$就可以了。

根据trusted region边界的大小$\delta$,$\beta$会存在一个相应的最大值。

首先说一下如何在trusted region内进行优化:

图中的$\eta$实际上就是$J$;i代表旧参数,i+1代表新参数。从前面的推导中,可以得出这几个结论:

  1. $M(\pi_{old})=J(\pi_{old})$
  2. $M$是$J$的下界
  3. L与J在$\theta_{old}$处一阶导一样,且由于KL的一阶导是0,导致M与J的局部一阶导也一样。

更新$\theta$的方法是这样的:计算当前M曲线的最优解,之后将下一步的$\theta$放在这里,更新参数后计算新M以及它的最优解。重复以上步骤即可。每一次计算$\theta$的更新方向并更新之前,起点$\theta_{old}$会被设置为当前的$\theta$,即先让$\theta_{old}=\theta_{cur}$,之后再计算$g \mid _{\theta=\theta_{old}}$ 与 $\nabla_{\theta}^{2} D \mid _{\theta=\theta_{old}}$

将$L$与$D$分别展开到至少一项与$\theta$相关的非0项,即

$L(\theta) \approx L(\theta_{old}) + g^T(\theta-\theta_{old})$

$D_{KL}^{mean} [\pi_{old} || \pi] \approx \frac{1}{2} (\theta-\theta_{old})^T \boldsymbol{H} (\theta-\theta_{old})$

 

优化目标就变成了:

$$\theta_{new} = \mathop{\arg \max}_{\theta} g^T(\theta-\theta_{old}) $$

subject to $\frac{1}{2} (\theta-\theta_{old})^T \boldsymbol{H} (\theta-\theta_{old}) \le \delta$

加上lagrange multiplier转化成无bound优化

$Gain = g^T(\theta-\theta_{old}) – \lambda \left( \frac{1}{2} (\theta-\theta_{old})^T \boldsymbol{H} (\theta-\theta_{old}) – \delta \right) $

让$\beta \boldsymbol{s} = \theta – \theta_{old}$

$0=\nabla_{\boldsymbol{s}} G = \boldsymbol{g} – \lambda \boldsymbol{H} \boldsymbol{s}$

$\boldsymbol{s}$的方向是$\boldsymbol{s}^* = \boldsymbol{H}^{-1} \boldsymbol{g}$,且要让$\lambda = 1$时,$s^*Hs^*=\delta$,$s^*$是M的最优解,那么需要的比例$\beta$就是$\beta = \sqrt{ \frac {2\delta} {\boldsymbol{s}^T \boldsymbol{H} \boldsymbol{s}}}$

优化方式则是沿着这个方向进行line search,首先从$M$的理论最优解$s^* = \boldsymbol{H}^{-1} \boldsymbol{g}$开始,如果loss不出问题,kl距离在允许范围内,且代理损失$M(\pi_{new}) \ge M(\pi_{old})$,则这次更新是有效的。否则需要回滚到$\theta_{old}$,然后尝试缩减步长,使用更小的步长,例如$\theta_{new}^{\text{next_try}} = \theta_{old} + 0.5 * \beta \boldsymbol{s}^*$,再尝试更新。


3. 数学技巧

前面基本上完成了对TRPO的理论讲解,但实现TRPO时还需要解决两个问题。一是计算KL的二阶导Hessian的开销非常大,因为它不是$\frac {\partial scalar} {\partial vector}$,而是一个矩阵:

$H_{ij} = \frac {\partial} {\partial \theta_j} \frac {\partial D_{KL}} {\partial \theta_i} $

不过我们并不需要单独计算出$\boldsymbol{H}$本身,而是都会让它后面跟一个向量,例如$\boldsymbol{x}$,所以只需要计算$\boldsymbol{y} = \boldsymbol{H} \boldsymbol{x}$即可。至于求inverse的操作,$\boldsymbol{H}^{-1} \boldsymbol{g}$,可以交给conjugate gradient来解决,也不需要直接算出$\boldsymbol{H}^{-1}$。

这样$y_k = \sum\limits_{j} H_{kj} x_j = \sum\limits_j \frac {\partial} {\partial \theta_j} \frac {\partial D_{KL}} {\partial \theta_k} x_j $

$ = \frac {\partial} {\partial \theta_k} \sum\limits_j \frac {\partial D_{KL}} {\partial \theta_j} x_j$

这样只需要先求KL的一阶导,与后接的向量$\boldsymbol{x}$进行点乘,再对参数$\boldsymbol{\theta}$求一次导就可以了。这个步骤就是fisher vector product。

除了fisher vector product中的求导部分,TRPO并不需要使用反向传播,所以可以在numpy里计算更新。另外注意对梯度求导时,pytorch默认梯度张量是no_grad的,所以无法求二阶导。要解决这个问题,需要在用autograd求1阶导时设置retain_graph=True,create_graph=True

关于conjugate gradient方法,参见这里

发表评论

邮箱地址不会被公开。 必填项已用*标注