Proximal Policy Optimization(PPO)和文本生成

ChatGPT的RLHF步使用了强化学习PPO算法。
PPO是一种策略梯度方法，其交替地进行与环境交互采样数据和使用随机梯度上升优化“代理”目标函数。标准策略梯度方法对每个数据样本执行一次梯度更新，而PPO可以采样一批数据后，对模型进行多次梯度更新。

策略梯度

策略梯度（Policy Gradient）方法梯度的计算如下：
$${\mathbb{E}}_{(a_t,s_t)\in {\pi }_{\theta }}[{\hat{A}}_t{\nabla }_{\theta }\log {\pi }_{\theta }(a_t|s_t)]$$${\hat{A}}_t$是优势函数（advantage function）$A_t$的估计。
$$A_t=Q(s_t,a_t)-V(s_t)$$优势函数计算的是，在该状态下采取这个行动的奖励与在该状态下的平均奖励的差值。
上面的导数可以通过对下面的目标求导获得：
$$L^{PG}(\theta )={\mathbb{E}}_{(a_t,s_t)\in {\pi }_{\theta }}[{\hat{A}}_t\log {\pi }_{\theta }(a_t|s_t)]$$

PPO（Proximal Policy Optimization）

PPO有两个形式，其中一种形式PPO_CLIP的优化目标函数是：
$$L^{CLIP}(\theta )={\mathbb{E}}_{(a_t,s_t)\in {\pi }_{{\theta }_{old}}}[\min (r_t(\theta ){\hat{A}}_t,clip(r_t(\theta ),1-ϵ,1+ϵ){\hat{A}}_t)]\text{\hspace{1em}(1)}$$其中$r_t(\theta )=\frac{{\pi }_{\theta }(a_t|s_t)}{{\pi }_{{\theta }_{old}}(a_t|s_t)}$。
PPO算法中的advantage用下面的公式估计：
$${\hat{A}}_t={\delta }_t+(\gamma \lambda ){\delta }_{t+1}+\cdots +(\gamma \lambda {)}^{T-t+1}{\delta }_{T-1}$$其中 ${\delta }_t=r_t+\gamma V(s_{t+1})-V(s_t)$，$r_t$是reward，$V(s_t)$是状态价值。
通常情况下，我们用一个网络同时学习策略和价值函数，这样策略和价值函数能共享参数，那么就需要结合策略代理和价值函数误差项的损失函数。再加上熵奖励（entropy bonus）来确保足够的探索，优化目标变为：
$$L^{CLIP+VF+S}(\theta )={\mathbb{E}}_{(a_t,s_t)\in {\pi }_{\theta }}[L_t^{CLIP}(\theta )-c_1{L}_t^{VF}(\theta )+c_{2}S[{\pi }_{\theta }](s_t)]$$其中$L_t^{VF}(\theta )=(V_{\theta }(s_t)-V_t^{targ}{)}^2$是价值函数的误差项，$S[{\pi }_{\theta }]$是entropy bonus。

完整的PPO算法如下。可以看到每个循环中，先采样N个T个时间步的数据，然后用采样的数据进行K个epoch的优化。

文本生成

在文本生成的情况下，给一个prompt，生成完整的response，是一个episode。动作空间是vocabulary。每生成一个词是一个时间步。

公式(1)需要advantage的估计，为了计算advantage，我们需要定义奖励（reward）$r$和估计状态价值函数$V(s)$。

用于强化学习的reward计算如下：
$$R(x,y)=r(x,y)-\beta \log \frac{\pi (y|x)}{\rho (y|x)}$$x是prompt，y是response，$r(x,y)$是reward model的输出，也就是下面代码中的score。注意这里reward model的输出称之为score，送入强化学习部分的才称为reward。$\pi (y|x)$是要学习的生成模型，$\rho (y|x)$是参数固定的原始生成模型。

在trl库中reward的计算如下。只将reward model的score添加到最后一个token的reward上，其他token的reward来自当前模型和 原始生成模型之间KL散度。这么做是为了减轻奖励模型的过度优化问题。

   def compute_rewards(
       self,
       scores: torch.FloatTensor,
       logprobs: torch.FloatTensor,
       ref_logprobs: torch.FloatTensor,
       masks: torch.LongTensor,
   ):
       """
       Compute per token rewards from scores and KL-penalty.

       Args:
           scores (`torch.FloatTensor`):
               Scores from the reward model, shape (`batch_size`)
           logprobs (`torch.FloatTensor`):
               Log probabilities of the model, shape (`batch_size`, `response_length`)
           ref_logprobs (`torch.FloatTensor`):
               Log probabilities of the reference model, shape (`batch_size`, `response_length`)
       """
       rewards, non_score_rewards = [], []
       for score, logprob, ref_logprob, mask in zip(scores, logprobs, ref_logprobs, masks):
           # compute KL penalty (from difference in logprobs)
           kl = self._kl_penalty(logprob, ref_logprob)
           non_score_reward = -self.kl_ctl.value * kl
           non_score_rewards.append(non_score_reward)
           reward = non_score_reward.clone()
           last_non_masked_index = mask.nonzero()[-1]

           # reward is preference model score + KL penalty
           reward[last_non_masked_index] += score
           rewards.append(reward)
       return torch.stack(rewards), torch.stack(non_score_rewards)

在trl库中用一个网络AutoModelForCausalLMWithValueHead学习策略${\pi }_{\theta }(s)$和状态价值函数$V(s)$。AutoModelForCausalLMWithValueHead在普通AutoModelForCausalLM模型上了一个线性层nn.Linear(hidden_size, 1)，用于估计状态价值函数$V(s)$。
普通AutoModelForCausalLM模型估计token概率即可作为策略${\pi }_{\theta }(s)$。

在trl库中advantage的计算如下。compute_advantage函数返回的returns作为$V_t^{targ}$用于学习状态价值函数$V(s)$。

    def compute_advantages(
        self: torch.FloatTensor,
        values: torch.FloatTensor, # AutoModelForCausalLMWithValueHead输出的状态价值估计V
        rewards: torch.FloatTensor, # compute_rewards函数计算得到的rewards
        mask: torch.FloatTensor,
    ):
        lastgaelam = 0
        advantages_reversed = []
        gen_len = rewards.shape[-1]

        values = values * mask
        rewards = rewards * mask

        for t in reversed(range(gen_len)):
            nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0
            delta = rewards[:, t] + self.config.gamma * nextvalues - values[:, t]
            lastgaelam = delta + self.config.gamma * self.config.lam * lastgaelam
            advantages_reversed.append(lastgaelam)
        advantages = torch.stack(advantages_reversed[::-1]).transpose(0, 1)

        returns = advantages + values
        advantages = masked_whiten(advantages, mask)
        advantages = advantages.detach()
        return values, advantages, returns

完整的loss计算可以参看ppo_trainer.py中的loss函数，因为函数较长，这里不粘贴了。
因为在reward的计算中可以考虑了entropy bonus，所以在更新网络的时候，没有再用$S[{\pi }_{\theta }]$的loss。

    def loss(
        self,
        old_logprobs: torch.FloatTensor,
        values: torch.FloatTensor,
        logits: torch.FloatTensor,
        vpreds: torch.FloatTensor,
        logprobs: torch.FloatTensor,
        mask: torch.LongTensor,
        advantages: torch.FloatTensor,
        returns: torch.FloatTensor,
    )

Reference

《Proximal Policy Optimization Algorithms》
《Fine-Tuning Language Models from Human Preferences》
《Training language models to follow instructions with human feedback》
https://github.com/huggingface/trl