rllm中的训练流程

1. 三种优势计算方法

在rllm框架中，实现了三种优势方法，分别是轨迹级别的优势（只针对最后的结果奖励计算优势）；广播模式下的逐步优势（计算结果奖励后，将结果奖励广播到每一步中）；单步模式下的逐步优势（计算结果奖励后，使用蒙特卡洛方法估计每一步的奖励）
下面，我使用一个例子来分别说明一下这几种优势的具体计算方法。

假设用户输入一个问题，Agent调用Python工具进行解决，同时，对于这个问题，产生两条轨迹：

• 轨迹 1（正确）：
• 第 1 轮：用户输入问题，LLM进行思考，并调用 Python 工具
• 第 2 轮：执行 Python 工具
• 第 3 轮：根据Python工具的执行结果，输出答案，返回 "4" → 最终奖励 1.0 ✅
• 轨迹 2（错误）：
• 第 1 轮：用户输入问题，LLM进行思考，并调用 Python 工具
• 第 2 轮：执行 Python 工具
• 第 3 轮：根据Python工具的执行结果，输出答案，返回 "5" → 最终奖励 0.0 ❌

方法 1：轨迹级优势（Trajectory-Level）

根据结果奖励（轨迹1的奖励为1.0，轨迹2的奖励为0），计算两条轨迹的优势（假设轨迹1的优势为+0.5，轨迹2的优势为-0.5），然后使用GRPO进行反向传播。

方法 2：逐步优势 - 广播模式（Stepwise Broadcast）

根据结果奖励计算两条轨迹的优势（和方法1相同），然后将优势广播到每一条轨迹的每一步中，这样一来，就产生了六个样本：

轨迹 1 中每一步的优势：
	- 第 1 轮：+0.5
	- 第 2 轮：+0.5
	- 第 3 轮：+0.5
轨迹 2 中每一步的优势：
	- 第 1 轮：-0.5
	- 第 2 轮：-0.5
	- 第 3 轮：-0.5

然后进行GRPO优化

方法 3：逐步优势 - 单步模式（Stepwise Per-Step）

根据结果奖励，使用蒙特卡罗方法计算每一步的奖励

轨迹 1 中每一步的奖励：
	第 3 轮：1.0
	第 2 轮：0.95 × 1.0 = 0.95
	第 1 轮：0.95 × 0.95 = 0.9025
轨迹 2 中每一步的奖励：
	第 3 轮：0.0
	第 2 轮：0.0
	第 1 轮：0.0

然后根据每一步的奖励，为每一步单独计算优势。
对于这种优势计算的方法来说，同样也有六条样本。

三种方法的比较

从方法1到方法3，逐渐变得更加复杂，对于中间步骤的奖励分配越来越精细。但是同样的，如果奖励定义不好，噪声会更大；并且计算起来更加困难。

2. 训练器代码讲解

定义在rllm/trainer/verl/agent_ppo_trainer.py

class AgentPPOTrainer(RayPPOTrainer):

    def fit_agent(self):
        # 一次训练循环
        for epoch in range(self.config.trainer.total_epochs):
            for batch_dict in self.train_dataloader:
                batch: DataProto = DataProto.from_single_dict(batch_dict)
                
                # ========== 1. 前向传播：生成轨迹 ==========
                # 会调用AgentExecutionEngine的run_agent_trajectory方法（上一篇文章提到过，是Agent推理的主要方法）
                
                # 如果使用逐步奖励，那么会输出所有步骤
                if self.config.rllm.stepwise_advantage.enable:
                    final_gen_batch_output = self.generate_agent_steps(
                        timing_raw=timing_raw, 
                        meta_info=batch.meta_info, 
                        uids=batch.non_tensor_batch["uid"]
                    )
                # 如果只使用结果奖励，那么只输出最终轨迹
                else:
                    final_gen_batch_output, generate_metrics = self.generate_agent_trajectory(
                        timing_raw=timing_raw, 
                        meta_info=batch.meta_info
                    )
                
                # ========== 2. 计算奖励 ==========
                # 2.1 使用奖励模型打分
                if self.use_rm:
                    reward_tensor = self.rm_wg.compute_rm_score(batch)
                    batch = batch.union(reward_tensor)
                
                # 2.2 使用奖励函数打分
                if "token_level_scores" not in batch.batch:
                    reward_tensor = self.reward_fn(batch)
                    batch.batch["token_level_scores"] = reward_tensor
                # 2.3  环境已经提供了奖励（从 trajectory.reward 转换而来) 
                # rllm框架默认会走这一条
                else:
                    reward_tensor = batch.batch["token_level_scores"]
                
                # ========== 3. 拒绝采样（可选）==========
                # 过滤掉一个组内全部答对或者全部答错的样本
                # （无论是逐步奖励还是结果奖励都是如此）
                # 【注意】：对于 stepwise_advantage，也需要只看最后一步的奖励
                if self.config.rllm.rejection_sample.enable:
                    ...
                
                # ========== 4. 准备优势计算的奖励 ==========
                # 默认：使用结果奖励
                batch.batch["token_level_rewards"] = batch.batch["token_level_scores"]
				         
				        # 如果是 per_step 方法，使用蒙特卡洛回报估计每一步的奖励
                if self.config.rllm.stepwise_advantage.enable: 
                    if self.config.rllm.stepwise_advantage.mode == "per_step":
                        batch.batch["token_level_rewards"] = batch.batch["mc_returns"]
            
                
                # ========== 5. 计算优势 ==========
                batch = compute_advantage(...)
                
                # ========== 6. 广播优势（broadcast 模式）==========
                if self.config.rllm.stepwise_advantage.enable and \
                   self.config.rllm.stepwise_advantage.mode == "broadcast":
                    # 将最后一步的优势广播到所有中间步骤
                    self._stepwise_advantage_broadcast(batch, other_step_batch=other_step_batch)
          
                
                # ========== 7. 模型更新 ==========

                actor_output = self.actor_rollout_wg.update_actor(batch)

其实，我们可以发现，当AgentExecutionEngine中的run_agent_trajectory方法跑通后，整体的训练代码就可以跑通。

因此，我们在自定义环境，Agent，工具，奖励时，只需要把推理脚本调试通，训练脚本自然就可以跑通，不需要在训练脚本上进行调试，这大大减少了debug的成本