nemo_automodel.components.loss.decision

Post-training objectives for decision models (per-question distributions over a finite option set).

Every question of every record is one decision: the model emits one logit per allowed option and the answer distribution is a softmax over them. Questions have different numbers of options, so a batch is represented as a padded [questions, options] logit matrix plus a boolean option mask, where questions counts every question of every record in the batch and options is the largest option count among them.

The configured decision objective supports two stages:

• SFT - label-smoothed cross-entropy on the gold option, plus a Brier term (calibration) and a ranked-probability score (squared EMD) for ordinal score questions. The optional ranked-probability term applies only to ordinal questions.

• RL - the decision is a one-step contextual bandit, so the policy gradient can be computed two ways:

  • expected_reward: maximise E_{a~pi}[r(a)] = sum_a pi(a) r(a) exactly. With a finite option set this is the exact policy gradient without option-sampling variance sum_a pi(a) (r(a) - E[r]) grad log pi(a).
  • grpo: sample num_samples options per question, use the group-mean (optionally std-normalised) reward as the baseline and apply a grouped REINFORCE surrogate for single-step decisions.

  Both add a KL penalty to a frozen reference policy (normally the SFT model) and an optional proper-scoring anchor (log score of the gold option). These objectives do not guarantee calibration.

The same module owns the tutorial-specific group-masked bidirectional InfoNCE objective and experimental Gaussian-proposal RLCR continuation.

Module Contents

Classes

Functions

Data

IGNORE_INDEX

SCORE_TYPE

API

class nemo_automodel.components.loss.decision.DecisionLoss(
    config: nemo_automodel.components.loss.decision.DecisionLossConfig
)

nemo_automodel.components.loss.decision.DecisionLoss._effective_weights(
    targets: nemo_automodel.components.loss.decision_batch.DecisionTargets
) -> torch.Tensor

nemo_automodel.components.loss.decision.DecisionLoss._rewards(
    targets: nemo_automodel.components.loss.decision_batch.DecisionTargets,
    onehot: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.loss.decision.DecisionLoss.forward(
    logits: torch.Tensor,
    targets: nemo_automodel.components.loss.decision_batch.DecisionTargets,
    generator: torch.Generator | None = None
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]

nemo_automodel.components.loss.decision.DecisionLoss.normalization_weight(
    targets: nemo_automodel.components.loss.decision_batch.DecisionTargets
) -> torch.Tensor

class nemo_automodel.components.loss.decision.DecisionLossConfig(
    mode: typing.Literal['sft', 'rl'] = 'sft',
    label_smoothing: float = 0.05,
    brier_weight: float = 1.0,
    ordinal_weight: float = 0.5,
    rl_objective: typing.Literal['expected_reward', 'grpo'] = 'expected_reward',
    num_samples: int = 8,
    sample_temperature: float = 1.0,
    normalize_advantages: bool = True,
    kl_coef: float = 0.05,
    anchor_weight: float = 0.1,
    question_weights: dict[str, float] = dict()
)

nemo_automodel.components.loss.decision.DecisionLossConfig.build() -> 'DecisionLoss'

class nemo_automodel.components.loss.decision.RLCRLossOutput(
    loss: torch.Tensor,
    policy_loss: torch.Tensor,
    supervised_loss: torch.Tensor,
    mean_reward: torch.Tensor,
    brier: torch.Tensor
)

nemo_automodel.components.loss.decision._masked_mean(
    values: torch.Tensor,
    weights: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.loss.decision.group_masked_bidirectional_infonce(
    state_projections: torch.Tensor,
    action_projections: torch.Tensor,
    logit_scale: torch.Tensor,
    group_ids: torch.Tensor | None = None
) -> torch.Tensor

nemo_automodel.components.loss.decision.masked_log_softmax(
    logits: torch.Tensor,
    mask: torch.Tensor
) -> torch.Tensor

nemo_automodel.components.loss.decision.rlcr_loss(
    logits: torch.Tensor,
    targets: torch.Tensor,
    group_size: int,
    exploration_std: float,
    policy_weight: float,
    supervised_weight: float,
    generator: torch.Generator | None = None
) -> nemo_automodel.components.loss.decision.RLCRLossOutput

nemo_automodel.components.loss.decision.IGNORE_INDEX = -100

nemo_automodel.components.loss.decision.SCORE_TYPE = 2