`nemo_rl.experience.rollouts`#

Module Contents#

Classes#

`EffortLevelsConfig`	Controls length-based reward shaping for low-effort prompts.
`_EffortShapingMetrics`
`AsyncNemoGymRolloutResult`

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`
`_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.
`run_async_multi_turn_rollout`	Run multi-turn rollouts with sample-level processing.
`_tensorize_by_key`
`_calculate_single_metric`
`get_nemo_gym_thinking_tags`	Return thinking tags used by the Gym-side detector.
`_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.
`run_async_nemo_gym_rollout`	Run multi-turn rollouts with NeMo-Gym. Please refer to the `run_async_multi_turn_rollout` docs for more information on the parameters.

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#

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)

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 multi-turn rollouts with sample-level processing.

Each sample in the batch proceeds through its interaction independently. Async generation is used internally when available but the function is synchronous.

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_seq_len – Maximum sequence length allowed
max_rollout_turns – Maximum number of agent-environment interaction turns
greedy – Whether to use greedy decoding

Returns:

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

Return type:

Tuple containing

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

class nemo_rl.experience.rollouts.AsyncNemoGymRolloutResult#

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

nemo_rl.experience.rollouts._calculate_single_metric( values: collections.abc.Sequence[float | int], batch_size: int, key_name: str, ) → dict#

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._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.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, 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, ) → nemo_rl.experience.rollouts.AsyncNemoGymRolloutResult#: Run multi-turn rollouts with NeMo-Gym. Please refer to the run_async_multi_turn_rollout docs for more information on the parameters.

nemo_rl.experience.rollouts#

Module Contents#

Classes#

Functions#

Data#

API#

`nemo_rl.experience.rollouts`#