diff --git a/docs/source/tutorials/sft-warmup.md b/docs/source/tutorials/sft-warmup.md new file mode 100644 index 000000000..fd0f827a2 --- /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/reasoning_gym_env/harness.py b/envs/reasoning_gym_env/harness.py new file mode 100644 index 000000000..27ed6d341 --- /dev/null +++ b/envs/reasoning_gym_env/harness.py @@ -0,0 +1,119 @@ +# 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 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: + """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/sft_warmup.ipynb b/examples/sft_warmup.ipynb new file mode 100644 index 000000000..b50a24b70 --- /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 +} \ No newline at end of file diff --git a/src/openenv/cli/commands/collect.py b/src/openenv/cli/commands/collect.py index 1bb329180..192326c1b 100644 --- a/src/openenv/cli/commands/collect.py +++ b/src/openenv/cli/commands/collect.py @@ -34,7 +34,7 @@ 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. +# 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] @@ -44,6 +44,17 @@ 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]] = { @@ -137,7 +148,10 @@ def _build_llm_model_step( 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 @@ -147,7 +161,7 @@ def _build_llm_model_step( port=llm_port, model=model, api_key=os.getenv("OPENAI_API_KEY") or "not-needed", - system_prompt=_SYSTEM_PROMPT, + system_prompt=effective_system_prompt, temperature=temperature, max_tokens=max_tokens, ) @@ -156,35 +170,58 @@ def _build_llm_model_step( provider=provider, model=model, api_key=_resolve_api_key(provider), - system_prompt=_SYSTEM_PROMPT, + system_prompt=effective_system_prompt, temperature=temperature, max_tokens=max_tokens, ) - return build_model_step(client, system_prompt=_SYSTEM_PROMPT) + 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. -def _build_session_factory(env_spec: str, base_url: str): - """Dispatch an env spec ``"openspiel:tic_tac_toe"`` to a session factory.""" - env_name, _, game = env_spec.partition(":") - if env_name != "openspiel": + 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"Unknown env {env_spec!r}. Supported: openspiel:", + f"Missing variant. Use e.g. {env_name}:chain_sum", param_hint="ENV", ) - if not game: - raise typer.BadParameter( - "Missing game name. Use e.g. openspiel:tic_tac_toe", - 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 OpenSpielEnv is None or OpenSpielSessionFactory is None: - raise typer.BadParameter( - "openspiel_env is not importable. Ensure envs/ is on PYTHONPATH.", - param_hint="ENV", + + 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, ) - return OpenSpielSessionFactory( - lambda: OpenSpielEnv(base_url=base_url), - game_name=game, + + raise typer.BadParameter( + f"Unknown env {env_spec!r}. Supported: openspiel:, reasoning_gym:", + param_hint="ENV", ) @@ -289,6 +326,24 @@ def collect( 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) @@ -300,7 +355,19 @@ def collect( param_hint="--model", ) - factory = _build_session_factory(env, base_url) + 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( { @@ -324,6 +391,7 @@ def collect( llm_port=llm_port, temperature=temperature, max_tokens=max_tokens, + system_prompt=system_prompt, ) else: model_step = _build_scripted_model_step() diff --git a/src/openenv/core/harness/collect.py b/src/openenv/core/harness/collect.py index 7ffbd4901..3bd3f9fe8 100644 --- a/src/openenv/core/harness/collect.py +++ b/src/openenv/core/harness/collect.py @@ -25,6 +25,13 @@ from typing import Any, Callable, Iterable, Iterator from huggingface_hub import HfApi +from rich.progress import ( + BarColumn, + MofNCompleteColumn, + Progress, + TextColumn, + TimeElapsedColumn, +) from ..env_server.mcp_types import Tool from ..llm_client import LLMClient @@ -256,40 +263,60 @@ def run( num_dropped = 0 rewards: list[float] = [] - 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 - continue - - session = self._session_factory.create(task=task, episode_id=episode_id) - try: - 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, - ) - finally: - session.close() + 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 + ) - if should_keep is not None and not should_keep(record): - num_dropped += 1 - continue + 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 - self._serializer.write_episode(record) - num_collected += 1 - rewards.append(record.reward) + session = self._session_factory.create(task=task, episode_id=episode_id) + try: + 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, + ) + finally: + 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 = ( diff --git a/src/openenv/core/llm_client.py b/src/openenv/core/llm_client.py index 9df2ff27a..1a609ba2a 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, )