【强化学习】15 —— TRPO（Trust Region Policy Optimization）

前言

之前介绍的基于策略的方法包括策略梯度算法和 Actor-Critic 算法。这些方法虽然简单、直观，但在实际应用过程中会遇到训练不稳定的情况。回顾一下基于策略的方法：参数化智能体的策略，并设计衡量策略好坏的目标函数，通过梯度上升的方法来最大化这个目标函数，使得策略最优。这种算法有一个明显的缺点：当策略网络是深度模型时，沿着策略梯度更新参数，很有可能由于步长太长，策略突然显著变差，进而影响训练效果；同时由于采集到的数据的分布会随策略的更新而变化，步长太长，更新的策略可能变差，导致采集到的数据变差，进而陷入恶性循环之中。本文将要介绍的信任区域策略优化（trust region policy optimization，TRPO）算法，在理论上能够保证策略学习的性能单调性，并在实际应用中取得了比策略梯度算法更好的效果。

TRPO特点

• 是一种on-policy的算法
• 可以用于连续或离散的动作空间
• Open AI Spinning Up 中的TRPO可以达到并行运行的效果
• TRPO以一种on-policy的方式训练了一个随机策略。这意味着它通过根据其随机策略的最新版本进行动作采样来进行探索。动作选择中的随机量取决于初始条件和训练过程。在训练过程中，策略通常会逐渐变得不那么随机，因为更新规则鼓励它利用已经发现的奖励。这可能导致策略陷入局部最优。

策略梯度的优化目标

假设当前策略为${\pi }_{\theta }$，其参数为$\theta$。对于策略梯度的优化目标，我们可以将其写成累计的折扣回报$J(\theta )={\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}[{\sum }_t{\gamma }^{t}r(s_t,a_t)]$，基于轨迹$\tau$，我们可以采样到累计的折扣回报。同时，因为$V^{{\pi }_{\theta }}(s)$相应的期望可以写成累计折扣回报的形式,$V^{{\pi }_{\theta }}(s)={\mathbb{E}}_{a\sim {\pi }_{\theta }(s)}[Q^{{\pi }_{\theta }}(s,a)]={\mathbb{E}}_{a\sim {\pi }_{\theta }(s)}[{\mathbb{E}}_{\tau \sim p_{\theta }(\tau )}[{\sum }_{s_k=s,a_k=a}{\sum }_{t=k}^{\infty }{\gamma }^{t-k}r(s_t,a_t)]]$，所以我们可以得到第二种优化目标（基于起始点的形式）：$J(\theta )={\mathbb{E}}_{s_0\sim p_{\theta }(s_0)}[V^{{\pi }_{\theta }}(s_0)]$。两者是等效的。

在策略优化中，我们考虑如何借助当前的$\theta$找到一个更优的参数${\theta }^{\prime }$，使得$J({\theta }^{\prime })\ge J(\theta )$。由于初始状态$s_0$的分布和策略无关（与起始的环境相关），因此上述策略下的优化目标$J(\theta )$可以写成在新策略${\pi }_{{\theta }^{\prime }}$的期望形式：$$\begin{array}{rl}J(\theta )& ={\mathbb{E}}_{s_0}[V^{{\pi }_{\theta }}(s_0)]\\ & ={\mathbb{E}}_{{\pi }_{{\theta }^{\prime }}}\left[\sum _{t=0}^{\infty }{\gamma }^t{V}^{{\pi }_{\theta }}(s_t)-\sum _{t=1}^{\infty }{\gamma }^t{V}^{{\pi }_{\theta }}(s_t)\right]\\ & =-{\mathbb{E}}_{{\pi }_{{\theta }^{\prime }}}\left[\sum _{t=0}^{\infty }{\gamma }^t\left(\gamma V^{{\pi }_{\theta }}(s_{t+1})-V^{{\pi }_{\theta }}(s_t)\right)\right]\end{array}$$

基于以上等式，我们可以推导新旧策略的目标函数之间的差距：$$\begin{array}{rl}J({\theta }^{\prime })-J(\theta )& ={\mathbb{E}}_{s_0}\left[V^{{\pi }_{{\theta }^{\prime }}}(s_0)\right]-{\mathbb{E}}_{s_0}\left[V^{{\pi }_{\theta }}(s_0)\right]\\ & ={\mathbb{E}}_{{\pi }_{{\theta }^{\prime }}}\left[\sum _{t=0}^{\infty }{\gamma }^{t}r(s_t,a_t)\right]+{\mathbb{E}}_{{\pi }_{{\theta }^{\prime }}}\left[\sum _{t=0}^{\infty }{\gamma }^t\left(\gamma V^{{\pi }_{\theta }}(s_{t+1})-V^{{\pi }_{\theta }}(s_t)\right)\right]\\ & ={\mathbb{E}}_{{\pi }_{{\theta }^{\prime }}}\left[\sum _{t=0}^{\infty }{\gamma }^t\left[r(s_t,a_t)+\gamma V^{{\pi }_{\theta }}(s_{t+1})-V^{{\pi }_{\theta }}(s_t)\right]\right]\end{array}$$

很自然的，红色标注的部分和我们之前介绍过的优势函数形式一致，因此，等式可以变为：$$={\mathbb{E}}_{{\pi }_{{\theta }^{\prime }}}\left[\sum _{t=0}^{\infty }{\gamma }^t{A}^{{\pi }_{\theta }}(s_t,a_t)\right]$$

使用重要性采样

直接对上式进行采样是困难的，因为${\pi }_{{\theta }^{\prime }}$是我们需要求解的策略，但我们又要用它来收集样本。把所有可能的新策略都拿来收集数据，然后判断哪个策略满足上述条件的做法显然是不现实的。因此，我们需要用到重要性采样（Importance Sampling）。$$\begin{array}{rl}& J({\theta }^{\prime })-J(\theta )\\ & ={\mathbb{E}}_{\tau \sim p_{{\theta }^{\prime }}(\tau )}\left[\sum _{t=0}^{\infty }{\gamma }^t{A}^{{\pi }_{\theta }}(s_t,a_t)\right]\end{array}$$

${\mathbb{E}}_{\tau \sim p_{{\theta }^{\prime }}(\tau )}$可以拆分成先采样出状态$s_t$，再基于${\pi }_{{\theta }^{\prime }}$采样出动作$a_t$，于是有：

$$=\sum _t{\mathbb{E}}_{s_t\sim p_{{\theta }^{\prime }}(s_t)}[{\mathbb{E}}_{a_t\sim {\pi }_{{\theta }^{\prime }}(a_t|s_t)}[{\gamma }^t{A}^{{\pi }_{\theta }}(s_t,a_t)]]$$

再利用重要性采样的方法，可得

$$=\sum _t{\mathbb{E}}_{s_t\sim p_{{\theta }^{\prime }}(s_t)}[{\mathbb{E}}_{a_t\sim {\pi }_{\theta }(a_t|s_t)}[\frac{{\pi }_{{\theta }^{\prime }}(a_t|s_t)}{{\pi }_{\theta }(a_t|s_t)}{\gamma }^t{A}^{{\pi }_{\theta }}(s_t,a_t)]]$$

忽略状态分布的差异

不过，此时前面依然是${\mathbb{E}}_{s_t\sim p_{{\theta }^{\prime }}(s_t)}$，$s_t$需要基于${\pi }_{{\theta }^{\prime }}$进行采样。对此，TRPO忽略状态分布的差异.当策略更新前后的变化较小时，可以近似地令${\pi }_{\theta }(s_t)\approx {\pi }_{{\theta }^{\prime }}(s_t)$。现在，假设使用确定性策略，且当${\pi }_{{\theta }^{\prime }}(s_t)\ne {\pi }_{\theta }({\mathrm{s}}_{\mathrm{t}})$的概率小于时；或者假设使用随机性策略，且当$a^{\prime }\sim {\pi }_{{\theta }^{\prime }}(\cdot |s_t)\ne a\sim {\pi }_{\theta }(\cdot |s_t)$的概率小于时，则新策略可以由与旧策略相等的部分和不相等的部分之和组成：$$p_{{\theta }^{\prime }}(s_t)=(1-ϵ{)}^t{p}_{\theta }(s_t)+(1-(1-ϵ{)}^t)p_{mistake}(s_t)$$

移项可得，

$$\begin{array}{rl}|p_{{\theta }^{\prime }}(s_t)-p_{\theta }(s_t)|& =(1-(1-ϵ{)}^t)|p_{mistake}(s_t)-p_{\theta }(s_t)|\\ & \le 2(1-(1-ϵ{)}^t)\\ & \le 2ϵt\end{array}$$

⚠️You Should Know
1.$|p_{mistake}(s_t)-p_{\theta }(s_t)|\le 2$，概率上下差值最大为2
2.$(1-ϵ{)}^t\ge 1-ϵt\mathrm{for}ϵ\in [0,1]$泰勒展开

最后，我们可以让${\pi }_{\theta }(s_t)\approx {\pi }_{{\theta }^{\prime }}(s_t)$$$J({\theta }^{\prime })-J(\theta )\approx \sum _t{\mathbb{E}}_{s_t\sim p_{\theta }(s_t)}[{\mathbb{E}}_{a_t\sim {\pi }_{\theta }(a_t|s_t)}[\frac{{\pi }_{{\theta }^{\prime }}(a_t|s_t)}{{\pi }_{\theta }(a_t|s_t)}{\gamma }^t{A}^{{\pi }_{\theta }}(s_t,a_t)]]$$

约束策略的变化

为了保证新旧策略足够接近，TRPO 使用了库尔贝克-莱布勒（Kullback-Leibler，KL）散度来约束策略更新的幅度。$$\begin{array}{rl}& {\theta }^{\prime }\leftarrow \arg \underset{{\theta }^{\prime }}{\max }\sum _t{\mathbb{E}}_{s_t\sim p_{\theta }(s_t)}[{\mathbb{E}}_{a_t\sim {\pi }_{\theta }(a_t|s_t)}[\frac{{\pi }_{{\theta }^{\prime }}(a_t|s_t)}{{\pi }_{\theta }(a_t|s_t)}{\gamma }^t{A}^{{\pi }_{\theta }}(s_t,a_t)]]\\ & \mathrm{such}\mathrm{that}{\mathbb{E}}_{s_t\sim p(s_t)}[D_{KL}({\pi }_{{\theta }^{\prime }}(a_t|s_t)\parallel {\pi }_{\theta }(a_t|s_t))]\le ϵ\end{array}$$

⚠️You Should Know
1.当$\theta ={\theta }^{\prime }$时，上述目标函数以及其约束都为0；此外，当$\theta ={\theta }^{\prime }$时，约束对$\theta$的梯度也为0.
2.实际多使用constraint violate as penalty进行求解$$\begin{array}{r}{\theta }^{\prime }\leftarrow \arg \underset{{\theta }^{\prime }}{\max }\sum _t{\mathbb{E}}_{s_t\sim p_{\theta }(s_t)}[{\mathbb{E}}_{a_t\sim {\pi }_{\theta }(a_t|s_t)}[\frac{{\pi }_{{\theta }^{\prime }}(a_t|s_t)}{{\pi }_{\theta }(a_t|s_t)}{\gamma }^t{A}^{{\pi }_{\theta }}(s_t,a_t)]]\\ -\lambda (D_{KL}({\pi }_{{\theta }^{\prime }}(a_t|s_t)\parallel {\pi }_{\theta }(a_t|s_t))-ϵ)\end{array}$$
2.1优化上式，更新${\theta }^{\prime }$
2.2更新惩罚项$\lambda \leftarrow \lambda +\alpha (D_{KL}({\pi }_{{\theta }^{\prime }}(a_t|s_t)\parallel {\pi }_{\theta }(a_t|s_t))-ϵ)$（若KL散度超出$ϵ$，则进行惩罚，超出越多，惩罚越多）

近似求解

上述式子的求解还是比较困难的，因此，TRPO进行了近似求解。为了方便描述，我们将目标函数和约束转化为以下形式：$$\begin{array}{r}{\theta }_{k+1}=\arg \underset{\theta }{\max }\mathcal{L}({\theta }_k,\theta )\\ s.t.{\bar{D}}_{KL}(\theta ||{\theta }_k)\le \delta \end{array}$$

其中，$\mathcal{L}({\theta }_k,\theta )=\underset{s,a\sim {\pi }_{{\theta }_k}}{E}\left[\frac{{\pi }_{\theta }(a|s)}{{\pi }_{{\theta }_k}(a|s)}{A}^{{\pi }_{{\theta }_k}}(s,a)\right],$${\bar{D}}_{KL}(\theta ||{\theta }_k)={\mathrm{E}}_{s\sim {\pi }_{{\theta }_k}}[D_{KL}\left({\pi }_{\theta }(\cdot |s)||{\pi }_{{\theta }_k}(\cdot |s))\right].$

对其基于${\theta }_k$进行泰勒展开：$$\begin{array}{rl}\mathcal{L}({\theta }_k,\theta )& \approx g^T(\theta -{\theta }_k)\\ {\bar{D}}_{KL}(\theta ||{\theta }_k)& \approx \frac{1}{2}(\theta -{\theta }_k{)}^{T}H(\theta -{\theta }_k)\end{array}$$

其中$g$表示目标函数的梯度${\nabla }_{\theta }J({\pi }_{\theta })$，$H$表示策略之间平均 KL 距离的黑塞矩阵（Hessian matrix）。因此，我们可以将问题转化为一个近似的优化问题：$$\begin{array}{r}{\theta }_{k+1}=\arg \underset{\theta }{\max }{g}^T(\theta -{\theta }_k)\\ \text{s.t. }\frac{1}{2}(\theta -{\theta }_k{)}^{T}H(\theta -{\theta }_k)\le \delta .\end{array}$$

此时，我们可以用卡罗需-库恩-塔克（Karush-Kuhn-Tucker，KKT）条件直接导出上述问题的解：
$${\theta }_{k+1}={\theta }_k+\sqrt{\frac{2\delta }{g^T{H}^{-1}g}}{H}^{-1}g$$

⚠️You Should Know
1.进一步的求解需要用到共轭梯度法（conjugate gradient method）或者拉格朗日对偶法（Lagrangian duality）。共轭梯度法（TRPO所用）的具体求解参考《动手学强化学习》第11.4节；拉格朗日对偶法（Lagrangian duality）则可参考Boyd的凸优化 —— Convex Optimization by Boyd and Vandenberghe, especially chapters 2 through 5.

2.一般来说，用神经网络表示的策略函数的参数数量都是成千上万的，计算和存储黑塞矩阵$H$的逆矩阵会耗费大量的内存资源和时间。TRPO 通过共轭梯度法（conjugate gradient method）回避了这个问题，它的核心思想是直接计算$Hx$向量,$(x=H^{-1}g)$ ，$x$即参数更新方向，而非存储黑塞矩阵$H$。公式描述如下：$$Hx={\nabla }_{\theta }\left({\left({\nabla }_{\theta }{\bar{D}}_{KL}(\theta ||{\theta }_k)\right)}^{T}x\right),$$

线性搜索

如果到此为止，计算出最终的结果，那么现在的算法和Natural Policy Gradient差不多。由于 TRPO 算法用到了泰勒展开的 1 阶和 2 阶近似，这并非精准求解，最终的策略结果并非能得到提升，并且未必能满足 KL 散度限制。因此，TRPO 在每次迭代的最后进行一次线性搜索（Line Search），以确保找到满足条件。具体来说，就是找到一个最小的非负整数$i$，使得按照$${\theta }_{k+1}={\theta }_k+{\alpha }^i\sqrt{\frac{2\delta }{x^{T}Hx}}x$$

其中$\alpha \in (0,1)$是一个决定线性搜索长度的超参数。求出的${\theta }_{k+1}$依然满足最初的 KL 散度限制，并且确实能够提升目标函数$\mathcal{L}({\theta }_k,\theta )$。

算法伪代码

广义优势估计

现在，我们还缺少一种合适的方法用来估计优势函数$A$。目前比较常用的一种方法为广义优势估计（Generalized Advantage Estimation，GAE），接下来我们简单介绍一下 GAE 的做法。首先，用${\delta }_t=r_t+\gamma V(s_{t+1})-V(s_t)$表示时序差分误差，其中$V$是一个已经学习的状态价值函数。于是，根据多步时序差分的思想，有：$$\begin{array}{rlrlrlrl}& A_t^{(1)}& & ={\delta }_t& & & & =-V(s_t)+r_t+\gamma V(s_{t+1})\\ & A_t^{(2)}& & ={\delta }_t+\gamma {\delta }_{t+1}& & & & =-V(s_t)+r_t+\gamma r_{t+1}+{\gamma }^{2}V(s_{t+2})\\ & A_t^{(3)}& & ={\delta }_t+\gamma {\delta }_{t+1}+{\gamma }^2{\delta }_{t+2}& & & & =-V(s_t)+r_t+\gamma r_{t+1}+{\gamma }^2{r}_{t+2}+{\gamma }^{3}V(s_{t+3})\\ & & ⋮& & & & ⋮& \\ & A_t^{(k)}& & =\sum _{l=0}^{k-1}{\gamma }^l{\delta }_{t+l}& & & & =-V(s_t)+r_t+\gamma r_{t+1}+\dots +{\gamma }^{k-1}{r}_{t+k-1}+{\gamma }^{k}V(s_{t+k})\end{array}$$
然后，GAE 将这些不同步数的优势估计进行指数加权平均：$$\begin{array}{rl}{A}_t^{GAE}& =(1-\lambda )(A_t^{(1)}+\lambda A_t^{(2)}+{\lambda }^2{A}_t^{(3)}+\cdots )\\ & =(1-\lambda )({\delta }_t+\lambda ({\delta }_t+\gamma {\delta }_{t+1})+{\lambda }^2({\delta }_t+\gamma {\delta }_{t+1}+{\gamma }^2{\delta }_{t+2})+\cdots )\\ & =(1-\lambda )(\delta (1+\lambda +{\lambda }^2+\cdots )+\gamma {\delta }_{t+1}(\lambda +{\lambda }^2+{\lambda }^3+\cdots )+\gamma \\ & =(1-\lambda )\left({\delta }_t\frac{1}{1-\lambda }+\gamma {\delta }_{t+1}\frac{\lambda }{1-\lambda }+{\gamma }^2{\delta }_{t+2}\frac{{\lambda }^2}{1-\lambda }+\cdots \right)\\ & =\sum _{l=0}^{\infty }(\gamma \lambda {)}^l{\delta }_{t+l}\end{array}$$

其中，$\lambda \in [0,1]$是在 GAE 中额外引入的一个超参数。当时$\lambda =0$，$A_t^{GAE}={\delta }_t=r_t+\gamma V(s_{t+1})-V(s_t)$，也即是仅仅只看一步差分得到的优势；当时$\lambda =1$，$A_t^{GAE}={\sum }_{l=0}^{\infty }{\gamma }^l{\delta }_{t+l}={\sum }_{l=0}^{\infty }{\gamma }^l{r}_{t+l}-V(s_t).$，则是看每一步差分得到优势的完全平均值。

由于广义优势估计(GAE)的实现有很多变化，下面是一种可能的实现方法：

def calculate_gae(rewards, values, gamma=0.99, lambda_=0.95):
    """
    计算广义优势估计(GAE)的值

    :param rewards: 一组连续的状态动作奖励(reward)，list
    :param values:  一组连续的状态的价值(value)， list
    :param gamma: 衰减系数，float
    :param lambda_: GAE超参数，float
    :return: 估计的优势值(advantages)，list
    """
    # 计算 TD 计算误差 delta
    deltas = [r + gamma * v_next - v for r, v, v_next in zip(rewards, values, values[1:] + [0])]
    
    # 计算 GAE
    advantages = []
    advantage = 0
    for delta in reversed(deltas):
        advantage = gamma * lambda_ * advantage + delta
        advantages.append(advantage)
    advantages.reverse()

    return advantages

这个函数的输入是同步的状态动作奖励和状态价值列表，以及GAE的超参数。它通过计算TD误差(delta)和广义优势估计(GAE)，输出估计的优势值(advantages)。

在训练过程中，我们会调用这个函数来计算advantages。例子伪代码如下：

for _ in range(num_iterations):
    #执行策略，获取state,action,reward
    states, actions, rewards = run_policy(env, model, num_steps)
    #计算所有状态的值
    values = calculate_values(env, model, states)
    #计算所有状态的advantages
    advantages = calculate_gae(rewards, values)
    #更新策略
    update_policy(model, states, actions, advantages)

这些伪代码只是一个粗略的概念，实际实现中需要很多细节。

代码实践

离散动作空间

class PolicyNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return F.softmax(self.fc2(x), dim=1)

# 输入是某个状态，输出则是状态的价值。
class ValueNet(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim):
        super(ValueNet, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc2 = torch.nn.Linear(hidden_dim, 1)

    def forward(self, x):
        x = F.relu(self.fc1(x))
        return self.fc2(x)

class TRPO:
    def __init__(self, state_dim, hidden_dim, action_dim, lambda_, gamma, alpha,
                 kl_constraint, critic_lr, device, numOfEpisodes, env):
        self.actor = PolicyNet(state_dim, hidden_dim, action_dim).to(device)
        self.critic = ValueNet(state_dim, hidden_dim).to(device)
        # 策略网络参数不需要优化器更新
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=critic_lr)
        self.gamma = gamma
        # GAE参数
        self.lambda_ = lambda_
        # KL距离最大限制
        self.kl_constraint = kl_constraint
        # 线性搜索参数
        self.alpha = alpha
        self.device = device
        self.env = env
        self.numOfEpisodes = numOfEpisodes

    def take_action(self, state):
        state = torch.tensor(np.array([state]), dtype=torch.float).to(self.device)
        probs = self.actor(state)
        action_dist = torch.distributions.Categorical(probs)
        return action_dist.sample().item()

    def cal_advantage(self, gamma, lambda_, td_delta):
        td_delta = td_delta.detach().numpy()
        advantages = []
        advantage = 0.0
        for delta in reversed(td_delta):
            advantage = gamma * lambda_ * advantage + delta
            advantages.append(advantage)
        advantages.reverse()
        return torch.FloatTensor(np.array(advantages))

    # compute surrogate object function
    def compute_surrogate_obj(self, states, actions, advantage, old_log_probs, actor):
        log_probs = torch.log(actor(states).gather(1, actions))
        ratio = torch.exp(log_probs - old_log_probs)
        return torch.mean(ratio * advantage)

    def hessian_matrix_vector_product(self, states, old_action_dists, vector):
        # 计算黑塞矩阵和一个向量的乘积
        new_action_dists = torch.distributions.Categorical(self.actor(states))
        kl = torch.mean(
            torch.distributions.kl.kl_divergence(old_action_dists, new_action_dists))  # 计算平均KL距离
        kl_grad = torch.autograd.grad(kl, self.actor.parameters(), create_graph=True)
        kl_grad_vector = torch.cat([grad.view(-1) for grad in kl_grad])
        # KL距离的梯度先和向量进行点积运算
        kl_grad_vector_product = torch.dot(kl_grad_vector, vector)
        grad2 = torch.autograd.grad(kl_grad_vector_product, self.actor.parameters())
        grad2_vector = torch.cat([grad.view(-1) for grad in grad2])
        return grad2_vector

    # 共轭梯度法求解方程
    def conjugate_gradient(self, grad, states, old_action_dists):
        x = torch.zeros_like(grad)
        r = grad.clone()
        p = grad.clone()
        rdotr = torch.dot(r, r)
        for i in range(10):  # 共轭梯度主循环
            Hp = self.hessian_matrix_vector_product(states, old_action_dists,
                                                    p)
            alpha = rdotr / torch.dot(p, Hp)
            x += alpha * p
            r -= alpha * Hp
            new_rdotr = torch.dot(r, r)
            if new_rdotr < 1e-10:
                break
            beta = new_rdotr / rdotr
            p = r + beta * p
            rdotr = new_rdotr
        return x

    def line_search(self, states, actions, advantage, old_log_probs,
                    old_action_dists, max_vec):  # 线性搜索
        old_para = torch.nn.utils.convert_parameters.parameters_to_vector(
            self.actor.parameters())
        old_obj = self.compute_surrogate_obj(states, actions, advantage,
                                             old_log_probs, self.actor)
        for i in range(15):  # 线性搜索主循环
            coef = self.alpha**i
            new_para = old_para + coef * max_vec
            new_actor = copy.deepcopy(self.actor)
            torch.nn.utils.convert_parameters.vector_to_parameters(
                new_para, new_actor.parameters())
            new_action_dists = torch.distributions.Categorical(
                new_actor(states))
            kl_div = torch.mean(
                torch.distributions.kl.kl_divergence(old_action_dists,
                                                     new_action_dists))
            new_obj = self.compute_surrogate_obj(states, actions, advantage,
                                                 old_log_probs, new_actor)
            if new_obj > old_obj and kl_div < self.kl_constraint:
                return new_para
        return old_para

    def update(self, transition_dict):
        states = torch.tensor(np.array(transition_dict['states']), dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(np.array(transition_dict['next_states']), dtype=torch.float).to(self.device)
        terminateds = torch.tensor(transition_dict['terminateds'], dtype=torch.float).view(-1, 1).to(self.device)
        truncateds = torch.tensor(transition_dict['truncateds'], dtype=torch.float).view(-1, 1).to(self.device)
        td_target = rewards + self.gamma * self.critic(next_states) * (1 - terminateds + truncateds)
        td_delta = td_target - self.critic(states)
        # estimate advantage using GAE
        advantage = self.cal_advantage(self.gamma, self.lambda_, td_delta.cpu()).to(self.device)

        old_log_probs = torch.log(self.actor(states).gather(1, actions)).detach()
        old_action_dists = torch.distributions.Categorical(self.actor(states).detach())

        # update Value function
        critic_loss = torch.mean(F.mse_loss(self.critic(states), td_target.detach()))
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()
        # estimate policy gradient
        # compute surrogate object function
        surrogate_obj = self.compute_surrogate_obj(states, actions, advantage, old_log_probs, self.actor)
        grads = torch.autograd.grad(surrogate_obj, self.actor.parameters())
        obj_grad = torch.cat([grad.view(-1) for grad in grads]).detach()
        # use the conjugate gradient algorithm to compute x = H^(-1)g
        descent_direction = self.conjugate_gradient(obj_grad, states, old_action_dists)
        Hd = self.hessian_matrix_vector_product(states, old_action_dists,
                                                descent_direction)
        max_coef = torch.sqrt(2 * self.kl_constraint / (torch.dot(descent_direction, Hd) + 1e-8))
        # update the policy by backtracking line search
        new_para = self.line_search(states, actions, advantage, old_log_probs,
                                    old_action_dists,
                                    descent_direction * max_coef)
        # 用线性搜索后的参数更新策略
        torch.nn.utils.convert_parameters.vector_to_parameters(new_para, self.actor.parameters())

    def TRPOrun(self):
        returnList = []
        for i in range(10):
            with tqdm(total=int(self.numOfEpisodes / 10), desc='Iteration %d' % i) as pbar:
                for episode in range(int(self.numOfEpisodes / 10)):
                    # initialize state
                    state, info = self.env.reset()
                    terminated = False
                    truncated = False
                    episodeReward = 0
                    transition_dict = {
                        'states': [], 'actions': [], 'next_states': [], 'rewards': [], 'terminateds': [], 'truncateds': []}
                    # Loop for each step of episode:
                    while 1:
                        action = self.take_action(state)
                        next_state, reward, terminated, truncated, info = self.env.step(action)
                        transition_dict['states'].append(state)
                        transition_dict['actions'].append(action)
                        transition_dict['next_states'].append(next_state)
                        transition_dict['rewards'].append(reward)
                        transition_dict['terminateds'].append(terminated)
                        transition_dict['truncateds'].append(truncated)
                        state = next_state
                        episodeReward += reward
                        if terminated or truncated:
                            break
                    self.update(transition_dict)
                    returnList.append(episodeReward)
                    if (episode + 1) % 10 == 0:  # 每10条序列打印一下这10条序列的平均回报
                        pbar.set_postfix({
                            'episode':
                                '%d' % (self.numOfEpisodes / 10 * i + episode + 1),
                            'return':
                                '%.3f' % np.mean(returnList[-10:])
                        })
                    pbar.update(1)
        return returnList

从结果中可以看到，TRPO收敛速度相对较快,性能优异。

超参数参考：

    agent = TRPO(state_dim=env.observation_space.shape[0],
                 hidden_dim=256,
                 action_dim=2,
                 lambda_=0.95,
                 gamma=0.99,
                 alpha=0.5,
                 kl_constraint=0.0005,
                 critic_lr=1e-2,
                 device=device,
                 numOfEpisodes=1000,
                 env=env)

连续动作空间

下面代码仅列出了需要修改部分的代码。

class PolicyNetContinuous(torch.nn.Module):
    def __init__(self, state_dim, hidden_dim, action_dim):
        super(PolicyNetContinuous, self).__init__()
        self.fc1 = torch.nn.Linear(state_dim, hidden_dim)
        self.fc_mu = torch.nn.Linear(hidden_dim, action_dim)
        self.fc_std = torch.nn.Linear(hidden_dim, action_dim)

    # 高斯分布的均值和标准差
    def forward(self, x):
        x = F.relu(self.fc1(x))
        mu = 2.0 * torch.tanh(self.fc_mu(x))
        std = F.softplus(self.fc_std(x))
        return mu, std

class TRPO:
    def __init__(self, state_dim, hidden_dim, action_dim, lambda_, gamma, alpha,
                 kl_constraint, critic_lr, device, numOfEpisodes, env):
        self.actor = PolicyNetContinuous(state_dim, hidden_dim, action_dim).to(device)
        self.critic = ValueNet(state_dim, hidden_dim).to(device)
        # 策略网络参数不需要优化器更新
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=critic_lr)
        self.gamma = gamma
        # GAE参数
        self.lambda_ = lambda_
        # KL距离最大限制
        self.kl_constraint = kl_constraint
        # 线性搜索参数
        self.alpha = alpha
        self.device = device
        self.env = env
        self.numOfEpisodes = numOfEpisodes

    def take_action(self, state):
        state = torch.tensor(np.array([state]), dtype=torch.float).to(self.device)
        mu, std = self.actor(state)
        action_dist = torch.distributions.Normal(mu, std)
        action = action_dist.sample()
        return [action.item()]
        
    # compute surrogate object function
    def compute_surrogate_obj(self, states, actions, advantage, old_log_probs, actor):
        mu, std = actor(states)
        action_dists = torch.distributions.Normal(mu, std)
        log_probs = action_dists.log_prob(actions)
        ratio = torch.exp(log_probs - old_log_probs)
        return torch.mean(ratio * advantage)
        
    def hessian_matrix_vector_product(self, states, old_action_dists, vector, damping=0.1):
        # 计算黑塞矩阵和一个向量的乘积
        mu, std = self.actor(states)
        new_action_dists = torch.distributions.Normal(mu, std)
        kl = torch.mean(
            torch.distributions.kl.kl_divergence(old_action_dists, new_action_dists))  # 计算平均KL距离
        kl_grad = torch.autograd.grad(kl, self.actor.parameters(), create_graph=True)
        kl_grad_vector = torch.cat([grad.view(-1) for grad in kl_grad])
        # KL距离的梯度先和向量进行点积运算
        kl_grad_vector_product = torch.dot(kl_grad_vector, vector)
        grad2 = torch.autograd.grad(kl_grad_vector_product, self.actor.parameters())
        grad2_vector = torch.cat([grad.view(-1) for grad in grad2])
        return grad2_vector + damping * vector
        
    def line_search(self, states, actions, advantage, old_log_probs,
                    old_action_dists, max_vec):
        old_para = torch.nn.utils.convert_parameters.parameters_to_vector(
            self.actor.parameters())
        old_obj = self.compute_surrogate_obj(states, actions, advantage,
                                             old_log_probs, self.actor)
        for i in range(15):
            coef = self.alpha**i
            new_para = old_para + coef * max_vec
            new_actor = copy.deepcopy(self.actor)
            torch.nn.utils.convert_parameters.vector_to_parameters(
                new_para, new_actor.parameters())
            mu, std = new_actor(states)
            new_action_dists = torch.distributions.Normal(mu, std)
            kl_div = torch.mean(
                torch.distributions.kl.kl_divergence(old_action_dists,
                                                     new_action_dists))
            new_obj = self.compute_surrogate_obj(states, actions, advantage,
                                                 old_log_probs, new_actor)
            if new_obj > old_obj and kl_div < self.kl_constraint:
                return new_para
        return old_para

    def update(self, transition_dict):
        states = torch.tensor(np.array(transition_dict['states']), dtype=torch.float).to(self.device)
        actions = torch.tensor(transition_dict['actions']).view(-1, 1).to(self.device)
        rewards = torch.tensor(transition_dict['rewards'], dtype=torch.float).view(-1, 1).to(self.device)
        next_states = torch.tensor(np.array(transition_dict['next_states']), dtype=torch.float).to(self.device)
        terminateds = torch.tensor(transition_dict['terminateds'], dtype=torch.float).view(-1, 1).to(self.device)
        truncateds = torch.tensor(transition_dict['truncateds'], dtype=torch.float).view(-1, 1).to(self.device)
        rewards = (rewards + 8.0) / 8.0  # 对奖励进行修改,方便训练

        td_target = rewards + self.gamma * self.critic(next_states) * (1 - terminateds + truncateds)
        td_delta = td_target - self.critic(states)
        # estimate advantage using GAE
        advantage = self.cal_advantage(self.gamma, self.lambda_, td_delta.cpu()).to(self.device)

        mu, std = self.actor(states)
        old_action_dists = torch.distributions.Normal(mu.detach(), std.detach())
        old_log_probs = old_action_dists.log_prob(actions)

        # update Value function
        critic_loss = torch.mean(F.mse_loss(self.critic(states), td_target.detach()))
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()
        # estimate policy gradient
        # compute surrogate object function
        surrogate_obj = self.compute_surrogate_obj(states, actions, advantage, old_log_probs, self.actor)
        grads = torch.autograd.grad(surrogate_obj, self.actor.parameters())
        obj_grad = torch.cat([grad.view(-1) for grad in grads]).detach()
        # use the conjugate gradient algorithm to compute x = H^(-1)g
        descent_direction = self.conjugate_gradient(obj_grad, states, old_action_dists)
        Hd = self.hessian_matrix_vector_product(states, old_action_dists,
                                                descent_direction)
        max_coef = torch.sqrt(2 * self.kl_constraint / (torch.dot(descent_direction, Hd) + 1e-8))
        # update the policy by backtracking line search
        new_para = self.line_search(states, actions, advantage, old_log_probs,
                                    old_action_dists,
                                    descent_direction * max_coef)
        # 用线性搜索后的参数更新策略
        torch.nn.utils.convert_parameters.vector_to_parameters(new_para, self.actor.parameters())

用 TRPO 在与连续动作交互的倒立摆环境中能够取得非常不错的效果，这说明 TRPO 中的信任区域优化方法在离散和连续动作空间都能有效工作。

超参数参考：

    agent = TRPO(state_dim=env.observation_space.shape[0],
                 hidden_dim=256,
                 action_dim=env.action_space.shape[0],
                 lambda_=0.85,
                 gamma=0.90,
                 alpha=0.5,
                 kl_constraint=0.0005,
                 critic_lr=1e-2,
                 device=device,
                 numOfEpisodes=2000,
                 env=env)

参考

[1] 伯禹AI
[2] https://www.davidsilver.uk/teaching/
[3] 动手学强化学习
[4] Reinforcement Learning
[5] SCHULMAN J, LEVINE S, ABBEEL P, et al. Trust region policy optimization [C]// International conference on machine learning, PMLR, 2015:1889-1897.
[6] Schulman J. Optimizing expectations: From deep reinforcement learning to stochastic computation graphs[D]. UC Berkeley, 2016.
[7] https://spinningup.openai.com/en/latest/algorithms/trpo.html