【Agentic RL】Dressage Staleness 控制深度分析

【Agentic RL】Dressage Staleness 控制深度分析

0x00 概要

https://github.com/Accio-Lab/Dressage 是非常有特色的 Agent RL 系统。本文就学习 Dressage 的 Staleness 实现。

0x01. Staleness 是什么问题?

1.1 问题背景

异步强化学习中,Rollout(数据采集)与 Training(模型更新)是并行执行的。Trainer 一边用当前策略更新权重,一边让 Rollout 工作器用旧权重持续采样。当一个 Rollout 批次最终送入 Trainer 时,生产这批样本所用的权重可能已更新数轮——这就是 Staleness(陈旧性)。

时间轴:
  t1: 权重 v1 → Rollout 开始用 v1 采样
  t2: 权重 v2 ← Trainer 完成一次更新
  t3: 权重 v3 ← Trainer 又完成一次更新
  t4: Rollout 采样结束 → 样本送入 Trainer
      但样本基于 v1 生成,当前权重已是 v3
                                ↑
                         这就是 staleness gap

1.2 两个不同来源的 Staleness

Staleness 在异步链路中有两个不同来源,必须分开讨论:

来源 粒度 成因 后果
单条长轨迹内部 Trajectory 级 Partial rollout 下一条轨迹被反复中断续写,前后 token 来自不同权重版本 陈旧前缀决定状态分布,即使只在最新 token 上算 loss,样本仍从陈旧上下文续写
Batch 组装时 Batch 级 异步采样与训练并行,Rollout 持续产出,已完成的轨迹在队列中等待时权重已更新多次 Train-Inference Gap 扩大,梯度方向偏置,极端情况下训练发散

1.3 核心矛盾

异步带来吞吐,但吞吐带来了陈旧性。 吞吐越高(Rollout 与 Training 的重叠越多),Staleness 就越严重。Dressage 的目标是在不牺牲吞吐的前提下,在两个粒度上分别控制版本的偏移。


0x02. 总体设计

整体设计哲学:从"防"到"治"

Dressage 不是用一个单一的开关解决 Staleness,而是构建了多层次防御体系:

Dressage-防御体系

两个粒度,正交约束

Dressage 按两个粒度分别设计了控制逻辑:

Dressage-控制逻辑


0x03. Trajectory 粒度:约束单条轨迹的版本跨度

3.1 前提:Partial Rollout 在 Proxy 层的透明中断续写

这一级控制建立在 Partial Rollout 之上。我们先理解 Partial Rollout 的设计。

当一次长生成正遇上权重更新,整条轨迹作废太浪费,硬用旧权重跑到底又会混入过时输出。Partial Rollout 的做法是在 token 边界暂停、改用新权重续写。难点在于:黑盒、不可控的多步 Agent(如 OpenCode 这类跑着自有 harness 的外部进程),框架并不掌控其执行流程,插不进暂停点。

Dressage 把这件事下沉到 Proxy 层,前提是所有模型调用都经过 Proxy(Agent 从不直接连推理引擎),于是 Proxy 成了唯一的收口:

Dressage-proxy

关键设计决策:

  • Proxy 按 request-id 中止 SGLang 在途生成,接住已产出的 partial token
  • Agent 的 HTTP 调用原地挂起,在 Agent 看来它只是在等待模型返回,对后台的权重切换毫不知情
  • 新权重就绪后,Proxy 从"原有前缀 + 已生成 partial"续写,最后把多段拼成完整响应交还给 Agent
  • 整个过程对 Agent 完全透明——不依赖 Agent 自身配合,对任意白盒/黑盒 Agent、任意本地或远程 Sandbox 都成立

3.2 问题:反复中断带来 Trajectory 粒度的 Staleness

Partial Rollout 让长轨迹不被浪费,但也带来新问题:一条轨迹若被反复中断,就会横跨太多版本。更麻烦的是,它的前缀、工具调用历史与环境状态,很大程度上是由更旧的 policy 决定的。即便训练时只在最后一个版本的 token 上算 loss,样本仍是从一个陈旧前缀和陈旧状态分布上续写而来,信息量更低,还会给当前策略更新引入偏差和噪声。

这里实际上有两个层面在同时起作用,分工很明确:

  • 训练样本构造阶段管的是"这条轨迹里到底训哪些 token"——通过 _mask_nonlast_version_tokens() 只让最后一个真实版本产出的 token 参与 loss,其余屏蔽掉。这是过滤偏移的影响。
  • Rollout 阶段管的是"这条轨迹还能不能继续攒新版本"——直接限制它的版本跨度。这是从源头限制偏移的幅度。

本篇关注的是后者。只做 loss mask 不够,因为样本终究是从陈旧前缀续写的;必须让 rollout 侧在轨迹跨过太多版本时就停止。

3.3 控制逻辑:版本跨度上限

Dressage 为单条 Partial 轨迹设定一个版本跨度上限(max_partial_rollout_preempts),例如"最多切换 2 次版本"意味着一条轨迹最多横跨 3 个版本。

实现位于 dressage/proxy/server.py:125-168:

def _raise_if_partial_version_span_exceeded(
    *, session, candidate_versions, partial_rollout, max_partial_rollout_preempts
):
    if not partial_rollout or max_partial_rollout_preempts is None:
        return
    # 收集 session 历史 + 本次候选版本,去重后按观测顺序排列
    versions = _ordered_real_versions(
        [*_session_response_versions(session), *candidate_versions]
    )
    version_span = len(versions)
    version_switches = max(0, version_span - 1)
    if version_switches <= max_partial_rollout_preempts:
        return
    # 超限 → HTTP 502,当场终止
    raise HTTPException(status_code=502, detail={"error": "partial_rollout_staleness_exceeded", ...})

3.4 超限后的处理

超限的轨迹走统一失败重试路径:

  1. Proxy 返回 HTTP 502,error 名为 partial_rollout_staleness_exceeded
  2. 黑盒服务器(rollout_llm_proxy.py)识别该错误名 → 记录到 turn scope → 返回合成空响应给 Agent(Agent 不崩溃)
  3. Rollout 收集器检测到状态标记 → 将 Group 标记为失败 → 默认重试 DRESSAGE_ROLLOUT_MAX_RETRIES 次
  4. 重试耗尽后,整组丢弃
  5. Rollout 日志持续汇报版本跨度分布(staleness/version_span_mean/max/min)和拒绝比例(staleness/partial_rollout_rejected_groups),判断这道控制是否频繁触发

"软拒绝"机制的核心价值:Agent 永远不会因为陈旧性而崩溃,最多是当前这条轨迹被丢弃重试。


0x04. Batch 粒度:约束进入 batch 的版本窗口

4.1 对外接口

这一级控制只暴露一个参数:一个训练 Batch 中最多允许包含最近观察到的若干个模型版本(配置项 dressage_staleness_keep_versions)。

# examples/scripts/default/dressage_staleness.yaml
dressage_staleness_keep_versions: 0   # 0=关闭(默认), N>0=开启
取值 含义
0 关闭 batch 级控制(默认)
N>0 只保留最近 N 个观测到的版本。如 keep_versions=2 则 batch 中可同时包含当前最新版本与上一版本,更早版本的轨迹被丢弃

最核心的权衡:窗口越窄数据越新鲜,但被丢弃的样本也越多。

4.2 版本判定是轨迹级的

Dressage 在生成时记录了 token 级的版本信息(full_versions[]),单条轨迹内部完全可能横跨多个版本。但在做 batch 级判定时,它以整条轨迹最后一个"可训练 token"的版本作为该轨迹新鲜度的代表:

# staleness.py:65-84
def trajectory_version_infos(group):
    latest_by_key = {}
    for order, sample in enumerate(group):
        version = real_version(metadata.get("dressage_end_token_version"))
        key = trajectory_key(sample)
        # 只保留同一条轨迹中 segment_index 最大的 segment 的版本
        position = _segment_index(sample, order)
        if key not in latest_by_key or latest_by_key[key][0] <= position:
            latest_by_key[key] = (position, version)
    return [TrajectoryVersionInfo(key=k, version=v) for k, (_, v) in latest_by_key.items()]

这是一个有意为之的选择:即便一条轨迹早期的若干片段由旧版本生成,只要它最终收敛于一个足够新的版本,就不会因"开头偏旧"而被整条误判丢弃。

4.3 核心机制:推进窗口 + 回溯复查

Batch 级控制的逻辑由两个核心动作配合完成,实现位于 dressage/rollout/staleness.py:

Dressage-Batch 级控制

动作一:推进窗口。每当一个新完成的组带来一个此前未出现的新版本,StalenessTracker 将版本追加到观测列表,current_version_index 和 cutoff_version_index 随之前移:

# staleness.py:111-119
def observe_group(self, group):
    previous_count = len(self.versions)
    for info in trajectory_version_infos(group):
        if info.version not in self.versions:
            self.versions.append(info.version)
    return len(self.versions) != previous_count  # 是否推进了窗口

动作二:回溯复查。

Batch 收集流程采用了"先缓存、触发新版本时再重过滤"的惰性策略。这意味着:不会在每批新版本到来时立即丢弃陈旧 Group,而是等有新版本到来时批量重检——这是一个吞吐和时效性的权衡设计。

窗口前移后,此前已进入候选 batch、但现已落到窗口之外的组会被重新筛查并移除:

# fully_async_rollout.py:446-451
advanced_version = staleness_filter.observe_completed(completed_groups)
if advanced_version and data:
    data = staleness_filter.filter_pending(data, logger)

同时,每个新加入的组在正式纳入 batch 前也会再次对照窗口检查,只要组内存在任一版本过旧的轨迹,即整组丢弃:

# staleness.py:121-129
def should_drop_group(self, group):
    cutoff = self.cutoff_version_index
    if cutoff is None:
        return False
    return any(self.version_index(info.version) < cutoff
               for info in trajectory_version_infos(group))

4.4 两个关键设计选择

选择一:版本号按观测顺序排列

Dressage 的一个关键设计选择:版本号按"被观测到的顺序"排列,而不是按版本标签的字典序或数值序。

# staleness.py:117
if info.version not in self.versions:
    self.versions.append(info.version)  # 追加到末尾

为什么这么做?因为模型版本号通常是 Git commit hash、时间戳或递增 ID。如果按字典序排序,可能出现"v10 < v2"这种反直觉的结果。按观测顺序排列保证了版本窗口的正确性。

选择二:整 GRPO Group 丢弃

GRPO 的 advantage normalization 依赖完整的 Group 结构。如果只删除 Group 内部分 trajectory,Group 的 reward 均值和标准差会偏移,破坏 advantage 计算。所以丢弃粒度是整组。

# staleness.py:152-158
def keep_group(self, group_id: int, group: list[Any], logger: Any) -> bool:
    ...
    if not self.tracker.should_drop_group(group):
        return True
    self._drop_group(group_id, logger)
    return False

丢整个 Group 而不是 Group 内的部分 Sample。这是因为 GRPO 的 advantage normalization 依赖完整的 Group 结构,如果只删部分 trajectory,Group 统计量会偏移,引入新的训练偏置。

4.5 不完整元数据的容错

当 parent_traj_id 缺失或版本号为 unknown/-1/none/空字符串时,该样本既不推进版本追踪器,也不触发丢弃:

# staleness.py:44-50
def real_version(value):
    if value is None or str(value).strip().lower() in _NON_REAL_VERSIONS:
        return None
    return str(value).strip()

这确保了陈旧性控制只作用于可靠的版本信息上。


0x05. 两个粒度的关系与互补

维度 Trajectory 级控制 Batch 级控制
作用位置 Proxy 侧(生成时) Rollout 侧(收集时)
控制对象 单条轨迹内部 整批轨迹集合
触发时机 每次生成续写前检查 Batch 组装循环中
超限后果 502 → 整组重试 整组丢弃
核心参数 max_partial_rollout_preempts keep_versions
对应函数 _raise_if_partial_version_span_exceeded StalenessGroupFilter

二者可同时启用:trajectory 控制约束单条轨迹能积累多少版本偏移,batch 控制约束进入 trainer 的整体版本新鲜度。

此外,还有其它机制:

5.1 训练侧的 Token-Level Loss Mask

位置:dressage/rollout/artifacts/samples.py:197-213

当一条轨迹确实跨越了多个版本时,Dressage 在构造训练样本时只保留最后一个真实版本对应 token 的 loss,将非最终版本 token 的 loss mask 置为 0。这确保了梯度只来自最新策略版本的 token,不被陈旧 token 污染。这是"治"——过滤偏移的影响,与 rollout 阶段的"防"——从源头限制偏移幅度,形成互补。

Token 级别 Loss Mask 流程(训练样本构造时)

Token 级别 Loss Mask 流程

5.2 配套机制

  • Proxy 暂停/恢复机制(dressage/training/train_async_with_rollout_pause.py):在 Trainer 更新权重前,通过 HTTP API 暂停 Proxy 上的活跃 SGLang 生成请求,等待安全边界后再更新权重,更新完成后恢复生成。这是 Partial Rollout 和 epoch staleness 检查的基础设施。
  • 黑盒 Agent 的 Staleness 错误处理(blackbox_server/proxy/rollout_llm_proxy.py):当 Proxy 返回因陈旧拒绝(HTTP 502)时,黑盒服务器识别特定错误名并返回合成空响应给 Agent,避免 Agent 崩溃。
  • 版本号全链路记录:从 Proxy 到 Rollout 到 Sample 构建,链路上每个环节都记录 token_version,确保可追溯。

0x06. 目前缺点与优化方向

6.1 当前缺点

缺点 说明
N/M 是静态超参数 keep_versions 和 max_preempts 训练启动时固定,不能自适应不同阶段。训练早期需更大窗口,后期需更小
Group 级丢弃代价高 丢弃整组意味着丢掉所有 trajectory(含新鲜的),尤其当组内版本分布不均匀时浪费大
二元丢弃策略 要么全留、要么全扔。没有软权重(如版本越旧在 loss 中权重越低)的平滑过渡
版本号的语义仅止于标签 版本号不含实际策略差异信息(KL 散度、参数变化量)。两个版本差异极小也一样丢弃
黑盒 Agent 空操作浪费 合成空响应虽不 crash,但 Agent 会继续执行空操作直到被 Terminate,浪费计算
Partial Rollout 重试延迟 Staleness 拒绝后整组重试,高压力场景下可能反复失败,影响吞吐

6.2 优化方向

方向 具体建议
自适应版本窗口 根据 KL 散度或 logprob 变化速率动态调整 keep_versions。变化大时缩小,稳定时扩大
Trajectory 级选择性过滤 Group 内只标记陈旧 trajectory 为 no_grad 而非丢弃整组,需配套改造 GRPO advantage 计算
Soft Staleness Weighting 引入 0-1 连续权重,版本越旧的轨迹 loss 权重越低,实现平滑过渡
基于策略差异的版本管理 用参数更新量或 KL 散度度量版本差异,变化小时放宽限制
阶段化配置 将 keep_versions 和 max_preempts 设为 YAML 的阶段化参数,不同 rollout_id 区间不同取值
Partial Rollout 智能预判 Proxy 根据训练节奏预估"这次生成大概多久会被中断",预测性调整 chunk 大小

0x07. 架构图

7.1 整体架构

整体架构

7.2 Staleness 组件关系

Staleness 组件关系

7.3 数据流图(Staleness 相关)

批次版本过滤流程(Batch-Level Staleness):

数据流图

0x08. 总结

Dressage 对 Staleness 问题的处理体现了两个粒度、正交约束的设计思想:

  1. Trajectory 粒度:基于 Proxy 层透明的 Partial Rollout 机制,限制单条轨迹的版本跨度上限。在 rollout 侧当场终止超限轨迹,从源头限制偏移幅度。

  2. Batch 粒度:通过"推进窗口 + 回溯复查"两个核心动作,将进入 Trainer 的轨迹约束在最近几个版本内。版本判定以整条轨迹最后一个可训练 token 为准,丢弃粒度为整个 GRPO Group。

  3. 配套机制:Token 级 Loss Mask 过滤偏移影响,Proxy 暂停/恢复提供基础设施,黑盒合成空响应保证容错。

  4. 可观测:通过 wandb 暴露 staleness/* 指标(版本跨度分布、丢弃组数、版本 gap),让训练者实时监控 Staleness 状况。

两组实验(AlfWorld + DAPO/OpenCode)一致表明:控制住版本偏移后,reward 学习更快,rollout 与当前 policy 的 logprob 差距显著缩小。

当前主要局限在于静态配置和二元丢弃策略,未来方向包括自适应窗口、软权重衰减和基于策略差异的版本管理。


TransFormer-封面

0xFF 参考

Dressage 的 True Staleness Control

posted @ 2026-10-20 21:07  罗西的思考  阅读(11)  评论(0)    收藏  举报