fix tau-bench token deltas for reasoning templates

Co-authored-by: Yuki <70699639+Y-L-LIU@users.noreply.github.com>

Author: DongzhuoranZhou

examples/tau-bench/token_delta.py

@@ -0,0 +1,35 @@
+from typing import Any
+
+
+def get_token_delta(tokenizer: Any, messages: list[dict[str, Any]]) -> tuple[list[int], list[int]]:
+    """Return the tokens and loss mask contributed by the last chat message."""
+    if not messages:
+        raise ValueError("Cannot calculate a token delta for an empty conversation")
+
+    is_assistant = messages[-1]["role"] == "assistant"
+    curr = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False)
+    prev = tokenizer.apply_chat_template(
+        messages[:-1],
+        add_generation_prompt=is_assistant,
+        tokenize=False,
+    )
+
+    if curr.startswith(prev):
+        new_text = curr[len(prev) :]
+    elif messages[-1]["role"] == "user":
+        # Reasoning templates such as Qwen3 can rewrite history when a new user
+        # message arrives. Render that message independently instead of slicing
+        # the rewritten conversation at the old conversation length.
+        new_text = tokenizer.apply_chat_template(
+            [messages[-1]],
+            add_generation_prompt=False,
+            tokenize=False,
+        )
+        if not curr.endswith(new_text):
+            raise ValueError("The latest user message is not a standalone suffix of the rendered conversation")
+    else:
+        raise ValueError("The chat template rewrote history while calculating a non-user token delta")
+
+    new_tokens = tokenizer.encode(new_text, add_special_tokens=False)
+    loss_value = 1 if is_assistant else 0
+    return new_tokens, [loss_value] * len(new_tokens)

examples/tau-bench/trainable_agents.py

@@ -8,6 +8,7 @@
 from tau_bench.agents.base import Agent
 from tau_bench.agents.tool_calling_agent import RESPOND_ACTION_NAME, ToolCallingAgent
 from tau_bench.types import Action, RunConfig
+from token_delta import get_token_delta
 from transformers import AutoTokenizer
 
 from slime.rollout.sglang_rollout import GenerateState
@@ -347,24 +348,7 @@ def _get_token_delta(self, tokenizer: AutoTokenizer, messages: list[dict]) -> tu
         Returns:
             Tuple of (token_ids, loss_mask)
         """
-        curr = tokenizer.apply_chat_template(messages, add_generation_prompt=False, tokenize=False)
-        token_ids = []
-        loss_mask = []
-
-        # Case 1: last message is an assistant response
-        if messages[-1]["role"] == "assistant":
-            prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=True, tokenize=False)
-            new_tokens = tokenizer.encode(curr[len(prev) :], add_special_tokens=False)
-            token_ids += new_tokens
-            loss_mask += [1] * len(new_tokens)  # Mask only the new assistant tokens
-        else:
-            # Case 2: last message is a tool response or environment observation
-            prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=False, tokenize=False)
-            new_tokens = tokenizer.encode(curr[len(prev) :], add_special_tokens=False)
-            token_ids += new_tokens
-            loss_mask += [0] * len(new_tokens)  # Don't mask environment/tool tokens
-
-        return token_ids, loss_mask
+        return get_token_delta(tokenizer, messages)
 
     def _build_final_result(
         self,

tests/test_tau_bench_token_delta.py

@@ -0,0 +1,80 @@
+import importlib.util
+import re
+from pathlib import Path
+
+import pytest
+
+
+TOKEN_DELTA_PATH = Path(__file__).parents[1] / "examples" / "tau-bench" / "token_delta.py"
+
+
+def _load_get_token_delta():
+    if not TOKEN_DELTA_PATH.exists():
+        pytest.fail("tau-bench token_delta helper does not exist")
+
+    spec = importlib.util.spec_from_file_location("tau_bench_token_delta", TOKEN_DELTA_PATH)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    spec.loader.exec_module(module)
+    return module.get_token_delta
+
+
+class HistoryRewritingTokenizer:
+    """Minimal chat template that hides old reasoning after a new user turn."""
+
+    @staticmethod
+    def _strip_reasoning(content: str) -> str:
+        return re.sub(r"<think>.*?</think>", "", content, flags=re.DOTALL)
+
+    def apply_chat_template(self, messages, *, add_generation_prompt, tokenize):
+        assert tokenize is False
+        last_user = max((i for i, message in enumerate(messages) if message["role"] == "user"), default=-1)
+        rendered = []
+        for i, message in enumerate(messages):
+            content = message["content"]
+            if message["role"] == "assistant" and i < last_user:
+                content = self._strip_reasoning(content)
+            rendered.append(f'<{message["role"]}>{content}</{message["role"]}>')
+        if add_generation_prompt:
+            rendered.append("<assistant>")
+        return "".join(rendered)
+
+    @staticmethod
+    def encode(text, *, add_special_tokens):
+        assert add_special_tokens is False
+        return list(text.encode())
+
+    @staticmethod
+    def decode(token_ids):
+        return bytes(token_ids).decode()
+
+
+@pytest.mark.unit
+def test_new_user_delta_survives_history_rewrite():
+    get_token_delta = _load_get_token_delta()
+    tokenizer = HistoryRewritingTokenizer()
+    messages = [
+        {"role": "user", "content": "first user"},
+        {"role": "assistant", "content": "<think>first reasoning</think>first answer"},
+        {"role": "user", "content": "second user must remain complete"},
+    ]
+
+    token_ids, loss_mask = get_token_delta(tokenizer, messages)
+
+    assert tokenizer.decode(token_ids) == "<user>second user must remain complete</user>"
+    assert loss_mask == [0] * len(token_ids)
+
+
+@pytest.mark.unit
+def test_assistant_delta_keeps_existing_append_only_behavior():
+    get_token_delta = _load_get_token_delta()
+    tokenizer = HistoryRewritingTokenizer()
+    messages = [
+        {"role": "user", "content": "user question"},
+        {"role": "assistant", "content": "<think>reasoning</think>answer"},
+    ]
+
+    token_ids, loss_mask = get_token_delta(tokenizer, messages)
+
+    assert tokenizer.decode(token_ids) == "<think>reasoning</think>answer</assistant>"
+    assert loss_mask == [1] * len(token_ids)