面向长时程工具使用智能体任务的高效强化学习

# 面向长程工具使用智能体任务的高效强化学习
Source: https://arxiv.org/html/2608.10357

###### 摘要

长程工具使用智能体必须对用户目标、领域策略、工具调用、模拟器状态以及延迟的可验证奖励进行推理。强化学习（RL）天然适合这一场景，但多轮在线策略展开会产生长上下文，而特定于模型的注意力层可能需要自定义掩码和习得的汇点归一化。我们提出 SinkFlex-RL，一种面向双控制工具使用环境的模块化训练系统。该系统结合了 Gymnasium 兼容的环境封装器、VERL 风格的展开数据流、无需独立价值模型的组相对策略优化，以及一条汇点感知的 FlexAttention 路径，旨在因果掩码和滑动窗口掩码下保持特定于模型的汇点缩放。在初步的 τ²-Bench 零售实验中，验证奖励（mean@1）从训练早期的 0.25 上升到观察窗口后期的 0.44，同时训练得分和轨迹奖励代理指标也呈上升趋势。在固定配置的内存基准测试中，优化后的注意力路径在 4096 token 时将峰值 VRAM 从 28.06GB 降至 22.52GB，降低 19.7%，并在 8192 token 配置下以 25.53GB 运行（而 eager 基线已内存不足）。这些结果说明了将环境接口、RL 数据流和注意力内核设计集成在一起，对于内存可行的长程智能体训练的价值。

## 1 引言

语言模型智能体越来越多地运行在交互式环境中：它们与用户对话，通过工具检查和更新状态，遵循领域策略，并且只有在长轨迹之后才能获得稀疏的成功信号。诸如 τ-Bench 和 τ²-Bench 之类的基准通过在带有工具 API、策略约束和模拟用户的客户服务类领域中评估智能体，将这一场景形式化（[Yao et al., 2024](https://arxiv.org/html/2608.10357#bib.bib10)；[Barres et al., 2025](https://arxiv.org/html/2608.10357#bib.bib1)）。这些场景同时考验推理质量和训练系统能力。单条轨迹可能包含许多对话轮次、工具调用和环境观测；在线策略 RL 通过要求新的展开、奖励检查以及重复的策略更新，使这一成本成倍增加。

本文研究内存可行的智能体 RL 的系统层面问题，针对一个大型开放权重混合专家（MoE）Transformer。MoE 架构通过将 token 路由到部分专家，可以以较低的每 token 前馈成本增加模型容量（[Fedus et al., 2022](https://arxiv.org/html/2608.10357#bib.bib4)），但并未消除长上下文注意力瓶颈。随着多轮轨迹的增长，eager 注意力可能在反向传播之前就耗尽高带宽内存。我们目前的系统评估覆盖到 8192 token 的序列长度；更长的目标负载是本设计的动机，但不在本次评估范围内。现有的融合注意力内核改善了内存移动和吞吐量（[Dao et al., 2022](https://arxiv.org/html/2608.10357#bib.bib2)；[Dao, 2024](https://arxiv.org/html/2608.10357#bib.bib3)），但生产模型可能需要具有习得汇点参数、异构掩码或固定内核接口未暴露的反向行为的注意力变体。

因此，我们将智能体 RL 视为一个系统集成问题，并做出三项贡献。首先，我们将双控制环境封装在 Gymnasium 风格接口之后，使展开、工具、用户模拟器和奖励检查器能够连接到标准 RL 数据流。其次，我们使用组相对策略优化（GRPO）从轨迹组中更新 actor，而无需训练单独的价值模型，遵循 [DeepSeekMath](https://arxiv.org/html/2608.10357#bib.bib7)（[Shao et al., 2024](https://arxiv.org/html/2608.10357#bib.bib7)）引入的基于内存考虑的动机，以及 VERL 风格的后训练数据流（[Sheng et al., 2024](https://arxiv.org/html/2608.10357#bib.bib15)）。第三，我们实现了一条汇点感知的 FlexAttention 路径，该路径将因果掩码和滑动窗口掩码与可微的汇点缩放组合在一起，同时避免 eager 注意力状态的物化。我们通过初步的 τ²-Bench 零售训练运行和一项覆盖到 8192 token 的峰值内存研究来评估集成后的流水线。本工作的核心贡献在于这些组件的集成以及由此带来的内存可行性。

## 2 相关工作

智能体基准。将语言模型评估为交互式智能体，需要模型能够发出 API 调用、管理变化的状态并处理多轮反馈的环境。AgentBench（[Liu et al., 2024](https://arxiv.org/html/2608.10357#bib.bib11)）推动了广泛的交互式评估，而 WebArena（[Zhou et al., 2024](https://arxiv.org/html/2608.10357#bib.bib12)）、GAIA（[Mialon et al., 2024](https://arxiv.org/html/2608.10357#bib.bib13)）和 SWE-bench（[Jimenez et al., 2024](https://arxiv.org/html/2608.10357#bib.bib14)）分别针对 Web 导航、通用助手能力和软件工程。对于有状态的工具使用，τ-Bench 评估的是可执行的状态变化，而非仅基于文本的判断（[Yao et al., 2024](https://arxiv.org/html/2608.10357#bib.bib10)）。我们使用 τ²-Bench（[Barres et al., 2025](https://arxiv.org/html/2608.10357#bib.bib1)），其双控制设置允许智能体和模拟用户都影响任务进度。这为收集缺失信息、从无效动作中恢复以及在多轮交互中保持一致性的能力提供了一个有代表性的测试平台。

多步工具使用的 RL 微调。在复杂的工具使用环境中，标注每个中间动作的成本很高，这促使使用程序化可验证的回合级奖励进行训练（[Zheng et al., 2025](https://arxiv.org/html/2608.10357#bib.bib16)；[Dong et al., 2026](https://arxiv.org/html/2608.10357#bib.bib17)）。PPO（[Schulman et al., 2017](https://arxiv.org/html/2608.10357#bib.bib6)）通常使用一个习得的价值函数，这为大型语言模型增加了内存和计算开销。GRPO（[Shao et al., 2024](https://arxiv.org/html/2608.10357#bib.bib7)）去除了独立的价值模型，并在样本输出组内对奖励进行归一化。诸如 ver](https://arxiv.org/html/2608.10357#bib.bib15)（[Sheng et al., 2024](https://arxiv.org/html/2608.10357#bib.bib15)）之类的框架将 rollout 生成、奖励计算和策略更新组织为分布式后训练数据流。我们的工作采用了这些算法和系统组件，并重点将它们与长上下文、模型兼容的注意力集成起来，用于多轮工具使用展开。

长上下文训练的注意力机制。Transformer 使用注意力作为其核心的序列混合操作（[Vaswani et al., 2017](https://arxiv.org/html/2608.10357#bib.bib8)），但密集注意力在序列长度上具有二次方的内存和计算复杂度。FlashAttention 和 FlashAttention-2 通过分块计算精确注意力并避免物化注意力矩阵来减少内存流量（[Dao et al., 2022](https://arxiv.org/html/2608.10357#bib.bib2)；[Dao, 2024](https://arxiv.org/html/2608.10357#bib.bib3)）。StreamingLLM 表明，保留注意力汇点有助于稳定窗口化的长上下文推理（[Xiao et al., 2024](https://arxiv.org/html/2608.10357#bib.bib9)）。PyTorch FlexAttention 提供了可编程接口来组合掩码和分数修改（[PyTorch Contributors, 2026](https://arxiv.org/html/2608.10357#bib.bib5)）。SinkFlex-RL 结合了这些思想，将高效注意力既视为内存问题，也视为模型兼容性问题，因为汇点感知归一化和异构掩码必须保留在可微的训练路径中。

## 3 智能体 RL 设置

##### 双控制环境。

在单控制工具使用基准中，智能体通常是唯一通过工具改变环境的角色，而用户被动地提供信息。在双控制环境中，双方都能影响世界状态：智能体选择对话和 API 动作，而模拟用户可以确认选择、提供缺失信息或采取用户侧动作。一个任务实例可以表示为

E=(g,πSOP,Atool,s0,u),(1)

其中 g 是用户目标，πSOP 表示领域策略或标准操作程序约束，Atool 是智能体工具集，s0 是初始共享状态，u 是用户模拟器。一条轨迹为

τ={(ot,at,rt,dt,it)}t=1T,(2)

其中 ot 是观测，at 是智能体动作，rt 是奖励或诊断信号，dt 是终止标志，it 包含元数据。训练器消费由基准的程序化检查器产生的轨迹级奖励 Ri。其他仪表盘诊断信息被单独记录，并非来自单独训练的 critic 网络的输出。

##### 为什么这一场景对训练系统造成压力。

环境带来三种系统压力。首先，奖励是延迟的，通常只有在最终状态被检查后才能验证，这增加了每个有用梯度所需的采样 token 数量。其次，轨迹是多轮且工具密集的，因此策略必须在上下文中保留用户消息、工具输出和领域策略约束。第三，系统必须在保持回合状态和策略版本一致的同时，协调环境执行、rollout 生成、奖励检查和策略优化。

[Figure 1](https://arxiv.org/html/2608.10357#S3.F1) 总结了本工作使用的双控制任务结构：任务规范初始化共享世界状态，用户和智能体交替执行动作，生成的轨迹在 rollout 组用于 GRPO 风格更新之前由程序化检查器评分。

Refer to caption图 1：一个双控制智能体 RL 回合。任务规范初始化环境，用户–智能体交互产生多轮轨迹，程序化检查器对生成的 rollout 组进行评分，以进行 GRPO 风格策略更新。

## 4 训练流水线

[Figure 2](https://arxiv.org/html/2608.10357#S4.F2) 总结了训练数据流。我们抽象地描述各组件，以避免依赖特定部署的基础设施。

Refer to caption图 2：训练系统数据流。Actor 通过 rollout 工作器和 Gymnasium 兼容封装器采样多轮轨迹；轨迹奖励和元数据流向 GRPO 训练器，后者更新并重新分发策略权重。

##### 环境封装器。

每个基准领域都通过一个通用的 reset/step 接口暴露。在 reset 时，封装器采样一个任务，初始化共享状态，并构造第一个策略观测。在 step 时，它将模型输出路由到三个处理器之一：自然语言响应、工具调用或终止。封装器随后推进用户模拟器和状态后端，并返回下一个观测以及奖励元数据。这种分离将特定于基准的用户模拟、工具行为和检查逻辑排除在训练器之外，并旨在通过同一接口支持额外的可执行环境。

##### Rollout 工作器。

rollout 工作器拥有智能体循环。在每一轮，它格式化观测，从当前策略 πθ 中采样，将响应解析为动作，并将转换追加到轨迹缓冲区。工作器重复这一过程，直到环境终止或达到最大轮数预算。训练器接收 token 级对数概率、动作掩码、轨迹奖励和回合元数据，而不直接依赖特定于基准的解析或工具实现。

##### 策略更新。

给定一个提示或任务上下文 x 以及从旧策略采样的 G 条 rollout 组 y1,...,yG，GRPO 计算组归一化的轨迹优势

A^i=Ri−μ(R1:G)σ(R1:G)+εA,(3)

其中 Ri 是程序化计算的轨迹奖励，εA 是一个小的数值常数。在结果级监督下，相同的归一化优势被分配给 rollout yi 中的每个被优化 token。对于 token yi,t，重要性采样比率为

ρi,t(θ)=πθ(yi,t∣x,yi,<t)πθold(yi,t∣x,yi,<t).(4)

定义裁剪比率为

ρ̄i,t(θ)=clip⁡(ρi,t(θ),1−εc,1+εc).(5)

训练器最小化的损失为

Ji,t(θ)=min⁡(ρi,t(θ)A^i,ρ̄i,t(θ)A^i),(6)

LGRPO(θ)=−1∑iTi∑i=1G∑t=1TiJi,t(θ)+βDKL(πθ∥πref),(7)

其中 Ti 是 rollout i 的被优化 token 长度，εc 是裁剪半径，KL 项将策略正则化到参考模型。这是通常 PPO/GRPO 裁剪最大化目标的最小化形式（[Schulman et al., 2017](https://arxiv.org/html/2608.10357#bib.bib6)）。该基线设计不训练单独的 critic 或价值网络。

## 5 汇点感知 FlexAttention

##### 内存墙。

当以 eager 方式实现时，长度为 n 的轨迹会对每个相关的批次/头组件产生一个 n×n 的注意力分数结构。因此，分数位置的数量随上下文长度呈二次方增长；例如，n=20,000 对应 4×10^8 个位置，这还未计入激活、优化器状态、MoE 路由或 rollout 批处理。融合注意力方法通过分块计算注意力并减少高带宽内存流量来解决这一问题（[Dao et al., 2022](https://arxiv.org/html/2608.10357#bib.bib2)）。然而，特定于模型的注意力层可能还需要自定义分数修改、混合因果和滑动窗口掩码，或习得的汇点逻辑。

##### 注意力汇点。

注意力汇点是吸取注意力质量的 token 或习得机制，能够稳定长上下文行为。先前工作表明，在流式推理的窗口化注意力下，保留初始汇点 token 可以恢复质量（[Xiao et al., 2024](https://arxiv.org/html/2608.10357#bib.bib9)）。我们的场景与仅推理的 KV 缓存管理不同：汇点行为是可训练注意力计算的一部分，因此必须保持与前向和反向路径的连接。

##### 零值汇点等价。

为了避免在键和值缓存中物化显式汇点 token，我们使用零值汇点的代数形式。令 sη 为一个习得的汇点 logit，其值向量为 vsink=0。对于查询 q，带显式汇点的注意力输出为

Osink=∑iexp⁡(q⋅ki)vi+exp⁡(sη)0∑iexp⁡(q⋅ki)+exp⁡(sη)