大模型强化学习——PPO项目实战

【PPO算法介绍】

PPO（Proximal Policy Optimization）是一种强化学习算法，它的目标是找到一个策略，使得根据这个策略采取行动可以获得最大的累积奖励。

PPO的主要思想是在更新策略时，尽量让新策略不要偏离旧策略太远。这是通过在目标函数中添加一个额外的项来实现的，这个额外的项会惩罚新策略和旧策略之间的差异。这样可以避免在更新策略时出现过大的跳跃，从而提高学习的稳定性。

举一个通俗的例子，假设你正在玩一个游戏，你已经找到了一个还不错的策略，可以让你获得一定的分数。现在，你想要改进你的策略，以便获得更高的分数。但是，你不希望新的策略和旧的策略差距太大，因为这可能会让你的分数大幅度下降。所以，你会尽量让新的策略和旧的策略相近，这就是PPO的主要思想。

在实际操作中，PPO通过使用一种叫做“裁剪目标函数”的技术来实现这个思想。这个技术会限制新策略和旧策略之间的差异，从而避免更新过程中的大幅度跳跃。

项目简介

RL阶段实战，通过强化学习PPO算法对SFT模型进行优化，帮助读者深入理解ChatGPT模型在RL阶段的任务流程。

项目主要结构如下：

• data 存放数据的文件夹
  • ppo_train.json 用于强化学习的文档数据
• rm_model RM阶段训练完成模型的文件路径
  • config.json
  • pytorch_model.bin
  • vocab.txt
• sft_model SFT阶段训练完成模型的文件路径
  • config.json
  • pytorch_model.bin
  • vocab.txt
• ppo_model PPO阶段训练完成模型的文件路径
  • config.json
  • pytorch_model.bin
  • vocab.txt
• data_set.py 模型所需数据类文件
• model.py 模型文件
• train.py 模型训练文件
• predict.py 模型推理文件

注意：由于GitHub不方便放模型文件，因此sft_model文件夹、rm_model文件夹和ppo_model文件夹中的模型bin文件，请从百度云盘中下载。

环境配置

模型训练或推理所需环境，请参考requirements.txt文件。

模型训练

模型的训练流程如下所示：

模型训练需要运行train.py文件，会自动生成output_dir文件夹，存放每个epoch保存的模型文件。

命令如下：

python3 train.py --device 0 \
                 --train_file_path "data/ppo_train.json" \
                 --max_len 768 \
                 --query_len 64 \
                 --batch_size 16

注意：当服务器资源不同或读者更换数据等时，可以在模型训练时修改响应参数，详细参数说明见代码或阅读书9.3.3小节。

模型训练示例如下：

模型训练阶段损失值变化如下：

模型推理

模型训练需要运行predict.py文件，可以采用项目中以提供的模型，也可以采用自己训练后的模型。

命令如下：

python3 predict.py --device 0 --model_path ppo_model

注意：如果修改模型路径，请修改--model_path参数。

模型推理示例如下：

样例1：
输入的正文为：大莱龙铁路位于山东省北部环渤海地区，西起位于益羊铁路的潍坊大家洼车站，向东经海化、寿光、寒亭、昌邑、平度、莱州、招远、终到龙口，连接山东半岛羊角沟、潍坊、莱州、龙口四个港口，全长175公里，工程建设概算总投资11.42亿元。铁路西与德大铁路、黄大铁路在大家洼站接轨，东与龙烟铁路相连。
生成的第1个问题为：该项目建成后对于做什么？
生成的第2个问题为：该铁路线建成后会带动什么方面？
样例2：
输入的正文为：椰子猫（学名：'），又名椰子狸，为分布于南亚及东南亚的一种麝猫。椰子猫平均重3.2公斤，体长53厘米，尾巴长48厘米。它们的毛粗糙，一般呈灰色，脚、耳朵及吻都是黑色的。它们的身体上有三间黑色斑纹，面部的斑纹则像浣熊，尾巴没有斑纹。椰子猫是夜间活动及杂食性的。它们在亚洲的生态位与在北美洲的浣熊相近。牠们栖息在森林、有树木的公园及花园之内。它们的爪锋利，可以用来攀爬。椰子猫分布在印度南部、斯里兰卡、东南亚及中国南部。
生成的第1个问题为：椰子猫是什么族群？
生成的第2个问题为：椰子猫到底是什么？

注意

需要36GB内存才可以训练。我自己单机无法运行：

  File "C:\Python311\Lib\site-packages\transformers\pytorch_utils.py", line 106, in forward
    x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: [enforce fail at ..\c10\core\impl\alloc_cpu.cpp:72] data. DefaultCPUAllocator: not enough memory: you tried to allocate 35684352 bytes.

【核心代码解读】　　

# -*- coding:utf-8 -*-
# @project: ChatGPT
# @filename: train
# @author: 刘聪NLP
# @zhihu: https://www.zhihu.com/people/LiuCongNLP
# @contact: [email protected]
# @time: 2023/4/14 17:40
"""
    文件说明：
            
"""

from utils import get_advantages_and_returns, actor_loss_function, critic_loss_function
from model import ActorModel, RewardModel, CriticModel
import argparse
from data_set import ExamplesSampler
import os
from transformers.models.bert import BertTokenizer
from torch.optim import Adam
import random
import numpy as np
import torch

try:
    from torch.utils.tensorboard import SummaryWriter
except ImportError:
    from tensorboard import SummaryWriter


def make_experience(args, actor_model, critic_model, ori_model, reward_model, input_ids, generate_kwargs):
    """获取经验数据"""
    actor_model.eval()
    critic_model.eval()
    with torch.no_grad():
        # 获取prompt内容长度
        prompt_length = input_ids.shape[1]
        # 使用动作模型通过已有提示生成指定内容，其中：seq_outputs为返回序列，包含prompt+生成的answer
        seq_outputs, attention_mask = actor_model.generate(input_ids, **generate_kwargs)
        # 通过动作模型和原始模型同时计算生成结果对应的log_probs
        action_log_probs = actor_model(seq_outputs, attention_mask)
        base_action_log_probs = ori_model(seq_outputs, attention_mask)
        # 通过评判模型计算生成的answer的分值
        value, _ = critic_model(seq_outputs, attention_mask, prompt_length)
        value = value[:, :-1]
        # 通过奖励模型计算生成奖励值，并对奖励值进行裁剪
        _, reward_score = reward_model.forward(seq_outputs, attention_mask, prompt_length=prompt_length)
        reward_clip = torch.clamp(reward_score, -args.reward_clip_eps, args.reward_clip_eps)
        # reward_clip = reward_score
        # 对动作模型和原始模型的log_probs进行kl散度计算，防止动作模型偏离原始模型
        kl_divergence = -args.kl_coef * (action_log_probs - base_action_log_probs)
        rewards = kl_divergence
        start_ids = input_ids.shape[1] - 1
        action_mask = attention_mask[:, 1:]
        ends_ids = start_ids + action_mask[:, start_ids:].sum(1)
        batch_size = action_log_probs.shape[0]
        # 将奖励值加到生成的answer最后一个token上
        for j in range(batch_size):
            rewards[j, start_ids:ends_ids[j]][-1] += reward_clip[j]
        # 通过奖励值计算优势函数
        advantages, returns = get_advantages_and_returns(value, rewards, start_ids, args.gamma, args.lam)

    experience = {"input_ids": input_ids, "seq_outputs": seq_outputs, "attention_mask": attention_mask,
                  "action_log_probs": action_log_probs, "value": value, "reward_score": reward_score,
                  "advantages": advantages, "returns": returns}
    return experience


def update_model(args, experience_list, actor_model, actor_optimizer, critic_model, critic_optimizer, tb_write,
                 ppo_step):
    """模型更新"""
    # 根据强化学习训练轮数，进行模型更新
    for _ in range(args.ppo_epoch):
        # 随机打乱经验池中的数据，并进行数据遍历
        random.shuffle(experience_list)
        for i_e, experience in enumerate(experience_list):
            ppo_step += 1
            start_ids = experience["input_ids"].size()[-1] - 1

            # 获取actor模型的log_probs
            action_log_probs = actor_model(experience["seq_outputs"], experience["attention_mask"])
            action_mask = experience["attention_mask"][:, 1:]
            # 计算actor模型损失值
            actor_loss = actor_loss_function(action_log_probs[:, start_ids:],
                                             experience["action_log_probs"][:, start_ids:], experience["advantages"],
                                             action_mask[:, start_ids:], args.policy_clip_eps)
            # actor模型梯度回传，梯度更新
            actor_loss.backward()
            tb_write.add_scalar("actor_loss", actor_loss.item(), ppo_step)
            torch.nn.utils.clip_grad_norm_(actor_model.parameters(), args.max_grad_norm)
            actor_optimizer.step()
            actor_optimizer.zero_grad()

            # 计算critic模型的value
            value, _ = critic_model(experience["seq_outputs"], experience["attention_mask"],
                                    experience["input_ids"].size()[-1])
            value = value[:, :-1]
            # 计算critic模型损失值
            critic_loss = critic_loss_function(value[:, start_ids:], experience["value"][:, start_ids:],
                                               experience["returns"], action_mask[:, start_ids:], args.value_clip_eps)
            # actor模型梯度回传，梯度更新
            critic_loss.backward()
            tb_write.add_scalar("critic_loss", critic_loss.item(), ppo_step)
            torch.nn.utils.clip_grad_norm_(critic_model.parameters(), args.max_grad_norm)
            critic_optimizer.step()
            critic_optimizer.zero_grad()
    return ppo_step


def train(args, ori_model, actor_model, reward_model, critic_model, tokenizer, dataset, device, tb_write):
    """模型训练"""
    # 根据actor模型和critic模型构建actor优化器和critic优化器
    actor_optimizer = Adam(actor_model.parameters(), lr=args.learning_rate, eps=args.adam_epsilon)
    critic_optimizer = Adam(critic_model.parameters(), lr=args.learning_rate, eps=args.adam_epsilon)

    cnt_timesteps = 0
    ppo_step = 0
    experience_list = []
    mean_reward = []
    # 训练
    for i in range(args.num_episodes):
        for timestep in range(args.max_timesteps):
            cnt_timesteps += 1
            # 从数据集中随机抽取batch_size大小数据
            prompt_list = dataset.sample(args.batch_size)
            # 生成模型所需的input_ids
            input_ids = tokenizer.batch_encode_plus(prompt_list, return_tensors="pt",
                                                    max_length=args.max_len - args.query_len - 3,
                                                    truncation=True, padding=True)["input_ids"]
            input_ids = input_ids.to(device)
            generate_kwargs = {
                "min_length": -1,
                "max_length": input_ids.shape[1] + args.query_len,
                "top_p": args.top_p,
                "repetition_penalty": args.repetition_penalty,
                "do_sample": args.do_sample,
                "pad_token_id": tokenizer.pad_token_id,
                "eos_token_id": tokenizer.eos_token_id,
                "num_return_sequences": args.num_return_sequences}
            # 生成经验数据，并添加到经验池中
            experience = make_experience(args, actor_model, critic_model, ori_model, reward_model, input_ids,
                                         generate_kwargs)
            experience_list.append(experience)
            # 记录数据中的奖励值
            mean_reward.extend(experience["reward_score"].detach().cpu().numpy().tolist())

            # 当到达更新步数，进行模型更新
            if (cnt_timesteps % args.update_timesteps == 0) and (cnt_timesteps != 0):
                # 打印并记录平均奖励值
                mr = np.mean(np.array(mean_reward))
                tb_write.add_scalar("mean_reward", mr, cnt_timesteps)
                print("mean_reward", mr)
                actor_model.train()
                critic_model.train()
                # 模型更新
                ppo_step = update_model(args, experience_list, actor_model, actor_optimizer, critic_model,
                                        critic_optimizer, tb_write, ppo_step)
                # 模型更新后，将经验池清空
                experience_list = []
                mean_reward = []
        # 模型保存
        actor_model.save_pretrained(os.path.join(args.output_dir, "checkpoint-{}".format(ppo_step)))
        tokenizer.save_pretrained(os.path.join(args.output_dir, "checkpoint-{}".format(ppo_step)))
        print("save model")


def set_args():
    """设置训练模型所需参数"""
    parser = argparse.ArgumentParser()
    parser.add_argument('--device', default='2', type=str, help='设置训练或测试时使用的显卡')
    parser.add_argument('--train_file_path', default='data/ppo_train.json',
                        type=str, help='训练数据集')
    parser.add_argument('--ori_model_path', default='sft_model/', type=str, help='SFT模型')
    parser.add_argument('--reward_model_path', default='rm_model/', type=str, help='奖励模型路径')
    parser.add_argument('--max_len', default=768, type=int, help='模型最大长度')
    parser.add_argument('--query_len', default=64, type=int, help='生成问题的最大长度')
    parser.add_argument('--batch_size', default=16, type=int, help='批次大小')
    parser.add_argument('--num_episodes', default=3, type=int, help='循环次数')
    parser.add_argument('--max_timesteps', default=80, type=int, help='单次训练最大步骤')
    parser.add_argument('--update_timesteps', default=20, type=int, help='模型更新步数')
    parser.add_argument('--kl_coef', default=0.02, type=float, help='kl散度概率')
    parser.add_argument('--ppo_epoch', default=2, type=int, help='强化学习训练轮数')
    parser.add_argument('--policy_clip_eps', default=0.2, type=float, help='策略裁剪')
    parser.add_argument('--value_clip_eps', default=0.2, type=float, help='值裁剪')
    parser.add_argument('--top_p', default=1.0, type=float, help='解码Top-p概率')
    parser.add_argument('--repetition_penalty', default=1.4, type=float, help='重复惩罚率')
    parser.add_argument('--do_sample', default=True, type=bool, help='随机解码')
    parser.add_argument('--num_return_sequences', default=1, type=int, help='生成内容个数')
    parser.add_argument('--max_grad_norm', default=1.0, type=float, help='')
    parser.add_argument('--reward_clip_eps', default=5.0, type=float, help='奖励值裁剪')
    parser.add_argument('--gamma', default=1.0, type=float, help='优势函数gamma值')
    parser.add_argument('--lam', default=0.95, type=float, help='优势函数lambda值')
    parser.add_argument('--learning_rate', default=1e-5, type=float, help='学习率')
    parser.add_argument('--adam_epsilon', default=1e-5, type=float, help='Adam优化器的epsilon值')
    parser.add_argument('--output_dir', default="output_dir", type=str, help='模型保存路径')
    parser.add_argument('--seed', default=2048, type=int, help='')
    return parser.parse_args()


def main():
    # 设置模型训练参数
    args = set_args()
    # 设置显卡信息
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["CUDA_VISIBLE_DEVICES"] = args.device
    # 获取device信息，用于模型训练
    device = torch.device("cuda" if torch.cuda.is_available() and int(args.device) >= 0 else "cpu")
    # 设置随机种子，方便模型复现
    if args.seed:
        torch.manual_seed(args.seed)
        random.seed(args.seed)
        np.random.seed(args.seed)
    if not os.path.exists(args.output_dir):
        os.mkdir(args.output_dir)
    tb_write = SummaryWriter()
    # 实例化原始模型、Actor模型、Reward模型和Critic模型
    ori_model = ActorModel(args.ori_model_path)
    ori_model.to(device)
    actor_model = ActorModel(args.ori_model_path)
    actor_model.to(device)

    reward_model = RewardModel.from_pretrained(args.reward_model_path)
    reward_model.to(device)

    critic_model = CriticModel.from_pretrained(args.reward_model_path)
    critic_model.to(device)

    # 实例化tokenizer
    tokenizer = BertTokenizer.from_pretrained(args.ori_model_path, padding_side='left')
    tokenizer.eos_token_id = tokenizer.sep_token_id
    # 加载训练数据
    dataset = ExamplesSampler(args.train_file_path)
    print("数据量：{}".format(dataset.__len__()))
    # 开始训练
    train(args, ori_model, actor_model, reward_model, critic_model, tokenizer, dataset, device, tb_write)


if __name__ == '__main__':
    main()

代码的主要功能和流程：

1. 导入所需的库和模块：这包括用于强化学习的工具函数，如get_advantages_and_returns，actor_loss_function和critic_loss_function，以及模型类，如ActorModel，RewardModel和CriticModel。

2. 定义函数make_experience：这个函数用于生成经验数据，它使用actor模型和critic模型来生成新的序列，并计算相关的奖励和优势。

3. 定义函数update_model：这个函数用于更新模型的参数。它使用PPO算法来更新actor模型和critic模型的参数。

4. 定义函数train：这个函数是训练循环的主体。它首先初始化模型和优化器，然后对每个训练步骤生成经验数据，并在适当的时候更新模型的参数。

5. 定义函数set_args：这个函数用于设置训练参数。

6. 定义函数main：这个函数是脚本的入口点。它首先设置训练参数，然后加载模型和数据，最后开始训练。

这个脚本的主要流程是：首先，它会加载模型和数据，然后开始训练循环。在每个训练步骤，它会生成新的经验数据，并在适当的时候更新模型的参数。训练完成后，它会保存模型的参数。

【PPO算法体现】

在这段代码中，PPO（Proximal Policy Optimization）算法主要体现在update_model函数中。这个函数负责更新模型的参数。

在update_model函数中，首先对经验池中的数据进行随机打乱，然后遍历每个经验数据，计算actor模型和critic模型的损失值，然后进行梯度回传和参数更新。

对于actor模型，损失函数是actor_loss_function，它计算的是策略的损失。这个损失函数中实现了PPO的核心思想，即限制新策略和旧策略之间的差异。

对于critic模型，损失函数是critic_loss_function，它计算的是值函数的损失。

【PPO算法步骤】

PPO（Proximal Policy Optimization）算法的实现步骤大致如下：

1. 初始化：初始化策略网络和价值网络。策略网络用于生成动作，价值网络用于估计每个状态的价值。

2. 采样：使用当前的策略网络进行多次采样，每次采样会生成一个轨迹，轨迹包含了状态、动作、奖励等信息。

3. 计算优势函数：对于每个轨迹中的每个状态，使用价值网络和奖励信息计算优势函数。优势函数表示采取某个动作比按照价值网络的预测采取动作好多少。

4. 更新策略网络：使用优势函数和PPO的目标函数来更新策略网络。PPO的目标函数会限制新策略和旧策略之间的差异，防止更新过程中出现大的跳跃。

5. 更新价值网络：使用轨迹中的奖励信息和价值网络的预测来更新价值网络。

6. 重复：重复上述步骤，直到满足停止条件，如达到最大迭代次数或策略网络的性能达到预设的阈值。

以上是PPO算法的基本步骤，实际的实现可能会有所不同，例如可能会添加一些额外的步骤来提高性能，如使用多个并行的环境进行采样，或使用更复杂的网络结构。

【策略网络】

策略网络（Policy Network）是强化学习中的一个重要概念，它是一个函数，用于在给定状态下生成动作。

在深度强化学习中，策略网络通常由神经网络来实现。这个神经网络的输入是状态，输出是在该状态下采取每个可能动作的概率。

策略网络有两种形式：

1. 确定性策略：在给定状态下，总是生成同一个动作。这种策略网络的输出是一个动作，而不是动作的概率。

2. 随机性策略：在给定状态下，生成每个可能动作的概率。这种策略网络的输出是一个概率分布。

在强化学习的训练过程中，我们通常会使用一种叫做策略梯度（Policy Gradient）的方法来更新策略网络，使得它能生成更好的动作。这个过程通常涉及到一种叫做优势函数（Advantage Function）的概念，它用于衡量在某个状态下，采取某个动作相比于按照当前策略采取动作的优势有多大。

【价值网络】

价值网络（Value Network）是强化学习中的一个重要概念，它是一个函数，用于估计在某个状态下，按照当前策略采取行动能获得的预期回报。

在深度强化学习中，价值网络通常由神经网络来实现。这个神经网络的输入是状态，输出是该状态的价值。

价值网络有两种形式：

1. 状态价值函数（V-function）：V(s)表示在状态s下，按照当前策略采取行动能获得的预期回报。

2. 动作价值函数（Q-function）：Q(s, a)表示在状态s下，采取动作a后能获得的预期回报。

在强化学习的训练过程中，我们通常会使用一种叫做值迭代（Value Iteration）的方法来更新价值网络，使得它能更准确地估计状态或动作的价值。这个过程通常涉及到一种叫做贝尔曼方程（Bellman Equation）的递归公式。

【优势函数】

势函数（Advantage Function）在强化学习中是一个非常重要的概念，它用于衡量在某个状态下，采取某个动作相比于按照当前策略采取动作的优势有多大。

优势函数A(s, a)的定义如下：

A(s, a) = Q(s, a) - V(s)

其中，Q(s, a)是动作价值函数，表示在状态s下采取动作a后能获得的预期回报；V(s)是状态价值函数，表示在状态s下按照当前策略π采取动作能获得的预期回报。

优势函数的值如果为正，表示采取动作a比按照当前策略的预期回报要高，反之则低。因此，优势函数可以用来指导策略的更新，即倾向于增大优势函数为正的动作的概率，减小优势函数为负的动作的概率。

【常用的优势函数】

在强化学习中，优势函数（Advantage Function）是一个重要的概念，它衡量在某个状态下，采取某个动作相比于按照当前策略采取动作的优势有多大。优势函数的计算方法有很多种，以下是一些常用的优势函数计算方法：

1. 基础优势函数：最基础的优势函数定义是 A(s, a) = Q(s, a) - V(s)，其中 Q(s, a) 是动作价值函数，V(s) 是状态价值函数。

2. 蒙特卡洛（Monte Carlo）估计：这种方法直接使用实际的回报减去状态价值函数来估计优势函数，即 A(s, a) = G_t - V(s)，其中 G_t 是从状态 s 开始的实际回报。

3. 时序差分（Temporal Difference）估计：这种方法使用一步的预期回报减去状态价值函数来估计优势函数，即 A(s, a) = r + γV(s') - V(s)，其中 r 是即时奖励，γ 是折扣因子，s' 是下一个状态。

4. 广义优势估计（Generalized Advantage Estimation，GAE）：这种方法是一种折衷的优势函数估计方法，它通过一个参数 λ 来控制蒙特卡洛估计和时序差分估计的权重，以达到方差和偏差之间的平衡。
PPO（Proximal Policy Optimization）算法通常使用的是广义优势估计（Generalized Advantage Estimation，GAE）作为优势函数。

GAE是一种平衡偏差和方差的优势函数估计方法。它通过引入一个衰减因子λ，结合了蒙特卡洛（Monte Carlo）方法和时序差分（Temporal Difference）方法的优点。λ的值在0和1之间，当λ为0时，GAE等同于时序差分估计；当λ为1时，GAE等同于蒙特卡洛估计。

GAE的计算公式如下：

A^GAE_t = δ_t + (γλ)δ_{t+1} + (γλ)^2 δ_{t+2} + ... + (γλ)^{T-t+1} δ_{T-1}

其中，δ_t = r_t + γV(s_{t+1}) - V(s_t)，r_t是即时奖励，γ是折扣因子，V(s)是状态价值函数。

在PPO算法中，使用GAE作为优势函数可以有效地减小估计的方差，提高学习的稳定性。

PPO算法使用示例：https://github.com/ericyangyu/PPO-for-Beginners

 注意gym和pyglet的版本，最新的无法运行！