Google 推出 Tunix：基于 JAX 的高吞吐智能体后训练库

大型语言模型（LLM）对齐的重点已迅速从静态聊天机器人对齐转向动态自治工作流。如今的模型不仅会对话——它们还能执行多步推理、调用外部 API 并与复杂环境互动。

训练推理型智能体面临特殊的挑战和瓶颈。智能体强化学习（agentic RL）训练的最新发展将过程从单轮对齐转向多轮决策，其中涉及复杂的环境互动和工具使用。这一转变在基础设施方面提出了新的挑战，影响展开性能和效率；当智能体暂停执行代码、查询数据库或等待网络搜索时，昂贵的 AI 加速器利用率会急剧下降，因为 TPU 在等待环境步骤时处于空闲状态。

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

除了编排之外，智能体强化学习还需要专门的可观测性。虽然标准分析器如 XProf 提供深度操作级别的轨迹，但其高开销限制了只能进行短时、零散的捕获。Tunix 引入了围绕特定领域 RL 指标直接构建的连续轻量级检测工具。通过将这些高层循环指标与 TPU 时间线关联，开发者可以获得执行效率的全局视图，从而快速发现并解决系统瓶颈。

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

实现硬件峰值吞吐量意味着需要让 TPU 不断保持繁忙状态。Tunix 通过结合异步 rollout 以消除执行空洞和滞后者，以及一个连续向训练器流式传输数据的解耦流水线来实现这一目标。

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

Tunix 通过异步轨迹收集引擎（Asynchronous Trajectory Collector Engine）解决了这个问题。

这种架构将模型推理、工具执行和奖励计算完全重叠，从而保持高硬件利用率。

尽管异步 rollout 解决了轨迹生成的瓶颈，但端到端 RL 工作流中硬件效率的另一个关键挑战是将动态、长度可变并且可能是长尾的 rollout 与严格同步的训练循环相连接。一个简单的做法是依赖一个同步点，这会强制加速器等待整个批次的轨迹完成后才能启动训练步骤，从而使训练 TPU 处于空闲状态。Tunix 通过将 rollout 和训练解耦成连续的生产者-消费者管道（如下图所示）消除了这一瓶颈：

轨迹组一旦完成，就会被后处理、评分，并直接流入训练器。这个管道确保同步训练器不断得到数据输入，从而最大化端到端吞吐量。

RL 框架中的一个主要摩擦点是算法与环境循环的严格耦合。修改代码库以支持新的开源软件（OSS）基准，例如 SWE-bench、WebArena 或自定义游戏引擎，通常需要大规模重写。

Tunix 通过解耦、可组合的架构解决了这个问题。通过公开干净的 API 边界，Tunix 自动化步骤调用和生命周期管理，使你可以完全专注于核心交互逻辑。

重要性：你可以在几分钟内接入任何开源 RL 环境。由于代理和环境逻辑与训练工作流完全解耦，用一个单步数学验证器替换为交互式 bash 终端无需修改任何训练代码。为了展示这种可组合设计的强大功能，我们接下来展示几个新代理、模型或环境能够多么轻松地使用的示例。你可以在我们的示例配方中找到更多定制化 Agent/Env 的详细示例：https://github.com/google/tunix/tree/main/examples。

Tunix 提供了内置类，如 ModelAgent 和 ToolAgent，可通过配置立即工作。

或者，您可以构建自己的自定义Agent，并添加具体逻辑来处理模型响应。Tunix会在端到端训练工作流中自动连接此Agent。例如：SWEAgent，FrozenLakeAgent

类似于Agents，Tunix提供了许多预构建的环境，包括TaskEnvironment，ToolEnvironment。或者，您也可以通过简单地实现几个主要API来引入自己的自定义环境，包括任何开源环境，例如下面的Gymnasium示例。

在大规模运行异步agentic训练时，传统日志记录显得不足。您需要细粒度但特定领域的可见性来识别效率问题：瓶颈出现在生成阶段吗？工具调用耗时过多吗？还是数据加载器过慢？

像XProf这样的标准分析器：https://openxla.org/xprof 提供详细的操作级追踪，以了解如内核和模型执行等微观性能。然而，使用这些工具捕获长跨度追踪通常成本高昂，并且在低级数据噪声中识别宏观瓶颈仍然困难。对于agentic RL的复杂工作流，开发者需要一种基于领域特定指标的轻量级宏观视图，这些指标能够直接映射到RL阶段。

Tunix通过仔细跟踪一组关键的RL特定指标提供了这一整体视图，这些指标既代表全球管道（如rollout、训练和权重同步阶段的交互），也代表重要子步骤（每次模型调用、环境交互等）。由于这些指标轻量，它们可以在整个训练任务中持续运行。用户可以快速识别工作流的全局停滞位置，然后部署像XProf这样的工具进行有针对性的进一步调试。

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

这种分阶段RL管道的宏观追踪使您能够：

Tunix 将分布式多轮强化学习（RL）的“黑箱”变成了训练任务可透明、可优化的时间线。

如果你正在评估用于智能体强化学习（Agentic RL）的框架，Tunix 的优势如下：

无论你是在复现最先进的推理模型、微调 Gemma 或 Qwen 系列模型以“思考”，还是部署复杂的多智能体系统，Tunix 都提供了下一代推理智能体所需的高性能基础。从今天开始构建吧！

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

The focus of LLM alignment has rapidly shifted from static chatbot alignment to dynamic agentic workflows . Today’s models don't just talk—they execute multi-step reasoning, call external APIs, and interact with complex environments.

Training reasoning agents encounters special challenges and bottlenecks. The recent evolution of agentic RL training shifts the process from single-turn alignment to multi-turn decision-making with complex environment interactions and tool usage. This shift raises new challenges on the infrastructure side for rollout performance and efficiency; when an agent pauses to execute code, query a database, or wait on a web search, the expensive AI accelerator utilization plummets as TPUs sit idle waiting for environment steps.

Tunix —Google’s post-training library—natively solves this bottleneck in its latest release, introducing an efficient, composable framework for training LLM agents at scale. Tunix keeps accelerators fully utilized on two fronts:

Beyond orchestration, agentic RL requires specialized observability. While standard profilers like XProf offer deep operator-level traces, their high overhead limits them to short, sporadic captures. Tunix introduces continuous, lightweight instrumentation built directly around domain-specific RL metrics. By correlating these high-level loop metrics with TPU timelines, developers get a global view of execution efficiency to quickly spot and resolve system bottlenecks.

Ultimately, Tunix is built to maximize TPU throughput, keep environments modular, and make multi-turn training efficiency fully transparent. Here is how it works under the hood.

Achieving peak hardware throughput means keeping TPUs constantly busy. Tunix accomplishes this by combining asynchronous rollouts to eliminate execution bubbles and stragglers with a decoupled pipeline that continuously streams data to the trainer.

In Agentic RL, trajectory generation (rollout) is the most time-consuming phase. However, the traditional synchronous rollout architecture creates two major problems, as depicted in the figure below.

Tunix solves this with an Asynchronous Trajectory Collector Engine .

This architecture completely overlaps model inference, tool execution, and reward computation, preserving high hardware utilization.

While async rollouts solve trajectory generation bottlenecks, another critical challenge for hardware efficiency in the end-to-end RL workflow is bridging dynamic, variable-length, and potentially long-tail rollouts with a strictly synchronous training loop. A naive approach relies on a synchronization point that forces the accelerator to wait until an entire batch of trajectories is complete before initiating the training step, starving the trainer TPU. Tunix eliminates this bottleneck by decoupling rollout and training into a continuous producer-consumer pipeline (illustrated in the diagram below):

The moment a trajectory group is complete, it is post-processed, scored, and streamed directly into the trainer. This pipeline ensures the synchronous trainer is constantly fed, maximizing end-to-end throughput.

A major friction point in RL frameworks is the rigid coupling of the algorithm to the environment loop. Modifying a codebase to support a new open-source software (OSS) benchmark like SWE-bench, WebArena, or a custom game engine often requires a massive rewrite.

Tunix resolves this with a decoupled, composable architecture. By exposing a clean API boundary, Tunix automates step invocation and lifecycle management so you can focus entirely on core interaction logic.

Why it matters: You can onboard any open-source RL environment in minutes. Because the agent and environment logic are completely decoupled from the training workflow, swapping a single-turn math verifier for an interactive bash terminal requires zero modifications to your training code. To demonstrate the power of this composable design, we next showcase a few examples of how easily new agents, models, or environments can be used. You can find more detailed examples of customized Agent/Env in our recipes：https://github.com/google/tunix/tree/main/examples.

Tunix offers built-in classes like ModelAgent and ToolAgent that work immediately via configuration.

Alternatively, you can build your own custom Agent and add specific logic on how to process the model responses. Tunix will automatically wire this agent in the end to end training workflow. E.g. SWEAgent , FrozenLakeAgent

Similar to Agents, Tunix offers a number of pre-built environments including TaskEnvironment , ToolEnvironment . Alternatively, you can also bring your own custom environment by simply implementing a few main APIs, including any open source environment such as the Gymnasium example below.

When running asynchronous agentic training at scale, traditional logging falls short. You need granular yet domain-specific visibility to identify efficiency problems: Is the bottleneck in the generation phase? Is the tool call taking too much time? Or is the data-loader too slow?

Standard profilers like XProf：https://openxla.org/xprof provide detailed, op-level traces to understand micro-level performance like kernel and model execution. However, capturing long-spanning traces with these tools is typically cost-prohibitive, and identifying macro-level bottlenecks within the noise of low-level data remains difficult. For the complex workflows of agentic RL, developers need a lightweight, macro-level view built on domain-specific metrics that map directly to RL stages.

Tunix delivers this big picture by carefully tracking a minimal set of critical RL-specific metrics that represent both the global pipeline (how rollout, training, and weight sync phases interact) and important sub-steps (each model call, environment interaction, etc.). Because they are lightweight, these metrics run continuously throughout the entire training job. Users can quickly identify where the workflow is stalling globally, and then deploy a tool like XProf for targeted, further debugging.

The figure above illustrates a Perfetto trace captured from a multi-turn agentic training job, detailing the staged execution timelines across CPU threads and TPU devices. As demonstrated by the trace, TPU device utilization is much higher than that of the CPU threads, whose idle time is primarily due to environment execution latency.

This macro-level tracing of staged RL pipelines enables you to:

Tunix turns the "black box" of distributed multi-turn RL into a transparent, optimizable timeline for the training job.

If you are evaluating frameworks for Agentic RL, here is how Tunix stands out:

Whether you are reproducing SOTA reasoning models, fine-tuning the Gemma or Qwen families to "think," or deploying complex multi-agent systems, Tunix provides the high-performance foundation needed for the next generation of reasoning agents. Start building today!

Tunix is under active open-source development by Google and the wider community. If you are building the next generation of reasoning agents, drop by our GitHub Issues and let us know what environments you are plugging in.