The hard part of training multi-turn agents is not always the model computation. It is the waiting: tool calls, code execution, database queries, web search, reward calculation, and environment steps can leave expensive accelerators idle. Google Tunix, a JAX-native post-training library, is aimed at that infrastructure gap. Its latest release focuses on keeping TPUs busy while reasoning agents interact with slow, variable environments.

Google frames the problem as a shift from static chatbot alignment to dynamic agent workflows. In a traditional synchronous rollout design, the system waits for an environment to initialize or return a state and reward before progressing. That creates execution bubbles on accelerators. Batched generation also suffers from stragglers, where one long trajectory dictates latency for the whole group. Tunix uses a Python asyncio-based RolloutOrchestrator to manage large pools of concurrent agent-environment interactions. When one agent pauses for host-side tool execution, the inference engine can generate tokens for another active trajectory. Tunix also integrates with async vLLM-TPU and SGLang-Jax.

The second change is decoupling rollout from training. A naive pipeline waits until an entire batch of trajectories is complete, which can starve the trainer TPU. Tunix instead sends completed trajectories into a high-throughput queue. AgenticRLLearner consumes that stream, dynamically grouping asynchronous trajectories for algorithms such as GRPO, where multiple reasoning paths per prompt are needed to calculate group advantages. Once a group is complete, it is processed, scored, and streamed to the trainer.

The release also adds composable agent and environment abstractions. Built-in ModelAgent, ToolAgent, TaskEnvironment, and ToolEnvironment classes cover common cases, while developers can attach custom agents or environments such as SWE-bench, WebArena, a game engine, or an internal simulator without rewriting the training loop. For observability, Tunix tracks lightweight RL-specific metrics across rollout, training, weight synchronization, tool latency, and environment latency. The useful test now is adoption: whether Tunix recipes for coding agents, math, and games produce measurable benchmark gains for Gemma, Qwen, and other open model families.

#tunix

Google Tunix targets the idle-TPU problem in agentic RL training