diff --git a/docs/source/tutorials/browsergym-harness.md b/docs/source/tutorials/browsergym-harness.md new file mode 100644 index 0000000000..3bc624266f --- /dev/null +++ b/docs/source/tutorials/browsergym-harness.md @@ -0,0 +1,129 @@ +# BrowserGym Harness Rollouts + +This tutorial shows how to drive BrowserGym through the OpenEnv harness runtime +when a trainer needs to keep token sampling, logprobs, and reward assignment +inside the training loop. + +```{note} +Use this pattern for tool-driven BrowserGym rollouts. For a standard +`reset()` / `step()` GRPO flow, keep using the Wordle GRPO tutorial. +``` + +## What You'll Build + +- A BrowserGym session factory that creates one environment client per rollout. +- A harness rollout function that TRL can call during training. +- A model-step wrapper that converts generated BrowserGym action text into + structured tool calls. + +## Install Dependencies + +Install OpenEnv, TRL, and the BrowserGym environment package: + +```bash +pip install -U "trl[vllm]" peft trackio kernels +pip install -U git+https://github.com/meta-pytorch/OpenEnv.git +pip install -U "openenv-browsergym @ git+https://huggingface.co/spaces/openenv/browsergym_env" +``` + +## Build The Session Factory + +`BrowserGymSessionFactory` adapts a BrowserGym client into the harness +`ResourceSession` interface. If your training setup already has an +`environment_factory`, pass that factory as `client_factory` so every rollout +gets a fresh environment instance. + +```python +from browsergym_env import BrowserGymEnv +from browsergym_env.harness import BrowserGymSessionFactory + +space_url = "https://openenv-browsergym-env.hf.space" + + +def environment_factory(): + return BrowserGymEnv(base_url=space_url) + + +session_factory = BrowserGymSessionFactory( + client_factory=environment_factory, + default_task="click-test", +) +``` + +The session exposes BrowserGym actions such as `click`, `fill`, `send_keys`, +`scroll`, and `noop` as MCP-style tools while still executing the corresponding +BrowserGym action strings under the hood. + +## Wrap TRL Generation + +The harness calls a `model_step` function for each turn. The model step should +use the trainer-owned generation path, then return a `ModelStepResult` with the +completion text, token ids, logprobs, and exactly one BrowserGym tool call. + +```python +from browsergym_env.harness import build_browsergym_action_tool_call +from openenv.core.harness import ModelStepResult +from openenv.core.llm_client import LLMResponse +from trl.experimental.openenv import generate_rollout_completions + + +def build_trl_browsergym_model_step(trainer, tokenizer): + def model_step(messages, tools, sampling): + del tools, sampling + prompt_text = tokenizer.apply_chat_template( + messages, + add_generation_prompt=True, + tokenize=False, + ) + rollout_output = generate_rollout_completions(trainer, [prompt_text])[0] + completion_text = rollout_output.get("text") or tokenizer.decode( + rollout_output["completion_ids"], + skip_special_tokens=True, + ) + tool_call = build_browsergym_action_tool_call(completion_text) + return ModelStepResult( + response=LLMResponse(content=completion_text, tool_calls=[tool_call]), + prompt_ids=list(rollout_output["prompt_ids"]), + completion_ids=list(rollout_output["completion_ids"]), + logprobs=list(rollout_output["logprobs"]), + ) + + return model_step +``` + +In practice, you should add a small parser around the completion text so common +outputs like `Action: click('13')` are normalized before calling +`build_browsergym_action_tool_call`. + +## Create The Rollout Function + +Pass the session factory, white-box harness adapter, and model-step builder to +`build_harness_rollout_func`: + +```python +from openenv.core.harness import ( + HarnessRunLimits, + MCPHarnessAdapter, + build_harness_rollout_func, +) + +rollout_func = build_harness_rollout_func( + session_factory=session_factory, + harness_adapter=MCPHarnessAdapter(), + model_step_builder=lambda trainer, session: build_trl_browsergym_model_step( + trainer, + tokenizer, + ), + limits=HarnessRunLimits(max_turns=10), +) +``` + +The returned function accepts TRL prompts and a trainer, runs one harness-backed +BrowserGym episode per prompt, and returns `prompt_ids`, `completion_ids`, +`logprobs`, `env_reward`, and `verify_metrics`. + +## Full Example + +See [`examples/browsergym_harness.py`](https://github.com/meta-pytorch/OpenEnv/blob/main/examples/browsergym_harness.py) +for a complete TRL-oriented helper that includes action normalization and a +ready-to-use `build_browsergym_rollout_func`. diff --git a/docs/source/tutorials/index.md b/docs/source/tutorials/index.md index 09a66cdae6..33c30de829 100644 --- a/docs/source/tutorials/index.md +++ b/docs/source/tutorials/index.md @@ -36,6 +36,7 @@ end-to-end-walkthrough mcp-environment rubrics wordle-grpo +browsergym-harness rl-training-2048 evaluation-inspect ``` diff --git a/docs/source/tutorials/sft-warmup.md b/docs/source/tutorials/sft-warmup.md new file mode 100644 index 0000000000..fd0f827a28 --- /dev/null +++ b/docs/source/tutorials/sft-warmup.md @@ -0,0 +1,427 @@ +# Collecting rollouts with OpenEnv for supervised training + +[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/meta-pytorch/OpenEnv/blob/main/examples/sft_warmup.ipynb) + +OpenEnv environments are not only useful for RL training — they are also a natural tool for **collecting +rollouts that become supervised training data**. The environment handles episode management, automatic scoring, +and reproducibility, so you get a reward-labeled dataset without writing any of that infrastructure yourself. + +This tutorial shows the full pipeline: + +1. Run a strong teacher model inside an OpenEnv environment to collect rollouts. +2. Use the environment's reward signal to filter out incorrect examples automatically. +3. Train a smaller student model on the filtered rollouts with TRL's `SFTTrainer`. + +As a concrete application, the resulting checkpoint is used as a warm-start for GRPO: once the student +reliably produces valid tool calls, GRPO's `reward_std` is non-zero from the first batch and the reward +curve climbs immediately. + +## Why use an environment to collect training data + +Building a supervised dataset usually means writing a custom collection loop, a scorer, and episode +bookkeeping. An OpenEnv environment gives you all three out of the box: + +- **Automatic scoring** — every `step()` returns a reward. Filter by `reward == 1.0` and you have a + clean, correct dataset with no manual labelling. +- **Reproducible episodes** — `reset(seed=42, size=N)` produces the same sequence of problems every + run. Anyone can regenerate the exact dataset. +- **Configurable difficulty** — adjust `DATASET_CONFIG` to control problem complexity without changing + any collection code. +- **Portable across environments** — the same collect → filter → train pipeline works for any OpenEnv + environment. Swap the env and the tool definition; everything else stays the same. + +## What you'll use + +| | | +|---|---| +| **Student model** | [`Qwen/Qwen3-1.7B`](https://huggingface.co/Qwen/Qwen3-1.7B) | +| **Teacher model** | `gpt-5-mini` via the OpenAI API | +| **Environment** | [`reasoning_gym_env`](https://github.com/meta-pytorch/OpenEnv/tree/main/envs/reasoning_gym_env) / `chain_sum` | +| **SFT trainer** | [TRL `SFTTrainer`](https://huggingface.co/docs/trl/main/en/sft_trainer) | +| **Next step** | [End-to-end walkthrough with GRPO](https://meta-pytorch.org/OpenEnv/tutorials/end-to-end-walkthrough.html) | + +--- + +## 1. Install dependencies + +```python +!pip install -q openai trl +!pip install -q openenv-core +!pip install -q --no-deps git+https://huggingface.co/spaces/sergiopaniego/reasoning_gym +!pip install -Uq "transformers>=5.3.0" +``` + +--- + +## 2. Set your credentials + +```python +import getpass, os + +if "OPENAI_API_KEY" not in os.environ: + os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API key: ") +``` + +You'll also need a Hugging Face login to download the base model and push both the collected dataset +and the fine-tuned checkpoint: + +```python +from huggingface_hub import notebook_login + +notebook_login() +``` + +```python +YOUR_HF_USERNAME = "your-username" # replace with your Hugging Face username +assert YOUR_HF_USERNAME != "your-username", "Replace YOUR_HF_USERNAME with your Hugging Face username" +``` + +--- + +## 3. Define the system prompt + +Use the same prompt as the [GRPO tutorial](https://meta-pytorch.org/OpenEnv/tutorials/end-to-end-walkthrough.html) +so the SFT-trained model is a drop-in replacement when you continue with GRPO. + +```python +SYSTEM_PROMPT = """You are a careful arithmetic assistant. + +You will be given a chain of integer additions. Compute the result and submit it as a single number. + +Rules: +1. Read the question carefully. +2. Use the tool `answer` exactly once with your final number. +3. The answer must be a single integer with no units or explanation. +""" +``` + +--- + +## 4. Configure data collection + +`DATASET_CONFIG` controls the difficulty of the `chain_sum` problems the environment generates: +`min_terms`/`max_terms` set how many integers are added together, and `min_digits`/`max_digits` set +how many digits each integer has. At these settings each problem is a sum of 2–3 two-digit numbers +— easy enough for `gpt-5-mini` to answer correctly ~90% of the time, which gives a clean training +signal after filtering. + +`N_EPISODES` is the number of problems to collect. 300 is enough to get ~270 correct examples after +filtering, which is sufficient for format compliance training. + +```python +DATASET_CONFIG = { + "min_terms": 2, + "max_terms": 3, + "min_digits": 2, + "max_digits": 2, +} + +N_EPISODES = 300 +``` + +--- + +## 5. Collect rollouts with `openenv collect` + +`openenv collect` runs the teacher model inside the environment and records every episode — the +environment's `step()` reward is written alongside the messages, so filtering by correctness requires +no additional scoring code. + +```python +import json, shlex + +dataset_config_arg = shlex.quote(json.dumps(DATASET_CONFIG)) +system_prompt_arg = shlex.quote(SYSTEM_PROMPT) +hub_repo_arg = shlex.quote(f"{YOUR_HF_USERNAME}/chain-sum-rollouts") + +!openenv collect reasoning_gym:chain_sum \ + --base-url https://sergiopaniego-reasoning-gym.hf.space \ + --provider openai \ + --model gpt-5-mini \ + --num-episodes {N_EPISODES} \ + --max-tokens 1024 \ + --dataset-config {dataset_config_arg} \ + --system-prompt {system_prompt_arg} \ + --push-to-hub {hub_repo_arg} \ + --output-dir ./rollouts +``` + +The command prints a live progress summary and pushes the collected episodes to the Hub as +`{YOUR_HF_USERNAME}/chain-sum-rollouts`. Pull them back to start filtering: + +```python +from datasets import load_dataset + +ds = load_dataset(f"{YOUR_HF_USERNAME}/chain-sum-rollouts", split="train") +raw_rollouts = list(ds) +print(f"Collected {len(raw_rollouts)} episodes") +``` + +The `messages` field stores the full conversation in standard OpenAI format (assistant messages have +a `tool_calls` list). Convert to Qwen3's `` text format before training — GRPOTrainer +produces this same format during RL, so the SFT checkpoint becomes a direct drop-in: + +```python +def to_qwen3_messages(record): + converted = [] + for msg in record["messages"]: + if msg["role"] == "tool": + continue # strip environment responses; SFT only needs the assistant turn + if msg["role"] == "assistant" and msg.get("tool_calls"): + tc = msg["tool_calls"][0] + args = json.loads(tc["function"]["arguments"]) + answer_str = args.get("answer", "") + tool_call_text = ( + "\n" + + json.dumps({"name": "answer", "arguments": {"answer": answer_str}}) + + "\n" + ) + converted.append({"role": "assistant", "content": tool_call_text}) + else: + converted.append(msg) + return {"messages": converted, "reward": record["reward"]} + +rollouts = [to_qwen3_messages(r) for r in raw_rollouts] +``` + +--- + +## 6. Filter the dataset + +Keep only episodes where the teacher answered correctly. The environment's reward signal does the +labelling — no manual annotation needed. + +```python +correct = [r for r in rollouts if r["reward"] == 1.0] +print(f"Correct: {len(correct)} / {len(rollouts)} ({len(correct)/len(rollouts):.1%})") +``` + +`gpt-5-mini` typically scores above 90% on `chain_sum` at this difficulty, so you should end up with +~270 examples from 300 rollouts. + +--- + +## 7. Inspect the dataset before training + +Always look at your data before training. Automated collection can introduce unexpected patterns that the +student model will learn to imitate. + +```python +import random + +for row in random.sample(correct, 3): + question = row["messages"][0]["content"] + response = row["messages"][1]["content"] + print(f"Q: {question}") + print(f"A: {response}") + print() +``` + +Things to check: + +- Does every response contain a valid `` block? +- Are the answers integers with no extra text? +- Is there any reasoning in the assistant message that you don't want the student to learn? + (For example: an internal monologue, disclaimers, or repeated phrasing that the teacher leaked + from its own system prompt.) + +--- + +## 8. Measure token lengths + +Set `max_length` in `SFTConfig` to cover nearly all examples without wasting GPU memory on padding. +The 99th percentile is a good target: you truncate fewer than 1% of examples while keeping batches tight. + +```python +import numpy as np +from transformers import AutoTokenizer + +tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-1.7B") + +lengths = [] +for row in correct: + text = tokenizer.apply_chat_template( + row["messages"], tokenize=False, add_generation_prompt=False + ) + ids = tokenizer.encode(text) + lengths.append(len(ids)) + +lengths = np.array(lengths) +MAX_SEQ_LEN = int(np.percentile(lengths, 99)) + 16 + +print( + f"p50={np.percentile(lengths, 50):.0f} " + f"p95={np.percentile(lengths, 95):.0f} " + f"p99={np.percentile(lengths, 99):.0f} " + f"max={lengths.max()}" +) +print(f"Setting MAX_SEQ_LEN = {MAX_SEQ_LEN}") +``` + +--- + +## 9. Fine-tune with SFTTrainer + +`assistant_only_loss=True` in `SFTConfig` masks the prompt tokens so the loss is computed only on the +assistant response — the `` block. This is more efficient than full-sequence training and avoids +accidentally reinforcing the system prompt wording. + +```python +from datasets import Dataset +from transformers import AutoModelForCausalLM +from trl import SFTConfig, SFTTrainer + +dataset = Dataset.from_list([{"messages": r["messages"]} for r in correct]) + +model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen3-1.7B") + +sft_config = SFTConfig( + output_dir="reasoning-gym-chain-sum-Qwen3-1.7B-sft", + max_length=MAX_SEQ_LEN, + num_train_epochs=3, + per_device_train_batch_size=4, + gradient_accumulation_steps=2, + learning_rate=2e-5, + warmup_steps=10, + lr_scheduler_type="cosine", + logging_steps=5, + save_strategy="no", + assistant_only_loss=True, +) + +trainer = SFTTrainer( + model=model, + train_dataset=dataset, + processing_class=tokenizer, + args=sft_config, +) + +trainer.train() +trainer.push_to_hub(commit_message="SFT warm-up on reasoning_gym chain_sum") +``` + +```{note} +Training ~270 examples for 3 epochs takes around 5 minutes on a single A100 (40 GB). The goal is format +compliance, not task mastery — a handful of epochs is enough. Mastery comes from GRPO. +``` + +--- + +## 10. Evaluate: before vs after + +Run both the base model and the SFT checkpoint on a held-out set and compare. The key metric for a +warm-up evaluation is **format compliance** — how often the model uses `` correctly — as +well as overall accuracy. + +```python +import re +from transformers import pipeline +from reasoning_gym_env.client import ReasoningGymEnv +from reasoning_gym_env.models import ReasoningGymAction + + +async def evaluate_model(model_name, n_eval=50, seed=999): + gen = pipeline( + "text-generation", + model=model_name, + tokenizer=model_name, + device_map="auto", + dtype="auto", + ) + gen.model.generation_config.max_length = None + tok = AutoTokenizer.from_pretrained(model_name) + env = ReasoningGymEnv(base_url="https://sergiopaniego-reasoning-gym.hf.space") + + obs = await env.reset( + dataset_name="chain_sum", + dataset_config=DATASET_CONFIG, + seed=seed, + size=n_eval, + ) + + rewards, format_hits = [], 0 + + for i in range(n_eval): + if i > 0: + obs = await env.reset() + + question = obs.observation.question + messages = [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": question}, + ] + prompt = tok.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + completion = gen(prompt, max_new_tokens=64)[0]["generated_text"][len(prompt):] + + m = re.search(r'"answer"\s*:\s*"?(\d+)"?', completion) + if m: + format_hits += 1 + answer = m.group(1) + else: + nums = re.findall(r"\b(\d+)\b", completion) + answer = nums[-1] if nums else "0" + + result = await env.step(ReasoningGymAction(answer=answer)) + rewards.append(float(result.observation.score or 0.0)) + + await env.close() + del gen # free GPU memory before loading the next model + + return { + "accuracy": sum(rewards) / len(rewards), + "format_compliance": format_hits / n_eval, + } + + +base_metrics = await evaluate_model("Qwen/Qwen3-1.7B") +sft_metrics = await evaluate_model(f"{YOUR_HF_USERNAME}/reasoning-gym-chain-sum-Qwen3-1.7B-sft") + +print(f"\n{'Metric':<25} {'Base model':>12} {'After SFT':>12} {'Delta':>10}") +print("-" * 62) +for key, label in [("format_compliance", "Format compliance"), ("accuracy", "Accuracy")]: + b, s = base_metrics[key], sft_metrics[key] + print(f"{label:<25} {b:>12.1%} {s:>12.1%} {(s - b) * 100:>+9.1f} pp") +``` + +A successful warm-up looks like this: + +| Metric | Base model | After SFT | Delta | +|---|---|---|---| +| Format compliance | ~0% | ~68% | +68 pp | +| Accuracy | ~4% | ~68% | +64 pp | + +Format compliance should jump sharply from near-zero — that's the primary goal. `Qwen3-1.7B` produces +essentially no valid `` blocks out of the box. After SFT on ~270 examples, the model reliably +uses the format, and accuracy rises in lockstep because correct format is a prerequisite for the +environment's scorer to award any credit. + +--- + +## 11. Where to go next: GRPO + +The SFT checkpoint is ready to use as the starting model for GRPO. In the +[end-to-end walkthrough](https://meta-pytorch.org/OpenEnv/tutorials/end-to-end-walkthrough.html), +change one line in section 8: + +```python +# Before (cold-start from the base model): +MODEL_NAME = "Qwen/Qwen3-1.7B" + +# After (warm-start from your SFT checkpoint): +MODEL_NAME = f"{YOUR_HF_USERNAME}/reasoning-gym-chain-sum-Qwen3-1.7B-sft" +``` + +With format compliance already near 100%, GRPO's `reward_std` will be non-zero from the very first +batch and the reward curve will climb immediately — no cold-start stall. + +**Other directions:** + +- **Harder tasks.** Increase `max_terms` or `max_digits` in `DATASET_CONFIG` and collect a new SFT set. + Once the student handles easier examples reliably, a harder GRPO phase can push further. +- **Different environments.** The same pipeline — teacher collects → filter → SFT → GRPO — applies to + any OpenEnv environment. Swap `reasoning_gym_env` and the `answer` tool definition for your env's + tool surface. +- **Larger teacher.** `gpt-5` or `claude-opus-4` as teacher will yield higher-quality examples, + especially for tasks where `gpt-5-mini` struggles. diff --git a/envs/browsergym_env/README.md b/envs/browsergym_env/README.md index 93a46d7ca9..8c92eccda7 100644 --- a/envs/browsergym_env/README.md +++ b/envs/browsergym_env/README.md @@ -73,6 +73,38 @@ for episode in range(1000): env.close() ``` +## Harness Sessions for TRL + +If you want BrowserGym to participate in a tool-driven harness instead of a +hand-written `env.reset()` / `env.step()` loop, use the BrowserGym session +factory: + +```python +from browsergym_env import BrowserGymEnv +from browsergym_env.harness import BrowserGymSessionFactory +from openenv.core.harness import ( + HarnessRunLimits, + MCPHarnessAdapter, + build_harness_rollout_func, +) + +session_factory = BrowserGymSessionFactory( + client_factory=lambda: BrowserGymEnv(base_url="https://openenv-browsergym-env.hf.space"), +) + +rollout_func = build_harness_rollout_func( + session_factory=session_factory, + harness_adapter=MCPHarnessAdapter(), + model_step_builder=..., # trainer-owned model sampling + limits=HarnessRunLimits(max_turns=10), +) +``` + +BrowserGym exposes `click`, `fill`, `send_keys`, `scroll`, and `noop` as MCP-style +tools while still translating them back into the underlying `BrowserGymAction` +strings. See [examples/browsergym_harness.py](https://github.com/meta-pytorch/OpenEnv/blob/main/examples/browsergym_harness.py) +for a full TRL-oriented example. + ### Available Tasks by Benchmark #### MiniWoB++ Tasks (Training - 100+ tasks) diff --git a/envs/browsergym_env/__init__.py b/envs/browsergym_env/__init__.py index ac4bda82b5..23cd8b2c95 100644 --- a/envs/browsergym_env/__init__.py +++ b/envs/browsergym_env/__init__.py @@ -61,12 +61,45 @@ ``` """ +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING + from .client import BrowserGymEnv from .models import BrowserGymAction, BrowserGymObservation, BrowserGymState +if TYPE_CHECKING: + from .harness import BrowserGymSessionFactory, build_browsergym_action_tool_call + __all__ = [ "BrowserGymEnv", "BrowserGymAction", "BrowserGymObservation", "BrowserGymState", + "BrowserGymSessionFactory", + "build_browsergym_action_tool_call", ] + +_LAZY_ATTRS = { + "BrowserGymSessionFactory": (".harness", "BrowserGymSessionFactory"), + "build_browsergym_action_tool_call": ( + ".harness", + "build_browsergym_action_tool_call", + ), +} + + +def __getattr__(name: str): + if name not in _LAZY_ATTRS: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + module_path, attr_name = _LAZY_ATTRS[name] + module = import_module(module_path, __name__) + value = getattr(module, attr_name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals().keys()) | set(__all__)) diff --git a/envs/browsergym_env/harness.py b/envs/browsergym_env/harness.py new file mode 100644 index 0000000000..186f260f49 --- /dev/null +++ b/envs/browsergym_env/harness.py @@ -0,0 +1,290 @@ +"""Harness-oriented BrowserGym session adapters.""" + +from __future__ import annotations + +import ast +import json +from typing import Any + +from openenv.core.env_server.mcp_types import Tool +from openenv.core.harness import ( + ResourceSessionFactory, + StepEnvSessionAdapter, + ToolResult, + VerifyResult, +) +from openenv.core.llm_client import ToolCall + +from .models import BrowserGymAction + +_BROWSERGYM_TOOLS = [ + Tool( + name="click", + description="Click an element by BrowserGym bid.", + input_schema={ + "type": "object", + "properties": { + "bid": {"type": "string"}, + }, + "required": ["bid"], + }, + ), + Tool( + name="fill", + description="Fill an input field by bid.", + input_schema={ + "type": "object", + "properties": { + "bid": {"type": "string"}, + "text": {"type": "string"}, + }, + "required": ["bid", "text"], + }, + ), + Tool( + name="send_keys", + description="Send keyboard input to the page.", + input_schema={ + "type": "object", + "properties": { + "text": {"type": "string"}, + }, + "required": ["text"], + }, + ), + Tool( + name="scroll", + description="Scroll the page up or down.", + input_schema={ + "type": "object", + "properties": { + "direction": {"type": "string", "enum": ["up", "down"]}, + }, + "required": ["direction"], + }, + ), + Tool( + name="noop", + description="Take no action on the current page.", + input_schema={"type": "object", "properties": {}}, + ), +] + + +def _quote(value: str) -> str: + return json.dumps(value, ensure_ascii=False) + + +def build_browsergym_action_str(tool_name: str, arguments: dict[str, Any]) -> str: + """Convert a BrowserGym tool call into the action string the env expects.""" + + if tool_name == "click": + return f"click({_quote(str(arguments['bid']))})" + if tool_name == "fill": + return ( + f"fill({_quote(str(arguments['bid']))}, {_quote(str(arguments['text']))})" + ) + if tool_name == "send_keys": + return f"send_keys({_quote(str(arguments['text']))})" + if tool_name == "scroll": + return f"scroll({_quote(str(arguments['direction']))})" + if tool_name == "noop": + return "noop()" + + raise KeyError(f"Unsupported BrowserGym tool: {tool_name}") + + +def _format_browsergym_prompt(observation: Any, task: Any) -> str: + goal = getattr(observation, "goal", "") or (task or "") + page_text = getattr(observation, "axtree_txt", "") or getattr( + observation, "text", "" + ) + error = getattr(observation, "error", "") + + parts = [] + if goal: + parts.append(f"Goal: {goal}") + if error: + parts.append(f"Previous action error: {error}") + if page_text: + parts.append(f"Page structure:\n{page_text}") + parts.append("Choose the next browser action.") + return "\n\n".join(parts) + + +def _build_browsergym_tool_result( + task: Any, +): + def builder( + tool_name: str, + arguments: dict[str, Any], + result: Any, + state: Any, + ) -> ToolResult: + observation = result.observation + data = { + "tool_name": tool_name, + "arguments": dict(arguments), + "goal": getattr(observation, "goal", "") or (task or ""), + "observation_text": getattr(observation, "axtree_txt", "") + or getattr(observation, "text", ""), + "url": getattr(observation, "url", ""), + "error": getattr(observation, "error", ""), + "last_action_error": getattr(observation, "last_action_error", False), + "reward": result.reward, + "done": result.done, + } + metadata = { + "reward": result.reward, + "state": state.model_dump() if hasattr(state, "model_dump") else state, + } + return ToolResult(data=data, done=bool(result.done), metadata=metadata) + + return builder + + +def _build_browsergym_verify( + transcript: list[dict[str, Any]], + final_state: Any | None, + last_result: Any | None, + state: Any, +) -> VerifyResult: + reward = None if last_result is None else last_result.reward + done = False if last_result is None else bool(last_result.done) + metrics = { + "step_count": getattr(state, "step_count", 0), + "cum_reward": getattr(state, "cum_reward", reward or 0.0), + "benchmark": getattr(state, "benchmark", ""), + "task_name": getattr(state, "task_name", ""), + } + artifacts = { + "final_state": state.model_dump() if hasattr(state, "model_dump") else state, + "final_rollout": final_state, + "transcript_length": len(transcript), + } + return VerifyResult( + env_reward=reward, + done=done, + metrics=metrics, + artifacts=artifacts, + ) + + +class BrowserGymSessionFactory(ResourceSessionFactory): + """Create BrowserGym-backed resource sessions from client factories.""" + + def __init__(self, client_factory, *, default_task: str | None = None): + self._client_factory = client_factory + self._default_task = default_task + + def create( + self, + task: Any, + seed: int | None = None, + episode_id: str | None = None, + ) -> StepEnvSessionAdapter: + session_task = task if task is not None else self._default_task + client = self._client_factory() + + reset_kwargs = {} + if session_task is not None: + reset_kwargs["task_name"] = session_task + + return StepEnvSessionAdapter( + client=client, + task=session_task, + seed=seed, + episode_id=episode_id, + tool_specs=list(_BROWSERGYM_TOOLS), + action_builder=lambda name, arguments: BrowserGymAction( + action_str=build_browsergym_action_str(name, arguments) + ), + initial_messages_builder=lambda result, current_task: [ + { + "role": "user", + "content": _format_browsergym_prompt( + result.observation, + current_task, + ), + } + ], + tool_result_builder=_build_browsergym_tool_result(session_task), + verify_builder=_build_browsergym_verify, + reset_kwargs=reset_kwargs, + ) + + +def _parse_action_call(action_text: str) -> tuple[str, list[Any]]: + try: + expression = ast.parse(action_text.strip(), mode="eval").body + except SyntaxError as exc: + raise ValueError(f"Unsupported BrowserGym action: {action_text}") from exc + + if not isinstance(expression, ast.Call) or not isinstance( + expression.func, ast.Name + ): + raise ValueError(f"Unsupported BrowserGym action: {action_text}") + if expression.keywords: + raise ValueError("BrowserGym action arguments must be positional") + + args: list[Any] = [] + for arg in expression.args: + try: + args.append(ast.literal_eval(arg)) + except (SyntaxError, ValueError) as exc: + raise ValueError("BrowserGym action arguments must be literals") from exc + + return expression.func.id, args + + +def _expect_str(value: Any, argument_name: str) -> str: + if not isinstance(value, str): + raise ValueError(f"BrowserGym {argument_name} argument must be a string") + return value + + +def build_browsergym_action_tool_call(action_text: str) -> ToolCall: + """Parse a text BrowserGym action into a structured tool call.""" + + name, args = _parse_action_call(action_text) + if name == "click" and len(args) == 1: + return ToolCall( + id="browsergym-click", + name="click", + args={"bid": _expect_str(args[0], "bid")}, + ) + if name == "fill" and len(args) == 2: + return ToolCall( + id="browsergym-fill", + name="fill", + args={ + "bid": _expect_str(args[0], "bid"), + "text": _expect_str(args[1], "text"), + }, + ) + if name == "send_keys" and len(args) == 1: + return ToolCall( + id="browsergym-send_keys", + name="send_keys", + args={"text": _expect_str(args[0], "text")}, + ) + if name == "scroll" and len(args) == 1: + direction = _expect_str(args[0], "direction") + if direction not in {"up", "down"}: + raise ValueError("BrowserGym scroll direction must be 'up' or 'down'") + return ToolCall( + id="browsergym-scroll", + name="scroll", + args={"direction": direction}, + ) + if name == "noop" and not args: + return ToolCall(id="browsergym-noop", name="noop", args={}) + + raise ValueError(f"Unsupported BrowserGym action: {action_text}") + + +__all__ = [ + "BrowserGymSessionFactory", + "build_browsergym_action_str", + "build_browsergym_action_tool_call", +] diff --git a/envs/browsergym_env/models.py b/envs/browsergym_env/models.py index 452404374d..33667c24d3 100644 --- a/envs/browsergym_env/models.py +++ b/envs/browsergym_env/models.py @@ -21,7 +21,7 @@ class BrowserGymAction(Action): - "click('Submit button')" - "fill('username', 'john@example.com')" - "goto('https://example.com')" - - "scroll(down)" + - "scroll('down')" - "send_keys('Enter')" """ diff --git a/envs/openspiel_env/harness.py b/envs/openspiel_env/harness.py new file mode 100644 index 0000000000..c789b59c5a --- /dev/null +++ b/envs/openspiel_env/harness.py @@ -0,0 +1,169 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Harness-oriented OpenSpiel session adapters. + +Follows the pattern introduced by ``browsergym_env.harness``: exposes an +OpenSpiel client as a ``ResourceSession`` driven through MCP-style tools, +so it can be consumed by ``openenv.core.harness`` adapters (e.g. the TTT +teacher rollouts used by the collect pipeline). +""" + +from __future__ import annotations + +from typing import Any, Callable + +from openenv.core.env_server.mcp_types import Tool +from openenv.core.harness import ( + ResourceSessionFactory, + StepEnvSessionAdapter, + ToolResult, +) + +from .client import OpenSpielEnv +from .models import OpenSpielAction + +_OPENSPIEL_TOOLS: list[Tool] = [ + Tool( + name="play_move", + description=( + "Play a move by its integer action_id. The action_id must come " + "from the `legal_actions` list shown in the latest observation." + ), + input_schema={ + "type": "object", + "properties": { + "action_id": { + "type": "integer", + "description": "Action id drawn from legal_actions.", + }, + }, + "required": ["action_id"], + }, + ), +] + + +def render_tic_tac_toe_board(info_state: list[float]) -> str: + """Render a TTT ``info_state`` (27 floats, 3 planes) as a 3x3 grid. + + OpenSpiel's TTT info_state stores an empty/X/O one-hot for each cell + across three consecutive 9-cell planes. Empty cells are rendered with + their action_id so the prompt doubles as an action legend. + """ + if len(info_state) != 27: + return "" + + rows: list[str] = [] + for r in range(3): + cells: list[str] = [] + for c in range(3): + idx = r * 3 + c + x_plane = info_state[9 + idx] + o_plane = info_state[18 + idx] + if x_plane > 0.5: + cells.append("X") + elif o_plane > 0.5: + cells.append("O") + else: + cells.append(str(idx)) + rows.append(" | ".join(cells)) + return "\n---------\n".join(rows) + + +def _format_initial_prompt(observation: Any, game_name: str) -> str: + legal = list(getattr(observation, "legal_actions", []) or []) + info_state = list(getattr(observation, "info_state", []) or []) + lines = [ + f"You are playing {game_name}.", + "Call the play_move tool with an action_id drawn from legal_actions.", + f"Legal actions: {legal}", + ] + if game_name == "tic_tac_toe": + board = render_tic_tac_toe_board(info_state) + if board: + lines.extend(["", "Board:", board]) + return "\n".join(lines) + + +def _build_tool_result(game_name: str) -> Callable[..., ToolResult]: + def builder( + tool_name: str, + arguments: dict[str, Any], + result: Any, + state: Any, + ) -> ToolResult: + observation = result.observation + legal = list(getattr(observation, "legal_actions", []) or []) + info_state = list(getattr(observation, "info_state", []) or []) + data: dict[str, Any] = { + "action_id": arguments.get("action_id"), + "legal_actions": legal, + "reward": result.reward, + "done": result.done, + "current_player_id": getattr(observation, "current_player_id", None), + } + if game_name == "tic_tac_toe": + board = render_tic_tac_toe_board(info_state) + if board: + data["board"] = board + return ToolResult( + data=data, + done=bool(result.done), + metadata={ + "reward": result.reward, + "state": state.model_dump() if hasattr(state, "model_dump") else state, + }, + ) + + return builder + + +class OpenSpielSessionFactory(ResourceSessionFactory): + """Create OpenSpiel-backed resource sessions for harness rollouts.""" + + def __init__( + self, + client_factory: Callable[[], OpenSpielEnv], + *, + game_name: str = "tic_tac_toe", + ): + self._client_factory = client_factory + self._game_name = game_name + + def create( + self, + task: Any = None, + seed: int | None = None, + episode_id: str | None = None, + ) -> StepEnvSessionAdapter: + client = self._client_factory() + game = self._game_name + + return StepEnvSessionAdapter( + client=client, + task=task, + seed=seed, + episode_id=episode_id, + tool_specs=list(_OPENSPIEL_TOOLS), + action_builder=lambda name, arguments: OpenSpielAction( + action_id=int(arguments["action_id"]), + game_name=game, + ), + initial_messages_builder=lambda result, current_task: [ + { + "role": "user", + "content": _format_initial_prompt(result.observation, game), + } + ], + tool_result_builder=_build_tool_result(game), + ) + + +__all__ = [ + "OpenSpielSessionFactory", + "render_tic_tac_toe_board", +] diff --git a/envs/reasoning_gym_env/harness.py b/envs/reasoning_gym_env/harness.py new file mode 100644 index 0000000000..36f1c89d89 --- /dev/null +++ b/envs/reasoning_gym_env/harness.py @@ -0,0 +1,123 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Harness-oriented Reasoning Gym session adapters. + +Follows the pattern introduced by ``openspiel_env.harness``: exposes a +ReasoningGymEnv client as a ``ResourceSession`` driven through MCP-style +tools, so it can be consumed by ``openenv.core.harness`` adapters (e.g. +the collect pipeline). +""" + +from __future__ import annotations + +import random +from typing import Any, Callable + +from openenv.core.env_server.mcp_types import Tool +from openenv.core.harness import ( + ResourceSessionFactory, + StepEnvSessionAdapter, + ToolResult, +) + +from .client import ReasoningGymEnv +from .models import ReasoningGymAction + +_REASONING_GYM_TOOLS: list[Tool] = [ + Tool( + name="answer", + description="Submit the final answer for the current question.", + input_schema={ + "type": "object", + "properties": { + "answer": { + "type": "string", + "description": "The answer to submit.", + }, + }, + "required": ["answer"], + }, + ), +] + + +def _build_tool_result() -> Callable[..., ToolResult]: + def builder( + tool_name: str, + arguments: dict[str, Any], + result: Any, + state: Any, + ) -> ToolResult: + observation = result.observation + return ToolResult( + data={ + "answer": arguments.get("answer"), + "score": observation.score, + "correct_answer": observation.correct_answer, + "reward": result.reward, + "done": result.done, + }, + done=bool(result.done), + metadata={ + "reward": result.reward, + "state": state.model_dump() if hasattr(state, "model_dump") else state, + }, + ) + + return builder + + +class ReasoningGymSessionFactory(ResourceSessionFactory): + """Create ReasoningGym-backed resource sessions for harness rollouts.""" + + def __init__( + self, + client_factory: Callable[[], ReasoningGymEnv], + *, + dataset_name: str, + dataset_config: dict[str, Any] | None = None, + ): + self._client_factory = client_factory + self._dataset_name = dataset_name + self._dataset_config = dataset_config or {} + + def create( + self, + task: Any = None, + seed: int | None = None, + episode_id: str | None = None, + ) -> StepEnvSessionAdapter: + client = self._client_factory() + effective_seed = seed if seed is not None else random.randint(0, 2**31 - 1) + + return StepEnvSessionAdapter( + client=client, + task=task, + seed=effective_seed, + episode_id=episode_id, + tool_specs=list(_REASONING_GYM_TOOLS), + action_builder=lambda name, arguments: ReasoningGymAction( + answer=str(arguments["answer"]), + ), + initial_messages_builder=lambda result, current_task: [ + { + "role": "user", + "content": result.observation.question, + } + ], + tool_result_builder=_build_tool_result(), + reset_kwargs={ + "dataset_name": self._dataset_name, + "dataset_config": self._dataset_config, + "size": 1, + }, + ) + + +__all__ = [ + "ReasoningGymSessionFactory", +] diff --git a/examples/browsergym_codex_eval.py b/examples/browsergym_codex_eval.py new file mode 100644 index 0000000000..7841a6c7a1 --- /dev/null +++ b/examples/browsergym_codex_eval.py @@ -0,0 +1,279 @@ +"""Experimental BrowserGym evaluation example using Codex CLI as the harness. + +This example is eval-only. Codex runs as a black-box CLI agent and connects to +BrowserGym through a local HTTP MCP bridge. It is not suitable for white-box RL +training or any setup that needs token IDs or logprobs from each model step. + +Manual prerequisites: +- `codex` installed and authenticated +- BrowserGym Docker image available locally, or a running BrowserGym server +""" + +from __future__ import annotations + +import argparse +import os +import shutil +import subprocess +import sys +import tempfile +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT / "src")) +sys.path.insert(0, str(_REPO_ROOT / "envs")) + +from browsergym_harness_eval_common import ( + build_browsergym_session_factory, + DEFAULT_BENCHMARK, + DEFAULT_BROWSERGYM_IMAGE, + DEFAULT_MAX_STEPS, + DEFAULT_TASK_NAME, + format_episode_summary, + run_black_box_episode, + SessionMCPHttpServer, + start_browsergym_runtime, + summarize_episodes, +) +from openenv.core.harness import ( + CLIHarnessAdapter, + HarnessRolloutResult, + HarnessRunLimits, + RolloutEvent, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Experimental BrowserGym evaluation through Codex CLI + MCP.", + ) + parser.add_argument( + "--base-url", + default=None, + help="Use an already-running BrowserGym server instead of starting Docker.", + ) + parser.add_argument( + "--image", + default=DEFAULT_BROWSERGYM_IMAGE, + help="Docker image to start when --base-url is not provided.", + ) + parser.add_argument( + "--benchmark", + default=DEFAULT_BENCHMARK, + help="BrowserGym benchmark to run when starting Docker.", + ) + parser.add_argument( + "--task-name", + default=DEFAULT_TASK_NAME, + help="BrowserGym task name to evaluate.", + ) + parser.add_argument( + "--episodes", + type=int, + default=1, + help="Number of episodes to evaluate.", + ) + parser.add_argument( + "--max-steps", + type=int, + default=DEFAULT_MAX_STEPS, + help="Maximum browser actions Codex should take.", + ) + parser.add_argument( + "--codex-bin", + default="codex", + help="Path to the Codex CLI binary.", + ) + parser.add_argument( + "--artifact-dir", + default=None, + help="Directory where stdout and the last-message artifact will be written.", + ) + parser.add_argument( + "--timeout-seconds", + type=int, + default=300, + help="Timeout for each Codex run.", + ) + return parser.parse_args() + + +def build_codex_prompt( + initial_messages: list[dict[str, object]], max_steps: int +) -> str: + initial_context = "\n\n".join( + str(message.get("content", "")).strip() + for message in initial_messages + if str(message.get("content", "")).strip() + ) + return ( + "You are evaluating a BrowserGym task.\n" + "Use only the MCP tools exposed by the 'browsergym' server.\n" + "Do not use shell commands, files, or any other tool.\n" + f"Take at most {max_steps} browser actions.\n" + "Stop once the task is complete or no useful progress is possible.\n\n" + "Initial BrowserGym context:\n" + f"{initial_context}" + ) + + +def build_codex_runner( + *, + codex_bin: str, + artifact_root: Path, + timeout_seconds: int, +): + def runner(bridge, session, limits): + if shutil.which(codex_bin) is None: + raise RuntimeError(f"Codex binary not found: {codex_bin}") + + initial_messages = session.initial_messages() + prompt_text = build_codex_prompt(initial_messages, limits.max_turns) + + artifact_root.mkdir(parents=True, exist_ok=True) + prompt_path = artifact_root / "codex_prompt.txt" + stdout_path = artifact_root / "codex_exec.jsonl" + stderr_path = artifact_root / "codex_exec.stderr" + last_message_path = artifact_root / "codex_last_message.txt" + prompt_path.write_text(prompt_text, encoding="utf-8") + + with tempfile.TemporaryDirectory(prefix="browsergym-codex-home-") as codex_home: + with tempfile.TemporaryDirectory( + prefix="browsergym-codex-cwd-" + ) as codex_cwd: + env = dict(os.environ) + env["CODEX_HOME"] = codex_home + + with SessionMCPHttpServer(bridge).start() as mcp_server: + add_cmd = [ + codex_bin, + "mcp", + "add", + "browsergym", + "--url", + mcp_server.url, + ] + add_result = subprocess.run( + add_cmd, + capture_output=True, + text=True, + env=env, + check=False, + timeout=timeout_seconds, + ) + if add_result.returncode != 0: + raise RuntimeError( + "Failed to register BrowserGym MCP server with Codex:\n" + f"{add_result.stderr.strip() or add_result.stdout.strip()}" + ) + + exec_cmd = [ + codex_bin, + "exec", + "--json", + "--skip-git-repo-check", + "-s", + "read-only", + "-C", + codex_cwd, + "--output-last-message", + str(last_message_path), + prompt_text, + ] + exec_result = subprocess.run( + exec_cmd, + capture_output=True, + text=True, + env=env, + check=False, + timeout=timeout_seconds, + ) + + stdout_path.write_text(exec_result.stdout, encoding="utf-8") + stderr_path.write_text(exec_result.stderr, encoding="utf-8") + + last_message = "" + if last_message_path.exists(): + last_message = last_message_path.read_text(encoding="utf-8").strip() + + return HarnessRolloutResult( + messages=[ + {"role": "user", "content": prompt_text}, + {"role": "assistant", "content": last_message}, + ], + events=[ + RolloutEvent( + type="codex_exec", + payload={ + "returncode": exec_result.returncode, + "artifacts": { + "prompt": str(prompt_path), + "stdout": str(stdout_path), + "stderr": str(stderr_path), + "last_message": str(last_message_path), + }, + }, + ) + ], + done=exec_result.returncode == 0, + metrics={ + "mode": "black_box", + "codex_returncode": exec_result.returncode, + "artifact_dir": str(artifact_root), + }, + ) + + return runner + + +def main() -> None: + args = parse_args() + artifact_root = Path( + args.artifact_dir or tempfile.mkdtemp(prefix="browsergym-codex-artifacts-") + ).resolve() + + with start_browsergym_runtime( + base_url=args.base_url, + image=args.image, + benchmark=args.benchmark, + task_name=args.task_name, + ) as runtime: + print(f"BrowserGym server: {runtime.base_url}") + print(f"Artifact dir: {artifact_root}") + + session_factory = build_browsergym_session_factory( + base_url=runtime.base_url, + task_name=args.task_name, + ) + episodes = [] + limits = HarnessRunLimits(max_turns=args.max_steps) + for index in range(1, args.episodes + 1): + episode_artifact_root = artifact_root / f"episode-{index}" + episode_adapter = CLIHarnessAdapter( + runner=build_codex_runner( + codex_bin=args.codex_bin, + artifact_root=episode_artifact_root, + timeout_seconds=args.timeout_seconds, + ) + ) + episode = run_black_box_episode( + session_factory=session_factory, + harness_adapter=episode_adapter, + limits=limits, + episode_id=f"browsergym-codex-eval-{index}", + ) + episodes.append(episode) + print(format_episode_summary(index, episode)) + + summary = summarize_episodes(episodes) + print( + "Aggregate: " + f"episodes={int(summary['episodes'])} " + f"avg_reward={summary['avg_reward']:.2f} " + f"success_rate={summary['success_rate']:.2%} " + f"avg_steps={summary['avg_steps']:.2f}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/browsergym_harness.py b/examples/browsergym_harness.py new file mode 100644 index 0000000000..fffbab9fe0 --- /dev/null +++ b/examples/browsergym_harness.py @@ -0,0 +1,154 @@ +# Copyright 2020-2026 The HuggingFace Team. 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. + +# /// script +# dependencies = [ +# "trl[vllm]", +# "peft", +# "trackio", +# "kernels", +# "openenv-browsergym @ git+https://huggingface.co/spaces/openenv/browsergym_env", +# ] +# /// + +"""Harness-oriented BrowserGym rollout example for TRL. + +This mirrors the usual BrowserGym + TRL setup, but moves environment interaction +behind the ResourceSession / HarnessAdapter abstractions so the same session +surface can be reused by white-box training and black-box evaluation harnesses. +""" + +from __future__ import annotations + +import re + +from browsergym_env import BrowserGymEnv +from browsergym_env.harness import ( + BrowserGymSessionFactory, + build_browsergym_action_tool_call, +) +from transformers import AutoTokenizer + +from openenv.core.harness import ( + HarnessRunLimits, + MCPHarnessAdapter, + ModelStepResult, + build_harness_rollout_func, +) +from openenv.core.llm_client import LLMResponse +from trl.experimental.openenv import generate_rollout_completions + +SYSTEM_PROMPT = """You control a web browser through BrowserGym actions. + +Reply with exactly one BrowserGym action such as: +- click('13') +- fill('42', 'hello world') +- send_keys('Enter') +- scroll('down') +- noop() +""" + +_ACTION_PREFIX_RE = re.compile(r"^(action|next action)\s*[:-]\s*", re.IGNORECASE) +_ACTION_PATTERN = re.compile( + r"(click|fill|send_keys|scroll|noop)\s*\(.*?\)", + re.IGNORECASE | re.DOTALL, +) + + +def _extract_browsergym_action(response_text: str) -> str: + if not response_text: + return "noop()" + + for raw_line in response_text.splitlines(): + line = _ACTION_PREFIX_RE.sub("", raw_line.strip().strip("`")) + if not line: + continue + if match := _ACTION_PATTERN.search(line): + action = re.sub(r"\s+", " ", match.group(0).strip()).rstrip(";") + open_paren = action.find("(") + if open_paren > 0: + return f"{action[:open_paren].lower()}{action[open_paren:]}" + return action + + if match := _ACTION_PATTERN.search(response_text): + action = re.sub(r"\s+", " ", match.group(0).strip()).rstrip(";") + open_paren = action.find("(") + if open_paren > 0: + return f"{action[:open_paren].lower()}{action[open_paren:]}" + return action + return "noop()" + + +def _build_safe_browsergym_tool_call(response_text: str): + try: + return build_browsergym_action_tool_call( + _extract_browsergym_action(response_text), + ) + except ValueError: + return build_browsergym_action_tool_call("noop()") + + +def build_trl_browsergym_model_step(trainer, tokenizer: AutoTokenizer): + """Wrap TRL generation so the harness can drive BrowserGym via tool calls.""" + + def model_step(messages, tools, sampling): + prompt_text = tokenizer.apply_chat_template( + [{"role": "system", "content": SYSTEM_PROMPT.strip()}, *messages], + add_generation_prompt=True, + tokenize=False, + ) + rollout_output = generate_rollout_completions(trainer, [prompt_text])[0] + completion_text = rollout_output.get("text") or tokenizer.decode( + rollout_output["completion_ids"], + skip_special_tokens=True, + ) + tool_call = _build_safe_browsergym_tool_call(completion_text) + return ModelStepResult( + response=LLMResponse( + content=completion_text, + tool_calls=[tool_call], + ), + prompt_ids=list(rollout_output["prompt_ids"]), + completion_ids=list(rollout_output["completion_ids"]), + logprobs=list(rollout_output["logprobs"]), + ) + + return model_step + + +def build_browsergym_rollout_func( + space_url: str, tokenizer: AutoTokenizer, max_steps: int +): + """Create a TRL rollout function backed by BrowserGym sessions.""" + + session_factory = BrowserGymSessionFactory( + client_factory=lambda: BrowserGymEnv(base_url=space_url), + ) + + return build_harness_rollout_func( + session_factory=session_factory, + harness_adapter=MCPHarnessAdapter(), + model_step_builder=lambda trainer, session: build_trl_browsergym_model_step( + trainer, + tokenizer, + ), + limits=HarnessRunLimits(max_turns=max_steps), + ) + + +__all__ = [ + "SYSTEM_PROMPT", + "build_browsergym_rollout_func", + "build_trl_browsergym_model_step", +] diff --git a/examples/browsergym_harness_eval.py b/examples/browsergym_harness_eval.py new file mode 100644 index 0000000000..4895714756 --- /dev/null +++ b/examples/browsergym_harness_eval.py @@ -0,0 +1,146 @@ +"""White-box BrowserGym evaluation example driven through the harness runtime. + +This is a plain inference/evaluation script. It runs BrowserGym locally by +default, uses an OpenAI-compatible chat model to propose BrowserGym actions, +and executes those actions through the harness/session layer. + +Example: + python examples/browsergym_harness_eval.py --model gpt-4.1-mini +""" + +from __future__ import annotations + +import argparse +import os +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT / "src")) +sys.path.insert(0, str(_REPO_ROOT / "envs")) + +from openenv.core.harness import HarnessRunLimits + +from browsergym_harness_eval_common import ( + DEFAULT_BENCHMARK, + DEFAULT_BROWSERGYM_IMAGE, + DEFAULT_MAX_STEPS, + DEFAULT_MODEL, + DEFAULT_TASK_NAME, + build_browsergym_session_factory, + build_openai_model_step, + create_openai_client, + format_episode_summary, + run_white_box_episode, + start_browsergym_runtime, + summarize_episodes, +) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Evaluate a BrowserGym task through the harness runtime.", + ) + parser.add_argument( + "--base-url", + default=None, + help="Use an already-running BrowserGym server instead of starting Docker.", + ) + parser.add_argument( + "--image", + default=DEFAULT_BROWSERGYM_IMAGE, + help="Docker image to start when --base-url is not provided.", + ) + parser.add_argument( + "--benchmark", + default=DEFAULT_BENCHMARK, + help="BrowserGym benchmark to run when starting Docker.", + ) + parser.add_argument( + "--task-name", + default=DEFAULT_TASK_NAME, + help="BrowserGym task name to evaluate.", + ) + parser.add_argument( + "--model", + default=DEFAULT_MODEL, + help="OpenAI-compatible chat model name.", + ) + parser.add_argument( + "--api-base-url", + default=os.getenv("OPENAI_BASE_URL"), + help="Optional OpenAI-compatible API base URL.", + ) + parser.add_argument( + "--episodes", + type=int, + default=1, + help="Number of episodes to evaluate.", + ) + parser.add_argument( + "--max-steps", + type=int, + default=DEFAULT_MAX_STEPS, + help="Maximum tool/model turns per episode.", + ) + parser.add_argument( + "--temperature", + type=float, + default=0.2, + help="Sampling temperature for the chat model.", + ) + parser.add_argument( + "--max-tokens", + type=int, + default=200, + help="Maximum completion tokens for each model step.", + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + client = create_openai_client(api_base_url=args.api_base_url) + + with start_browsergym_runtime( + base_url=args.base_url, + image=args.image, + benchmark=args.benchmark, + task_name=args.task_name, + ) as runtime: + print(f"BrowserGym server: {runtime.base_url}") + session_factory = build_browsergym_session_factory( + base_url=runtime.base_url, + task_name=args.task_name, + ) + model_step = build_openai_model_step( + client=client, + model=args.model, + temperature=args.temperature, + max_tokens=args.max_tokens, + ) + + episodes = [] + limits = HarnessRunLimits(max_turns=args.max_steps) + for index in range(1, args.episodes + 1): + episode = run_white_box_episode( + session_factory=session_factory, + model_step=model_step, + limits=limits, + episode_id=f"browsergym-eval-{index}", + ) + episodes.append(episode) + print(format_episode_summary(index, episode)) + + summary = summarize_episodes(episodes) + print( + "Aggregate: " + f"episodes={int(summary['episodes'])} " + f"avg_reward={summary['avg_reward']:.2f} " + f"success_rate={summary['success_rate']:.2%} " + f"avg_steps={summary['avg_steps']:.2f}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/browsergym_harness_eval_common.py b/examples/browsergym_harness_eval_common.py new file mode 100644 index 0000000000..cc4dded6e9 --- /dev/null +++ b/examples/browsergym_harness_eval_common.py @@ -0,0 +1,539 @@ +"""Shared helpers for BrowserGym harness inference and evaluation examples.""" + +from __future__ import annotations + +import json +import os +import re +import threading +from contextlib import AbstractContextManager +from dataclasses import dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from typing import Any + +from browsergym_env import BrowserGymEnv +from browsergym_env.harness import ( + BrowserGymSessionFactory, + build_browsergym_action_tool_call, +) +from openai import OpenAI +from openenv.core.containers.runtime.providers import LocalDockerProvider +from openenv.core.env_server.mcp_types import JsonRpcErrorCode, JsonRpcResponse +from openenv.core.harness import ( + _resolve_env_reward, + CLIHarnessAdapter, + HarnessRolloutResult, + HarnessRunLimits, + MCPHarnessAdapter, + ModelStepResult, + SessionMCPBridge, + VerifyResult, +) +from openenv.core.llm_client import LLMResponse + +DEFAULT_BROWSERGYM_IMAGE = "browsergym-env:latest" +DEFAULT_BENCHMARK = "miniwob" +DEFAULT_TASK_NAME = "click-test" +DEFAULT_MAX_STEPS = 8 +DEFAULT_MODEL = os.getenv("OPENAI_MODEL", "gpt-4.1-mini") +DEFAULT_TEMPERATURE = 0.2 +DEFAULT_MAX_TOKENS = 200 +FALLBACK_ACTION = "noop()" + +SYSTEM_PROMPT = """You control a web browser through BrowserGym. + +Reply with exactly one BrowserGym action string, such as: +- click('13') +- fill('42', 'hello world') +- send_keys('Enter') +- scroll('down') +- noop() + +Use only one action. Do not explain your reasoning. +If you are unsure, reply with noop(). +""" + +ACTION_PREFIX_RE = re.compile(r"^(action|next action)\s*[:-]\s*", re.IGNORECASE) +ACTION_PATTERN = re.compile( + r"(click|fill|send_keys|scroll|noop)\s*\(.*?\)", + re.IGNORECASE | re.DOTALL, +) + + +def _serialize_json(value: Any) -> str: + if value is None: + return "" + if isinstance(value, str): + return value + return json.dumps(value, sort_keys=True, default=str) + + +def extract_browsergym_action( + response_text: str, + fallback_action: str = FALLBACK_ACTION, +) -> str: + """Extract a BrowserGym action from free-form model output.""" + + if not response_text: + return fallback_action + + for raw_line in response_text.splitlines(): + line = raw_line.strip().strip("`") + if not line: + continue + line = ACTION_PREFIX_RE.sub("", line) + match = ACTION_PATTERN.search(line) + if match: + action = re.sub(r"\s+", " ", match.group(0).strip()) + open_paren = action.find("(") + if open_paren > 0: + action = f"{action[:open_paren].lower()}{action[open_paren:]}" + return action.rstrip(";") + + match = ACTION_PATTERN.search(response_text) + if match: + action = re.sub(r"\s+", " ", match.group(0).strip()) + open_paren = action.find("(") + if open_paren > 0: + action = f"{action[:open_paren].lower()}{action[open_paren:]}" + return action.rstrip(";") + + return fallback_action + + +def build_browsergym_tool_call_from_output(response_text: str): + """Convert free-form model output into a BrowserGym tool call.""" + + action = extract_browsergym_action(response_text) + try: + return build_browsergym_action_tool_call(action) + except ValueError: + return build_browsergym_action_tool_call(FALLBACK_ACTION) + + +def build_browsergym_chat_messages( + messages: list[dict[str, Any]], + *, + system_prompt: str = SYSTEM_PROMPT, +) -> list[dict[str, Any]]: + """Convert harness messages into a plain chat-completions transcript.""" + + chat_messages: list[dict[str, Any]] = [ + {"role": "system", "content": system_prompt.strip()} + ] + + for message in messages: + role = message.get("role") + content = message.get("content") or "" + + if role == "tool": + tool_name = message.get("name", "tool") + chat_messages.append( + { + "role": "user", + "content": f"Observation from {tool_name}:\n{content}", + } + ) + continue + + if role not in {"user", "assistant", "system"}: + role = "user" + chat_messages.append({"role": role, "content": str(content)}) + + return chat_messages + + +def build_openai_model_step( + client: OpenAI, + model: str, + *, + system_prompt: str = SYSTEM_PROMPT, + temperature: float = DEFAULT_TEMPERATURE, + max_tokens: int = DEFAULT_MAX_TOKENS, +): + """Build a white-box model step for BrowserGym inference/evaluation.""" + + def model_step(messages, tools, sampling): + del tools + completion = client.chat.completions.create( + model=model, + messages=build_browsergym_chat_messages( + messages, + system_prompt=system_prompt, + ), + temperature=sampling.get("temperature", temperature), + max_tokens=sampling.get("max_tokens", max_tokens), + stream=False, + ) + response_text = completion.choices[0].message.content or "" + tool_call = build_browsergym_tool_call_from_output(response_text) + return ModelStepResult( + response=LLMResponse(content=response_text, tool_calls=[tool_call]), + ) + + return model_step + + +def create_openai_client( + *, + api_base_url: str | None = None, + api_key: str | None = None, +) -> OpenAI: + """Create a sync OpenAI-compatible client for the examples.""" + + client_kwargs: dict[str, Any] = {} + if api_base_url: + client_kwargs["base_url"] = api_base_url + + resolved_key = api_key or os.getenv("OPENAI_API_KEY") or os.getenv("API_KEY") + if resolved_key: + client_kwargs["api_key"] = resolved_key + + return OpenAI(**client_kwargs) + + +@dataclass +class BrowserGymRuntime(AbstractContextManager["BrowserGymRuntime"]): + """Runtime information for a BrowserGym server used by the examples.""" + + base_url: str + provider: LocalDockerProvider | None = None + + def close(self) -> None: + if self.provider is not None: + self.provider.stop_container() + self.provider = None + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self.close() + + +def start_browsergym_runtime( + *, + base_url: str | None = None, + image: str = DEFAULT_BROWSERGYM_IMAGE, + benchmark: str = DEFAULT_BENCHMARK, + task_name: str = DEFAULT_TASK_NAME, +) -> BrowserGymRuntime: + """Start or attach to a BrowserGym server for evaluation examples.""" + + if base_url is not None: + return BrowserGymRuntime(base_url=base_url) + + provider = LocalDockerProvider() + env_vars = { + "BROWSERGYM_BENCHMARK": benchmark, + "BROWSERGYM_TASK_NAME": task_name, + } + runtime_base_url = provider.start_container(image=image, env_vars=env_vars) + provider.wait_for_ready(runtime_base_url) + return BrowserGymRuntime(base_url=runtime_base_url, provider=provider) + + +def build_browsergym_session_factory( + *, + base_url: str, + task_name: str = DEFAULT_TASK_NAME, +) -> BrowserGymSessionFactory: + """Create a BrowserGym session factory for a running environment server.""" + + return BrowserGymSessionFactory( + client_factory=lambda: BrowserGymEnv(base_url=base_url), + default_task=task_name, + ) + + +def _rollout_final_state(rollout: HarnessRolloutResult) -> dict[str, Any]: + return { + "done": rollout.done, + "metrics": dict(rollout.metrics), + "events": [ + { + "type": event.type, + "payload": dict(event.payload), + } + for event in rollout.events + ], + "tool_trace": [ + { + "tool_name": entry.tool_name, + "arguments": dict(entry.arguments), + "result": { + "data": entry.result.data, + "done": entry.result.done, + "metadata": dict(entry.result.metadata), + "error": entry.result.error, + }, + } + for entry in rollout.tool_trace + ], + } + + +@dataclass +class EpisodeEvaluation: + """Normalized result from one inference/evaluation episode.""" + + rollout: HarnessRolloutResult + verify: VerifyResult + reward: float + done: bool + success: bool + step_count: int + + +def _finalize_episode( + session, + rollout: HarnessRolloutResult, +) -> EpisodeEvaluation: + verify = session.verify( + transcript=rollout.messages, + final_state=_rollout_final_state(rollout), + ) + reward = _resolve_env_reward(rollout, verify) + done = bool(verify.done or rollout.done) + step_count = int(verify.metrics.get("step_count", len(rollout.tool_trace))) + success = reward > 0.0 + return EpisodeEvaluation( + rollout=rollout, + verify=verify, + reward=reward, + done=done, + success=success, + step_count=step_count, + ) + + +def run_white_box_episode( + *, + session_factory: BrowserGymSessionFactory, + model_step, + limits: HarnessRunLimits, + task: Any = None, + seed: int | None = None, + episode_id: str | None = None, +) -> EpisodeEvaluation: + """Run one BrowserGym episode through the white-box harness path.""" + + session = session_factory.create(task=task, seed=seed, episode_id=episode_id) + try: + rollout = MCPHarnessAdapter().run_white_box( + model_step=model_step, + session=session, + limits=limits, + ) + return _finalize_episode(session, rollout) + finally: + session.close() + + +def run_black_box_episode( + *, + session_factory: BrowserGymSessionFactory, + harness_adapter: CLIHarnessAdapter, + limits: HarnessRunLimits, + task: Any = None, + seed: int | None = None, + episode_id: str | None = None, +) -> EpisodeEvaluation: + """Run one BrowserGym episode through the black-box harness path.""" + + session = session_factory.create(task=task, seed=seed, episode_id=episode_id) + try: + rollout = harness_adapter.run_black_box( + session=session, + limits=limits, + ) + return _finalize_episode(session, rollout) + finally: + session.close() + + +def summarize_episodes(episodes: list[EpisodeEvaluation]) -> dict[str, float]: + """Aggregate reward and completion metrics across episodes.""" + + if not episodes: + return { + "episodes": 0.0, + "avg_reward": 0.0, + "success_rate": 0.0, + "avg_steps": 0.0, + } + + episode_count = len(episodes) + total_reward = sum(episode.reward for episode in episodes) + total_success = sum(1 for episode in episodes if episode.success) + total_steps = sum(episode.step_count for episode in episodes) + return { + "episodes": float(episode_count), + "avg_reward": total_reward / episode_count, + "success_rate": total_success / episode_count, + "avg_steps": total_steps / episode_count, + } + + +def format_episode_summary(index: int, episode: EpisodeEvaluation) -> str: + """Render one episode summary line for CLI output.""" + + return ( + f"Episode {index}: reward={episode.reward:.2f} " + f"success={episode.success} done={episode.done} steps={episode.step_count}" + ) + + +def _normalize_mcp_call_result(result: dict[str, Any]) -> dict[str, Any]: + text_payload = result.get("error") or _serialize_json(result.get("data")) + return { + "content": [{"type": "text", "text": text_payload}], + "structuredContent": { + "data": result.get("data"), + "done": result.get("done", False), + "metadata": result.get("metadata", {}), + "error": result.get("error"), + }, + "data": result.get("data"), + "isError": bool(result.get("error")), + } + + +class SessionMCPHttpServer(AbstractContextManager["SessionMCPHttpServer"]): + """Tiny HTTP wrapper that exposes SessionMCPBridge on /mcp.""" + + def __init__( + self, + bridge: SessionMCPBridge, + *, + host: str = "127.0.0.1", + port: int = 0, + ): + self._bridge = bridge + self._host = host + self._port = port + self._server: ThreadingHTTPServer | None = None + self._thread: threading.Thread | None = None + + @property + def url(self) -> str: + if self._server is None: + raise RuntimeError("HTTP bridge is not running") + host, port = self._server.server_address[:2] + return f"http://{host}:{port}/mcp" + + def start(self) -> "SessionMCPHttpServer": + if self._server is not None: + return self + + bridge = self._bridge + + class Handler(BaseHTTPRequestHandler): + def _send_json(self, payload: dict[str, Any], status: int = 200) -> None: + body = json.dumps(payload).encode("utf-8") + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: Any) -> None: + del format, args + + def do_GET(self) -> None: # noqa: N802 + if self.path != "/mcp": + self.send_error(404) + return + self._send_json({"status": "ok", "transport": "jsonrpc"}) + + def do_POST(self) -> None: # noqa: N802 + if self.path != "/mcp": + self.send_error(404) + return + + content_length = int(self.headers.get("Content-Length", "0")) + raw_body = self.rfile.read(content_length) + try: + request = json.loads(raw_body or b"{}") + except json.JSONDecodeError: + self._send_json( + JsonRpcResponse.error_response( + JsonRpcErrorCode.PARSE_ERROR, + message="Parse error", + ).model_dump() + ) + return + + try: + response = bridge.handle_request(request) + except KeyError as exc: + response = JsonRpcResponse.error_response( + JsonRpcErrorCode.METHOD_NOT_FOUND, + message=str(exc.args[0]) if exc.args else "Method not found", + request_id=request.get("id"), + ).model_dump() + except ValueError as exc: + response = JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_PARAMS, + message=str(exc), + request_id=request.get("id"), + ).model_dump() + except Exception as exc: + response = JsonRpcResponse.error_response( + JsonRpcErrorCode.INTERNAL_ERROR, + message=str(exc), + request_id=request.get("id"), + ).model_dump() + + if ( + request.get("method") == "tools/call" + and "result" in response + and isinstance(response["result"], dict) + ): + response["result"] = _normalize_mcp_call_result(response["result"]) + + self._send_json(response) + + self._server = ThreadingHTTPServer((self._host, self._port), Handler) + self._thread = threading.Thread( + target=self._server.serve_forever, + name="browsergym-session-mcp-http", + daemon=True, + ) + self._thread.start() + return self + + def close(self) -> None: + if self._server is None: + return + self._server.shutdown() + self._server.server_close() + self._server = None + if self._thread is not None: + self._thread.join(timeout=5) + self._thread = None + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + self.close() + + +__all__ = [ + "BrowserGymRuntime", + "DEFAULT_BENCHMARK", + "DEFAULT_BROWSERGYM_IMAGE", + "DEFAULT_MAX_STEPS", + "DEFAULT_MODEL", + "DEFAULT_TASK_NAME", + "EpisodeEvaluation", + "FALLBACK_ACTION", + "SYSTEM_PROMPT", + "SessionMCPHttpServer", + "build_browsergym_chat_messages", + "build_browsergym_session_factory", + "build_browsergym_tool_call_from_output", + "build_openai_model_step", + "create_openai_client", + "extract_browsergym_action", + "format_episode_summary", + "run_white_box_episode", + "run_black_box_episode", + "start_browsergym_runtime", + "summarize_episodes", +] diff --git a/examples/sft_warmup.ipynb b/examples/sft_warmup.ipynb new file mode 100644 index 0000000000..c0bd6cd4ed --- /dev/null +++ b/examples/sft_warmup.ipynb @@ -0,0 +1,180 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": "# Collecting rollouts with OpenEnv for supervised training\n\n[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/meta-pytorch/OpenEnv/blob/main/examples/sft_warmup.ipynb)\n\nOpenEnv environments are not only useful for RL training — they are also a natural tool for **collecting rollouts that become supervised training data**. The environment handles episode management, automatic scoring, and reproducibility, so you get a reward-labeled dataset without writing any of that infrastructure yourself.\n\n**What you'll build:**\n- ~300 rollouts collected from `gpt-5-mini` running inside `reasoning_gym_env`, scored automatically by the environment\n- A filtered, inspected SFT dataset with token-length analysis\n- A fine-tuned `Qwen/Qwen3-1.7B` checkpoint trained on the collected rollouts\n- A before/after eval table showing the improvement from training on environment-collected data\n\n**Next step:** use the checkpoint as a warm-start for GRPO in the [End-to-end walkthrough](https://meta-pytorch.org/OpenEnv/tutorials/end-to-end-walkthrough.html)." + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 1. Install dependencies" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "!pip install -q openai trl openenv-core\n!pip install -q --no-deps git+https://huggingface.co/spaces/sergiopaniego/reasoning_gym\n!pip install -Uq \"transformers>=5.3.0\"" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 2. Set your credentials" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "import getpass, os\n\nif \"OPENAI_API_KEY\" not in os.environ:\n os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"OpenAI API key: \")" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "from huggingface_hub import notebook_login\n\nnotebook_login()\n\nYOUR_HF_USERNAME = \"your-username\" # replace with your Hugging Face username\nassert YOUR_HF_USERNAME != \"your-username\", \"Replace YOUR_HF_USERNAME with your Hugging Face username\"" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 3. Define the system prompt\n\nUse the same prompt as the GRPO tutorial so the SFT checkpoint is a drop-in when you continue with GRPO." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "SYSTEM_PROMPT = \"\"\"You are a careful arithmetic assistant.\n\nYou will be given a chain of integer additions. Compute the result and submit it as a single number.\n\nRules:\n1. Read the question carefully.\n2. Use the tool `answer` exactly once with your final number.\n3. The answer must be a single integer with no units or explanation.\n\"\"\"" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 4. Configure data collection\n\n`DATASET_CONFIG` controls the difficulty of the `chain_sum` problems the environment generates:\n`min_terms`/`max_terms` set how many integers are added together, and `min_digits`/`max_digits` set\nhow many digits each integer has. At these settings each problem is a sum of 2–3 two-digit numbers —\neasy enough for `gpt-5-mini` to answer correctly ~90% of the time, which gives a clean training signal after filtering.\n\n`N_EPISODES` is the number of problems to collect. 300 is enough to get ~270 correct examples after filtering, which is sufficient for format compliance training." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "DATASET_CONFIG = {\n \"min_terms\": 2,\n \"max_terms\": 3,\n \"min_digits\": 2,\n \"max_digits\": 2,\n}\n\nN_EPISODES = 300" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 5. Collect rollouts\n\nThe `openenv collect` pipeline runs the teacher model inside the environment and records every episode.\nThe environment's `step()` reward is written alongside the messages, so filtering by correctness requires\nno additional scoring code.\n\nThe same pipeline is available as a Python API — the CLI wraps these exact calls:\n\n> **CLI equivalent:** `openenv collect reasoning_gym:chain_sum --provider openai --model gpt-5-mini --num-episodes 300 --push-to-hub /chain-sum-rollouts ...`" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "import os\nfrom openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter\nfrom openenv.core.harness.collect import build_model_step, CollectRunner, RolloutSerializer\nfrom openenv.core.llm_client import create_llm_client\nfrom reasoning_gym_env.client import ReasoningGymEnv\nfrom reasoning_gym_env.harness import ReasoningGymSessionFactory\n\nclient = create_llm_client(\n provider=\"openai\",\n model=\"gpt-5-mini\",\n api_key=os.environ[\"OPENAI_API_KEY\"],\n max_tokens=1024,\n)\nmodel_step = build_model_step(client, system_prompt=SYSTEM_PROMPT)\n\nfactory = ReasoningGymSessionFactory(\n lambda: ReasoningGymEnv(base_url=\"https://sergiopaniego-reasoning-gym.hf.space\"),\n dataset_name=\"chain_sum\",\n dataset_config=DATASET_CONFIG,\n)\n\nserializer = RolloutSerializer(\"./rollouts\")\nrunner = CollectRunner(\n session_factory=factory,\n harness_adapter=MCPHarnessAdapter(),\n serializer=serializer,\n limits=HarnessRunLimits(max_turns=9),\n)\n\nresult = runner.run(model_step=model_step, num_episodes=N_EPISODES)\nprint(f\"Collected={result.num_collected} dropped={result.num_dropped} avg_reward={result.avg_reward:.3f} success_rate={result.success_rate:.0%}\")" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "from openenv.core.harness.collect import push_to_hf_hub\n\nurl = push_to_hf_hub(\n output_dir=\"./rollouts\",\n repo_id=f\"{YOUR_HF_USERNAME}/chain-sum-rollouts\",\n)\nprint(f\"Dataset at: {url}\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "The collected `messages` use the standard OpenAI `tool_calls` format. Convert to Qwen3's\n`` text format — GRPOTrainer produces this same format during RL, making the\nSFT checkpoint a direct drop-in." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "from datasets import load_dataset\nimport json\n\nds = load_dataset(f\"{YOUR_HF_USERNAME}/chain-sum-rollouts\", split=\"train\")\nraw_rollouts = list(ds)\nprint(f\"Collected {len(raw_rollouts)} episodes\")\n\ndef to_qwen3_messages(record):\n converted = []\n for msg in record[\"messages\"]:\n if msg[\"role\"] == \"tool\":\n continue # strip env responses; SFT only needs the assistant turn\n if msg[\"role\"] == \"assistant\" and msg.get(\"tool_calls\"):\n tc = msg[\"tool_calls\"][0]\n args = json.loads(tc[\"function\"][\"arguments\"])\n answer_str = args.get(\"answer\", \"\")\n tool_call_text = (\n \"\\n\"\n + json.dumps({\"name\": \"answer\", \"arguments\": {\"answer\": answer_str}})\n + \"\\n\"\n )\n converted.append({\"role\": \"assistant\", \"content\": tool_call_text})\n else:\n converted.append(msg)\n return {\"messages\": converted, \"reward\": record[\"reward\"]}\n\nrollouts = [to_qwen3_messages(r) for r in raw_rollouts]" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 6. Filter the dataset\n\nKeep only episodes where the teacher answered correctly. The environment's reward signal does the\nlabelling — no manual annotation needed." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "correct = [r for r in rollouts if r[\"reward\"] == 1.0]\nprint(f\"Correct: {len(correct)} / {len(rollouts)} ({len(correct)/len(rollouts):.1%})\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 7. Inspect the dataset before training\n\nAlways look at your data before training. Automated collection can introduce unexpected patterns that the\nstudent model will learn to imitate." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "import random\n\nfor row in random.sample(correct, 3):\n question = row[\"messages\"][0][\"content\"]\n response = row[\"messages\"][1][\"content\"]\n print(f\"Q: {question}\")\n print(f\"A: {response}\")\n print()" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 8. Measure token lengths\n\nSet `max_seq_length` to the 99th percentile so fewer than 1% of examples are truncated." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "import numpy as np\nfrom transformers import AutoTokenizer\n\ntokenizer = AutoTokenizer.from_pretrained(\"Qwen/Qwen3-1.7B\")\n\nlengths = []\nfor row in correct:\n text = tokenizer.apply_chat_template(\n row[\"messages\"], tokenize=False, add_generation_prompt=False\n )\n ids = tokenizer.encode(text)\n lengths.append(len(ids))\n\nlengths = np.array(lengths)\nMAX_SEQ_LEN = int(np.percentile(lengths, 99)) + 16\n\nprint(\n f\"p50={np.percentile(lengths, 50):.0f} \"\n f\"p95={np.percentile(lengths, 95):.0f} \"\n f\"p99={np.percentile(lengths, 99):.0f} \"\n f\"max={lengths.max()}\"\n)\nprint(f\"Setting MAX_SEQ_LEN = {MAX_SEQ_LEN}\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 9. Fine-tune with SFTTrainer\n\n`assistant_only_loss=True` in `SFTConfig` masks the prompt tokens so the loss is computed only on the assistant's `` response — more efficient than full-sequence training and avoids accidentally reinforcing the system prompt wording." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "from datasets import Dataset\nfrom transformers import AutoModelForCausalLM\nfrom trl import SFTConfig, SFTTrainer\n\ndataset = Dataset.from_list([{\"messages\": r[\"messages\"]} for r in correct])\n\nmodel = AutoModelForCausalLM.from_pretrained(\"Qwen/Qwen3-1.7B\")\n\nsft_config = SFTConfig(\n output_dir=\"reasoning-gym-chain-sum-Qwen3-1.7B-sft\",\n max_length=MAX_SEQ_LEN,\n num_train_epochs=3,\n per_device_train_batch_size=4,\n gradient_accumulation_steps=2,\n learning_rate=2e-5,\n warmup_steps=10,\n lr_scheduler_type=\"cosine\",\n logging_steps=5,\n save_strategy=\"no\",\n assistant_only_loss=True,\n)\n\ntrainer = SFTTrainer(\n model=model,\n train_dataset=dataset,\n processing_class=tokenizer,\n args=sft_config,\n)\n\ntrainer.train()\ntrainer.push_to_hub(commit_message=\"SFT warm-up on reasoning_gym chain_sum\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 10. Evaluate: before vs after\n\nRun both the base model and the SFT checkpoint on a held-out set with a different seed.\nThe key metric is **format compliance** — how often the model uses `` correctly." + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "import re\nfrom transformers import pipeline\nfrom reasoning_gym_env.client import ReasoningGymEnv\nfrom reasoning_gym_env.models import ReasoningGymAction\n\nasync def evaluate_model(model_name, n_eval=50, seed=999):\n gen = pipeline(\n \"text-generation\",\n model=model_name,\n tokenizer=model_name,\n device_map=\"auto\",\n dtype=\"auto\",\n )\n gen.model.generation_config.max_length = None\n tok = AutoTokenizer.from_pretrained(model_name)\n env = ReasoningGymEnv(base_url=\"https://sergiopaniego-reasoning-gym.hf.space\")\n\n obs = await env.reset(\n dataset_name=\"chain_sum\",\n dataset_config=DATASET_CONFIG,\n seed=seed,\n size=n_eval,\n )\n\n rewards, format_hits = [], 0\n\n for i in range(n_eval):\n if i > 0:\n obs = await env.reset()\n\n question = obs.observation.question\n messages = [\n {\"role\": \"system\", \"content\": SYSTEM_PROMPT},\n {\"role\": \"user\", \"content\": question},\n ]\n prompt = tok.apply_chat_template(\n messages, tokenize=False, add_generation_prompt=True\n )\n completion = gen(prompt, max_new_tokens=64)[0][\"generated_text\"][len(prompt):]\n\n m = re.search(r'\"answer\"\\s*:\\s*\"?(\\d+)\"?', completion)\n if m:\n format_hits += 1\n answer = m.group(1)\n else:\n nums = re.findall(r\"\\b(\\d+)\\b\", completion)\n answer = nums[-1] if nums else \"0\"\n\n result = await env.step(ReasoningGymAction(answer=answer))\n rewards.append(float(result.observation.score or 0.0))\n\n await env.close()\n del gen\n\n return {\n \"accuracy\": sum(rewards) / len(rewards),\n \"format_compliance\": format_hits / n_eval,\n }" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "base_metrics = await evaluate_model(\"Qwen/Qwen3-1.7B\")\nsft_metrics = await evaluate_model(f\"{YOUR_HF_USERNAME}/reasoning-gym-chain-sum-Qwen3-1.7B-sft\")\n\nprint(f\"\\n{'Metric':<25} {'Base model':>12} {'After SFT':>12} {'Delta':>10}\")\nprint(\"-\" * 62)\nfor key, label in [(\"format_compliance\", \"Format compliance\"), (\"accuracy\", \"Accuracy\")]:\n b, s = base_metrics[key], sft_metrics[key]\n print(f\"{label:<25} {b:>12.1%} {s:>12.1%} {(s - b) * 100:>+9.1f} pp\")" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "---\n\n## 11. Where to go next: GRPO\n\nThe SFT checkpoint is ready to use as the warm-start for GRPO. In the\n[end-to-end walkthrough](https://meta-pytorch.org/OpenEnv/tutorials/end-to-end-walkthrough.html),\nchange one line in the trainer setup:\n\n```python\n# Before (cold-start from the base model):\nMODEL_NAME = \"Qwen/Qwen3-1.7B\"\n\n# After (warm-start from your SFT checkpoint):\nMODEL_NAME = f\"{YOUR_HF_USERNAME}/reasoning-gym-chain-sum-Qwen3-1.7B-sft\"\n```\n\nWith format compliance already near 100%, GRPO's `reward_std` will be non-zero from the first batch\nand the reward curve will climb immediately." + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.11.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/examples/ttt_collect_demo.py b/examples/ttt_collect_demo.py new file mode 100644 index 0000000000..c879fd26c3 --- /dev/null +++ b/examples/ttt_collect_demo.py @@ -0,0 +1,225 @@ +"""Smoke-test demo for the harness collect pipeline on Tic-Tac-Toe. + +Two run modes: + +1. Scripted mode (default, no setup): a scripted OpenSpiel client and a + deterministic teacher exercise the full pipeline (session → harness → + verify → serialize) and leave a real ``results.jsonl`` + + ``metadata.json`` on disk so you can inspect the schema downstream + tools will consume. + +2. Remote mode (``--base-url``): talks to a real OpenSpiel server — + either a local Docker container or a deployed Hugging Face Space. + Requires ``OPENSPIEL_GAME=tic_tac_toe`` set on the server. + +Usage: + # Scripted (no Docker, no API): + PYTHONPATH=src:envs uv run python examples/ttt_collect_demo.py \\ + --num-episodes 10 --output-dir /tmp/ttt-collect-demo + + # Against a Hugging Face Space: + PYTHONPATH=src:envs uv run python examples/ttt_collect_demo.py \\ + --base-url https://-.hf.space \\ + --num-episodes 10 --output-dir /tmp/ttt-collect-hf +""" + +from __future__ import annotations + +import argparse +import json +import sys +from pathlib import Path +from typing import Any + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT / "src")) +sys.path.insert(0, str(_REPO_ROOT / "envs")) + +from openenv.core.client_types import StepResult +from openenv.core.harness import ( + HarnessRunLimits, + MCPHarnessAdapter, + ModelStepResult, +) +from openenv.core.harness.collect import CollectRunner, RolloutSerializer +from openenv.core.llm_client import LLMResponse, ToolCall + +from openspiel_env.client import OpenSpielEnv +from openspiel_env.harness import OpenSpielSessionFactory +from openspiel_env.models import OpenSpielAction, OpenSpielObservation, OpenSpielState + + +class ScriptedTTTClient: + """A tic_tac_toe client that ends after `turns_to_win` plays.""" + + def __init__(self, turns_to_win: int = 1, winning: bool = True): + self._turns_to_win = turns_to_win + self._winning = winning + self._played = 0 + + def reset(self, **kwargs: Any) -> StepResult[OpenSpielObservation]: + self._played = 0 + return StepResult( + observation=OpenSpielObservation( + info_state=[1.0] * 9 + [0.0] * 18, + legal_actions=list(range(9)), + current_player_id=0, + game_phase="playing", + ), + reward=0.0, + done=False, + ) + + def step(self, action: OpenSpielAction) -> StepResult[OpenSpielObservation]: + self._played += 1 + done = self._played >= self._turns_to_win + reward = 1.0 if (done and self._winning) else 0.0 + info_state = [0.0] * 27 + info_state[9] = 1.0 + for i in range(1, 9): + info_state[i] = 1.0 + return StepResult( + observation=OpenSpielObservation( + info_state=info_state, + legal_actions=[] if done else list(range(1, 9)), + current_player_id=0, + game_phase="terminal" if done else "playing", + ), + reward=reward, + done=done, + ) + + def state(self) -> OpenSpielState: + return OpenSpielState(step_count=self._played, game_name="tic_tac_toe") + + def close(self) -> None: + return None + + +def _extract_legal_actions(messages: list[dict[str, Any]]) -> list[int]: + """Scan messages for the most recent ``legal_actions`` list. + + The prompt and tool results both surface ``legal_actions`` — we look + at both so the teacher stays legal across turns. + """ + import json as _json + import re as _re + + for message in reversed(messages): + content = message.get("content") or "" + if not isinstance(content, str): + continue + + # Tool results serialize as JSON strings containing legal_actions. + try: + payload = _json.loads(content) + except (_json.JSONDecodeError, ValueError): + payload = None + if isinstance(payload, dict) and "legal_actions" in payload: + legal = payload["legal_actions"] + if isinstance(legal, list): + return [int(a) for a in legal] + + # The initial prompt uses the textual form ``Legal actions: [0, 1, ...]``. + match = _re.search(r"Legal actions:\s*\[([^\]]*)\]", content) + if match: + raw = match.group(1).strip() + if not raw: + return [] + return [int(x) for x in raw.split(",")] + + return [] + + +def scripted_teacher(messages, tools, sampling): + """Pick the first legal action seen in the conversation so far.""" + del tools, sampling + legal = _extract_legal_actions(messages) + action_id = legal[0] if legal else 0 + return ModelStepResult( + response=LLMResponse( + content=f"Playing cell {action_id} (first legal).", + tool_calls=[ + ToolCall( + id=f"t-{action_id}", + name="play_move", + args={"action_id": action_id}, + ) + ], + ), + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--num-episodes", type=int, default=5) + parser.add_argument("--output-dir", type=Path, default=Path("/tmp/ttt-collect-demo")) + parser.add_argument("--max-turns", type=int, default=9) + parser.add_argument( + "--base-url", + default=None, + help=( + "OpenSpiel server URL (local Docker or Hugging Face Space). " + "When omitted, runs against the built-in scripted client." + ), + ) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + if args.base_url: + print(f"Using remote OpenSpiel server: {args.base_url}") + client_factory = lambda: OpenSpielEnv(base_url=args.base_url) # noqa: E731 + teacher_name = "scripted-legal-action" + else: + print("Using built-in scripted client (no Docker required)") + client_factory = lambda: ScriptedTTTClient( # noqa: E731 + turns_to_win=1, winning=True + ) + teacher_name = "scripted-demo" + + factory = OpenSpielSessionFactory(client_factory, game_name="tic_tac_toe") + serializer = RolloutSerializer(args.output_dir) + serializer.write_metadata( + { + "env_id": "tic_tac_toe", + "teacher_model": teacher_name, + "num_episodes": args.num_episodes, + "base_url": args.base_url, + } + ) + + runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=serializer, + limits=HarnessRunLimits(max_turns=args.max_turns), + ) + + result = runner.run( + model_step=scripted_teacher, + num_episodes=args.num_episodes, + episode_id_prefix="ttt", + ) + + print( + f"Collected: {result.num_collected} " + f"Skipped: {result.num_skipped} " + f"Dropped: {result.num_dropped}" + ) + print(f"Avg reward: {result.avg_reward:.3f} Success rate: {result.success_rate:.0%}") + print() + print(f"results.jsonl: {serializer.results_path}") + print(f"metadata.json: {serializer.metadata_path}") + print() + + with serializer.results_path.open("r", encoding="utf-8") as handle: + first_line = handle.readline().strip() + print("First episode (pretty):") + print(json.dumps(json.loads(first_line), indent=2)[:1500]) + + +if __name__ == "__main__": + main() diff --git a/examples/ttt_collect_with_llm.py b/examples/ttt_collect_with_llm.py new file mode 100644 index 0000000000..cf24bb1ff7 --- /dev/null +++ b/examples/ttt_collect_with_llm.py @@ -0,0 +1,229 @@ +"""Collect Tic-Tac-Toe rollouts with a real LLM as the teacher. + +Provider-agnostic via ``openenv.core.llm_client``: works with OpenAI, +Anthropic, and any OpenAI-compatible endpoint (vLLM, TGI, Ollama, +HuggingFace Inference, Together, Groq, Fireworks, ...) without any +code change. + +The OpenSpiel server can be local (Docker) or a Hugging Face Space — +pass the URL via ``--base-url``. The LLM endpoint is independent: +``--llm-endpoint`` + ``--llm-port`` for OpenAI-compatible servers, or +``--provider openai|anthropic`` for hosted APIs. + +Examples: + # Hosted OpenAI against a HF Space: + OPENAI_API_KEY=sk-... python examples/ttt_collect_with_llm.py \\ + --base-url https://-.hf.space \\ + --provider openai --model gpt-5-mini \\ + --num-episodes 20 \\ + --output-dir /tmp/ttt-openai + + # Local vLLM (OpenAI-compatible) with Qwen: + python examples/ttt_collect_with_llm.py \\ + --base-url http://localhost:8000 \\ + --llm-endpoint http://localhost --llm-port 8001 \\ + --model Qwen/Qwen2.5-7B-Instruct \\ + --num-episodes 20 + + # Anthropic against a HF Space: + ANTHROPIC_API_KEY=sk-ant-... python examples/ttt_collect_with_llm.py \\ + --base-url https://-.hf.space \\ + --provider anthropic --model claude-sonnet-4-6 \\ + --num-episodes 20 +""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parent.parent +sys.path.insert(0, str(_REPO_ROOT / "src")) +sys.path.insert(0, str(_REPO_ROOT / "envs")) + +from openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter +from openenv.core.harness.collect import ( + CollectRunner, + RolloutSerializer, + build_model_step, +) +from openenv.core.llm_client import LLMClient, OpenAIClient, create_llm_client + +from openspiel_env.client import OpenSpielEnv +from openspiel_env.harness import OpenSpielSessionFactory + +SYSTEM_PROMPT = """You are playing tic-tac-toe against a random opponent. + +For each turn, call the `play_move` tool with an ``action_id`` drawn from +the ``legal_actions`` list shown in the latest observation. Numbers on +the board are empty cells; `X` and `O` are played cells. You play as X. +Choose moves that win or block the opponent. +""" + + +def _env_api_key(provider: str) -> str: + mapping = { + "openai": ["OPENAI_API_KEY", "API_KEY"], + "anthropic": ["ANTHROPIC_API_KEY", "API_KEY"], + } + for name in mapping.get(provider, []): + value = os.getenv(name) + if value: + return value + raise RuntimeError( + f"No API key found for provider={provider!r}. " + f"Set one of: {mapping.get(provider, [])}" + ) + + +def build_llm_client(args: argparse.Namespace) -> LLMClient: + """Pick the right LLMClient subclass based on CLI flags.""" + if args.llm_endpoint: + # OpenAI-compatible self-hosted endpoint (vLLM, TGI, Ollama, etc.). + return OpenAIClient( + endpoint=args.llm_endpoint, + port=args.llm_port, + model=args.model, + api_key=os.getenv("OPENAI_API_KEY") or "not-needed", + system_prompt=SYSTEM_PROMPT, + temperature=args.temperature, + max_tokens=args.max_tokens, + ) + + # Hosted provider. + return create_llm_client( + provider=args.provider, + model=args.model, + api_key=_env_api_key(args.provider), + system_prompt=SYSTEM_PROMPT, + temperature=args.temperature, + max_tokens=args.max_tokens, + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + + parser.add_argument( + "--base-url", + required=True, + help="OpenSpiel server URL (local Docker or HF Space).", + ) + parser.add_argument("--num-episodes", type=int, default=10) + parser.add_argument("--max-turns", type=int, default=9) + parser.add_argument("--output-dir", type=Path, default=Path("/tmp/ttt-collect-llm")) + parser.add_argument( + "--episode-id-prefix", + default="ttt", + help="Prefix for serialized episode ids (stable across resumes).", + ) + + # LLM selection: either a hosted provider or a self-hosted endpoint. + parser.add_argument( + "--provider", + choices=["openai", "anthropic"], + default="openai", + help="Hosted provider (ignored when --llm-endpoint is set).", + ) + parser.add_argument( + "--model", + default="gpt-5-mini", + help="Model name to pass to the LLM API.", + ) + parser.add_argument( + "--llm-endpoint", + default=None, + help="OpenAI-compatible endpoint (e.g. http://localhost for vLLM).", + ) + parser.add_argument("--llm-port", type=int, default=8000) + parser.add_argument("--temperature", type=float, default=0.2) + parser.add_argument("--max-tokens", type=int, default=200) + + parser.add_argument( + "--keep-losses", + action="store_true", + help="Keep losing rollouts. Default keeps only wins + draws (reward >= 0).", + ) + + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + llm_client = build_llm_client(args) + teacher_label = ( + f"{args.llm_endpoint}:{args.llm_port}/{args.model}" + if args.llm_endpoint + else f"{args.provider}/{args.model}" + ) + + factory = OpenSpielSessionFactory( + lambda: OpenSpielEnv(base_url=args.base_url), + game_name="tic_tac_toe", + ) + serializer = RolloutSerializer(args.output_dir) + serializer.write_metadata( + { + "env_id": "tic_tac_toe", + "env_base_url": args.base_url, + "teacher": teacher_label, + "num_episodes_requested": args.num_episodes, + "temperature": args.temperature, + "keep_losses": args.keep_losses, + } + ) + + runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=serializer, + limits=HarnessRunLimits(max_turns=args.max_turns), + ) + + model_step = build_model_step(llm_client, system_prompt=SYSTEM_PROMPT) + + should_keep = None if args.keep_losses else (lambda record: record.reward >= 0.0) + + print(f"Env server: {args.base_url}") + print(f"Teacher: {teacher_label}") + print(f"Output dir: {args.output_dir}") + print() + + result = runner.run( + model_step=model_step, + num_episodes=args.num_episodes, + episode_id_prefix=args.episode_id_prefix, + should_keep=should_keep, + ) + + print( + f"Collected: {result.num_collected} " + f"Skipped: {result.num_skipped} " + f"Dropped: {result.num_dropped}" + ) + print( + f"Avg reward: {result.avg_reward:.3f} " + f"Success rate: {result.success_rate:.0%}" + ) + print() + print(f"results.jsonl: {serializer.results_path}") + print(f"metadata.json: {serializer.metadata_path}") + + if serializer.results_path.exists(): + with serializer.results_path.open("r", encoding="utf-8") as handle: + first_line = handle.readline().strip() + if first_line: + payload = json.loads(first_line) + print() + print( + f"First episode: id={payload['episode_id']} " + f"reward={payload['reward']} turns={payload['metrics'].get('turns')}" + ) + + +if __name__ == "__main__": + main() diff --git a/src/openenv/cli/__main__.py b/src/openenv/cli/__main__.py index 6b457cb7e1..b80e5b9fd9 100644 --- a/src/openenv/cli/__main__.py +++ b/src/openenv/cli/__main__.py @@ -14,7 +14,16 @@ import sys import typer -from openenv.cli.commands import build, fork, init, push, serve, skills, validate +from openenv.cli.commands import ( + build, + collect, + fork, + init, + push, + serve, + skills, + validate, +) # Create the main CLI app app = typer.Typer( @@ -47,6 +56,10 @@ name="skills", help="Manage OpenEnv skills for AI assistants", ) +app.command( + name="collect", + help="Collect rollouts from a deployed OpenEnv environment", +)(collect.collect) # Entry point for setuptools diff --git a/src/openenv/cli/commands/collect.py b/src/openenv/cli/commands/collect.py new file mode 100644 index 0000000000..aedcb9f420 --- /dev/null +++ b/src/openenv/cli/commands/collect.py @@ -0,0 +1,443 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Run a rollout collection job against a deployed OpenEnv environment. + +Currently ships with a built-in registry for ``openspiel_env``. Adding +other envs is a matter of writing a thin session factory in the env +package and registering it here. +""" + +from __future__ import annotations + +import json +import os +import re +import sys +from pathlib import Path +from typing import Annotated, Any + +import typer +from openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter, ModelStepResult +from openenv.core.harness.collect import ( + build_model_step, + CollectRunner, + push_to_hf_hub, + RolloutSerializer, +) +from openenv.core.llm_client import create_llm_client, LLMClient, LLMResponse, ToolCall + +from .._cli_utils import console + +# Imported eagerly so tests can monkeypatch these names on this module. +# The real imports only happen when the selected --env is openspiel/reasoning_gym. +try: + from openspiel_env.client import OpenSpielEnv # type: ignore[import-not-found] + from openspiel_env.harness import ( # type: ignore[import-not-found] + OpenSpielSessionFactory, + ) +except ImportError: # pragma: no cover - openspiel env optional at import time + OpenSpielEnv = None # type: ignore[assignment] + OpenSpielSessionFactory = None # type: ignore[assignment] + +try: + from reasoning_gym_env.client import ( # type: ignore[import-not-found] + ReasoningGymEnv, + ) + from reasoning_gym_env.harness import ( # type: ignore[import-not-found] + ReasoningGymSessionFactory, + ) +except ImportError: # pragma: no cover - reasoning_gym env optional at import time + ReasoningGymEnv = None # type: ignore[assignment] + ReasoningGymSessionFactory = None # type: ignore[assignment] + +app = typer.Typer(help="Collect rollouts from a deployed OpenEnv environment.") + +_PROVIDER_API_KEY_ENVS: dict[str, list[str]] = { + "openai": ["OPENAI_API_KEY", "API_KEY"], + "anthropic": ["ANTHROPIC_API_KEY", "API_KEY"], +} + +_SYSTEM_PROMPT = """You are playing a game through a tool-calling interface. +For each turn, inspect the observation and call the provided tool with a +valid argument. Only call one tool per turn.""" + + +def _uses_llm_teacher(provider: str, llm_endpoint: str | None) -> bool: + """Return whether this collect job should use an LLM-backed teacher.""" + return llm_endpoint is not None or provider != "scripted" + + +def _teacher_provider_label(provider: str, llm_endpoint: str | None) -> str: + """Normalize the provider label shown in metadata and CLI output.""" + if llm_endpoint is not None: + return "openai-compatible" + return provider + + +def _resolve_api_key(provider: str) -> str: + for name in _PROVIDER_API_KEY_ENVS.get(provider, []): + value = os.getenv(name) + if value: + return value + raise typer.BadParameter( + f"No API key found for provider={provider!r}. " + f"Set one of: {_PROVIDER_API_KEY_ENVS.get(provider, [])}", + param_hint="--provider", + ) + + +def _build_scripted_model_step(): + """Default teacher: pick the first legal action from the latest observation. + + Matches the scripted teacher in ``examples/ttt_collect_demo.py``. Lets + users smoke-test the CLI without any API key configured. + """ + + def _extract_legal_actions(messages: list[dict[str, Any]]) -> list[int]: + for message in reversed(messages): + content = message.get("content") or "" + if not isinstance(content, str): + continue + try: + payload = json.loads(content) + except (json.JSONDecodeError, ValueError): + payload = None + if isinstance(payload, dict) and "legal_actions" in payload: + legal = payload["legal_actions"] + if isinstance(legal, list): + return [int(a) for a in legal] + match = re.search(r"Legal actions:\s*\[([^\]]*)\]", content) + if match: + raw = match.group(1).strip() + if not raw: + return [] + return [int(x) for x in raw.split(",")] + return [] + + def model_step(messages, tools, sampling): + del sampling + legal = _extract_legal_actions(messages) + action_id = legal[0] if legal else 0 + tool_name = tools[0].name if tools else "play_move" + return ModelStepResult( + response=LLMResponse( + content=f"Playing action_id={action_id} (first legal).", + tool_calls=[ + ToolCall( + id=f"scripted-{action_id}", + name=tool_name, + args={"action_id": action_id}, + ) + ], + ), + ) + + return model_step + + +def _build_llm_model_step( + provider: str, + model: str, + *, + llm_endpoint: str | None, + llm_port: int, + temperature: float, + max_tokens: int, + system_prompt: str | None = None, +): + effective_system_prompt = system_prompt or _SYSTEM_PROMPT + + if llm_endpoint: + # Self-hosted OpenAI-compatible endpoint (vLLM, TGI, Ollama, ...). + from openenv.core.llm_client import OpenAIClient + + client: LLMClient = OpenAIClient( + endpoint=llm_endpoint, + port=llm_port, + model=model, + api_key=os.getenv("OPENAI_API_KEY") or "not-needed", + system_prompt=effective_system_prompt, + temperature=temperature, + max_tokens=max_tokens, + ) + else: + client = create_llm_client( + provider=provider, + model=model, + api_key=_resolve_api_key(provider), + system_prompt=effective_system_prompt, + temperature=temperature, + max_tokens=max_tokens, + ) + + return build_model_step(client, system_prompt=effective_system_prompt) + + +def _build_session_factory( + env_spec: str, + base_url: str, + dataset_config: dict[str, Any] | None = None, +): + """Dispatch an env spec to a session factory. + + Supported specs: + - ``openspiel:`` e.g. ``openspiel:tic_tac_toe`` + - ``reasoning_gym:`` e.g. ``reasoning_gym:chain_sum`` + """ + env_name, _, variant = env_spec.partition(":") + if not variant: + raise typer.BadParameter( + f"Missing variant. Use e.g. {env_name}:chain_sum", + param_hint="ENV", + ) + + if env_name == "openspiel": + if OpenSpielEnv is None or OpenSpielSessionFactory is None: + raise typer.BadParameter( + "openspiel_env is not importable. Ensure envs/ is on PYTHONPATH.", + param_hint="ENV", + ) + return OpenSpielSessionFactory( + lambda: OpenSpielEnv(base_url=base_url), + game_name=variant, + ) + + if env_name == "reasoning_gym": + if ReasoningGymEnv is None or ReasoningGymSessionFactory is None: + raise typer.BadParameter( + "reasoning_gym_env is not importable. Ensure envs/ is on PYTHONPATH.", + param_hint="ENV", + ) + return ReasoningGymSessionFactory( + lambda: ReasoningGymEnv(base_url=base_url), + dataset_name=variant, + dataset_config=dataset_config, + ) + + raise typer.BadParameter( + f"Unknown env {env_spec!r}. Supported: openspiel:, reasoning_gym:", + param_hint="ENV", + ) + + +@app.command() +def collect( + env: Annotated[ + str, + typer.Argument( + help="Env spec in 'family:variant' form (e.g. openspiel:tic_tac_toe).", + ), + ], + base_url: Annotated[ + str, + typer.Option( + "--base-url", + help="Env server URL (local Docker or Hugging Face Space).", + ), + ], + output_dir: Annotated[ + Path, + typer.Option( + "--output-dir", + "-o", + help="Directory to write results.jsonl + metadata.json.", + ), + ], + num_episodes: Annotated[ + int, + typer.Option("--num-episodes", "-n", help="Number of episodes to collect."), + ] = 10, + max_turns: Annotated[ + int, typer.Option("--max-turns", help="Max tool/model turns per episode.") + ] = 9, + episode_id_prefix: Annotated[ + str, + typer.Option("--episode-id-prefix", help="Prefix for serialized episode ids."), + ] = "ep", + resume: Annotated[ + bool, + typer.Option( + "--resume/--no-resume", + help="Skip episodes already present in results.jsonl.", + ), + ] = True, + provider: Annotated[ + str, + typer.Option( + "--provider", + help="Teacher provider: scripted | openai | anthropic.", + ), + ] = "scripted", + model: Annotated[ + str | None, + typer.Option( + "--model", + help="Model id (required when provider != scripted).", + ), + ] = None, + llm_endpoint: Annotated[ + str | None, + typer.Option( + "--llm-endpoint", + help="OpenAI-compatible endpoint URL (for self-hosted vLLM/TGI/Ollama).", + ), + ] = None, + llm_port: Annotated[ + int, + typer.Option("--llm-port", help="Port for self-hosted LLM endpoint."), + ] = 8000, + temperature: Annotated[ + float, typer.Option("--temperature", help="Sampling temperature.") + ] = 0.2, + max_tokens: Annotated[ + int, typer.Option("--max-tokens", help="Max completion tokens.") + ] = 200, + keep_losses: Annotated[ + bool, + typer.Option( + "--keep-losses", + help="Keep losing rollouts (default: filter rollouts with reward < 0).", + ), + ] = False, + push_to_hub: Annotated[ + str | None, + typer.Option( + "--push-to-hub", + "-H", + help="Destination dataset repo id ('user/name'). Uploads after collect.", + ), + ] = None, + private: Annotated[ + bool, + typer.Option( + "--private", + help="Create the Hub dataset repo as private.", + ), + ] = False, + commit_message: Annotated[ + str | None, + typer.Option( + "--commit-message", + help="Commit message for the Hub upload.", + ), + ] = None, + dataset_config: Annotated[ + str | None, + typer.Option( + "--dataset-config", + help=( + "JSON string of dataset config for envs that support it " + "(e.g. reasoning_gym). Example: " + '\'{"min_terms": 2, "max_terms": 3}\'' + ), + ), + ] = None, + system_prompt: Annotated[ + str | None, + typer.Option( + "--system-prompt", + help="Custom system prompt for the teacher model.", + ), + ] = None, +) -> None: + """Collect rollouts from a deployed OpenEnv environment.""" + uses_llm_teacher = _uses_llm_teacher(provider, llm_endpoint) + teacher_provider = _teacher_provider_label(provider, llm_endpoint) + + if uses_llm_teacher and not model: + raise typer.BadParameter( + "--model is required when using a hosted provider or --llm-endpoint.", + param_hint="--model", + ) + + parsed_dataset_config: dict[str, Any] | None = None + if dataset_config is not None: + try: + parsed_dataset_config = json.loads(dataset_config) + except json.JSONDecodeError as exc: + raise typer.BadParameter( + f"--dataset-config must be valid JSON: {exc}", + param_hint="--dataset-config", + ) + + factory = _build_session_factory( + env, base_url, dataset_config=parsed_dataset_config + ) + serializer = RolloutSerializer(output_dir) + serializer.write_metadata( + { + "env": env, + "env_base_url": base_url, + "provider": teacher_provider, + "model": model, + "llm_endpoint": llm_endpoint, + "llm_port": llm_port if llm_endpoint else None, + "num_episodes_requested": num_episodes, + "temperature": temperature, + "keep_losses": keep_losses, + } + ) + + if uses_llm_teacher: + model_step = _build_llm_model_step( + provider=provider, + model=model, # type: ignore[arg-type] # validated above + llm_endpoint=llm_endpoint, + llm_port=llm_port, + temperature=temperature, + max_tokens=max_tokens, + system_prompt=system_prompt, + ) + else: + model_step = _build_scripted_model_step() + + should_keep = None if keep_losses else (lambda record: record.reward >= 0.0) + + collect_runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=serializer, + limits=HarnessRunLimits(max_turns=max_turns), + ) + + console.print(f"[cyan]Env server:[/cyan] {base_url}") + console.print( + f"[cyan]Teacher:[/cyan] {teacher_provider}" + (f"/{model}" if model else "") + ) + console.print(f"[cyan]Output:[/cyan] {output_dir}") + + result = collect_runner.run( + model_step=model_step, + num_episodes=num_episodes, + episode_id_prefix=episode_id_prefix, + resume=resume, + should_keep=should_keep, + ) + + console.print( + f"[green]Collected[/green]={result.num_collected} " + f"[yellow]skipped[/yellow]={result.num_skipped} " + f"[red]dropped[/red]={result.num_dropped} " + f"[red]failed[/red]={result.num_failed}" + ) + console.print( + f"avg_reward={result.avg_reward:.3f} success_rate={result.success_rate:.0%}" + ) + + if push_to_hub: + console.print(f"[cyan]Pushing to Hub:[/cyan] {push_to_hub}") + url = push_to_hf_hub( + output_dir=output_dir, + repo_id=push_to_hub, + private=private, + commit_message=commit_message, + ) + console.print(f"[green]Dataset at:[/green] {url}") + + +def _invoke_via_entry_point() -> None: # pragma: no cover - wired in __main__ + sys.exit(app()) diff --git a/src/openenv/core/README.md b/src/openenv/core/README.md index 5d153f1e4f..5698a79010 100644 --- a/src/openenv/core/README.md +++ b/src/openenv/core/README.md @@ -28,6 +28,32 @@ Core components for OpenEnv - a framework for building HTTP-based agentic enviro - **Container Providers**: Pluggable architecture for running containers (Docker, Kubernetes, etc.) - **Type System**: Strongly-typed Action/Observation/State interfaces - **Web Interface**: Optional web UI for interacting with environments +- **Experimental Harness Helpers**: `openenv.core.harness` provides + `ResourceSessionFactory`, `StepEnvSessionAdapter`, and `HarnessAdapter` for + MCP-first training/evaluation harnesses + +## Experimental Harness Helpers + +OpenEnv now includes an additive harness-facing layer for cases where a rollout +driver should interact with environment resources through tools instead of +calling `reset()` / `step()` directly. These helpers are currently +experimental while RFC 005 remains in review, so import them from +`openenv.core.harness`. + +Core pieces: + +- `ResourceSessionFactory.create(task, seed, episode_id)` creates an isolated + per-rollout session +- `StepEnvSessionAdapter` wraps an existing client and exposes session tools +- `SessionMCPBridge` exposes a session through MCP-style JSON-RPC +- `MCPHarnessAdapter` drives white-box rollouts where the trainer still owns + model sampling, token IDs, and logprobs +- `CLIHarnessAdapter` supports opaque evaluation harnesses over the same session + surface +- `build_harness_rollout_func(...)` builds a TRL-compatible `rollout_func` + +This layer is additive. Existing environment clients and MCP environments keep +their current control-plane APIs unchanged. ## Installation diff --git a/src/openenv/core/harness/README.md b/src/openenv/core/harness/README.md new file mode 100644 index 0000000000..d45199e81a --- /dev/null +++ b/src/openenv/core/harness/README.md @@ -0,0 +1,104 @@ +# `openenv.core.harness` + +Session-based runtime for driving rollouts against an OpenEnv environment +through an LLM + MCP-style tool loop. + +- `__init__.py` — runtime primitives (from RFC 005 / PR #471): + `ResourceSession`, `ResourceSessionFactory`, `StepEnvSessionAdapter`, + `HarnessAdapter`, `MCPHarnessAdapter`, `CLIHarnessAdapter`, + `build_harness_rollout_func`, etc. +- `collect.py` — synthetic dataset generation on top of the runtime: + `EpisodeRecord`, `RolloutSerializer`, `CollectRunner`, `build_model_step`, + `push_to_hf_hub`. + +## Quick start + +The `openenv collect` CLI wires everything together. See `openenv collect --help`. + +```bash +# Smoke test (no API key, no Docker — scripted teacher picks the first legal action) +openenv collect openspiel:tic_tac_toe \ + --base-url https://-.hf.space \ + --output-dir /tmp/ttt-scripted \ + -n 10 --provider scripted + +# Real LLM teacher + push the dataset to the Hub +OPENAI_API_KEY=... openenv collect openspiel:tic_tac_toe \ + --base-url https://-.hf.space \ + --output-dir /tmp/ttt-sft-v1 \ + -n 200 --provider openai --model gpt-5-mini \ + --push-to-hub /ttt-sft-v1 +``` + +Programmatic use: + +```python +from openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter +from openenv.core.harness.collect import ( + CollectRunner, + RolloutSerializer, + build_model_step, + push_to_hf_hub, +) +from openenv.core.llm_client import create_llm_client +from openspiel_env.client import OpenSpielEnv +from openspiel_env.harness import OpenSpielSessionFactory + +factory = OpenSpielSessionFactory( + lambda: OpenSpielEnv(base_url="https://-.hf.space"), + game_name="tic_tac_toe", +) +llm = create_llm_client(provider="openai", model="gpt-5-mini", api_key=...) +runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=RolloutSerializer("/tmp/ttt-sft-v1"), + limits=HarnessRunLimits(max_turns=9), +) + +result = runner.run( + model_step=build_model_step(llm), + num_episodes=200, + should_keep=lambda r: r.reward >= 0.0, +) +push_to_hf_hub("/tmp/ttt-sft-v1", "/ttt-sft-v1") +``` + +## Output schema + +`results.jsonl` — one JSON object per episode: + +| field | type | description | +| --- | --- | --- | +| `episode_id` | string | Stable id for resume (`-`) | +| `messages` | list | Chat transcript (user/assistant/tool). TRL `SFTTrainer`-compatible | +| `reward` | float | Env reward (derived via `_resolve_env_reward`, never synthesized) | +| `done` | bool | Whether the episode terminated | +| `tool_trace` | list | Structured tool calls with args & results | +| `metrics` | struct | Harness-level metrics (turns, tool_calls) | +| `verify_metrics` | struct | Session-level metrics (step_count, …) | +| `artifacts` | struct | Verification artifacts (final state, etc.) | +| `task` | any | Optional task spec passed by `CollectRunner.tasks` | +| `extra` | struct | Free-form annotations supplied by the caller | + +`metadata.json` — run-level info (env, model, n, temperature, filter flags). + +`README.md` is auto-generated by `push_to_hf_hub` with YAML front-matter so the +HF Dataset Viewer renders `results.jsonl` as `split=train`. + +## Design notes + +- **Thin envs** — `OpenSpielSessionFactory` and friends live in the env package + (`envs//harness.py`) and wrap the existing `EnvClient`. The collector + treats the env as a black box. +- **Rewards in env** — `EpisodeRecord.from_rollout` reuses the runtime's + `_resolve_env_reward` so any mismatch between the reward emitted by a tool + result and the one forwarded by `session.verify()` raises immediately. +- **Provider-agnostic teacher** — `build_model_step` adapts any + `openenv.core.llm_client.LLMClient` (OpenAI, Anthropic, OpenAI-compatible + self-hosted endpoints) into a `ModelStep`. No schema code per provider. +- **Resume** — `RolloutSerializer.collected_episode_ids()` reads prior runs + from disk; `CollectRunner.run(resume=True)` skips ids already serialized. +- **Filtering** — `should_keep` is a `Callable[[EpisodeRecord], bool]`. Defaults + to keeping everything. The CLI filters `reward >= 0` unless `--keep-losses` + is set. diff --git a/src/openenv/core/harness/__init__.py b/src/openenv/core/harness/__init__.py new file mode 100644 index 0000000000..8c1412ca49 --- /dev/null +++ b/src/openenv/core/harness/__init__.py @@ -0,0 +1,723 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Experimental harness helpers for training and evaluation. + +These helpers live outside the stable ``openenv.core`` package surface while +RFC 005 is still under review. Import them from ``openenv.core.harness``. +""" + +from __future__ import annotations + +import json +import math +from abc import ABC, abstractmethod +from dataclasses import dataclass, field +from typing import Any, Callable, Protocol + +from ..client_types import StepResult +from ..env_server.mcp_types import JsonRpcErrorCode, JsonRpcResponse, Tool +from ..env_server.types import State +from ..llm_client import LLMResponse + +Message = dict[str, Any] +RESERVED_TOOL_NAMES = frozenset({"reset", "step", "state", "close"}) + + +@dataclass +class ToolResult: + """Normalized result from a resource session tool invocation.""" + + data: Any = None + done: bool = False + metadata: dict[str, Any] = field(default_factory=dict) + error: str | None = None + + +@dataclass +class VerifyResult: + """Final rollout data produced after a rollout completes. + + ``env_reward`` must forward reward already produced inside the environment. + It must not synthesize a new reward in the orchestration layer. + """ + + env_reward: float | None = None + done: bool = False + metrics: dict[str, Any] = field(default_factory=dict) + artifacts: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class RolloutEvent: + """Event emitted while a harness drives a rollout.""" + + type: str + payload: dict[str, Any] = field(default_factory=dict) + + +@dataclass +class ToolTraceEntry: + """Record of a tool call issued by the harness.""" + + tool_name: str + arguments: dict[str, Any] + result: ToolResult + + +@dataclass +class ModelStepResult: + """Structured output from one white-box model sampling step.""" + + response: LLMResponse + prompt_ids: list[int] = field(default_factory=list) + completion_ids: list[int] = field(default_factory=list) + logprobs: list[float] = field(default_factory=list) + + +@dataclass +class HarnessRolloutResult: + """Complete result of a harness-driven rollout.""" + + messages: list[Message] = field(default_factory=list) + tool_trace: list[ToolTraceEntry] = field(default_factory=list) + events: list[RolloutEvent] = field(default_factory=list) + done: bool = False + metrics: dict[str, Any] = field(default_factory=dict) + prompt_ids: list[int] = field(default_factory=list) + completion_ids: list[int] = field(default_factory=list) + logprobs: list[float] = field(default_factory=list) + + +@dataclass +class HarnessRunLimits: + """Execution limits for harness-driven rollouts.""" + + max_turns: int = 10 + max_tool_calls_per_turn: int | None = None + max_total_tool_calls: int | None = None + sampling: dict[str, Any] = field(default_factory=dict) + + +class ModelStep(Protocol): + """Callable used by white-box harnesses to sample the next model turn.""" + + def __call__( + self, + messages: list[Message], + tools: list[Tool], + sampling: dict[str, Any], + ) -> ModelStepResult: ... + + +class ResourceSession(ABC): + """Per-rollout environment/resource session exposed to harnesses.""" + + @abstractmethod + def initial_messages(self) -> list[Message]: + """Return initial prompt messages for the harness.""" + + @abstractmethod + def list_tools(self) -> list[Tool]: + """Return the MCP-style tool manifest exposed by the session.""" + + @abstractmethod + def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: + """Invoke a session tool.""" + + @abstractmethod + def verify( + self, + transcript: list[Message], + final_state: Any | None = None, + ) -> VerifyResult: + """Finalize a rollout after the harness has stopped. + + This hook may add metrics or artifacts and may forward the final reward + already produced by the environment. + """ + + @abstractmethod + def close(self) -> None: + """Release session resources.""" + + +class ResourceSessionFactory(ABC): + """Factory for producing isolated per-rollout sessions.""" + + @abstractmethod + def create( + self, + task: Any, + seed: int | None = None, + episode_id: str | None = None, + ) -> ResourceSession: + """Create one isolated resource session for a rollout.""" + + +class HarnessAdapter(ABC): + """Interface implemented by harness drivers.""" + + @abstractmethod + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + """Run a rollout while the trainer owns model sampling.""" + + @abstractmethod + def run_black_box( + self, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + """Run a rollout with an opaque harness (evaluation-only path).""" + + +def _serialize_for_message(value: Any) -> str: + """Convert structured tool data to a stable text payload.""" + + if value is None: + return "" + if isinstance(value, str): + return value + return json.dumps(value, sort_keys=True, default=str) + + +def _state_to_data(state: Any) -> Any: + """Convert state objects to plain data for metrics and artifacts.""" + + if state is None: + return None + if hasattr(state, "model_dump"): + return state.model_dump() + return state + + +def _tool_result_reward(tool_result: ToolResult) -> float | None: + """Extract the environment reward already emitted by a tool result.""" + + reward = tool_result.metadata.get("reward") + if reward is None and isinstance(tool_result.data, dict): + reward = tool_result.data.get("reward") + if reward is None: + return None + return float(reward) + + +def _resolve_env_reward( + rollout: HarnessRolloutResult, + verify: VerifyResult, +) -> float: + """Resolve the final environment reward without allowing external synthesis.""" + + trace_reward: float | None = None + for entry in reversed(rollout.tool_trace): + trace_reward = _tool_result_reward(entry.result) + if trace_reward is not None: + break + + verify_reward = None if verify.env_reward is None else float(verify.env_reward) + if ( + trace_reward is not None + and verify_reward is not None + and not math.isclose( + verify_reward, + trace_reward, + rel_tol=1e-9, + abs_tol=1e-6, + ) + ): + raise ValueError( + "verify.env_reward must forward the environment reward from the rollout" + ) + + if trace_reward is not None: + return trace_reward + if verify_reward is not None: + return verify_reward + raise ValueError("rollout did not produce an environment reward") + + +class StepEnvSessionAdapter(ResourceSession): + """Expose an existing step/reset/state client as a resource session.""" + + def __init__( + self, + client: Any, + *, + task: Any = None, + seed: int | None = None, + episode_id: str | None = None, + tool_specs: list[Tool], + action_builder: Callable[[str, dict[str, Any]], Any], + initial_messages_builder: Callable[[StepResult[Any], Any], list[Message]], + tool_result_builder: Callable[ + [str, dict[str, Any], StepResult[Any], Any], + ToolResult, + ] + | None = None, + verify_builder: Callable[ + [list[Message], Any | None, StepResult[Any] | None, Any], + VerifyResult, + ] + | None = None, + reset_kwargs: dict[str, Any] | None = None, + ): + if hasattr(client, "sync") and callable(client.sync): + self._client = client.sync() + else: + self._client = client + + self._task = task + self._tool_specs = list(tool_specs) + self._action_builder = action_builder + self._initial_messages_builder = initial_messages_builder + self._tool_result_builder = tool_result_builder or self._default_tool_result + self._verify_builder = verify_builder or self._default_verify + self._closed = False + self._tools_by_name = {tool.name: tool for tool in self._tool_specs} + reserved = sorted(set(self._tools_by_name) & RESERVED_TOOL_NAMES) + if reserved: + raise ValueError( + "Tool names are reserved for orchestration controls: " + + ", ".join(reserved) + ) + + reset_payload = dict(reset_kwargs or {}) + if seed is not None: + reset_payload.setdefault("seed", seed) + if episode_id is not None: + reset_payload.setdefault("episode_id", episode_id) + + self._last_result: StepResult[Any] | None = None + self._last_state = None + try: + self._initial_result: StepResult[Any] = self._client.reset(**reset_payload) + self._last_state = self._read_state() + except Exception: + self.close() + raise + + def _read_state(self) -> Any: + if hasattr(self._client, "state") and callable(self._client.state): + return self._client.state() + return None + + def _default_tool_result( + self, + tool_name: str, + arguments: dict[str, Any], + result: StepResult[Any], + state: Any, + ) -> ToolResult: + return ToolResult( + data={ + "tool_name": tool_name, + "arguments": dict(arguments), + "observation": result.observation, + "reward": result.reward, + "done": result.done, + }, + done=bool(result.done), + metadata={ + "reward": result.reward, + "state": _state_to_data(state), + }, + ) + + def _default_verify( + self, + transcript: list[Message], + final_state: Any | None, + last_result: StepResult[Any] | None, + state: Any, + ) -> VerifyResult: + reward = None if last_result is None else last_result.reward + done = False if last_result is None else bool(last_result.done) + metrics = {} + if isinstance(state, State): + metrics["step_count"] = state.step_count + elif isinstance(state, dict) and "step_count" in state: + metrics["step_count"] = state["step_count"] + + return VerifyResult( + env_reward=reward, + done=done, + metrics=metrics, + artifacts={ + "final_state": _state_to_data(state), + "transcript_length": len(transcript), + }, + ) + + def initial_messages(self) -> list[Message]: + return list(self._initial_messages_builder(self._initial_result, self._task)) + + def list_tools(self) -> list[Tool]: + return list(self._tool_specs) + + def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: + if name in RESERVED_TOOL_NAMES: + raise KeyError(f"Reserved orchestration tool: {name}") + if name not in self._tools_by_name: + raise KeyError(f"Unknown tool: {name}") + + action = self._action_builder(name, dict(arguments)) + result = self._client.step(action) + state = self._read_state() + self._last_result = result + self._last_state = state + return self._tool_result_builder(name, dict(arguments), result, state) + + def verify( + self, + transcript: list[Message], + final_state: Any | None = None, + ) -> VerifyResult: + return self._verify_builder( + list(transcript), + final_state, + self._last_result, + self._last_state, + ) + + def close(self) -> None: + if self._closed: + return + self._closed = True + if hasattr(self._client, "close") and callable(self._client.close): + self._client.close() + + +class SessionMCPBridge: + """Expose a resource session through an in-process MCP JSON-RPC bridge.""" + + def __init__(self, session: ResourceSession): + self.session = session + + def handle_request(self, request: dict[str, Any]) -> dict[str, Any]: + request_id = request.get("id") + method = request.get("method") + params = request.get("params", {}) or {} + + if request.get("jsonrpc") != "2.0": + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_REQUEST, + message="Invalid Request", + request_id=request_id, + ).model_dump() + + if method == "tools/list": + tools = [ + { + "name": tool.name, + "description": tool.description, + "inputSchema": tool.input_schema, + } + for tool in self.session.list_tools() + ] + return JsonRpcResponse.success( + {"tools": tools}, + request_id=request_id, + ).model_dump() + + if method == "tools/call": + if "name" not in params: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_PARAMS, + message="Missing tool name", + request_id=request_id, + ).model_dump() + + tool_name = params["name"] + if tool_name in RESERVED_TOOL_NAMES: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_PARAMS, + message=f"Reserved orchestration tool: {tool_name}", + request_id=request_id, + ).model_dump() + + try: + result = self.session.call_tool( + tool_name, + dict(params.get("arguments", {})), + ) + except KeyError as exc: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.METHOD_NOT_FOUND, + message=str(exc.args[0]) if exc.args else "Method not found", + request_id=request_id, + ).model_dump() + except ValueError as exc: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INVALID_PARAMS, + message=str(exc), + request_id=request_id, + ).model_dump() + except Exception as exc: + return JsonRpcResponse.error_response( + JsonRpcErrorCode.INTERNAL_ERROR, + message=str(exc), + request_id=request_id, + ).model_dump() + return JsonRpcResponse.success( + { + "data": result.data, + "done": result.done, + "metadata": dict(result.metadata), + "error": result.error, + }, + request_id=request_id, + ).model_dump() + + return JsonRpcResponse.error_response( + JsonRpcErrorCode.METHOD_NOT_FOUND, + message="Method not found", + request_id=request_id, + ).model_dump() + + +class MCPHarnessAdapter(HarnessAdapter): + """White-box harness that follows an MCP tool-calling loop.""" + + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + run_limits = limits or HarnessRunLimits() + messages = list(session.initial_messages()) + tools = session.list_tools() + result = HarnessRolloutResult(messages=list(messages)) + total_tool_calls = 0 + + for turn_index in range(run_limits.max_turns): + step_result = model_step(messages, tools, dict(run_limits.sampling)) + assistant_message = step_result.response.to_message_dict() + messages.append(assistant_message) + + result.prompt_ids.extend(step_result.prompt_ids) + result.completion_ids.extend(step_result.completion_ids) + result.logprobs.extend(step_result.logprobs) + result.events.append( + RolloutEvent( + type="model_response", + payload={ + "turn": turn_index, + "content": step_result.response.content, + "tool_calls": [ + { + "id": tool_call.id, + "name": tool_call.name, + "arguments": dict(tool_call.args), + } + for tool_call in step_result.response.tool_calls + ], + }, + ) + ) + + if not step_result.response.tool_calls: + result.done = True + break + + tool_calls = step_result.response.tool_calls + if run_limits.max_tool_calls_per_turn is not None: + tool_calls = tool_calls[: run_limits.max_tool_calls_per_turn] + + for tool_call in tool_calls: + if ( + run_limits.max_total_tool_calls is not None + and total_tool_calls >= run_limits.max_total_tool_calls + ): + result.metrics["truncated"] = True + result.messages = list(messages) + return result + + tool_result = session.call_tool(tool_call.name, dict(tool_call.args)) + total_tool_calls += 1 + + trace_entry = ToolTraceEntry( + tool_name=tool_call.name, + arguments=dict(tool_call.args), + result=tool_result, + ) + result.tool_trace.append(trace_entry) + result.events.append( + RolloutEvent( + type="tool_call", + payload={ + "tool_name": tool_call.name, + "arguments": dict(tool_call.args), + "done": tool_result.done, + }, + ) + ) + + messages.append( + { + "role": "tool", + "tool_call_id": tool_call.id, + "name": tool_call.name, + "content": _serialize_for_message(tool_result.data), + } + ) + + if tool_result.done: + result.done = True + break + + if result.done: + break + + result.messages = list(messages) + result.metrics.setdefault( + "turns", + len([event for event in result.events if event.type == "model_response"]), + ) + result.metrics.setdefault("tool_calls", len(result.tool_trace)) + return result + + def run_black_box( + self, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + raise NotImplementedError( + "MCPHarnessAdapter is a white-box harness. " + "Use CLIHarnessAdapter for opaque evaluation harnesses." + ) + + +class CLIHarnessAdapter(HarnessAdapter): + """Thin black-box adapter for opaque CLI-style harnesses.""" + + def __init__( + self, + runner: Callable[ + [SessionMCPBridge, ResourceSession, HarnessRunLimits], + HarnessRolloutResult, + ], + ): + self._runner = runner + + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + raise NotImplementedError( + "CLIHarnessAdapter only supports black-box evaluation rollouts." + ) + + def run_black_box( + self, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + run_limits = limits or HarnessRunLimits() + bridge = SessionMCPBridge(session) + return self._runner(bridge, session, run_limits) + + +def build_harness_rollout_func( + *, + session_factory: Any, + harness_adapter: HarnessAdapter, + model_step_builder: Callable[[Any, ResourceSession], ModelStep], + limits: HarnessRunLimits | None = None, + reward_key: str = "env_reward", +) -> Callable[[list[Any], Any], dict[str, list[Any]]]: + """Build a TRL-compatible rollout function from sessions and harnesses.""" + + def rollout_func(prompts: list[Any], trainer: Any) -> dict[str, list[Any]]: + all_prompt_ids: list[list[int]] = [] + all_completion_ids: list[list[int]] = [] + all_logprobs: list[list[float]] = [] + rewards: list[float] = [] + verify_metrics: list[dict[str, Any]] = [] + + for prompt in prompts: + session = session_factory.create(task=prompt) + try: + model_step = model_step_builder(trainer, session) + rollout = harness_adapter.run_white_box( + model_step=model_step, + session=session, + limits=limits, + ) + verify = session.verify( + transcript=rollout.messages, + final_state={ + "done": rollout.done, + "metrics": dict(rollout.metrics), + "events": [ + { + "type": event.type, + "payload": dict(event.payload), + } + for event in rollout.events + ], + "tool_trace": [ + { + "tool_name": entry.tool_name, + "arguments": dict(entry.arguments), + "result": { + "data": entry.result.data, + "done": entry.result.done, + "metadata": dict(entry.result.metadata), + "error": entry.result.error, + }, + } + for entry in rollout.tool_trace + ], + }, + ) + + all_prompt_ids.append(list(rollout.prompt_ids)) + all_completion_ids.append(list(rollout.completion_ids)) + all_logprobs.append(list(rollout.logprobs)) + rewards.append(_resolve_env_reward(rollout, verify)) + verify_metrics.append(dict(verify.metrics)) + finally: + session.close() + + return { + "prompt_ids": all_prompt_ids, + "completion_ids": all_completion_ids, + "logprobs": all_logprobs, + reward_key: rewards, + "verify_metrics": verify_metrics, + } + + return rollout_func + + +__all__ = [ + "CLIHarnessAdapter", + "HarnessAdapter", + "HarnessRolloutResult", + "HarnessRunLimits", + "MCPHarnessAdapter", + "Message", + "ModelStep", + "ModelStepResult", + "RESERVED_TOOL_NAMES", + "ResourceSession", + "ResourceSessionFactory", + "RolloutEvent", + "SessionMCPBridge", + "StepEnvSessionAdapter", + "ToolResult", + "ToolTraceEntry", + "VerifyResult", + "build_harness_rollout_func", +] diff --git a/src/openenv/core/harness/collect.py b/src/openenv/core/harness/collect.py new file mode 100644 index 0000000000..0cab2eadf7 --- /dev/null +++ b/src/openenv/core/harness/collect.py @@ -0,0 +1,586 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Rollout collection helpers for synthetic dataset generation. + +Complements ``openenv.core.harness`` with: + +- ``EpisodeRecord`` — serializable view of one rollout + its verification. +- ``RolloutSerializer`` — append-only JSONL writer with metadata sidecar. +- ``CollectRunner`` — orchestrates repeated rollouts, supports resume, and + optional per-record filtering via ``should_keep``. + +The serialized schema is designed to be consumed directly by TRL's +``SFTTrainer`` (``messages`` column) or by ``datasets.Dataset.from_json``. +""" + +from __future__ import annotations + +import json +import warnings +from dataclasses import asdict, dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Callable, Iterable, Iterator + +from ..env_server.mcp_types import Tool +from ..llm_client import LLMClient +from ..utils import run_async_safely +from . import ( + _resolve_env_reward, + HarnessAdapter, + HarnessRolloutResult, + HarnessRunLimits, + ModelStep, + ModelStepResult, + ResourceSessionFactory, + VerifyResult, +) + +_SUPPORTED_SAMPLING_KEYS = frozenset({"temperature", "max_tokens", "top_p"}) + +RESULTS_FILENAME = "results.jsonl" +METADATA_FILENAME = "metadata.json" + + +def _tool_trace_to_plain(rollout: HarnessRolloutResult) -> list[dict[str, Any]]: + return [ + { + "tool_name": entry.tool_name, + "arguments": dict(entry.arguments), + "result": { + "data": entry.result.data, + "done": entry.result.done, + "metadata": dict(entry.result.metadata), + "error": entry.result.error, + }, + } + for entry in rollout.tool_trace + ] + + +@dataclass +class EpisodeRecord: + """Serializable view of one collected episode.""" + + episode_id: str + messages: list[dict[str, Any]] + reward: float + done: bool + tool_trace: list[dict[str, Any]] + metrics: dict[str, Any] + verify_metrics: dict[str, Any] + artifacts: dict[str, Any] + task: Any = None + extra: dict[str, Any] = field(default_factory=dict) + + @classmethod + def from_rollout( + cls, + episode_id: str, + rollout: HarnessRolloutResult, + verify: VerifyResult, + task: Any = None, + extra: dict[str, Any] | None = None, + ) -> "EpisodeRecord": + """Build a record from a harness rollout and its verification. + + Uses ``_resolve_env_reward`` so that any disagreement between the + reward emitted inside the environment and the one forwarded by + ``verify()`` raises — preserving the "rewards in env" invariant. + """ + reward = _resolve_env_reward(rollout, verify) + return cls( + episode_id=episode_id, + messages=[dict(m) for m in rollout.messages], + reward=reward, + done=bool(rollout.done or verify.done), + tool_trace=_tool_trace_to_plain(rollout), + metrics=dict(rollout.metrics), + verify_metrics=dict(verify.metrics), + artifacts=dict(verify.artifacts), + task=task, + extra=dict(extra or {}), + ) + + def to_dict(self) -> dict[str, Any]: + return json.loads(json.dumps(asdict(self), default=str)) + + +class RolloutSerializer: + """Append-only JSONL writer for collected episodes.""" + + def __init__(self, output_dir: Path | str): + self._output_dir = Path(output_dir) + + @property + def output_dir(self) -> Path: + return self._output_dir + + @property + def results_path(self) -> Path: + return self._output_dir / RESULTS_FILENAME + + @property + def metadata_path(self) -> Path: + return self._output_dir / METADATA_FILENAME + + def write_episode(self, record: EpisodeRecord) -> None: + self._output_dir.mkdir(parents=True, exist_ok=True) + line = json.dumps(record.to_dict(), default=str) + with self.results_path.open("a", encoding="utf-8") as handle: + handle.write(line + "\n") + + def write_metadata(self, metadata: dict[str, Any]) -> None: + self._output_dir.mkdir(parents=True, exist_ok=True) + self.metadata_path.write_text( + json.dumps(metadata, indent=2, default=str), + encoding="utf-8", + ) + + def collected_episode_ids(self) -> set[str]: + """Return episode ids already persisted on disk. Used for resume.""" + if not self.results_path.exists(): + return set() + + ids: set[str] = set() + with self.results_path.open("r", encoding="utf-8") as handle: + for raw in handle: + raw = raw.strip() + if not raw: + continue + try: + payload = json.loads(raw) + except json.JSONDecodeError: + continue + ep_id = payload.get("episode_id") + if isinstance(ep_id, str): + ids.add(ep_id) + return ids + + +def _rollout_final_state(rollout: HarnessRolloutResult) -> dict[str, Any]: + """Build the structured payload passed to ``session.verify()``. + + Mirrors the shape used by ``build_harness_rollout_func`` so verifiers + can rely on the same contract regardless of caller (TRL vs. collector). + """ + return { + "done": rollout.done, + "metrics": dict(rollout.metrics), + "events": [ + {"type": event.type, "payload": dict(event.payload)} + for event in rollout.events + ], + "tool_trace": [ + { + "tool_name": entry.tool_name, + "arguments": dict(entry.arguments), + "result": { + "data": entry.result.data, + "done": entry.result.done, + "metadata": dict(entry.result.metadata), + "error": entry.result.error, + }, + } + for entry in rollout.tool_trace + ], + } + + +@dataclass +class CollectResult: + """Summary returned by ``CollectRunner.run()``.""" + + num_collected: int + num_skipped: int + num_dropped: int + num_failed: int + episode_ids: list[str] + avg_reward: float + success_rate: float + + +class CollectRunner: + """Drive repeated rollouts into an append-only JSONL dataset. + + The runner is stateless across ``run()`` calls — resume state lives on + disk via the ``RolloutSerializer``. A single runner instance can be + reused across runs targeting different output dirs. + """ + + def __init__( + self, + *, + session_factory: ResourceSessionFactory, + harness_adapter: HarnessAdapter, + serializer: RolloutSerializer, + tasks: Iterable[Any] | None = None, + limits: HarnessRunLimits | None = None, + ): + self._session_factory = session_factory + self._harness_adapter = harness_adapter + self._serializer = serializer + self._tasks: Iterator[Any] | None = iter(tasks) if tasks is not None else None + self._limits = limits + + def _next_task(self) -> Any: + if self._tasks is None: + return None + try: + return next(self._tasks) + except StopIteration: + return None + + def run( + self, + *, + model_step: ModelStep, + num_episodes: int, + episode_id_prefix: str = "ep", + resume: bool = True, + should_keep: Callable[[EpisodeRecord], bool] | None = None, + ) -> CollectResult: + planned_ids = [f"{episode_id_prefix}-{i:06d}" for i in range(num_episodes)] + + if resume: + already = self._serializer.collected_episode_ids() + else: + self._reset_results_file() + already = set() + + num_skipped = 0 + num_collected = 0 + num_dropped = 0 + num_failed = 0 + rewards: list[float] = [] + + from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + TextColumn, + TimeElapsedColumn, + ) + + progress = Progress( + TextColumn("[cyan]{task.description}"), + BarColumn(), + MofNCompleteColumn(), + TextColumn( + "collected={task.fields[collected]} reward={task.fields[avg_reward]:.2f}" + ), + TimeElapsedColumn(), + ) + task_id = progress.add_task( + "Collecting", total=len(planned_ids), collected=0, avg_reward=0.0 + ) + + with progress: + for episode_id in planned_ids: + # Preserve task-to-episode alignment even when resume skips ids. + task = self._next_task() + if episode_id in already: + num_skipped += 1 + progress.advance(task_id) + continue + + session = None + try: + session = self._session_factory.create( + task=task, + episode_id=episode_id, + ) + rollout = self._harness_adapter.run_white_box( + model_step=model_step, + session=session, + limits=self._limits, + ) + verify = session.verify( + transcript=rollout.messages, + final_state=_rollout_final_state(rollout), + ) + record = EpisodeRecord.from_rollout( + episode_id=episode_id, + rollout=rollout, + verify=verify, + task=task, + ) + except Exception as exc: + num_failed += 1 + progress.console.print( + f"[red]failed[/red] {episode_id}: {type(exc).__name__}: {exc}" + ) + progress.advance(task_id) + continue + finally: + if session is not None: + session.close() + + if should_keep is not None and not should_keep(record): + num_dropped += 1 + progress.advance(task_id) + continue + + self._serializer.write_episode(record) + num_collected += 1 + rewards.append(record.reward) + avg = sum(rewards) / len(rewards) + progress.update( + task_id, advance=1, collected=num_collected, avg_reward=avg + ) + + avg_reward = sum(rewards) / len(rewards) if rewards else 0.0 + success_rate = ( + sum(1 for r in rewards if r > 0) / len(rewards) if rewards else 0.0 + ) + + return CollectResult( + num_collected=num_collected, + num_skipped=num_skipped, + num_dropped=num_dropped, + num_failed=num_failed, + episode_ids=planned_ids, + avg_reward=avg_reward, + success_rate=success_rate, + ) + + def _reset_results_file(self) -> None: + if self._serializer.results_path.exists(): + results_path = self._serializer.results_path + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + backup_path = results_path.with_name( + f"{results_path.stem}.{timestamp}.bak{results_path.suffix}" + ) + counter = 1 + while backup_path.exists(): + backup_path = results_path.with_name( + f"{results_path.stem}.{timestamp}.{counter}.bak" + f"{results_path.suffix}" + ) + counter += 1 + results_path.replace(backup_path) + warnings.warn( + f"resume=False moved existing results file to {backup_path}", + RuntimeWarning, + stacklevel=2, + ) + + +def _tool_to_mcp_dict(tool: Tool) -> dict[str, Any]: + return { + "name": tool.name, + "description": tool.description, + "inputSchema": tool.input_schema, + } + + +def build_model_step( + llm_client: LLMClient, + *, + system_prompt: str | None = None, +) -> ModelStep: + """Adapt any OpenEnv ``LLMClient`` into a ``ModelStep`` for the harness. + + Works with every provider the ``LLMClient`` hierarchy already covers — + OpenAI and any OpenAI-compatible endpoint (vLLM, TGI, Ollama, HF + Inference, Together, Groq, Fireworks, ...), Anthropic natively, and + any future subclass that implements ``complete_with_tools``. Provider + -specific schema conversion lives inside the client; this helper only + handles the sync/async adaptation and tool dict shape. + """ + + def model_step( + messages: list[dict[str, Any]], + tools: list[Tool], + sampling: dict[str, Any], + ) -> ModelStepResult: + effective_messages = list(messages) + if system_prompt and not any( + m.get("role") == "system" for m in effective_messages + ): + effective_messages.insert(0, {"role": "system", "content": system_prompt}) + + tool_dicts = [_tool_to_mcp_dict(t) for t in tools] + dropped_keys = sorted(set(sampling) - _SUPPORTED_SAMPLING_KEYS) + if dropped_keys: + warnings.warn( + "Dropping unsupported sampling keys: " + ", ".join(dropped_keys), + RuntimeWarning, + stacklevel=2, + ) + filtered_sampling = { + k: v for k, v in sampling.items() if k in _SUPPORTED_SAMPLING_KEYS + } + + response = run_async_safely( + llm_client.complete_with_tools( + messages=effective_messages, + tools=tool_dicts, + **filtered_sampling, + ) + ) + return ModelStepResult(response=response) + + return model_step + + +def _count_episodes(results_path: Path) -> int: + if not results_path.exists(): + return 0 + with results_path.open("r", encoding="utf-8") as handle: + return sum(1 for line in handle if line.strip()) + + +README_FILENAME = "README.md" + +_DATASET_CARD_FRONTMATTER = """--- +configs: +- config_name: default + data_files: + - split: train + path: "results.jsonl" +--- +""" + + +def build_dataset_readme(output_dir: Path) -> str: + """Build the README.md written alongside ``results.jsonl`` on the Hub. + + The YAML front-matter pins ``results.jsonl`` as the only dataset file + so the HF Dataset Viewer does not try to cast ``metadata.json`` into + the same schema. + """ + output_dir = Path(output_dir) + results_path = output_dir / RESULTS_FILENAME + metadata_path = output_dir / METADATA_FILENAME + + metadata: dict[str, Any] = {} + if metadata_path.exists(): + try: + metadata = json.loads(metadata_path.read_text(encoding="utf-8")) + except (json.JSONDecodeError, OSError): + metadata = {} + + lines: list[str] = [ + _DATASET_CARD_FRONTMATTER.strip(), + "", + "# OpenEnv rollouts", + "", + "Collected with [OpenEnv](https://github.com/meta-pytorch/OpenEnv) " + "(`openenv collect`).", + "", + f"- **Episodes:** {_count_episodes(results_path)}", + ] + if metadata: + lines.append("- **Run metadata:**") + lines.append("") + lines.append(" | key | value |") + lines.append(" | --- | --- |") + for key, value in metadata.items(): + lines.append(f" | `{key}` | `{value}` |") + + lines.extend( + [ + "", + "## Schema", + "", + "Each line of `results.jsonl` is one episode:", + "", + "- `episode_id` (string)", + "- `messages` (chat transcript; TRL `SFTTrainer`-compatible)", + "- `reward` (float)", + "- `done` (bool)", + "- `tool_trace` (list of `{tool_name, arguments, result}`)", + "- `metrics`, `verify_metrics`, `artifacts`", + "- `task`, `extra` (optional call-site annotations)", + "", + "## Load", + "", + "```python", + "from datasets import load_dataset", + 'ds = load_dataset("/", split="train")', + "```", + "", + ] + ) + + return "\n".join(lines) + + +def push_to_hf_hub( + output_dir: Path | str, + repo_id: str, + *, + private: bool = False, + commit_message: str | None = None, + token: str | None = None, +) -> str: + """Upload a collected rollout directory to the Hugging Face Hub as a dataset. + + Uploads ``results.jsonl``, ``metadata.json`` (if present), and an + auto-generated ``README.md`` whose YAML front-matter tells the HF + Dataset Viewer to treat ``results.jsonl`` as the sole dataset file. + Consumers load the result with ``load_dataset("/", + split="train")``. + + Args: + output_dir: Directory previously populated by ``RolloutSerializer``. + repo_id: Destination repo in ``"user/name"`` form. + private: Create the dataset repo as private. + commit_message: Override the default commit message. + token: HF token. Defaults to the token resolved by ``huggingface_hub``. + + Returns: + The public URL of the dataset on the Hub. + + Raises: + FileNotFoundError: If ``results.jsonl`` is missing from ``output_dir``. + """ + output_dir = Path(output_dir) + results_path = output_dir / RESULTS_FILENAME + if not results_path.exists(): + raise FileNotFoundError( + f"No {RESULTS_FILENAME} found in {output_dir}. " + "Run a collect job before pushing to the Hub." + ) + + readme_path = output_dir / README_FILENAME + readme_path.write_text(build_dataset_readme(output_dir), encoding="utf-8") + + num_episodes = _count_episodes(results_path) + message = commit_message or f"Add {num_episodes} collected episode(s)" + + from huggingface_hub import HfApi + + api = HfApi(token=token) + api.create_repo( + repo_id=repo_id, + repo_type="dataset", + private=private, + exist_ok=True, + ) + api.upload_folder( + folder_path=str(output_dir), + repo_id=repo_id, + repo_type="dataset", + commit_message=message, + ) + return f"https://huggingface.co/datasets/{repo_id}" + + +__all__ = [ + "build_dataset_readme", + "build_model_step", + "CollectResult", + "CollectRunner", + "EpisodeRecord", + "push_to_hf_hub", + "RolloutSerializer", + "METADATA_FILENAME", + "README_FILENAME", + "RESULTS_FILENAME", +] diff --git a/src/openenv/core/llm_client.py b/src/openenv/core/llm_client.py index 9df2ff27ae..1a609ba2aa 100644 --- a/src/openenv/core/llm_client.py +++ b/src/openenv/core/llm_client.py @@ -134,6 +134,9 @@ class OpenAIClient(LLMClient): system_prompt: Optional system message prepended to every request. temperature: Default sampling temperature. max_tokens: Default max tokens in the response. + use_max_completion_tokens: Use max_completion_tokens instead of + max_tokens. Required for newer OpenAI models (gpt-5-mini, o1, o3). + Not supported by self-hosted OpenAI-compatible endpoints. """ def __init__( @@ -145,12 +148,17 @@ def __init__( system_prompt: str | None = None, temperature: float = 0.0, max_tokens: int = 256, + use_max_completion_tokens: bool = False, ): super().__init__(endpoint, port) self.model = model self.system_prompt = system_prompt self.temperature = temperature self.max_tokens = max_tokens + self._tokens_param = ( + "max_completion_tokens" if use_max_completion_tokens else "max_tokens" + ) + self._omit_temperature = use_max_completion_tokens self._client = AsyncOpenAI( base_url=f"{self.base_url}/v1", @@ -172,12 +180,14 @@ async def complete(self, prompt: str, **kwargs) -> str: messages.append({"role": "system", "content": self.system_prompt}) messages.append({"role": "user", "content": prompt}) - response = await self._client.chat.completions.create( - model=self.model, - messages=messages, - temperature=kwargs.get("temperature", self.temperature), - max_tokens=kwargs.get("max_tokens", self.max_tokens), - ) + call_kwargs: dict[str, Any] = { + "model": self.model, + "messages": messages, + self._tokens_param: kwargs.get("max_tokens", self.max_tokens), + } + if not self._omit_temperature: + call_kwargs["temperature"] = kwargs.get("temperature", self.temperature) + response = await self._client.chat.completions.create(**call_kwargs) return response.choices[0].message.content or "" async def complete_with_tools( @@ -189,9 +199,10 @@ async def complete_with_tools( create_kwargs: dict[str, Any] = { "model": self.model, "messages": messages, - "temperature": kwargs.get("temperature", self.temperature), - "max_tokens": kwargs.get("max_tokens", self.max_tokens), + self._tokens_param: kwargs.get("max_tokens", self.max_tokens), } + if not self._omit_temperature: + create_kwargs["temperature"] = kwargs.get("temperature", self.temperature) openai_tools = _mcp_tools_to_openai(tools) if openai_tools: create_kwargs["tools"] = openai_tools @@ -315,6 +326,13 @@ async def complete_with_tools( "anthropic": ("https://api.anthropic.com", 443, AnthropicClient), } +# Models that require max_completion_tokens instead of max_tokens and do not +# accept an explicit temperature parameter. Checked by prefix to cover versioned +# names such as "o1-2024-12-17" or "gpt-5-mini-2026-01-15". +_MAX_COMPLETION_TOKENS_PREFIXES: frozenset[str] = frozenset( + {"gpt-5-mini", "o1", "o3", "o4-mini"} +) + def create_llm_client( provider: str, @@ -345,6 +363,11 @@ def create_llm_client( f"Supported: {sorted(_HOSTED_PROVIDERS)}" ) endpoint, port, cls = _HOSTED_PROVIDERS[key] + extra: dict[str, Any] = {} + if cls is OpenAIClient and any( + model.startswith(prefix) for prefix in _MAX_COMPLETION_TOKENS_PREFIXES + ): + extra["use_max_completion_tokens"] = True return cls( endpoint, port, @@ -353,6 +376,7 @@ def create_llm_client( system_prompt=system_prompt, temperature=temperature, max_tokens=max_tokens, + **extra, ) diff --git a/tests/core/test_harness_collect.py b/tests/core/test_harness_collect.py new file mode 100644 index 0000000000..f9ab94e239 --- /dev/null +++ b/tests/core/test_harness_collect.py @@ -0,0 +1,835 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for rollout collection, serialization, and resume.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from openenv.core.env_server.mcp_types import Tool +from openenv.core.harness import ( + HarnessAdapter, + HarnessRolloutResult, + HarnessRunLimits, + ModelStep, + ModelStepResult, + ResourceSession, + ResourceSessionFactory, + ToolResult, + ToolTraceEntry, + VerifyResult, +) +from openenv.core.harness.collect import ( + build_dataset_readme, + build_model_step, + CollectResult, + CollectRunner, + EpisodeRecord, + push_to_hf_hub, + README_FILENAME, + RolloutSerializer, +) +from openenv.core.llm_client import LLMClient, LLMResponse, ToolCall + + +def _fake_rollout( + *, + reward: float = 1.0, + done: bool = True, + turns: int = 1, +) -> HarnessRolloutResult: + messages: list[dict[str, Any]] = [ + {"role": "user", "content": "play"}, + ] + trace: list[ToolTraceEntry] = [] + for i in range(turns): + messages.append({"role": "assistant", "content": f"move {i}"}) + trace.append( + ToolTraceEntry( + tool_name="play_move", + arguments={"cell": i}, + result=ToolResult( + data={"board": "...", "reward": 0.0 if i < turns - 1 else reward}, + done=(i == turns - 1) and done, + metadata={"reward": 0.0 if i < turns - 1 else reward}, + ), + ) + ) + messages.append( + { + "role": "tool", + "name": "play_move", + "content": f"cell={i}", + } + ) + + return HarnessRolloutResult( + messages=messages, + tool_trace=trace, + events=[], + done=done, + metrics={"turns": turns, "tool_calls": turns}, + prompt_ids=[], + completion_ids=[], + logprobs=[], + ) + + +def _fake_verify(reward: float = 1.0) -> VerifyResult: + return VerifyResult( + env_reward=reward, + done=True, + metrics={"step_count": 1}, + artifacts={"final_state": {"winner": "X"}}, + ) + + +class TestEpisodeRecord: + def test_from_rollout_populates_core_fields(self): + rollout = _fake_rollout(reward=1.0) + verify = _fake_verify(reward=1.0) + + record = EpisodeRecord.from_rollout("ep1", rollout, verify) + + assert record.episode_id == "ep1" + assert record.messages == rollout.messages + assert record.reward == 1.0 + assert record.done is True + assert record.metrics == {"turns": 1, "tool_calls": 1} + + def test_from_rollout_flattens_tool_trace(self): + rollout = _fake_rollout(turns=2) + verify = _fake_verify() + + record = EpisodeRecord.from_rollout("ep1", rollout, verify) + + assert len(record.tool_trace) == 2 + first = record.tool_trace[0] + assert first["tool_name"] == "play_move" + assert first["arguments"] == {"cell": 0} + assert "result" in first + assert first["result"]["done"] is False + + def test_from_rollout_uses_resolve_env_reward(self): + """Reward must be derived through _resolve_env_reward to respect the + 'rewards in env' invariant enforced by the harness runtime.""" + rollout = _fake_rollout(reward=1.0) + # Verify carries a mismatching reward — must raise, NOT be silently ignored. + verify = VerifyResult(env_reward=0.0) + + with pytest.raises(ValueError): + EpisodeRecord.from_rollout("ep1", rollout, verify) + + def test_from_rollout_accepts_task_and_extra(self): + rollout = _fake_rollout() + verify = _fake_verify() + + record = EpisodeRecord.from_rollout( + "ep1", + rollout, + verify, + task={"game": "tic_tac_toe", "seed": 42}, + extra={"teacher": "gpt-5-mini"}, + ) + + assert record.task == {"game": "tic_tac_toe", "seed": 42} + assert record.extra == {"teacher": "gpt-5-mini"} + + def test_to_dict_returns_json_safe_payload(self): + record = EpisodeRecord.from_rollout( + "ep1", + _fake_rollout(), + _fake_verify(), + task=Path("task.json"), + extra={"artifact": Path("artifact.txt")}, + ) + + payload = record.to_dict() + + json.dumps(payload) + assert payload["task"] == "task.json" + assert payload["extra"]["artifact"] == "artifact.txt" + + +class TestRolloutSerializer: + def test_writes_one_jsonl_line_per_episode(self, tmp_path: Path): + serializer = RolloutSerializer(tmp_path) + + serializer.write_episode( + EpisodeRecord.from_rollout("ep1", _fake_rollout(), _fake_verify()) + ) + serializer.write_episode( + EpisodeRecord.from_rollout("ep2", _fake_rollout(), _fake_verify()) + ) + + lines = serializer.results_path.read_text().strip().splitlines() + assert [json.loads(line)["episode_id"] for line in lines] == ["ep1", "ep2"] + + def test_line_contains_trl_ready_messages_column(self, tmp_path: Path): + serializer = RolloutSerializer(tmp_path) + rollout = _fake_rollout(turns=2) + + serializer.write_episode( + EpisodeRecord.from_rollout("ep1", rollout, _fake_verify()) + ) + + payload = json.loads(serializer.results_path.read_text().strip()) + assert payload["messages"] == rollout.messages + assert payload["messages"][0]["role"] == "user" + + def test_record_schema_has_required_keys(self, tmp_path: Path): + serializer = RolloutSerializer(tmp_path) + serializer.write_episode( + EpisodeRecord.from_rollout("ep1", _fake_rollout(), _fake_verify()) + ) + + payload = json.loads(serializer.results_path.read_text().strip()) + for key in ( + "episode_id", + "messages", + "reward", + "done", + "tool_trace", + "metrics", + "verify_metrics", + "artifacts", + ): + assert key in payload, f"missing {key}" + + def test_writes_metadata_sidecar(self, tmp_path: Path): + serializer = RolloutSerializer(tmp_path) + serializer.write_metadata( + {"env_id": "tic_tac_toe", "teacher_model": "gpt-5-mini", "num_episodes": 10} + ) + + metadata = json.loads(serializer.metadata_path.read_text()) + assert metadata["env_id"] == "tic_tac_toe" + assert metadata["teacher_model"] == "gpt-5-mini" + + def test_collected_episode_ids_from_prior_run(self, tmp_path: Path): + serializer = RolloutSerializer(tmp_path) + serializer.write_episode( + EpisodeRecord.from_rollout("ep1", _fake_rollout(), _fake_verify()) + ) + serializer.write_episode( + EpisodeRecord.from_rollout("ep2", _fake_rollout(), _fake_verify()) + ) + + reopened = RolloutSerializer(tmp_path) + assert reopened.collected_episode_ids() == {"ep1", "ep2"} + + def test_collected_episode_ids_empty_when_no_file(self, tmp_path: Path): + serializer = RolloutSerializer(tmp_path / "brand-new") + assert serializer.collected_episode_ids() == set() + + def test_creates_output_dir_on_first_write(self, tmp_path: Path): + target = tmp_path / "nested" / "out" + serializer = RolloutSerializer(target) + + serializer.write_episode( + EpisodeRecord.from_rollout("ep1", _fake_rollout(), _fake_verify()) + ) + + assert target.exists() + assert serializer.results_path.exists() + + def test_append_mode_survives_new_serializer_instance(self, tmp_path: Path): + """Serializer must append, not truncate, when a previous run exists.""" + first = RolloutSerializer(tmp_path) + first.write_episode( + EpisodeRecord.from_rollout("ep1", _fake_rollout(), _fake_verify()) + ) + + second = RolloutSerializer(tmp_path) + second.write_episode( + EpisodeRecord.from_rollout("ep2", _fake_rollout(), _fake_verify()) + ) + + lines = (tmp_path / "results.jsonl").read_text().strip().splitlines() + ids = [json.loads(line)["episode_id"] for line in lines] + assert ids == ["ep1", "ep2"] + + +class _FakeSession(ResourceSession): + """Minimal session that returns a fixed reward on verify().""" + + def __init__(self, *, task: Any = None, reward: float = 1.0): + self.task = task + self._reward = reward + self.closed = False + + def initial_messages(self) -> list[dict[str, Any]]: + return [{"role": "user", "content": f"task={self.task}"}] + + def list_tools(self) -> list[Tool]: + return [] + + def call_tool(self, name: str, arguments: dict[str, Any]) -> ToolResult: + return ToolResult( + data={"ack": name}, + done=True, + metadata={"reward": self._reward}, + ) + + def verify( + self, + transcript: list[dict[str, Any]], + final_state: Any | None = None, + ) -> VerifyResult: + return VerifyResult( + env_reward=self._reward, + done=True, + metrics={"step_count": 1}, + artifacts={"task": self.task}, + ) + + def close(self) -> None: + self.closed = True + + +class _FakeFactory(ResourceSessionFactory): + """Records create() calls and hands out deterministic sessions.""" + + def __init__(self, reward_map: dict[Any, float] | None = None): + self.created: list[dict[str, Any]] = [] + self._reward_map = reward_map or {} + self.sessions: list[_FakeSession] = [] + + def create( + self, + task: Any, + seed: int | None = None, + episode_id: str | None = None, + ) -> _FakeSession: + self.created.append({"task": task, "seed": seed, "episode_id": episode_id}) + reward = self._reward_map.get(task, 1.0) if self._reward_map else 1.0 + session = _FakeSession(task=task, reward=reward) + self.sessions.append(session) + return session + + +class _FakeAdapter(HarnessAdapter): + """Produces a scripted rollout keyed off the session's task.""" + + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + fake_session = session # type: ignore[assignment] + reward = fake_session._reward # type: ignore[attr-defined] + messages = list(session.initial_messages()) + messages.append({"role": "assistant", "content": "move"}) + tool_result = session.call_tool("play_move", {"cell": 0}) + trace = [ + ToolTraceEntry( + tool_name="play_move", + arguments={"cell": 0}, + result=tool_result, + ) + ] + messages.append({"role": "tool", "name": "play_move", "content": "ok"}) + return HarnessRolloutResult( + messages=messages, + tool_trace=trace, + events=[], + done=True, + metrics={"turns": 1, "reward": reward}, + ) + + def run_black_box( + self, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + raise NotImplementedError + + +def _noop_model_step( + messages: list[dict[str, Any]], + tools: list[Tool], + sampling: dict[str, Any], +) -> ModelStepResult: + return ModelStepResult(response=LLMResponse(content="ok", tool_calls=[])) + + +class TestCollectRunner: + def test_writes_requested_number_of_episodes(self, tmp_path: Path): + runner = CollectRunner( + session_factory=_FakeFactory(), + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + ) + + result = runner.run(model_step=_noop_model_step, num_episodes=3) + + assert result.num_collected == 3 + assert result.num_skipped == 0 + lines = (tmp_path / "results.jsonl").read_text().strip().splitlines() + assert len(lines) == 3 + + def test_assigns_deterministic_episode_ids(self, tmp_path: Path): + runner = CollectRunner( + session_factory=_FakeFactory(), + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + ) + + result = runner.run( + model_step=_noop_model_step, + num_episodes=2, + episode_id_prefix="ttt", + ) + + assert result.episode_ids == ["ttt-000000", "ttt-000001"] + + def test_resume_skips_already_collected(self, tmp_path: Path): + """Running twice on the same dir with resume=True (default) must + pick up where the last run left off.""" + factory = _FakeFactory() + runner = CollectRunner( + session_factory=factory, + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + ) + first = runner.run(model_step=_noop_model_step, num_episodes=2) + assert first.num_collected == 2 + + # Second runner instance, same output dir, asks for 5 episodes total. + factory2 = _FakeFactory() + runner2 = CollectRunner( + session_factory=factory2, + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + ) + second = runner2.run(model_step=_noop_model_step, num_episodes=5) + + assert second.num_collected == 3 + assert second.num_skipped == 2 + # Only the 3 new episodes should have triggered a factory.create(). + assert len(factory2.created) == 3 + + def test_resume_false_forces_full_rerun(self, tmp_path: Path): + runner = CollectRunner( + session_factory=_FakeFactory(), + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + ) + runner.run(model_step=_noop_model_step, num_episodes=2) + + # Same output dir, resume=False should overwrite and collect 2 again. + with pytest.warns(RuntimeWarning, match="moved existing results file"): + runner.run( + model_step=_noop_model_step, + num_episodes=2, + resume=False, + ) + + lines = (tmp_path / "results.jsonl").read_text().strip().splitlines() + assert len(lines) == 2 + backups = list(tmp_path.glob("results.*.bak.jsonl")) + assert len(backups) == 1 + assert len(backups[0].read_text().strip().splitlines()) == 2 + + def test_closes_session_after_each_episode(self, tmp_path: Path): + factory = _FakeFactory() + runner = CollectRunner( + session_factory=factory, + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + ) + + runner.run(model_step=_noop_model_step, num_episodes=2) + + assert all(session.closed for session in factory.sessions) + + def test_should_keep_filter_drops_episodes(self, tmp_path: Path): + factory = _FakeFactory(reward_map={0: 1.0, 1: 0.0, 2: 1.0}) + runner = CollectRunner( + session_factory=factory, + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + tasks=iter([0, 1, 2]), + ) + + result = runner.run( + model_step=_noop_model_step, + num_episodes=3, + should_keep=lambda record: record.reward > 0.0, + ) + + assert result.num_collected == 2 + assert result.num_dropped == 1 + lines = (tmp_path / "results.jsonl").read_text().strip().splitlines() + assert len(lines) == 2 + + def test_passes_tasks_from_iterable_to_factory(self, tmp_path: Path): + factory = _FakeFactory() + runner = CollectRunner( + session_factory=factory, + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + tasks=iter(["a", "b", "c"]), + ) + + runner.run(model_step=_noop_model_step, num_episodes=3) + + assert [call["task"] for call in factory.created] == ["a", "b", "c"] + + def test_resume_consumes_task_slots_for_skipped_episodes(self, tmp_path: Path): + first_runner = CollectRunner( + session_factory=_FakeFactory(), + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + tasks=iter(["a", "b"]), + ) + first_runner.run(model_step=_noop_model_step, num_episodes=2) + + resumed_factory = _FakeFactory() + resumed_runner = CollectRunner( + session_factory=resumed_factory, + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + tasks=iter(["a", "b", "c", "d"]), + ) + resumed_runner.run(model_step=_noop_model_step, num_episodes=4) + + assert [call["task"] for call in resumed_factory.created] == ["c", "d"] + payloads = [ + json.loads(line) + for line in (tmp_path / "results.jsonl").read_text().strip().splitlines() + ] + assert payloads[2]["episode_id"] == "ep-000002" + assert payloads[2]["task"] == "c" + assert payloads[3]["episode_id"] == "ep-000003" + assert payloads[3]["task"] == "d" + + def test_collect_result_reports_success_rate(self, tmp_path: Path): + factory = _FakeFactory(reward_map={0: 1.0, 1: 0.0, 2: 1.0, 3: 1.0}) + runner = CollectRunner( + session_factory=factory, + harness_adapter=_FakeAdapter(), + serializer=RolloutSerializer(tmp_path), + tasks=iter([0, 1, 2, 3]), + ) + + result = runner.run(model_step=_noop_model_step, num_episodes=4) + + assert isinstance(result, CollectResult) + assert result.num_collected == 4 + assert result.success_rate == pytest.approx(0.75) + assert result.avg_reward == pytest.approx(0.75) + + def test_episode_failure_is_counted_and_collection_continues( + self, + tmp_path: Path, + ): + class FailingAdapter(_FakeAdapter): + def run_white_box( + self, + model_step: ModelStep, + session: ResourceSession, + limits: HarnessRunLimits | None = None, + ) -> HarnessRolloutResult: + if getattr(session, "task") == "bad": + raise RuntimeError("episode failed") + return super().run_white_box(model_step, session, limits) + + factory = _FakeFactory() + runner = CollectRunner( + session_factory=factory, + harness_adapter=FailingAdapter(), + serializer=RolloutSerializer(tmp_path), + tasks=iter(["ok-1", "bad", "ok-2"]), + ) + + result = runner.run(model_step=_noop_model_step, num_episodes=3) + + assert result.num_collected == 2 + assert result.num_failed == 1 + assert all(session.closed for session in factory.sessions) + payloads = [ + json.loads(line) + for line in (tmp_path / "results.jsonl").read_text().strip().splitlines() + ] + assert [payload["task"] for payload in payloads] == ["ok-1", "ok-2"] + + +class _RecordingLLMClient(LLMClient): + """Captures the args passed to complete_with_tools and returns a canned response.""" + + def __init__(self, response: LLMResponse): + super().__init__(endpoint="http://recorder", port=0) + self._response = response + self.calls: list[dict[str, Any]] = [] + + async def complete(self, prompt: str, **kwargs: Any) -> str: # pragma: no cover + raise NotImplementedError + + async def complete_with_tools( + self, + messages: list[dict[str, Any]], + tools: list[dict[str, Any]], + **kwargs: Any, + ) -> LLMResponse: + self.calls.append( + {"messages": messages, "tools": tools, "sampling": dict(kwargs)} + ) + return self._response + + +class TestBuildModelStep: + def _tool(self) -> Tool: + return Tool( + name="play_move", + description="Play a move.", + input_schema={ + "type": "object", + "properties": {"action_id": {"type": "integer"}}, + "required": ["action_id"], + }, + ) + + def test_returns_callable_matching_model_step_protocol(self): + client = _RecordingLLMClient(LLMResponse(content="ok", tool_calls=[])) + step = build_model_step(client) + + result = step([{"role": "user", "content": "hi"}], [], {}) + + assert isinstance(result, ModelStepResult) + assert result.response.content == "ok" + + def test_converts_tools_to_mcp_dict_shape(self): + client = _RecordingLLMClient(LLMResponse(content="", tool_calls=[])) + step = build_model_step(client) + + step([{"role": "user", "content": "go"}], [self._tool()], {}) + + tools = client.calls[0]["tools"] + assert tools[0]["name"] == "play_move" + assert tools[0]["description"] == "Play a move." + assert tools[0]["inputSchema"]["properties"]["action_id"]["type"] == "integer" + + def test_forwards_messages_verbatim(self): + client = _RecordingLLMClient(LLMResponse(content="", tool_calls=[])) + step = build_model_step(client) + + messages = [ + {"role": "user", "content": "play 0"}, + {"role": "assistant", "content": "ok"}, + ] + step(messages, [], {}) + + assert client.calls[0]["messages"] == messages + + def test_prepends_system_prompt_when_missing(self): + client = _RecordingLLMClient(LLMResponse(content="", tool_calls=[])) + step = build_model_step(client, system_prompt="You are a TTT expert.") + + step([{"role": "user", "content": "play"}], [], {}) + + sent = client.calls[0]["messages"] + assert sent[0] == {"role": "system", "content": "You are a TTT expert."} + assert sent[1]["role"] == "user" + + def test_does_not_double_prepend_system_prompt(self): + client = _RecordingLLMClient(LLMResponse(content="", tool_calls=[])) + step = build_model_step(client, system_prompt="A") + + step( + [ + {"role": "system", "content": "pre-existing"}, + {"role": "user", "content": "play"}, + ], + [], + {}, + ) + + sent = client.calls[0]["messages"] + system_roles = [m for m in sent if m["role"] == "system"] + assert len(system_roles) == 1 + assert system_roles[0]["content"] == "pre-existing" + + def test_forwards_sampling_kwargs(self): + client = _RecordingLLMClient(LLMResponse(content="", tool_calls=[])) + step = build_model_step(client) + + with pytest.warns( + RuntimeWarning, + match="Dropping unsupported sampling keys: unknown", + ): + step( + [{"role": "user", "content": "go"}], + [], + {"temperature": 0.7, "max_tokens": 50, "unknown": "drop"}, + ) + + sampling = client.calls[0]["sampling"] + assert sampling["temperature"] == 0.7 + assert sampling["max_tokens"] == 50 + # Unknown keys should NOT leak to the client (provider SDKs reject them). + assert "unknown" not in sampling + + def test_passes_through_tool_calls(self): + response = LLMResponse( + content="Playing cell 0.", + tool_calls=[ToolCall(id="abc", name="play_move", args={"action_id": 0})], + ) + client = _RecordingLLMClient(response) + step = build_model_step(client) + + result = step([{"role": "user", "content": "play"}], [self._tool()], {}) + + assert len(result.response.tool_calls) == 1 + assert result.response.tool_calls[0].name == "play_move" + assert result.response.tool_calls[0].args == {"action_id": 0} + + @pytest.mark.asyncio + async def test_runs_safely_inside_existing_event_loop(self): + client = _RecordingLLMClient(LLMResponse(content="ok", tool_calls=[])) + step = build_model_step(client) + + result = step([{"role": "user", "content": "play"}], [], {}) + + assert result.response.content == "ok" + assert client.calls[0]["messages"] == [{"role": "user", "content": "play"}] + + +def _populated_output_dir(tmp_path: Path) -> Path: + """Write a minimal results.jsonl + metadata.json into tmp_path.""" + serializer = RolloutSerializer(tmp_path) + serializer.write_metadata({"env_id": "tic_tac_toe", "num_episodes": 1}) + serializer.write_episode( + EpisodeRecord.from_rollout("ep1", _fake_rollout(), _fake_verify()) + ) + return tmp_path + + +class TestPushToHfHub: + def test_creates_dataset_repo_and_uploads(self, tmp_path: Path): + output = _populated_output_dir(tmp_path) + mock_api = MagicMock() + + with patch("huggingface_hub.HfApi", return_value=mock_api): + url = push_to_hf_hub(output, "user/ttt-sft-v1") + + mock_api.create_repo.assert_called_once() + create_kwargs = mock_api.create_repo.call_args.kwargs + assert create_kwargs["repo_id"] == "user/ttt-sft-v1" + assert create_kwargs["repo_type"] == "dataset" + assert create_kwargs["exist_ok"] is True + + mock_api.upload_folder.assert_called_once() + upload_kwargs = mock_api.upload_folder.call_args.kwargs + assert upload_kwargs["repo_id"] == "user/ttt-sft-v1" + assert upload_kwargs["repo_type"] == "dataset" + assert Path(upload_kwargs["folder_path"]) == output + + assert url == "https://huggingface.co/datasets/user/ttt-sft-v1" + + def test_private_flag_propagates(self, tmp_path: Path): + output = _populated_output_dir(tmp_path) + mock_api = MagicMock() + + with patch("huggingface_hub.HfApi", return_value=mock_api): + push_to_hf_hub(output, "user/private-ds", private=True) + + assert mock_api.create_repo.call_args.kwargs["private"] is True + + def test_custom_commit_message(self, tmp_path: Path): + output = _populated_output_dir(tmp_path) + mock_api = MagicMock() + + with patch("huggingface_hub.HfApi", return_value=mock_api): + push_to_hf_hub(output, "user/ds", commit_message="seed: 200 TTT games") + + assert ( + mock_api.upload_folder.call_args.kwargs["commit_message"] + == "seed: 200 TTT games" + ) + + def test_default_commit_message_mentions_episodes(self, tmp_path: Path): + output = _populated_output_dir(tmp_path) + mock_api = MagicMock() + + with patch("huggingface_hub.HfApi", return_value=mock_api): + push_to_hf_hub(output, "user/ds") + + msg = mock_api.upload_folder.call_args.kwargs["commit_message"] + # Default message includes episode count so the Hub history is informative. + assert "1" in msg and "episode" in msg.lower() + + def test_errors_when_results_jsonl_missing(self, tmp_path: Path): + # Empty dir — no results.jsonl written. + with pytest.raises(FileNotFoundError): + push_to_hf_hub(tmp_path, "user/ds") + + def test_token_forwarded_to_hfapi(self, tmp_path: Path): + output = _populated_output_dir(tmp_path) + mock_api = MagicMock() + api_cls = MagicMock(return_value=mock_api) + + with patch("huggingface_hub.HfApi", api_cls): + push_to_hf_hub(output, "user/ds", token="hf_xxx") + + assert api_cls.call_args.kwargs["token"] == "hf_xxx" + + def test_writes_readme_with_viewer_frontmatter(self, tmp_path: Path): + """HF Dataset Viewer needs YAML front-matter to know which file is the + dataset — otherwise it tries to cast ``metadata.json`` into the same + schema as ``results.jsonl`` and errors out.""" + output = _populated_output_dir(tmp_path) + mock_api = MagicMock() + + with patch("huggingface_hub.HfApi", return_value=mock_api): + push_to_hf_hub(output, "user/ds") + + readme = (output / README_FILENAME).read_text(encoding="utf-8") + assert readme.startswith("---") + assert "configs:" in readme + assert 'path: "results.jsonl"' in readme + assert "split: train" in readme + + +class TestBuildDatasetReadme: + def test_includes_episode_count(self, tmp_path: Path): + _populated_output_dir(tmp_path) + + readme = build_dataset_readme(tmp_path) + + assert "Episodes:** 1" in readme + + def test_includes_metadata_table(self, tmp_path: Path): + _populated_output_dir(tmp_path) + # overwrite metadata with a known payload + (tmp_path / "metadata.json").write_text( + json.dumps({"provider": "openai", "model": "gpt-5-mini"}) + ) + + readme = build_dataset_readme(tmp_path) + + assert "`provider`" in readme + assert "`openai`" in readme + assert "`model`" in readme + assert "`gpt-5-mini`" in readme + + def test_resilient_to_missing_metadata(self, tmp_path: Path): + serializer = RolloutSerializer(tmp_path) + serializer.write_episode( + EpisodeRecord.from_rollout("ep1", _fake_rollout(), _fake_verify()) + ) + # No metadata.json written. + + readme = build_dataset_readme(tmp_path) + + assert "OpenEnv rollouts" in readme + assert "Schema" in readme diff --git a/tests/core/test_harness_runtime.py b/tests/core/test_harness_runtime.py new file mode 100644 index 0000000000..388f708456 --- /dev/null +++ b/tests/core/test_harness_runtime.py @@ -0,0 +1,702 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for the harness/session runtime layer.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any + +import pytest +from openenv.core.client_types import StepResult +from openenv.core.env_server.mcp_types import Tool +from openenv.core.env_server.types import State +from openenv.core.harness import ( + build_harness_rollout_func, + CLIHarnessAdapter, + HarnessRolloutResult, + HarnessRunLimits, + MCPHarnessAdapter, + ModelStepResult, + RESERVED_TOOL_NAMES, + RolloutEvent, + SessionMCPBridge, + StepEnvSessionAdapter, + ToolResult, + VerifyResult, +) +from openenv.core.llm_client import LLMResponse, ToolCall + + +class FakeSyncEnv: + """Simple sync environment used to test session adapters.""" + + def __init__(self): + self.closed = False + self.step_count = 0 + self.reset_calls: list[dict[str, Any]] = [] + self.step_calls: list[str] = [] + + def reset(self, **kwargs: Any) -> StepResult[dict[str, Any]]: + self.reset_calls.append(kwargs) + self.step_count = 0 + return StepResult( + observation={"text": "Ready", "task": kwargs.get("task")}, + reward=0.0, + done=False, + ) + + def step(self, action: str) -> StepResult[dict[str, Any]]: + self.step_calls.append(action) + self.step_count += 1 + done = action == "finish" + reward = 1.0 if done else 0.25 + return StepResult( + observation={ + "text": f"ran:{action}", + "done": done, + "reward": reward, + }, + reward=reward, + done=done, + ) + + def state(self) -> State: + return State(episode_id="fake-episode", step_count=self.step_count) + + def close(self) -> None: + self.closed = True + + +@dataclass +class RecordingSessionFactory: + """Factory used to test rollout helper ordering.""" + + session: Any + created_tasks: list[Any] = field(default_factory=list) + + def create( + self, + task: Any, + seed: int | None = None, + episode_id: str | None = None, + ) -> Any: + self.created_tasks.append((task, seed, episode_id)) + return self.session + + +class TestStepEnvSessionAdapter: + """Tests for the generic step-based resource session adapter.""" + + def test_session_exposes_tools_and_lifecycle(self): + env = FakeSyncEnv() + tools = [ + Tool( + name="advance", + description="Advance the fake environment", + input_schema={ + "type": "object", + "properties": { + "value": {"type": "string"}, + }, + "required": ["value"], + }, + ), + ] + + session = StepEnvSessionAdapter( + client=env, + task="test-task", + tool_specs=tools, + action_builder=lambda name, arguments: arguments["value"], + initial_messages_builder=lambda result, task: [ + {"role": "user", "content": f"{task}:{result.observation['text']}"} + ], + ) + + assert session.initial_messages() == [ + {"role": "user", "content": "test-task:Ready"} + ] + assert session.list_tools() == tools + + tool_result = session.call_tool("advance", {"value": "next"}) + + assert tool_result.data["observation"]["text"] == "ran:next" + assert tool_result.done is False + assert tool_result.metadata["reward"] == 0.25 + + verify_result = session.verify( + transcript=[{"role": "assistant", "content": ""}] + ) + assert verify_result.env_reward == 0.25 + assert verify_result.done is False + assert verify_result.metrics["step_count"] == 1 + + session.close() + assert env.closed is True + + def test_session_uses_custom_verify_builder(self): + env = FakeSyncEnv() + seen: dict[str, Any] = {} + + session = StepEnvSessionAdapter( + client=env, + task="verify-me", + tool_specs=[], + action_builder=lambda name, arguments: None, + initial_messages_builder=lambda result, task: [], + verify_builder=lambda transcript, final_state, last_result, state: ( + seen.update( + { + "transcript": transcript, + "final_state": final_state, + "last_result": last_result, + "state": state, + } + ) + or VerifyResult(env_reward=9.0, done=True, metrics={"kind": "custom"}) + ), + ) + + result = session.verify( + transcript=[{"role": "assistant", "content": "done"}], + final_state={"terminal": True}, + ) + + assert result.env_reward == 9.0 + assert result.done is True + assert result.metrics == {"kind": "custom"} + assert seen["final_state"] == {"terminal": True} + + def test_session_rejects_reserved_tool_names(self): + env = FakeSyncEnv() + + with pytest.raises(ValueError, match="reserved"): + StepEnvSessionAdapter( + client=env, + task="reserved-tool", + tool_specs=[ + Tool( + name="reset", + description="Must remain an orchestration control", + input_schema={"type": "object", "properties": {}}, + ) + ], + action_builder=lambda name, arguments: name, + initial_messages_builder=lambda result, task: [], + ) + + def test_session_closes_client_when_reset_fails(self): + class FailingResetEnv(FakeSyncEnv): + def reset(self, **kwargs: Any) -> StepResult[dict[str, Any]]: + raise RuntimeError("reset failed") + + env = FailingResetEnv() + + with pytest.raises(RuntimeError, match="reset failed"): + StepEnvSessionAdapter( + client=env, + task="bad-reset", + tool_specs=[], + action_builder=lambda name, arguments: None, + initial_messages_builder=lambda result, task: [], + ) + + assert env.closed is True + + +class TestSessionMCPBridge: + """Tests for exposing sessions through an MCP-style JSON-RPC bridge.""" + + def test_tools_list_and_tools_call(self): + env = FakeSyncEnv() + session = StepEnvSessionAdapter( + client=env, + task="bridge-task", + tool_specs=[ + Tool( + name="finish", + description="Finish the fake task", + input_schema={ + "type": "object", + "properties": {}, + }, + ) + ], + action_builder=lambda name, arguments: name, + initial_messages_builder=lambda result, task: [], + ) + bridge = SessionMCPBridge(session) + + list_response = bridge.handle_request( + {"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}} + ) + assert list_response["result"]["tools"][0]["name"] == "finish" + + call_response = bridge.handle_request( + { + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "finish", "arguments": {}}, + } + ) + assert call_response["result"]["data"]["reward"] == 1.0 + assert call_response["result"]["done"] is True + + def test_unknown_method_returns_jsonrpc_error(self): + env = FakeSyncEnv() + session = StepEnvSessionAdapter( + client=env, + task="bridge-task", + tool_specs=[], + action_builder=lambda name, arguments: None, + initial_messages_builder=lambda result, task: [], + ) + bridge = SessionMCPBridge(session) + + response = bridge.handle_request( + {"jsonrpc": "2.0", "id": 7, "method": "resources/list", "params": {}} + ) + + assert response["error"]["message"] == "Method not found" + assert response["id"] == 7 + + def test_unknown_tool_returns_jsonrpc_error(self): + env = FakeSyncEnv() + session = StepEnvSessionAdapter( + client=env, + task="bridge-task", + tool_specs=[], + action_builder=lambda name, arguments: None, + initial_messages_builder=lambda result, task: [], + ) + bridge = SessionMCPBridge(session) + + response = bridge.handle_request( + { + "jsonrpc": "2.0", + "id": 9, + "method": "tools/call", + "params": {"name": "missing", "arguments": {}}, + } + ) + + assert response["error"]["message"] == "Unknown tool: missing" + assert response["id"] == 9 + + def test_reserved_tool_returns_jsonrpc_error(self): + class ReservedSession: + def list_tools(self): + return [] + + def call_tool(self, name, arguments): + raise AssertionError("reserved names must be blocked by the bridge") + + bridge = SessionMCPBridge(ReservedSession()) + + for name in RESERVED_TOOL_NAMES: + response = bridge.handle_request( + { + "jsonrpc": "2.0", + "id": 12, + "method": "tools/call", + "params": {"name": name, "arguments": {}}, + } + ) + + assert ( + response["error"]["message"] == f"Reserved orchestration tool: {name}" + ) + assert response["id"] == 12 + + def test_value_error_returns_invalid_params_error(self): + class ValueErrorSession: + def list_tools(self): + return [] + + def call_tool(self, name, arguments): + raise ValueError("bad arguments") + + bridge = SessionMCPBridge(ValueErrorSession()) + + response = bridge.handle_request( + { + "jsonrpc": "2.0", + "id": 10, + "method": "tools/call", + "params": {"name": "finish", "arguments": {}}, + } + ) + + assert response["error"]["message"] == "bad arguments" + assert response["id"] == 10 + + def test_unexpected_tool_error_returns_internal_error(self): + class FailingSession: + def list_tools(self): + return [] + + def call_tool(self, name, arguments): + raise RuntimeError("bridge exploded") + + bridge = SessionMCPBridge(FailingSession()) + + response = bridge.handle_request( + { + "jsonrpc": "2.0", + "id": 11, + "method": "tools/call", + "params": {"name": "finish", "arguments": {}}, + } + ) + + assert response["error"]["message"] == "bridge exploded" + assert response["id"] == 11 + + +class TestMCPHarnessAdapter: + """Tests for the white-box MCP-first harness adapter.""" + + def test_run_white_box_accumulates_model_tokens_and_tool_trace(self): + env = FakeSyncEnv() + session = StepEnvSessionAdapter( + client=env, + task="white-box-task", + tool_specs=[ + Tool( + name="finish", + description="Finish the task", + input_schema={"type": "object", "properties": {}}, + ) + ], + action_builder=lambda name, arguments: name, + initial_messages_builder=lambda result, task: [ + {"role": "user", "content": "Solve the task"} + ], + ) + + adapter = MCPHarnessAdapter() + + model_calls: list[tuple[list[dict[str, Any]], list[Tool], dict[str, Any]]] = [] + + def model_step(messages, tools, sampling): + model_calls.append((messages, tools, sampling)) + return ModelStepResult( + response=LLMResponse( + content="Calling finish", + tool_calls=[ToolCall(id="tool-1", name="finish", args={})], + ), + prompt_ids=[11, 12], + completion_ids=[21, 22], + logprobs=[-0.1, -0.2], + ) + + result = adapter.run_white_box( + model_step=model_step, + session=session, + limits=HarnessRunLimits(max_turns=2), + ) + + assert result.done is True + assert result.prompt_ids == [11, 12] + assert result.completion_ids == [21, 22] + assert result.logprobs == [-0.1, -0.2] + assert len(result.tool_trace) == 1 + assert result.tool_trace[0].tool_name == "finish" + assert result.tool_trace[0].result.done is True + assert model_calls[0][2] == {} + + def test_run_white_box_does_not_verify_session(self): + class VerifyingSession: + def __init__(self): + self.verified = False + + def initial_messages(self): + return [{"role": "user", "content": "task"}] + + def list_tools(self): + return [] + + def call_tool(self, name, arguments): + raise AssertionError("No tools expected") + + def verify(self, transcript, final_state=None): + self.verified = True + return VerifyResult(env_reward=1.0) + + def close(self): + pass + + session = VerifyingSession() + adapter = MCPHarnessAdapter() + + adapter.run_white_box( + model_step=lambda messages, tools, sampling: ModelStepResult( + response=LLMResponse(content="Final answer", tool_calls=[]), + prompt_ids=[1], + completion_ids=[2], + logprobs=[-0.5], + ), + session=session, + limits=HarnessRunLimits(max_turns=1), + ) + + assert session.verified is False + + +class TestBuildHarnessRolloutFunc: + """Tests for the TRL rollout helper built on sessions/harnesses.""" + + def test_rollout_func_verifies_after_harness_completes(self): + order: list[str] = [] + + class RecordingSession: + def __init__(self): + self.closed = False + + def initial_messages(self): + return [{"role": "user", "content": "task"}] + + def list_tools(self): + return [ + Tool( + name="finish", + description="Finish the task", + input_schema={"type": "object", "properties": {}}, + ) + ] + + def call_tool(self, name, arguments): + order.append("tool") + return ToolResult(data={"reward": 1.0}, done=True) + + def verify(self, transcript, final_state=None): + order.append("verify") + return VerifyResult( + env_reward=1.0, + done=True, + metrics={"verified": True}, + ) + + def close(self): + order.append("close") + self.closed = True + + session = RecordingSession() + factory = RecordingSessionFactory(session=session) + adapter = MCPHarnessAdapter() + + def model_step_builder(trainer, session): + return lambda messages, tools, sampling: ( + order.append("model") + or ModelStepResult( + response=LLMResponse( + content="calling tool", + tool_calls=[ToolCall(id="call-1", name="finish", args={})], + ), + prompt_ids=[101], + completion_ids=[202], + logprobs=[-0.25], + ) + ) + + rollout_func = build_harness_rollout_func( + session_factory=factory, + harness_adapter=adapter, + model_step_builder=model_step_builder, + limits=HarnessRunLimits(max_turns=2), + ) + + result = rollout_func(["task-a"], trainer=object()) + + assert result["prompt_ids"] == [[101]] + assert result["completion_ids"] == [[202]] + assert result["logprobs"] == [[-0.25]] + assert result["env_reward"] == [1.0] + assert result["verify_metrics"] == [{"verified": True}] + assert factory.created_tasks == [("task-a", None, None)] + assert order == ["model", "tool", "verify", "close"] + + def test_rollout_func_rejects_verify_reward_mismatch(self): + class RecordingSession: + def initial_messages(self): + return [{"role": "user", "content": "task"}] + + def list_tools(self): + return [ + Tool( + name="finish", + description="Finish the task", + input_schema={"type": "object", "properties": {}}, + ) + ] + + def call_tool(self, name, arguments): + return ToolResult( + data={"reward": 1.0}, + done=True, + metadata={"reward": 1.0}, + ) + + def verify(self, transcript, final_state=None): + return VerifyResult(env_reward=0.5, done=True) + + def close(self): + pass + + factory = RecordingSessionFactory(session=RecordingSession()) + rollout_func = build_harness_rollout_func( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + model_step_builder=lambda trainer, session: ( + lambda messages, tools, sampling: ModelStepResult( + response=LLMResponse( + content="calling tool", + tool_calls=[ToolCall(id="call-1", name="finish", args={})], + ), + ) + ), + limits=HarnessRunLimits(max_turns=1), + ) + + with pytest.raises( + ValueError, + match="verify.env_reward must forward the environment reward", + ): + rollout_func(["task-a"], trainer=object()) + + def test_rollout_func_accepts_float_rounding_in_reward_check(self): + class RecordingSession: + def initial_messages(self): + return [{"role": "user", "content": "task"}] + + def list_tools(self): + return [ + Tool( + name="finish", + description="Finish the task", + input_schema={"type": "object", "properties": {}}, + ) + ] + + def call_tool(self, name, arguments): + return ToolResult( + data={"reward": 0.3}, + done=True, + metadata={"reward": 0.3}, + ) + + def verify(self, transcript, final_state=None): + return VerifyResult(env_reward=0.1 + 0.2, done=True) + + def close(self): + pass + + rollout_func = build_harness_rollout_func( + session_factory=RecordingSessionFactory(session=RecordingSession()), + harness_adapter=MCPHarnessAdapter(), + model_step_builder=lambda trainer, session: ( + lambda messages, tools, sampling: ModelStepResult( + response=LLMResponse( + content="calling tool", + tool_calls=[ToolCall(id="call-1", name="finish", args={})], + ), + ) + ), + limits=HarnessRunLimits(max_turns=1), + ) + + result = rollout_func(["task-a"], trainer=object()) + + assert result["env_reward"] == [0.3] + + def test_rollout_func_rejects_missing_reward(self): + class RecordingSession: + def initial_messages(self): + return [{"role": "user", "content": "task"}] + + def list_tools(self): + return [] + + def call_tool(self, name, arguments): + raise AssertionError("No tools expected") + + def verify(self, transcript, final_state=None): + return VerifyResult(env_reward=None, done=True) + + def close(self): + pass + + rollout_func = build_harness_rollout_func( + session_factory=RecordingSessionFactory(session=RecordingSession()), + harness_adapter=MCPHarnessAdapter(), + model_step_builder=lambda trainer, session: ( + lambda messages, tools, sampling: ModelStepResult( + response=LLMResponse(content="done", tool_calls=[]), + ) + ), + limits=HarnessRunLimits(max_turns=1), + ) + + with pytest.raises(ValueError, match="did not produce an environment reward"): + rollout_func(["task-a"], trainer=object()) + + +class TestCLIHarnessAdapter: + """Tests for black-box CLI-style evaluation harnesses.""" + + def test_black_box_runner_can_use_session_mcp_bridge(self): + env = FakeSyncEnv() + session = StepEnvSessionAdapter( + client=env, + task="cli-task", + tool_specs=[ + Tool( + name="finish", + description="Finish the task", + input_schema={"type": "object", "properties": {}}, + ) + ], + action_builder=lambda name, arguments: name, + initial_messages_builder=lambda result, task: [], + ) + + def runner(bridge, current_session, limits): + list_response = bridge.handle_request( + {"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}} + ) + call_response = bridge.handle_request( + { + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "finish", "arguments": {}}, + } + ) + return HarnessRolloutResult( + messages=[{"role": "assistant", "content": "done"}], + tool_trace=[], + events=[ + RolloutEvent( + type="cli_runner", + payload={ + "tools": list_response["result"]["tools"], + "result": call_response["result"], + }, + ) + ], + done=True, + metrics={"mode": "black_box"}, + ) + + adapter = CLIHarnessAdapter(runner=runner) + result = adapter.run_black_box(session=session, limits=HarnessRunLimits()) + + assert result.done is True + assert result.metrics["mode"] == "black_box" + assert result.events[0].payload["result"]["done"] is True diff --git a/tests/envs/test_browsergym_harness.py b/tests/envs/test_browsergym_harness.py new file mode 100644 index 0000000000..5d66247fad --- /dev/null +++ b/tests/envs/test_browsergym_harness.py @@ -0,0 +1,198 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for BrowserGym's harness-oriented session adapter.""" + +from __future__ import annotations + +from typing import Any + +import pytest +from browsergym_env import BrowserGymAction, BrowserGymObservation, BrowserGymState +from browsergym_env.harness import ( + BrowserGymSessionFactory, + build_browsergym_action_str, + build_browsergym_action_tool_call, +) +from openenv.core.client_types import StepResult +from openenv.core.harness import ( + build_harness_rollout_func, + HarnessRunLimits, + MCPHarnessAdapter, + ModelStepResult, + ResourceSessionFactory, +) +from openenv.core.llm_client import LLMResponse + + +class FakeBrowserGymClient: + """Small BrowserGym-like client used for harness tests.""" + + def __init__(self): + self.closed = False + self._step_count = 0 + self._cum_reward = 0.0 + self.step_actions: list[str] = [] + + def reset(self, **kwargs: Any) -> StepResult[BrowserGymObservation]: + self._step_count = 0 + self._cum_reward = 0.0 + return StepResult( + observation=BrowserGymObservation( + goal="Click the highlighted button", + axtree_txt="[13] button 'Continue'", + text="[13] button 'Continue'", + url="http://example.test", + done=False, + reward=0.0, + ), + reward=0.0, + done=False, + ) + + def step(self, action: BrowserGymAction) -> StepResult[BrowserGymObservation]: + self.step_actions.append(action.action_str) + self._step_count += 1 + done = action.action_str == 'click("13")' + reward = 1.0 if done else 0.0 + self._cum_reward += reward + return StepResult( + observation=BrowserGymObservation( + goal="Click the highlighted button", + axtree_txt="[13] button 'Continue'", + text="[13] button 'Continue'", + url="http://example.test/after", + done=done, + reward=reward, + last_action_error=False, + error="", + ), + reward=reward, + done=done, + ) + + def state(self) -> BrowserGymState: + return BrowserGymState( + episode_id="browsergym-episode", + step_count=self._step_count, + benchmark="miniwob", + task_name="click-test", + goal="Click the highlighted button", + current_url="http://example.test/after", + cum_reward=self._cum_reward, + ) + + def close(self) -> None: + self.closed = True + + +def test_browsergym_session_factory_exposes_expected_tools(): + factory = BrowserGymSessionFactory(client_factory=FakeBrowserGymClient) + assert isinstance(factory, ResourceSessionFactory) + + session = factory.create(task="ignored-task") + + tool_names = [tool.name for tool in session.list_tools()] + assert tool_names == ["click", "fill", "send_keys", "scroll", "noop"] + + messages = session.initial_messages() + assert len(messages) == 1 + assert "Click the highlighted button" in messages[0]["content"] + assert "[13] button 'Continue'" in messages[0]["content"] + + session.close() + + +def test_browsergym_tool_calls_match_plain_env_rewards(): + plain_client = FakeBrowserGymClient() + plain_client.reset() + plain_result = plain_client.step(BrowserGymAction(action_str='click("13")')) + + factory = BrowserGymSessionFactory(client_factory=FakeBrowserGymClient) + session = factory.create(task="ignored-task") + tool_result = session.call_tool("click", {"bid": "13"}) + verify_result = session.verify(transcript=[{"role": "assistant", "content": ""}]) + + assert tool_result.done is True + assert tool_result.metadata["reward"] == plain_result.reward + assert verify_result.env_reward == plain_result.reward + assert verify_result.done == plain_result.done + session.close() + + +def test_browsergym_action_text_parser(): + click = build_browsergym_action_tool_call("click('13')") + fill = build_browsergym_action_tool_call("fill('42', 'hello')") + scroll = build_browsergym_action_tool_call("scroll('down')") + noop = build_browsergym_action_tool_call("noop()") + + assert click.name == "click" + assert click.args == {"bid": "13"} + assert fill.name == "fill" + assert fill.args == {"bid": "42", "text": "hello"} + assert scroll.name == "scroll" + assert scroll.args == {"direction": "down"} + assert noop.name == "noop" + assert noop.args == {} + + +def test_browsergym_fill_parser_supports_embedded_quote(): + fill = build_browsergym_action_tool_call('fill("42", "O\'Brien")') + + assert fill.name == "fill" + assert fill.args == {"bid": "42", "text": "O'Brien"} + + +def test_browsergym_fill_parser_supports_single_quoted_embedded_quote(): + fill = build_browsergym_action_tool_call("fill('42', 'O\\'Brien')") + + assert fill.name == "fill" + assert fill.args == {"bid": "42", "text": "O'Brien"} + + +def test_browsergym_action_builder_quotes_round_trip_through_parser(): + action = build_browsergym_action_str( + "fill", + {"bid": 'button"42', "text": "O'Brien\nNext"}, + ) + tool_call = build_browsergym_action_tool_call(action) + + assert action == 'fill("button\\"42", "O\'Brien\\nNext")' + assert tool_call.name == "fill" + assert tool_call.args == {"bid": 'button"42', "text": "O'Brien\nNext"} + + +def test_browsergym_action_parser_rejects_malformed_strings(): + with pytest.raises(ValueError, match="Unsupported BrowserGym action"): + build_browsergym_action_tool_call("fill('42', 'O'Brien')") + + +def test_browsergym_session_factory_works_with_generic_rollout_helper(): + factory = BrowserGymSessionFactory(client_factory=FakeBrowserGymClient) + adapter = MCPHarnessAdapter() + + def model_step_builder(trainer, session): + tool_call = build_browsergym_action_tool_call("click('13')") + return lambda messages, tools, sampling: ModelStepResult( + response=LLMResponse(content="click", tool_calls=[tool_call]), + prompt_ids=[3, 4], + completion_ids=[5, 6], + logprobs=[-0.3, -0.4], + ) + + rollout_func = build_harness_rollout_func( + session_factory=factory, + harness_adapter=adapter, + model_step_builder=model_step_builder, + limits=HarnessRunLimits(max_turns=3), + ) + + result = rollout_func(["browsergym prompt"], trainer=object()) + + assert result["prompt_ids"] == [[3, 4]] + assert result["completion_ids"] == [[5, 6]] + assert result["logprobs"] == [[-0.3, -0.4]] + assert result["env_reward"] == [1.0] diff --git a/tests/envs/test_openspiel_collect.py b/tests/envs/test_openspiel_collect.py new file mode 100644 index 0000000000..099210a494 --- /dev/null +++ b/tests/envs/test_openspiel_collect.py @@ -0,0 +1,197 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""End-to-end smoke test for the rollout collect pipeline on OpenSpiel TTT. + +Wires together the real components that a teacher-rollout job would use: + +- ``OpenSpielSessionFactory`` (the env-side adapter for ``tic_tac_toe``) +- ``MCPHarnessAdapter`` (the white-box ReAct harness from #471) +- ``CollectRunner`` + ``RolloutSerializer`` (this PR) + +The only fake here is the OpenSpiel *client* itself — a deterministic +scripted game — so we can verify the full pipeline without Docker. +""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +from openenv.core.client_types import StepResult +from openenv.core.harness import HarnessRunLimits, MCPHarnessAdapter, ModelStepResult +from openenv.core.harness.collect import CollectRunner, RolloutSerializer +from openenv.core.llm_client import LLMResponse, ToolCall +from openspiel_env.harness import OpenSpielSessionFactory +from openspiel_env.models import OpenSpielAction, OpenSpielObservation, OpenSpielState + + +class ScriptedOpenSpielClient: + """A tic_tac_toe client that wins after ``moves_until_win`` plays.""" + + def __init__(self, moves_until_win: int = 1): + self._moves_until_win = moves_until_win + self._played = 0 + self.closed = False + + def reset(self, **kwargs: Any) -> StepResult[OpenSpielObservation]: + self._played = 0 + empty_board = [1.0] * 9 + [0.0] * 18 + return StepResult( + observation=OpenSpielObservation( + info_state=empty_board, + legal_actions=list(range(9)), + current_player_id=0, + game_phase="playing", + ), + reward=0.0, + done=False, + ) + + def step(self, action: OpenSpielAction) -> StepResult[OpenSpielObservation]: + self._played += 1 + done = self._played >= self._moves_until_win + reward = 1.0 if done else 0.0 + info_state = [0.0] * 27 + # mark cell 0 as X just to have a non-empty board + info_state[9] = 1.0 + for i in range(1, 9): + info_state[i] = 1.0 + return StepResult( + observation=OpenSpielObservation( + info_state=info_state, + legal_actions=[] if done else list(range(1, 9)), + current_player_id=0, + game_phase="terminal" if done else "playing", + ), + reward=reward, + done=done, + ) + + def state(self) -> OpenSpielState: + return OpenSpielState(step_count=self._played, game_name="tic_tac_toe") + + def close(self) -> None: + self.closed = True + + +def _teacher_model_step(messages, tools, sampling): + """Deterministic teacher: always call play_move with the first legal action. + + Stands in for a hosted LLM (e.g. gpt-5-mini). Real teachers will parse + the prompt to pick an action_id. + """ + del messages, tools, sampling + return ModelStepResult( + response=LLMResponse( + content="Playing the first legal cell.", + tool_calls=[ + ToolCall( + id="teacher-0", + name="play_move", + args={"action_id": 0}, + ), + ], + ), + ) + + +def test_collect_writes_trl_ready_jsonl_dataset(tmp_path: Path): + factory = OpenSpielSessionFactory( + lambda: ScriptedOpenSpielClient(moves_until_win=1), + game_name="tic_tac_toe", + ) + runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=RolloutSerializer(tmp_path), + limits=HarnessRunLimits(max_turns=4), + ) + + result = runner.run( + model_step=_teacher_model_step, + num_episodes=5, + episode_id_prefix="ttt", + ) + + assert result.num_collected == 5 + assert result.num_skipped == 0 + assert result.success_rate == 1.0 + + lines = (tmp_path / "results.jsonl").read_text().strip().splitlines() + assert len(lines) == 5 + first = json.loads(lines[0]) + + # Schema sanity — enough to feed TRL SFTTrainer directly. + assert first["episode_id"] == "ttt-000000" + assert first["reward"] == 1.0 + assert first["done"] is True + assert any(m["role"] == "user" for m in first["messages"]) + assert any(m["role"] == "assistant" for m in first["messages"]) + assert any(m["role"] == "tool" for m in first["messages"]) + assert first["tool_trace"][0]["tool_name"] == "play_move" + assert first["tool_trace"][0]["arguments"] == {"action_id": 0} + + +def test_collect_resume_skips_prior_episodes(tmp_path: Path): + factory = OpenSpielSessionFactory( + lambda: ScriptedOpenSpielClient(moves_until_win=1), + ) + serializer = RolloutSerializer(tmp_path) + + first_runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=serializer, + limits=HarnessRunLimits(max_turns=4), + ) + first = first_runner.run(model_step=_teacher_model_step, num_episodes=2) + assert first.num_collected == 2 + + second_runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=RolloutSerializer(tmp_path), + limits=HarnessRunLimits(max_turns=4), + ) + second = second_runner.run(model_step=_teacher_model_step, num_episodes=4) + + assert second.num_collected == 2 + assert second.num_skipped == 2 + assert len((tmp_path / "results.jsonl").read_text().strip().splitlines()) == 4 + + +def test_collect_should_keep_filters_losing_rollouts(tmp_path: Path): + # Alternate winning/losing clients on each create(). + wins = [True, False, True, True] + client_iter = iter(wins) + + def client_factory(): + winning = next(client_iter) + return ScriptedOpenSpielClient(moves_until_win=1 if winning else 10) + + factory = OpenSpielSessionFactory(client_factory) + runner = CollectRunner( + session_factory=factory, + harness_adapter=MCPHarnessAdapter(), + serializer=RolloutSerializer(tmp_path), + # 3 turns is enough for the winning cases but not the losing one. + limits=HarnessRunLimits(max_turns=3), + ) + + result = runner.run( + model_step=_teacher_model_step, + num_episodes=4, + should_keep=lambda record: record.reward > 0.0, + ) + + assert result.num_collected == 3 + assert result.num_dropped == 1 + lines = (tmp_path / "results.jsonl").read_text().strip().splitlines() + assert len(lines) == 3 + for line in lines: + assert json.loads(line)["reward"] == 1.0 diff --git a/tests/envs/test_openspiel_harness.py b/tests/envs/test_openspiel_harness.py new file mode 100644 index 0000000000..b052cf61da --- /dev/null +++ b/tests/envs/test_openspiel_harness.py @@ -0,0 +1,140 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for the OpenSpiel harness session adapter.""" + +from __future__ import annotations + +from typing import Any + +from openenv.core.client_types import StepResult +from openspiel_env.harness import OpenSpielSessionFactory, render_tic_tac_toe_board +from openspiel_env.models import OpenSpielAction, OpenSpielObservation, OpenSpielState + + +class FakeOpenSpielClient: + """In-memory OpenSpiel client that plays a scripted TTT game.""" + + def __init__(self, winning: bool = True): + self._winning = winning + self.reset_calls: list[dict[str, Any]] = [] + self.step_calls: list[OpenSpielAction] = [] + self.closed = False + + def reset(self, **kwargs: Any) -> StepResult[OpenSpielObservation]: + self.reset_calls.append(kwargs) + return StepResult( + observation=OpenSpielObservation( + info_state=[1.0] * 9 + [0.0] * 18, + legal_actions=list(range(9)), + game_phase="playing", + current_player_id=0, + ), + reward=0.0, + done=False, + ) + + def step(self, action: OpenSpielAction) -> StepResult[OpenSpielObservation]: + self.step_calls.append(action) + done = len(self.step_calls) >= 1 + reward = 1.0 if (done and self._winning) else 0.0 + info_state = ( + [0.0] * 9 + [1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0] + [0.0] * 9 + ) + return StepResult( + observation=OpenSpielObservation( + info_state=info_state, + legal_actions=[] if done else list(range(1, 9)), + game_phase="terminal" if done else "playing", + current_player_id=0, + ), + reward=reward, + done=done, + ) + + def state(self) -> OpenSpielState: + return OpenSpielState( + step_count=len(self.step_calls), + game_name="tic_tac_toe", + ) + + def close(self) -> None: + self.closed = True + + +class TestRenderTicTacToeBoard: + def test_renders_empty_board_with_action_ids(self): + empty = [1.0] * 9 + [0.0] * 18 + board = render_tic_tac_toe_board(empty) + assert "0" in board and "8" in board + + def test_renders_x_and_o_pieces(self): + info_state = list([0.0] * 27) + # cell 0 -> X + info_state[9 + 0] = 1.0 + # cell 4 -> O + info_state[18 + 4] = 1.0 + # remaining empties + for idx in (1, 2, 3, 5, 6, 7, 8): + info_state[idx] = 1.0 + board = render_tic_tac_toe_board(info_state) + assert "X" in board + assert "O" in board + + def test_empty_string_when_info_state_wrong_shape(self): + assert render_tic_tac_toe_board([0.0] * 5) == "" + + +class TestOpenSpielSessionFactory: + def test_initial_message_includes_legal_actions(self): + client = FakeOpenSpielClient() + factory = OpenSpielSessionFactory(lambda: client, game_name="tic_tac_toe") + + session = factory.create(episode_id="ttt-000000") + try: + initial = session.initial_messages() + finally: + session.close() + + assert len(initial) == 1 + content = initial[0]["content"] + assert "tic_tac_toe" in content + assert "Legal actions: [0, 1, 2, 3, 4, 5, 6, 7, 8]" in content + assert "Board:" in content + + def test_play_move_tool_call_invokes_step_with_action_id(self): + client = FakeOpenSpielClient() + factory = OpenSpielSessionFactory(lambda: client) + + session = factory.create(episode_id="ttt-000001") + try: + tool_result = session.call_tool("play_move", {"action_id": 4}) + finally: + session.close() + + assert client.step_calls[0].action_id == 4 + assert client.step_calls[0].game_name == "tic_tac_toe" + assert tool_result.done is True + assert tool_result.metadata["reward"] == 1.0 + assert "board" in tool_result.data + + def test_session_closes_underlying_client(self): + client = FakeOpenSpielClient() + factory = OpenSpielSessionFactory(lambda: client) + + session = factory.create() + session.close() + + assert client.closed is True + + def test_reset_forwards_seed_and_episode_id(self): + client = FakeOpenSpielClient() + factory = OpenSpielSessionFactory(lambda: client) + + factory.create(seed=7, episode_id="ttt-000002") + + assert client.reset_calls[0]["seed"] == 7 + assert client.reset_calls[0]["episode_id"] == "ttt-000002" diff --git a/tests/scripts/test_browsergym_harness_eval_examples.py b/tests/scripts/test_browsergym_harness_eval_examples.py new file mode 100644 index 0000000000..ee46fa99ae --- /dev/null +++ b/tests/scripts/test_browsergym_harness_eval_examples.py @@ -0,0 +1,222 @@ +"""Smoke tests for the BrowserGym harness evaluation examples.""" + +from __future__ import annotations + +import os +import sys +from types import SimpleNamespace +from typing import Any + +import pytest +import requests + +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +sys.path.insert(0, os.path.join(_REPO_ROOT, "src")) +sys.path.insert(0, os.path.join(_REPO_ROOT, "envs")) +sys.path.insert(0, os.path.join(_REPO_ROOT, "examples")) + +from browsergym_env import BrowserGymAction, BrowserGymObservation, BrowserGymState +from browsergym_env.harness import BrowserGymSessionFactory +from browsergym_harness_eval_common import ( + _finalize_episode, + build_openai_model_step, + create_openai_client, + extract_browsergym_action, + run_white_box_episode, + SessionMCPHttpServer, + summarize_episodes, +) +from openenv.core.client_types import StepResult +from openenv.core.harness import ( + HarnessRolloutResult, + HarnessRunLimits, + SessionMCPBridge, + VerifyResult, +) + + +class FakeBrowserGymClient: + """Small BrowserGym-like client used for example tests.""" + + def __init__(self): + self.closed = False + self._step_count = 0 + self._cum_reward = 0.0 + self.step_actions: list[str] = [] + + def reset(self, **kwargs: Any) -> StepResult[BrowserGymObservation]: + del kwargs + self._step_count = 0 + self._cum_reward = 0.0 + return StepResult( + observation=BrowserGymObservation( + goal="Click the highlighted button", + axtree_txt="[13] button 'Continue'", + text="[13] button 'Continue'", + url="http://example.test", + done=False, + reward=0.0, + ), + reward=0.0, + done=False, + ) + + def step(self, action: BrowserGymAction) -> StepResult[BrowserGymObservation]: + self.step_actions.append(action.action_str) + self._step_count += 1 + done = action.action_str == 'click("13")' + reward = 1.0 if done else 0.0 + self._cum_reward += reward + return StepResult( + observation=BrowserGymObservation( + goal="Click the highlighted button", + axtree_txt="[13] button 'Continue'", + text="[13] button 'Continue'", + url="http://example.test/after", + done=done, + reward=reward, + last_action_error=False, + error="", + ), + reward=reward, + done=done, + ) + + def state(self) -> BrowserGymState: + return BrowserGymState( + episode_id="browsergym-episode", + step_count=self._step_count, + benchmark="miniwob", + task_name="click-test", + goal="Click the highlighted button", + current_url="http://example.test/after", + cum_reward=self._cum_reward, + ) + + def close(self) -> None: + self.closed = True + + +class FakeOpenAIClient: + """Minimal fake OpenAI client for build_openai_model_step tests.""" + + def __init__(self, outputs: list[str]): + self._outputs = list(outputs) + self.calls: list[dict[str, Any]] = [] + self.chat = SimpleNamespace(completions=self) + + def create(self, **kwargs: Any): + self.calls.append(kwargs) + content = self._outputs.pop(0) + return SimpleNamespace( + choices=[SimpleNamespace(message=SimpleNamespace(content=content))] + ) + + +def test_extract_browsergym_action_tolerates_common_outputs(): + assert extract_browsergym_action("click('13')") == "click('13')" + assert extract_browsergym_action("Action: click('13')") == "click('13')" + assert extract_browsergym_action("```text\nclick('13')\n```") == "click('13')" + assert extract_browsergym_action("I am not sure what to do.") == "noop()" + + +def test_session_http_bridge_lists_tools_and_calls_tool(): + factory = BrowserGymSessionFactory(client_factory=FakeBrowserGymClient) + session = factory.create(task=None) + bridge = SessionMCPBridge(session) + + try: + with SessionMCPHttpServer(bridge).start() as server: + list_response = requests.post( + server.url, + json={"jsonrpc": "2.0", "id": 1, "method": "tools/list", "params": {}}, + timeout=5, + ).json() + call_response = requests.post( + server.url, + json={ + "jsonrpc": "2.0", + "id": 2, + "method": "tools/call", + "params": {"name": "click", "arguments": {"bid": "13"}}, + }, + timeout=5, + ).json() + missing_response = requests.post( + server.url, + json={ + "jsonrpc": "2.0", + "id": 3, + "method": "tools/call", + "params": {"name": "missing", "arguments": {}}, + }, + timeout=5, + ).json() + finally: + session.close() + + assert list_response["result"]["tools"][0]["name"] == "click" + assert call_response["result"]["structuredContent"]["done"] is True + assert call_response["result"]["data"]["reward"] == 1.0 + assert missing_response["error"]["message"] == "Unknown tool: missing" + + +def test_white_box_example_loop_runs_one_episode_and_aggregates_metrics(): + fake_client = FakeOpenAIClient(["Action: click('13')"]) + session_factory = BrowserGymSessionFactory( + client_factory=FakeBrowserGymClient, + default_task="click-test", + ) + model_step = build_openai_model_step(fake_client, model="fake-model") + + episode = run_white_box_episode( + session_factory=session_factory, + model_step=model_step, + limits=HarnessRunLimits(max_turns=3), + episode_id="example-episode", + ) + summary = summarize_episodes([episode]) + + assert episode.reward == 1.0 + assert episode.success is True + assert episode.done is True + assert episode.step_count == 1 + assert summary["avg_reward"] == 1.0 + assert summary["success_rate"] == 1.0 + assert summary["avg_steps"] == 1.0 + assert fake_client.calls[0]["messages"][0]["role"] == "system" + + +def test_finalize_episode_rejects_missing_env_reward(): + class MissingRewardSession: + def verify(self, transcript, final_state=None): + return VerifyResult(env_reward=None, done=True) + + with pytest.raises(ValueError, match="did not produce an environment reward"): + _finalize_episode( + MissingRewardSession(), + HarnessRolloutResult( + messages=[{"role": "assistant", "content": "done"}], + done=True, + ), + ) + + +def test_create_openai_client_does_not_fallback_to_hf_token(monkeypatch): + import browsergym_harness_eval_common as common + + captured_kwargs: dict[str, Any] = {} + + def fake_openai(**kwargs): + captured_kwargs.update(kwargs) + return object() + + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.delenv("API_KEY", raising=False) + monkeypatch.setenv("HF_TOKEN", "hf-token") + monkeypatch.setattr(common, "OpenAI", fake_openai) + + client = create_openai_client() + + assert client is not None + assert "api_key" not in captured_kwargs diff --git a/tests/test_cli/test_collect.py b/tests/test_cli/test_collect.py new file mode 100644 index 0000000000..5e5a9fad3f --- /dev/null +++ b/tests/test_cli/test_collect.py @@ -0,0 +1,273 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Tests for the ``openenv collect`` CLI command.""" + +from __future__ import annotations + +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest +from openenv.cli.__main__ import app +from typer.testing import CliRunner + +runner = CliRunner() + + +@pytest.fixture +def mock_pipeline(): + """Patch the collect pipeline internals so tests don't hit network or disk.""" + with ( + patch("openenv.cli.commands.collect.OpenSpielEnv") as env_cls, + patch("openenv.cli.commands.collect.OpenSpielSessionFactory") as factory_cls, + patch("openenv.cli.commands.collect.CollectRunner") as runner_cls, + patch("openenv.cli.commands.collect.RolloutSerializer") as serializer_cls, + patch("openenv.cli.commands.collect.MCPHarnessAdapter") as adapter_cls, + ): + collect_result = MagicMock() + collect_result.num_collected = 3 + collect_result.num_skipped = 0 + collect_result.num_dropped = 0 + collect_result.num_failed = 0 + collect_result.avg_reward = 0.6 + collect_result.success_rate = 0.6 + collect_result.episode_ids = ["ep-000000", "ep-000001", "ep-000002"] + + runner_instance = MagicMock() + runner_instance.run.return_value = collect_result + runner_cls.return_value = runner_instance + + yield { + "env_cls": env_cls, + "factory_cls": factory_cls, + "runner_cls": runner_cls, + "runner_instance": runner_instance, + "serializer_cls": serializer_cls, + "adapter_cls": adapter_cls, + } + + +def test_scripted_provider_requires_no_llm_client(tmp_path: Path, mock_pipeline): + with patch("openenv.cli.commands.collect.create_llm_client") as create_client: + result = runner.invoke( + app, + [ + "collect", + "openspiel:tic_tac_toe", + "--base-url", + "https://example.hf.space", + "--output-dir", + str(tmp_path), + "-n", + "3", + "--provider", + "scripted", + ], + ) + + assert result.exit_code == 0, result.output + create_client.assert_not_called() + mock_pipeline["runner_instance"].run.assert_called_once() + + +def test_openai_provider_builds_llm_client(tmp_path: Path, mock_pipeline): + llm_client = MagicMock() + with ( + patch( + "openenv.cli.commands.collect.create_llm_client", return_value=llm_client + ) as create_client, + patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}), + ): + result = runner.invoke( + app, + [ + "collect", + "openspiel:tic_tac_toe", + "--base-url", + "https://example.hf.space", + "--output-dir", + str(tmp_path), + "-n", + "5", + "--provider", + "openai", + "--model", + "gpt-5-mini", + ], + ) + + assert result.exit_code == 0, result.output + create_client.assert_called_once() + kwargs = create_client.call_args.kwargs + assert kwargs["provider"] == "openai" + assert kwargs["model"] == "gpt-5-mini" + assert kwargs["api_key"] == "sk-test" + + +def test_llm_endpoint_uses_llm_teacher_with_default_provider( + tmp_path: Path, mock_pipeline +): + llm_step = MagicMock(name="llm_step") + scripted_step = MagicMock(name="scripted_step") + + with ( + patch( + "openenv.cli.commands.collect._build_llm_model_step", + return_value=llm_step, + ) as build_llm, + patch( + "openenv.cli.commands.collect._build_scripted_model_step", + return_value=scripted_step, + ) as build_scripted, + ): + result = runner.invoke( + app, + [ + "collect", + "openspiel:tic_tac_toe", + "--base-url", + "https://example.hf.space", + "--output-dir", + str(tmp_path), + "--llm-endpoint", + "http://localhost", + "--model", + "Qwen/Qwen2.5-7B-Instruct", + ], + ) + + assert result.exit_code == 0, result.output + build_llm.assert_called_once() + build_scripted.assert_not_called() + assert ( + mock_pipeline["runner_instance"].run.call_args.kwargs["model_step"] is llm_step + ) + metadata = mock_pipeline[ + "serializer_cls" + ].return_value.write_metadata.call_args.args[0] + assert metadata["provider"] == "openai-compatible" + assert metadata["llm_endpoint"] == "http://localhost" + + +def test_push_to_hub_triggers_upload(tmp_path: Path, mock_pipeline): + with ( + patch("openenv.cli.commands.collect.push_to_hf_hub") as push_mock, + ): + push_mock.return_value = "https://huggingface.co/datasets/user/ttt" + result = runner.invoke( + app, + [ + "collect", + "openspiel:tic_tac_toe", + "--base-url", + "https://example.hf.space", + "--output-dir", + str(tmp_path), + "-n", + "3", + "--provider", + "scripted", + "--push-to-hub", + "user/ttt", + "--private", + ], + ) + + assert result.exit_code == 0, result.output + push_mock.assert_called_once() + push_kwargs = push_mock.call_args.kwargs + assert push_kwargs["repo_id"] == "user/ttt" + assert push_kwargs["private"] is True + + +def test_unknown_env_id_exits_nonzero(tmp_path: Path): + result = runner.invoke( + app, + [ + "collect", + "unknown:foo", + "--base-url", + "http://example", + "--output-dir", + str(tmp_path), + "--provider", + "scripted", + ], + ) + + assert result.exit_code != 0 + assert ( + "unknown" in result.output.lower() or "not supported" in result.output.lower() + ) + + +def test_openai_provider_errors_without_model(tmp_path: Path, mock_pipeline): + with patch.dict("os.environ", {"OPENAI_API_KEY": "sk-test"}): + result = runner.invoke( + app, + [ + "collect", + "openspiel:tic_tac_toe", + "--base-url", + "https://example.hf.space", + "--output-dir", + str(tmp_path), + "--provider", + "openai", + ], + ) + + assert result.exit_code != 0 + assert "model" in result.output.lower() + + +def test_keep_losses_disables_filter(tmp_path: Path, mock_pipeline): + result = runner.invoke( + app, + [ + "collect", + "openspiel:tic_tac_toe", + "--base-url", + "https://example.hf.space", + "--output-dir", + str(tmp_path), + "--provider", + "scripted", + "--keep-losses", + ], + ) + + assert result.exit_code == 0, result.output + run_kwargs = mock_pipeline["runner_instance"].run.call_args.kwargs + # When --keep-losses is passed, no filter should be installed. + assert run_kwargs.get("should_keep") is None + + +def test_default_filters_losing_rollouts(tmp_path: Path, mock_pipeline): + result = runner.invoke( + app, + [ + "collect", + "openspiel:tic_tac_toe", + "--base-url", + "https://example.hf.space", + "--output-dir", + str(tmp_path), + "--provider", + "scripted", + ], + ) + + assert result.exit_code == 0, result.output + run_kwargs = mock_pipeline["runner_instance"].run.call_args.kwargs + should_keep = run_kwargs.get("should_keep") + assert should_keep is not None + # Sanity: a winning record should be kept, a losing one dropped. + winning = MagicMock(reward=1.0) + losing = MagicMock(reward=-1.0) + assert should_keep(winning) is True + assert should_keep(losing) is False