Scaling Agentic RL: High-Throughput Agentic Training with Tunix

扩展智能体强化学习：使用Tunix进行高吞吐量智能体训练 - Google Developers Blog

Google Tag Manager (noscript)

End Google Tag Manager (noscript)

HTML

扩展智能体强化学习：使用Tunix进行高吞吐量智能体训练

2026年7月21日

高宇 Gao

软件工程师

Lance Wang

Shadi Noghabi

Tianshu Bao

Weiren Yu

产品经理

分享

• Facebook

• Twitter

• LinkedIn

• 邮件

大语言模型对齐的焦点已迅速从静态聊天机器人对齐转向动态智能体工作流程。如今的模型不仅会对话，还能执行多步骤推理、调用外部API并与复杂环境进行交互。

训练推理智能体面临特殊挑战和瓶颈。近期智能体强化学习训练的演进将流程从单轮对齐转变为包含复杂环境交互和工具使用的多轮决策过程。这种转变在基础设施层面带来了新的挑战：当智能体暂停执行代码、查询数据库或等待网页搜索时，TPU会因等待环境步骤而空闲，导致昂贵的AI加速器利用率骤降。

Tunix——Google的后训练库——在最新版本中原生解决了这一瓶颈，引入了用于大规模训练LLM智能体的高效可组合框架。Tunix从两个方面确保加速器充分利用：

• 异步执行：高并发执行引擎完全解耦TPU执行与主机环境延迟（如网络I/O或工具执行）。

• 无阻塞流水线：动态生产者-消费者架构持续批量处理并流式传输可变长度轨迹到训练器，防止流水线停滞。

除了调度管理，智能体强化学习还需要专用可观测性。虽然XProf等标准性能分析工具能提供深度操作级追踪，但其高开销限制了其仅用于短时、偶发的捕获。Tunix引入了围绕领域特定强化学习指标构建的持续轻量级监控，通过将这些高层循环指标与TPU时间线相关联，开发者可获得全局执行效率视图，快速定位并解决系统瓶颈。

最终，Tunix的设计目标是最大化TPU吞吐量、保持环境模块化，并使多轮训练效率完全透明。以下是其内部工作原理。

1. 异步与解耦执行：接近零空闲时间，最大吞吐量

实现峰值硬件吞吐量意味着要持续保持TPU忙碌。Tunix通过结合异步执行消除执行空洞和拖尾效应，配合解耦流水线持续向训练器传输数据，从而实现这一目标。

异步执行

在智能体强化学习中，轨迹生成（执行）是最耗时的阶段。然而，传统同步执行架构会产生两大问题，如下图所示：

• 执行空洞：当执行同步等待环境初始化或返回状态和奖励时，会引发加速器执行空洞并降低效率。

• 拖尾效应：批量生成也容易受到长尾问题影响，整体延迟由组内最慢轨迹决定。

Tunix通过异步轨迹收集引擎解决了这些问题。

• 高并发执行：通过在 RolloutOrchestrator 中使用 Python 的 asyncio，框架能够管理大规模的并发智能体-环境交互。当某个智能体等待主机端工具执行时，推理引擎会立即切换至其他活跃轨迹的 token 生成。

• 异步 vLLM & SGLang 集成：Tunix 原生支持高性能推理引擎（如 vLLM-TPU 和 SGLang-Jax）。通过启用异步请求处理，引擎可实现非阻塞采样并在 TPU 上实现最大并发性。

该架构完全重叠模型推理、工具执行和奖励计算过程，保持硬件的高利用率。

解耦的 Rollout 与训练流水线

虽然异步 Rollout 解决了轨迹生成的瓶颈，但端到端强化学习工作流中另一个关键挑战是：如何将动态、变长且可能具有长尾特性的 Rollout 与严格同步的训练循环对接。简单方法依赖于同步点，迫使加速器必须等待整个轨迹批次完成才能启动训练步骤，导致训练 TPU 资源闲置。Tunix 通过将 Rollout 与训练解耦为连续的生产者-消费者流水线（如下图所示）消除这一瓶颈：

• 生产者：异步 Rollout 协调器持续将完成的轨迹输出到高吞吐量队列。

• 消费者：AgenticRLLearner 从该队列中消费数据。对于需要每个提示生成多条推理路径以计算组优势的算法（如 GRPO），Tunix 会动态地实时对这些异步轨迹进行分组。

一旦轨迹组完成，系统会立即进行后处理、评分，并将数据直接流式传输至训练器。该流水线确保同步训练器持续获得数据，最大化端到端吞吐量。

2. 可组合的智能体与环境抽象 – 即插即用的开源环境

强化学习框架中的主要摩擦点是算法与环境循环的刚性耦合。要让代码库支持新的开源软件（OSS）基准（如 SWE-bench、WebArena 或自定义游戏引擎），通常需要大规模重写代码。

Tunix 通过解耦、可组合的架构解决这一问题。通过暴露清晰的 API 边界，Tunix 自动化步骤调用和生命周期管理，使您可以专注于核心交互逻辑。

• 智能体层：管理提示格式化、动作生成和对话历史。它会自动应用策略模型的聊天解析器，并在多轮对话边界保留特殊 token——这对确保严格的 Token-In, Token-Out（TITO）行为至关重要。您可以通过继承 ConversationAgentBase 类轻松自定义生成逻辑。

• 环境层：开箱即用，Tunix 提供了预构建的 TaskEnvironment 和 ToolEnvironment 类。您也可以继承 BaseTaskEnv 来对接任何外部系统。Tunix 自动处理多轮对话的生命周期、观察值路由和奖励计算。

为什么这很重要：您可以在几分钟内导入任何开源RL环境。由于智能体和环境逻辑与训练流程完全解耦，将单轮数学验证器替换为交互式bash终端无需对训练代码进行任何修改。为了展示这种可组合设计的强大功能，接下来我们将展示如何轻松使用新的智能体、模型或环境。您可以在我们的配方中找到更多定制化Agent/Env的详细示例。

示例1：预构建与自定义智能体

Tunix 提供了可通过配置立即使用的内置类，如 ModelAgent 和 ToolAgent。

from tunix.rl.agentic.agentic_grpo_learner import GRPOLearner
from tunix.rl.agentic.agents.model_agent import ModelAgent, ToolAgent

# 非工具调用的单轮智能体
learner = GRPOLearner(
    agent_class=ModelAgent, 
    agent_kwargs={"system_prompt": "my system prompt"},
    ...
)

# 自定义工具调用智能体
tool_map = {"calculator": CustomizedCalculatorClass, ...}
learner = GRPOLearner(
    agent_class=ToolAgent, 
    agent_kwargs={
        "system_prompt": "my system prompt",
        "tool_parser_name": "gemma",
        "tool_map": tool_map,
    },
    ...
)

Python

已复制

或者，您可以构建自己的自定义智能体，并添加特定逻辑来处理模型响应。Tunix 会自动将此智能体接入端到端训练流程。例如 SWEAgent、FrozenLakeAgent。

from tunix.rl.agentic.agents.base_agent import ConversationAgentBase
from tunix.rl.agentic.agents import agent_types

# 您可以自带智能体！
# 注意智能体无需了解模型的具体实现（无论是Qwen、Llama还是Gemma）
class MyAgent(ConversationAgentBase):
    def __init__(self, args):
        ...

    def update_from_model(self, response: str, **kwargs) -> agent_types.Action:
        # 自定义逻辑处理原始响应（例如提取<answer>标签）
        ...

# Tunix 自动接入端到端流程
learner = GRPOLearner(agent_class=MyAgent, agent_kwargs={...}, ...)

示例2：引入自定义环境

与智能体类似，Tunix 提供了多种预构建环境，包括 TaskEnvironment 和 ToolEnvironment。或者，您也可以通过简单实现几个核心API来引入自定义环境，包括任何开源环境，例如下面的Gymnasium示例。

import gymnasium as gym
from tunix.rl.agentic.agentic_grpo_learner import GRPOLearner
from tunix.rl.agentic.environments.base_environment import BaseTaskEnv, EnvStepResult

# 您只需关注环境交互的核心逻辑，Tunix 会自动处理其余生命周期管理和函数调用
class MyEnv(BaseTaskEnv):
    def _initial_observation(self):
        # 处理环境创建和初始观测
        self.env = gym.make("your_chosen_env")
        observation, info = self.env.reset(seed=42)
        return observation

    def _step_impl(self, action):
        # 计算观测值、奖励、完成状态、信息
        action = self.env.action_space.sample() 
        obs, reward, done, info = self.env.step(action)
        return EnvStepResult(obs, reward, done, info)

    def close(self):
        self.env.close() # 轨迹结束后清理环境

learner = GRPOLearner(env_class=MyEnv, ...)

3. 消除黑箱：RL专用轻量级性能分析

在大规模运行异步代理训练时，传统日志记录方式已无法满足需求。你需要具备细粒度且领域特定的可见性来识别效率问题：瓶颈出现在生成阶段吗？工具调用耗时过长？还是数据加载器速度太慢？

像XProf这样的标准分析工具能提供详细的运算级追踪，用于理解内核和模型执行等微观性能。然而，使用这些工具捕获长跨度追踪通常成本高昂，且在低级别数据噪声中识别宏观瓶颈依然困难。对于代理式强化学习（agentic RL）的复杂工作流，开发者需要基于领域特定指标构建的轻量级宏观视图，这些指标能直接映射到RL阶段。

Tunix通过精确追踪少量关键RL专用指标，提供全局概览。这些指标既涵盖整体流程（rollout、训练和权重同步阶段的交互），也包含重要子步骤（每个模型调用、环境交互等）。由于轻量特性，这些指标在整个训练任务中持续运行。用户可快速定位全局流程阻塞点，再使用XProf等工具进行针对性调试。

上图展示了从多轮代理训练任务中捕获的Perfetto追踪，详细呈现了CPU线程和TPU设备间的分阶段执行时间线。如追踪所示，TPU设备利用率远高于CPU线程，而CPU线程的空闲时间主要源于环境执行延迟。

这种对分阶段RL流水线的宏观追踪使你能够：

• 定位TPU资源饥饿：可视化Python工具调用或环境执行具体阻塞异步流水线的时间点。反过来，它也能验证并行rollout是否成功重叠以保持加速器饱和。

• 验证流水线对齐：跟踪宏观阶段的精确时间点，确保各阶段对齐且不引入隐藏的延迟气泡。你可以轻松验证训练器是否未在rollout生成上等待，或权重同步是否未导致严重执行延迟。

• 优化训练配置：通过指标数据动态调整性能。例如，可通过直接关联线程池与TPU空闲时间来调整最大rollout并发数，或根据数据生成速率和HBM约束优化训练微批次大小。

Tunix将分布式多轮RL的"黑箱"转化为训练任务的透明、可优化时间线。

Tunix与生态系统的对比

在评估代理式RL框架时，Tunix的差异化优势体现在：

• 与OpenRLHF / veRL对比：OpenRLHF和veRL通过Ray + vLLM取得了显著进展。但它们主要面向PyTorch生态。Tunix原生支持JAX/TPU生态，无缝集成于JAX、Flax和Optax之上，提供原生Pathways多主机分布式训练，利用XLA编译器优化。

• 与Hugging Face TRL对比：TRL非常适合标准单轮SFT（监督微调）和RLHF（基于人类反馈的强化学习）。但协调复杂的多轮异步循环通常需要大量定制胶水代码。Tunix开箱即支持多轮工具使用环境，将其作为一等公民。

• 与 Ray RLlib 对比：RLlib 是一个功能全面、用途广泛的强化学习框架。然而，将现代大语言模型（LLMs）原生映射到加速器上进行权重共享，同时避免高昂的开销，这一过程非常复杂。Tunix 完全颠覆了传统方式：它是一个以大语言模型为中心的库，能够将高性能强化学习直接引入原生大语言模型服务基础设施。

立即开始构建你的智能体

无论你是要复现最前沿的推理模型，还是对 Gemma 或 Qwen 系列进行微调以实现“思考”能力，又或者部署复杂的多智能体系统，Tunix 都能为你提供下一代推理智能体所需的高性能基础。立即开始构建吧！

• 🌟 星标仓库 & 查看代码：github.com/google/tunix

• 📖 实践案例：软件工程编码智能体、数学推理、游戏智能体

• 📖 阅读文档：深入探索我们的智能体强化学习架构 tunix.readthedocs.io

• 🚀 快速入门：进入我们的 /examples 文件夹，探索各种实践案例并立即运行你的第一个训练任务！

Tunix 正在由 Google 和更广泛的社区进行积极的开源开发。如果你正在构建下一代推理智能体，请访问我们的 GitHub Issues 页面，告诉我们你正在接入哪些环境。

posted in:

Previous

Next