Tunix Adds High-Throughput Agentic RL Training on TPUs

Google's JAX-native post-training library Tunix now supports high-throughput agentic reinforcement learning, targeting multi-turn tool-use training.

axonn bots
axonn bots
·3 min read

Training a language model to reason well on a single-turn math problem is a fundamentally easier task than training it to act well across dozens of turns as an autonomous agent calling tools, reading results, and adjusting its plan. Google's Tunix library, a JAX-native framework for post-training large language models, has been extending its reinforcement learning support specifically to handle that harder, multi-turn agentic case at meaningful scale on TPUs.

Tunix was built to streamline the post-training stage of the LLM lifecycle: supervised fine-tuning, standard reinforcement learning from human feedback, and now agentic RL, all within a single JAX-native codebase. Its selling point relative to heavier alternatives is architectural restraint. Rather than layering abstraction on top of abstraction, Tunix exposes a comparatively thin, customizable training loop that lets researchers modify the parts they care about without fighting a framework designed for a different use case. It integrates with Flax NNX for modeling and leans on high-performance inference engines like vLLM and SGLang-JAX to handle the rollout phase, where the model actually generates the trajectories that reinforcement learning trains against.

Agentic reinforcement learning raises a specific engineering problem that standard RLHF doesn't: an agent's rollout isn't one forward pass, it's a sequence of interactions with an environment or a set of tools, and each of those interactions needs to happen efficiently and asynchronously if training is going to run at any reasonable throughput. Tunix's agentic RL support handles multi-turn agent-environment interaction, tool usage, and asynchronous rollout as first-class capabilities, rather than treating them as an awkward workaround bolted onto a single-turn training loop.

On the algorithm side, Tunix supports several reinforcement learning approaches suited to different tradeoffs. Its PPO learner implements Proximal Policy Optimization, the actor-critic method long considered a gold standard for RLHF and well suited to complex, sequential tasks like tool-using agentic workflows. Its GRPO learner offers Group Relative Policy Optimization, a critic-free approach that normalizes rewards across a group of generated responses to guide the model without the overhead of training a separate critic network, and a GSPO variant adjusts how token-level advantage is computed for better stability across long multi-turn trajectories.

An independent research group at UC San Diego, working on a project called GRL, put this agentic RL support to a concrete test by using video games as a source of verifiable, multi-turn training environments. Most RL post-training gains to date have come from domains like math and code, where correctness is easy to check automatically; games offered the team a similarly checkable but much richer multi-turn setting. Using Tunix's multi-host, multi-turn RL support alongside Google's TPU Research Cloud, the team reported measurable generalization beyond the games themselves, including gains in planning tasks and more modest, less consistent improvements on computer-use benchmarks, suggesting that skills learned in a game environment can transfer, at least partially, to more open-ended digital tasks.

The broader significance is less about any single benchmark and more about what Tunix's agentic RL support signals for where post-training research is heading. As more production AI systems shift from single-shot question answering toward agents that plan, call tools, and operate over many turns, the infrastructure used to train those systems needs to treat multi-turn interaction as the default case rather than a special one. A lightweight, JAX-native library that keeps that complexity manageable, rather than requiring researchers to build custom rollout infrastructure from scratch, lowers the bar for smaller teams to do serious agentic RL work on TPU hardware rather than leaving it exclusively to labs with dedicated infrastructure teams.