nemo_rl.algorithms.single_controller_utils.utils#

Helpers used by SingleControllerActor.

Module Contents#

Functions#

aggregate_step_metrics

Reduce per-microbatch metric lists into step-level scalars.

reduce_advantage_pump_metrics

Reduce per-step accumulators from _advantage_stage into step scalars.

_reduce_seq_logprob_error_metrics

Reduce sequence-error metrics across streaming chunks.

apply_message_level_advantage_penalties

Overwrite flagged token advantages while leaving valid tokens unchanged.

tensor_field

Read a tensor column from a TensorDict, depadding if nested.

squeeze_trailing_unit_dim

Drop a trailing dim of size 1 if present.

fields_for_put

Pack tensors for DataPlane put, re-nesting jagged rows when needed.

Data#

_MB_METRIC_MIN

_MB_METRIC_MAX

_MB_METRIC_MEAN

API#

nemo_rl.algorithms.single_controller_utils.utils._MB_METRIC_MIN: frozenset[str]#

‘frozenset(…)’

nemo_rl.algorithms.single_controller_utils.utils._MB_METRIC_MAX: frozenset[str]#

‘frozenset(…)’

nemo_rl.algorithms.single_controller_utils.utils._MB_METRIC_MEAN: frozenset[str]#

‘frozenset(…)’

nemo_rl.algorithms.single_controller_utils.utils.aggregate_step_metrics(
train_result: dict[str, Any],
) → dict[str, Any]#

Reduce per-microbatch metric lists into step-level scalars.

Parameters:

train_result – Output of TQPolicy.finish_train_step.

Returns:

Flat dict of step-level scalars ready for logging.

nemo_rl.algorithms.single_controller_utils.utils.reduce_advantage_pump_metrics(
rewards: list[torch.Tensor],
masked_advantages: list[torch.Tensor],
sequence_lengths: list[int],
*,
sample_masks: list[torch.Tensor] | None = None,
seq_logprob_error_metrics: list[dict[str, float]] | None = None,
num_mask_sample_filtered: list[int] | None = None,
num_invalid_tool_calls: list[int] | None = None,
num_malformed_thinking: list[int] | None = None,
num_assistant_messages: list[int] | None = None,
num_routed_experts_backfilled: list[int] | None = None,
) → dict[str, float]#

Reduce per-step accumulators from _advantage_stage into step scalars.

Parameters:

  • rewards – One tensor per advantage_stage call; each row a sample reward.

  • masked_advantages – Token-masked advantages, one tensor per call.

  • sequence_lengths – All input_lengths trained on this step.

  • sample_masks – Row validity for each rewards entry (token-capture placeholders and mask_sample/overlong/seq-logprob-error rows carry 0). Weights reward so it averages over trained rows only, matching what the advantage estimator’s baseline already excludes. None keeps the legacy unweighted mean.

  • seq_logprob_error_metrics – Sequence-error metrics and their aggregation counts, one record per streaming chunk.

  • num_mask_sample_filtered – Environment-flagged sample counts, one per streaming chunk.

  • num_invalid_tool_calls – Per-sample invalid tool-call counts.

  • num_malformed_thinking – Per-sample malformed-thinking counts.

  • num_assistant_messages – Per-sample assistant message counts (rate denominator).

Returns:

Step-level reward, advantage, token-count, optional sequence log-probability error metrics, the num_mask_sample_filtered count, and per-sample violation counts.

nemo_rl.algorithms.single_controller_utils.utils._reduce_seq_logprob_error_metrics(
records: list[dict[str, float]],
) → dict[str, float]#

Reduce sequence-error metrics across streaming chunks.

nemo_rl.algorithms.single_controller_utils.utils.apply_message_level_advantage_penalties(
advantages: torch.Tensor,
*,
invalid_tool_call_mask: torch.Tensor,
malformed_thinking_mask: torch.Tensor,
invalid_tool_call_advantage: float | None,
malformed_thinking_advantage: float | None,
) → torch.Tensor#

Overwrite flagged token advantages while leaving valid tokens unchanged.

Invalid-tool-call penalties take precedence when both masks select the same token, matching the legacy GRPO message-level implementation.

Parameters:

  • advantages – Per-token advantages of shape (batch, seq).

  • invalid_tool_call_mask – Bool mask of the same shape as advantages; True at tokens produced by an invalid tool call.

  • malformed_thinking_mask – Bool mask of the same shape as advantages; True at tokens produced by malformed thinking.

  • invalid_tool_call_advantage – Value to overwrite flagged tokens with, or None to leave the invalid-tool-call branch disabled.

  • malformed_thinking_advantage – Value to overwrite flagged tokens with, or None to leave the malformed-thinking branch disabled.

Returns:

New tensor with penalties applied. The original advantages object is returned unchanged when both advantages are None.

Raises:

ValueError – If either mask shape does not match advantages.

nemo_rl.algorithms.single_controller_utils.utils.tensor_field(
data: tensordict.TensorDict,
field_name: str,
) → torch.Tensor#

Read a tensor column from a TensorDict, depadding if nested.

Parameters:

  • data – TensorDict returned by the data plane.

  • field_name – Column name to fetch.

Returns:

Dense tensor (nested columns are padded with zeros).

nemo_rl.algorithms.single_controller_utils.utils.squeeze_trailing_unit_dim(value: torch.Tensor) → torch.Tensor#

Drop a trailing dim of size 1 if present.

Parameters:

value – Input tensor.

Returns:

Tensor without the trailing unit dim.

nemo_rl.algorithms.single_controller_utils.utils.fields_for_put(
meta: nemo_rl.data_plane.KVBatchMeta,
fields: dict[str, torch.Tensor],
) → tensordict.TensorDict#

Pack tensors for DataPlane put, re-nesting jagged rows when needed.

Parameters:

  • meta – Batch meta whose sequence_lengths drive the nesting.

  • fields – Field name to dense tensor.

Returns:

TensorDict shaped for dp_client.put_samples.