feat: Add Tau bench environment

examples/configs/grpo_tau_bench.yaml

@@ -0,0 +1,113 @@
+# GRPO training on tau-bench: multi-turn customer service tasks with LLM-as-judge rewards.
+#
+# The agent must complete realistic retail or airline service tasks by calling domain
+# tools across multiple turns.  Rewards blend tau-bench's deterministic task-completion
+# check with an optional LLM judge score for policy compliance.
+#
+# Data preparation: run scripts/prepare_tau_bench_data.py to generate the JSONL files
+# before starting training.  The script calls tau-bench's env.reset() for each task
+# so that the prompts (tool definitions, domain wiki, initial user message) are
+# exactly consistent with the environment state.
+#
+# Tool call format expected from the policy:
+#   <tool_call>{"name": "tool_name", "arguments": {"key": "value"}}</tool_call>
+# Text without a <tool_call> block is treated as a respond action to the user.
+
+defaults: "grpo_math_1B_megatron.yaml"
+
+grpo:
+  num_prompts_per_step: 16
+  num_generations_per_prompt: 8
+  # tau-bench episodes can be up to ~50 tool calls; set headroom above max_steps.
+  max_rollout_turns: 56
+  max_num_steps: 10000
+  max_num_epochs: 1
+
+checkpointing:
+  checkpoint_dir: "results/grpo-tau-bench"
+  metric_name: "val:tau_bench/task_completion_rate"
+  higher_is_better: true
+
+policy:
+  model_name: "Qwen/Qwen2.5-7B-Instruct"
+  max_total_sequence_length: 8192
+  train_micro_batch_size: 1
+  logprob_batch_size: 1
+  train_global_batch_size: 128
+
+  megatron_cfg:
+    enabled: true
+    tensor_model_parallel_size: 2
+    pipeline_model_parallel_size: 1
+    context_parallel_size: 1
+    pipeline_dtype: ${policy.precision}
+
+  generation:
+    backend: "vllm"
+    max_new_tokens: 512   # per turn; tool calls are short
+    temperature: 1.0
+    top_p: 1.0
+    stop_strings: null    # stop strings are supplied per-turn by the environment
+    stop_token_ids: null
+    vllm_cfg:
+      async_engine: false
+      tensor_parallel_size: 2
+      gpu_memory_utilization: 0.6
+      max_model_len: ${policy.max_total_sequence_length}
+
+cluster:
+  gpus_per_node: 8
+  num_nodes: 1
+
+data:
+  train:
+    dataset_name: "tau_bench"
+    tau_bench_env_name: "retail"   # "retail" or "airline"
+    split: "train"
+  validation:
+    dataset_name: "tau_bench"
+    tau_bench_env_name: "retail"
+    split: "dev"
+  default:
+    prompt_file: null
+    system_prompt_file: null
+    processor: "tau_bench_data_processor"
+    env_name: "tau_bench"
+
+env:
+  tau_bench:
+    num_workers: 2
+    env_name: "retail"       # "retail" or "airline"
+    task_split: "train"
+    user_strategy: "mock"
+    # User simulator model: any OpenAI-compatible endpoint.
+    user_model: null # unused when user_strategy="mock". Required if user_strategy="llm"
+    user_base_url: null # unused when user_strategy="mock". Required if user_strategy="llm"
+    user_api_key: null # unused when user_strategy="mock"
+    max_steps: 50
+    # LLM-as-judge: scores policy compliance and communication quality.
+    # Set judge_model to null to disable and use pure tau-bench rewards.
+    judge_model: "mock"
+    judge_base_url: null # unused when judge_model="mock"
+    judge_api_key: null # unused when judge_model="mock"
+    judge_weight: 0.3   # 70% task-completion + 30% judge score
+    mock_latency_s: 0.1 # how much time to wait before returning mock user/judge generations
+    mock_stop_prob: 0.4
+
+logger:
+  log_dir: "logs"
+  wandb_enabled: false
+  tensorboard_enabled: false
+  mlflow_enabled: false
+  swanlab_enabled: false
+  monitor_gpus: true
+  wandb:
+    project: "ashors-grpo-dev"
+    name: "grpo-tau-bench"
+  swanlab:
+    project: "grpo-dev"
+    name: "grpo-tau-bench"
+  mlflow:
+    experiment_name: "grpo-dev"
+    run_name: "grpo-tau-bench"
+    tracking_uri: "http://localhost:5000"

nemo_rl/data/datasets/response_datasets/__init__.py

@@ -42,6 +42,7 @@
 from nemo_rl.data.datasets.response_datasets.refcoco import RefCOCODataset
 from nemo_rl.data.datasets.response_datasets.response_dataset import ResponseDataset
 from nemo_rl.data.datasets.response_datasets.squad import SquadDataset
+from nemo_rl.data.datasets.response_datasets.tau_bench_dataset import TauBenchDataset
 from nemo_rl.data.datasets.response_datasets.tulu3 import Tulu3SftMixtureDataset
 
 DATASET_REGISTRY = {
@@ -63,6 +64,7 @@
     "tulu3_sft_mixture": Tulu3SftMixtureDataset,
     "gsm8k": GSM8KDataset,
     "Nemotron-Cascade-2-SFT-Math": NemotronCascade2SFTMathDataset,
+    "tau_bench": TauBenchDataset,
     # load from local JSONL file or HuggingFace
     "openai_format": OpenAIFormatDataset,
     "NemoGymDataset": NemoGymDataset,
@@ -116,5 +118,6 @@ def load_response_dataset(data_config: ResponseDatasetConfig):
     "ResponseDataset",
     "SquadDataset",
     "Tulu3SftMixtureDataset",
+    "TauBenchDataset",
     "load_response_dataset",
 ]

nemo_rl/data/datasets/response_datasets/tau_bench_dataset.py

@@ -0,0 +1,168 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""TauBench dataset for nemo-rl GRPO training.
+
+Each sample contains:
+  - messages: [system_message, initial_user_message]
+  - extra_env_info: {task_index, episode_id: null, step_count: 0}
+  - task_name: "tau_bench"
+
+The system message is built from the environment's tool definitions, domain
+wiki, and policy rules, exactly matching what TauBenchEnvironment initialises
+at training time.
+"""
+
+from typing import Any
+
+import ray
+from datasets import Dataset
+
+from nemo_rl.data.datasets.raw_dataset import RawDataset
+from nemo_rl.distributed.virtual_cluster import PY_EXECUTABLES
+
+
+_TOOL_CALL_INSTRUCTIONS = """\
+To call a tool, wrap a JSON object in <tool_call> tags:
+
+    <tool_call>{"name": "tool_name", "arguments": {"param": "value"}}</tool_call>
+
+Call one tool at a time and wait for the result before proceeding. \
+To send a final message to the customer (ending the interaction), \
+respond with plain text only — no <tool_call> tags."""
+
+
+def _format_tool(tool_info: dict[str, Any]) -> str:
+    lines = [f"### {tool_info['name']}"]
+    if "description" in tool_info:
+        lines.append(f"Description: {tool_info['description']}")
+    params = tool_info.get("parameters", {})
+    if isinstance(params, dict):
+        props = params.get("properties", {})
+        required = set(params.get("required", []))
+        if props:
+            lines.append("Parameters:")
+            for name, schema in props.items():
+                req = " (required)" if name in required else " (optional)"
+                ptype = schema.get("type", "any")
+                desc = schema.get("description", "")
+                lines.append(f"  - {name} ({ptype}{req}): {desc}")
+    return "\n".join(lines)
+
+
+def _build_system_prompt(env: Any, env_name: str) -> str:
+    tool_infos = []
+    for t in env.tools_info or []:
+        # tau_bench wraps each tool spec as {"type": "function", "function": {...}}
+        info = t["function"] if "function" in t else t
+        tool_infos.append(info)
+
+    tools_section = "\n\n".join(_format_tool(t) for t in tool_infos)
+    rules_section = "\n".join(f"- {r}" for r in (env.rules or []))
+    wiki = (env.wiki or "").strip()
+
+    return "\n\n".join(
+        filter(
+            None,
+            [
+                f"You are a customer service agent for a {env_name} company. "
+                "Complete the customer's request using the tools below.",
+                f"## Tool Usage\n\n{_TOOL_CALL_INSTRUCTIONS}",
+                f"## Available Tools\n\n{tools_section}" if tools_section else None,
+                f"## Domain Knowledge\n\n{wiki}" if wiki else None,
+                f"## Policies\n\n{rules_section}" if rules_section else None,
+            ],
+        )
+    )
+
+
+def _build_records(tau_bench_env_name: str, split: str) -> list[dict[str, Any]]:
+    """Build the list of dataset records from a tau-bench environment.
+
+    This function is invoked via Ray remote with PY_EXECUTABLES.TAU_BENCH so
+    that tau_bench is only required in the TAU_BENCH Ray environment, not in
+    the driver process.
+    """
+    from tau_bench.envs import get_env
+
+    # Use the "human" strategy so no LLM provider is required at data-load
+    # time. We read task.instruction directly — env.reset() is intentionally
+    # NOT called here because HumanUserSimulationEnv.reset() blocks on
+    # input(), which would hang the process for every task.
+    env = get_env(
+        env_name=tau_bench_env_name,
+        user_strategy="human",
+        user_model="",
+        task_split=split,
+    )
+
+    system_prompt = _build_system_prompt(env, tau_bench_env_name)
+
+    # The actual customer opening message comes from env.reset() (via the LLM
+    # user simulator) at training time.  TauBenchWorker.execute() returns it as
+    # the first observation (a "pre-step") before any agent action is processed.
+    # The placeholder below just gives the model a neutral starting point for
+    # the wasted generation on that pre-step turn; the real conversation begins
+    # once the agent sees the customer's actual first message.
+    placeholder_user_msg = "Hi! How can I help you today?"
+    return [
+        {
+            "messages": [
+                {"role": "system", "content": system_prompt},
+                {"role": "user", "content": placeholder_user_msg},
+            ],
+            "extra_env_info": {
+                "task_index": task_index,
+                "episode_id": None,
+                "step_count": 0,
+            },
+            "task_name": "tau_bench",
+        }
+        for task_index, task in enumerate(env.tasks)
+    ]
+
+
+class TauBenchDataset(RawDataset):
+    """Dataset built from a tau-bench environment.
+
+    Calls tau_bench.envs.get_env() at construction time and iterates every
+    task in the requested split to produce (system_prompt, initial_user_message)
+    pairs.  No offline data-preparation step is needed.
+
+    Args:
+        tau_bench_env_name: tau-bench domain — "retail" or "airline".
+        split: Dataset split — "train", "dev", or "test" (default: "train").
+        repeat: Number of times to repeat the dataset (default: 1).
+    """
+
+    def __init__(
+        self,
+        tau_bench_env_name: str,
+        split: str = "train",
+        repeat: int = 1,
+        **kwargs: Any,
+    ) -> None:
+        self.task_name = "tau_bench"
+
+        # Run record construction in the TAU_BENCH Ray environment so that
+        # tau_bench is not required in the driver process.
+        records = ray.get(
+            ray.remote(_build_records)
+            .options(runtime_env={"py_executable": PY_EXECUTABLES.TAU_BENCH})
+            .remote(tau_bench_env_name, split)
+        )
+
+        self.dataset = Dataset.from_list(records)
+
+        if repeat > 1:
+            self.dataset = self.dataset.repeat(repeat)

nemo_rl/data/processors.py

@@ -707,6 +707,62 @@ def nemo_gym_data_processor(
     return output
 
 
+def tau_bench_data_processor(
+    datum_dict: dict[str, Any],
+    task_data_spec: TaskDataSpec,
+    tokenizer: TokenizerType,
+    max_seq_length: int,
+    idx: int,
+) -> DatumSpec:
+    """Process a tau-bench datum into a DatumSpec for multi-turn GRPO training.
+
+    Expects datum_dict produced by ``examples/prepare_tau_bench_data.py``:
+      - messages: [system_message, initial_user_message]
+      - extra_env_info: {task_index, episode_id, step_count}
+      - task_name: "tau_bench"
+
+    The full system + user prompt is collapsed into a single message_log entry
+    (role "user") whose content is the chat-template-formatted string.  This is
+    the same layout used by math_hf_data_processor and is required because the
+    multi-turn rollout loop appends subsequent turns on top of this entry.
+    """
+    messages = datum_dict["messages"]
+    extra_env_info = datum_dict["extra_env_info"]
+
+    formatted: str = tokenizer.apply_chat_template(  # type: ignore[assignment]
+        messages,
+        tokenize=False,
+        add_generation_prompt=True,
+        add_special_tokens=False,
+    )
+    token_ids = tokenizer(
+        formatted,
+        return_tensors="pt",
+        add_special_tokens=False,
+    )["input_ids"][0]
+
+    message_log: LLMMessageLogType = [
+        {"role": "user", "content": formatted, "token_ids": token_ids}
+    ]
+
+    length = len(token_ids)
+    loss_multiplier = 1.0
+    if length >= max_seq_length:
+        for msg in message_log:
+            msg["token_ids"] = msg["token_ids"][: min(4, max_seq_length // len(message_log))]
+        loss_multiplier = 0.0
+
+    output: DatumSpec = {
+        "message_log": message_log,
+        "length": length,
+        "extra_env_info": extra_env_info,
+        "loss_multiplier": loss_multiplier,
+        "idx": idx,
+        "task_name": datum_dict.get("task_name", "tau_bench"),
+    }
+    return output
+
+
 # Processor registry. Key is the processor name, value is the processor function.
 # Note: We cast the literal dict to Dict[str, TaskDataProcessFnCallable] because
 # type checkers see each concrete function's signature as a distinct callable type.
@@ -724,6 +780,7 @@ def nemo_gym_data_processor(
         "sft_processor": sft_processor,
         "vlm_hf_data_processor": vlm_hf_data_processor,
         "nemo_gym_data_processor": nemo_gym_data_processor,
+        "tau_bench_data_processor": tau_bench_data_processor,
     },
 )
 

nemo_rl/distributed/ray_actor_environment_registry.py

@@ -47,6 +47,7 @@
     "nemo_rl.algorithms.async_utils.ReplayBuffer": PY_EXECUTABLES.VLLM,
     "nemo_rl.environments.tools.retriever.RAGEnvironment": PY_EXECUTABLES.SYSTEM,
     "nemo_rl.environments.nemo_gym.NemoGym": PY_EXECUTABLES.NEMO_GYM,
+    "nemo_rl.environments.tau_bench_environment.TauBenchEnvironment": PY_EXECUTABLES.TAU_BENCH,
 }
 
 

nemo_rl/distributed/virtual_cluster.py

@@ -61,6 +61,9 @@ class PY_EXECUTABLES:
     # Use NeMo-Gym dependencies
     NEMO_GYM = f"uv run --locked --extra nemo_gym --directory {git_root}"
 
+    # Use tau-bench dependencies
+    TAU_BENCH = f"uv run --locked --extra tau_bench --directory {git_root}"
+
     # Use NeMo-RL direct dependencies and SGLang.
     SGLANG = f"uv run --locked --extra sglang --directory {git_root}"