# Google Tunix Agentic RL: Async Rollouts on JAX and TPU

> Design Tunix agentic RL pipelines with async rollouts, grouped trajectories, composable tool environments, TPU meshes, and Perfetto traces.

- **Published**: 2026-07-23
- **Category**: AI Infrastructure
- **URL**: https://agentpedia.codes/blog/google-tunix-agentic-rl-guide

---

> **Important callout**

**Bottom line:** Tunix gives JAX and TPU teams a first-party agentic RL path with concurrent multi-turn environments, grouped trajectory queues, separate rollout and training meshes, and continuous Perfetto traces. Its architecture can hide tool and environment latency, but Google's July 21, 2026 post does not publish a standardized throughput comparison. Treat "near-zero idle time" as a design goal to verify on your own workload.

Google published its [high-throughput agentic training guide](https://developers.googleblog.com/scaling-agentic-rl-high-throughput-agentic-training-with-tunix/) on July 21, 2026. This article checks the launch claims against Tunix's current documentation and source code, then turns them into a deployment design, starter configuration, failure checklist, and measurement plan. The research cutoff is July 23, 2026.

## What Google published

[Tunix](https://github.com/google/tunix) is Google's Apache-2.0, JAX-based library for LLM post-training. It covers supervised fine-tuning, preference optimization, reinforcement learning, and agentic RL. The current README labels the project a **V2 release** under active development. At the research cutoff, the `main` branch identified `google-tunix` 0.1.7 as **Alpha**, matching the latest GitHub release, v0.1.7. V2 describes the project's architecture generation, while Alpha is the package maturity classifier teams should use for risk planning.

The July release makes four parts of agentic RL explicit:

| Component | Current Tunix behavior | Practical implication |
| --- | --- | --- |
| Trajectory collection | `RolloutOrchestrator` runs multiple agent-environment episodes with `asyncio` | One slow tool call does not have to block every other active episode |
| Grouping | `GroupQueueManager` releases a group only after the configured number of trajectories for that key is complete | GRPO keeps its same-prompt comparison requirement without waiting for unrelated prompt groups |
| Pipeline placement | Rollout, actor, and reference roles can share a mesh or use distinct meshes | Separate rollout and trainer hardware can overlap work; a shared mesh executes phases sequentially |
| Observability | RL stages and sub-steps can be exported as Perfetto timelines | Teams can inspect rollout, environment, training, data-loading, and weight-sync gaps across a full job |

Google describes the result as asynchronous rollouts plus "barrier-free" pipelining. The useful interpretation is narrower: Tunix removes the **whole input-batch completion barrier** by streaming ready trajectory groups. It does not remove every dependency in the algorithm or runtime.

## How the asynchronous pipeline works

A Tunix agentic GRPO run has two levels of concurrency.

First, each `TrajectoryCollectEngine` owns one agent-environment episode. It repeatedly asks the model for an action, steps the environment, records observations and rewards, and stops on termination, timeout, token limit, or the episode step limit. The orchestrator keeps up to `max_concurrency` of these episodes active. When one episode waits on a calculator, shell sandbox, search service, or remote judge, another can submit model work.

Second, completed trajectories flow through a grouped queue. For GRPO, Tunix creates `num_generations` agent-environment pairs for each source prompt and gives them the same group ID. A group becomes consumable only when all of its generations are present. Ready groups can reach downstream reward, advantage, log-probability, and training work while unrelated long-tail groups are still running.

```text
prompts
  └─> agent + environment pairs (up to max_concurrency active)
        ├─> model call ─> tool/environment step ─> next model call
        ├─> model call ─> tool/environment step ─> finish
        └─> model call ─> finish
              │
              v
       grouped trajectory queue
       (num_generations per prompt)
              │
              v
       reward + advantage + log probabilities
              │
              v
       trainer update ─> weight synchronization ─> next policy version
```

This design attacks the straggler problem without changing GRPO's group-relative calculation. It also explains why queue depth alone is not a success metric: a growing ready queue can mean rollout is healthy while the trainer is under-provisioned, and many half-filled groups can consume memory without producing trainable batches.

## Which barriers remain

Three synchronization points still matter.

1. **GRPO group completion.** Tunix cannot compute group-relative advantages until every required generation for a prompt is available. A single pathological episode can still delay its own group.
2. **Weight synchronization.** When the trainer and rollout model do not share weights, Tunix requests a weight-sync lock. New rollout starts pause while updated parameters are transferred and the policy version advances.
3. **Hardware placement.** Tunix enables rollout/training overlap only when actor and rollout use different meshes. If the roles share a mesh, interleaving would create resource contention, so the pipeline remains sequential across those phases.

Tunix also exposes `off_policy_steps`. Raising it permits older trajectories after a policy update, which can keep a pipeline moving but changes the freshness boundary. Start at zero, measure weight-sync stalls, and relax the setting only with an explicit numerical-quality experiment.

> **Warning callout**

Do not translate Google's "barrier-free" wording into "fully asynchronous optimization." The trainer remains synchronous, complete GRPO groups remain required, and weight synchronization deliberately excludes new rollout starts.

## Tools, agents and environments

Tunix separates the policy-facing agent from the task-facing environment:

- `ModelAgent` treats one model response as the final answer for single-turn tasks.
- `ToolAgent` parses tool calls, invokes registered tools, and appends tool results to conversation history.
- `TaskEnvironment` supplies a single-turn task and final reward.
- `ToolEnvironment` supports repeated tool calls until the agent invokes `finish` or reaches `max_steps`.
- Custom agents inherit `ConversationAgentBase`; custom environments inherit `BaseTaskEnv` and implement initialization, stepping, and cleanup behavior.

A tool implements a JSON schema plus synchronous `apply()` or asynchronous `apply_async()` logic. `ToolManager` can execute multiple calls in parallel. The abstraction is useful, but framework-level concurrency does not make a tool safe. Every environment still needs bounded time, output size, credentials, network access, filesystem access, and cleanup.

A practical environment contract should make these outcomes unambiguous:

| Outcome | Required behavior |
| --- | --- |
| Normal step | Return an observation, scalar reward contribution, termination state, and useful metadata |
| Tool timeout | Cancel work, release resources, and emit a machine-readable failure status |
| Invalid action | Return a deterministic observation or terminate according to the task policy |
| Episode limit | Preserve the truncation reason instead of presenting the result as a normal completion |
| Cleanup | Close sandboxes, clients, files, and subprocesses even after cancellation |

For a broader application-side view of tool permissions and handoffs, compare AgentPedia's [multi-agent orchestration guide](/blog/antigravity-agent-orchestration-multi-agent) and [agent security guide](/blog/antigravity-security-guide). Tunix controls training flow; it does not replace sandboxing or least privilege.

## Choose a deployment shape

Tunix supports colocated and disaggregated role placement.

| Shape | Placement | What happens | Best first use |
| --- | --- | --- | --- |
| Colocated | Actor, reference, and rollout share one mesh | Components use the TPU mesh sequentially; optional CPU offload can reduce HBM pressure | Small experiments where the model needs most available memory |
| Disaggregated two-way | Trainer/reference and rollout use different meshes | Rollout and training can overlap; weights must move between meshes | Throughput tests with enough devices to dedicate rollout capacity |
| Disaggregated three-way | Trainer, reference, and rollout have separate allocations | More independent capacity and more placement/synchronization complexity | Large runs after the two-way topology is understood |

Google's current demo script exposes `colocated`, `disaggregated-2-way`, and `disaggregated-3-way` cluster setups. The [performance guide](https://tunix.readthedocs.io/en/latest/performance.html) recommends disaggregation when global throughput matters, but it does not supply a universal split. Measure generation rate, training consumption rate, HBM, and synchronization time before changing the ratio.

The easiest mistake is to reserve too little HBM for the rollout engine. vLLM and SGLang-JAX need model weights, runtime buffers, and KV cache; the trainer needs parameters, optimizer state, gradients, and activations. A topology that fits at startup can still fail on longer prompts or more concurrent sequences.

## Starter configuration

Tunix requires Python 3.11 or newer. The official quickstart recommends a virtual environment and `google-tunix[prod]` for TPU-backed JAX:

```bash
python3 -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install "google-tunix[prod]"

python - <<'PY'
import jax
import tunix
print("backend:", jax.default_backend())
print("devices:", jax.devices())
print("tunix:", tunix.__file__)
PY
```

Run the repository's agentic example before adapting it. The current [Gemma agentic GRPO notebook](https://github.com/google/tunix/blob/main/examples/agentic/gemma_grpo_demo_nb.py) demonstrates the full actor, reference, tokenizer, rewards, cluster, and learner setup. Accelerated vLLM and SGLang-JAX engines require separate, version-sensitive installation; they are not bundled by the basic package install.

The following fragment shows the important configuration relationships rather than a complete training program:

```python
from tunix.rl import rl_cluster
from tunix.rl.agentic.agentic_grpo_learner import GRPOConfig
from tunix.rl.rollout import base_rollout

algo = GRPOConfig(
    num_generations=4,
    max_concurrency=32,
    max_response_length=768,
    episode_timeout=600.0,
    off_policy_steps=0,
)

rollout = base_rollout.RolloutConfig(
    max_prompt_length=256,
    max_tokens_to_generate=768,  # must match max_response_length
    kv_cache_size=1280,          # vanilla: prompt + generation + headroom
    temperature=0.9,
    top_p=1.0,
    top_k=50,
    return_logprobs=True,
)

cluster = rl_cluster.ClusterConfig(
    role_to_mesh={
        rl_cluster.Role.ACTOR: train_mesh,
        rl_cluster.Role.REFERENCE: train_mesh,
        rl_cluster.Role.ROLLOUT: rollout_mesh,
    },
    rollout_engine="vanilla",
    training_config=training_config,
    rollout_config=rollout,
)
```

The current `AgenticRLLearner` validates two easy-to-miss constraints: `max_tokens_to_generate` must equal `max_response_length`, and rollout log probabilities must be enabled when `use_rollout_logps=True`. Agentic vLLM also requires server mode in the current source. Treat source examples and pinned requirements as the compatibility contract, not a loosely remembered set of package versions.

For vLLM or SGLang-JAX, start from the official [`grpo_demo_llama3_qwen2.py`](https://github.com/google/tunix/blob/main/scripts/grpo_demo_llama3_qwen2.py). It includes rollout-engine selection, model mappings, mesh construction, and colocated/disaggregated paths. Copying only the engine name omits required model mapping, memory, and weight-sync configuration.

## Continuous Perfetto observability

Tunix includes the `perfetto` package and instruments agentic-loop stages with lightweight RL-specific spans. The July launch distinguishes this continuous macro view from short, operator-level XProf captures. Current documentation limits performance tracing to the GRPO main entry point; its v2 path is experimental, and trace writing must be enabled. Google does not publish tracing-overhead measurements, so "lightweight" is a design claim to test rather than a free guarantee.

Use the timeline to answer operational questions, not to decorate a report:

| Trace pattern | Likely constraint | Next experiment |
| --- | --- | --- |
| Long gaps between model calls inside episodes | Environment or tool latency | Separate queue, network, sandbox startup, and reward timing |
| Rollout devices busy while trainer waits | Too few complete groups or insufficient rollout rate | Increase safe concurrency; inspect half-filled groups and stragglers |
| Ready data accumulates while trainer stays busy | Training is the consumer bottleneck | Adjust trainer mesh, micro-batch, or model memory strategy |
| Repeated weight-sync spans dominate steps | Transfer/resharding or update cadence | Measure LoRA-only sync, topology, and update frequency |
| Many shape compilations or slow first batches | Variable shapes and JAX compilation | Stabilize prompt/generation buckets and warm representative shapes |

Keep throughput counters beside the trace: completed trajectories per minute, complete groups per minute, trainer examples per minute, policy age, timeout rate, queue occupancy, HBM high-water mark, and accepted task reward. High accelerator utilization with a collapsing solve rate is not a successful optimization.

## Tunix vs veRL, OpenRLHF and TRL

This is a design comparison, not a benchmark. The projects evolve quickly, and all four now document agent or multi-turn paths.

| Framework | Primary stack | Current agentic path | Placement and rollout emphasis | Choose it when |
| --- | --- | --- | --- | --- |
| Tunix | JAX, Flax NNX, Optax; TPU-oriented | Built-in agents, tools, environments, concurrent trajectory collection, grouped queues | JAX meshes; vanilla, vLLM-TPU, and SGLang-JAX rollout; Pathways multi-host support | Your models and operations already live in JAX/TPU and you want those abstractions in one stack |
| veRL | PyTorch with FSDP/FSDP2 or Megatron; Ray-based control | Multi-turn tool calling plus the separate `uni-agent` project; Agent Loop API is labeled Alpha | Flexible GPU placement with vLLM, SGLang, or Transformers rollout | You need an established PyTorch/GPU RL stack with broad algorithm and engine coverage |
| OpenRLHF | PyTorch, Ray, vLLM, DeepSpeed or Molt | Unified single- and multi-turn agent executors; async and partial-rollout paths | GPU hybrid/colocated and distributed Ray workflows | You want a Ray + vLLM RLHF stack with explicit agent-executor interfaces |
| TRL | Transformers, Accelerate, PyTorch trainers | `GRPOTrainer` supports stateful environments and multi-environment routing; those hooks and AsyncGRPO remain experimental | High-level Trainer/Accelerate path; optional vLLM and distributed backends | You value the Hugging Face model/dataset ecosystem and a smaller trainer-level integration surface |

The July Google post contrasts Tunix with older characterizations of its peers. Current official materials matter more: veRL documents multi-turn tool calling, OpenRLHF documents multi-turn and async agent workflows, and TRL's current GRPO surface includes environment-owned rewards. Tunix's differentiator is therefore not "only framework with agentic RL." It is the combination of agentic abstractions with JAX-native training and TPU placement.

## Failure modes and guardrails

### Concurrency overwhelms the environment

`max_concurrency` limits active episodes, not every downstream resource. Add separate semaphores and rate limits for model requests, sandboxes, databases, judges, and external APIs. Measure queue wait and service saturation before raising concurrency.

### A group never completes

Enforce episode timeouts, maximum steps, maximum response tokens, and explicit cancellation. Track incomplete group age and define whether a failed member invalidates, retries, or replaces the group. Silent replacement can bias the training distribution.

### Rollout and trainer policies drift

Keep `off_policy_steps=0` for the baseline. Log each trajectory's policy version and reject stale examples according to a documented rule. If allowing older trajectories improves utilization, compare reward, KL, clipping, and reproducibility against the strict baseline.

### Weight sync becomes the hidden barrier

Trace sync duration and frequency. LoRA can reduce transfer volume because Tunix synchronizes only adapter parameters when enabled, but that benefit depends on the actual model and engine path. Verify the rollout output after every sync strategy change.

### Token boundaries change across turns

Use the model's intended chat parser and tokenizer, preserve tool-message boundaries, and test round trips with real multi-turn transcripts. Reward functions must score the intended assistant action rather than tool output or formatting artifacts.

### Restart restores weights but not the whole pipeline

Tunix documents Orbax-based model, optimizer, and step checkpointing when `checkpoint_root_directory` is configured. Agentic jobs also have queues, environment state, external sandboxes, policy versions, and side effects. Run an interruption test and verify exactly which work replays, disappears, or duplicates.

## Practical test plan

Use a small model and a synthetic environment whose latency distribution you control.

1. **Correctness baseline:** one concurrent episode, vanilla rollout, colocated mesh, fixed seed, short prompts, and deterministic tools. Verify trajectory text, masks, group IDs, rewards, advantages, and policy versions.
2. **Straggler experiment:** inject 50 ms, 500 ms, and 5 s environment delays. Compare synchronous-looking concurrency of one with increasing `max_concurrency`. Record complete groups per minute, not just raw requests.
3. **Grouping test:** use four generations per prompt and force one member to time out. Confirm group release, retry, and cleanup behavior match the policy.
4. **Topology test:** compare colocated and disaggregated meshes with the same total device budget when possible. Record rollout rate, trainer rate, sync time, HBM, and end-to-end step time.
5. **Engine test:** establish vanilla correctness first, then add vLLM or SGLang-JAX at a pinned revision. Compare tokens, stop conditions, log probabilities, mapping, and post-sync outputs before comparing speed.
6. **Backpressure test:** deliberately slow the trainer, then the environment. Set alerts for ready queue growth, half-filled buckets, oldest group age, and host memory.
7. **Recovery test:** interrupt during rollout, training, checkpoint write, and weight sync. Verify restart step, duplicated external actions, stale trajectories, and resource cleanup.
8. **Quality gate:** run a fixed held-out task set after each throughput change. Reject a performance gain that changes reward distribution, truncation rate, policy age, or solve quality beyond the agreed tolerance.

Google has not published a standardized Tunix-versus-veRL/OpenRLHF/TRL end-to-end benchmark for this release. A fair internal comparison must hold model, prompts, generation settings, algorithm, group size, hardware class, environment latency, and quality threshold constant.

## Adoption checklist

Before a long-running job, confirm:

- Python, JAX, Tunix, rollout engine, model, tokenizer, and requirement revisions are pinned.
- The selected model has tested trainer-to-rollout parameter mappings.
- Prompt length, response length, KV cache, HBM fractions, and mesh sizes fit worst-case inputs.
- Environment calls have limits, cancellation, idempotency, least privilege, and complete cleanup.
- Group timeout and retry behavior cannot silently skew the dataset.
- Policy version, stale-example handling, and weight-sync cadence are visible.
- Checkpoint and restart behavior has been exercised under interruption.
- Perfetto traces and throughput/quality counters cover rollout, environment, queues, training, and sync.
- A held-out evaluation blocks throughput changes that damage task quality.
- Operators can stop the job and revoke environment credentials without waiting for a clean training step.

For application-facing observability patterns, AgentPedia's [OpenAI Agents Python guide](/blog/openai-agents-python-guide) is useful companion reading. Its SDK differs from Tunix, but the same principle applies: every tool call and handoff needs inspectable boundaries.

## Use it or skip it

**Use Tunix as the leading candidate** when your training stack is already JAX/Flax on TPU, multi-turn tool latency is the bottleneck, and your team can validate an actively changing Alpha-class package. Its agent/environment boundary, group-aware queues, mesh placement, and RL-specific traces address real systems problems.

**Run a bounded comparison first** when your production stack is PyTorch/GPU or already standardized on veRL, OpenRLHF, or TRL. Migration cost, model support, operational familiarity, and engine compatibility can outweigh an architectural feature list.

**Wait or contribute upstream** if you need a stable API contract, a supported accelerated-engine combination not covered by current pinned requirements, independently reproduced throughput results, or complete queue/environment restart semantics out of the box.

The most defensible claim is not that Tunix eliminates idle time. It gives JAX/TPU teams the controls and traces needed to find where idle time moves, overlap work where hardware permits, and measure whether the change improves useful training throughput.

## FAQ

### What is Google Tunix?

Tunix is Google's Apache-2.0, JAX-based library for LLM post-training. It supports supervised fine-tuning, reinforcement learning and agentic RL, with TPU-oriented rollout, training and weight-synchronization paths.

### Does Tunix eliminate every synchronization barrier in agentic RL?

No. Tunix removes the whole-batch rollout barrier by streaming completed trajectories and can overlap rollout with training on separate meshes. GRPO still needs a complete trajectory group for each prompt, and rollouts pause when weight synchronization takes the rollout lock.

### Does Google prove that Tunix has near-zero TPU idle time?

Google uses near-zero idle time as an architecture claim in its July 21, 2026 launch article, but it does not publish a standardized end-to-end benchmark against veRL, OpenRLHF or TRL. Teams should verify utilization and throughput with their own environment-latency distribution.

### Which rollout engines does Tunix support?

Current Tunix documentation and source support its vanilla sampler plus vLLM on TPU and SGLang-JAX integrations. The accelerated engines require separately installed, version-sensitive dependencies.

### Is Tunix production-ready?

The repository describes a V2 release and active development, while the current main branch identifies google-tunix 0.1.7 as Alpha, matching the latest GitHub release, v0.1.7. Treat APIs and dependency pins as changeable, pin a reviewed revision and validate restart, observability and numerical behavior before production use.

### When should a team choose Tunix over veRL, OpenRLHF or TRL?

Tunix is the clearest fit when the team already uses JAX, Flax and TPU infrastructure and needs multi-turn environments in the same stack. veRL and OpenRLHF are stronger defaults for established PyTorch GPU and Ray-based deployments, while TRL offers a higher-level Transformers and Accelerate workflow with current multi-environment agentic GRPO support.


---

[Join the Agentpedia newsletter](https://agentpedia.codes/blog)

[Browse related Agentpedia articles](https://agentpedia.codes/blog)

## Official sources

### Google and Tunix

- [Google Developers Blog: Scaling Agentic RL with Tunix](https://developers.googleblog.com/scaling-agentic-rl-high-throughput-agentic-training-with-tunix/), July 21, 2026
- [Tunix repository and README](https://github.com/google/tunix), [package metadata](https://github.com/google/tunix/blob/main/pyproject.toml), and [releases](https://github.com/google/tunix/releases)
- [Tunix Agentic RL documentation](https://tunix.readthedocs.io/en/latest/agentic_rl.html), [rollout engines](https://tunix.readthedocs.io/en/latest/rollout.html), [performance](https://tunix.readthedocs.io/en/latest/performance.html), and [reliability](https://tunix.readthedocs.io/en/latest/reliability.html)
- [Agentic learner source](https://github.com/google/tunix/blob/main/tunix/rl/agentic/agentic_rl_learner.py), [rollout orchestrator](https://github.com/google/tunix/blob/main/tunix/rl/agentic/pipeline/rollout_orchestrator.py), and [trajectory engine](https://github.com/google/tunix/blob/main/tunix/rl/agentic/trajectory/trajectory_collect_engine.py)
- [Agentic Gemma GRPO example](https://github.com/google/tunix/blob/main/examples/agentic/gemma_grpo_demo_nb.py) and [Llama/Qwen GRPO script](https://github.com/google/tunix/blob/main/scripts/grpo_demo_llama3_qwen2.py)

### Comparison projects

- [veRL repository and current feature list](https://github.com/verl-project/verl)
- [OpenRLHF repository, architecture, and multi-turn agent documentation](https://github.com/OpenRLHF/OpenRLHF)
- [Hugging Face TRL GRPOTrainer](https://huggingface.co/docs/trl/main/en/grpo_trainer) and [distributed training guide](https://huggingface.co/docs/trl/main/en/distributing_training)


---

- [All articles](https://agentpedia.codes/blog)