feat(harness): rollout collection + openenv collect CLI

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",
+]

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://<user>-<repo>.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()