diff --git a/examples/tau-bench/README.md b/examples/tau-bench/README.md
new file mode 100644
index 000000000..85611d128
--- /dev/null
+++ b/examples/tau-bench/README.md
@@ -0,0 +1,65 @@
+# Tau bench
+This example shows vime training in an agentic multi-turn tool use environment.
+
+
+## Environment Setup
+This example assumes a vime container image. Install tau-bench dependencies:
+
+```bash
+cd /root/
+git clone https://github.com/JD-ETH/tau-bench.git
+cd tau-bench
+git checkout feature/litellm-retry
+pip install -e . --no-deps
+pip install litellm
+```
+
+Use the following script to generate task index jsonl for training:
+
+```bash
+cd /root/vime/examples/tau-bench
+python tau1_mock.py --local_dir /root/tau-bench/
+```
+
+Initialize the Qwen3-4B-Instruct-2507 model needed for tool use:
+
+```bash
+# hf checkpoint
+hf download Qwen/Qwen3-4B-Instruct-2507 --local-dir /root/Qwen3-4B-Instruct-2507
+
+# mcore checkpoint
+cd /root/vime
+source scripts/models/qwen3-4B-Instruct-2507.sh
+PYTHONPATH=/root/Megatron-LM python tools/convert_hf_to_torch_dist.py \
+ ${MODEL_ARGS[@]} \
+ --hf-checkpoint /root/Qwen3-4B-Instruct-2507 \
+ --save /root/Qwen3-4B-Instruct-2507_torch_dist
+```
+
+## Running the Script
+
+You need to configure your litellm API in generate_with_tau.py for user simulation:
+
+TAU_CONFIGS = {
+ "env": "retail", # Select between ["retail", "airline"]
+ "agent_strategy": "tool-calling", # Select between ["tool-calling", "act", "react", "few-shot"], only tool-calling implemented for now
+ "user_model": "gemini-2.0-flash-lite", # Cheap Model for user simulator
+ "user_model_provider": "gemini",
+ "task_split": "train", # Select between ["train", "test", "dev"] for retail, ["test"] for airline
+ "user_strategy": "llm", # Select between ["llm", "react", "verify", "reflection"]
+ "model_provider": "auto_router", # Unused, required
+ "model": "qwen3-4b", # Unused, reqired
+}
+# Replace with your actual API key for user sim
+GEMINI_API_KEY = "YOUR KEY"
+
+Multi-turn limit: set env `TAU_MAX_TURNS` (default 10) or pass `--max-turns` to train.py.
+
+Agent rollout always uses vLLM (`/inference/v1/generate`); only `TAU_CONFIGS` controls the user simulator.
+
+And run:
+
+```bash
+cd /root/vime
+bash examples/tau-bench/run_qwen3_4B.sh
+```
diff --git a/examples/tau-bench/generate_with_tau.py b/examples/tau-bench/generate_with_tau.py
new file mode 100644
index 000000000..3c331794e
--- /dev/null
+++ b/examples/tau-bench/generate_with_tau.py
@@ -0,0 +1,105 @@
+"""Tau-bench multi-turn custom rollout for vime (vLLM render + generate)."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import os
+from typing import Any
+
+from tau_bench.types import RunConfig
+from trainable_agents import TrainableTauBenchAgent, agent_factory, patch_tau_user_retries
+
+from vime.utils.types import Sample
+
+logger = logging.getLogger(__name__)
+
+_TAU_DEFAULT_MAX_TURNS = 10
+
+_inflight_sem: asyncio.Semaphore | None = None
+
+# Tau-bench user-simulator configuration (edit TAU_CONFIGS below).
+# Agent rollout uses vLLM; only user_model / user_model_provider affect the user simulator here.
+TAU_CONFIGS = {
+ "env": "retail", # Select between ["retail", "airline"]
+ "agent_strategy": "tool-calling", # Select between ["tool-calling", "act", "react", "few-shot"]
+ # Default: local vLLM user sim (no external API). For Gemini API user sim, switch to:
+ # "user_model": "gemini-2.5-flash-lite", "user_model_provider": "gemini",
+ "user_model": "openai/local-qwen3-4b",
+ "user_model_provider": "openai",
+ "task_split": "train", # Select between ["train", "test", "dev"] for retail
+ "user_strategy": "llm", # Select between ["llm", "react", "verify", "reflection"]
+ "model_provider": "auto_router", # Unused, required
+ "model": "qwen3-4b", # Unused, required
+}
+# Replace with your actual API key when user_model_provider is gemini.
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY", "NONE")
+os.environ["GEMINI_API_KEY"] = GEMINI_API_KEY
+tau_config = RunConfig(**TAU_CONFIGS)
+
+
+def _get_inflight_sem() -> asyncio.Semaphore:
+ global _inflight_sem
+ if _inflight_sem is None:
+ _inflight_sem = asyncio.Semaphore(int(os.environ.get("TAU_MAX_INFLIGHT", "8")))
+ return _inflight_sem
+
+
+patch_tau_user_retries()
+
+
+def _ensure_tau_args(args: Any) -> None:
+ if getattr(args, "max_turns", None) is None:
+ env_max = os.environ.get("TAU_MAX_TURNS")
+ args.max_turns = int(env_max) if env_max is not None else _TAU_DEFAULT_MAX_TURNS
+
+
+def resolve_tau_config(args: Any) -> RunConfig:
+ """Build RunConfig from TAU_CONFIGS, with optional local-vLLM user-sim routing."""
+ user_model = tau_config.user_model
+ user_model_provider = tau_config.user_model_provider
+
+ if user_model_provider == "openai" and "local" in user_model:
+ vllm_router_host = getattr(args, "vllm_router_ip", "127.0.0.1")
+ vllm_router_port = getattr(args, "vllm_router_port", 3250)
+ vllm_model_name = getattr(args, "vllm_model_name", getattr(args, "hf_checkpoint", ""))
+ os.environ["OPENAI_API_KEY"] = os.environ.get("OPENAI_API_KEY", "dummy")
+ os.environ["OPENAI_API_BASE"] = f"http://{vllm_router_host}:{vllm_router_port}/v1"
+ user_model = vllm_model_name
+
+ return RunConfig(
+ env=tau_config.env,
+ agent_strategy=tau_config.agent_strategy,
+ user_model=user_model,
+ user_model_provider=user_model_provider,
+ task_split=tau_config.task_split,
+ user_strategy=tau_config.user_strategy,
+ model_provider=tau_config.model_provider,
+ model=tau_config.model,
+ )
+
+
+async def batched_tau_bench_rm(args, samples, **kwargs) -> list[float] | float:
+ if isinstance(samples, Sample):
+ return samples.reward if samples.reward is not None else 0.0
+ rewards = [s.reward if s.reward is not None else 0.0 for s in samples]
+ max_r = max(rewards) if rewards else 1.0
+ if max_r > 0:
+ rewards = [r / max_r for r in rewards]
+ return rewards
+
+
+async def generate(args: Any, sample: Sample, sampling_params) -> Sample:
+ assert not args.partial_rollout, "Partial rollout is not supported for tau-bench interactions."
+ _ensure_tau_args(args)
+ args.tau_bench_config = resolve_tau_config(args)
+
+ task_index = sample.prompt
+ logger.info(f"Starting agent-environment interaction for task {task_index}")
+
+ async with _get_inflight_sem():
+ agent: TrainableTauBenchAgent = agent_factory()
+ result = await agent.asolve(args, sample, sampling_params)
+
+ logger.info(f"Finished agent-environment interaction for task {task_index}")
+ return result
diff --git a/examples/tau-bench/openai_tool_adapter.py b/examples/tau-bench/openai_tool_adapter.py
new file mode 100644
index 000000000..53c744f08
--- /dev/null
+++ b/examples/tau-bench/openai_tool_adapter.py
@@ -0,0 +1,67 @@
+import logging
+from dataclasses import dataclass, field
+from typing import Any
+
+try:
+ from .vllm_tool_parser import parse_tools
+except ImportError:
+ from vllm_tool_parser import parse_tools
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class OpenAIToolCall:
+ id: str
+ type: str = "function"
+ function: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class OpenAIAssistantMessage:
+ role: str = "assistant"
+ content: str | None = None
+ tool_calls: list[OpenAIToolCall] | None = None
+
+
+class OpenAICompatibleToolCallAdapter:
+ def __init__(self, tools_info: list[dict[str, Any]], parser_type: str = "qwen25"):
+ self.tools_info = tools_info
+ self.parser_type = parser_type
+
+ def parse_response_to_openai_format(self, response: str) -> dict[str, Any]:
+ try:
+ parsed = parse_tools(response, self.tools_info, self.parser_type)
+ normal_text = parsed["normal_text"]
+ calls = parsed["calls"]
+ openai_message = self._convert_to_openai_message(normal_text, calls)
+ return {"openai_message": openai_message, "parsed_result": parsed, "success": True}
+ except Exception as e:
+ logger.warning(f"Parsing failed with error: {e}")
+ return {"openai_message": None, "parsed_result": None, "success": False, "error": str(e)}
+
+ def _convert_to_openai_message(self, normal_text: str, calls: list[dict[str, Any]]) -> OpenAIAssistantMessage:
+ if not calls:
+ return OpenAIAssistantMessage(role="assistant", content=normal_text, tool_calls=None)
+
+ openai_tool_calls = []
+ for i, call in enumerate(calls):
+ openai_tool_calls.append(
+ OpenAIToolCall(
+ id=f"call_{i}_{call.get('name', 'unknown')}",
+ type="function",
+ function={"name": call.get("name", ""), "arguments": call.get("parameters", "{}")},
+ )
+ )
+
+ return OpenAIAssistantMessage(
+ role="assistant",
+ content=normal_text if normal_text.strip() else None,
+ tool_calls=openai_tool_calls,
+ )
+
+
+def create_openai_adapter(
+ tools_info: list[dict[str, Any]], parser_type: str = "qwen25"
+) -> OpenAICompatibleToolCallAdapter:
+ return OpenAICompatibleToolCallAdapter(tools_info, parser_type)
diff --git a/examples/tau-bench/run_qwen3_4B_npu.sh b/examples/tau-bench/run_qwen3_4B_npu.sh
new file mode 100644
index 000000000..c5bad23fc
--- /dev/null
+++ b/examples/tau-bench/run_qwen3_4B_npu.sh
@@ -0,0 +1,148 @@
+#!/bin/bash
+
+if grep -q $'\r' "$0" 2>/dev/null; then
+ exec bash <(sed 's/\r$//' "$0") "$@"
+fi
+
+# for rerun the task
+pkill -9 vllm 2>/dev/null || true
+pkill -9 VLLM 2>/dev/null || true
+sleep 3
+ray stop --force 2>/dev/null || true
+pkill -9 ray 2>/dev/null || true
+pkill -9 -f 'python3 train.py' 2>/dev/null || true
+sleep 3
+
+set -ex
+
+export PYTHONUNBUFFERED=1
+export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15
+export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1
+export CUDA_DEVICE_MAX_CONNECTIONS=1
+export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050
+export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050
+export HYDRA_FULL_ERROR=1
+export VLLM_ASCEND_ENABLE_NZ=0
+export VLLM_USE_AOT_COMPILE=0
+export VIME_VLLM_SERVER_HEALTH_TIMEOUT_SEC=900
+
+unset PYTORCH_CUDA_ALLOC_CONF PYTORCH_ALLOC_CONF
+unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY
+
+SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" &>/dev/null && pwd)"
+source "${SCRIPT_DIR}/../../scripts/models/qwen3-4B-Instruct-2507.sh"
+
+export PYTHONPATH="${SCRIPT_DIR}:/root/Megatron-Bridge/src:/root/Megatron-LM:${PYTHONPATH:-}"
+
+DATA_ROOT="${DATA_ROOT:-/root}"
+TAU_BENCH_ROOT="${TAU_BENCH_ROOT:-/root/tau-bench}"
+
+CKPT_ARGS=(
+ --hf-checkpoint ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/
+ --load ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/
+ --ref-load ${DATA_ROOT}/weights/Qwen3-4B-Instruct-2507/
+ --megatron-to-hf-mode bridge
+)
+
+ROLLOUT_ARGS=(
+ --prompt-data ${TAU_BENCH_ROOT}/retail_train_tasks.jsonl
+ --input-key index
+ --rollout-shuffle
+ --num-rollout 500
+ --rollout-batch-size 32
+ --n-samples-per-prompt 8
+ --rollout-max-response-len 4096
+ --rollout-max-context-len 16384
+ --rollout-temperature 0.7
+ --global-batch-size 256
+ --dynamic-sampling-filter-path vime.rollout.filter_hub.dynamic_sampling_filters.check_reward_nonzero_std
+ --balance-data
+)
+
+EVAL_ARGS=(
+ --eval-interval 5
+ --eval-prompt-data retail-dev ${TAU_BENCH_ROOT}/retail_dev_tasks.jsonl
+ --n-samples-per-eval-prompt 1
+ --eval-max-response-len 4096
+ --eval-top-k 1
+)
+
+PERF_ARGS=(
+ --tensor-model-parallel-size 2
+ --sequence-parallel
+ --pipeline-model-parallel-size 1
+ --context-parallel-size 1
+ --expert-model-parallel-size 1
+ --expert-tensor-parallel-size 1
+ --recompute-granularity full
+ --recompute-method uniform
+ --recompute-num-layers 1
+ --use-dynamic-batch-size
+ --max-tokens-per-gpu 9216
+)
+
+GRPO_ARGS=(
+ --advantage-estimator grpo
+ --use-kl-loss
+ --kl-loss-coef 0.001
+ --kl-loss-type low_var_kl
+ --entropy-coef 0.01
+ --eps-clip 0.2
+ --eps-clip-high 0.28
+)
+
+OPTIMIZER_ARGS=(
+ --optimizer adam
+ --lr 5e-6
+ --lr-decay-style constant
+ --weight-decay 0.1
+ --adam-beta1 0.9
+ --adam-beta2 0.98
+)
+
+VLLM_ARGS=(
+ --rollout-num-gpus-per-engine 1
+ --vllm-gpu-memory-utilization 0.7
+ --vllm-max-model-len 16384
+)
+
+MISC_ARGS=(
+ --attention-dropout 0.0
+ --hidden-dropout 0.0
+ --accumulate-allreduce-grads-in-fp32
+ --attention-softmax-in-fp32
+ --attention-backend flash
+ --use-flash-attn
+ --no-gradient-accumulation-fusion
+)
+
+CUSTOM_ARGS=(
+ --custom-generate-function-path generate_with_tau.generate
+ --custom-rm-path generate_with_tau.batched_tau_bench_rm
+)
+
+export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
+
+ray start --head \
+ --node-ip-address "${MASTER_ADDR}" \
+ --disable-usage-stats \
+ --dashboard-host=0.0.0.0 \
+ --dashboard-port=8265
+
+ray job submit --address="http://127.0.0.1:8265" \
+ -- python3 train.py \
+ --train-backend megatron \
+ --actor-num-nodes 1 \
+ --actor-num-gpus-per-node 4 \
+ --rollout-num-gpus 4 \
+ "${MODEL_ARGS[@]}" \
+ "${CKPT_ARGS[@]}" \
+ "${ROLLOUT_ARGS[@]}" \
+ "${EVAL_ARGS[@]}" \
+ "${OPTIMIZER_ARGS[@]}" \
+ "${GRPO_ARGS[@]}" \
+ "${PERF_ARGS[@]}" \
+ "${VLLM_ARGS[@]}" \
+ "${CUSTOM_ARGS[@]}" \
+ "${MISC_ARGS[@]}"
+
diff --git a/examples/tau-bench/tau1_mock.py b/examples/tau-bench/tau1_mock.py
new file mode 100644
index 000000000..4be62d629
--- /dev/null
+++ b/examples/tau-bench/tau1_mock.py
@@ -0,0 +1,39 @@
+import argparse
+import json
+import os
+
+from tau_bench.envs import get_env
+from tau_bench.types import RunConfig
+
+ALL_DATA_MAPPINGS = {"retail": ["train", "test", "dev"], "airline": ["test"]}
+
+
+def main():
+ parser = argparse.ArgumentParser(description="Tau1 Mock Script")
+ parser.add_argument("--local_dir", required=True, help="Path to the local directory")
+ args = parser.parse_args()
+
+ local_dir = args.local_dir
+ if not os.path.isdir(local_dir):
+ os.makedirs(local_dir)
+ config = RunConfig(model_provider="mock", user_model_provider="mock", user_strategy="human", model="mock")
+ for env, split in ALL_DATA_MAPPINGS.items():
+ for s in split:
+ config.env = env
+ config.task_split = s
+ env_instance = get_env(
+ env_name=config.env,
+ user_strategy=config.user_strategy,
+ user_model=config.user_model,
+ task_split=config.task_split,
+ )
+ output_path = os.path.join(local_dir, f"{env}_{s}_tasks.jsonl")
+ with open(output_path, "w") as f:
+ for i, task in enumerate(env_instance.tasks):
+ row = {"index": i, "metadata": task.model_dump()}
+ f.write(json.dumps(row) + "\n")
+ print(f"Saved preprocessed task indices for {env} ({s}) to {output_path}")
+
+
+if __name__ == "__main__":
+ main()
diff --git a/examples/tau-bench/trainable_agents.py b/examples/tau-bench/trainable_agents.py
new file mode 100644
index 000000000..6e2043f9b
--- /dev/null
+++ b/examples/tau-bench/trainable_agents.py
@@ -0,0 +1,581 @@
+"""Trainable tau-bench agent for vime vLLM rollout."""
+
+from __future__ import annotations
+
+import base64
+import io
+import json
+import logging
+import os
+import uuid
+from typing import Any
+
+import numpy as np
+from openai_tool_adapter import create_openai_adapter
+from tau_bench.agents.tool_calling_agent import RESPOND_ACTION_NAME
+from tau_bench.envs import get_env
+from tau_bench.types import Action, RunConfig
+
+from vime.rollout.vllm_rollout import (
+ GenerateState,
+ _build_inference_sampling_params,
+ _coerce_flat_int_token_ids,
+ _mm_render_response_to_generate_body,
+)
+from vime.utils.http_utils import post
+from vime.utils.types import Sample
+
+logger = logging.getLogger(__name__)
+
+
+def patch_tau_user_retries() -> None:
+ """Reduce circuit-breaker fatality: more LiteLLM retries with backoff."""
+ try:
+ import tau_bench.envs.user as user_mod
+
+ user_mod.MAX_RETRIES = int(os.environ.get("TAU_USER_LITELLM_RETRIES", "30"))
+ user_mod.RETRY_DELAY_SECONDS = float(os.environ.get("TAU_USER_LITELLM_RETRY_DELAY", "2"))
+ except Exception:
+ pass
+
+
+def _parse_choice_tokens_and_logprobs(choice: dict[str, Any]) -> tuple[list[int], list[float]]:
+ """Parse token_ids + logprobs from vLLM /inference/v1/generate choice."""
+ tids_raw = choice.get("token_ids")
+ if not (isinstance(tids_raw, list) and tids_raw and all(isinstance(x, int) for x in tids_raw)):
+ return [], []
+ tids = [int(x) for x in tids_raw]
+ lp = choice.get("logprobs")
+ if not isinstance(lp, dict):
+ return tids, [0.0] * len(tids)
+ content = lp.get("content")
+ if isinstance(content, list) and content:
+ log_probs = [
+ float(content[i].get("logprob", 0.0)) if i < len(content) and isinstance(content[i], dict) else 0.0
+ for i in range(len(tids))
+ ]
+ return tids, log_probs
+ return tids, [0.0] * len(tids)
+
+
+def _maybe_apply_routed_experts(args: Any, sample: Sample, choice: dict[str, Any]) -> None:
+ if choice.get("routed_experts") is None:
+ return
+ raw = base64.b64decode(choice["routed_experts"].encode("ascii"), validate=True)
+ arr = np.load(io.BytesIO(raw), allow_pickle=False)
+ sample.rollout_routed_experts = np.ascontiguousarray(arr.astype(np.int32, copy=True)).reshape(
+ len(sample.tokens) - 1,
+ args.num_layers,
+ args.moe_router_topk,
+ )
+
+
+class TauBenchEnv:
+ def __init__(
+ self,
+ *,
+ tau_config: RunConfig,
+ task_index: int | None = None,
+ max_turns: int = 30,
+ ):
+ self.tau_config = tau_config
+ self.task_index = task_index
+ self.max_turns = max_turns
+ self.turn = 0
+ self.total_reward = 0.0
+ self.info: dict[str, Any] = {}
+ self.env = None
+ self.openai_adapter = None
+ self.successful_tool_calls = 0
+ self.total_tool_calls = 0
+ self.format_correct_calls = 0
+
+ def reset(self):
+ self.turn = 0
+ self.total_reward = 0.0
+ self.info = {}
+ self.successful_tool_calls = 0
+ self.total_tool_calls = 0
+ self.format_correct_calls = 0
+
+ self.env = get_env(
+ env_name=self.tau_config.env,
+ user_strategy=self.tau_config.user_strategy,
+ user_model=self.tau_config.user_model,
+ user_provider=self.tau_config.user_model_provider,
+ task_split=self.tau_config.task_split,
+ task_index=self.task_index,
+ )
+
+ self.openai_adapter = create_openai_adapter(
+ tools_info=self.env.tools_info,
+ parser_type="qwen25",
+ )
+
+ env_reset_res = self.env.reset(task_index=self.task_index) if self.task_index is not None else self.env.reset()
+ observation = env_reset_res.observation
+ self.info = self._to_dict(env_reset_res.info)
+
+ return {
+ "obs_str": observation,
+ "role": "user",
+ "wiki": self.env.wiki,
+ "tools_info": self.env.tools_info,
+ }
+
+ def step(self, response_text: str):
+ self.turn += 1
+ is_final_turn = self.turn >= self.max_turns
+
+ openai_result = self.openai_adapter.parse_response_to_openai_format(response_text)
+
+ if not openai_result["success"]:
+ logger.warning(f"Tool parsing failed: {openai_result.get('error')}")
+ return (
+ {
+ "obs_str": "Failed to parse tool call. Please try again.",
+ "role": "tool",
+ },
+ is_final_turn,
+ {"tool_executed": False, "parse_error": openai_result.get("error")},
+ )
+
+ parsed = openai_result["parsed_result"]
+ agent_content, calls = parsed["normal_text"], parsed["calls"]
+
+ if calls:
+ self.format_correct_calls += 1
+
+ action = self._call_to_action(calls, agent_content)
+
+ is_tool_call = action.name != RESPOND_ACTION_NAME
+ if is_tool_call:
+ self.total_tool_calls += 1
+
+ try:
+ env_response = self.env.step(action)
+ except Exception as e:
+ logger.warning(f"Environment step failed: {e}")
+ return (
+ {
+ "obs_str": f"Environment error: {e}",
+ "role": "tool",
+ },
+ True,
+ {"tool_executed": False, "env_error": str(e)},
+ )
+
+ self.total_reward = env_response.reward
+ self.info.update(self._to_dict(env_response.info))
+
+ obs_lower = env_response.observation.lower() if env_response.observation else ""
+ if is_tool_call and not obs_lower.startswith(("error", "failed", "invalid", "not found")):
+ self.successful_tool_calls += 1
+
+ if action.name != RESPOND_ACTION_NAME:
+ obs_role = "tool"
+ else:
+ obs_role = "user"
+ obs_content = env_response.observation
+
+ done = env_response.done or is_final_turn
+
+ return (
+ {
+ "obs_str": obs_content,
+ "role": obs_role,
+ "reward": env_response.reward,
+ },
+ done,
+ {"tool_executed": True, "action": action.name},
+ )
+
+ def _call_to_action(self, calls: list[Any], text_response: str) -> Action:
+ action = Action(name=RESPOND_ACTION_NAME, kwargs={"content": text_response})
+ if calls:
+ if len(calls) > 1:
+ logger.debug("Multiple tool calls identified, only taking first.")
+ tool_call = calls[0]
+ try:
+ params = (
+ json.loads(tool_call["parameters"])
+ if isinstance(tool_call["parameters"], str)
+ else tool_call["parameters"]
+ )
+ if not isinstance(params, dict):
+ logger.warning(f"{params} does not follow dict structure for action")
+ else:
+ action = Action(name=tool_call["name"], kwargs=params)
+ except json.JSONDecodeError as e:
+ logger.warning(f"Failed to parse parameters as JSON: {e}")
+ return action
+
+ def close(self):
+ pass
+
+ def format_observation(self, observation: dict) -> dict:
+ observation = observation or {}
+ content = observation.get("obs_str", "")
+ return {
+ "role": observation.get("role", "user"),
+ "content": content + "\n/no_think" if observation.get("role") == "user" else content,
+ }
+
+ @staticmethod
+ def _to_dict(info: Any) -> dict:
+ if hasattr(info, "model_dump"):
+ return info.model_dump()
+ if isinstance(info, dict):
+ return info
+ return {}
+
+
+def build_env(sample: Sample | None = None, args: Any | None = None, **_: Any) -> TauBenchEnv:
+ tau_bench_config = getattr(args, "tau_bench_config", None)
+ if tau_bench_config is None:
+ raise RuntimeError("args.tau_bench_config is missing; generate_with_tau.generate must set it from TAU_CONFIGS")
+
+ task_index = None
+ if sample is not None and sample.prompt is not None:
+ try:
+ task_index = int(sample.prompt)
+ except (ValueError, TypeError):
+ pass
+
+ max_turns = getattr(args, "max_turns", 30)
+ if max_turns is None:
+ max_turns = 30
+
+ return TauBenchEnv(
+ tau_config=tau_bench_config,
+ task_index=task_index,
+ max_turns=max_turns,
+ )
+
+
+def compute_process_reward(env: TauBenchEnv, base_reward: float) -> float:
+ reward = 0.0
+ if base_reward > 0:
+ reward += 1.0
+ if env.successful_tool_calls > 0:
+ reward += 0.1 * env.successful_tool_calls
+ if env.format_correct_calls > 0:
+ reward += 0.05 * env.format_correct_calls
+ reward = min(reward, 1.5)
+ return reward
+
+
+def _build_tools_section(tools_info: list[dict]) -> str:
+ if not tools_info:
+ return ""
+ tools_json = json.dumps(tools_info, ensure_ascii=False)
+ parts = [
+ "",
+ "",
+ "# Tools",
+ "",
+ "You may call one or more functions to assist with the user query.",
+ "",
+ "You are provided with function signatures within XML tags:",
+ "",
+ tools_json,
+ "",
+ "",
+ "For each function call, return a json object with function name and arguments within",
+ "",
+ '{"name": , "arguments": }',
+ "",
+ "XML tags.",
+ ]
+ return "\n".join(parts)
+
+
+def _messages_for_render(messages: list[dict]) -> list[dict]:
+ out: list[dict] = []
+ for msg in messages:
+ out.append({"role": msg["role"], "content": msg.get("content", "")})
+ return out
+
+
+class TrainableTauBenchAgent:
+ """Trainable tau-bench agent using vLLM render + /inference/v1/generate."""
+
+ async def asolve(self, args: Any, sample: Sample, sampling_params) -> Sample:
+ assert not args.partial_rollout, "Partial rollout is not supported for tau-bench interactions."
+
+ state = GenerateState(args)
+ base_url = f"http://{args.vllm_router_ip}:{args.vllm_router_port}"
+
+ sample.metadata = sample.metadata or {}
+
+ headers = None
+ if getattr(args, "router_policy", None) == "consistent_hash":
+ sample.session_id = sample.session_id or str(uuid.uuid4())
+ headers = {"x-session-id": sample.session_id}
+
+ try:
+ env = build_env(sample=sample, args=args)
+ except TypeError:
+ env = build_env(sample, args)
+
+ initial_obs = env.reset()
+ wiki = initial_obs.get("wiki", "")
+ short_wiki_chars = os.environ.get("TAU_SHORT_WIKI")
+ if short_wiki_chars is not None:
+ wiki = wiki[: int(short_wiki_chars)]
+ tools_info = initial_obs.get("tools_info", [])
+ if os.environ.get("TAU_SHORT_TOOLS") == "1":
+ keep = {"find_user_id_by_email", "find_user_id_by_name_zip", "get_user_details", "get_order_details"}
+ tools_info = [tool for tool in tools_info if tool.get("function", {}).get("name") in keep]
+ tools_section = _build_tools_section(tools_info)
+
+ messages: list[dict] = [
+ {"role": "system", "content": wiki + tools_section + "\n/no_think"},
+ {"role": "user", "content": initial_obs.get("obs_str", "")},
+ ]
+
+ response_tokens: list[int] = []
+ sample.loss_mask = sample.loss_mask or []
+ sample.rollout_log_probs = sample.rollout_log_probs or []
+ sample.tokens = list(sample.tokens) if sample.tokens else []
+
+ sampling_params = sampling_params.copy()
+ sampling_params["repetition_penalty"] = 1.1
+ sampling_params.setdefault("stop", [""])
+ inference_sampling_params = _build_inference_sampling_params(sampling_params)
+ max_response_budget = sampling_params.get("max_new_tokens")
+
+ def remaining_budget() -> int | None:
+ return None if max_response_budget is None else max_response_budget - sample.response_length
+
+ def _ensure_trainable_skeleton() -> None:
+ """Megatron padding needs total_length >= 1 (F.pad uses prompt_length - 1)."""
+ if not sample.tokens:
+ eos_id = getattr(state.tokenizer, "eos_token_id", None)
+ pad_id = getattr(state.tokenizer, "pad_token_id", None)
+ token_id = eos_id if eos_id is not None else pad_id if pad_id is not None else 0
+ sample.tokens = [int(token_id)]
+ sample.loss_mask = sample.loss_mask or []
+ sample.rollout_log_probs = sample.rollout_log_probs or []
+ if len(sample.rollout_log_probs) < len(sample.loss_mask):
+ sample.rollout_log_probs.extend([0.0] * (len(sample.loss_mask) - len(sample.rollout_log_probs)))
+ sample.response_length = len(sample.loss_mask)
+
+ def _mark_truncated(reward_base: float = 0.0, *, remove: bool = True) -> Sample:
+ _ensure_trainable_skeleton()
+ if remove:
+ sample.remove_sample = True
+ sample.reward = compute_process_reward(env, reward_base)
+ sample.status = Sample.Status.TRUNCATED
+ return sample
+
+ async def safe_render() -> dict | None:
+ try:
+ render_messages = _messages_for_render(messages)
+ payload = {"model": args.hf_checkpoint, "messages": render_messages}
+ render_data = await post(
+ f"{base_url}/v1/chat/completions/render", payload, headers=headers, max_retries=3
+ )
+ return _mm_render_response_to_generate_body(render_data, args.hf_checkpoint)
+ except Exception as exc:
+ logger.warning("render failed, skipping task: %s", exc)
+ return None
+
+ def append_response_window(
+ token_ids: list[int],
+ loss_mask: list[int],
+ log_probs: list[float] | None = None,
+ ) -> None:
+ if not token_ids:
+ return
+ if len(loss_mask) != len(token_ids):
+ raise ValueError(f"loss_mask length {len(loss_mask)} != token_ids length {len(token_ids)}")
+ sample.tokens.extend(token_ids)
+ sample.loss_mask.extend(loss_mask)
+ sample.rollout_log_probs.extend(log_probs if log_probs is not None else [0.0] * len(token_ids))
+ sample.response_length += len(token_ids)
+
+ def sampling_params_for_turn() -> dict | None:
+ params = dict(inference_sampling_params)
+ max_tokens = remaining_budget()
+ if max_tokens is None:
+ return params
+ if max_tokens <= 0:
+ return None
+ params["max_tokens"] = max_tokens
+ return params
+
+ try:
+ pending_obs_offset: int | None = None
+ rendered_body = await safe_render()
+ if rendered_body is None:
+ return _mark_truncated()
+ prompt_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids"))
+ if not sample.tokens:
+ sample.tokens = list(prompt_ids)
+ if args.rollout_max_context_len is not None:
+ context_budget = max(0, args.rollout_max_context_len - len(sample.tokens))
+ if max_response_budget is None:
+ max_response_budget = context_budget
+ else:
+ max_response_budget = min(max_response_budget, context_budget)
+
+ vllm_max_len = getattr(args, "vllm_max_model_len", 16384) or 16384
+ if len(prompt_ids) >= vllm_max_len - 64:
+ logger.info(f"prompt too long ({len(prompt_ids)} tokens >= {vllm_max_len - 64}), skipping task")
+ return _mark_truncated()
+
+ for turn_idx in range(args.max_turns):
+ input_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids"))
+
+ if pending_obs_offset is not None:
+ obs_tokens = input_ids[pending_obs_offset:]
+ remaining = remaining_budget()
+ if remaining is not None and len(obs_tokens) > remaining:
+ append_response_window(obs_tokens[: max(remaining, 0)], [0] * max(remaining, 0))
+ sample.status = Sample.Status.TRUNCATED
+ break
+ append_response_window(obs_tokens, [0] * len(obs_tokens))
+ pending_obs_offset = None
+
+ current_sampling_params = sampling_params_for_turn()
+ if current_sampling_params is None:
+ sample.status = Sample.Status.TRUNCATED
+ break
+
+ body = dict(rendered_body)
+ body["sampling_params"] = current_sampling_params
+ output = await post(f"{base_url}/inference/v1/generate", body, headers=headers)
+ choice = output["choices"][0]
+ finish_reason = choice.get("finish_reason") or "stop"
+ new_tokens, new_logprobs = _parse_choice_tokens_and_logprobs(choice)
+
+ if not new_tokens:
+ if finish_reason in ("abort", "cancelled"):
+ sample.status = Sample.Status.ABORTED
+ break
+
+ response_text = state.tokenizer.decode(new_tokens, skip_special_tokens=False) if new_tokens else ""
+ train_tokens = list(new_tokens)
+ train_logprobs = list(new_logprobs)
+ train_loss_mask = [1] * len(train_tokens)
+
+ stop = current_sampling_params.get("stop")
+ if not stop:
+ stop = [""]
+ stop_strings = (stop,) if isinstance(stop, str) else tuple(stop) if stop else ()
+ hit_stop_str = None
+ hit_stop_pos = len(response_text)
+ if stop_strings:
+ for ss in stop_strings:
+ pos = response_text.find(ss)
+ if pos != -1 and pos < hit_stop_pos:
+ hit_stop_str = ss
+ hit_stop_pos = pos
+ if hit_stop_str is not None:
+ truncated_text = response_text[: hit_stop_pos + len(hit_stop_str)]
+ trunc_token_count = 0
+ for t in range(1, len(train_tokens) + 1):
+ partial = state.tokenizer.decode(train_tokens[:t], skip_special_tokens=False)
+ if partial >= truncated_text:
+ trunc_token_count = t
+ break
+ if trunc_token_count == 0:
+ trunc_token_count = len(train_tokens)
+ train_tokens = train_tokens[:trunc_token_count]
+ train_logprobs = train_logprobs[:trunc_token_count]
+ train_loss_mask = train_loss_mask[:trunc_token_count]
+ response_text = truncated_text
+ finish_reason = "stop"
+
+ eos_token_id = getattr(state.tokenizer, "eos_token_id", None)
+ append_stop_eos = (
+ stop
+ and eos_token_id is not None
+ and getattr(args, "append_eos_token_after_stop_str_in_multi_turn", True)
+ )
+ if append_stop_eos:
+ already_has_eos = bool(train_tokens and train_tokens[-1] == eos_token_id)
+ if stop_strings and response_text.endswith(stop_strings) and not already_has_eos:
+ if getattr(args, "use_rollout_routing_replay", False):
+ raise RuntimeError(
+ "Routing replay is not supported when appending an artificial EOS after a stop string, "
+ "because vLLM does not return routed experts for that extra token."
+ )
+ train_tokens.append(int(eos_token_id))
+ train_logprobs.append(0.0)
+ train_loss_mask.append(0)
+
+ response_tokens.extend(new_tokens)
+ append_response_window(train_tokens, train_loss_mask, train_logprobs)
+ _maybe_apply_routed_experts(args, sample, choice)
+
+ messages.append({"role": "assistant", "content": response_text})
+
+ if finish_reason == "length":
+ sample.status = Sample.Status.TRUNCATED
+ break
+ if finish_reason in ("abort", "cancelled"):
+ sample.status = Sample.Status.ABORTED
+ break
+
+ observation, done, step_info = env.step(response_text)
+
+ if done:
+ base_reward = observation.get("reward", 0.0)
+ sample.reward = compute_process_reward(env, base_reward)
+ sample.status = Sample.Status.COMPLETED
+ break
+
+ next_user_message = env.format_observation(observation)
+ messages.append(next_user_message)
+
+ if turn_idx + 1 >= args.max_turns:
+ sample.reward = compute_process_reward(env, 0.0)
+ sample.status = Sample.Status.TRUNCATED
+ break
+
+ pending_obs_offset = len(input_ids) + len(train_tokens)
+ max_ctx = args.rollout_max_context_len or 8192
+ if len(sample.tokens) >= max_ctx - 64:
+ logger.info(
+ f"[turn={turn_idx}] context overflow: {len(sample.tokens)} tokens >= {max_ctx - 64}, truncating"
+ )
+ sample.reward = compute_process_reward(env, 0.0)
+ sample.status = Sample.Status.TRUNCATED
+ break
+ rendered_body = await safe_render()
+ if rendered_body is None:
+ return _mark_truncated()
+ rendered_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids"))
+ is_prefix_stable = rendered_ids[:pending_obs_offset] == sample.tokens[:pending_obs_offset]
+ sample.metadata["multiturn_render"] = {
+ "prefix_stable": is_prefix_stable,
+ "prefix_len": pending_obs_offset,
+ "sample_len": len(sample.tokens),
+ "rendered_len": len(rendered_ids),
+ "turn": turn_idx + 1,
+ }
+ if getattr(args, "strict_multiturn_render_token_match", False) and not is_prefix_stable:
+ raise RuntimeError(
+ "Full conversation render is not prefix-stable with the generated token stream: "
+ f"{sample.metadata['multiturn_render']}"
+ )
+
+ sample.response = state.tokenizer.decode(response_tokens, skip_special_tokens=False)
+ sample.response_length = len(sample.loss_mask)
+ _ensure_trainable_skeleton()
+ if sample.status == Sample.Status.PENDING:
+ sample.status = Sample.Status.COMPLETED
+ if sample.reward is None or sample.reward == 0.0:
+ sample.reward = compute_process_reward(env, getattr(env, "total_reward", 0.0))
+ return sample
+ finally:
+ try:
+ env.close()
+ except Exception:
+ pass
+
+
+def agent_factory() -> TrainableTauBenchAgent:
+ return TrainableTauBenchAgent()
diff --git a/examples/tau-bench/vllm_tool_parser.py b/examples/tau-bench/vllm_tool_parser.py
new file mode 100644
index 000000000..603c0f360
--- /dev/null
+++ b/examples/tau-bench/vllm_tool_parser.py
@@ -0,0 +1,95 @@
+"""Local tool-call parser for vLLM rollout."""
+
+import json
+import re
+from typing import Any
+
+
+def parse_tools(response: str, tools: list[dict[str, Any]], parser: str = "qwen25") -> dict[str, Any]:
+ if parser == "qwen25":
+ return _parse_qwen25_tools(response)
+ return _parse_qwen25_tools(response)
+
+
+def _try_parse_json_tool_call(text: str) -> dict[str, Any] | None:
+ try:
+ parsed = json.loads(text)
+ if isinstance(parsed, dict) and "name" in parsed:
+ name = parsed.get("name", "")
+ parameters = parsed.get("arguments", parsed.get("parameters", {}))
+ if isinstance(parameters, str):
+ try:
+ parameters = json.loads(parameters)
+ except json.JSONDecodeError:
+ pass
+ return {"name": name, "parameters": parameters}
+ except (json.JSONDecodeError, TypeError):
+ pass
+ return None
+
+
+def _parse_qwen25_tools(response: str) -> dict[str, Any]:
+ call_open = chr(60) + "tool_call" + chr(62)
+ call_close = chr(60) + "/tool_call" + chr(62)
+ call_open_alt = chr(60) + "call" + chr(62)
+ call_close_alt = chr(60) + "/call" + chr(62)
+ pattern = r"(?:" + call_open + "|" + call_open_alt + r")\s*(.*?)\s*(?:" + call_close + "|" + call_close_alt + ")"
+ tool_call_pattern = re.compile(pattern, re.DOTALL)
+ matches = tool_call_pattern.findall(response)
+
+ if matches:
+ parts = tool_call_pattern.split(response)
+ normal_text = parts[0].strip() if parts else ""
+ calls = []
+ for match in matches:
+ match = match.strip()
+ parsed_call = _try_parse_json_tool_call(match)
+ if parsed_call:
+ calls.append(parsed_call)
+ else:
+ try:
+ json_match = re.search(r"\{.*\}", match, re.DOTALL)
+ if json_match:
+ parsed_call = _try_parse_json_tool_call(json_match.group())
+ if parsed_call:
+ calls.append(parsed_call)
+ else:
+ calls.append({"name": match, "parameters": {}})
+ else:
+ calls.append({"name": match, "parameters": {}})
+ except (json.JSONDecodeError, AttributeError):
+ calls.append({"name": match, "parameters": {}})
+
+ return {
+ "normal_text": normal_text,
+ "calls": calls,
+ }
+
+ cleaned = re.sub(r"<\|im_end\|>", "", response).strip()
+ parsed_call = _try_parse_json_tool_call(cleaned)
+ if parsed_call:
+ return {
+ "normal_text": "",
+ "calls": [parsed_call],
+ }
+
+ json_pattern = re.compile(r'\{[^{}]*"name"\s*:\s*"[^"]+?"[^{}]*\}', re.DOTALL)
+ json_matches = json_pattern.findall(response)
+ if json_matches:
+ calls = []
+ for jm in json_matches:
+ parsed_call = _try_parse_json_tool_call(jm)
+ if parsed_call:
+ calls.append(parsed_call)
+ if calls:
+ normal_text = json_pattern.sub("", response).strip()
+ normal_text = re.sub(r"<\|im_end\|>", "", normal_text).strip()
+ return {
+ "normal_text": normal_text,
+ "calls": calls,
+ }
+
+ return {
+ "normal_text": response,
+ "calls": [],
+ }