integrations.nemo_gym.rollout

integrations.nemo_gym.rollout

Multi-turn NeMo Gym rollouts for async GRPO.

A NeMo Gym agent server runs a whole rollout behind its /run endpoint: it generates against our vLLM server through its model proxy, executes tool calls on resource servers, verifies the outcome and returns the prompt and generated token IDs and logprobs of every model call along with the reward. NemoGymRolloutLoop runs in async GRPO’s rollout process and sends each rollout to an agent instead of generating itself. It rebuilds training rows from the returned model calls with TRL’s multi-turn reconciler and scores the agent’s reward as a reward column.

Classes

Name	Description
NemoGymRolloutLoop	Rolls out each prompt through a NeMo Gym agent’s /run (see module doc).

NemoGymRolloutLoop

integrations.nemo_gym.rollout.NemoGymRolloutLoop(
    agent_servers,
    run_timeout=300.0,
    **kwargs,
)

Rolls out each prompt through a NeMo Gym agent’s /run (see module doc).

Parameters

Name	Type	Description	Default
agent_servers	dict[str, str]	Agent name to base URL. A row’s agent_ref.name picks its agent; rows without one go to the first.	required
run_timeout	float	Seconds a /run may go without sending data.	300.0

Functions

Name	Description
agent_request	The /run body for one rollout of a NeMo Gym dataset row.
messages_from_agent_output	The rollout as chat messages, for reward functions and completion logs.
turns_from_agent_output	One turn per model call: NeMo Gym puts a call’s token IDs on its last item.

agent_request

integrations.nemo_gym.rollout.agent_request(
    row,
    *,
    max_tokens,
    temperature,
    top_p,
)

The /run body for one rollout of a NeMo Gym dataset row.

messages_from_agent_output

integrations.nemo_gym.rollout.messages_from_agent_output(output)

The rollout as chat messages, for reward functions and completion logs.

turns_from_agent_output

integrations.nemo_gym.rollout.turns_from_agent_output(output)

One turn per model call: NeMo Gym puts a call’s token IDs on its last item.