此外,我们知道在强化学习中通常还需要去让智能体能够具有 exploratory 的表现,这样才能挖掘出更多具有高价值的行为,所以 PPO 算法在训练策略网络时还在目标函数中加上了和熵相关的一项奖励 :
这里以离散动作为例,可以看到如果实施动作的概率分散到不同的动作上将具有更大的熵。
前面介绍策略梯度时知道策略梯度中还涉及价值函数/优势函数的估计,在 PPO 算法中也是采用神经网络来估计状态价值函数,训练价值网络的目标函数通常仅需要最小化价值网络的预测和目标的平方误差就可以了
综上所述,PPO 算法完整的优化目标函数由三部分组成,为
PPO 算法代码解读了解完原理后,我们来看一下强化的代码一般怎么写的。
训练智能体前一般需要定义一个环境
envs = xxxx
环境需要具有两个主要的功能函数,一个是 step,它的输入是动作,输出是下一步的观测、奖励、以及表示环境是否结束等额外信息
next_obs, reward, done, info = envs.step(action)
另一个是 reset,主要用来重置环境。
然后需要定义一个智能体,智能体包含策略网络和价值网络两部分(也就是演员和评论家),get_value 函数使用价值网络评估状态的价值,get_action_and_value 函数使用策略网络给出了某个状态下动作的概率分布(以及对数概率)、概率分布采样得到的动作,概率分布的熵、以及状态的价值。
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 的完整代码可以参考https://github.com/huggingface/deep-rl-class/blob/main/unit8/unit8.ipynb
也建议大家有空可以阅读一下这篇文章,有助于加强对强化的理解:
https://fse.studenttheses.ub.rug.nl/25709/1/mAI_2021_BickD.pdf
利用奖励模型强化 GPT:迈向 ChatGPT看到了这里,我们已经学会啦如何训练一个 GPT,训练一个奖励模型,以及强化学习算法 PPO 的训练过程,将它们组合起来,就能够做出一个 ChatGPT 了!
这里我们先定义两个 GPT 模型,一个用于强化,一个用于参考(因为我们通常也不希望强化的模型完全朝着更高的奖励去优化,所以可以约束一下强化的模型和原始模型仍然具有一定的“相似性”),参考模型不用于优化,所以设为 eval 模式即可。
model = GPT()
ref_model = GPT()
ref_model.eval()
需要注意的是为了适应强化这个框架,我们还需要对原来的 GPT 模型进行一定的封装,其实主要就是加一层 value head(线性层),让它预测每一个 token 的价值(理解为值函数,将 token 的隐藏状态转化为一个标量值)。
为了让大家更加清楚地了解训练的过程,我们以一条样本为例来展示说明数据的传递过程。
假设我们有一个 query(就是送入 gpt 的输入),前面说过模型的输入一般经过分词器转化为一串 id,比如:
送入 GPT 后,生成模型会接着输入进行“续写”,得到一条回答 response:
然后我们可以把这 query 和 response 拼接起来并解码为文本后送入奖励模型,得到一个 rewards(可以理解 chatgpt 的回答好,就意味这把提问和模型的回答拼在一起看是合理的),接下来就可以考虑使用 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
接下来进行模型的前向传播(这里是需要梯度的),并索引出 response 部分的动作概率和值函数:
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 了,也就是 ChatGPT 实现的核心思想。
可以看上面的通过强化 GPT 再用于生成的一个效果:
以及将结果再用奖励模型打分的前后对比:
mean:
rewards (before) 0.156629
rewards (after) 1.686487
median:
rewards (before) -0.547091
rewards (after) 2.479868
可以看到 GPT 生成结果的整体得分都有比较大的提升,这也说明了如果奖励模型训练得好,那用来做 ChatGPT 的效果自然也就能够大幅度提高。
本章的完整代码可以参考开源项目:
https://github.com/lvwerra/trl
写在最后,其实个人觉得 ChatGPT 能够这么火(或者说效果能够做得这么好)主要还是来自于大量高质量的人工标注数据(第一步的监督,第二步的奖励模型的训练),验证了一句话,人工智能就是需要“人工”才能走向“智能”哈哈。
Reference
- https://d4mucfpksywv.cloudfront.net/better-language-models/language-models.pdf
- https://s3-us-west-2.amazonaws.com/openai-assets/research-covers/language-unsupervised/language_understanding_paper.pdf
- https://arxiv.org/pdf/2005.14165.pdf
- https://openai.com/blog/chatgpt/
- https://huggingface.co/docs/transformers/model_doc/gpt2
- https://huggingface.co/blog/rlhf
- https://huggingface.co/blog/deep-rl-ppo
- https://fse.studenttheses.ub.rug.nl/25709/1/mAI_2021_BickD.pdf
- https://github.com/lvwerra/trl
- https://github.com/wangshusen/DRL