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

0 阅读14分钟

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

0x00 概要

github.com/Accio-Lab/D… 是非常有特色的 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 按两个粒度分别设计了控制逻辑:

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

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 级控制

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_preemptskeep_versions
对应函数_raise_if_partial_version_span_exceededStalenessGroupFilter

二者可同时启用: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 流程

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 组件关系

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-封面

TransFormer-封面

0xFF 参考

Dressage 的 True Staleness Control

本文使用 markdown.com.cn 排版