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.

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],
seq_logprob_error_metrics: list[dict[str, float]] | 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.

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

Returns:

Step-level reward, advantage, token-count, and optional sequence log-probability error metrics.

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.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.