feat: log toolcall and thinktag violation rate into single controller

nemo_rl/algorithms/single_controller.py

@@ -82,6 +82,7 @@
 from nemo_rl.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.environments.nemo_gym import should_use_nemo_gym
 from nemo_rl.experience.failures import RolloutStall
+from nemo_rl.experience.payload import VIOLATION_TAG_KEYS
 from nemo_rl.experience.rollout_manager import RolloutOutcome
 from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration
 from nemo_rl.models.generation.vllm import VllmGeneration
@@ -312,6 +313,7 @@ def __init__(
             "masked_advantages": [],
             "sequence_lengths": [],
             "seq_logprob_error_metrics": [],
+            **{key: [] for key in VIOLATION_TAG_KEYS},
         }
         self._opd_stat_sum = 0.0
         self._opd_stat_sumsq = 0.0
@@ -1960,6 +1962,10 @@ async def _advantage_stage(self, meta: KVBatchMeta) -> tuple[KVBatchMeta, bool]:
             The updated batch metadata and whether the batch contains at least
             one valid training token.
         """
+        for tag in meta.tags or []:
+            for key in VIOLATION_TAG_KEYS:
+                self._step_log_dict.setdefault(key, []).append(int(tag.get(key, 0)))
+
         if self._advantage_estimator is None:
             return meta, True
         adv_cfg = self._advantage_cfg

nemo_rl/algorithms/single_controller_utils/utils.py

@@ -94,7 +94,11 @@ def 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,
+    num_invalid_tool_calls: list[int] | None = None,
+    num_malformed_thinking: list[int] | None = None,
+    num_assistant_messages: list[int] | None = None,
 ) -> dict[str, float]:
     """Reduce per-step accumulators from _advantage_stage into step scalars.
 
@@ -104,10 +108,14 @@ def reduce_advantage_pump_metrics(
         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.
+        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, and optional sequence
-        log-probability error metrics.
+        Step-level reward, advantage, token-count, optional sequence
+        log-probability error metrics, and per-sample violation counts.
+
     """
     out: dict[str, float] = {}
     if rewards:
@@ -126,6 +134,15 @@ def reduce_advantage_pump_metrics(
         out["total_num_tokens"] = float(sum(sequence_lengths))
     if seq_logprob_error_metrics:
         out.update(_reduce_seq_logprob_error_metrics(seq_logprob_error_metrics))
+    n_asst = sum(num_assistant_messages or [])
+    if n_asst:
+        n_invalid = sum(num_invalid_tool_calls or [])
+        n_malformed = sum(num_malformed_thinking or [])
+        out["invalid_tool_call_rate"] = n_invalid / n_asst
+        out["malformed_thinking_rate"] = n_malformed / n_asst
+        out["num_invalid_tool_calls"] = float(n_invalid)
+        out["num_malformed_thinking"] = float(n_malformed)
+        out["num_assistant_messages"] = float(n_asst)
     return out
 
 

nemo_rl/experience/payload.py

@@ -21,12 +21,38 @@
 import torch
 from tensordict import TensorDict
 
+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.distributed.batched_data_dict import BatchedDataDict
 from nemo_rl.experience.interfaces import PromptGroupRecord
 
+VIOLATION_TAG_KEYS = (
+    "num_invalid_tool_calls",
+    "num_malformed_thinking",
+    "num_assistant_messages",
+)
+# Per-row violation counts ride ``tags`` rather than the tensor fields, so this
+# key is carried on the train batch and consumed by pack_payload.
+_VIOLATION_COUNTS_KEY = "violation_counts"
+
+
+def _violation_counts(
+    message_log: LLMMessageLogType | VLMMessageLogType,
+) -> dict[str, int]:
+    """Count invalid tool calls / malformed thinking over flagged assistant turns."""
+    counts = dict.fromkeys(VIOLATION_TAG_KEYS, 0)
+    for message in message_log:
+        if message["role"] != "assistant" or "generation_logprobs" not in message:
+            continue
+        counts["num_assistant_messages"] += 1
+        if message.get("is_invalid_tool_call", False):
+            counts["num_invalid_tool_calls"] += 1
+        if message.get("has_malformed_thinking", False):
+            counts["num_malformed_thinking"] += 1
+    return counts
+
 
 def record_to_train_batch(
     record: PromptGroupRecord,
@@ -41,7 +67,8 @@ def record_to_train_batch(
 
     Returns:
         BatchedDataDict with input_ids, input_lengths, generation_logprobs, token_mask,
-        sample_mask, prompt_ids_for_adv, total_reward, and optional routed_experts.
+        sample_mask, prompt_ids_for_adv, total_reward, violation counts, and optional
+        routed_experts.
     """
     # Lazy imports: grpo and llm_message_utils transitively pull
     # experience.rollouts, so importing at module top risks a cycle.
@@ -90,6 +117,7 @@ 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],
     }
     if ROUTED_EXPERTS_FIELD in flat:
         train_data[ROUTED_EXPERTS_FIELD] = flat[ROUTED_EXPERTS_FIELD]
@@ -110,7 +138,8 @@ def pack_payload(
         group_id: Per-group identifier used as the sample_id prefix; the caller owns uniqueness.
 
     Returns:
-        sample_ids of the form {group_id}_g{i}, a jagged-packed TensorDict, and per-row tags.
+        sample_ids of the form {group_id}_g{i}, a jagged-packed TensorDict, and per-row
+        tags carrying weight_version plus any per-row violation counts.
     """
     lengths = train_batch["input_lengths"]
     n = int(lengths.shape[0])
@@ -124,5 +153,6 @@ def pack_payload(
         tensor_fields, lengths=lengths, token_aligned_fields=TOKEN_ALIGNED_FIELDS
     )
     sample_ids = [f"{group_id}_g{i}" for i in range(n)]
-    tags = [{"weight_version": weight_version} for _ in range(n)]
+    violations = train_batch.get(_VIOLATION_COUNTS_KEY, [{}] * n)
+    tags = [{"weight_version": weight_version, **violations[i]} for i in range(n)]
     return sample_ids, fields_td, tags

tests/unit/experience/test_payload.py

@@ -132,7 +132,15 @@ def test_record_to_train_batch_preserves_routed_experts_in_tq_payload() -> None:
     packed_rows = list(packed_routes.unbind())
     assert torch.equal(packed_rows[0], expected_routes[0])
     assert torch.equal(packed_rows[1], expected_routes[1])
-    assert tags == [{"weight_version": 3}, {"weight_version": 3}]
+    no_violations = {
+        "num_invalid_tool_calls": 0,
+        "num_malformed_thinking": 0,
+        "num_assistant_messages": 1,
+    }
+    assert tags == [
+        {"weight_version": 3, **no_violations},
+        {"weight_version": 3, **no_violations},
+    ]
 
 
 def test_record_to_train_batch_omits_routed_experts_when_absent() -> None:
@@ -190,3 +198,42 @@ def test_record_to_train_batch_backfills_routes_for_failed_completion() -> None:
     _, fields, _ = pack_payload(train_batch, weight_version=3, group_id="group")
     assert "routed_experts" in fields
     assert list(fields["routed_experts"].unbind())[1].shape == (2, 2, 2)
+
+
+def test_pack_payload_stamps_violation_counts_on_tags() -> None:
+    """Each flag lands in its own counter; a row that never generated counts zero."""
+    completions = [
+        _completion(route_start=10, reward=1.0),
+        _completion(route_start=30, reward=1.0),
+        _failed_completion(),
+    ]
+    completions[0].message_log[1]["is_invalid_tool_call"] = True
+    completions[1].message_log[1]["has_malformed_thinking"] = True
+
+    train_batch = record_to_train_batch(
+        _record(completions),
+        pad_value_dict={"token_ids": 0, "input_ids": 0},
+    )
+    _, fields, tags = pack_payload(train_batch, weight_version=7, group_id="g")
+
+    assert "violation_counts" not in fields
+    assert tags == [
+        {
+            "weight_version": 7,
+            "num_invalid_tool_calls": 1,
+            "num_malformed_thinking": 0,
+            "num_assistant_messages": 1,
+        },
+        {
+            "weight_version": 7,
+            "num_invalid_tool_calls": 0,
+            "num_malformed_thinking": 1,
+            "num_assistant_messages": 1,
+        },
+        {
+            "weight_version": 7,
+            "num_invalid_tool_calls": 0,
+            "num_malformed_thinking": 0,
+            "num_assistant_messages": 0,
+        },
+    ]

tests/unit/single_controller/test_rollout_pump.py

@@ -1143,5 +1143,10 @@ def test_rollout_pump_writes_expected_tq_data(
     )
     for tag in tags:
         assert tag["weight_version"] == 0
-        # Slim tag schema: weight_version is the only field producers stamp.
-        assert set(tag) == {"weight_version"}
+        # Tag schema: weight_version plus per-row violation counts.
+        assert set(tag) == {
+            "weight_version",
+            "num_invalid_tool_calls",
+            "num_malformed_thinking",
+            "num_assistant_messages",
+        }

tests/unit/single_controller/test_utils.py

@@ -204,6 +204,34 @@ def test_seq_logprob_error_metrics_are_reduced_across_streaming_chunks(
         assert out["num_masked_seqs_by_logprob_error"] == 3
         assert out["masked_correct_pct"] == pytest.approx(1.0 / 3)
 
+    def test_violation_rates_from_per_sample_counts(self) -> None:
+        out = reduce_advantage_pump_metrics(
+            rewards=[],
+            masked_advantages=[],
+            sequence_lengths=[],
+            num_invalid_tool_calls=[1, 0, 1],
+            num_malformed_thinking=[0, 1, 0],
+            num_assistant_messages=[2, 1, 1],
+        )
+        assert out["invalid_tool_call_rate"] == pytest.approx(0.5)
+        assert out["malformed_thinking_rate"] == pytest.approx(0.25)
+        assert out["num_invalid_tool_calls"] == pytest.approx(2.0)
+        assert out["num_malformed_thinking"] == pytest.approx(1.0)
+        assert out["num_assistant_messages"] == pytest.approx(4.0)
+
+    def test_no_assistant_messages_omits_violation_metrics(self) -> None:
+        assert (
+            reduce_advantage_pump_metrics(
+                [],
+                [],
+                [],
+                num_invalid_tool_calls=[0],
+                num_malformed_thinking=[0],
+                num_assistant_messages=[0],
+            )
+            == {}
+        )
+
 
 class TestFieldsForPut:
     def test_no_sequence_lengths_packs_contiguous(self) -> None: