world-model-optimizer 0.2.0__py3-none-any.whl

wmo/env/base.py

@@ -0,0 +1,121 @@
+"""The `Env` protocol and its world-model backend.
+
+An `Env` is one episode's worth of environment: `reset` starts it, `step` advances it. Real
+environments (a benchmark harness, a coded oracle app, a simulator) implement the same protocol in
+their example folders, which is what makes "iterate in the world model, validate in the real env"
+a one-line swap instead of two agent loops.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+from wmo.core.types import Action, EnvState, Observation
+from wmo.engine.world_model import WorldModel
+from wmo.optimize.reward import EpisodeScore
+from wmo.tracking import RunRecord
+
+
+@runtime_checkable
+class Env(Protocol):
+    """One episode of an environment an agent steps against."""
+
+    def reset(self, task: str | None = None, seed_state: EnvState | None = None) -> EnvState:
+        """Start a fresh episode; returns the environment's LIVE state view.
+
+        Contract: the returned `EnvState` is the env's current state object, updated in place as
+        the episode advances — callers that need a point-in-time snapshot must copy it
+        (`state.model_copy(deep=True)`), which is what `run_episode` does per recorded step.
+        """
+        ...
+
+    def step(self, action: Action) -> Observation:
+        """Apply `action` to the current episode and return the environment's response."""
+        ...
+
+    def close(self) -> None:
+        """Release episode resources (sessions, containers, sim handles). Idempotent."""
+        ...
+
+
+class WorldModelEnv:
+    """`Env` backed by a `WorldModel` session.
+
+    Each `reset` opens a new session (ending any previous one); `step` delegates to
+    `WorldModel.step`. `close` ends the session in the world model — freeing its history and
+    metering — and keeps the final token/cost record available as `usage`.
+
+    RL rollouts need the episode judged before the session's history is freed, and `run_episode`
+    closes the env in its `finally` — so with `score_on_close=True`, `close` scores the session
+    (`WorldModel.score_session`) right before ending it and keeps the result as `last_score`:
+
+        env = WorldModelEnv(wm, score_on_close=True)
+        result = run_episode(env, agent, task)
+        reward = env.last_score.reward  # scalar for GRPO/PPO/REINFORCE++; .critique for SDPO
+    """
+
+    def __init__(self, world_model: WorldModel, *, score_on_close: bool = False) -> None:
+        self._world_model = world_model
+        self._score_on_close = score_on_close
+        self._session_id: str | None = None
+        self._usage: RunRecord | None = None
+        self._last_score: EpisodeScore | None = None
+        self._score_error: Exception | None = None
+
+    @property
+    def session_id(self) -> str:
+        if self._session_id is None:
+            raise RuntimeError("WorldModelEnv has no active episode; call reset() first")
+        return self._session_id
+
+    @property
+    def usage(self) -> RunRecord | None:
+        """Token/cost/time of the current episode (live) or the last closed one (final)."""
+        if self._session_id is not None:
+            return self._world_model.session_usage(self._session_id)
+        return self._usage
+
+    @property
+    def last_score(self) -> EpisodeScore:
+        """The episode score captured by the most recent scoring `close`.
+
+        Raises if no scored episode has completed yet — either the env was built without
+        `score_on_close=True`, or `close` hasn't run. If the judge call itself failed during
+        `close` (throttle, network), the ORIGINAL failure is re-raised here: `run_episode`
+        deliberately swallows teardown errors, and a batch caller must see the real cause
+        instead of a misleading "no scored episode".
+        """
+        if self._score_error is not None:
+            raise RuntimeError("scoring failed during close; see cause") from self._score_error
+        if self._last_score is None:
+            raise RuntimeError(
+                "no scored episode yet; construct WorldModelEnv(wm, score_on_close=True) "
+                "and complete an episode (run_episode closes — and thus scores — for you)"
+            )
+        return self._last_score
+
+    def reset(self, task: str | None = None, seed_state: EnvState | None = None) -> EnvState:
+        self.close()  # a leftover session would otherwise leak in the world model
+        session = self._world_model.new_session(task=task, seed_state=seed_state)
+        self._session_id = session.id
+        return session.state
+
+    def step(self, action: Action) -> Observation:
+        return self._world_model.step(self.session_id, action)
+
+    def close(self) -> None:
+        if self._session_id is None:
+            return
+        try:
+            if self._score_on_close:
+                self._last_score = None
+                self._score_error = None
+                try:
+                    self._last_score = self._world_model.score_session(self._session_id)
+                except Exception as exc:  # noqa: BLE001 - preserved and re-raised by last_score
+                    # A judge failure must not leak the session or masquerade as "unscored";
+                    # the session still ends below and last_score re-raises this cause.
+                    self._score_error = exc
+        finally:
+            self._usage = self._world_model.end_session(self._session_id)
+            self._session_id = None

wmo/env/closed_loop.py

@@ -0,0 +1,229 @@
+"""Closed-loop evaluation: roll every pool candidate over a scenario set, collect the matrix.
+
+Each (candidate model, scenario, episode) cell runs one `run_episode` with the candidate driving
+`LLMAgent` against the env, then reads the env's episode score (VERIFY). The result is the
+`OutcomeMatrix` the routing optimizer fits on and the improvement report cites.
+
+Measurement notes:
+- Latency is measured per POLICY CALL (the candidate's own completions), not per episode: episode
+  wall time is dominated by the world model's simulation latency, which production traffic never
+  pays, so quoting it would flatter nobody honestly.
+- Cost is the candidate side only, priced by its own pool entry; the env's serve/judge cost is
+  metered separately by the world model (D12 cost split).
+- Every raw candidate reply is stored (`ScenarioOutcome.replies`): that is the future
+  distillation feed. Providers do not yet surface separated thinking blocks; when they do, the
+  capture point is `_TimedProvider.complete`.
+- A cell is recorded SCORED only when the episode also ran clean. An episode that errored
+  mid-flight (provider throttle, agent crash) is unscored even if the env still produced a
+  score, because that score grades a run the candidate never got to finish: counting it as
+  reward would read infrastructure failure as incapability. The salvaged critique is kept on
+  the row, prefixed, so the diagnostic text is not lost.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import time
+from typing import TYPE_CHECKING, Protocol, cast
+
+from wmo.env.base import Env
+from wmo.env.episode import run_episode
+from wmo.env.llm_agent import LLMAgent
+from wmo.env.scenarios import Scenario
+from wmo.optimize.outcomes import OutcomeMatrix, ScenarioOutcome
+from wmo.optimize.reward import EpisodeScore
+from wmo.providers.base import (
+    DEFAULT_MAX_TOKENS,
+    Completion,
+    Message,
+    Provider,
+    ProviderConfig,
+    TokenUsage,
+    VerifyResult,
+)
+from wmo.providers.pool import ModelPool, PoolEntry, pool_provider
+
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+logger = logging.getLogger(__name__)
+
+
+class _TimedProvider:
+    """Wraps the candidate's provider to record per-call seconds, usage, and raw replies."""
+
+    def __init__(self, provider: Provider) -> None:
+        self._provider = provider
+        self.call_seconds: list[float] = []
+        self.replies: list[str] = []
+        self.usage = TokenUsage()
+
+    @property
+    def config(self) -> ProviderConfig:
+        return self._provider.config
+
+    def complete(
+        self,
+        system: str,
+        messages: list[Message],
+        *,
+        temperature: float = 0.7,
+        max_tokens: int = DEFAULT_MAX_TOKENS,
+    ) -> Completion:
+        started = time.monotonic()
+        completion = self._provider.complete(
+            system, messages, temperature=temperature, max_tokens=max_tokens
+        )
+        self.call_seconds.append(time.monotonic() - started)
+        self.replies.append(completion.text)
+        self.usage = TokenUsage(
+            input_tokens=self.usage.input_tokens + completion.usage.input_tokens,
+            cached_input_tokens=self.usage.cached_input_tokens
+            + completion.usage.cached_input_tokens,
+            output_tokens=self.usage.output_tokens + completion.usage.output_tokens,
+        )
+        return completion
+
+    def embed(self, texts: list[str]) -> list[list[float]]:
+        return self._provider.embed(texts)
+
+    def verify(self) -> VerifyResult:
+        return self._provider.verify()
+
+
+_NEEDS_SCORING_ENV = (
+    "env produced no episode score; evaluate_pool needs a scoring env "
+    "(e.g. WorldModelEnv(world_model, score_on_close=True))"
+)
+
+# Prefix on a critique salvaged from an episode that errored: the text is diagnostic, the
+# verdict is not evidence (see the module docstring).
+_SALVAGE_PREFIX = "salvage-judged despite error: "
+
+
+class _ScoringEnv(Protocol):
+    """The scoring surface `evaluate_pool` needs on top of `Env`; `WorldModelEnv` provides it."""
+
+    @property
+    def last_score(self) -> EpisodeScore: ...
+
+
+def _read_episode_score(env: Env) -> tuple[EpisodeScore | None, str | None]:
+    """Read `env.last_score` defensively, returning (score, reason it is missing).
+
+    The two failure modes are NOT the same failure. A missing attribute means the caller handed
+    `evaluate_pool` a non-scoring env, so no cell of the sweep can ever be evidence and the
+    error is fatal. A `last_score` that RAISES (WorldModelEnv re-raises a close-time scoring
+    failure, e.g. a throttled judge call) is per-episode: that one cell is unscored with the
+    reason recorded, and the sweep keeps its other completed cells.
+    """
+    scoring = cast("_ScoringEnv", env)
+    try:
+        raw: object = scoring.last_score
+    except AttributeError as exc:
+        raise ValueError(_NEEDS_SCORING_ENV) from exc
+    except RuntimeError as exc:
+        return None, f"episode scoring failed: {exc}"
+    if raw is None:
+        return None, None
+    if not isinstance(raw, EpisodeScore):
+        # An unscored row must always say WHY (the outcomes contract); a wrong-typed score
+        # silently becoming reward=None/error=None would violate it.
+        return None, (
+            f"env last_score is {type(raw).__name__}, not EpisodeScore; episode left unscored"
+        )
+    return raw, None
+
+
+def scenario_id(scenario: Scenario) -> str:
+    """Stable id for a scenario: its first provenance trace id, else a hash of the task.
+
+    Provisional until wm-create's generate contract ships first-class scenario ids
+    (DECISIONS.md 2026-07-23); both forms are deterministic across runs.
+    """
+    if scenario.provenance:
+        return scenario.provenance[0]
+    return hashlib.sha256(scenario.task.encode("utf-8")).hexdigest()[:12]
+
+
+def evaluate_pool(
+    env_factory: Callable[[], Env],
+    pool: ModelPool,
+    scenarios: list[Scenario],
+    *,
+    episodes_per_scenario: int = 1,
+    max_steps: int = 20,
+    agent_temperature: float = 0.0,
+    tools_hint: str | None = None,
+    provider_factory: Callable[[PoolEntry], Provider] = pool_provider,
+    on_outcome: Callable[[ScenarioOutcome], None] | None = None,
+) -> OutcomeMatrix:
+    """Run every pool candidate over `scenarios`, one fresh env per episode.
+
+    The env must score episodes on close (`WorldModelEnv(..., score_on_close=True)`): a matrix
+    without verified rewards is not evidence. Episodes that error, and episodes whose scoring
+    itself fails, are recorded unscored (`reward=None`, `error` set) rather than defaulted to 0,
+    and never abort the sweep. `on_outcome` fires after each cell for progress display; a
+    callback that raises is logged and ignored, since a broken progress pipe must not throw away
+    the cells already paid for.
+    """
+    outcomes: list[ScenarioOutcome] = []
+    for entry in pool.models:
+        for scenario in scenarios:
+            sid = scenario_id(scenario)
+            for episode in range(episodes_per_scenario):
+                timed = _TimedProvider(provider_factory(entry))
+                agent = LLMAgent(timed, temperature=agent_temperature, tools_hint=tools_hint)
+                env = env_factory()
+                result = run_episode(env, agent, scenario.task, max_steps=max_steps)
+                score, score_error = _read_episode_score(env)
+                error = result.error
+                if score_error is not None:
+                    error = f"{error}; {score_error}" if error else score_error
+                if score is None and error is None:
+                    raise ValueError(_NEEDS_SCORING_ENV)
+                critique = score.critique if score else ""
+                if score is not None and result.error is not None:
+                    # The episode broke mid-flight, so this verdict grades an unfinished run:
+                    # keep the text, drop the reward (see the module docstring).
+                    critique = f"{_SALVAGE_PREFIX}{critique}" if critique else ""
+                    score = None
+                outcome = ScenarioOutcome(
+                    scenario_id=sid,
+                    task=scenario.task,
+                    model=entry.name,
+                    episode=episode,
+                    reward=score.reward if score else None,
+                    success=score.success if score else False,
+                    critique=critique,
+                    steps=len(result.steps),
+                    stop_reason=str(result.stop_reason),
+                    usage=timed.usage,
+                    cost_usd=entry.cost_usd(timed.usage),
+                    call_seconds=timed.call_seconds,
+                    replies=timed.replies,
+                    error=error,
+                )
+                outcomes.append(outcome)
+                if on_outcome is not None:
+                    try:
+                        on_outcome(outcome)
+                    except Exception:  # noqa: BLE001 - a broken progress pipe never costs cells
+                        logger.warning(
+                            "on_outcome callback failed for %s on %s ep%d; sweep continues",
+                            entry.name,
+                            sid,
+                            episode,
+                            exc_info=True,
+                        )
+                logger.info(
+                    "closed-loop %s on %s ep%d: reward=%s cost=$%.5f steps=%d",
+                    entry.name,
+                    sid,
+                    episode,
+                    "unscored" if outcome.reward is None else f"{outcome.reward:.2f}",
+                    outcome.cost_usd,
+                    outcome.steps,
+                )
+    return OutcomeMatrix(pool=pool.models, outcomes=outcomes)

wmo/env/episode.py

@@ -0,0 +1,107 @@
+"""`run_episode`: the one agent-vs-environment rollout loop.
+
+Every workstream that "runs an agent against an environment for N steps and scores the result"
+uses this loop, so episode records are comparable across the world model and real backends.
+"""
+
+from __future__ import annotations
+
+import traceback
+from enum import StrEnum
+from typing import Protocol, runtime_checkable
+
+from pydantic import BaseModel, Field
+
+from wmo.core.types import Action, ActionKind, EnvState, Step
+from wmo.env.base import Env
+
+# An agent signals it is finished by returning a MESSAGE action whose content equals this.
+DONE_SIGNAL = "<DONE>"
+
+
+@runtime_checkable
+class Agent(Protocol):
+    """Anything that maps the episode so far to the next action.
+
+    `act` sees the task, the env's live state view (see `Env.reset`), and the full history of
+    steps taken this episode. Return an `Action` to continue, or a MESSAGE action whose content
+    is `DONE_SIGNAL` to stop.
+    """
+
+    def act(self, task: str | None, state: EnvState, history: list[Step]) -> Action: ...
+
+
+class StopReason(StrEnum):
+    AGENT_DONE = "agent_done"  # the agent returned DONE_SIGNAL
+    MAX_STEPS = "max_steps"  # the step budget ran out
+    ENV_ERROR = "env_error"  # env.step raised; episode recorded up to the failure
+    AGENT_ERROR = "agent_error"  # agent.act raised; episode recorded up to the failure
+
+
+class EpisodeResult(BaseModel):
+    """One completed rollout: what happened, why it stopped."""
+
+    task: str | None = None
+    steps: list[Step] = Field(default_factory=list)
+    stop_reason: StopReason
+    error: str | None = None  # set on ENV_ERROR/AGENT_ERROR; names the failure
+    error_traceback: str | None = None  # full traceback of that failure, for debugging batches
+
+
+def run_episode(
+    env: Env,
+    agent: Agent,
+    task: str | None = None,
+    *,
+    seed_state: EnvState | None = None,
+    max_steps: int = 20,
+) -> EpisodeResult:
+    """Roll one episode of `agent` against `env`, bounded by `max_steps`.
+
+    The env's `reset`/`close` bracket the episode; each turn the agent proposes an action from the
+    accumulated history and the env answers with an observation. Each recorded step's
+    `state_before` is a deep copy of the env's state at that moment (the live state object keeps
+    mutating; see `Env.reset`). An env exception is recorded (not raised) so batch runs survive a
+    flaky backend; callers inspect `stop_reason`/`error`.
+    """
+    if max_steps < 1:
+        raise ValueError(f"max_steps must be >= 1, got {max_steps}")
+    try:
+        state = env.reset(task=task, seed_state=seed_state)
+        history: list[Step] = []
+        for _ in range(max_steps):
+            try:
+                action = agent.act(task, state, history)
+            except Exception as exc:  # noqa: BLE001 - batch runs must survive one bad episode
+                return EpisodeResult(
+                    task=task,
+                    steps=history,
+                    stop_reason=StopReason.AGENT_ERROR,
+                    error=f"{type(exc).__name__}: {exc} (in agent.act)",
+                    error_traceback="".join(traceback.format_exception(exc)),
+                )
+            if action.kind is ActionKind.MESSAGE and action.content == DONE_SIGNAL:
+                return EpisodeResult(task=task, steps=history, stop_reason=StopReason.AGENT_DONE)
+            state_before = state.model_copy(deep=True)
+            try:
+                observation = env.step(action)
+            except Exception as exc:  # noqa: BLE001 - batch runs must survive one bad episode
+                return EpisodeResult(
+                    task=task,
+                    steps=history,
+                    stop_reason=StopReason.ENV_ERROR,
+                    error=(
+                        f"{type(exc).__name__}: {exc} "
+                        f"(while executing {action.kind.value} {action.name or action.content!r})"
+                    ),
+                    error_traceback="".join(traceback.format_exception(exc)),
+                )
+            history.append(
+                Step(action=action, observation=observation, state_before=state_before, task=task)
+            )
+        return EpisodeResult(task=task, steps=history, stop_reason=StopReason.MAX_STEPS)
+    finally:
+        try:
+            env.close()
+        except Exception:  # noqa: BLE001, S110 - a teardown failure must not mask the result
+            pass

wmo/env/llm_agent.py

@@ -0,0 +1,93 @@
+"""A minimal LLM agent for rollouts: one tool call (or DONE) per turn, JSON-formatted.
+
+This is the reusable counterpart of the throwaway agent inside `wmo demo`: it implements the
+`Agent` protocol so scenario verification and research runs can roll real episodes against a
+world model without every caller re-writing the same prompt-and-parse loop. It is deliberately
+simple — no planning scaffold — because its role is "a competent baseline agent", not SOTA.
+"""
+
+from __future__ import annotations
+
+import json
+
+from pydantic import BaseModel, ValidationError
+
+from wmo.core.parsing import extract_json_object
+from wmo.core.types import Action, ActionKind, EnvState, JsonObject, Step
+from wmo.env.episode import DONE_SIGNAL
+from wmo.providers.base import Message, Provider
+
+AGENT_SYSTEM = """You are an agent operating in a tool environment to complete a task.
+
+Each turn, respond with ONLY a JSON object, no prose around it — one of:
+{"tool": "<tool name>", "arguments": {...}}         to act,
+{"done": true, "summary": "<what you achieved>"}    when the task is complete or impossible.
+
+Choose tool names and arguments consistent with the environment's responses so far. Work
+efficiently: no redundant calls, finish as soon as the task is done."""
+
+_MAX_HISTORY_CHARS = 500
+
+
+class _AgentReply(BaseModel):
+    """Lenient view of the agent's JSON reply."""
+
+    tool: str | None = None
+    arguments: JsonObject = {}
+    done: bool = False
+    summary: str = ""
+
+
+class LLMAgent:
+    """`Agent`-protocol adapter around a provider: history in, one JSON tool call out."""
+
+    def __init__(
+        self, provider: Provider, *, temperature: float = 0.0, tools_hint: str | None = None
+    ) -> None:
+        self._provider = provider
+        self._temperature = temperature
+        # Corpus-derived tool surface (names + argument keys observed in the traces). Without
+        # it, capable models honestly refuse to invent tools while weaker ones hallucinate
+        # them, and closed-loop rewards measure affordance-guessing instead of capability.
+        self._system = AGENT_SYSTEM if not tools_hint else f"{AGENT_SYSTEM}\n\n{tools_hint}"
+
+    def act(self, task: str | None, state: EnvState, history: list[Step]) -> Action:
+        prompt = _render_turn(task, state, history)
+        completion = self._provider.complete(
+            self._system,
+            [Message(role="user", content=prompt)],
+            temperature=self._temperature,
+            max_tokens=1024,
+        )
+        raw = extract_json_object(completion.text)
+        if raw is not None:
+            try:
+                reply = _AgentReply.model_validate_json(raw)
+            except ValidationError:
+                reply = None
+            if reply is not None:
+                if reply.done or reply.tool is None:
+                    return Action(kind=ActionKind.MESSAGE, content=DONE_SIGNAL)
+                return Action(kind=ActionKind.TOOL_CALL, name=reply.tool, arguments=reply.arguments)
+        # Unparseable reply: surface it as a message action; the env will answer and the episode
+        # continues rather than crashing the batch.
+        return Action(kind=ActionKind.MESSAGE, content=completion.text.strip()[:_MAX_HISTORY_CHARS])
+
+
+def _render_turn(task: str | None, state: EnvState, history: list[Step]) -> str:
+    lines = [f"TASK: {task or '(none)'}"]
+    if state.scratchpad:
+        lines.append(f"ENVIRONMENT NOTES: {state.scratchpad}")
+    if history:
+        lines.append("EPISODE SO FAR:")
+        for index, step in enumerate(history):
+            action = step.action
+            if action.kind is ActionKind.TOOL_CALL:
+                call = f"{action.name}({json.dumps(action.arguments, default=str)})"
+            else:
+                call = f"message: {action.content}"
+            observation = step.observation.content[:_MAX_HISTORY_CHARS]
+            error_mark = " [ERROR]" if step.observation.is_error else ""
+            lines.append(f"{index}. {call} -> {observation}{error_mark}")
+    lines.append("Your next move (JSON only):")
+    return "\n".join(lines)

wmo/env/scenarios.py

@@ -0,0 +1,73 @@
+"""Scenarios: the task prompts an agent trains and is evaluated on, derived from traces.
+
+v1 scenario creation is deliberately minimal (the prompts already recorded in the corpus's traces
+ARE the scenarios): `scenarios_from_traces` extracts one `Scenario` per unique task from the given
+traces. Callers control leakage by choosing which traces to pass — extract training scenarios from
+the train split and held-out scenarios from the test split, and the two can never overlap because
+the whole-trace split already separated them.
+
+Principled scenario *generation* (coverage, difficulty calibration, counterfactuals) is a later
+layer that will produce the same `Scenario` type.
+"""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+from wmo.core.types import ActionKind, Trace
+
+
+class Scenario(BaseModel):
+    """One task an agent can attempt against an environment."""
+
+    task: str  # the prompt handed to the agent (and to Env.reset)
+    provenance: list[str] = Field(default_factory=list)  # trace_ids this scenario came from
+
+
+def scenarios_from_traces(traces: list[Trace]) -> list[Scenario]:
+    """Extract the unique task prompts from `traces`, in first-seen order.
+
+    Traces without a task (or with a whitespace-only one) are skipped: a scenario is exactly "a
+    prompt we can hand to an agent", and an empty prompt isn't one. Duplicate tasks collapse into
+    a single scenario whose `provenance` lists every contributing trace.
+    """
+    by_task: dict[str, Scenario] = {}
+    for trace in traces:
+        task = _trace_task(trace)
+        if task is None:
+            continue
+        scenario = by_task.get(task)
+        if scenario is None:
+            by_task[task] = Scenario(task=task, provenance=[trace.trace_id])
+        elif trace.trace_id not in scenario.provenance:
+            scenario.provenance.append(trace.trace_id)
+    return list(by_task.values())
+
+
+def _trace_task(trace: Trace) -> str | None:
+    """The trace's task prompt: first non-empty per-step task (steps carry it in this corpus)."""
+    for step in trace.steps:
+        if step.task and step.task.strip():
+            return step.task.strip()
+    return None
+
+
+def tools_hint_from_traces(traces: list[Trace]) -> str:
+    """Summarize the corpus's observed tool surface for the rollout agent's system prompt.
+
+    One line per tool: name + the union of argument keys seen in the traces (the D28
+    tools.json derivation, inlined). Empty string for tool-less corpora. This is the
+    trace-derived stopgap for the scenario tool-surface contract: agents told what exists
+    stop being scored on their willingness to hallucinate affordances.
+    """
+    args_by_tool: dict[str, set[str]] = {}
+    for trace in traces:
+        for step in trace.steps:
+            action = step.action
+            if action.kind is not ActionKind.TOOL_CALL or not action.name:
+                continue
+            args_by_tool.setdefault(action.name, set()).update(action.arguments)
+    if not args_by_tool:
+        return ""
+    lines = [f"- {name}({', '.join(sorted(keys))})" for name, keys in sorted(args_by_tool.items())]
+    return "AVAILABLE TOOLS (observed in this environment):\n" + "\n".join(lines)

wmo/evals/__init__.py

@@ -0,0 +1,52 @@
+"""World-model evaluation: one interface, an open-loop and a closed-loop implementation.
+
+- `base` — the general `Evaluation`/`EvalResult` interface: mode-specific inputs are bound at
+  construction; `run()` returns a report with a one-line `summary()` and a `headline` score.
+- `open_loop` — teacher-forced replay of held-out trace steps, scored for per-step reconstruction
+  fidelity (the default `wmo eval` mode).
+- `closed_loop` — a live agent runs tasks with the world model as its environment, gold-judged for
+  end-to-end task success over k=3 passes (`wmo eval --mode closed-loop`).
+- `agreement` — compare two closed-loop reports task-by-task (e.g. simulated vs real): the
+  outcome-agreement validity check.
+- `gold` / `tasks` — the gold-assertion judge and the task specs closed-loop eval scores against.
+"""
+
+from wmo.evals.agreement import AgreementReport, compute_agreement
+from wmo.evals.base import EvalResult, Evaluation
+from wmo.evals.closed_loop import (
+    ClosedLoopEval,
+    ClosedLoopReport,
+    TaskOutcome,
+    WorldModelEnvironment,
+    evaluate_closed_loop,
+    evaluate_with_env,
+)
+from wmo.evals.gold import GoldJudge, GoldVerdict
+from wmo.evals.grid import GridCell, GridResult, ModelSpec, merge_results, run_grid
+from wmo.evals.open_loop import EvalReport, OpenLoopEval, evaluate_files
+from wmo.evals.tasks import TaskSpec, load_tasks
+
+__all__ = [
+    "AgreementReport",
+    "ClosedLoopEval",
+    "ClosedLoopReport",
+    "EvalReport",
+    "EvalResult",
+    "Evaluation",
+    "GoldJudge",
+    "GoldVerdict",
+    "GridCell",
+    "GridResult",
+    "ModelSpec",
+    "OpenLoopEval",
+    "TaskOutcome",
+    "TaskSpec",
+    "WorldModelEnvironment",
+    "compute_agreement",
+    "evaluate_closed_loop",
+    "evaluate_files",
+    "evaluate_with_env",
+    "load_tasks",
+    "merge_results",
+    "run_grid",
+]