PPO算法介绍

什么是PPO算法？

近端策略优化(Proximal Policy Optimization, PPO)是一种基于策略梯度的强化学习算法，由OpenAI在2017年提出。PPO算法在保持训练稳定性的同时，能够获得较好的样本效率和性能表现。PPO的核心思想是通过限制策略更新的幅度，避免过大的策略变化导致性能崩溃。

PPO算法有两种主要变体：PPO-Penalty和PPO-Clip。在实际应用中，PPO-Clip因其实现简单且性能优越而被广泛采用。

PPO算法原理

1. 策略目标函数

传统的策略梯度方法中，在强化学习中，目标函数 $J(\theta)$ 表示策略 $\pi_\theta$ 的期望累积回报。具体形式如下：

这里：

• $\tau = (s_0, a_0, s_1, a_1, \dots, s_T)$ 是轨迹（trajectory），由状态和动作序列组成。
• $\pi_\theta(a_t | s_t)$ 是策略，表示在状态 $s_t$ 下选择动作 $a_t$ 的概率分布。
• $R(s_t, a_t)$ 是即时奖励（reward）。
• $J(\theta)$ 是策略的性能指标。
• $\mathbb{E}_{\tau \sim \pi_\theta}$ 表示对所有可能轨迹的期望。
• $\sum_{t=0}^{T} R(s_t, a_t)$ 是轨迹的累积奖励。

虽然这个函数可以通过梯度上升进行优化，但存在两个关键问题：

• 步长（学习率）难以确定：太小导致训练过慢，太大导致训练不稳定
• 样本效率低：每个样本只能使用一次

2. 近端策略优化 - PPO-Clip

PPO-Clip通过引入一个新的目标函数解决了这些问题，称为"裁剪代理目标函数"(Clipped Surrogate Objective)：

其中：

• $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$ 是重要性采样比率
• $A_t$ 是优势函数估计
• $\epsilon$ 是一个超参数，用于控制新策略与旧策略的差异范围（通常设为0.1或0.2）

这个目标函数的关键特点是：它通过剪裁比率$r_t(\theta)$，限制了新旧策略之间的差异。这确保了策略更新不会太大，从而提高了训练的稳定性。

3. 情景分析

为了更好地理解PPO-Clip的工作原理，我们分析几种不同情况：

当优势$A_t > 0$时：

• 如果$r_t(\theta) < 1+\epsilon$，目标函数为$r_t(\theta)A_t$，鼓励增加这个动作的概率
• 如果$r_t(\theta) > 1+\epsilon$，目标函数为$(1+\epsilon)A_t$，不再鼓励增加概率

当优势$A_t < 0$时：

• 如果$r_t(\theta) > 1-\epsilon$，目标函数为$r_t(\theta)A_t$，鼓励减少这个动作的概率
• 如果$r_t(\theta) < 1-\epsilon$，目标函数为$(1-\epsilon)A_t$，不再鼓励减少概率

通过这种方式，PPO算法实现了策略的保守更新，避免了性能的剧烈波动。

PPO在项目中的实现

在LLaMA-Factory项目中，PPO算法被实现为CustomPPOTrainer类，继承自trl库的PPOTrainer和Trainer类。下面是算法实现的关键部分：

1. PPO配置

python
ppo_config = PPOConfig(
    model_name=model_args.model_name_or_path,
    learning_rate=training_args.learning_rate,
    mini_batch_size=training_args.per_device_train_batch_size,
    batch_size=backward_batch_size * finetuning_args.ppo_buffer_size,
    gradient_accumulation_steps=training_args.gradient_accumulation_steps,
    ppo_epochs=finetuning_args.ppo_epochs,
    max_grad_norm=training_args.max_grad_norm,
    seed=training_args.seed,
    optimize_device_cache=True,
    target=finetuning_args.ppo_target,
    use_score_scaling=finetuning_args.ppo_score_norm,
    use_score_norm=finetuning_args.ppo_score_norm,
    whiten_rewards=finetuning_args.ppo_whiten_rewards,
    accelerator_kwargs={"step_scheduler_with_optimizer": False},
    log_with=training_args.report_to[0] if training_args.report_to else None,
    project_kwargs={"logging_dir": training_args.logging_dir},
)

2. PPO训练循环

python
def ppo_train(self, resume_from_checkpoint: Optional[str] = None) -> None:
    """Implement training loop for the PPO stage, like _inner_training_loop() in Huggingface's Trainer."""
    
    # ... 准备阶段省略 ...
    
    for step in tqdm(range(max_steps), disable=not self.is_local_process_zero()):
        try:
            batch = next(dataiter)
        except StopIteration:
            dataiter = iter(self.dataloader)
            batch = next(dataiter)

        # 获取输入
        self.model.eval()
        queries, responses, rewards = [], [], []
        for idx in range(0, self.config.batch_size, self.config.mini_batch_size):
            mini_batch = {
                "input_ids": batch["input_ids"][idx : idx + self.config.mini_batch_size],
                "attention_mask": batch["attention_mask"][idx : idx + self.config.mini_batch_size],
            }
            mini_batch_queries, mini_batch_responses = self.get_inputs(mini_batch)
            mini_batch_rewards = self.get_rewards(mini_batch_queries, mini_batch_responses)
            queries.extend(mini_batch_queries)
            responses.extend(mini_batch_responses)
            rewards.extend(mini_batch_rewards)

        # 执行PPO步骤
        self.model.train()
        stats = self.step(queries, responses, rewards)
        loss_meter.update(float(stats["ppo/loss/total"]), n=len(rewards))
        reward_meter.update(torch.stack(rewards).mean().item(), n=len(rewards))
        
        # ... 后续步骤省略 ...

3. 获取奖励

python
@torch.no_grad()
def get_rewards(self, queries: list["torch.Tensor"], responses: list["torch.Tensor"]) -> list["torch.Tensor"]:
    """Compute scores using given reward model."""
    if self.finetuning_args.reward_model_type == "api":
        token_ids = [torch.cat((q, r), dim=-1).tolist() for q, r in zip(queries, responses)]
        messages = self.tokenizer.batch_decode(token_ids, skip_special_tokens=False)
        return get_rewards_from_server(self.reward_model, messages)

    batch: dict[str, torch.Tensor] = self.prepare_model_inputs(queries, responses)
    unwrapped_model: AutoModelForCausalLMWithValueHead = self.accelerator.unwrap_model(self.model)

    if self.finetuning_args.reward_model_type == "lora":
        replace_model(unwrapped_model, target="reward")
        reward_model = self.model
    else:
        reward_model = self.reward_model

    with unwrap_model_for_generation(reward_model, self.accelerator), self.amp_context:
        values: torch.Tensor = reward_model(**batch, return_dict=True, use_cache=False)[-1]

    if self.finetuning_args.reward_model_type == "lora":
        replace_model(unwrapped_model, target="default")

    rewards = values.gather(dim=-1, index=(batch["attention_mask"].sum(dim=-1, keepdim=True) - 1))
    return rewards.float().detach()

PPO算法的数学公式总结

1. 重要性采样比率:

   $r_t(\theta) = \frac{\pi_\theta(a_t|s_t)}{\pi_{\theta_{old}}(a_t|s_t)}$

2. 裁剪代理目标函数:

   $L^{CLIP}(\theta) = \mathbb{E}_t \left[ \min(r_t(\theta) A_t, \text{clip}(r_t(\theta), 1-\epsilon, 1+\epsilon) A_t) \right]$

3. 完整的PPO目标函数 (包含值函数损失和熵正则化):

   $L^{TOTAL}(\theta) = \mathbb{E}_t \left[ L^{CLIP}(\theta) - c_1 L^{VF}(\theta) + c_2 S[\pi_\theta](s_t) \right]$

   其中:

   • $L^{VF}$ 是值函数损失: $(V_\theta(s_t) - V_t^{target})^2$
   • $S$ 是策略的熵
   • $c_1$ 和 $c_2$ 是系数

4. 广义优势估计 (GAE):

   $A_t = \delta_t + (1-\lambda)\delta_{t+1} + ... + (1-\lambda)^{T-t-1}\delta_{T-1}$

   其中 $\delta_t = r_t + \gamma V(s_{t+1}) - V(s_t)$

项目中PPO的实现代码

下面是LLaMA-Factory项目中PPO算法的核心实现部分：

python
class CustomPPOTrainer(PPOTrainer, Trainer):
    """继承自PPOTrainer的自定义实现"""

    def __init__(
        self,
        model_args: "ModelArguments",
        training_args: "Seq2SeqTrainingArguments",
        finetuning_args: "FinetuningArguments",
        generating_args: "GeneratingArguments",
        callbacks: Optional[list["TrainerCallback"]],
        model: "AutoModelForCausalLMWithValueHead",
        reward_model: Optional["AutoModelForCausalLMWithValueHead"],
        ref_model: Optional["AutoModelForCausalLMWithValueHead"],
        tokenizer: "PreTrainedTokenizer",
        processor: Optional["ProcessorMixin"],
        data_collator: "DataCollatorWithPadding",
        train_dataset: Optional["Dataset"] = None,
        eval_dataset: Optional["Dataset"] = None,
    ) -> None:
        # ... 初始化阶段省略 ...

    def ppo_train(self, resume_from_checkpoint: Optional[str] = None) -> None:
        """PPO训练循环的实现"""
        # ... 训练循环实现省略 ...

    @torch.no_grad()
    def get_inputs(self, batch: dict[str, "torch.Tensor"]) -> tuple[list["torch.Tensor"], list["torch.Tensor"]]:
        """获取输入和输出，用于PPO训练"""
        # ... 实现省略 ...

    @torch.no_grad()
    def get_rewards(self, queries: list["torch.Tensor"], responses: list["torch.Tensor"]) -> list["torch.Tensor"]:
        """计算奖励"""
        # ... 实现省略 ...

    @override
    @PPODecorators.empty_device_cache()
    def batched_forward_pass(
        self,
        model: "AutoModelForCausalLMWithValueHead",
        queries: "torch.Tensor",
        responses: "torch.Tensor",
        model_inputs: dict[str, Any],
        return_logits: bool = False,
        response_masks: Optional["torch.Tensor"] = None,
    ) -> tuple["torch.Tensor", Optional["torch.Tensor"], "torch.Tensor", "torch.Tensor"]:
        """批量前向传播"""
        # ... 实现省略 ...

PPO算法的优势

1. 稳定性更好：通过限制策略更新的幅度，避免了大的策略变化导致的性能崩溃。

2. 实现简单：相比于TRPO等其他信任区域方法，PPO算法实现更加简单，只需要一阶优化方法。

3. 样本效率高：通过重要性采样和多次迭代，PPO可以多次利用同一批样本，提高了样本效率。

4. 性能优越：在多种强化学习任务中，PPO展示了与SOTA算法相当甚至更好的性能。

5. 适用性广：PPO适用于连续动作空间和离散动作空间，能够应用于各种不同类型的强化学习问题。

结论

PPO是一种强大而实用的强化学习算法，通过限制新旧策略之间的差异，实现了训练过程的稳定性和高效性。在LLaMA-Factory项目中，PPO被实现为CustomPPOTrainer类，用于训练大型语言模型。该实现包括PPO配置、训练循环、奖励计算等关键组件，提供了一个完整的PPO训练框架。

通过近端策略优化，模型可以在保证稳定性的同时，有效地从经验中学习改进策略，最终达到更好的性能表现。