diff --git a/envs/openspiel_env/harness.py b/envs/openspiel_env/harness.py new file mode 100644 index 0000000000..a51604bb01 --- /dev/null +++ b/envs/openspiel_env/harness.py @@ -0,0 +1,165 @@ +# 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 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: + """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/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..1ea72c0e2b --- /dev/null +++ b/src/openenv/cli/commands/collect.py @@ -0,0 +1,376 @@ +# 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.env_server.mcp_types import Tool # noqa: F401 (re-exported for users) +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. +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] + +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, +): + 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=_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=_SYSTEM_PROMPT, + temperature=temperature, + max_tokens=max_tokens, + ) + + return build_model_step(client, system_prompt=_SYSTEM_PROMPT) + + +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": + raise typer.BadParameter( + f"Unknown env {env_spec!r}. Supported: openspiel:", + param_hint="ENV", + ) + if not game: + raise typer.BadParameter( + "Missing game name. Use e.g. openspiel:tic_tac_toe", + param_hint="ENV", + ) + 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=game, + ) + + +@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, +) -> 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", + ) + + factory = _build_session_factory(env, base_url) + 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, + ) + 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}" + ) + 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/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.py b/src/openenv/core/harness/__init__.py similarity index 99% rename from src/openenv/core/harness.py rename to src/openenv/core/harness/__init__.py index 38ce1ab3a1..62be80055f 100644 --- a/src/openenv/core/harness.py +++ b/src/openenv/core/harness/__init__.py @@ -17,10 +17,10 @@ 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 +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] diff --git a/src/openenv/core/harness/collect.py b/src/openenv/core/harness/collect.py new file mode 100644 index 0000000000..7ffbd4901d --- /dev/null +++ b/src/openenv/core/harness/collect.py @@ -0,0 +1,517 @@ +# 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 +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any, Callable, Iterable, Iterator + +from huggingface_hub import HfApi + +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 asdict(self) + + +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 + 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 + 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() + + if should_keep is not None and not should_keep(record): + num_dropped += 1 + continue + + self._serializer.write_episode(record) + num_collected += 1 + rewards.append(record.reward) + + 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, + episode_ids=planned_ids, + avg_reward=avg_reward, + success_rate=success_rate, + ) + + def _reset_results_file(self) -> None: + if self._serializer.results_path.exists(): + self._serializer.results_path.unlink() + + +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] + 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)" + + 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/tests/core/test_harness_collect.py b/tests/core/test_harness_collect.py new file mode 100644 index 0000000000..67392e335f --- /dev/null +++ b/tests/core/test_harness_collect.py @@ -0,0 +1,778 @@ +# 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"} + + +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. + 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 + + 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) + + +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) + + 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("openenv.core.harness.collect.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("openenv.core.harness.collect.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("openenv.core.harness.collect.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("openenv.core.harness.collect.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("openenv.core.harness.collect.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("openenv.core.harness.collect.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/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/test_cli/test_collect.py b/tests/test_cli/test_collect.py new file mode 100644 index 0000000000..a32ccd533b --- /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.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