HARTS: Efficient Agentic Reinforcement Learning for Hybrid-Attention Models over Arbitrary Rollout Trees
作者: Boyuan Meng, Peihua Bao, Hong Liu, Xiaowei Zhu, Chao Wang, Gen Li, Zhenxuan Pan
分类: cs.LG, cs.DC
发布日期: 2026-08-28
💡 一句话要点
提出HARTS以解决混合注意力模型中的高效强化学习问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 混合注意力模型 强化学习 回滚树 激活重计算 性能加速 微批次调度 数据并行
📋 核心要点
- 现有的强化学习方法在处理具有共享历史的回滚树时，往往需要独立重新计算共享前缀，效率低下。
- HARTS通过联合规划微批次和数据并行副本分配，采用非重放的紧凑标记工作，提升了混合注意力模型的执行效率。
- 在SWE-bench任务生成的Agentic RL工作负载上，HARTS实现了4.81到4.87倍的加速，展现了显著的性能提升。
📝 摘要（中文）
代理强化学习（RL）通常会产生具有共享历史的非规则回滚树。现有系统主要针对全注意力模型，缺乏与激活重计算兼容的密集可微混合注意力执行。本文提出HARTS（混合注意力RL在树结构上），通过前缀压缩后的非重放紧凑标记工作，联合规划微批次、数据并行副本分配和微批次槽调度。HARTS在每轮中将所有分支批量化为一个打包调用，支持激活重计算，并恢复每个标记的对数概率。HARTS在Agentic RL工作负载上实现了4.81到4.87倍的前向/反向/梯度加速，展示了在真实混合注意力模型上任意回滚树前缀共享的加速效果。
🔬 方法详解
问题定义：本文旨在解决现有强化学习方法在处理非规则回滚树时的效率问题，尤其是共享历史的重计算所带来的性能瓶颈。
核心思路：HARTS通过压缩前缀和非重放的紧凑标记工作，联合规划微批次和数据并行副本分配，从而提高了混合注意力模型的执行效率。
技术框架：HARTS的整体架构包括前缀压缩、微批次调度、状态恢复与重放等模块，确保在每轮中将所有分支批量化为一个打包调用。
关键创新：HARTS是首个在真实混合注意力模型上展示任意回滚树前缀共享加速的系统，显著提升了执行效率。
关键设计：HARTS采用了线性时间算法协调块边界状态恢复与重放，支持激活重计算，并在每个调用中恢复每个标记的对数概率，确保了数值对齐。
🖼️ 关键图片
📊 实验亮点
HARTS在Agentic RL工作负载上实现了4.81到4.87倍的前向、反向和梯度加速，性能提升显著。与基线相比，其数值差异可与自重跑变异相媲美，且在训练的前120步中，奖励趋势与基线相似，表明其有效性。
🎯 应用场景
HARTS的研究成果在强化学习、自然语言处理和机器人控制等领域具有广泛的应用潜力。通过提高混合注意力模型的执行效率，HARTS能够加速训练过程，提升智能体在复杂环境中的决策能力，具有重要的实际价值和未来影响。
📄 摘要（原文）
Agentic reinforcement learning (RL) often produces irregular rollout trees with shared histories. Training root-to-leaf trajectories independently recomputes these shared prefixes. Existing systems primarily target full-attention models and lack dense, differentiable hybrid-attention execution compatible with activation recomputation. We present HARTS (Hybrid-Attention RL over Tree Structures). HARTS jointly plans microbatches, data-parallel (DP) replica assignments, and microbatch-slot schedules using non-replay compact-token work after prefix compression. For chunkwise linear attention, a linear-time algorithm coordinates chunk-boundary state recovery and replay and produces the minimum number of sequential linear-attention calls under our packed execution model. HARTS preserves the chunkwise state partitioning of trajectory-wise training: it does not repeat projections, MLP/MoE computation, or final outputs, and performs only bounded state replay for numerical alignment. Per round, HARTS batches all branches into one packed call, propagates gradients through differentiable state handoffs, supports activation recomputation, and restores per-token log-probabilities. For deterministic, no-token-drop top-$k$ MoE routing, semantic multiplicities restore MoE-objective token weights and load statistics. Existing RL objectives retain their interface. To our knowledge, HARTS is the first system to demonstrate arbitrary-rollout-tree prefix-sharing speedups on a real hybrid-attention model. On an Agentic RL workload generated from SWE-bench tasks, HARTS achieves $4.81$--$4.87\times$ forward/backward/gradient speedup with activation recomputation across multiple parallel configurations. Its numerical differences are comparable to baseline self-rerun variation, and its reward trend is similar to the baseline over the first 120 steps of $τ^3$-Bench training.