What does a trajectory collect engine do in agentic RL?

It owns one episode. Given an agent and an environment, it resets both, then loops: send the conversation so far to the model, let the agent parse the reply into an action, hand that action to the environment, and append the environment’s observation back into the conversation. The loop ends when the environment says done, when the step budget runs out, when the response token budget runs out, or when a deadline fires. After the loop it adds the environment’s final reward to the last step, computes discounted Monte Carlo returns backwards over the steps, sums the undiscounted rewards into a trajectory reward, and closes the environment. It returns the episode in one of four shapes: the full Trajectory object, per-step dicts, a flat token dict for training, or the raw conversation messages.

How does Tunix run many agentic rollouts at once?

Two ways, at two levels. Inside the engine file, the static method collect_multiple takes a list of agent and environment pairs, builds one engine per pair, launches them all as asyncio tasks, and yields each result through asyncio.as_completed, so a finished episode is handed back the moment it finishes rather than at the end of the batch. There is no concurrency cap in that method. The pipeline used in real training does cap it: RolloutOrchestrator in tunix/rl/agentic/pipeline/rollout_orchestrator.py keeps at most max_concurrency episodes alive, and each one takes a shared rollout lock for its whole duration so a weight sync can wait for a clean boundary instead of interrupting a live episode.

What happens if a tool call never returns?

The environment step is the only call in the file with a deadline on it. The engine computes remaining_time as the episode timeout minus the time already spent, then runs env.step in a thread pool under asyncio.wait_for with that value. When the tool blocks past it, wait_for raises asyncio.TimeoutError, the engine sets the trajectory status to ENV_TIMEOUT, logs which step hung, marks the current step done and returns, so collect breaks out of its loop and still computes rewards and closes the environment. The episode is unblocked but the thread is not: cancelling the future that wraps a running executor callable does not stop the callable, so the hung tool keeps holding a worker in a pool whose default size is capped at min(32, cpu_count + 4). The file admits the same failure for env.close, whose own 150-second timeout logs that the executor thread may be leaked and that this will starve the thread pool over time.

Which tokens in an agentic trajectory carry gradient?

Only the ones the policy generated. In Token mode the engine appends the assistant tokens with a mask of ones over the sampled tokens and zeros over any end-of-turn tokens the chat parser had to append, then appends the environment tokens with an all-zero mask, so tool output is present in the sequence for context but contributes nothing to the loss. Two details are easy to miss. Tokens, masks and logprobs are appended in lockstep, with zeros filled in for logprobs the sampler did not return, because a short logprob array would offset every step after it. And the environment message of a terminal step is never tokenized at all, guarded by not done and not step_timed_out, since nothing follows it for the policy to condition on.

What is the difference between TIMEOUT and ENV_TIMEOUT?

TIMEOUT means the episode as a whole outlived its wall clock: after a step completes, the engine compares elapsed time against the timeout it was constructed with, and if the step ran past it the trajectory is marked TIMEOUT and the loop stops. ENV_TIMEOUT is narrower and means one environment step hung and was killed by asyncio.wait_for before it returned. Both sit in the default filter_statuses set alongside MAX_STEPS_REACHED and MAX_CONTEXT_LIMIT_REACHED, so when overlong_filter is on, a trajectory that ended any of those four ways gets an all-zero token mask and trains on nothing while its conversation and reward are still reported.

The Trajectory Collect Engine, Line by Line

A normal RL rollout is one prompt in, one completion out. An agentic rollout is a conversation: the model writes a tool call, an environment runs it, the result comes back as a message, and the model writes again, for however many turns the task takes. This walk reads the 763-line file in Google’s tunix that owns exactly one of those episodes: how it counts steps, where the only deadline in the file actually sits, which tokens end up carrying gradient, and how dozens of these run at once and stream their finished trajectories to a learner that never stops training.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of The Trajectory Collect Engine, Line by Line — built to be understood, not skimmed.

Questions

What does a trajectory collect engine do in agentic RL?: It owns one episode. Given an agent and an environment, it resets both, then loops: send the conversation so far to the model, let the agent parse the reply into an action, hand that action to the environment, and append the environment’s observation back into the conversation. The loop ends when the environment says done, when the step budget runs out, when the response token budget runs out, or when a deadline fires. After the loop it adds the environment’s final reward to the last step, computes discounted Monte Carlo returns backwards over the steps, sums the undiscounted rewards into a trajectory reward, and closes the environment. It returns the episode in one of four shapes: the full Trajectory object, per-step dicts, a flat token dict for training, or the raw conversation messages.
How does Tunix run many agentic rollouts at once?: Two ways, at two levels. Inside the engine file, the static method collect_multiple takes a list of agent and environment pairs, builds one engine per pair, launches them all as asyncio tasks, and yields each result through asyncio.as_completed, so a finished episode is handed back the moment it finishes rather than at the end of the batch. There is no concurrency cap in that method. The pipeline used in real training does cap it: RolloutOrchestrator in tunix/rl/agentic/pipeline/rollout_orchestrator.py keeps at most max_concurrency episodes alive, and each one takes a shared rollout lock for its whole duration so a weight sync can wait for a clean boundary instead of interrupting a live episode.
What happens if a tool call never returns?: The environment step is the only call in the file with a deadline on it. The engine computes remaining_time as the episode timeout minus the time already spent, then runs env.step in a thread pool under asyncio.wait_for with that value. When the tool blocks past it, wait_for raises asyncio.TimeoutError, the engine sets the trajectory status to ENV_TIMEOUT, logs which step hung, marks the current step done and returns, so collect breaks out of its loop and still computes rewards and closes the environment. The episode is unblocked but the thread is not: cancelling the future that wraps a running executor callable does not stop the callable, so the hung tool keeps holding a worker in a pool whose default size is capped at min(32, cpu_count + 4). The file admits the same failure for env.close, whose own 150-second timeout logs that the executor thread may be leaked and that this will starve the thread pool over time.
Which tokens in an agentic trajectory carry gradient?: Only the ones the policy generated. In Token mode the engine appends the assistant tokens with a mask of ones over the sampled tokens and zeros over any end-of-turn tokens the chat parser had to append, then appends the environment tokens with an all-zero mask, so tool output is present in the sequence for context but contributes nothing to the loss. Two details are easy to miss. Tokens, masks and logprobs are appended in lockstep, with zeros filled in for logprobs the sampler did not return, because a short logprob array would offset every step after it. And the environment message of a terminal step is never tokenized at all, guarded by not done and not step_timed_out, since nothing follows it for the policy to condition on.
What is the difference between TIMEOUT and ENV_TIMEOUT?: TIMEOUT means the episode as a whole outlived its wall clock: after a step completes, the engine compares elapsed time against the timeout it was constructed with, and if the step ran past it the trajectory is marked TIMEOUT and the loop stops. ENV_TIMEOUT is narrower and means one environment step hung and was killed by asyncio.wait_for before it returned. Both sit in the default filter_statuses set alongside MAX_STEPS_REACHED and MAX_CONTEXT_LIMIT_REACHED, so when overlong_filter is on, a trajectory that ended any of those four ways gets an all-zero token mask and trains on nothing while its conversation and reward are still reported.

The Trajectory Collect Engine, Line by Line

Questions

Related explainers