nemo_rl.experience.dynamic_sampling#

Shared dynamic-sampling reward participation, filtering and step metrics.

Module Contents#

Classes#

StepRewardStats

Reward totals for the groups generated for one training step.

Functions#

mask_sample_flagged

Read both mask flags, since Gym may default the first-class flag to false.

rewards_carry_no_signal

Whether a group’s participating rewards leave every row zero advantage.

participating_rewards

The rewards that vote in this group’s baseline and std.

record_group_rewards

Add one finished group’s rewards to the tally of the step it was made for.

drain_step_reward_metrics

Return the reward metrics for step and forget that step and older ones.

API#

nemo_rl.experience.dynamic_sampling.mask_sample_flagged(extra: dict[str, Any] | None) → bool#

Read both mask flags, since Gym may default the first-class flag to false.

nemo_rl.experience.dynamic_sampling.rewards_carry_no_signal(
rewards: collections.abc.Sequence[float],
) → bool#

Whether a group’s participating rewards leave every row zero advantage.

True for at most one participating reward, or when they are all exactly equal. Exact equality rather than a std threshold because float roundoff can give a uniform group a tiny nonzero std.

rewards must already be restricted to the rows that vote in the group statistics. Passing every row instead keeps groups whose voters are uniform: at masked_reward_policy="exclude" with overlong_filtering, [0 (truncated), 1, 1, 1] has spread overall and none at all among the three rows the baseline actually sees.

nemo_rl.experience.dynamic_sampling.participating_rewards(
record: nemo_rl.experience.interfaces.PromptGroupRecord,
*,
masked_reward_policy: Optional[str],
overlong_filtering: bool,
) → list[float]#

The rewards that vote in this group’s baseline and std.

Mirrors v1’s _dynamic_sampling_valid_mask, which is the mask the group statistics are actually computed under. Filtering on the full reward list instead would keep groups whose voters are uniform, and drop groups whose only spread comes from rows that never vote.

A zero dataset weight is excluded under either policy: processors.py cuts a prompt that did not fit the sequence limit down to a few tokens and gives it loss_multiplier == 0 before any rollout, so its rewards are placeholders rather than measurements. That is group-level here, so such a group has no voters at all and is trivial by construction.

Under exclude the environment’s mask_sample rows drop out, and with overlong_filtering so do truncated ones – the rows the loss drops later in the step. Under include both still vote: their rewards are measurements of the policy even though they carry no loss.

class nemo_rl.experience.dynamic_sampling.StepRewardStats#

Reward totals for the groups generated for one training step.

Counts every group that finished generating, whether dynamic sampling kept or dropped it. The kept groups all have reward spread, so their mean alone says little about how well the policy is doing.

groups: int#

0

zero_variance_groups: int#

0

reward_sum: float#

0.0

samples: int#

0

first_attempt_reward_sum: float#

0.0

first_attempt_samples: int#

0

nemo_rl.experience.dynamic_sampling.record_group_rewards(
step_rewards: dict[int, nemo_rl.experience.dynamic_sampling.StepRewardStats],
target_step: Optional[int],
rewards: collections.abc.Sequence[float],
*,
replacement_index: int,
zero_variance: bool,
) → None#

Add one finished group’s rewards to the tally of the step it was made for.

Parameters:
  • target_step – Step the group was stamped for. None (sampler does not stamp steps) means there is nothing to report against, so skip.

  • rewards – Every row’s reward, voting or not. unfiltered_reward is the pre-filter progress signal, so it averages the whole group.

  • replacement_index – 0 for the originally scheduled prompt, N for its N-th replacement.

  • zero_variance – Whether the group’s voting rewards carry no signal. Passed in rather than recomputed here, so the count cannot disagree with the decision the filter actually made on the same group.

nemo_rl.experience.dynamic_sampling.drain_step_reward_metrics(
step_rewards: dict[int, nemo_rl.experience.dynamic_sampling.StepRewardStats],
step: int,
) → dict[str, float]#

Return the reward metrics for step and forget that step and older ones.

Groups count toward the step they were generated for, even if a later step borrowed them (promote_ready_group). With lookahead above zero that can differ from the step that trained on them.

Returns:

unfiltered_reward (mean over every generated group, kept or dropped), first_attempt_reward (same, but first attempts only), generated_prompt_groups, zero_variance_prompt_groups and zero_variance_fraction. Empty if nothing was generated for the step.