chess_env / src /core /harness /README.md
sergiopaniego's picture
sergiopaniego HF Staff
Upload folder using huggingface_hub
7aca0bc verified
|
Raw History Blame Contribute Delete
7.42 kB

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.
  • environment.py, adapter.py — HarnessEnvironment and AgenticHarnessAdapter (RFC 005), to run an external agentic harness inside an environment, one conversational turn per step(). Worked example: examples/claude_code_harness_eval.
  • 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.

# Smoke test (no API key, no Docker — scripted teacher picks the first legal action)
openenv collect openspiel:tic_tac_toe \
    --base-url https://<user>-<space>.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://<user>-<space>.hf.space \
    --output-dir /tmp/ttt-sft-v1 \
    -n 200 --provider openai --model gpt-5-mini \
    --push-to-hub <user>/ttt-sft-v1

# Self-hosted teacher (vLLM/TGI/Ollama) via any OpenAI-compatible base URL
openenv collect openspiel:tic_tac_toe \
    --base-url https://<user>-<space>.hf.space \
    --output-dir /tmp/ttt-sft-local \
    -n 200 --llm-endpoint http://localhost:8000 --model Qwen/Qwen3-1.7B

Programmatic use:

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://<user>-<space>.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", "<user>/ttt-sft-v1")

Output schema

results.jsonl — one JSON object per episode:

field type description
episode_id string Stable id for resume (<prefix>-<index>)
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/<env>/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.

Training captures

A loop-owning session can implement TrainableSession.fetch_training_trace() to return a TrainingTrace. Harbor implements this API. Raw fetch_proxy_trace() remains available for diagnostics and older consumers.

from openenv.core.harness import TrainingTrace, TrainingTurn

trace = TrainingTrace(turns=[TrainingTurn(
    node_id="call-0",
    prompt_token_ids=[10, 11],
    completion_token_ids=[20, 21, 22],
    per_token_logps=[-0.2, -0.4, -0.1],
    loss_mask=[0, 0, 1, 0, 1],
)])
restored = TrainingTrace.model_validate_json(trace.model_dump_json())

The mask covers prompt plus completion. Prompt positions are zero; completion positions select the sampled tokens that receive loss. A zero-masked turn stays available as context. Tokens, masks and finite nonpositive logprobs are validated when the object is built or deserialized. Missing masks and duplicate call IDs are errors. Tokens are never reconstructed from text.

The environment selects agent calls and excludes auxiliary calls and discarded retries before returning this object. Shared graph nodes appear once. Selection uses the existing capture classifier; its tool-manifest and retry heuristics are not a substitute for explicit harness evidence. The raw graph remains available for auditing. Changing a task's supervision means changing its producer mask, not adding a turn-selection callback in the trainer.

verify() remains the source of the task reward. None is ungraded and 0 is an incorrect answer. A single-call rollout produces a diagnostic warning, not a capture failure. Masks do not alter the verifier score. Trainers own advantages, loss weighting, packing and weight updates. Graph identity does not by itself normalize a rollout that expands into multiple training rows.

For training, construct the session factory with the trainer's sampling policy. Harbor applies it before inference, so the trainer need not compare sampling metadata after generation. When logprob entries identify their tokens (token_id:N), training export checks that each one pairs with its token. A disagreement on an agent turn rejects the training trace; raw capture stays available for inspection. Plain token text is not re-tokenized to infer an ID.

Migration from TRL selection hooks: use fetch_training_trace() instead of a raw trace, move whole-turn exclusions into the producer's loss_mask, and keep final answers eligible unless your task explicitly masks them. The old Harbor export continues to read legacy captures; the strict API requires explicit masks and node IDs. Evaluation-only captures cannot be used for training.