From 55e6d33623cabf17279856ce39f5b2ff5f221b4e Mon Sep 17 00:00:00 2001 From: DongzhuoranZhou <110855293+DongzhuoranZhou@users.noreply.github.com> Date: Thu, 16 Jul 2026 11:12:37 +0200 Subject: [PATCH 1/4] fix tau-bench token deltas for reasoning templates Co-authored-by: Yuki <70699639+Y-L-LIU@users.noreply.github.com> --- examples/tau-bench/token_delta.py | 35 +++++++++++ examples/tau-bench/trainable_agents.py | 20 +------ tests/test_tau_bench_token_delta.py | 80 ++++++++++++++++++++++++++ 3 files changed, 117 insertions(+), 18 deletions(-) create mode 100644 examples/tau-bench/token_delta.py create mode 100644 tests/test_tau_bench_token_delta.py diff --git a/examples/tau-bench/token_delta.py b/examples/tau-bench/token_delta.py new file mode 100644 index 0000000000..4ba04f0643 --- /dev/null +++ b/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) diff --git a/examples/tau-bench/trainable_agents.py b/examples/tau-bench/trainable_agents.py index 5cfc3f43d0..6685de5144 100644 --- a/examples/tau-bench/trainable_agents.py +++ b/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, diff --git a/tests/test_tau_bench_token_delta.py b/tests/test_tau_bench_token_delta.py new file mode 100644 index 0000000000..8bf971d8f2 --- /dev/null +++ b/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".*?", "", 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}') + if add_generation_prompt: + rendered.append("") + 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": "first reasoningfirst answer"}, + {"role": "user", "content": "second user must remain complete"}, + ] + + token_ids, loss_mask = get_token_delta(tokenizer, messages) + + assert tokenizer.decode(token_ids) == "second user must remain complete" + 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": "reasoninganswer"}, + ] + + token_ids, loss_mask = get_token_delta(tokenizer, messages) + + assert tokenizer.decode(token_ids) == "reasoninganswer" + assert loss_mask == [1] * len(token_ids) From 9af43e2d822cc6b956fbe022ee654f8af280347a Mon Sep 17 00:00:00 2001 From: DongzhuoranZhou <110855293+DongzhuoranZhou@users.noreply.github.com> Date: Fri, 17 Jul 2026 23:36:47 +0800 Subject: [PATCH 2/4] fix later tau-bench assistant generation prefixes Co-authored-by: Yuki <70699639+Y-L-LIU@users.noreply.github.com> --- examples/tau-bench/token_delta.py | 41 ++++++++++--- examples/tau-bench/trainable_agents.py | 20 ++++++- tests/test_tau_bench_token_delta.py | 83 ++++++++++++++++++++++++++ 3 files changed, 133 insertions(+), 11 deletions(-) diff --git a/examples/tau-bench/token_delta.py b/examples/tau-bench/token_delta.py index 4ba04f0643..deb14667c7 100644 --- a/examples/tau-bench/token_delta.py +++ b/examples/tau-bench/token_delta.py @@ -1,18 +1,44 @@ from typing import Any -def get_token_delta(tokenizer: Any, messages: list[dict[str, Any]]) -> tuple[list[int], list[int]]: +def get_token_delta( + tokenizer: Any, + messages: list[dict[str, Any]], + *, + include_generation_prompt: bool = False, +) -> 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 is_assistant: + prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=False, tokenize=False) + generation_prompt = tokenizer.apply_chat_template( + messages[:-1], add_generation_prompt=True, tokenize=False + ) + if not generation_prompt.startswith(prev): + raise ValueError("Adding the assistant generation prompt rewrote the rendered conversation") + if not curr.startswith(generation_prompt): + raise ValueError("The assistant response does not extend its generation prompt") + + generation_prompt_text = generation_prompt[len(prev) :] + if not include_generation_prompt: + new_text = curr[len(generation_prompt) :] + new_tokens = tokenizer.encode(new_text, add_special_tokens=False) + return new_tokens, [1] * len(new_tokens) + + new_text = curr[len(prev) :] + new_tokens = tokenizer.encode(new_text, add_special_tokens=False) + generation_prompt_length = len(tokenizer.encode(generation_prompt_text, add_special_tokens=False)) + masked_prefix_length = min(generation_prompt_length, len(new_tokens)) + loss_mask = [0] * masked_prefix_length + loss_mask.extend([1] * (len(new_tokens) - masked_prefix_length)) + return new_tokens, loss_mask + + prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=False, tokenize=False) if curr.startswith(prev): new_text = curr[len(prev) :] @@ -31,5 +57,4 @@ def get_token_delta(tokenizer: Any, messages: list[dict[str, Any]]) -> tuple[lis 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) + return new_tokens, [0] * len(new_tokens) diff --git a/examples/tau-bench/trainable_agents.py b/examples/tau-bench/trainable_agents.py index 6685de5144..effc5d7a66 100644 --- a/examples/tau-bench/trainable_agents.py +++ b/examples/tau-bench/trainable_agents.py @@ -275,7 +275,11 @@ async def asolve( # Add assistant response to conversation messages.append({"role": "assistant", "content": response}) - assistant_token_ids, assistant_loss_mask = self._get_token_delta(state.tokenizer, messages) + assistant_token_ids, assistant_loss_mask = self._get_token_delta( + state.tokenizer, + messages, + include_generation_prompt=bool(response_token_ids), + ) response_token_ids.extend(assistant_token_ids) loss_masks.extend(assistant_loss_mask) @@ -332,7 +336,13 @@ async def asolve( res, total_reward, info, messages, loss_masks, prompt_token_ids, response_token_ids ) - def _get_token_delta(self, tokenizer: AutoTokenizer, messages: list[dict]) -> tuple[list[int], list[int]]: + def _get_token_delta( + self, + tokenizer: AutoTokenizer, + messages: list[dict], + *, + include_generation_prompt: bool = False, + ) -> tuple[list[int], list[int]]: """ Calculate token delta for multi-turn conversations. @@ -348,7 +358,11 @@ def _get_token_delta(self, tokenizer: AutoTokenizer, messages: list[dict]) -> tu Returns: Tuple of (token_ids, loss_mask) """ - return get_token_delta(tokenizer, messages) + return get_token_delta( + tokenizer, + messages, + include_generation_prompt=include_generation_prompt, + ) def _build_final_result( self, diff --git a/tests/test_tau_bench_token_delta.py b/tests/test_tau_bench_token_delta.py index 8bf971d8f2..377fd26e6a 100644 --- a/tests/test_tau_bench_token_delta.py +++ b/tests/test_tau_bench_token_delta.py @@ -49,6 +49,25 @@ def decode(token_ids): return bytes(token_ids).decode() +class BoundaryMergingTokenizer(HistoryRewritingTokenizer): + """Tokenizer where the generation-prefix tail merges with a leading newline.""" + + @staticmethod + def encode(text, *, add_special_tokens): + assert add_special_tokens is False + raw = text.encode() + token_ids = [] + index = 0 + while index < len(raw): + if raw[index : index + 2] == b">\n": + token_ids.append(1000) + index += 2 + else: + token_ids.append(raw[index]) + index += 1 + return token_ids + + @pytest.mark.unit def test_new_user_delta_survives_history_rewrite(): get_token_delta = _load_get_token_delta() @@ -78,3 +97,67 @@ def test_assistant_delta_keeps_existing_append_only_behavior(): assert tokenizer.decode(token_ids) == "reasoninganswer" assert loss_mask == [1] * len(token_ids) + + +@pytest.mark.unit +def test_accumulated_multiturn_tokens_keep_every_assistant_generation_prefix(): + get_token_delta = _load_get_token_delta() + tokenizer = HistoryRewritingTokenizer() + messages = [ + {"role": "user", "content": "first user"}, + {"role": "assistant", "content": "first reasoningfirst answer"}, + {"role": "user", "content": "second user"}, + {"role": "assistant", "content": "second reasoningsecond answer"}, + ] + + initial_prompt = tokenizer.apply_chat_template( + messages[:1], add_generation_prompt=True, tokenize=False + ) + token_ids = tokenizer.encode(initial_prompt, add_special_tokens=False) + loss_mask = [0] * len(token_ids) + + for end in range(2, len(messages) + 1): + include_generation_prompt = messages[end - 1]["role"] == "assistant" and end > 2 + delta_ids, delta_mask = get_token_delta( + tokenizer, + messages[:end], + include_generation_prompt=include_generation_prompt, + ) + token_ids.extend(delta_ids) + loss_mask.extend(delta_mask) + + decoded = tokenizer.decode(token_ids) + assert decoded.count("") == 2 + assert decoded.count("") == 2 + assert "first reasoning" in decoded + assert "second reasoning" in decoded + assert "second user" in decoded + + second_prefix = decoded.rindex("") + assert loss_mask[second_prefix : second_prefix + len("")] == [0] * len("") + + +@pytest.mark.unit +def test_later_assistant_allows_bpe_merge_across_generation_prefix_boundary(): + get_token_delta = _load_get_token_delta() + tokenizer = BoundaryMergingTokenizer() + messages = [ + {"role": "user", "content": "first user"}, + {"role": "assistant", "content": "first reasoningfirst answer"}, + {"role": "user", "content": "second user"}, + {"role": "assistant", "content": "\nsecond answer"}, + ] + + token_ids, loss_mask = get_token_delta( + tokenizer, + messages, + include_generation_prompt=True, + ) + + expected_text = "\nsecond answer" + expected_ids = tokenizer.encode(expected_text, add_special_tokens=False) + generation_prefix_length = len(tokenizer.encode("", add_special_tokens=False)) + assert token_ids == expected_ids + assert loss_mask == [0] * generation_prefix_length + [1] * ( + len(expected_ids) - generation_prefix_length + ) From 11d3427ff7981678d957cb5073503cd9a554429d Mon Sep 17 00:00:00 2001 From: DongzhuoranZhou <110855293+DongzhuoranZhou@users.noreply.github.com> Date: Sat, 18 Jul 2026 00:37:26 +0800 Subject: [PATCH 3/4] format tau-bench token delta fix --- examples/tau-bench/token_delta.py | 4 +--- tests/test_tau_bench_token_delta.py | 8 ++------ 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/examples/tau-bench/token_delta.py b/examples/tau-bench/token_delta.py index deb14667c7..b78023d8ce 100644 --- a/examples/tau-bench/token_delta.py +++ b/examples/tau-bench/token_delta.py @@ -16,9 +16,7 @@ def get_token_delta( if is_assistant: prev = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=False, tokenize=False) - generation_prompt = tokenizer.apply_chat_template( - messages[:-1], add_generation_prompt=True, tokenize=False - ) + generation_prompt = tokenizer.apply_chat_template(messages[:-1], add_generation_prompt=True, tokenize=False) if not generation_prompt.startswith(prev): raise ValueError("Adding the assistant generation prompt rewrote the rendered conversation") if not curr.startswith(generation_prompt): diff --git a/tests/test_tau_bench_token_delta.py b/tests/test_tau_bench_token_delta.py index 377fd26e6a..7593f60857 100644 --- a/tests/test_tau_bench_token_delta.py +++ b/tests/test_tau_bench_token_delta.py @@ -110,9 +110,7 @@ def test_accumulated_multiturn_tokens_keep_every_assistant_generation_prefix(): {"role": "assistant", "content": "second reasoningsecond answer"}, ] - initial_prompt = tokenizer.apply_chat_template( - messages[:1], add_generation_prompt=True, tokenize=False - ) + initial_prompt = tokenizer.apply_chat_template(messages[:1], add_generation_prompt=True, tokenize=False) token_ids = tokenizer.encode(initial_prompt, add_special_tokens=False) loss_mask = [0] * len(token_ids) @@ -158,6 +156,4 @@ def test_later_assistant_allows_bpe_merge_across_generation_prefix_boundary(): expected_ids = tokenizer.encode(expected_text, add_special_tokens=False) generation_prefix_length = len(tokenizer.encode("", add_special_tokens=False)) assert token_ids == expected_ids - assert loss_mask == [0] * generation_prefix_length + [1] * ( - len(expected_ids) - generation_prefix_length - ) + assert loss_mask == [0] * generation_prefix_length + [1] * (len(expected_ids) - generation_prefix_length) From 6641542515e22972bc5d3eecd65a018243c3e006 Mon Sep 17 00:00:00 2001 From: DongzhuoranZhou <110855293+DongzhuoranZhou@users.noreply.github.com> Date: Sat, 18 Jul 2026 11:23:55 +0800 Subject: [PATCH 4/4] test tau-bench token deltas across all user turns Co-authored-by: Y-L-LIU --- tests/test_tau_bench_token_delta.py | 48 +++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tests/test_tau_bench_token_delta.py b/tests/test_tau_bench_token_delta.py index 7593f60857..2b88c41120 100644 --- a/tests/test_tau_bench_token_delta.py +++ b/tests/test_tau_bench_token_delta.py @@ -135,6 +135,54 @@ def test_accumulated_multiturn_tokens_keep_every_assistant_generation_prefix(): assert loss_mask[second_prefix : second_prefix + len("")] == [0] * len("") +@pytest.mark.unit +def test_accumulated_six_real_user_turns_keep_every_user_and_assistant_boundary(): + get_token_delta = _load_get_token_delta() + tokenizer = HistoryRewritingTokenizer() + messages = [{"role": "user", "content": "user 1"}] + + initial_prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + token_ids = tokenizer.encode(initial_prompt, add_special_tokens=False) + loss_mask = [0] * len(token_ids) + + for turn in range(1, 7): + messages.append( + { + "role": "assistant", + "content": f"reasoning {turn}answer {turn}", + } + ) + delta_ids, delta_mask = get_token_delta( + tokenizer, + messages, + include_generation_prompt=turn > 1, + ) + token_ids.extend(delta_ids) + loss_mask.extend(delta_mask) + + if turn < 6: + messages.append({"role": "user", "content": f"user {turn + 1}"}) + delta_ids, delta_mask = get_token_delta(tokenizer, messages) + token_ids.extend(delta_ids) + loss_mask.extend(delta_mask) + + decoded = tokenizer.decode(token_ids) + assert decoded.count("") == 6 + assert decoded.count("") == 6 + + for turn in range(1, 7): + user_span = f"user {turn}" + user_start = decoded.index(user_span) + assert loss_mask[user_start : user_start + len(user_span)] == [0] * len(user_span) + + assistant_span = f"reasoning {turn}answer {turn}" + assistant_start = decoded.index(assistant_span) + assert loss_mask[assistant_start : assistant_start + len(assistant_span)] == [1] * len(assistant_span) + + prefix_start = decoded.rfind("", 0, assistant_start) + assert loss_mask[prefix_start:assistant_start] == [0] * (assistant_start - prefix_start) + + @pytest.mark.unit def test_later_assistant_allows_bpe_merge_across_generation_prefix_boundary(): get_token_delta = _load_get_token_delta()