fix(verl): classify trajectory prefix mismatch diagnostics

agentlightning/verl/rollout_adapter.py

@@ -7,6 +7,7 @@
 import base64  # [multimodal-patch]
 import io
 import json
+import re
 import zipfile
 from typing import Any, cast
 
@@ -60,6 +61,39 @@ def _decode_token_ids(tokenizer: Any | None, ids: list[int]) -> str:
     return text
 
 
+def _classify_trace_merge_mismatch(
+    tokenizer: Any | None, previous_ids: list[int], prompt_ids: list[int]
+) -> tuple[bool, bool, bool]:
+    """Return diagnostic (template, retokenization, other) flags, never merge decisions.
+
+    Only decoded-prefix equality, whitespace and the explicit think/tool-call
+    delimiters are recognized. These are heuristics, not proof of the cause.
+    """
+    if tokenizer is None:
+        return False, False, True
+    try:
+        try:
+            previous_text = tokenizer.decode(previous_ids, skip_special_tokens=False)
+            prompt_text = tokenizer.decode(prompt_ids, skip_special_tokens=False)
+        except TypeError:
+            previous_text = tokenizer.decode(previous_ids)
+            prompt_text = tokenizer.decode(prompt_ids)
+    except Exception:
+        return False, False, True
+    if not isinstance(previous_text, str) or not isinstance(prompt_text, str):
+        return False, False, True
+    if previous_text and prompt_text.startswith(previous_text):
+        return False, True, False
+
+    # Keep all content, including reasoning/tool-call bodies and unknown tags.
+    # An empty normalized context cannot establish a meaningful prefix match.
+    previous_normalized = " ".join(re.sub(r"</?(?:think|tool_call)>", "", previous_text).split())
+    prompt_normalized = " ".join(re.sub(r"</?(?:think|tool_call)>", "", prompt_text).split())
+    if previous_normalized and prompt_normalized.startswith(previous_normalized):
+        return True, False, False
+    return False, False, True
+
+
 def _decode_trace_text(tokenizer: Any | None, ids: list[int]) -> str:
     text = _decode_token_ids(tokenizer, ids)
     if len(text) > _TRACE_MERGE_MISMATCH_TEXT_LIMIT:
@@ -593,16 +627,18 @@ def append_training_row(
                         continue
 
                     if len(merge_mismatch_rows) < _TRACE_MERGE_MISMATCH_WANDB_LIMIT:
+                        template_mismatch, retoken_mismatch, others_mismatch = _classify_trace_merge_mismatch(
+                            self.tokenizer, current_context, prompt_ids
+                        )
                         merge_mismatch_rows.append(
                             {
                                 "global_steps": global_steps,
                                 "rollout_id": rollout.rollout_id,
                                 "data_id": rollout.data_id,
                                 "turn_index": turn_index,
-                                # Token-prefix failures are classified as other mismatches.
-                                "template_mismatch": False,
-                                "retoken_mismatch": False,
-                                "others_mismatch": True,
+                                "template_mismatch": template_mismatch,
+                                "retoken_mismatch": retoken_mismatch,
+                                "others_mismatch": others_mismatch,
                                 "prompt_length": len(prompt_ids),
                                 "response_length": len(response_ids),
                                 "previous_trace_length": len(current_context),

docs/20-trainer-configuration.md

@@ -139,6 +139,14 @@ agentlightning:
 
 The aggregator automatically merges consecutive calls when the next prompt starts with the exact token sequence of the previous prompt and response. Tokens added between calls, such as tool observations, are retained as context but masked from the policy loss. If exact token-prefix continuity is broken, the aggregator starts a new training row instead of merging incompatible calls.
 
+For up to 100 such failures per batch, `training/trace_merge_mismatches` reports diagnostic flags:
+
+- `retoken_mismatch`: the decoded next prompt still starts with the decoded previous prompt and response;
+- `template_mismatch`: that prefix is recovered only after collapsing whitespace and removing the literal `<think>`, `</think>`, `<tool_call>` and `</tool_call>` delimiters;
+- `others_mismatch`: neither check matches, no tokenizer is available, or decoding fails.
+
+These flags are conservative heuristics, not proof of the underlying cause. Template normalization keeps reasoning/tool-call content and unknown tags; an empty normalized context is not a match. Classification never changes the strict token-prefix split or the resulting training tensors.
+
 In this mode:
 
 - `trajectory_max_prompt_length` limits the initial prompt in each merged training row;

tests/verl/test_rollout_adapter.py

@@ -336,6 +336,115 @@ def test_trajectory_prefix_mismatch_wandb_table_is_capped_at_100(
     assert table.rows[-1][3] == 100
 
 
+@pytest.mark.parametrize(
+    ("prompt", "response", "next_prompt", "use_tokenizer", "expected"),
+    [
+        ([1], [2], [3, 4], True, (False, True, False)),
+        ([1, 6], [7, 2], [1, 2, 4], True, (True, False, False)),
+        ([1, 10], [2, 11], [1, 2, 4], True, (True, False, False)),
+        ([9], [5], [3, 5, 4], True, (True, False, False)),
+        ([1], [2], [8, 4], True, (False, False, True)),
+        ([1, 6], [7, 2], [8, 4], True, (False, False, True)),
+        ([1, 12], [13, 2], [1, 2, 4], True, (False, False, True)),
+        ([6], [7], [8, 4], True, (False, False, True)),
+        ([-1], [2], [3, 4], True, (False, False, True)),
+        ([1], [2], [3, 4], False, (False, False, True)),
+    ],
+    ids=[
+        "retoken",
+        "think-delimiters",
+        "tool-call-delimiters",
+        "whitespace",
+        "content",
+        "template-and-content",
+        "unknown-delimiters",
+        "empty-normalized-context",
+        "decode-error",
+        "no-tokenizer",
+    ],
+)
+def test_trajectory_mismatch_classification_preserves_training_data(
+    monkeypatch: pytest.MonkeyPatch,
+    prompt: list[int],
+    response: list[int],
+    next_prompt: list[int],
+    use_tokenizer: bool,
+    expected: tuple[bool, bool, bool],
+) -> None:
+    from tokenizers import Tokenizer, models
+    from transformers import PreTrainedTokenizerFast
+
+    tokenizer = PreTrainedTokenizerFast(
+        tokenizer_object=Tokenizer(
+            models.WordLevel(
+                {
+                    "[UNK]": 0,
+                    "hello": 1,
+                    "world": 2,
+                    "hello world": 3,
+                    "obs": 4,
+                    "done": 5,
+                    "<think>": 6,
+                    "</think>": 7,
+                    "changed": 8,
+                    "hello\nworld": 9,
+                    "<tool_call>": 10,
+                    "</tool_call>": 11,
+                    "<custom>": 12,
+                    "</custom>": 13,
+                },
+                unk_token="[UNK]",
+            )
+        ),
+        unk_token="[UNK]",
+    )
+    logged: list[tuple[dict[str, FakeTable], int]] = []
+    _install_fake_wandb(monkeypatch, logged)
+    triplets = [_triplet(prompt, response), _triplet(next_prompt, [5])]
+    for triplet in triplets:
+        routes = np.arange(
+            (len(triplet.prompt["token_ids"]) + len(triplet.response["token_ids"])) * 4,
+            dtype=np.uint8,
+        ).reshape(-1, 2, 2)
+        triplet.response["routed_experts"] = _encode_routes(routes)
+    rollout = CompletedRollout(
+        rollout_id="classification",
+        data_id="data-1",
+        step=0,
+        sample_idx_in_step=0,
+        enqueue_time=0.0,
+        final_reward=0.75,
+        triplets=triplets,
+    )
+    original = rollout.model_copy(deep=True)
+    adapter = RolloutAdapter(
+        max_prompt_length=12,
+        max_response_length=12,
+        device=torch.device("cpu"),
+        pad_token_id=0,
+        trace_aggregator_level="trajectory",
+        require_routed_experts=True,
+    )
+    baseline, baseline_metrics = adapter.get_train_data_batch([rollout])
+    logged.clear()
+    adapter.tokenizer = tokenizer if use_tokenizer else None
+    batch, metrics = adapter.get_train_data_batch([rollout])
+
+    table, _ = _logged_table(logged, "training/trace_merge_mismatches")
+    row = dict(zip(table.columns, table.rows[0], strict=True))
+    assert (row["template_mismatch"], row["retoken_mismatch"], row["others_mismatch"]) == expected
+    assert len(table.rows) == 1
+    assert len(batch) == 2  # Diagnostic classification never relaxes the exact token-prefix split.
+    assert rollout == original
+    assert metrics == baseline_metrics
+    assert batch.batch is not None and baseline.batch is not None
+    assert set(batch.batch.keys()) == set(baseline.batch.keys())
+    for key, value in baseline.batch.items():
+        torch.testing.assert_close(batch.batch[key], value)
+    for key, value in baseline.non_tensor_batch.items():
+        np.testing.assert_array_equal(batch.non_tensor_batch[key], value)
+
+
 def test_training_step_uploads_24_compact_rollout_trajectories_to_wandb_zip(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: