跳到主要内容

数据截至 (上游 commit 20daedc47558)

StarPO 主循环：rollout 与 update

本章讲什么：StarPO（State-Thinking-Action-Reward Policy Optimization）是 RAGEN 的核心算法。看它怎么 把一整条多轮轨迹当成一个 RL 样本来优化。

4.1 一句话：StarPO 是什么​

要解决的小问题。 普通 RLHF 只优化“一问一答”。但 agent 任务是多轮的：推理→动作→看新状态→再推理… 怎么把这整条“推理+动作交错”的轨迹当一个整体来优化？

思路。 把交互建模为 MDP：状态和动作都是 token 序列，每轮输出 <think>…</think><answer>…</answer>。 StarPO 用重要性采样优化整条轨迹，可插 PPO（token 级优势 + value function）或 GRPO（整条轨迹拿归一化奖励）。

两个交替阶段。 全在 RayAgentTrainer.fit（agent_trainer.py:942）里：

for step in total_training_steps:
    rollout 阶段：agent_proxy.rollout()  → 一批轨迹（DataProto）
    update 阶段：过滤 → 算 log_prob → 算优势 → update_actor

4.2 rollout 阶段：一批轨迹怎么采出来​

核心是 LLMAgentProxy.rollout（agent_proxy.py:225）。它是个“全部环境齐步走 max_turn 轮”的循环：

# 示意，非源码（改编自 agent_proxy.py:LLMAgentProxy.rollout）
env_outputs = es_manager.reset()          # 初始状态
for i in range(max_turn):
    if len(env_outputs) == 0: break       # 全做完了就停
    lm_inputs = ctx_manager.get_lm_inputs(env_outputs, prepare_for_update=False)
    lm_inputs.meta_info["mode"] = mode    # multiturn-start/middle/end
    lm_outputs = self.generate_sequences(lm_inputs)   # 调 LLM
    env_inputs = ctx_manager.get_env_inputs(lm_outputs)  # 解析动作
    env_outputs = es_manager.step(env_inputs)            # 执行
rollout_states = es_manager.get_rollout_states()
rollouts = ctx_manager.formulate_rollouts(rollout_states)  # 打包成训练样本

两个细节值得看。

  • mode 标记轮次位置。 首轮=multiturn-start，末轮=multiturn-end（agent_proxy.py:252）。这个 mode 会传给下游 vLLM（推测用来控制 KV cache / early stop）。
  • “finalize”收尾。 若多轮还没到 end 就全部做完，rollout 会带 skip_generation=True 补一次 multiturn-end （agent_proxy.py:286）——让轨迹有个干净的结尾状态。

实例级熵顺手累加。 rollout 过程中把每轮的熵、token 数、轮数按 env_id 累加，最后除得到“每轨迹平均熵” （agent_proxy.py:298）。这是诊断指标之一。

4.3 update 阶段：从轨迹到梯度​

回到 fit（agent_trainer.py:1046 起的主循环）。一个 step 内，rollout 之后依次是：

1. （V2）collapse_detector.compute_collapse_metrics()   # 过滤前算，公平对比 → 第4章
2. rollout_filter.filter()                              # 按奖励方差筛组   → 第3章
3. adjust_batch()                                       # 补/删样本凑整除
4. actor_rollout_wg.compute_log_prob()                  # 重算 old_log_prob + 熵
5. compute_advantage()                                  # GRPO / PPO 算优势
6. update_critic() / update_actor()                     # veRL 做梯度下降

为什么要 adjust_batch。 过滤后样本数不规整，但 mini-batch / GPU 数 / 组数都要整除。adjust_batch （agent_trainer.py:62）按这三者的最小公倍数，用 copy（复制样本）或 delete（删样本）凑整 （agent_trainer.py:1194）。

关键：uid 就是 group_id。 看 agent_trainer.py:1276：batch.non_tensor_batch["uid"] = group_ids。这一行 决定了 GRPO “按谁分组算基线”——同一个 group_id（同一初始状态）的几条轨迹互为参照。

response_mask 就用 loss_mask。 agent_trainer.py:1279 直接把 loss_mask 赋给 response_mask——因为多轮 场景下“该算损失的 token”就是 loss_mask 决定的。

4.4 GRPO 优势：组内减均值，加 episode 去重​

GRPO 的核心思想：一个 prompt 采 N 条轨迹，用这 N 条的组内均值当基线，谁高于均值谁就是正优势。 RAGEN 在 ragen/trainer/core_algos.py:26 重写了 compute_grpo_outcome_advantage，加了一个关键的 episode 去重：

# 示意，非源码（改编自 core_algos.py:compute_grpo_outcome_advantage）
seen_pairs = set()
for i in range(bsz):
    if episode_ids is not None:
        pair = (index[i], episode_ids[i])   # (group_id, episode)
        if pair in seen_pairs: continue     # 同 episode 只计一次
        seen_pairs.add(pair)
    id2score[index[i]].append(scores[i])
# 组内减均值（可选除以方差）
scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + eps)

为什么要去重。 在 single_turn / limited_multi_turn 模式下，一条 episode 被拆成了多个样本（每轮一个）。 如果不去重，长轨迹会因为轮数多而在组内均值里占更大权重，偏差。去重让每条 episode 只贡献一次均值/方差。

norm_adv_by_std_in_grpo 决定是否除以方差（除了就是常规 GRPO，不除就是 Dr.GRPO 风格，core_algos.py:78）。 单样本组（只一条）特判：mean=0, std=1（core_algos.py:68），避免除零。

4.5 PPO / GAE 路径​

config/base.yaml:97 默认 adv_estimator: gae（PPO）。compute_advantage（agent_trainer.py:140）按 adv_estimator 分派到 veRL 的 GAE / GRPO / REINFORCE++ / RLOO / REMAX 等。

bi-level GAE（两层 MDP）。 RAGEN 自带一个 compute_bi_level_gae_advantage_return（core_algos.py:94）：高层是 “每轮一个 step”的 agentic MDP（用 high_level_gamma），低层是“每个 token”的 token MDP（用 gamma）。它先在 每轮的 eos 位置算高层优势/回报，再把该回报当“该轮的序列奖励”向 token 级下传。注意 train.py:162 断言： 开 bi_level_gae 时 use_turn_scores 必须为 False（轮分数路径还没接好, inferred from assertion comment）。

4.6 三种“提前收手”（early stopping）​

RAGEN 把“训练已经崩了”当作一类判据，主动停训（这是“诊断 agent RL”这个主题的体现）：

触发条件逻辑位置
过滤后连续空 batch连续 N 步过滤出 0 样本agent_trainer.py:1143
奖励方差崩塌连续 10 步方差 < 0.1×基线方差agent_trainer.py:1228
验证集成功率过低连续 5 步 success < 1%agent_trainer.py:1490

其中“基线方差”取前 10 个成功 step 的 rollout/in_group_reward_std 均值（agent_trainer.py:1217）。

4.7 边界与坑​

  • 奖励管理器是“dummy”的。 train.py:291 固定用 DummyRewardManager：奖励早在环境里算好了（non_tensor_batch['reward']）， 这里只是把它搬到最后一个 token（train.py:68）。naive/prime 等 RM 路径被注释掉了。
  • 奖励归一化责任有重叠。 奖励在 ctx_manager 已组内归一化一次，进 GRPO 时 uid=group_id 又会再走一次组内归一化 （agent_trainer.py:1271 附近注释偏乱）。
  • REMAX 未测试。 agent_trainer.py:1246 明确 exit()，暂不支持。

4.8 代码地图​

主题文件符号
主循环ragen/trainer/agent_trainer.pyRayAgentTrainer.fit
一批轨迹采样ragen/llm_agent/agent_proxy.pyLLMAgentProxy.rollout
batch 凑整ragen/trainer/agent_trainer.pyadjust_batch
优势分派ragen/trainer/agent_trainer.pycompute_advantage
GRPO + episode 去重ragen/trainer/core_algos.pycompute_grpo_outcome_advantage
bi-level GAEragen/trainer/core_algos.pycompute_bi_level_gae_advantage_return
early stoppingragen/trainer/agent_trainer.pyfit（consecutive_variances / consecutive_low_success）
配置校验train.pyadd_dependency_and_validate_config