diff --git a/agentlightning/verl/rollout_adapter.py b/agentlightning/verl/rollout_adapter.py
index a898176b0..36b5b4340 100644
--- a/agentlightning/verl/rollout_adapter.py
+++ b/agentlightning/verl/rollout_adapter.py
@@ -19,6 +19,7 @@
_TRACE_MERGE_MISMATCH_WANDB_LIMIT = 100
_TRACE_MERGE_MISMATCH_TEXT_LIMIT = 4000
+_TRACE_MERGE_MISMATCH_TOKEN_LIMIT = 128
_ROLLOUT_TRAJECTORY_WANDB_LIMIT = 24
_TRACE_MERGE_MISMATCH_COLUMNS = [
"global_steps",
@@ -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.
+
+ Compare complete short inputs only; long or unrecognized cases stay other.
+ Template causes are not inferred, and no extra full-history decoding is done.
+ """
+ if (
+ tokenizer is None
+ or not previous_ids
+ or not prompt_ids
+ or len(previous_ids) > _TRACE_MERGE_MISMATCH_TOKEN_LIMIT
+ or len(prompt_ids) > _TRACE_MERGE_MISMATCH_TOKEN_LIMIT
+ ):
+ 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
+
+ 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),
diff --git a/docs/20-trainer-configuration.md b/docs/20-trainer-configuration.md
index 63af662b4..cc7aa8757 100644
--- a/docs/20-trainer-configuration.md
+++ b/docs/20-trainer-configuration.md
@@ -139,6 +139,10 @@ 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 `retoken_mismatch` when the decoded next prompt still starts with the decoded previous prompt and response. This check only runs when both complete token sequences contain at most 128 tokens, limiting the extra decoding to short inputs. Longer inputs, absent tokenizers, decoding failures and other unrecognized cases remain `others_mismatch`.
+
+`template_mismatch` remains false: no template normalization or root-cause inference is attempted. These flags are diagnostic hints, not proof of a cause. Classification never changes the strict token-prefix split or the resulting training tensors; long trajectories may therefore have retokenization mismatches that are still reported as `other`.
+
In this mode:
- `trajectory_max_prompt_length` limits the initial prompt in each merged training row;
diff --git a/tests/verl/test_rollout_adapter.py b/tests/verl/test_rollout_adapter.py
index 5b4bca405..f76d57b36 100644
--- a/tests/verl/test_rollout_adapter.py
+++ b/tests/verl/test_rollout_adapter.py
@@ -9,7 +9,7 @@
import weakref
import zipfile
from types import SimpleNamespace
-from typing import ClassVar
+from typing import Any, ClassVar
import pytest
@@ -336,6 +336,218 @@ 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, (False, False, True)),
+ ([1, 10], [2, 11], [1, 2, 4], True, (False, False, True)),
+ ([9], [5], [3, 5, 4], True, (False, False, True)),
+ ([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)),
+ ([1], [2], [1], True, (False, False, True)),
+ ([1], [2] + [4] * 126, [3] + [4] * 126, True, (False, True, False)),
+ ([1], [2] + [4] * 127, [3] + [4] * 127, True, (False, False, True)),
+ ([1], [2], [3] + [4] * 128, True, (False, False, True)),
+ ],
+ ids=[
+ "retoken",
+ "think-delimiters",
+ "tool-call-delimiters",
+ "whitespace",
+ "content",
+ "template-and-content",
+ "unknown-delimiters",
+ "empty-text",
+ "decode-error",
+ "no-tokenizer",
+ "response-cannot-complete-prefix",
+ "previous-128",
+ "previous-129",
+ "prompt-129",
+ ],
+)
+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,
+ "": 6,
+ "": 7,
+ "changed": 8,
+ "hello\nworld": 9,
+ "": 10,
+ "": 11,
+ "": 12,
+ "": 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, [2])]
+ 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)
+
+
+@pytest.mark.parametrize(
+ ("previous_size", "prompt_size", "expected_calls"),
+ [(128, 128, [128, 128]), (129, 128, []), (128, 129, []), (4096, 4096, [])],
+)
+def test_mismatch_classification_bounds_decoding(
+ previous_size: int, prompt_size: int, expected_calls: list[int]
+) -> None:
+ from tokenizers import Tokenizer, models
+ from transformers import PreTrainedTokenizerFast
+
+ tokenizer = PreTrainedTokenizerFast(
+ tokenizer_object=Tokenizer(models.WordLevel({"[UNK]": 0, "hello": 1}, unk_token="[UNK]")),
+ unk_token="[UNK]",
+ )
+ calls: list[int] = []
+
+ class CountedTokenizer:
+ def decode(self, ids: list[int], **kwargs: Any) -> str:
+ calls.append(len(ids))
+ text = tokenizer.decode(ids, **kwargs)
+ assert isinstance(text, str)
+ return text
+
+ result = rollout_adapter_module._classify_trace_merge_mismatch(
+ CountedTokenizer(), [1] * previous_size, [1] * prompt_size
+ )
+ assert calls == expected_calls
+ assert result == ((False, True, False) if expected_calls else (False, False, True))
+
+
+def test_mismatch_classification_does_not_use_response_or_numeric_fallback() -> None:
+ assert rollout_adapter_module._classify_trace_merge_mismatch(None, [1], [1]) == (False, False, True)
+ # The next response could complete a decoded prefix; only the prompt is eligible.
+ assert rollout_adapter_module._classify_trace_merge_mismatch(FakeTokenizer(), [1, 2], [1]) == (
+ False,
+ False,
+ True,
+ )
+
+
+def test_mismatch_classification_keeps_legacy_decoder_and_rejects_invalid_text() -> None:
+ class LegacyTokenizer:
+ def decode(self, ids: list[int]) -> str:
+ return " ".join(str(token) for token in ids)
+
+ assert rollout_adapter_module._classify_trace_merge_mismatch(LegacyTokenizer(), [1], [1, 2]) == (
+ False,
+ True,
+ False,
+ )
+ invalid = SimpleNamespace(decode=lambda ids, **kwargs: None)
+ assert rollout_adapter_module._classify_trace_merge_mismatch(invalid, [1], [1, 2]) == (False, False, True)
+
+
+@pytest.mark.parametrize(("context_size", "expected_decode_calls"), [(2, 402), (129, 202)])
+def test_mismatch_decoding_stops_at_table_limit(
+ monkeypatch: pytest.MonkeyPatch, context_size: int, expected_decode_calls: int
+) -> None:
+ from tokenizers import Tokenizer, models
+ from transformers import PreTrainedTokenizerFast
+
+ tokenizer = PreTrainedTokenizerFast(
+ tokenizer_object=Tokenizer(models.WordLevel({"[UNK]": 0, "hello": 1, "world": 2}, unk_token="[UNK]")),
+ unk_token="[UNK]",
+ )
+ calls: list[int] = []
+
+ class CountedTokenizer:
+ def decode(self, ids: list[int], **kwargs: Any) -> str:
+ calls.append(len(ids))
+ text = tokenizer.decode(ids, **kwargs)
+ assert isinstance(text, str)
+ return text
+
+ logged: list[tuple[dict[str, FakeTable], int]] = []
+ _install_fake_wandb(monkeypatch, logged)
+ rollout = CompletedRollout(
+ rollout_id="bounded",
+ data_id="data-1",
+ step=0,
+ sample_idx_in_step=0,
+ enqueue_time=0.0,
+ final_reward=0.5,
+ triplets=[_triplet([1 + turn % 2] * (context_size - 1), [1]) for turn in range(102)],
+ )
+ adapter = _adapter()
+ adapter.tokenizer = CountedTokenizer()
+ adapter.get_train_data_batch([rollout])
+ table, _ = _logged_table(logged, "training/trace_merge_mismatches")
+ assert len(table.rows) == 100
+ # Each row decodes two traces; the trajectory summary decodes its final prompt/response.
+ # Only short rows add two classifier decodes, and none are added beyond the 100-row cap.
+ assert len(calls) == expected_decode_calls
+
+
def test_training_step_uploads_24_compact_rollout_trajectories_to_wandb_zip(
monkeypatch: pytest.MonkeyPatch,
) -> None: