feat(grpo): support message penalties in single controller

examples/configs/grpo_math_1B_megatron_single_controller.yaml

@@ -7,6 +7,10 @@ grpo:
   # SC reads async_rl instead, and the entrypoint raises if this block is set.
   async_grpo: null
   val_period: 0
+  # Optional NeMo-Gym message-level penalties. Values overwrite computed
+  # advantages on tokens belonging to the corresponding assistant message.
+  invalid_tool_call_advantage: null
+  malformed_thinking_advantage: null
 
 async_rl:
   sampler:

examples/configs/ppo_math_1B_megatron_single_controller.yaml

@@ -15,6 +15,10 @@ ppo:
     enabled: false
   reward_scaling:
     enabled: false
+  # Optional NeMo-Gym message-level penalties. Values overwrite computed
+  # advantages on tokens belonging to the corresponding assistant message.
+  invalid_tool_call_advantage: null
+  malformed_thinking_advantage: null
 
 async_rl:
   sampler:

nemo_rl/algorithms/async_utils/replay_buffer.py

@@ -981,11 +981,13 @@ def __init__(
         partition_id: str,
         *,
         pad_value_dict: Mapping[str, int],
+        include_message_violation_fields: bool,
         require_routed_experts: bool = False,
     ):
         self._dp_client = dp_client
         self._partition_id = partition_id
         self._pad_value_dict = dict(pad_value_dict)
+        self._include_message_violation_fields = include_message_violation_fields
         self._require_routed_experts = require_routed_experts
         self.meta_list: list[Optional[KVBatchMeta]] = []
         self.start_weight_list: list[int] = []
@@ -1090,7 +1092,11 @@ async def commit(
                 "TQReplayBuffer must be bound to the controller data-plane "
                 "checkpoint barrier before committing samples"
             )
-        train_batch = record_to_train_batch(record, pad_value_dict=self._pad_value_dict)
+        train_batch = record_to_train_batch(
+            record,
+            pad_value_dict=self._pad_value_dict,
+            include_message_violation_fields=self._include_message_violation_fields,
+        )
         sample_ids, fields, tags = pack_payload(
             train_batch,
             weight_version=start_weight_version,

nemo_rl/algorithms/ppo.py

@@ -205,6 +205,11 @@ class PPOConfig(BaseModel, extra="allow"):
     # Nullable sequence-level multiplicative probability-error threshold.
     # None logs metrics without masking; values above the threshold are excluded.
     seq_logprob_error_threshold: float | None = None
+    # Advantage value assigned to invalid-tool-call tokens; None disables it.
+    invalid_tool_call_advantage: float | None = None
+    # Advantage value assigned to malformed-thinking tokens; None disables it.
+    malformed_thinking_advantage: float | None = None
+
     # Asynchronous PPO uses a replay buffer with non-colocated generation.
     # Legacy async config block; SC reads its async knobs from `async_rl` instead.
     async_ppo: AsyncPPOConfig | None = Field(default_factory=AsyncPPOConfig)

nemo_rl/algorithms/single_controller.py

@@ -93,6 +93,7 @@
 from nemo_rl.algorithms.single_controller_utils.setup import SingleControllerActorArgs
 from nemo_rl.algorithms.single_controller_utils.utils import (
     aggregate_step_metrics,
+    apply_message_level_advantage_penalties,
     fields_for_put,
     reduce_advantage_pump_metrics,
     squeeze_trailing_unit_dim,
@@ -201,6 +202,10 @@ def __init__(
         self._is_ppo: bool = is_ppo_run(master_config)
         # GRPO has no epoch knob: it makes one optimizer step per RL step.
         self._ppo_epochs: int = self._algo_cfg.ppo_epochs if self._is_ppo else 1
+        self._message_level_advantage_penalties_enabled = (
+            self._algo_cfg.invalid_tool_call_advantage is not None
+            or self._algo_cfg.malformed_thinking_advantage is not None
+        )
 
         self._policy_logprobs_required = not (
             master_config.loss_fn.force_on_policy_ratio
@@ -3265,6 +3270,27 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
             if self._is_ppo:
                 returns = torch.zeros_like(mask)
 
+        if self._message_level_advantage_penalties_enabled:
+            # Sequence-error filtering and the pre-existing sample mask remain
+            # authoritative: a message penalty must not make a filtered token
+            # trainable again.
+            valid_tokens = mask.bool()
+            advantages = apply_message_level_advantage_penalties(
+                advantages,
+                invalid_tool_call_mask=(
+                    tensor_field(data, adv_cfg.invalid_tool_call_mask_field).bool()
+                    & valid_tokens
+                ),
+                malformed_thinking_mask=(
+                    tensor_field(data, adv_cfg.malformed_thinking_mask_field).bool()
+                    & valid_tokens
+                ),
+                invalid_tool_call_advantage=self._algo_cfg.invalid_tool_call_advantage,
+                malformed_thinking_advantage=(
+                    self._algo_cfg.malformed_thinking_advantage
+                ),
+            )
+
         response_advantages = torch.masked_select(advantages, mask.bool())
         self._step_log_dict["rewards"].append(rewards.detach().cpu())
         self._step_log_dict["masked_advantages"].append(
@@ -3306,6 +3332,13 @@ def _advantage_input_fields(self) -> list[str]:
             adv_cfg.sample_mask_field,
             *adv_cfg.repeated_batch_fields,
         ]
+        if self._message_level_advantage_penalties_enabled:
+            fields.extend(
+                [
+                    adv_cfg.invalid_tool_call_mask_field,
+                    adv_cfg.malformed_thinking_mask_field,
+                ]
+            )
         if self._policy_logprobs_required:
             fields.append(adv_cfg.policy_logprobs_field)
         if self._policy_logprobs_required:

nemo_rl/algorithms/single_controller_utils/config.py

@@ -42,7 +42,12 @@
 from nemo_rl.algorithms.ppo import PPOConfig
 from nemo_rl.data import DataConfig
 from nemo_rl.data_plane.interfaces import DataPlaneConfig
+from nemo_rl.data_plane.schema import (
+    INVALID_TOOL_CALL_MASK,
+    MALFORMED_THINKING_MASK,
+)
 from nemo_rl.distributed.virtual_cluster import ClusterConfig
+from nemo_rl.environments.nemo_gym import should_use_nemo_gym
 from nemo_rl.models.policy import PolicyConfig
 from nemo_rl.models.value import ValueConfig
 from nemo_rl.utils.checkpoint import CheckpointingConfig
@@ -1024,6 +1029,17 @@ def validate_single_controller_config(master_config: MasterConfig) -> None:
     # configs can omit it. Only apply rollout-path validation when it is present.
     env_config = getattr(master_config, "env", None)
 
+    penalties_enabled = (
+        algo_cfg.invalid_tool_call_advantage is not None
+        or algo_cfg.malformed_thinking_advantage is not None
+    )
+    if penalties_enabled and not should_use_nemo_gym(master_config):
+        raise ValueError(
+            "invalid_tool_call_advantage and malformed_thinking_advantage on the "
+            "active algorithm block require the NeMo-Gym rollout path "
+            "(env.should_use_nemo_gym=true) on SingleController."
+        )
+
     opd_enabled = opd_module.is_opd_enabled(master_config)
     if opd_enabled and is_ppo_run(master_config):
         raise ValueError(
@@ -1110,6 +1126,8 @@ class AdvantageConfig:
     reward_field: str = "total_reward"
     token_mask_field: str = "token_mask"
     sample_mask_field: str = "sample_mask"
+    invalid_tool_call_mask_field: str = INVALID_TOOL_CALL_MASK
+    malformed_thinking_mask_field: str = MALFORMED_THINKING_MASK
     repeated_batch_fields: list[str] = field(default_factory=list)
     policy_logprobs_field: str = "prev_logprobs"
     generation_logprobs_field: str = "generation_logprobs"

nemo_rl/algorithms/single_controller_utils/setup.py

@@ -1393,6 +1393,10 @@ def _build_generation_then_trainer(
         dp_client,
         partition_id=partition_id,
         pad_value_dict={"token_ids": pad_id, "input_ids": pad_id},
+        include_message_violation_fields=(
+            algo_cfg.invalid_tool_call_advantage is not None
+            or algo_cfg.malformed_thinking_advantage is not None
+        ),
         require_routed_experts=router_replay_enabled(policy_config),
     )
     rollout_manager = RolloutManager(

nemo_rl/algorithms/single_controller_utils/utils.py

@@ -198,6 +198,74 @@ def reduce_range(
     return reduced
 
 
+def 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.
+
+    Args:
+        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``.
+    """
+    if invalid_tool_call_mask.shape != advantages.shape:
+        raise ValueError(
+            "invalid_tool_call_mask shape "
+            f"{tuple(invalid_tool_call_mask.shape)} does not match advantages "
+            f"{tuple(advantages.shape)}"
+        )
+    if malformed_thinking_mask.shape != advantages.shape:
+        raise ValueError(
+            "malformed_thinking_mask shape "
+            f"{tuple(malformed_thinking_mask.shape)} does not match advantages "
+            f"{tuple(advantages.shape)}"
+        )
+
+    result = advantages
+    if malformed_thinking_advantage is not None:
+        result = torch.where(
+            malformed_thinking_mask.bool(),
+            torch.as_tensor(
+                malformed_thinking_advantage,
+                dtype=advantages.dtype,
+                device=advantages.device,
+            ),
+            result,
+        )
+    if invalid_tool_call_advantage is not None:
+        result = torch.where(
+            invalid_tool_call_mask.bool(),
+            torch.as_tensor(
+                invalid_tool_call_advantage,
+                dtype=advantages.dtype,
+                device=advantages.device,
+            ),
+            result,
+        )
+    return result
+
+
 def tensor_field(data: TensorDict, field_name: str) -> torch.Tensor:
     """Read a tensor column from a TensorDict, depadding if nested.
 

nemo_rl/data_plane/column_io.py

@@ -37,7 +37,12 @@
 from nemo_rl.data.llm_message_utils import attach_message_log_view
 from nemo_rl.data_plane.codec import materialize, pack_jagged_fields
 from nemo_rl.data_plane.interfaces import DataPlaneClient, KVBatchMeta
-from nemo_rl.data_plane.schema import GLOBAL_FORWARD_PAD_SEQLEN, Layout
+from nemo_rl.data_plane.schema import (
+    GLOBAL_FORWARD_PAD_SEQLEN,
+    INVALID_TOOL_CALL_MASK,
+    MALFORMED_THINKING_MASK,
+    Layout,
+)
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 
 TOKEN_ALIGNED_FIELDS = frozenset(
@@ -53,6 +58,8 @@
         "token_mask",
         "sample_mask",
         "routed_experts",
+        INVALID_TOOL_CALL_MASK,
+        MALFORMED_THINKING_MASK,
     }
 )
 

nemo_rl/data_plane/schema.py

@@ -30,6 +30,10 @@
 SAMPLE_MASK = "sample_mask"
 META_IDX = "meta_idx"
 
+# Token-aligned message-violation fields consumed by SingleController advantages.
+INVALID_TOOL_CALL_MASK = "invalid_tool_call_mask"
+MALFORMED_THINKING_MASK = "malformed_thinking_mask"
+
 # Tensor fields in the train partition. Rollout writes the input
 # subset on first put; later stages add prev_logprobs /
 # reference_policy_logprobs (workers) and advantages (driver).
@@ -56,6 +60,8 @@
     "values",
     "returns",
     "teacher_reference_logprobs",
+    INVALID_TOOL_CALL_MASK,
+    MALFORMED_THINKING_MASK,
 )
 
 # Subset fetched by logprob / ref-logprob workers.

nemo_rl/experience/payload.py

@@ -15,7 +15,7 @@
 """Producer-side payload helpers for the async-RL TQ path."""
 
 from collections.abc import Mapping
-from typing import Any
+from typing import Any, cast
 
 import numpy as np
 import torch
@@ -24,7 +24,11 @@
 from nemo_rl.data.interfaces import LLMMessageLogType, VLMMessageLogType
 from nemo_rl.data_plane.codec import pack_jagged_fields
 from nemo_rl.data_plane.column_io import TOKEN_ALIGNED_FIELDS
-from nemo_rl.data_plane.schema import ROUTED_EXPERTS_FIELD
+from nemo_rl.data_plane.schema import (
+    INVALID_TOOL_CALL_MASK,
+    MALFORMED_THINKING_MASK,
+    ROUTED_EXPERTS_FIELD,
+)
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.experience.interfaces import PromptGroupRecord
 
@@ -54,21 +58,53 @@ def _violation_counts(
     return counts
 
 
+def _add_message_violation_masks(
+    message_logs: list[LLMMessageLogType | VLMMessageLogType],
+) -> None:
+    """Attach token-aligned masks for generated assistant violations.
+
+    This must run before the generic message normalizer fills missing
+    ``generation_logprobs`` on prompt and environment messages, because field
+    presence distinguishes generated assistant turns.
+    """
+    for message_log in message_logs:
+        for message in message_log:
+            token_ids = cast(torch.Tensor, message["token_ids"])
+            is_generated_assistant = (
+                message["role"] == "assistant" and "generation_logprobs" in message
+            )
+            is_invalid = is_generated_assistant and bool(
+                message.get("is_invalid_tool_call", False)
+            )
+            is_malformed = is_generated_assistant and bool(
+                message.get("has_malformed_thinking", False)
+            )
+            message[INVALID_TOOL_CALL_MASK] = torch.full_like(
+                token_ids, is_invalid, dtype=torch.bool
+            )
+            message[MALFORMED_THINKING_MASK] = torch.full_like(
+                token_ids, is_malformed, dtype=torch.bool
+            )
+
+
 def record_to_train_batch(
     record: PromptGroupRecord,
     *,
     pad_value_dict: Mapping[str, int],
+    include_message_violation_fields: bool,
 ) -> BatchedDataDict[Any]:
     """Convert one prompt group's record into a packed BatchedDataDict of N rows.
 
     Args:
         record: Rollout's PromptGroupRecord with N completions to flatten into rows.
         pad_value_dict: Field-name → pad value used by batched_message_log_to_flat_message.
+        include_message_violation_fields: Whether to tensorize message violation
+            flags for configured advantage penalties.
 
     Returns:
         BatchedDataDict with input_ids, input_lengths, generation_logprobs, token_mask,
         sample_mask, prompt_ids_for_adv, total_reward, violation counts, and optional
-        routed_experts.
+        routed experts and message-violation masks.
     """
     # Lazy imports: grpo and llm_message_utils transitively pull
     # experience.rollouts, so importing at module top risks a cycle.
@@ -84,7 +120,10 @@ def record_to_train_batch(
     assert n > 0, "PromptGroupRecord has no completions"
 
     message_logs = [c.message_log for c in completions]
+    violation_counts = [_violation_counts(message_log) for message_log in message_logs]
     prompt_token_count = sum(len(m["token_ids"]) for m in record.prompt)
+    if include_message_violation_fields:
+        _add_message_violation_masks(message_logs)
     prompt_lengths = torch.full((n,), prompt_token_count, dtype=torch.long)
 
     # Must precede the prompt extraction: it reuses the same message dicts, so
@@ -117,10 +156,13 @@ def record_to_train_batch(
         "sample_mask": sample_mask,
         "prompt_ids_for_adv": prompt_flat["token_ids"],
         "total_reward": total_reward,
-        _VIOLATION_COUNTS_KEY: [_violation_counts(ml) for ml in message_logs],
+        _VIOLATION_COUNTS_KEY: violation_counts,
     }
     if ROUTED_EXPERTS_FIELD in flat:
         train_data[ROUTED_EXPERTS_FIELD] = flat[ROUTED_EXPERTS_FIELD]
+    if include_message_violation_fields:
+        train_data[INVALID_TOOL_CALL_MASK] = flat[INVALID_TOOL_CALL_MASK]
+        train_data[MALFORMED_THINKING_MASK] = flat[MALFORMED_THINKING_MASK]
     return BatchedDataDict[Any](train_data)