深度强化学习笔记之PPO实现细节（2）

本文主要参考于Coding PPO from Scratch with PyTorch系列，但本文并不会像该系列一样手把手讲解全部的实现细节，只是记录一下自己在实现过程中遇到的一些问题和思考。

下图是采用Clipped Surrogate Objective的PPO伪代码，本文的代码实现主要根据它来实现。

1.构建目标函数

PPO算法的实现重点，就是为了得到上图中的两个目标函数。也就是说，我们只要可以构造出式(1)和式(2)作为损失函数，基本就实现了PPO。
θ k + 1 = arg ⁡ max ⁡ θ 1 ∣ D k ∣ T ∑ τ ∈ D k ∑ t = 0 T min ⁡ ( π θ ( a t ∣ s t ) π θ k ( a t ∣ s t ) A π θ k ( s t , a t ) ,    g ( ϵ , A π θ k ( s t , a t ) ) ) (1) \\theta_{k+1}=\\arg \\max_{\\theta} \\frac{1}{|D_k|T} \\sum_{\\tau \\in D_k}\\sum^T_{t=0}\\min \\bigg( \\frac{\\pi_{\\theta}(a_t|s_t)}{\\pi_{\\theta_k}(a_t|s_t)}A^{\\pi_{\\theta_k}}(s_t,a_t), \\; g(\\epsilon, A^{\\pi_{\\theta_k}}(s_t,a_t)) \\bigg) \\tag1 θk+1​=argθmax​∣Dk​∣T1​τ∈Dk​∑​t=0∑T​min(πθk​​(at​∣st​)πθ​(at​∣st​)​Aπθk​​(st​,at​),g(ϵ,Aπθk​​(st​,at​)))(1)

ϕ k + 1 = arg ⁡ min ⁡ ϕ 1 ∣ D k ∣ T ∑ τ ∈ D k ∑ t = 0 T ( V ϕ ( s t ) − R ^ t ) 2 (2) \\phi_{k+1}=\\arg \\min_{\\phi} \\frac{1}{|D_k|T} \\sum_{\\tau \\in D_k}\\sum^T_{t=0} \\Big( V_{\\phi}(s_t)-\\hat{R}_t \\Big)^2 \\tag2 ϕk+1​=argϕmin​∣Dk​∣T1​τ∈Dk​∑​t=0∑T​(Vϕ​(st​)−R^t​)2(2)

在上式中，k是最外一层的主循环，可以理解为训练一轮PPO的次数，也就是PPO的更新次数； D k D_k Dk​表示第k轮得到的轨迹集合，也就是这一轮PPO的训练资料。

PPO算法会写成\"actor-critic\"的形式，在具体实现时，会采用两个神经网络分别作为actor和critic的参数形式。在上式中， θ \\theta θ表示actor的参数， ϕ \\phi ϕ表示critic的参数。在这里， π θ ( a t ∣ s t ) \\pi_{\\theta}(a_t|s_t) πθ​(at​∣st​)和 V ϕ ( s t ) V_{\\phi}(s_t) Vϕ​(st​)就是两个网络模型，输入为 s t s_t st​。

# actor-critic model
actor_model = FeedForwardNN(obs_dim, action_dim)  # At
critic_model = FeedForwardNN(obs_dim, 1)  # V(St)

若要构建式(1)和(2)，则需要知道以下5个函数值：

• π θ ( a t ∣ s t ) \\pi_{\\theta}(a_t|s_t) πθ​(at​∣st​)， π θ k ( a t ∣ s t ) \\pi_{\\theta_k}(a_t|s_t) πθk​​(at​∣st​)， V ϕ ( s t ) V_{\\phi}(s_t) Vϕ​(st​)， A π θ k ( s t , a t ) A^{\\pi_{\\theta_k}}(s_t,a_t) Aπθk​​(st​,at​)， R ^ t \\hat{R}_t R^t​

1.1.Actor

actor本质上就是policy，遵循policy gradient的准则进行策略更新。

如果action是连续的，那么actor会输出一个action的均值（维度与action相同），以此构建一个多维正态分布；如果action是离散的，那么actor则会输出每一个action的概率，以此构建一个Categorical分布。critic则只输出一个标量（表示输入observation的价值）。

from gym.spaces import Box, Discrete
from torch.distributions import MultivariateNormal, Categorical

## action连续: Box(low=-1, high=1, shape=(2, ))
act_mean = actor_model(obs)
# 创建一个多维正态分布
dist = MultivariateNormal(loc=act_mean, covariance_matrix=cov_mat)
# 从该正态分布中采样得到一个action
action = dist.sample()
# 计算这个action的log概率
log_prob = dist.log_prob(action)  # log(pi)

## action离散: Discrete(4)
act_probs = actor_model(obs)
dist = Categorical(act_probs)
# 从Categorical中采样得到一个action
action = dist.sample()
# 计算这个action的log概率
log_prob = dist.log_prob(action)  # log(pi)

以上程序描述了根据actor的输出计算action的概率（ π θ ( a t ∣ s t ) \\pi_{\\theta}(a_t|s_t) πθ​(at​∣st​)）的常规方法。可以发现，它并不是直接把actor的输出作为动作概率值的，而是以此构建了一个概率分布，然后从这个分布中采样并计算相应的log概率值。这里有两个问题需要注意：

• 为什么需要构建概率分布：为了平衡探索和利用的矛盾，作用与 ϵ − g r e e d y \\epsilon-greedy ϵ−greedy方法相同；

• 为什么输出log概率值：log函数可以把乘除法变成加减法，计算更加简单。
  log ⁡ π θ ( a t ∣ s t ) π θ k ( a t ∣ s t ) = log ⁡ π θ ( a t ∣ s t ) − log ⁡ π θ k ( a t ∣ s t ) \\log{\\frac{\\pi_{\\theta}(a_t|s_t)}{\\pi_{\\theta_k}(a_t|s_t)}}= \\log{\\pi_{\\theta}(a_t|s_t)}-\\log{\\pi_{\\theta_k}(a_t|s_t)} logπθk​​(at​∣st​)πθ​(at​∣st​)​=logπθ​(at​∣st​)−logπθk​​(at​∣st​)

1.2.Critic

V ϕ ( s t ) V_{\\phi}(s_t) Vϕ​(st​)的计算方式比较简单，就是critic的输出，表示状态 s t s_t st​的价值，不需要其它处理。

V = critic_model(batch_obs).squeeze()

1.3.Advantage function

Sometimes in RL, we don’t need to describe how good an action is in an absolute sense, but only how much better it is than others on average. That is to say, we want to know the relative advantage of that action. We make this concept precise with the advantage function.

在policy gradient中，优势函数 A π θ ( s t , a t ) A^{\\pi_{\\theta}}(s_t,a_t) Aπθ​(st​,at​)十分重要。它可以描述基于一个策略 π θ \\pi_{\\theta} πθ​，在状态 s t s_t st​采取一个特定动作 a t a_t at​的相对价值。数学形式可以表示为
A π θ ( s t , a t ) = Q π θ ( s t , a t ) − V ( s t ) A^{\\pi_{\\theta}}(s_t,a_t)=Q^{\\pi_{\\theta}}(s_t,a_t)-V(s_t) Aπθ​(st​,at​)=Qπθ​(st​,at​)−V(st​)
上式中， V ( s t ) V(s_t) V(st​)可以用critic估计得到， Q π θ ( s t , a t ) Q^{\\pi_{\\theta}}(s_t,a_t) Qπθ​(st​,at​)可以用下式计算得到
Q π θ ( s t , a t ) = r t + γ ⋅ r t + 1 + . . . + γ T − t ⋅ r T = r t + γ ⋅ ( r t + 1 + . . . + γ T − t − 1 ⋅ r T ) = r t + γ ⋅ Q π θ ( s t + 1 , a t + 1 ) (3) \\begin{aligned} Q^{\\pi_{\\theta}}(s_t,a_t)&=r_t+\\gamma \\cdot r_{t+1}+...+\\gamma^{T-t} \\cdot r_{T} \\\\ &=r_t+\\gamma \\cdot(r_{t+1}+...+\\gamma^{T-t-1} \\cdot r_{T}) \\\\ &=r_t+\\gamma \\cdot Q^{\\pi_{\\theta}}(s_{t+1},a_{t+1}) \\tag3 \\end{aligned} Qπθ​(st​,at​)​=rt​+γ⋅rt+1​+...+γT−t⋅rT​=rt​+γ⋅(rt+1​+...+γT−t−1⋅rT​)=rt​+γ⋅Qπθ​(st+1​,at+1​)​(3)
以下程序描述了如何计算Q值，输入是一个列表，其中的每一个元素都是一条完整的奖励序列 ( r 1 , . . . , r T ) (r_1,...,r_T) (r1​,...,rT​)。根据式(3)，在实际计算时，我们应该以从后向前的顺序计算Q值。

def compute_rtgs(self, batch_rews):
    \"\"\"
    Compute the Reward-To-Go of each timestep in a batch given the rewards.
    :param batch_rews: [[r_1,r_2,...,r_T], [], ..., []]
    :return:
    \"\"\"
    batch_rtgs = []
	# 因为batch_rtgs是以将新元素插入到起始位置的形式添加新元素的，
    # 所以这里需要反向列表来保持顺序的一致
    for ep_rews in reversed(batch_rews):

        discounted_reward = 0  # The discounted reward so far
        # 这里反向的原因源于式(3)
        for rew in reversed(ep_rews):
            discounted_reward = rew + discounted_reward * self.gamma
            # 在列表的起始位置插入元素
            batch_rtgs.insert(0, discounted_reward)  

    # Convert the rewards-to-go into a tensor
    batch_rtgs = torch.tensor(batch_rtgs, dtype=torch.float).to(self.device)

    return batch_rtgs

1.4. R ^ t \\hat{R}_t R^t​

在参考源码中， R ^ t \\hat{R}_t R^t​的取值与 Q π θ ( s t , a t ) Q^{\\pi_{\\theta}}(s_t,a_t) Qπθ​(st​,at​)相同。我不知道是否还有别的方法可以表示 R ^ t \\hat{R}_t R^t​。

2.收集数据

PPO会使用上一轮的actor（ π θ k \\pi_{\\theta_k} πθk​​）与环境互动来收集训练数据。那么，具体应该收集哪些数据呢？

根据前文内容，可以很容易的知道，我们收集的数据需要包含以下内容：

• 状态序列 ( s 1 , s 2 , . . . , s T − 1 ) (s_1,s_2,...,s_{T-1}) (s1​,s2​,...,sT−1​)；
• 动作序列 ( a 1 , a 2 , . . . , a T − 1 ) (a_1,a_2,...,a_{T-1}) (a1​,a2​,...,aT−1​)；
• 奖励序列 ( r 1 , r 2 , . . . , r T − 1 ) (r_1,r_2,...,r_{T-1}) (r1​,r2​,...,rT−1​)；

以上都只是一个episode的内容。实际上，我们会收集多个episode的数据，然后一起训练。
D k : { ( s 1 1 , a 1 1 , r 1 1 , . . . , s T 1 − 1 1 ) , . . . , ( s 1 n , a 1 n , r 1 n , . . . , s T n − 1 n ) } D_k: \\Big\\{(s^1_1,a^1_1,r^1_1,...,s^1_{T_1-1}), ..., (s^n_1,a^n_1,r^n_1,...,s^n_{T_n-1})\\Big\\} Dk​:{(s11​,a11​,r11​,...,sT1​−11​),...,(s1n​,a1n​,r1n​,...,sTn​−1n​)}
如此，一个batch的状态和动作数据可以表示为
state batch : { s 1 1 , s 2 1 , . . . , s T 1 − 1 1 , . . . , s 1 n , s 2 n , . . . , s T n − 1 n } action batch : { a 1 1 , a 2 1 , . . . , a T 1 − 1 1 , . . . , a 1 n , a 2 n , . . . , a T n − 1 n } reward-to-go batch : { g 1 1 , g 2 1 , . . . , g T 1 − 1 1 , . . . , g 1 n , g 2 n , . . . , g T n − 1 n } \\begin{aligned} \\text{state batch}&: \\Big\\{ s^1_1,s^1_2,...,s^1_{T_1-1},...,s^n_1,s^n_2,...,s^n_{T_n-1} \\Big\\} \\\\ \\text{action batch}&: \\Big\\{ a^1_1,a^1_2,...,a^1_{T_1-1},...,a^n_1,a^n_2,...,a^n_{T_n-1} \\Big\\} \\\\ \\text{reward-to-go batch}&: \\Big\\{ g^1_1,g^1_2,...,g^1_{T_1-1},...,g^n_1,g^n_2,...,g^n_{T_n-1} \\Big\\} \\\\ \\end{aligned} state batchaction batchreward-to-go batch​:{s11​,s21​,...,sT1​−11​,...,s1n​,s2n​,...,sTn​−1n​}:{a11​,a21​,...,aT1​−11​,...,a1n​,a2n​,...,aTn​−1n​}:{g11​,g21​,...,gT1​−11​,...,g1n​,g2n​,...,gTn​−1n​}​

• 假设动作是连续的，参数为 θ k \\theta_k θk​的actor接受state batch的输入，可以得到一个action的平均值，然后以此构建一个多维正态分布，得到 log ⁡ π θ k ( a t ∣ s t ) \\log{\\pi_{\\theta_k}(a_t|s_t)} logπθk​​(at​∣st​)；
• 参数为 ϕ k \\phi_{k} ϕk​的critic接受state batch的输入，可以得到一个状态的评价值 V ϕ k ( s t ) V_{\\phi_k}(s_t) Vϕk​​(st​)；
• 根据式(3)，可以由奖励序列计算得到reward-to-go batch，也就是 Q π θ ( s t , a t ) Q^{\\pi_{\\theta}}(s_t,a_t) Qπθ​(st​,at​)和 R ^ t \\hat{R}_t R^t​。前者可以用于计算 A π θ k ( s t , a t ) A^{\\pi_{\\theta_k}}(s_t,a_t) Aπθk​​(st​,at​)。

由此，我们已经得到了第k轮全部的训练资料。

# collect trajectories by the past actor
batch_obs, batch_acts, batch_log_probs, batch_rtgs, batch_lens = self.collect_trajectories()

# Calculate advantage at k-th iteration
V, _ = self.evaluate(batch_obs, batch_acts)
A_k = batch_rtgs - V.detach()  # ALG STEP 5

# update PPO
for _ in range(self.update_time_per_iteration):  # ALG STEP 6 & 7
    # Calculate V_phi and pi_theta(a_t | s_t)
    V, curr_log_probs = self.evaluate(batch_obs, batch_acts)

    # Calculate the ratio pi_theta(a_t | s_t) / pi_theta_k(a_t | s_t)
    ratios = torch.exp(curr_log_probs - batch_log_probs)

    # Calculate surrogate losses.
    surr1 = ratios * A_k
    surr2 = torch.clamp(ratios, 1 - self.clip, 1 + self.clip) * A_k

    actor_loss = (-torch.min(surr1, surr2)).mean()
    critic_loss = torch.nn.MSELoss()(V, batch_rtgs)

    self.actor_optimizer.zero_grad()
    actor_loss.backward()
    self.actor_optimizer.step()

    self.critic_optimizer.zero_grad()
    critic_loss.backward()
    self.critic_optimizer.step()

以上程序描述了完整的PPO的更新过程。可以发现，对于一笔训练资料，PPO并不是更新一次就结束了，而是会连续更新多次。这样可以加快训练速度，至于为什么可以这么做可以参考李宏毅老师关于PPO的视频，这里就不在赘述。