feat: Add Tau bench environment

[ashors1]

What does this PR do ?

This PR adds a Tau bench environment to serve as an example of a multi-turn environment with an option for LLM-as-a-judge. This can be used to analyze the performance of more complicated multi-turn environments natively in NeMo-RL.

Note that we provide the option to mock user and judge models using user_strategy='mock' and judge_model='mock'.

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[terrykong]

Review Summary

Nice work adding the tau-bench multi-turn environment — the architecture is clean and the mock mode is a thoughtful addition for testing without LLM credentials. Two critical issues need attention before merge:

1. _tokenize_env_observation silently breaks all existing multi-turn environments — the new function passes role="environment" to apply_chat_template, which silently drops the message on Qwen/Llama templates. Math, code, sliding_puzzle, vlm, and reward_model environments would lose all observation feedback in rollouts.

2. TestCallJudge is entirely broken — fixture missing _mock_mode attribute, and all tests mock requests.post instead of litellm.completion.

Additional notes (not tied to specific diff lines)

• Documentation: This is a new feature (feat:) but there is no user-facing documentation. docs/guides/environments.md does not mention tau-bench, and there is no guide analogous to docs/guides/grpo-sliding-puzzle.md. Consider adding a section.

• Design note: The pre-step mechanism wastes one generation turn per episode (~2% overhead with max_steps=50). The tradeoff is documented in the code and reasonable, but worth noting.

• Global litellm monkey-patch: Mock mode at tau_bench_environment.py:149 patches litellm.completion globally at the module level. Any other code in the same worker process that uses litellm.completion would also get the mock. Consider a more targeted patch.

• Env var clobbering: os.environ["OPENAI_API_KEY"] = user_api_key at tau_bench_environment.py:188 unconditionally overwrites any existing key (defaults to "dummy" from config fallback). Guard with if user_api_key:.

Generated by Claude Code

[terrykong]

test_tau_bench_environment.py:74-92 — Broken test fixture

_make_worker() never sets _mock_mode or _mock_latency_s. When _call_judge runs if self._mock_mode: (line 266), it raises AttributeError.

Additionally, all TestCallJudge tests mock requests.post but the implementation uses litellm.completion. The mocks are never hit — tests either fail or pass for the wrong reason (the broad except Exception returns 0.0).

Suggested fix: add worker._mock_mode = False and worker._mock_latency_s = 0.0 to _make_worker(), then replace mock.patch("requests.post", ...) with mock.patch("litellm.completion", ...) and build mocks matching litellm's return type (resp.choices[0].message.content).

[ashors1]

/ok to test 97b9b98

[ashors1]

/ok to test 82b27b5

[ashors1]

/ok to test 202821a

[ashors1]

/ok to test 8b434ac

[ashors1]

/ok to test 968b1ff

[ashors1]

/ok to test 968b1ff

[terrykong]

Re-review — tau-bench environment

Thanks for the thorough follow-ups — the entire prior review is addressed: obs_use_chat_template() cleanly gates the chat-template path (defaults False, so existing environments are untouched), the TestCallJudge fixture/mocks are fixed (31 tests pass locally), and the API-key / mock-gating / dead-code / dependency-pin / docs items are all resolved. The get_idx_grouping refactor is also correct — it fixes the multi-turn shared-prompt grouping case.

One blocker remains, plus a few low-severity suggestions inline:

• Blocker: tests/test_suites/disabled.txt points at a stale filename → 3 unit tests in test_recipes_and_test_suites.py fail (reproduced locally). Likely the cause of the blocked merge.
• Low-severity: judge error handling, ambient OPENAI_API_KEY clobber, a docstring fix, an uncovered token-slicing routine (with a ready-to-drop test), a mock-mode thread-safety note, and a metric label.

All 121 unit tests across the 4 changed test files pass, and pre-commit is clean on the PR's files.

Generated by Claude Code

[terrykong]

disabled.txt:12

CI-breaking. This disables grpo_tau_bench_local.sh, but the committed script and recipe are named grpo-qwen2.5-7b-instruct-16n8g-tau-bench-local — the real name is in no suite list. tests/unit/test_recipes_and_test_suites.py asserts exact set-equality between on-disk .sh files, recipe YAMLs, and suite entries, so test_all_test_scripts_accounted_for_in_test_suites, test_all_recipe_yamls_accounted_for_in_test_suites, and test_all_tests_can_find_config_if_dryrun all fail (reproduced locally). This is likely why the merge is blocked.

Action: point this entry at the actual filename:

Suggested change

	tests/test_suites/llm/grpo_tau_bench_local.sh
	tests/test_suites/llm/grpo-qwen2.5-7b-instruct-16n8g-tau-bench-local.sh

[terrykong]

tau_bench_environment.py:336

_call_judge collapses every error to 0.0 with no retry, which lumps two very different failures together:

• Transient / infra (timeout, connection refused, rate limit, 5xx, judge server not up yet) — not the agent's fault and not a real score of 0. With judge_weight: 0.3, a brief judge outage silently scores affected episodes 0.7*tau_reward and writes a fake 0.0 into judge_score → pollutes tau_bench/mean_judge_score, with nothing in the metrics to reveal it.
• Parse / content (json.JSONDecodeError, missing score) — the judge replied but the output is unusable. Since the judge runs at temperature=0.0, a retry reproduces the same output, so retry won't help here; but these should still be counted so an always-mis-formatting judge (all-zeros signal) is visible.

There's already precedent for this disambiguation in this same file: the tau_env.step retry loop branches on the litellm exception class, treating litellm.BadRequestError as non-retryable and retrying everything else 5×. The judge path just doesn't mirror it.

Action: catch by exception class (litellm re-exports the OpenAI-style typed exceptions, so no need to parse HTTP status codes):

for attempt in range(5):
    try:
        resp = litellm.completion(model=self._judge_model, ..., timeout=60)
        return float(json.loads(resp.choices[0].message.content)["score"])
    except (litellm.Timeout, litellm.APIConnectionError, litellm.RateLimitError,
            litellm.InternalServerError, litellm.ServiceUnavailableError) as e:
        last_exc = e                       # transient -> back off and retry
        time.sleep(2**attempt + random.uniform(0, 1))
    except (json.JSONDecodeError, KeyError, ValueError):
        self._judge_parse_failures += 1    # deterministic at temp=0; count, don't retry
        return None
return None  # retries exhausted -> a real outage, NOT a real 0

And the part that matters most for the training signal: on None, don't blend judge_weight*0 — fall back to reward = tau_reward for that sample (score it as if the judge were disabled this episode) and surface a judge_failures count in the metrics, so "judge was down" is never indistinguishable from "agent scored 0." (If a hard raise is too aggressive for an example env, the minimum bar is transient-retry + a logged failure counter.)

[terrykong]

tau_bench_environment.py:597

What to change: this line. The worker added an if user_api_key: guard (L216) to avoid clobbering an ambient OPENAI_API_KEY, but this call site passes cfg.get("user_api_key") or "dummy" — always truthy — so the guard always fires, the elif "OPENAI_API_KEY" not in os.environ fallback (L218) is dead code, and a user who omits user_api_key while relying on a real ambient OPENAI_API_KEY (real-OpenAI user simulator) has it overwritten with "dummy". (It's also a call-site default, which config-conventions forbids — the "dummy" default belongs in the worker's env fallback, not here.)

Action: drop the or "dummy" so the worker receives None when the key is unset and its existing elif "OPENAI_API_KEY" not in os.environ branch supplies "dummy" only when there's no ambient key:

Suggested change

	user_api_key=cfg.get("user_api_key") or "dummy",
	user_api_key=cfg.get("user_api_key"),

(You may also want to loosen the worker's user_api_key: str annotation to Optional[str] at L129.)

Note: judge_api_key on L601 does not have the clobber problem — it's passed straight to litellm.completion(api_key=...) (L332) and never written to os.environ. Its or "dummy" is only the milder call-site-default smell, so cleaning it up is optional.

[terrykong]

tau_bench_environment.py:65-69

These docstrings say judge_model/judge_base_url/judge_api_key are "Ignored when user_strategy is 'mock'", but the judge is mocked independently — gated on judge_model == "mock" (L155), not on user_strategy. So user_strategy="mock" + a real judge_model still issues real judge calls.

Action: update lines 66/68/69 to "Ignored when judge_model is 'mock'."

[terrykong]

rollouts.py:233

_tokenize_env_observation(use_chat_template=True) — the dummy-primed token-slicing routine central to tau-bench correctness — has no test coverage (grep-confirmed). Its own docstring lists real failure modes (loop.last separator drift, context-dependent BPE), so it's worth locking down.

Action: add these tests to tests/unit/experience/test_rollouts.py (verified passing locally against Qwen2.5-1.5B-Instruct via the existing rollout_tokenizer fixture):

from nemo_rl.experience.rollouts import _tokenize_env_observation


def test_tokenize_env_observation_raw(rollout_tokenizer):
    out = _tokenize_env_observation(
        rollout_tokenizer, "tool", "hello world", use_chat_template=False
    )
    expected = rollout_tokenizer(
        "hello world", return_tensors="pt", add_special_tokens=False
    ).input_ids[0]
    assert torch.equal(out, expected)


def test_tokenize_env_observation_tool_chat_template(rollout_tokenizer):
    out = _tokenize_env_observation(
        rollout_tokenizer, "tool", "ORDER_STATUS: shipped", use_chat_template=True
    )
    decoded = rollout_tokenizer.decode(out)
    assert "ORDER_STATUS: shipped" in decoded
    assert decoded.startswith("<|im_end|>")
    assert decoded.rstrip().endswith("<|im_start|>assistant")
    assert "<tool_response>" in decoded  # Qwen wraps tool results in a user turn


def test_tokenize_env_observation_user_chat_template(rollout_tokenizer):
    out = _tokenize_env_observation(
        rollout_tokenizer, "user", "Can you help me?", use_chat_template=True
    )
    decoded = rollout_tokenizer.decode(out)
    assert decoded.startswith("<|im_end|>")
    assert "<|im_start|>user" in decoded
    assert "Can you help me?" in decoded
    assert decoded.rstrip().endswith("<|im_start|>assistant")


def test_tokenize_env_observation_excludes_dummy_priming(rollout_tokenizer):
    # The dummy "x" priming message must not leak into the returned tokens.
    out = _tokenize_env_observation(
        rollout_tokenizer, "user", "unique_marker_zzz", use_chat_template=True
    )
    decoded = rollout_tokenizer.decode(out)
    assert "unique_marker_zzz" in decoded
    assert "<|im_start|>user\nx<|im_end|>" not in decoded

[terrykong]

tau_bench_environment.py:406

Non-blocking (only reachable with user_strategy="mock" + max_concurrent_api_requests>1, which no shipped recipe uses). _execute_one runs concurrently in a ThreadPoolExecutor (default 8), and each thread does with patch("litellm.completion", ...). unittest.mock.patch mutates a process-global and isn't thread-safe, so overlapping threads leak the patch — a mock-mode step can call the real litellm.completion, and the global can be left permanently patched (reproduced: 8 leak events at 8 workers). The inline comment's "other litellm users are not affected" doesn't hold under concurrency.

Action: avoid patching the global under concurrency — inject the mock into tau-bench's user simulator, guard the patched region with a lock, or force max_concurrent=1 in mock mode.

[terrykong]

tau_bench_environment.py:693

tau_bench/task_completion_rate is computed from the blended rewards*is_end, so when judge_weight>0 it reports the tau+judge blend, not the true completion rate (mean_tau_reward already reports the pure value).

Action: rename this to mean_reward, or compute it from the tau_reward metadata so the label matches its value.

[ashors1]

/ok to test 5499b99

[ashors1]

/ok to test a370039

[ashors1]

/ok to test ff60101

[ashors1]

/ok to test 09413fe