`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#

`attach_initial_nemo_gym_image_payloads`	Attach initial Gym image tensors once, before prompt repeat.
`_add_multimodal_generation_payload`	Attach one policy-ready or vLLM-ready media representation.
`_reattach_original_multimodal_payloads`	Restore exact prompt media omitted by a remote Gym rollout.
`_reattach_static_multimodal_payloads_to_result`	Restore static media to each Gym-authored message-log representation.
`attach_static_multimodal_payload`	Copy policy-ready media from static source turns to Gym-authored turns.
`_add_r3_fallback_metrics`
`_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.
`get_effort_config`	Resolve the optional Gym effort config using its centralized defaults.
`_terminal_completion_length`	Return inline terminal assistant length; capture is shaped after verification.
`_apply_effort_shaping`	Apply length-based reward shaping for low-effort prompts.
`_effort_shaping_metrics`	Build the rollout-metric entries for one group’s effort-shaping lists.
`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`
`compute_reward_penalty_metrics`	Return enabled penalty rates using the legacy NeMo-Gym metric names.
`_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.
`_bucket_nemo_gym_rows_by_instance`	Split already-stamped rows into one bucket per actor, by prompt group.
`_merge_nemo_gym_instance_streams`	Interleave K actor streams into one, yielding rows as they complete.
`_merge_nemo_gym_timing_metrics`	Combine per-instance actor timings into per-instance plus rolled-up metrics.
`_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.
`_single_group_reward_diagnostics`	Summarize reward variation within one prompt’s generations.
`_group_reward_diagnostics`	Compute diagnostics over the actual prompt groups in `results`.
`_nemo_gym_group_metrics`	Compute metrics shared by both NeMo-Gym rollout implementations.
`nemo_gym_rollout_hit_max_tokens`	Return whether a NeMo-Gym rollout was cut off by the token budget.
`_postprocess_single_nemo_gym_group`	Postprocess one complete prompt group from the NeMo-Gym stream.

Data#

`TokenizerType`
`_REWARD_PENALTY_METRICS`
`_SHARED_GENRM_METRICS`
`_LENGTH_LIMIT_INCOMPLETE_REASONS`

API#

nemo_rl.experience.rollouts.TokenizerType#: None

nemo_rl.experience.rollouts._REWARD_PENALTY_METRICS#: None

nemo_rl.experience.rollouts._SHARED_GENRM_METRICS#: (‘reward_score_raw’, ‘reward_rubric_aggregate_valid’, ‘reward_overall_score_raw’, ‘reward_overall_sc…

nemo_rl.experience.rollouts.attach_initial_nemo_gym_image_payloads( batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], processor: Any, *, env_config: collections.abc.Mapping[str, Any], ) → None#

Attach initial Gym image tensors once, before prompt repeat.

The NeMo Gym dataset deliberately carries only the Responses request in extra_env_info. Dedup-enabled GRPO calls this helper on the unrepeated prompt batch, allowing repeat_interleave(..., share_immutable_media=True) to retain one physical processor output per prompt. Flag-off runs never call this helper.

Takes master_config.env and resolves pad_dynamic_image_shapes itself, mirroring the per-turn attach inside the NeMo-Gym actor. Resolving here rather than at each call site means a caller cannot supply the wrong value – the divergence between the two attach paths that this helper previously had. The flag only matters for a turn carrying more than one image at differing resolutions, where the processor returns a ragged CHW list; without it the processor is asked to stack those and raises before the shapes are read.

nemo_rl.experience.rollouts._add_multimodal_generation_payload( generation_input_data: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.models.generation.interfaces.GenerationDatumSpec], flat_messages: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.FlatMessagesType], active_batch: nemo_rl.distributed.batched_data_dict.BatchedDataDict[nemo_rl.data.interfaces.DatumSpec], policy_generation: nemo_rl.models.generation.interfaces.GenerationInterface, *, deduplicate_multimodal_data: bool, ) → None#

Attach one policy-ready or vLLM-ready media representation.

The compact policy representation remains in message_log for later logprob/training construction. When every active row has a vLLM-ready prompt, sending that representation as well is redundant.

nemo_rl.experience.rollouts._reattach_original_multimodal_payloads( results: list[dict[str, Any]], original_message_logs: list[nemo_rl.data.interfaces.LLMMessageLogType | nemo_rl.data.interfaces.VLMMessageLogType], ) → None#

Restore exact prompt media omitted by a remote Gym rollout.

User turns are matched by their ordinal position. Only explicit PackedTensor values and named vLLM-ready media are restored, so arbitrary non-text metadata is never misclassified as media. Newly returned Gym media is left untouched unless it occupies the corresponding original prompt key.

nemo_rl.experience.rollouts._reattach_static_multimodal_payloads_to_result( result: dict[str, Any], source_message_log: list[dict[str, Any]], tokenizer: Any = None, ) → None#: Restore static media to each Gym-authored message-log representation.

nemo_rl.experience.rollouts.attach_static_multimodal_payload( target_message_log: list[dict[str, Any]], source_message_log: list[dict[str, Any]], processed_target_ids: set[int] | None = None, tokenizer: Any = None, ) → None#

Copy policy-ready media from static source turns to Gym-authored turns.

A shared processed_target_ids prevents aliased views from consuming a rollout-matched marker and then overwriting the repaired media. When a tokenizer is supplied, static video placeholder geometry is verified before tensors are copied.

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._mask_sample_flags( extras: collections.abc.Iterable[dict[str, Any] | None], ) → torch.Tensor#

Return True for samples the environment asks GRPO to mask from loss.

The flag is read from the verify response’s first-class mask_sample field (NeMo-Gym BaseVerifyResponse) or from the response-extras mapping instance_config.mask_sample. Both reads are consulted because Gym serializes mask_sample: false on every response, so an environment that flags only instance_config.mask_sample carries a false first-class field beside the true nested one.

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]], ) → list[int]#

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.

Returns:: how many of its messages were sentinel-filled. A replay_buffer-level “field entirely absent” guard cannot see this — it only catches total absence across the whole batch, not a message backfilled inside an otherwise-routed rollout.
Return type:: One count per entry in message_logs

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.

validate_length_bound() → nemo_rl.experience.rollouts.EffortLevelsConfig#

nemo_rl.experience.rollouts.get_effort_config( env: dict[str, Any], ) → Optional[nemo_rl.experience.rollouts.EffortLevelsConfig]#: Resolve the optional Gym effort config using its centralized defaults.

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._terminal_completion_length(result: dict) → Optional[int]#: Return inline terminal assistant length; capture is shaped after verification.

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._effort_shaping_metrics( shaping: nemo_rl.experience.rollouts._EffortShapingMetrics, ) → dict[str, float]#

Build the rollout-metric entries for one group’s effort-shaping lists.

Shared by the batched v1 path and the SingleController rollout manager so the two cannot drift apart.

Parameters:: shaping – Per-sample tracking lists returned by _apply_effort_shaping.
Returns:: Metric name to value. Empty only when shaping was disabled; callers update an existing dict, so an absent key leaves the metric unreported rather than reporting a zero.

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, deduplicate_multimodal_data: 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.
deduplicate_multimodal_data – Send only vLLM-ready media through the vLLM generation boundary while retaining compact policy media for logprob and training.

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, *, sample_multimodal_data: dict[str, Any] | None = None, deduplicate_multimodal_data: 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
sample_multimodal_data – vLLM-ready media fields for this sample.
deduplicate_multimodal_data – Avoid sending both vLLM-ready and policy-ready media through the async generation boundary.

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, deduplicate_multimodal_data: 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
deduplicate_multimodal_data – Avoid redundant media at generation boundaries while preserving compact policy media in the trajectory.

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, deduplicate_multimodal_data: 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, deduplicate_multimodal_data: 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, deduplicate_multimodal_data: 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, resolved_agent_ref: 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.compute_reward_penalty_metrics( penalty_counts: dict[str, int], num_results: int, reward_penalty_config: dict[str, Any] | pydantic.BaseModel | None, ) → dict[str, float]#: Return enabled penalty rates using the legacy NeMo-Gym metric names.

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, num_generations: int, ) → None#: Apply NeMo-RL sampling parameters and stable row indices in place.

nemo_rl.experience.rollouts._bucket_nemo_gym_rows_by_instance( rows: list[dict], shard_set: nemo_rl.environments.nemo_gym.NemoGymShardSet, num_generations: int, ) → list[tuple[str, Any, list[dict]]]#

Split already-stamped rows into one bucket per actor, by prompt group.

Buckets are per actor instance, not per shard: a shard with replicas gets one bucket per replica, which is how replication spreads a step’s load.

Bucketing is by group and never by row. A group split across instances would deadlock any verifier that scores rollouts against their peers: those buffer a whole group in one process and block until it is complete, so each half would wait forever for rollouts sitting in the other.

Rows keep the global _rowidx stamped over the whole batch, so the accumulator downstream never learns that shards exist.

Returns:: (instance_label, handle, rows) per actor that has work, in first-use order. Actors with no rows are omitted rather than sent an empty batch.

async nemo_rl.experience.rollouts._merge_nemo_gym_instance_streams( buckets: list[tuple[str, Any, list[dict]]], timer_prefix: str, deduplicate_multimodal_data: bool = False, debug_payload_metrics: bool = False, ) → collections.abc.AsyncGenerator[tuple[int, dict, dict, dict | None, str], None]#

Interleave K actor streams into one, yielding rows as they complete.

run_rollouts is a streaming generator, so there is nothing to gather: this replaces one stream with a merge of K and hands the merged rows to the accumulator that already exists. With a single bucket it degenerates to that one stream, which is what keeps the unsharded path unchanged.

Yields:: (row_index, resolved_agent_ref, result, timing_metrics, instance_label). timing_metrics is non-None only on the last row of each bucket.

nemo_rl.experience.rollouts._merge_nemo_gym_timing_metrics( timing_by_instance: dict[str, dict[str, Any]], timer_prefix: str, ) → dict[str, Any]#

Combine per-instance actor timings into per-instance plus rolled-up metrics.

Keyed by instance rather than by shard, so the replicas of one shard do not overwrite each other and report as though the shard had done a fraction of its work.

Every instance reports the same key names, so the roll-up takes the maximum rather than the sum: instances run concurrently and the step is gated by the slowest one, not by their total. That also means the rolled-up postprocess_results_pct is the worst instance’s share rather than a batch-wide average, which is the honest reading – averaging percentages over instances with different row counts would not mean anything.

With one instance this returns exactly what it reported, so the unsharded path is unchanged.

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, length_penalty_config: nemo_rl.utils.length_penalty.LengthPenaltyConfig | dict[str, Any] | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, returns_entire_batch: bool = False, routing_group_size: Optional[int] = None, sampling_params: Optional[nemo_rl.models.generation.interfaces.GenerationSamplingParams] = None, deduplicate_multimodal_data: bool = False, debug_payload_metrics: bool = False, identity_num_generations: Optional[int] = 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 for local postprocessing. The actor holds its own, installed once at spinup – see NemoGym.set_tokenizer.
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.
length_penalty_config – Optional grpo.length_penalty block; applied per prompt group of identity_num_generations rows.
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.
routing_group_size – Number of rows that must stay on one actor instance. Defaults to num_generations. Synchronous callers use the true per-prompt generation count while accumulating one complete batch.
sampling_params – Sampling profile stamped onto every NeMo-Gym row. None uses the train profile from generation_config; validation passes its own profile explicitly.
deduplicate_multimodal_data – Omit initial policy-ready media from the remote Gym return and restore the exact original payload locally.
debug_payload_metrics – Emit logical, physical, and serialized media payload metrics at the Gym Ray boundary.
identity_num_generations – Number of contiguous rows sharing one logical prompt-group identity. Defaults to num_generations. Synchronous callers set this independently because they collect the full batch as one result while preserving per-prompt GenRM cohort identities.

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 either generation count is not positive, the batch is empty or not divisible by the relevant generation count, 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, *, num_generations_per_prompt: int, 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, length_penalty_config: nemo_rl.utils.length_penalty.LengthPenaltyConfig | dict[str, Any] | 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, deduplicate_multimodal_data: bool = False, debug_payload_metrics: bool = False, ) → 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 for local postprocessing. The actor holds its own, installed once at spinup – see NemoGym.set_tokenizer.
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.
num_generations_per_prompt – Number of contiguous rows belonging to each logical prompt group. This controls Gym/GenRM cohort identity only; the synchronous API still collects and returns the entire input batch.
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.
length_penalty_config – Optional grpo.length_penalty block; applied per prompt group of num_generations_per_prompt rows.
thinking_tags – Optional opening and closing tags used by thinking penalties.
num_generations_per_prompt – Number of contiguous rows produced from each original prompt. Each such group stays on one actor instance.
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.
deduplicate_multimodal_data – Omit initial policy-ready media from the remote Gym return and restore it from the input batch.
debug_payload_metrics – Emit exact Gym Ray-boundary media payload metrics.

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._single_group_reward_diagnostics( results: list[dict], ) → dict[str, float]#: Summarize reward variation within one prompt’s generations.

nemo_rl.experience.rollouts._group_reward_diagnostics( results: list[dict], group_size: int | None = None, ) → dict[str, float]#: Compute diagnostics over the actual prompt groups in results.

nemo_rl.experience.rollouts._nemo_gym_group_metrics( results: list[dict], group_size: int | None = None, ) → dict[str, Any]#

Compute metrics shared by both NeMo-Gym rollout implementations.

Used by the batched NeMo-Gym rollout path (run_async_nemo_gym_rollout) and the single-controller path (AsyncNemoGymRolloutImpl). group_size identifies prompt-group boundaries; when omitted, all results are treated as one group.

nemo_rl.experience.rollouts._LENGTH_LIMIT_INCOMPLETE_REASONS#: ‘frozenset(…)’

nemo_rl.experience.rollouts.nemo_gym_rollout_hit_max_tokens( full_result: dict[str, Any] | None, *, total_tokens: int, max_total_tokens_per_sample: int, ) → bool#

Return whether a NeMo-Gym rollout was cut off by the token budget.

An exact total_tokens == max_total_tokens_per_sample match almost never fires for multi-turn rollouts: when the next prompt no longer fits, Gym’s vLLM model returns an empty finish_reason="length" completion that adds no tokens, so the chain stops short of the limit. Count the rollout as truncated if it reached the limit, or if the final response is incomplete for a length reason.

The length reason is deliberately not restricted to context exhaustion: a final response stopped by its per-turn max_output_tokens cap carries the same incomplete_details.reason and also counts as truncated, even when the accumulated token count is below max_total_tokens_per_sample. Its final generation was length-cut, so it belongs in truncation_rate and is a valid overlong_filtering target. Responses that are incomplete for another reason (e.g. content_filter), responses without incomplete_details, and absent or malformed metadata are not truncation.

Parameters:

full_result – The Gym rollout result; its response key holds the final Responses object when the agent produced one. Agents that return the model’s final Responses object carry incomplete_details; agents that synthesize their terminal response do not, and fall back to the token-count check.
total_tokens – Token count of the rollout’s full message log.
max_total_tokens_per_sample – The generation backend’s context limit.

Returns:

True if the rollout was cut off by a token limit.

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, length_penalty_config: nemo_rl.utils.length_penalty.LengthPenaltyConfig | dict[str, Any] | None = None, thinking_tags: list[str] | tuple[str, ...] | None = None, mask_env_flagged_samples: bool = True, group_size: Optional[int] = None, ) → nemo_rl.experience.rollouts.NemoGymRolloutResult#

Postprocess one complete prompt group from the NeMo-Gym stream.

group_size is the number of contiguous rows sharing one prompt identity (identity_num_generations); the synchronous path hands the whole batch in as one call, so it may hold several prompt groups. Defaults to the batch.

nemo_rl.experience.rollouts#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.experience.rollouts`#