Tunix, Google's JAX-native post-training library, addresses key infrastructure bottlenecks in agentic reinforcement learning. When training LLM agents, TPUs sit idle waiting for environment steps like code execution or web searches. Tunix solves this with an Asynchronous Trajectory Collector Engine that overlaps model inference, tool execution, and reward computation. A decoupled producer-consumer pipeline continuously feeds the synchronous trainer, eliminating starvation. The framework also exposes a composable API so developers can plug in any custom environment (including Gymnasium-compatible ones) or agent without modifying training code. For observability, Tunix provides lightweight, continuous RL-specific metrics that give a macro-level view of pipeline stages, complementing low-level profilers like XProf. Code examples show how to wire up built-in agents like ModelAgent and ToolAgent, or bring custom implementations, with minimal boilerplate.
