nemo_rl.experience.rollouts#

Module Contents#

Classes#

EffortLevelsConfig

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

_EffortShapingMetrics

AsyncNemoGymRolloutResult

Functions#

_add_r3_fallback_metrics

_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

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._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.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,
) → 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.