`nemo_rl.experience.rollouts`#

Module Contents#

Classes#

`EffortLevelsConfig`	Controls length-based reward shaping for low-effort prompts.
`_EffortShapingMetrics`
`RolloutGroupResult`	One prompt group’s rollout batch and metrics.
`NemoGymRolloutResult`	Processed NeMo-Gym rollouts for one prompt group or synchronous batch.
`_CompletedNemoGymGroup`	One complete Gym prompt group restored to input-row order.
`_NemoGymStreamAccumulator`	Validate streamed Gym rows and assemble complete prompt groups.

Functions#

`_add_r3_fallback_metrics`
`_extract_mask_sample_flags`	Return True for samples the environment asks GRPO to mask from loss.
`_attach_routed_experts_to_message_log_prefix`	Attach routed-expert slices to existing messages and return prefix length.
`_find_routed_experts_template`
`_dummy_routed_experts_for_tokens`
`backfill_missing_routed_experts`	Give every tokenized message a `routed_experts` row, in place.
`_apply_effort_shaping`	Apply length-based reward shaping for low-effort prompts.
`generate_responses`	Generate responses from policy using synchronous generation.
`generate_responses_async`	Async version of generate_responses that properly calls generate_async.
`calculate_rewards`	Calculate rewards for generated responses and get environment feedback.
`run_multi_turn_rollout`	Runs a multi-turn rollout loop, interacting with the environment.
`async_generate_response_for_sample_turn`	Generate a response for a single sample’s turn using async generation.
`run_sample_multi_turn_rollout`	Run a multi-turn rollout for a single sample.
`_aggregate_multi_turn_rollout_metrics`	Aggregate native rollout metrics over an arbitrary set of samples.
`_run_multi_turn_rollout_async`	Run one native rollout batch and retain metrics at sample granularity.
`run_async_multi_turn_rollout`	Run a complete native rollout batch from a synchronous call site.
`run_async_multi_turn_rollout_groups`	Run one native batch, then yield prompt groups with group-local metrics.
`_tensorize_by_key`
`get_nemo_gym_thinking_tags`	Return thinking tags used by the Gym-side detector.
`should_mask_flagged_samples`	Read `env.should_mask_flagged_samples`; absent means True.
`_get_reward_penalty_config_value`
`_get_reward_penalty_token_id`
`_get_required_reward_penalty_token_id`
`_get_reward_penalty_token_ids`
`_get_required_reward_penalty_token_ids`
`_infer_single_token_id`
`resolve_reward_penalty_config`	Resolve tokenizer-derived reward penalty fields.
`apply_reward_penalties`	Apply reward penalties to results, setting reward to 0.0 when triggered.
`_prepare_nemo_gym_rows`	Apply NeMo-RL sampling parameters and stable row indices in place.
`_tensorize_nemo_gym_result`	Convert token fields returned by the Gym actor back to tensors.
`run_async_nemo_gym_rollout`	Stream complete NeMo-Gym prompt groups in group-completion order.
`run_nemo_gym_rollout_sync`	Run and return one complete NeMo-Gym batch synchronously.
`_postprocess_single_nemo_gym_group`	Postprocess one complete prompt group from the NeMo-Gym stream.

Data#

TokenizerType

API#

nemo_rl.experience.rollouts.TokenizerType#: None

nemo_rl.experience.rollouts._add_r3_fallback_metrics( gen_metrics: dict[str, float | int], generation_outputs: nemo_rl.distributed.batched_data_dict.BatchedDataDict, ) → None#

nemo_rl.experience.rollouts._extract_mask_sample_flags( results: list[dict[str, Any]], ) → torch.Tensor#: Return True for samples the environment asks GRPO to mask from loss.

nemo_rl.experience.rollouts._attach_routed_experts_to_message_log_prefix( message_log: list[dict], routed_experts: torch.Tensor, ) → int#: Attach routed-expert slices to existing messages and return prefix length.

nemo_rl.experience.rollouts._find_routed_experts_template( message_log: list[dict], ) → Optional[torch.Tensor]#

nemo_rl.experience.rollouts._dummy_routed_experts_for_tokens( token_ids: torch.Tensor, template: torch.Tensor, ) → torch.Tensor#

nemo_rl.experience.rollouts.backfill_missing_routed_experts( message_logs: collections.abc.Sequence[list[dict]], ) → None#

Give every tokenized message a routed_experts row, in place.

Routes are attached only where generation ran, so a trajectory whose first turn raised (or a turn whose routes vLLM could not return) leaves messages without the field while its siblings have it. Flattening then either stacks ragged ranks or silently concatenates a short column, so fill the gaps with the all–1 missing-route sentinel: Megatron routes those tokens with its own router, which is the honest answer for tokens no capture covered.

No-op when the batch carries no routes at all — that is the router-replay-off case, and on the TQ paths the producer-side guard must still see the field missing so it can report a capture failure.

class nemo_rl.experience.rollouts.EffortLevelsConfig#

Bases: pydantic.BaseModel

Controls length-based reward shaping for low-effort prompts.

When a prompt contains low_string, the final reward is adjusted by a length-reward term that penalises overly long responses. The reward formula is::

length_reward = min(1, low_weight * (1 - response_len / low_ub))
new_reward    = orig_reward
              + orig_reward * max(length_reward, 0)
              + low_penalty * min(length_reward, 0)

Setting low_weight = 0 or leaving low_string empty disables the shaping entirely.

low_weight: float#

0.0

Weight applied to the length-reward term. Set to 0 to disable.

low_penalty: float#

1.0

Coefficient for the negative length-reward penalty.

low_ub: int#

64000

Response-length upper bound (in tokens) used to normalise the term.

low_string: str = <Multiline-String>#: Substring that must appear in the user prompt to trigger shaping.

class nemo_rl.experience.rollouts._EffortShapingMetrics#

length_rewards_low: list[float]#: None

rewards_low: list[float]#: None

low_lengths: list[int]#: None

high_lengths: list[int]#: None

nemo_rl.experience.rollouts._apply_effort_shaping( results: list[dict], nemo_gym_rows: list[dict], effort_config: Optional[nemo_rl.experience.rollouts.EffortLevelsConfig], ) → nemo_rl.experience.rollouts._EffortShapingMetrics#

Apply length-based reward shaping for low-effort prompts.

Modifies results[i]["full_result"]["reward"] in place for samples whose last user-turn prompt contains effort_config.low_string. Returns per-sample tracking lists used to populate rollout metrics.

No-ops (returns empty lists) when effort_config is None, low_weight is zero, or low_string is empty.

nemo_rl.experience.rollouts.generate_responses( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, generation_input_data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec], batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], tokenizer: nemo_rl.experience.rollouts.TokenizerType, input_lengths: torch.Tensor, include_logprobs: bool = True, greedy: bool = False, ) → tuple[nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], list[torch.Tensor], dict[str, float | int]]#: Generate responses from policy using synchronous generation.

async nemo_rl.experience.rollouts.generate_responses_async( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, generation_input_data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec], batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], tokenizer: nemo_rl.experience.rollouts.TokenizerType, input_lengths: torch.Tensor, include_logprobs: bool = True, greedy: bool = False, ) → tuple[nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], list[torch.Tensor], dict[str, float | int]]#: Async version of generate_responses that properly calls generate_async.

nemo_rl.experience.rollouts.calculate_rewards( batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], ) → nemo_rl.environments.interfaces.EnvironmentReturn#

Calculate rewards for generated responses and get environment feedback.

Parameters:

batch – Batch containing message_log (LLMMessageLogType) with generated responses
task_to_env – Dictionary mapping task names to their corresponding environments

Returns:

observations: List of observations from the environment for the next turn.
metadata: List of extracted metadata from the environment.
next_stop_strings: List of stop strings for the next generation step.
rewards: Tensor of rewards for the last turn.
terminateds: Tensor of booleans indicating if an episode ended naturally.

Return type:

EnvironmentReturn namedtuple containing

nemo_rl.experience.rollouts.run_multi_turn_rollout( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, input_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], tokenizer: nemo_rl.experience.rollouts.TokenizerType, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, ) → tuple[nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], dict[str, Any]]#

Runs a multi-turn rollout loop, interacting with the environment.

Parameters:

policy_generation – The generation interface (policy).
input_batch – The starting batch containing initial message logs.
tokenizer – The tokenizer.
task_to_env – Dictionary mapping task names to environment instances.
max_rollout_turns – Maximum number of agent-environment interaction turns.
max_seq_len – Maximum sequence length allowed.
greedy – Whether to use greedy decoding.

Returns:

BatchedDataDict with the full interaction history and accumulated rewards
Dictionary of rollout metrics

Return type:

Tuple containing

async nemo_rl.experience.rollouts.async_generate_response_for_sample_turn( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, sample_message_log: list[dict], sample_stop_strings: list[str] | None, tokenizer: nemo_rl.experience.rollouts.TokenizerType, max_seq_len: int, greedy: bool = False, ) → tuple[list[dict], torch.Tensor, torch.Tensor, dict[str, float]]#

Generate a response for a single sample’s turn using async generation.

Parameters:

policy_generation – The generation interface to use
sample_message_log – Message log for a single sample
sample_stop_strings – Stop strings for this sample
tokenizer – Tokenizer to use
max_seq_len – Maximum sequence length
greedy – Whether to use greedy decoding

Returns:

Tuple of (updated_message_log, generated_tokens, input_lengths, generation_metrics)

async nemo_rl.experience.rollouts.run_sample_multi_turn_rollout( sample_idx: int, initial_sample_state: dict, policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, tokenizer: nemo_rl.experience.rollouts.TokenizerType, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, ) → tuple[dict, dict[str, Any]]#

Run a multi-turn rollout for a single sample.

This function manages the complete lifecycle of one sample’s interaction. Async generation is used internally when available.

Parameters:

sample_idx – Index of this sample in the original batch
initial_sample_state – Initial state containing message_log, extra_env_info, etc.
policy_generation – The generation interface
tokenizer – Tokenizer to use
task_to_env – Environment mapping
max_seq_len – Maximum sequence length
max_rollout_turns – Maximum number of turns
greedy – Whether to use greedy decoding

Returns:

Tuple of (final_sample_state, sample_metrics)

class nemo_rl.experience.rollouts.RolloutGroupResult#

One prompt group’s rollout batch and metrics.

group_index: int#: None

final_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec]#: None

rollout_metrics: dict[str, Any]#: None

task_index: Optional[int]#: None

nemo_rl.experience.rollouts._aggregate_multi_turn_rollout_metrics( all_sample_metrics: collections.abc.Sequence[dict[str, Any]], ) → dict[str, Any]#: Aggregate native rollout metrics over an arbitrary set of samples.

async nemo_rl.experience.rollouts._run_multi_turn_rollout_async( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, input_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], tokenizer: nemo_rl.experience.rollouts.TokenizerType, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, ) → tuple[nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], list[dict[str, Any]]]#: Run one native rollout batch and retain metrics at sample granularity.

nemo_rl.experience.rollouts.run_async_multi_turn_rollout( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, input_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], tokenizer: nemo_rl.experience.rollouts.TokenizerType, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], max_seq_len: int, max_rollout_turns: int = 999999, greedy: bool = False, ) → tuple[nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], dict[str, Any]]#

Run a complete native rollout batch from a synchronous call site.

Each sample proceeds through its interaction independently. Generation is asynchronous internally, while this compatibility API returns only after the full batch and its aggregate metrics are ready.

Parameters:

policy_generation – Generation interface used to produce policy responses.
input_batch – Batch containing the initial message logs and environment data.
tokenizer – Tokenizer used to encode and decode rollout messages.
task_to_env – Mapping from task names to their environment implementations.
max_seq_len – Maximum total token length for each rollout sample.
max_rollout_turns – Maximum number of agent-environment interaction turns.
greedy – Whether policy generation should use greedy decoding.

Returns:

A tuple containing the completed rollout batch and metrics aggregated over every sample in that batch.

Raises:

RuntimeError – If an individual sample rollout fails.

async nemo_rl.experience.rollouts.run_async_multi_turn_rollout_groups( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, input_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], tokenizer: nemo_rl.experience.rollouts.TokenizerType, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], max_seq_len: int, num_generations: int, max_rollout_turns: int = 999999, greedy: bool = False, ) → collections.abc.AsyncGenerator[nemo_rl.experience.rollouts.RolloutGroupResult, None]#

Run one native batch, then yield prompt groups with group-local metrics.

This intentionally retains the native path’s full-batch completion barrier. The group iterator gives the collector a common interface with NeMo-Gym without changing native rollout scheduling semantics.

Parameters:

policy_generation – Generation interface used to produce policy responses.
input_batch – Batch containing prompts repeated contiguously by group.
tokenizer – Tokenizer used to encode and decode rollout messages.
task_to_env – Mapping from task names to their environment implementations.
max_seq_len – Maximum total token length for each rollout sample.
num_generations – Number of contiguous rollout samples in each prompt group.
max_rollout_turns – Maximum number of agent-environment interaction turns.
greedy – Whether policy generation should use greedy decoding.

Yields:

Complete prompt groups in input order. Each RolloutGroupResult contains exactly num_generations samples and metrics aggregated only over those samples.

Raises:

ValueError – If num_generations is not positive or the batch size is not divisible by num_generations.
RuntimeError – If an individual sample rollout fails.

nemo_rl.experience.rollouts._tensorize_by_key(message_logs: list, key: str)#

class nemo_rl.experience.rollouts.NemoGymRolloutResult#

Processed NeMo-Gym rollouts for one prompt group or synchronous batch.

input_ids: torch.Tensor#: None

final_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec]#: None

rollout_metrics: dict[str, Any]#: None

task_index: Optional[int]#: None

class nemo_rl.experience.rollouts._CompletedNemoGymGroup#

One complete Gym prompt group restored to input-row order.

group_index: int#: None

rows: list[dict]#: None

results: list[dict]#: None

class nemo_rl.experience.rollouts._NemoGymStreamAccumulator( rows: list[dict], num_generations: int, allow_mixed_agents: bool, )#

Validate streamed Gym rows and assemble complete prompt groups.

NeMo Gym returns rows in completion order. This accumulator owns all ordering and completeness rules so the rollout loop only needs to postprocess completed groups.

Initialization

property is_complete: bool#

add( row_index: int, result: dict, ) → nemo_rl.experience.rollouts._CompletedNemoGymGroup | None#: Add one streamed row and return its group when that group is complete.

finish() → None#: Raise when the stream ended before every expected row arrived.

nemo_rl.experience.rollouts.get_nemo_gym_thinking_tags( env_config: dict[str, Any], ) → list[str]#: Return thinking tags used by the Gym-side detector.

nemo_rl.experience.rollouts.should_mask_flagged_samples(env_config: dict[str, Any]) → bool#

Read env.should_mask_flagged_samples; absent means True.

True (the default): env-driven mask_sample flags are carried in the rollout batch and flagged samples are dropped from the loss.

Set false when the flags are too coarse to honor: for example, Gym flags rollouts that hit max iterations even when they solve the task, and those are samples worth training on. It also keeps batch composition deterministic for controlled benchmark runs — how many samples get flagged varies run to run.

nemo_rl.experience.rollouts._get_reward_penalty_config_value( reward_penalty_config: dict[str, Any] | pydantic.BaseModel | None, key: str, ) → Any#

nemo_rl.experience.rollouts._get_reward_penalty_token_id( reward_penalty_config: dict[str, Any] | pydantic.BaseModel, key: str, ) → int | None#

nemo_rl.experience.rollouts._get_required_reward_penalty_token_id( reward_penalty_config: dict[str, Any] | pydantic.BaseModel, key: str, ) → int#

nemo_rl.experience.rollouts._get_reward_penalty_token_ids( reward_penalty_config: dict[str, Any] | pydantic.BaseModel, key: str, ) → list[int] | None#

nemo_rl.experience.rollouts._get_required_reward_penalty_token_ids( reward_penalty_config: dict[str, Any] | pydantic.BaseModel, key: str, ) → list[int]#

nemo_rl.experience.rollouts._infer_single_token_id(tokenizer: Any, text: str) → int | None#

Resolve tokenizer-derived reward penalty fields.

User config must explicitly provide unwanted token IDs when penalize_unwanted_tokens is enabled. Think-tag IDs are inferred only when each configured tag is exactly one token.

nemo_rl.experience.rollouts.apply_reward_penalties( results: list[dict], reward_penalty_config: dict[str, Any] | pydantic.BaseModel | None, ) → dict[str, int]#

Apply reward penalties to results, setting reward to 0.0 when triggered.

All penalties are gated by reward_penalty_config flags. Returns a dict of penalty counts keyed by penalty name.

NOTE: These penalties assume Gym-path message_log structure where roles strictly alternate “user” → “assistant”. Tool responses are folded into user prompt tokens by _postprocess_nemo_gym_to_nemo_rl_result and never appear as separate message_log entries. Do not call from non-Gym rollout paths.

Penalties:

penalize_duplicated_reasoning (text-based) Checks response[“output”] items. If a “reasoning” item’s summary text exactly matches the next item’s content text (after strip), the model is copying its thinking into the final answer verbatim. Data: full_result[“response”][“output”] — reasoning has summary[0][“text”], message has content[0][“text”].
penalize_empty_final_answer (text-based) Walks response[“output”] in reverse to find the last message-type item. If no message item exists or its content text is empty, the model failed to produce a final answer. Skipped when the last output item is a function_call (model was mid-agentic-loop, not producing an empty answer). Data: full_result[“response”][“output”] — message items have content[0][“text”].
penalize_unwanted_tokens (token-based) Currently checks that none of the explicitly configured unwanted token IDs appear anywhere in an assistant generation, including as the terminal token. A turn may contain multiple unwanted tokens, so the whole assistant token sequence is checked rather than excluding the trailing position. Data: message_log[i][“token_ids”] where role == “assistant”.
penalize_malformed_think_tag (message flag + token/string fallback) Three complementary checks to catch malformed think tags: a) Existing Gym flag: honors assistant message has_malformed_thinking. b) Token ID check: when think tag IDs are resolved from config override or single-token tokenizer encodings, infers thinking mode from prompt token counts. If prompt has open==close: enable_thinking=False, expect 0 open and 0 close in generation. If prompt has open==close+1: enable_thinking=True, expect 0 open and 1 close in generation. Any other prompt pattern or mismatched generation counts is a violation. This fallback is skipped when the tags do not resolve to one token each. c) String check: the model can spell out thinking tags with piecemeal regular tokens (e.g. “<”, “/”, “thi”, “nk”, “>”) that bypass special token IDs. Checks generation_str (decoded generation text) per output item: open-tag count must be 0 (always in prompt, never generated), close-tag count must be 0 or 1. Data: message_log pairs for token IDs, full_result output items for strings.

nemo_rl.experience.rollouts._prepare_nemo_gym_rows( rows: list[dict], generation_config: nemo_rl.models.generation.interfaces.GenerationConfig, sampling_params: nemo_rl.models.generation.interfaces.GenerationSamplingParams, ) → None#: Apply NeMo-RL sampling parameters and stable row indices in place.

nemo_rl.experience.rollouts._tensorize_nemo_gym_result(result: dict) → None#: Convert token fields returned by the Gym actor back to tensors.

async nemo_rl.experience.rollouts.run_async_nemo_gym_rollout( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, input_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], tokenizer: nemo_rl.experience.rollouts.TokenizerType, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], generation_config: nemo_rl.models.generation.interfaces.GenerationConfig, num_generations: int, log_full_result_tables: bool, max_seq_len: Optional[int] = None, max_rollout_turns: Optional[int] = None, greedy: bool = False, effort_config: Optional[nemo_rl.experience.rollouts.EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | pydantic.BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, returns_entire_batch: bool = False, sampling_params: Optional[nemo_rl.models.generation.interfaces.GenerationSamplingParams] = None, ) → collections.abc.AsyncGenerator[nemo_rl.experience.rollouts.NemoGymRolloutResult, None]#

Stream complete NeMo-Gym prompt groups in group-completion order.

The actor streams individual rows in arbitrary completion order. Rows are validated and restored to input order within each num_generations group before the group is postprocessed and yielded. Synchronous call sites should use :func:run_nemo_gym_rollout_sync.

Parameters:

policy_generation – Generation interface whose configuration supplies the model’s maximum sequence length.
input_batch – Batch whose extra_env_info field contains NeMo-Gym rows.
tokenizer – Tokenizer used by the NeMo-Gym actor and local postprocessing.
task_to_env – Environment mapping containing the "nemo_gym" actor.
generation_config – Sampling parameters forwarded to every NeMo-Gym row.
num_generations – Number of contiguous rows belonging to each prompt group.
log_full_result_tables – Whether to include complete per-agent result payloads as W&B Tables in the rollout metrics.
max_seq_len – Policy sequence-length limit used for compatibility validation. NeMo-Gym still relies on the generation engine’s configured limit.
max_rollout_turns – Must be None because NeMo-Gym owns turn limits.
greedy – Must be False because this path does not support greedy mode.
effort_config – Optional configuration for effort-based reward shaping.
reward_penalty_config – Optional reward-penalty configuration.
thinking_tags – Optional opening and closing tags used by thinking penalties.
mask_env_flagged_samples – Whether to carry env-driven mask_sample flags in the rollout batch for loss masking.
returns_entire_batch – Whether to treat the input as one potentially heterogeneous group. This requires num_generations to equal the batch size and is used by synchronous callers.
sampling_params – Sampling profile stamped onto every NeMo-Gym row. None uses the train profile from generation_config; validation passes its own profile explicitly.

Yields:

NemoGymRolloutResult objects in prompt-group completion order. Rows inside each result are restored to input order. The final result also carries actor-wide and rollout-wide timing metrics.

Raises:

AssertionError – If an unsupported generation option is requested.
TypeError – If a row lacks a valid responses_create_params dictionary or the actor returns a non-integer row index.
ValueError – If num_generations is not positive, the batch is empty or not divisible by num_generations, returns_entire_batch has an incompatible size, a streamed row index is out of range or duplicated, a prompt group mixes agents, or its task indices disagree.
RuntimeError – If the actor fails, returns NaN generation logprobs, ends the stream before all expected rows arrive, or produces no final group.

nemo_rl.experience.rollouts.run_nemo_gym_rollout_sync( policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, input_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], tokenizer: nemo_rl.experience.rollouts.TokenizerType, task_to_env: dict[str, nemo_rl.environments.interfaces.EnvironmentInterface], generation_config: nemo_rl.models.generation.interfaces.GenerationConfig, log_full_result_tables: bool, max_seq_len: Optional[int] = None, max_rollout_turns: Optional[int] = None, greedy: bool = False, effort_config: Optional[nemo_rl.experience.rollouts.EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | pydantic.BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, sampling_params: Optional[nemo_rl.models.generation.interfaces.GenerationSamplingParams] = None, mask_env_flagged_samples: bool = True, ) → nemo_rl.experience.rollouts.NemoGymRolloutResult#

Run and return one complete NeMo-Gym batch synchronously.

This compatibility API drains :func:run_async_nemo_gym_rollout with the whole input treated as one heterogeneous group, restoring input order and returning only after every row is complete.

Parameters:

policy_generation – Generation interface whose configuration supplies the model’s maximum sequence length.
input_batch – Batch whose extra_env_info field contains NeMo-Gym rows.
tokenizer – Tokenizer used by the NeMo-Gym actor and local postprocessing.
task_to_env – Environment mapping containing the "nemo_gym" actor.
generation_config – Sampling parameters forwarded to every NeMo-Gym row.
log_full_result_tables – Whether to include complete per-agent result payloads as W&B Tables in the rollout metrics.
max_seq_len – Policy sequence-length limit used for compatibility validation.
max_rollout_turns – Must be None because NeMo-Gym owns turn limits.
greedy – Must be False because this path does not support greedy mode.
effort_config – Optional configuration for effort-based reward shaping.
reward_penalty_config – Optional reward-penalty configuration.
thinking_tags – Optional opening and closing tags used by thinking penalties.
sampling_params – Sampling profile stamped onto every NeMo-Gym row. None uses the train profile from generation_config; validation passes its own profile explicitly.
mask_env_flagged_samples – Whether to carry env-driven mask_sample flags in the rollout batch for loss masking.

Returns:

The fully postprocessed NeMo-Gym rollout batch in input-row order.

Raises:

AssertionError – If an unsupported generation option is requested.
TypeError – If a NeMo-Gym row or streamed row index has an invalid type.
ValueError – If streamed rows violate the ordering, uniqueness, grouping, or task-index invariants documented by :func:run_async_nemo_gym_rollout.
RuntimeError – If called from a running event loop, the actor or stream fails, or NeMo-Gym returns no complete rollout batch.

nemo_rl.experience.rollouts._postprocess_single_nemo_gym_group( nemo_gym_rows: list[dict], results: list[dict], timer: nemo_rl.utils.timer.Timer, timer_prefix: str, policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, input_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], tokenizer: nemo_rl.experience.rollouts.TokenizerType, log_full_result_tables: bool, effort_config: Optional[nemo_rl.experience.rollouts.EffortLevelsConfig] = None, reward_penalty_config: dict[str, Any] | pydantic.BaseModel | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, ) → nemo_rl.experience.rollouts.NemoGymRolloutResult#: Postprocess one complete prompt group from the NeMo-Gym stream.

nemo_rl.experience.rollouts#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.experience.rollouts`#