The bottleneck is TPU idling
Tunix is Google's new JAX-native post-training library for training LLM agents, and it exists because of a specific pain point. When you train a multi-turn, tool-using reasoning agent with reinforcement learning, the TPU spends most of its time waiting. The agent pauses to call a tool. The tool calls the network. The network is slow. The TPU sits idle. Multiply that across a batch of agents and you get a small fraction of the hardware utilization you expected.
Tunix is built to eliminate that idle time.
How it keeps the TPU busy
The library attacks the bottleneck on two fronts. The first is asynchronous rollouts. A high-concurrency rollout engine, built on Python's asyncio within the RolloutOrchestrator, decouples TPU execution from host-side environment latency. While one agent pauses for a tool call, the inference engine pivots to generate tokens for other active trajectories. The integration with vLLM-TPU and SGLang-Jax is native, and async request handling ensures non-blocking sampling and maximum concurrency on the TPU. Inference, tool execution, and reward computation all overlap, preserving hardware utilization.
The second is barrier-free pipelining. A dynamic producer-consumer architecture constantly batches and streams variable-length trajectories to the trainer, preventing pipeline stalls. The producer is the async rollout orchestrator, continuously yielding completed trajectories into a high-throughput queue. The consumer is the AgenticRLLearner, which pulls from that queue. For algorithms like GRPO that require multiple reasoning paths per prompt to compute group advantages, Tunix dynamically groups these asynchronous trajectories on the fly. The moment a trajectory group is complete, it is post-processed, scored, and streamed directly into the trainer. The pipeline ensures the synchronous trainer is constantly fed.
Continuous profiling, not sporadic deep dives
Beyond orchestration, agentic RL needs specialized observability. Standard profilers like XProf offer deep operator-level traces, but their 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 spot and resolve system bottlenecks quickly. The instrumentation is RL-stage-aware, so the metrics that matter for the rollout phase are different from the metrics that matter for the trainer phase.
What you can train with it
Tunix v0.1.6 added Agentic RL and VLM training. The earlier v0.1.0 release covered supervised fine-tuning (full weights fine-tuning, LoRA, DPO, ORPO), reinforcement learning (PPO, GRPO, GSPO-token, DAPO, DR-GRPO), and agentic RL with multi-turn tool use, asynchronous rollout for high-throughput trajectory collection, and trajectory batching and grouping. Models supported include Gemma, Llama, and Qwen families. The training stack integration covers Flax, Optax, Orbax, with MaxText and MaxDiffusion at the model layer.
Why it matters
The agentic RL ecosystem has been waiting for a library that treats the TPU idling problem as a first-class concern. Most agentic RL frameworks inherit the synchronous rollout model from language model RL, and the synchronous model is exactly wrong for tool-using agents. Tunix's bet is that the asynchronous model is correct, and the proof is the throughput numbers. The library is under active development, the V2 release is the first one with serious agentic focus, and the path forward is to keep adding more algorithms, more environments, and more integrations.