同样的我们需要对其进行分词，转为一串 id，定义分词器，并构造数据集：
print("Tokenizing train, validate, test text ")
tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-uncased')
train_encodings = tokenizer(train_texts, truncation=True, padding=True)
test_encodings = tokenizer(test_texts, truncation=True, padding=True)
class IMDbDataset(torch.utils.data.Dataset):
def __init__(self, encodings, labels):
self.encodings = encodings
self.labels = labels
def __getitem__(self, idx):
item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
item['labels'] = torch.tensor(self.labels[idx])
return item
def __len__(self):
return len(self.labels)
print("Loading tokenized text into Pytorch Datasets ")
train_dataset = IMDbDataset(train_encodings, train_labels)
test_dataset = IMDbDataset(test_encodings, test_labels)
train_loader = DataLoader(train_dataset, batch_size=10, shuffle=True)
接下来，定义一个用于预测得分的模型以及优化器。
print("Loading pre-trained DistilBERT model ")
model = DistilBertForSequenceClassification.from_pretrained('distilbert-base-uncased')
model.to(device)
model.train()
optim = AdamW(model.parameters(), lr=5e-5)
然后，开始模型的训练：
for epoch in range(3):
for (b_ix, batch) in enumerate(train_loader):
optim.zero_grad()
input_ids = batch['input_ids'].to(device)
attention_mask = batch['attention_mask'].to(device)
labels = batch['labels'].to(device)
outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
loss = outputs[0]
loss.backward()
optim.step()
训练完模型后，我们就可以对任意文本来进行情感的打分了，比如：
I like you,I love you
显然这是一句很 的话，模型的打分为：
详细的细节，可以参考James D. 的博客：。
模型的选择问题，这里作为演示采样了一个 作为奖励模型，但实际上 用的奖励模型的规模应该是和生成模型差不多大小的模型。
到目前为止，我们已经能够训练一个能够生成文本的 GPT，一个能够对文本打分的奖励模型，接下来就需要考虑如果用奖励模型来教会 GPT 生成高分文本了。考虑到很多同学可能没有强化学习的基础，这里将插入一章强化学习的介绍。
三、强化学习
强化学习的内容其实很多，背后涉及的数学也挺多，身边很多人在入门过程就被劝退，为了让大家更易于接受，这里主要介绍一些必备的基础知识，以及从策略梯度方法到 PPO 算法的演进。
3.1 强化术语
在强化学习中，智能体和环境进行交互并基于一定的奖励机制来提高自身的水平；智能体的决策，一般称为策略（）决定了在当前环境状态下智能体应该去实施什么动作；实施完动作后，智能体会随着环境进入下一个状态，并获得相应的奖励或惩罚。这里附上两页 帮助大家理解。
强化学习的最终目的就是要学会一个使得智能体能够最大化期望回报的 ，其中的回报就是对奖励进行衰减求和：
和期望回报紧密相关的还有两个概念，一个是动作状态价值函数（观测到状态 ，做完决策，选中动作
）：
另一个是状态价值函数（可以理解为比如下围棋时评价当前状态
的胜率）
了解了上面的定义后我们可以停下来思考一下，强化学习学什么？
主要想学的肯定就是策略 函数，也就是从状态到动作的一个映射，如果直接学习它，那就能够拿来使用了，这类方法也叫做基于 的方法；如果采用间接点的方式，也可以学习值函数，然后根据值的大小来选择动作，这类方法也叫做基于 value 的方法。当然，通常基于 的方法也会涉及到值函数的近似。
3.2 从策略梯度到 PPO
我们要学一个策略函数，但是并不知道策略函数长什样，怎么去定义它才是合适的。好在有了深度学习这一工具，我们可以无脑用一个神经网络来近似策略函数，然后通过优化神经网络参数的方式来学习得到一个策略函数。
优化神经网络的参数需要有个目标函数，如果一个策略很好，那么状态价值的均值应当很大，因此我们定义目标函数:
这个目标函数排除掉了状态S的因素，只依赖于策略网络π的参数Θ ；策略越好，则
越大。所以策略学习可以描述为这样一个优化问题：
我们希望通过对策略网络参数Θ的更新，使得目标函数
越来越大，也就意味着策略 网络越来越强。想要求解最大化问题，显然可以用梯度上升更新模型的参数。值得庆幸的是，策略函数的梯度还能被推导出来：
策略梯度定理的详细推导这里就不展开了，我们需要记住的是能计算出目标函数关于参数的梯度，那就能用来更新参数，也就能学习出策略函数了。
当然这里面还涉及动作价值函数Q的估计，如果用实际观测的回报来近似，那就是 算法，如果再用一个神经网络来近似这个价值函数，那就是演员-评论家算法。PS：在实际使用中，策略梯度中的 Q 有多种不同的替代形式，常见效果比较好的形式是采用优势函数A（状态动作值函数Q减去状态值函数V ）来替代。
传统的策略梯度算法的局限性在于它是 – 的，也就是说每次获取的训练数据只被用来更新一次模型的参数后就丢掉了，因此 PPO 算法的主要改进在于构造了新的目标函数（避免较大的参数变化），使得每次获取的训练数据能够被用于多次的参数更新。
其中，比值函数为当前策略和历史策略在状态St下实施动作At的概率的比值：
通过这一比值也就能够评估新旧策略的差异性，从而能够保证策略函数在更新参数时不会跟旧策略的差异太大。有时间的同学也可以对比值在不同区间时目标函数的情况进行考虑，也就是如下表的情况。
此外，我们知道在强化学习中通常还需要去让智能体能够具有 的表现，这样才能挖掘出更多具有高价值的行为，所以 PPO 算法在训练策略网络时还在目标函数中加上了和熵相关的一项奖励H ：
这里以离散动作为例，可以看到如果实施动作的概率分散到不同的动作上将具有更大的熵。
前面介绍策略梯度时知道策略梯度中还涉及价值函数/优势函数的估计，在 PPO 算法中也是采用神经网络来估计状态价值函数，训练价值网络的目标函数通常仅需要最小化价值网络的预测和目标的平方误差就可以了：
综上所述，PPO 算法完整的优化目标函数由三部分组成：
3.3 PPO 算法代码解读
了解完原理后，我们来看一下强化的代码一般怎么写的。训练智能体前一般需要定义一个环境：
envs = xxxx
环境需要具有两个主要的功能函数，一个是 step，它的输入是动作，输出是下一步的观测、奖励、以及表示环境是否结束等额外信息：
next_obs, reward, done, info = envs.step(action)
另一个是 reset，主要用来重置环境。
然后需要定义一个智能体，智能体包含策略网络和价值网络两部分（也就是演员和评论家）， 函数使用价值网络评估状态的价值， 函数使用策略网络给出了某个状态下动作的概率分布（以及对数概率）、概率分布采样得到的动作，概率分布的熵、以及状态的价值。
class Agent(nn.Module):
def __init__(self, envs):
super().__init__()
self.critic = nn.Sequential(
layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), 64)),
nn.Tanh(),
layer_init(nn.Linear(64, 64)),
nn.Tanh(),
layer_init(nn.Linear(64, 1), std=1.0),
)
self.actor = nn.Sequential(
layer_init(nn.Linear(np.array(envs.single_observation_space.shape).prod(), 64)),
nn.Tanh(),
layer_init(nn.Linear(64, 64)),
nn.Tanh(),
layer_init(nn.Linear(64, envs.single_action_space.n), std=0.01),
)
def get_value(self, x):
return self.critic(x)
def get_action_and_value(self, x, action=None):
logits = self.actor(x)
probs = Categorical(logits=logits)
if action is None:
action = probs.sample()
return action, probs.log_prob(action), probs.entropy(), self.critic(x)
接下来就只需要考虑如何收集数据和训练网络了。收集数据阶段主要包括两个部分，一部分是用智能体去和环境做交互，并保存相应的状态、动作等信息，另一部分主要是根据每一步的奖励来计算每一步的回报，从而计算用于评估动作好坏的优势函数值。
for step in range(0, args.num_steps):
obs[step] = next_obs
dones[step] = next_done
with torch.no_grad():
action, logprob, _, value = agent.get_action_and_value(next_obs)
values[step] = value.flatten()
actions[step] = action
logprobs[step] = logprob
next_obs, reward, done, info = envs.step(action.cpu().numpy())
rewards[step] = torch.tensor(reward).to(device).view(-1)
with torch.no_grad():
next_value = agent.get_value(next_obs).reshape(1, -1)
returns = torch.zeros_like(rewards).to(device)
for t in reversed(range(args.num_steps)):
if t == args.num_steps - 1:
nextnonterminal = 1.0 - next_done
next_return = next_value
else:
nextnonterminal = 1.0 - dones[t + 1]
next_return = returns[t + 1]
returns[t] = rewards[t] + args.gamma * nextnonterminal * next_return
advantages = returns - values
训练网络部分主要就是根据之前提到的三部分目标函数，依次计算新旧策略的差异从而计算策略网络的损失、价值网络的损失 以及 关于动作多样性的熵奖励；累加后再进行反向传播就可以更新网络的参数了。
for epoch in range(args.update_epochs):
np.random.shuffle(b_inds)
for start in range(0, args.batch_size, args.minibatch_size):
end = start + args.minibatch_size
mb_inds = b_inds[start:end]
# 计算新旧策略的差异
_, newlogprob, entropy, newvalue = agent.get_action_and_value(b_obs[mb_inds], b_actions.long()[mb_inds])
logratio = newlogprob - b_logprobs[mb_inds]
ratio = logratio.exp()
mb_advantages = b_advantages[mb_inds]
# 策略网络损失
pg_loss1 = -mb_advantages * ratio
pg_loss2 = -mb_advantages * torch.clamp(ratio, 1 - args.clip_coef, 1 + args.clip_coef)
pg_loss = torch.max(pg_loss1, pg_loss2).mean()
# 价值网络损失
newvalue = newvalue.view(-1)
v_loss = 0.5 * ((newvalue - b_returns[mb_inds]) ** 2).mean()
entropy_loss = entropy.mean()
loss = pg_loss - args.ent_coef * entropy_loss + v_loss * args.vf_coef
optimizer.zero_grad()
loss.backward()
nn.utils.clip_grad_norm_(agent.parameters(), args.max_grad_norm)
optimizer.step()
PPO 的完整代码可以参考
也建议大家有空可以阅读一下这篇文章，有助于加强对强化的理解：
四、利用奖励模型强化 GPT
看到了这里，我们已经学会啦如何训练一个 GPT，训练一个奖励模型，以及强化学习算法 PPO 的训练过程，将它们组合起来，就能够做出一个 了！
这里我们先定义两个 GPT 模型，一个用于强化，一个用于参考（因为我们通常也不希望强化的模型完全朝着更高的奖励去优化，所以可以约束一下强化的模型和原始模型仍然具有一定的“相似性”），参考模型不用于优化，所以设为 eval 模式即可。
model = GPT()
ref_model = GPT()
ref_model.eval()
需要注意的是为了适应强化这个框架，我们还需要对原来的 GPT 模型进行一定的封装，其实主要就是加一层 value head（线性层），让它预测每一个 token 的价值（理解为值函数，将 token 的隐藏状态转化为一个标量值）。
为了让大家更加清楚地了解训练的过程，我们以一条样本为例来展示说明数据的传递过程。假设我们有一个 query（就是送入 gpt 的输入），前面说过模型的输入一般经过分词器转化为一串 id，比如：
送入 GPT 后，生成模型会接着输入进行“续写”，得到一条回答 ：
然后我们可以把这 query 和 拼接起来并解码为文本后送入奖励模型，得到一个 （可以理解 的回答好，就意味这把提问和模型的回答拼在一起看是合理的），接下来就可以考虑使用 PPO 算法来对 GPT 进行强化了。
首先，第一步我们就是把“提问”和“回答”拼接起来，送给奖励模型，让它给 GPT 的回答打个分：
texts = [q + r for q, r in zip(query, response)]
score = reward_model(texts)[1]["score"]
有了这个得分后，就可以让 GPT 模型朝着尽可能高分的方向去优化了。
把拼接后的文本送入模型，并取出对应 token 的（对数）概率，这里并不需要梯度的传播，主要是获得“旧”的动作（模型的回答），以及用于计算每一个 token 的奖励：
with torch.no_grad():
logits, _, v = self.model(**input_kwargs)
ref_logits, _, _ = self.ref_model(**input_kwargs)
old_logprobs = logprobs_from_logits(logits[:, :-1, :], input_ids[:, 1:])
ref_logprobs = logprobs_from_logits(ref_logits[:, :-1, :], input_ids[:, 1:])
计算奖励这里考虑两个部分，一部分是用于强化的 GPT 模型的回答和参考模型回答的 KL 散度（如前面所说，不能让模型一味朝着奖励高的方向优化），另一部分就是奖励模型给出的评分：
kl = old_logprobs - ref_logprob
reward = -kl
reward[-1] += score
接下来进行模型的前向传播（这里是需要梯度的），并索引出 部分的动作概率和值函数：
logits, _, vpred = self.model(**input_kwargs)
logprob = logprobs_from_logits(logits[:, :-1, :], model_input[:, 1:])
logprob, vpred = logprob[:, -gen_len:], vpred[:, -gen_len:]
计算 loss 之前我们需要估计一下优势函数和回报，优势函数用于更新策略网络，回报用于更新价值网络，这里采用经典的 GAE 方法（可以有效降低策略梯度的估计方差）来估计优势函数：
lastgaelam = 0
advantages_reversed = []
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
计算策略网络的损失：
ratio = torch.exp(logprob - old_logprobs)
pg_losses = -advantages * torch.clamp(ratio, 1.0 - self.config.cliprange, 1.0 + self.config.cliprange)
pg_loss = torch.mean(pg_losses)
价值函数的损失：
vf_losses = (vpred - returns) ** 2
vf_loss = 0.5 * torch.mean(vf_losses)
整体误差反向传播：
loss = pg_loss + self.config.vf_coef * vf_loss
optimizer.zero_grad()
accelerator.backward(loss)
optimizer.step()
这样就完成了使用 PPO 来强化 GPT 的一个 step 了，也就是 实现的核心思想。
可以看上面的通过强化 GPT 再用于生成的一个效果：
以及将结果再用奖励模型打分的前后对比：
mean:
rewards (before) 0.156629
rewards (after) 1.686487
median:
rewards (before) -0.547091
rewards (after) 2.479868
可以看到 GPT 生成结果的整体得分都有比较大的提升，这也说明了如果奖励模型训练得好，那用来做 的效果自然也就能够大幅度提高。
本章的完整代码可以参考开源项目：
参考：
323AI导航网发布
相关文章
