Efficient Reinforcement Learning for Long-Horizon Tool-Use Agentic Tasks

📄 arXiv: 2608.10357v1 📥 PDF

作者: Zelei Cheng, Amritansh Mishra, Sambit Sahu, William Campbell

分类: cs.LG, cs.AI

发布日期: 2026-08-11

备注: Published at the COLM 2026 Workshop on Efficient Reasoning


💡 一句话要点

提出SINKFLEX-RL以解决长时间工具使用代理任务的挑战

🎯 匹配领域: 支柱二:RL算法与架构 (RL & Architecture)

关键词: 强化学习 长时间任务 工具使用 模块化设计 内存优化 注意力机制 代理系统

📋 核心要点

  1. 长时间工具使用代理面临多轮在线回放导致的长上下文问题,现有方法在处理复杂任务时效率低下。
  2. 提出SINKFLEX-RL,通过模块化设计,结合环境接口、RL数据流和注意力内核设计,提升长时间代理的训练效率。
  3. 实验结果显示,优化后的注意力路径在内存使用上显著降低,验证奖励从0.25提升至0.44,展示了方法的有效性。

📝 摘要(中文)

长时间工具使用代理需要考虑用户目标、领域策略、工具调用、模拟器状态和延迟可验证奖励。强化学习(RL)适合这种场景,但多轮在线回放会产生长上下文,而特定模型的注意力层可能需要自定义掩码和学习的汇聚归一化。本文提出了SINKFLEX-RL,一个模块化的RL训练系统,适用于双控制工具使用环境。该系统结合了兼容Gymnasium的环境包装器、VERL风格的回放数据流、无单独价值模型的组相对策略优化,以及旨在在因果和滑动窗口掩码下保持模型特定汇聚缩放的汇聚感知FlexAttention路径。在初步的Tau2Bench零售运行中,验证奖励从训练初期的0.25上升到后期的0.44,同时训练分数和轨迹奖励代理也呈上升趋势。

🔬 方法详解

问题定义:本文旨在解决长时间工具使用代理任务中的复杂性,现有方法在多轮在线回放中面临长上下文和内存使用效率低的问题。

核心思路:SINKFLEX-RL通过模块化设计,整合环境接口和RL数据流,采用汇聚感知的FlexAttention路径,以提高训练效率和内存使用。

技术框架:系统包括一个兼容Gymnasium的环境包装器、VERL风格的回放数据流、组相对策略优化模块,以及汇聚感知的FlexAttention路径,确保在因果和滑动窗口掩码下的有效性。

关键创新:最重要的创新在于汇聚感知的FlexAttention路径设计,能够在不牺牲性能的情况下,优化内存使用并减少模型特定的汇聚缩放需求。

关键设计:在参数设置上,优化后的注意力路径将峰值VRAM从28.06GB降低至22.52GB,减少了19.7%;同时在8192-token配置下,使用25.53GB内存,避免了基线方法的内存溢出。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果显示,SINKFLEX-RL在Tau2Bench零售运行中,验证奖励从0.25提升至0.44,训练分数和轨迹奖励也呈上升趋势。此外,优化后的注意力路径显著降低了内存使用,展示了该方法在长时间任务中的优势。

🎯 应用场景

该研究的潜在应用领域包括机器人操作、智能家居和自动化制造等场景,能够有效提升长时间任务的执行效率和可靠性。未来,该方法可能推动更复杂的代理系统的发展,增强其在动态环境中的适应能力。

📄 摘要(原文)

Long-horizon tool-using agents must reason over user goals, domain policies, tool calls, simulator state, and delayed verifiable rewards. Reinforcement learning (RL) is a natural fit for this setting, but multi-turn on-policy rollouts create long contexts, while model-specific attention layers may require custom masks and learned sink normalization. We present SINKFLEX-RL, a modular training system for RL in dual-control tool-use environments. The system combines a Gymnasium-compatible environment wrapper, a VERL-style rollout dataflow, group-relative policy optimization without a separate value model, and a sink-aware FlexAttention path designed to preserve model-specific sink scaling under causal and sliding-window masks. In a preliminary Tau2Bench retail run, validation reward (mean@1) rises from 0.25 early in training to $0.44$ later in the observed training window, while training-score and trajectory-reward proxies also trend upward. In a fixed-configuration memory benchmark, the optimized attention path reduces peak VRAM from 28.06GB to 22.52GB at 4096 tokens, a $19.7\%$ reduction, and runs the measured 8192-token configuration using $25.53$~GB where the eager baseline runs out of memory. These results illustrate the value of integrating environment interfaces, RL dataflow, and attention-kernel design for memory-feasible long-horizon agent training.