From a8df7932f700425dfdafcefc07d6ba169d77c4ea Mon Sep 17 00:00:00 2001 From: mikasenghaas Date: Mon, 29 Jun 2026 22:02:58 +0000 Subject: [PATCH 1/9] feat(graph): typed MessageNode.kind for replay resume points (Option B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Illustrative change for the replay-buffer design discussion: how a harness marks nodes a replay buffer can resume from. Two things are tagged because the typed graph can't express them — branch provenance (compaction) and tool failure status (failed_tool_call). Tool identity / "is this a tool call" stay intrinsic to the graph. Option B adds a typed `kind: NodeTag` field to MessageNode (a core schema change), so the tag rides with the node (no keying, validated, rides the wire + dump automatically). The program is the sensor for tool failures (the trace drops isError/exit); the harness is the writer, stamping node.kind on the finished graph post-launch. annotate.py reads it. Compare with Option A (trace.info side-channel, no schema change): exp/replay-node-tags-info. Co-Authored-By: Claude Opus 4.8 (1M context) --- environments/compact/compact/annotate.py | 86 ++++++++++++++++++++++++ environments/compact/compact/harness.py | 17 ++++- environments/compact/compact/program.py | 24 +++++-- verifiers/v1/graph.py | 15 ++++- 4 files changed, 135 insertions(+), 7 deletions(-) create mode 100644 environments/compact/compact/annotate.py diff --git a/environments/compact/compact/annotate.py b/environments/compact/compact/annotate.py new file mode 100644 index 0000000000..5116f64519 --- /dev/null +++ b/environments/compact/compact/annotate.py @@ -0,0 +1,86 @@ +"""Replay resume-point annotation + selection (Option B: typed ``MessageNode.kind`` field). + +During generation the harness tags interesting nodes; the replay-buffer selector reads the +tags back. Tool *identity* (which tool) and "is this a tool call" come from the typed graph, +so only two things need a tag: branch provenance (compaction/subagent) and tool failure +status (``failed_tool_call``), which the trace otherwise drops. + +Option B stores the tag on the node itself (``MessageNode.kind``; see ``verifiers/v1/graph.py``). +The label rides *with* the node — no keying, typed/validated, and visible anywhere traces are +inspected — at the cost of a change to the core ``MessageNode`` schema. +""" + +from __future__ import annotations + +from verifiers.v1.graph import MessageNode +from verifiers.v1.trace import Trace +from verifiers.v1.types import AssistantMessage, ToolMessage + + +def branch_start_nodes(trace: Trace) -> list[int]: + """First node of each forked branch. A node with >1 child is a fork point; ``children[0]`` + is the original line and ``children[1:]`` are the forks. The compacting harness rewrites + context every turn, so each fork start is a compaction boundary.""" + children: dict[int | None, list[int]] = {} + for nid, node in enumerate(trace.nodes): + children.setdefault(node.parent, []).append(nid) + starts: list[int] = [] + for kids in children.values(): + if len(kids) > 1: + starts.extend(kids[1:]) + return starts + + +def tool_nodes(trace: Trace) -> dict[str, int]: + """``tool_call_id -> node_id`` for every tool result — the key for joining the program's + per-call failure log back onto the graph.""" + return { + node.message.tool_call_id: nid + for nid, node in enumerate(trace.nodes) + if isinstance(node.message, ToolMessage) + } + + +def tool_name(trace: Trace, node: MessageNode) -> str | None: + """The tool/function name for a tool-result node: ``ToolMessage.name`` when the dialect + recovered it, else the issuing assistant's matching ``tool_calls[].name``.""" + m = node.message + if not isinstance(m, ToolMessage): + return None + if m.name: + return m.name + nid = node.parent + while nid is not None: + parent = trace.nodes[nid].message + if isinstance(parent, AssistantMessage): + for tc in parent.tool_calls or []: + if tc.id == m.tool_call_id: + return tc.name + break + nid = trace.nodes[nid].parent + return None + + +def get_tag(trace: Trace, node_id: int) -> str | None: + """Read a node's replay tag (Option B: from the typed ``MessageNode.kind`` field).""" + return trace.nodes[node_id].kind + + +def resume_points(trace: Trace, *, kinds: set[str]) -> list[dict]: + """Enumerate replay resume points, filtered to ``kinds``. Tool calls come from the typed + graph; compaction comes from the node tag. Each point carries its ``node`` id, ``kind``, + and for tool calls the ``tool`` name and whether it ``failed``.""" + points: list[dict] = [] + for nid, node in enumerate(trace.nodes): + if isinstance(node.message, ToolMessage): + points.append( + { + "node": nid, + "kind": "tool_call", + "tool": tool_name(trace, node), + "failed": get_tag(trace, nid) == "failed_tool_call", + } + ) + elif get_tag(trace, nid) == "compaction": + points.append({"node": nid, "kind": "compaction"}) + return [p for p in points if p["kind"] in kinds] diff --git a/environments/compact/compact/harness.py b/environments/compact/compact/harness.py index a75b8ab4a4..bb1ef83b81 100644 --- a/environments/compact/compact/harness.py +++ b/environments/compact/compact/harness.py @@ -16,6 +16,8 @@ from verifiers.v1.runtimes import ProgramResult, Runtime from verifiers.v1.trace import Trace +from compact.annotate import branch_start_nodes, tool_nodes + PROGRAM_SOURCE = (Path(__file__).resolve().parent / "program.py").read_text() @@ -39,10 +41,12 @@ async def launch( secret: str, mcp_urls: dict[str, str], ) -> ProgramResult: + tool_log_path = "/tmp/vf_compact_tool_log.json" env = { "OPENAI_BASE_URL": endpoint, "OPENAI_API_KEY": secret, "OPENAI_MODEL": ctx.model, + "TOOL_LOG": tool_log_path, } if mcp_urls: # The program connects to the tool servers over HTTP; hand it a standard @@ -51,4 +55,15 @@ async def launch( {"mcpServers": {name: {"url": url} for name, url in mcp_urls.items()}} ) program = await runtime.prepare_uv_script(PROGRAM_SOURCE, self.config.env) - return await runtime.run_program([*program, trace.task.prompt], env) + result = await runtime.run_program([*program, trace.task.prompt], env) + + # Tag replay resume points on the finished graph (Option B: typed MessageNode.kind). + # The harness is the only writer; the program was the sensor for tool failures. + for nid in branch_start_nodes(trace): # this harness rewrites context every turn + trace.nodes[nid].kind = "compaction" + by_id = tool_nodes(trace) + for rec in json.loads(await runtime.read(tool_log_path)): + nid = by_id.get(rec["tool_call_id"]) + if nid is not None: + trace.nodes[nid].kind = rec["tag"] + return result diff --git a/environments/compact/compact/program.py b/environments/compact/compact/program.py index 50687e3eaf..f79a7d7f96 100644 --- a/environments/compact/compact/program.py +++ b/environments/compact/compact/program.py @@ -27,9 +27,14 @@ import re import sys from contextlib import AsyncExitStack +from pathlib import Path from openai import AsyncOpenAI +# Where to surface per-tool-call failures the trace can't keep (the harness reads this back +# and tags the matching nodes). The harness sets it; absent when running the program standalone. +TOOL_LOG = os.environ.get("TOOL_LOG") + SYSTEM = ( "You solve a task across several turns; your NOTES are your only lasting memory. The " "first turn shows the task; after that you see only your notes — plus, the turn right " @@ -103,6 +108,9 @@ async def main() -> None: config = json.loads(os.environ.get("MCP_CONFIG", "{}")) notes: str | None = None # the durable memory carried across turns tool_output: str | None = None # the last tool result, kept for exactly one turn + tool_log: list[dict] = [] # per-tool-call status the trace drops; surfaced for the harness + if TOOL_LOG: + Path(TOOL_LOG).write_text("[]") # ensure the harness always finds the sensor's log async with AsyncExitStack() as stack: tools, dispatch = ( await connect_mcp(stack, config) if config.get("mcpServers") else ([], {}) @@ -129,13 +137,19 @@ async def main() -> None: results = [] for call in message.tool_calls: args = json.loads(call.function.arguments or "{}") - result = ( - await call_mcp(dispatch, call.function.name, args) - if call.function.name in dispatch - else f"error: unknown tool {call.function.name!r}" - ) + if call.function.name in dispatch: + try: + result, tag = await call_mcp(dispatch, call.function.name, args), None + except Exception as e: # the failure signal the trace can't keep + result, tag = f"error: {e}", "failed_tool_call" + else: + result, tag = f"error: unknown tool {call.function.name!r}", "failed_tool_call" + if tag: + tool_log.append({"tool_call_id": call.id, "tag": tag}) results.append(f"{call.function.name}:\n{result}") tool_output = "\n\n".join(results) + if TOOL_LOG: # the harness reads this back to tag failed-tool-call nodes + Path(TOOL_LOG).write_text(json.dumps(tool_log)) if __name__ == "__main__": diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index 371bd7174a..561f485a2a 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -21,7 +21,7 @@ import hashlib import json from dataclasses import dataclass -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, Literal import numpy as np from pydantic import ConfigDict, Field, field_serializer, field_validator @@ -60,6 +60,13 @@ def _decode_ndarray(d: dict) -> np.ndarray: return np.frombuffer(d["data"], dtype=np.dtype(d["dtype"])).reshape(d["shape"]) +NodeTag = Literal["compaction", "subagent", "failed_tool_call"] | None +"""Replay resume-point tag for a node: branch provenance (compaction/subagent) or tool failure +status (failed_tool_call) — the two things the typed graph can't otherwise express. Written by +the harness during generation, read by the replay-buffer resume-point selector. None for an +ordinary node.""" + + class MessageNode(StrictBaseModel): """One message in the graph: a message plus the tokens it adds to the cumulative sequence. Concatenating a root→leaf path's nodes reconstructs that branch's full token @@ -73,6 +80,12 @@ class MessageNode(StrictBaseModel): """True iff a model call produced this message (the response passed to `commit`); False for every prompt-supplied message — including assistant/tool messages fabricated as context the model never generated, which role alone can't tell apart from real turns.""" + kind: NodeTag = None + """Replay resume-point tag, written by the harness during generation: branch provenance + (compaction/subagent) or tool failure status (failed_tool_call); None for an ordinary node. + Read by the replay-buffer resume-point selector. Tool identity and "is this a tool call" + come from the typed message, so they need no tag. A plain str/None, so it rides the wire and + the JSON dump automatically (no `_NODE_DUMP_EXCLUDE` entry needed).""" token_ids: list[int] = Field(default_factory=list) """This message's delta contribution to the cumulative token sequence: its leading template scaffold + its own tokens — for an assistant, the generation-prompt scaffold From ab09c167a3ce60808c6f9d46a2dd5b402fd6e2fb Mon Sep 17 00:00:00 2001 From: mikasenghaas Date: Mon, 29 Jun 2026 23:54:16 +0000 Subject: [PATCH 2/9] feat(replay): compaction before/after resume points + snapshot plumbing + ReplayTaskset (Option B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Extends the Option-B node-field skeleton into a runnable replay-buffer slice: - Drop the failed-tool-call path (deprioritized): program.py reverted; NodeTag is now compaction-only (compaction_before/compaction_after/subagent) and authored purely from graph structure (no program sensor). - Two resume points per compaction: `compaction_after` (post-compaction branch start — continue from the compaction message) and `compaction_before` (pre-compaction branch leaf — the model regenerates the compaction, then continues). Stamped on MessageNode.kind. - Typed `MessageNode.snapshot_ref` for exec/sandbox replay + Runtime.snapshot()/restore() stubs (verifiers/v1/runtimes/base.py); durable-ref contract documented. Capture (per-turn) is a follow-up framework hook; restore is wired in ReplayTaskset.setup (no-op while None). - New `environments/replay` env: ReplayTaskset materializes one task per resume point from a buffer glob (offline), seeds task.prompt with the root->node prefix (run with the default harness), and stubs scoring (reuse the original verifier — TODO). Online/growing buffer noted as a rollout-time-sampling follow-up. Option B stores tags/snapshot refs as typed MessageNode fields (core schema change). Compare with Option A (trace.info side-channel): exp/replay-node-tags-info. Co-Authored-By: Claude Opus 4.8 (1M context) --- environments/compact/compact/annotate.py | 90 ++++++----------------- environments/compact/compact/harness.py | 20 ++--- environments/compact/compact/program.py | 24 ++---- environments/replay/pyproject.toml | 13 ++++ environments/replay/replay/__init__.py | 8 ++ environments/replay/replay/selector.py | 45 ++++++++++++ environments/replay/replay/taskset.py | 94 ++++++++++++++++++++++++ verifiers/v1/graph.py | 19 +++-- verifiers/v1/runtimes/base.py | 16 ++++ 9 files changed, 224 insertions(+), 105 deletions(-) create mode 100644 environments/replay/pyproject.toml create mode 100644 environments/replay/replay/__init__.py create mode 100644 environments/replay/replay/selector.py create mode 100644 environments/replay/replay/taskset.py diff --git a/environments/compact/compact/annotate.py b/environments/compact/compact/annotate.py index 5116f64519..bf20ae51b5 100644 --- a/environments/compact/compact/annotate.py +++ b/environments/compact/compact/annotate.py @@ -1,26 +1,29 @@ -"""Replay resume-point annotation + selection (Option B: typed ``MessageNode.kind`` field). +"""Producer-side helpers: locate the compaction resume points a harness should tag. -During generation the harness tags interesting nodes; the replay-buffer selector reads the -tags back. Tool *identity* (which tool) and "is this a tool call" come from the typed graph, -so only two things need a tag: branch provenance (compaction/subagent) and tool failure -status (``failed_tool_call``), which the trace otherwise drops. +A compacting rollout splits into branches (one per context rewrite). Each compaction exposes +two replay resume points: -Option B stores the tag on the node itself (``MessageNode.kind``; see ``verifiers/v1/graph.py``). -The label rides *with* the node — no keying, typed/validated, and visible anywhere traces are -inspected — at the cost of a change to the core ``MessageNode`` schema. +- ``compaction_after`` — the post-compaction branch start (the rewritten ``[system, user(notes)]``). + Resuming here, the model continues solving *from* the compaction message. +- ``compaction_before`` — the leaf of the branch that compaction summarized (the prior turn's + response). Resuming here, the model is back in the pre-compaction context and its continuation + *writes* the compaction itself (then keeps solving). + +This module only *finds* the nodes; where the tag is stored is the A/B decision (Option A: +``trace.info``; Option B: ``MessageNode.kind``) and lives at the harness write site. """ from __future__ import annotations -from verifiers.v1.graph import MessageNode +from verifiers.v1 import graph from verifiers.v1.trace import Trace -from verifiers.v1.types import AssistantMessage, ToolMessage -def branch_start_nodes(trace: Trace) -> list[int]: - """First node of each forked branch. A node with >1 child is a fork point; ``children[0]`` - is the original line and ``children[1:]`` are the forks. The compacting harness rewrites - context every turn, so each fork start is a compaction boundary.""" +def compaction_after_nodes(trace: Trace) -> list[int]: + """Post-compaction branch starts: the first node of each forked branch. A node with >1 + child is a fork point; ``children[0]`` is the original line, ``children[1:]`` are the + rewritten (post-compaction) branches. The compacting harness rewrites every turn, so each + is a compaction boundary.""" children: dict[int | None, list[int]] = {} for nid, node in enumerate(trace.nodes): children.setdefault(node.parent, []).append(nid) @@ -31,56 +34,9 @@ def branch_start_nodes(trace: Trace) -> list[int]: return starts -def tool_nodes(trace: Trace) -> dict[str, int]: - """``tool_call_id -> node_id`` for every tool result — the key for joining the program's - per-call failure log back onto the graph.""" - return { - node.message.tool_call_id: nid - for nid, node in enumerate(trace.nodes) - if isinstance(node.message, ToolMessage) - } - - -def tool_name(trace: Trace, node: MessageNode) -> str | None: - """The tool/function name for a tool-result node: ``ToolMessage.name`` when the dialect - recovered it, else the issuing assistant's matching ``tool_calls[].name``.""" - m = node.message - if not isinstance(m, ToolMessage): - return None - if m.name: - return m.name - nid = node.parent - while nid is not None: - parent = trace.nodes[nid].message - if isinstance(parent, AssistantMessage): - for tc in parent.tool_calls or []: - if tc.id == m.tool_call_id: - return tc.name - break - nid = trace.nodes[nid].parent - return None - - -def get_tag(trace: Trace, node_id: int) -> str | None: - """Read a node's replay tag (Option B: from the typed ``MessageNode.kind`` field).""" - return trace.nodes[node_id].kind - - -def resume_points(trace: Trace, *, kinds: set[str]) -> list[dict]: - """Enumerate replay resume points, filtered to ``kinds``. Tool calls come from the typed - graph; compaction comes from the node tag. Each point carries its ``node`` id, ``kind``, - and for tool calls the ``tool`` name and whether it ``failed``.""" - points: list[dict] = [] - for nid, node in enumerate(trace.nodes): - if isinstance(node.message, ToolMessage): - points.append( - { - "node": nid, - "kind": "tool_call", - "tool": tool_name(trace, node), - "failed": get_tag(trace, nid) == "failed_tool_call", - } - ) - elif get_tag(trace, nid) == "compaction": - points.append({"node": nid, "kind": "compaction"}) - return [p for p in points if p["kind"] in kinds] +def compaction_before_nodes(trace: Trace) -> list[int]: + """Pre-compaction points: every branch leaf except the final-answer branch. Each such leaf + is the turn whose output was summarized into the next branch's compaction message, so + resuming there puts the model right before it writes a compaction.""" + leaves = sorted(graph.leaves(trace)) + return leaves[:-1] if len(leaves) > 1 else [] diff --git a/environments/compact/compact/harness.py b/environments/compact/compact/harness.py index bb1ef83b81..41008879e0 100644 --- a/environments/compact/compact/harness.py +++ b/environments/compact/compact/harness.py @@ -16,7 +16,7 @@ from verifiers.v1.runtimes import ProgramResult, Runtime from verifiers.v1.trace import Trace -from compact.annotate import branch_start_nodes, tool_nodes +from compact.annotate import compaction_after_nodes, compaction_before_nodes PROGRAM_SOURCE = (Path(__file__).resolve().parent / "program.py").read_text() @@ -41,12 +41,10 @@ async def launch( secret: str, mcp_urls: dict[str, str], ) -> ProgramResult: - tool_log_path = "/tmp/vf_compact_tool_log.json" env = { "OPENAI_BASE_URL": endpoint, "OPENAI_API_KEY": secret, "OPENAI_MODEL": ctx.model, - "TOOL_LOG": tool_log_path, } if mcp_urls: # The program connects to the tool servers over HTTP; hand it a standard @@ -57,13 +55,11 @@ async def launch( program = await runtime.prepare_uv_script(PROGRAM_SOURCE, self.config.env) result = await runtime.run_program([*program, trace.task.prompt], env) - # Tag replay resume points on the finished graph (Option B: typed MessageNode.kind). - # The harness is the only writer; the program was the sensor for tool failures. - for nid in branch_start_nodes(trace): # this harness rewrites context every turn - trace.nodes[nid].kind = "compaction" - by_id = tool_nodes(trace) - for rec in json.loads(await runtime.read(tool_log_path)): - nid = by_id.get(rec["tool_call_id"]) - if nid is not None: - trace.nodes[nid].kind = rec["tag"] + # Tag compaction resume points on the finished graph (Option B: typed MessageNode.kind). + # Two points per compaction: the post-compaction branch start (resume to continue from + # the compaction message) and the pre-compaction branch leaf (resume to regenerate it). + for nid in compaction_after_nodes(trace): + trace.nodes[nid].kind = "compaction_after" + for nid in compaction_before_nodes(trace): + trace.nodes[nid].kind = "compaction_before" return result diff --git a/environments/compact/compact/program.py b/environments/compact/compact/program.py index f79a7d7f96..50687e3eaf 100644 --- a/environments/compact/compact/program.py +++ b/environments/compact/compact/program.py @@ -27,14 +27,9 @@ import re import sys from contextlib import AsyncExitStack -from pathlib import Path from openai import AsyncOpenAI -# Where to surface per-tool-call failures the trace can't keep (the harness reads this back -# and tags the matching nodes). The harness sets it; absent when running the program standalone. -TOOL_LOG = os.environ.get("TOOL_LOG") - SYSTEM = ( "You solve a task across several turns; your NOTES are your only lasting memory. The " "first turn shows the task; after that you see only your notes — plus, the turn right " @@ -108,9 +103,6 @@ async def main() -> None: config = json.loads(os.environ.get("MCP_CONFIG", "{}")) notes: str | None = None # the durable memory carried across turns tool_output: str | None = None # the last tool result, kept for exactly one turn - tool_log: list[dict] = [] # per-tool-call status the trace drops; surfaced for the harness - if TOOL_LOG: - Path(TOOL_LOG).write_text("[]") # ensure the harness always finds the sensor's log async with AsyncExitStack() as stack: tools, dispatch = ( await connect_mcp(stack, config) if config.get("mcpServers") else ([], {}) @@ -137,19 +129,13 @@ async def main() -> None: results = [] for call in message.tool_calls: args = json.loads(call.function.arguments or "{}") - if call.function.name in dispatch: - try: - result, tag = await call_mcp(dispatch, call.function.name, args), None - except Exception as e: # the failure signal the trace can't keep - result, tag = f"error: {e}", "failed_tool_call" - else: - result, tag = f"error: unknown tool {call.function.name!r}", "failed_tool_call" - if tag: - tool_log.append({"tool_call_id": call.id, "tag": tag}) + result = ( + await call_mcp(dispatch, call.function.name, args) + if call.function.name in dispatch + else f"error: unknown tool {call.function.name!r}" + ) results.append(f"{call.function.name}:\n{result}") tool_output = "\n\n".join(results) - if TOOL_LOG: # the harness reads this back to tag failed-tool-call nodes - Path(TOOL_LOG).write_text(json.dumps(tool_log)) if __name__ == "__main__": diff --git a/environments/replay/pyproject.toml b/environments/replay/pyproject.toml new file mode 100644 index 0000000000..7d409d74fb --- /dev/null +++ b/environments/replay/pyproject.toml @@ -0,0 +1,13 @@ +[project] +name = "replay" +version = "0.1.0" +description = "replay — a replay-buffer taskset that resumes old rollouts from compaction points." +requires-python = ">=3.10" +dependencies = ["verifiers"] + +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[tool.hatch.build.targets.wheel] +packages = ["replay"] diff --git a/environments/replay/replay/__init__.py b/environments/replay/replay/__init__.py new file mode 100644 index 0000000000..f31a689957 --- /dev/null +++ b/environments/replay/replay/__init__.py @@ -0,0 +1,8 @@ +"""replay — resume old rollouts from tagged compaction points (a replay buffer). + +Option B: compaction tags are read from the typed ``MessageNode.kind`` field (see +``replay/selector.py``).""" + +from replay.taskset import ReplayTask, ReplayTaskset, ReplayTasksetConfig + +__all__ = ["ReplayTaskset", "ReplayTasksetConfig", "ReplayTask"] diff --git a/environments/replay/replay/selector.py b/environments/replay/replay/selector.py new file mode 100644 index 0000000000..266088e99b --- /dev/null +++ b/environments/replay/replay/selector.py @@ -0,0 +1,45 @@ +"""Consumer-side resume-point selection for the replay buffer (Option B: ``MessageNode`` fields). + +Reads the compaction tags the producing harness stamped on ``MessageNode.kind`` and turns them +into resume points the ReplayTaskset can sample. ``seed_messages`` builds the replay prefix (the +root->node path); ``snapshot_ref_of`` returns the durable sandbox handle to restore (None until +per-turn snapshot capture is wired). +""" + +from __future__ import annotations + +from verifiers.v1.trace import Trace +from verifiers.v1.types import Messages + + +def get_tag(trace: Trace, node_id: int) -> str | None: + """Read a node's replay tag (Option B: from the typed ``MessageNode.kind`` field).""" + return trace.nodes[node_id].kind + + +def snapshot_ref_of(trace: Trace, node_id: int) -> str | None: + """Durable sandbox snapshot ref for a node (Option B: from ``MessageNode.snapshot_ref``). + None when snapshotting was off/unsupported (the skeleton never captures one yet).""" + return trace.nodes[node_id].snapshot_ref + + +def seed_messages(trace: Trace, node_id: int) -> Messages: + """The replay prefix: messages along root->node_id, in order.""" + path: Messages = [] + nid: int | None = node_id + while nid is not None: + path.append(trace.nodes[nid].message) + nid = trace.nodes[nid].parent + path.reverse() + return path + + +def resume_points(trace: Trace, *, kinds: set[str]) -> list[dict]: + """Compaction resume points whose tag is in ``kinds``. Each: ``node`` id and ``kind`` + ("compaction_before" | "compaction_after").""" + points: list[dict] = [] + for nid in range(len(trace.nodes)): + tag = get_tag(trace, nid) + if tag in kinds: + points.append({"node": nid, "kind": tag}) + return points diff --git a/environments/replay/replay/taskset.py b/environments/replay/replay/taskset.py new file mode 100644 index 0000000000..940f9abc18 --- /dev/null +++ b/environments/replay/replay/taskset.py @@ -0,0 +1,94 @@ +"""A minimal replay-buffer taskset (skeleton). + +Materializes replay tasks from a buffer of stored rollouts (offline mode): each compaction +resume point in each stored trace becomes one task whose ``prompt`` is the replay prefix +(root->resume node). Run it with the default harness (``--harness.id default``), which accepts +a Messages prompt and continues the rollout from the prefix: + +- ``compaction_after`` -> the prefix ends at the compaction message; the model continues solving. +- ``compaction_before`` -> the prefix is the pre-compaction context; the model first writes the + compaction (the compact system prompt instructs it), then continues. + +Scoring reuses the ORIGINAL task's verifier (skeleton stub below). Exec/sandbox envs restore the +resume point's snapshot in ``setup`` before the harness runs (skeleton: ref is None -> no-op). + +NOTE (online/growing buffer): ``load_tasks`` runs once at env-server start, so this offline +materialization sees only the buffer present at startup. The online variant samples a stored +trace + resume point per rollout in a custom replay harness (rollout-time), keeping num_tasks +virtual; left as a follow-up. +""" + +from __future__ import annotations + +import glob +import json + +from verifiers.v1.decorators import reward +from verifiers.v1.runtimes import Runtime +from verifiers.v1.task import Task +from verifiers.v1.taskset import Taskset, TasksetConfig +from verifiers.v1.trace import Trace, WireTrace + +from replay.selector import resume_points, seed_messages, snapshot_ref_of + + +class ReplayTask(Task): + """A replayed prefix plus the provenance needed to score it with the original verifier.""" + + source_trace_id: str = "" + resume_node: int = -1 + resume_kind: str = "" + snapshot_ref: str | None = None + original_task: dict = {} + original_reward: float = 0.0 + + +class ReplayTasksetConfig(TasksetConfig): + buffer_glob: str = "" + """Glob of stored-rollout JSONL files, e.g. ``.../rollouts/step_*/train_rollouts.jsonl``.""" + kinds: list[str] = ["compaction_after", "compaction_before"] + """Which resume-point kinds to materialize tasks from.""" + + +class ReplayTaskset(Taskset[ReplayTask, ReplayTasksetConfig]): + def load_tasks(self) -> list[ReplayTask]: + tasks: list[ReplayTask] = [] + kinds = set(self.config.kinds) + for path in sorted(glob.glob(self.config.buffer_glob)): + with open(path) as f: + for line in f: + line = line.strip() + if not line: + continue + src: Trace = WireTrace.model_validate(json.loads(line)) + for pt in resume_points(src, kinds=kinds): + tasks.append( + ReplayTask( + idx=len(tasks), + prompt=seed_messages(src, pt["node"]), + source_trace_id=src.id, + resume_node=pt["node"], + resume_kind=pt["kind"], + snapshot_ref=snapshot_ref_of(src, pt["node"]), + original_task=src.task.model_dump(), + original_reward=src.reward, + ) + ) + return tasks + + async def setup(self, task: ReplayTask, runtime: Runtime) -> None: + # Exec/sandbox replay: restore the sandbox to the resume point before the harness runs. + # Skeleton: snapshot capture isn't wired yet, so ref is None -> nothing to restore. + if task.snapshot_ref is not None: + await runtime.restore(task.snapshot_ref) + + @reward + def replay_reward(self, trace: Trace, task: ReplayTask) -> float: + # TODO: reuse the ORIGINAL env's verifier — reconstruct its taskset + the original Task + # from ``task.original_task`` and delegate (``inner.score(trace, runtime)``); for the + # judge-style mode, compare the model's verdict against ``task.original_reward``. + # Skeleton placeholder so the env is runnable end-to-end. + return 0.0 + + +__all__ = ["ReplayTaskset", "ReplayTasksetConfig", "ReplayTask"] diff --git a/verifiers/v1/graph.py b/verifiers/v1/graph.py index 561f485a2a..15c8c5044c 100644 --- a/verifiers/v1/graph.py +++ b/verifiers/v1/graph.py @@ -60,9 +60,10 @@ def _decode_ndarray(d: dict) -> np.ndarray: return np.frombuffer(d["data"], dtype=np.dtype(d["dtype"])).reshape(d["shape"]) -NodeTag = Literal["compaction", "subagent", "failed_tool_call"] | None -"""Replay resume-point tag for a node: branch provenance (compaction/subagent) or tool failure -status (failed_tool_call) — the two things the typed graph can't otherwise express. Written by +NodeTag = Literal["compaction_before", "compaction_after", "subagent"] | None +"""Replay resume-point tag for a node — branch provenance the typed graph can't otherwise +express. ``compaction_before``/``compaction_after`` mark the two replay points around a +compaction (regenerate it vs. continue from it); ``subagent`` marks a subagent fork. Written by the harness during generation, read by the replay-buffer resume-point selector. None for an ordinary node.""" @@ -82,10 +83,14 @@ class MessageNode(StrictBaseModel): the model never generated, which role alone can't tell apart from real turns.""" kind: NodeTag = None """Replay resume-point tag, written by the harness during generation: branch provenance - (compaction/subagent) or tool failure status (failed_tool_call); None for an ordinary node. - Read by the replay-buffer resume-point selector. Tool identity and "is this a tool call" - come from the typed message, so they need no tag. A plain str/None, so it rides the wire and - the JSON dump automatically (no `_NODE_DUMP_EXCLUDE` entry needed).""" + (compaction_before/compaction_after/subagent); None for an ordinary node. Read by the + replay-buffer resume-point selector. A plain str/None, so it rides the wire and the JSON + dump automatically (no `_NODE_DUMP_EXCLUDE` entry needed).""" + snapshot_ref: str | None = None + """Durable handle to the sandbox state captured after this node's turn (registry tag / + modal snapshot id / object-store key); restore it to resume a replay from here. Written by + the per-turn snapshot hook during generation; None when snapshotting was off/unsupported. + A plain str/None, so it rides the wire and the JSON dump automatically.""" token_ids: list[int] = Field(default_factory=list) """This message's delta contribution to the cumulative token sequence: its leading template scaffold + its own tokens — for an assistant, the generation-prompt scaffold diff --git a/verifiers/v1/runtimes/base.py b/verifiers/v1/runtimes/base.py index 8029deda4e..bb517dfb4d 100644 --- a/verifiers/v1/runtimes/base.py +++ b/verifiers/v1/runtimes/base.py @@ -258,6 +258,22 @@ async def read(self, path: str) -> bytes: async def write(self, path: str, data: bytes) -> None: """Write a file into the runtime's workspace, creating parent dirs.""" + # --- snapshot / restore (replay buffer) --- + + async def snapshot(self) -> str | None: + """Persist the current sandbox state and return a DURABLE ref (e.g. a registry tag, + modal snapshot id, or object-store key). The ref must outlive this runtime instance: + it is read back by a *different*, later rollout — possibly on another machine, possibly + replaying an offline buffer from a prior run. Override per runtime when the sandbox + supports checkpointing. Skeleton: returns None (no snapshot taken).""" + return None + + async def restore(self, ref: str) -> None: + """Restore the sandbox to a previously snapshotted state before continuing a replay + from mid-rollout. Override per runtime. Skeleton: refuse a real ref rather than + silently resuming on the wrong state (a None ref is a no-op handled by the caller).""" + raise NotImplementedError(f"{self.type}: snapshot/restore not available yet") + # --- networking --- @property From 4f7e8203166dbca53d29f118b0e76a304be994dd Mon Sep 17 00:00:00 2001 From: mikasenghaas Date: Tue, 30 Jun 2026 00:01:36 +0000 Subject: [PATCH 3/9] feat(replay): reuse original verifier for scoring + online rollout-time sampling (Option B) - Scoring delegation: ReplayTaskset.score now reuses the ORIGINAL env's verifier (config.inner) by running its rewards/metrics over the replay trace with the original task swapped in. Original task + reward ride on the replay task (offline) or in trace.info["replay"] (online). Empty inner.id => no-op (skeleton-safe). - Online buffer: config.mode="online" returns pool_size virtual slots; new ReplayHarness samples a stored trace + resume point from the LIVE buffer per rollout (re-globs each time), restores the snapshot, seeds the default chat loop with the root->node prefix via INITIAL_MESSAGES, and stashes provenance. Empty buffer => stop("replay_buffer_empty") (warmup). Offline path (materialize in load_tasks) unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --- environments/replay/replay/__init__.py | 14 ++- environments/replay/replay/harness.py | 114 +++++++++++++++++++ environments/replay/replay/taskset.py | 152 ++++++++++++++++--------- 3 files changed, 223 insertions(+), 57 deletions(-) create mode 100644 environments/replay/replay/harness.py diff --git a/environments/replay/replay/__init__.py b/environments/replay/replay/__init__.py index f31a689957..35b8b6974a 100644 --- a/environments/replay/replay/__init__.py +++ b/environments/replay/replay/__init__.py @@ -1,8 +1,18 @@ """replay — resume old rollouts from tagged compaction points (a replay buffer). -Option B: compaction tags are read from the typed ``MessageNode.kind`` field (see +Offline (``ReplayTaskset`` materializes tasks from the buffer) or online (``ReplayHarness`` +samples the live buffer per rollout). Scoring reuses the original env's verifier. + +Option B: compaction tags + snapshot refs are read from typed ``MessageNode`` fields (see ``replay/selector.py``).""" +from replay.harness import ReplayHarness, ReplayHarnessConfig from replay.taskset import ReplayTask, ReplayTaskset, ReplayTasksetConfig -__all__ = ["ReplayTaskset", "ReplayTasksetConfig", "ReplayTask"] +__all__ = [ + "ReplayTaskset", + "ReplayTasksetConfig", + "ReplayTask", + "ReplayHarness", + "ReplayHarnessConfig", +] diff --git a/environments/replay/replay/harness.py b/environments/replay/replay/harness.py new file mode 100644 index 0000000000..2c196c76d8 --- /dev/null +++ b/environments/replay/replay/harness.py @@ -0,0 +1,114 @@ +"""ReplayHarness — rollout-time sampling for an online (growing) replay buffer. + +The offline path materializes tasks in ``ReplayTaskset.load_tasks``, which runs once at +env-server start and so can't see rollouts written later. This harness instead samples a stored +trace + resume point from the *live* buffer on each rollout, so this run's own rollouts are +replayed as the buffer fills. It restores the resume point's sandbox snapshot, seeds the default +chat loop with the replay prefix (``root->node``), and stashes provenance in ``trace.info`` so +``ReplayTaskset.score`` can reuse the original verifier. + +It reuses the default harness's program (a growing-message-list chat loop) and seeds it via the +same ``INITIAL_MESSAGES`` channel the default harness uses for a Messages prompt. +""" + +from __future__ import annotations + +import glob +import json +import random +from pathlib import Path + +from verifiers.v1.clients import RolloutContext +from verifiers.v1.dialects.chat import message_to_wire +from verifiers.v1.harnesses.default.harness import ( + PROGRAM_SOURCE, + DefaultHarness, + DefaultHarnessConfig, +) +from verifiers.v1.runtimes import ProgramResult, Runtime +from verifiers.v1.trace import Trace, WireTrace + +from replay.selector import resume_points, seed_messages, snapshot_ref_of + + +class ReplayHarnessConfig(DefaultHarnessConfig): + buffer_glob: str = "" + """Glob of stored-rollout JSONL files to sample from (the live, possibly growing, buffer).""" + kinds: list[str] = ["compaction_after", "compaction_before"] + """Which resume-point kinds to sample.""" + + +class ReplayHarness(DefaultHarness): + """Subclasses the default harness (its chat-loop program) but seeds from a sampled buffer + rollout instead of ``task.prompt``.""" + + SUPPORTS_MESSAGE_PROMPT = True + + async def launch( + self, + ctx: RolloutContext, + trace: Trace, + runtime: Runtime, + endpoint: str, + secret: str, + mcp_urls: dict[str, str], + ) -> ProgramResult: + rng = random.Random(trace.id) # deterministic per rollout, varies across rollouts + sample = self._sample(rng) + if sample is None: # buffer empty (warmup) or no matching resume points yet + trace.stop("replay_buffer_empty") + return ProgramResult(exit_code=0, stdout="", stderr="") + src, point = sample + + ref = snapshot_ref_of(src, point["node"]) + if ref is not None: # exec/sandbox replay; skeleton refs are None -> skip + await runtime.restore(ref) + + # Stash provenance so ReplayTaskset.score can reuse the original verifier. + trace.info["replay"] = { + "source_id": src.id, + "resume_node": point["node"], + "kind": point["kind"], + "original_task": src.task.model_dump(), + "original_reward": src.reward, + } + + # Seed the default chat loop with the replay prefix (mirror DefaultHarness.launch). + seed = seed_messages(src, point["node"]) + env = {**self.config.env} + env["INITIAL_MESSAGES"] = json.dumps([message_to_wire(m) for m in seed]) + args = [f"--base-url={endpoint}", f"--api-key={secret}", f"--model={ctx.model}"] + if mcp_urls: + args.append( + "--mcp-config=" + + json.dumps( + {"mcpServers": {name: {"url": url} for name, url in mcp_urls.items()}} + ) + ) + program = await runtime.prepare_uv_script(PROGRAM_SOURCE, self.config.env) + return await runtime.run_program([*program, *args], env) + + def _sample(self, rng: random.Random) -> tuple[Trace, dict] | None: + """Scan the live buffer in random order; return the first (trace, resume point) found. + Re-globs every rollout, so files written after env-server start are included.""" + kinds = set(self.config.kinds) + files = sorted(glob.glob(self.config.buffer_glob)) + rng.shuffle(files) + for path in files: + try: + lines = Path(path).read_text().splitlines() + except OSError: + continue + rng.shuffle(lines) + for line in lines: + line = line.strip() + if not line: + continue + src = WireTrace.model_validate(json.loads(line)) + points = resume_points(src, kinds=kinds) + if points: + return src, rng.choice(points) + return None + + +__all__ = ["ReplayHarness", "ReplayHarnessConfig"] diff --git a/environments/replay/replay/taskset.py b/environments/replay/replay/taskset.py index 940f9abc18..22f63f7bee 100644 --- a/environments/replay/replay/taskset.py +++ b/environments/replay/replay/taskset.py @@ -1,31 +1,37 @@ -"""A minimal replay-buffer taskset (skeleton). - -Materializes replay tasks from a buffer of stored rollouts (offline mode): each compaction -resume point in each stored trace becomes one task whose ``prompt`` is the replay prefix -(root->resume node). Run it with the default harness (``--harness.id default``), which accepts -a Messages prompt and continues the rollout from the prefix: - -- ``compaction_after`` -> the prefix ends at the compaction message; the model continues solving. -- ``compaction_before`` -> the prefix is the pre-compaction context; the model first writes the - compaction (the compact system prompt instructs it), then continues. - -Scoring reuses the ORIGINAL task's verifier (skeleton stub below). Exec/sandbox envs restore the -resume point's snapshot in ``setup`` before the harness runs (skeleton: ref is None -> no-op). - -NOTE (online/growing buffer): ``load_tasks`` runs once at env-server start, so this offline -materialization sees only the buffer present at startup. The online variant samples a stored -trace + resume point per rollout in a custom replay harness (rollout-time), keeping num_tasks -virtual; left as a follow-up. +"""The replay-buffer taskset: resume old rollouts from tagged compaction points. + +Two sourcing modes (``config.mode``): + +- ``offline`` — ``load_tasks`` materializes one task per compaction resume point found in the + buffer at env-server start (snapshot of a prior run's, or this run's, rollouts). The replay + prefix (``root->node``) is the task ``prompt``; run with the default harness. +- ``online`` — ``load_tasks`` returns ``pool_size`` virtual slots and the bundled + :class:`ReplayHarness` samples a stored trace + resume point from the *live* buffer at rollout + time, so a growing buffer (this run's own rollouts) is picked up as it fills. + +Each compaction exposes two resume points: +- ``compaction_after`` -> prefix ends at the compaction message; the model continues solving. +- ``compaction_before`` -> prefix is the pre-compaction context; the model writes the compaction + itself, then continues. + +Scoring reuses the ORIGINAL env's verifier: ``config.inner`` names the taskset the rollouts came +from; ``score`` runs its ``@reward``/``@metric`` over the replay trace with the original task +swapped in (so e.g. a math verifier checks the replayed continuation's final answer against the +original ground truth). The original task + reward ride on the replay task (offline) or in +``trace.info["replay"]`` (online). """ from __future__ import annotations import glob import json +from typing import Literal + +from pydantic import SerializeAsAny -from verifiers.v1.decorators import reward +from verifiers.v1.loaders import load_taskset from verifiers.v1.runtimes import Runtime -from verifiers.v1.task import Task +from verifiers.v1.task import Task, WireTask from verifiers.v1.taskset import Taskset, TasksetConfig from verifiers.v1.trace import Trace, WireTrace @@ -33,7 +39,8 @@ class ReplayTask(Task): - """A replayed prefix plus the provenance needed to score it with the original verifier.""" + """A replayed prefix plus the provenance needed to score it with the original verifier + (offline mode; online mode carries the same provenance in ``trace.info["replay"]``).""" source_trace_id: str = "" resume_node: int = -1 @@ -44,51 +51,86 @@ class ReplayTask(Task): class ReplayTasksetConfig(TasksetConfig): + mode: Literal["offline", "online"] = "offline" buffer_glob: str = "" """Glob of stored-rollout JSONL files, e.g. ``.../rollouts/step_*/train_rollouts.jsonl``.""" kinds: list[str] = ["compaction_after", "compaction_before"] - """Which resume-point kinds to materialize tasks from.""" + """Which resume-point kinds to replay from.""" + pool_size: int = 1024 + """Online mode: number of virtual task slots (num_tasks); the harness samples per rollout.""" + inner: SerializeAsAny[TasksetConfig] = TasksetConfig() + """The ORIGINAL env's taskset (the verifier to reuse). Empty id => scoring is a no-op.""" + + +def _iter_traces(buffer_glob: str): + for path in sorted(glob.glob(buffer_glob)): + with open(path) as f: + for line in f: + line = line.strip() + if line: + yield WireTrace.model_validate(json.loads(line)) class ReplayTaskset(Taskset[ReplayTask, ReplayTasksetConfig]): + def __init__(self, config: ReplayTasksetConfig) -> None: + super().__init__(config) + # Reuse the original env's verifier when configured (skeleton: empty id => no scoring). + self._inner: Taskset | None = load_taskset(config.inner) if config.inner.id else None + def load_tasks(self) -> list[ReplayTask]: - tasks: list[ReplayTask] = [] + if self.config.mode == "online": + # Virtual slots; ReplayHarness samples the live buffer per rollout. + return [ReplayTask(idx=i, prompt=None) for i in range(self.config.pool_size)] kinds = set(self.config.kinds) - for path in sorted(glob.glob(self.config.buffer_glob)): - with open(path) as f: - for line in f: - line = line.strip() - if not line: - continue - src: Trace = WireTrace.model_validate(json.loads(line)) - for pt in resume_points(src, kinds=kinds): - tasks.append( - ReplayTask( - idx=len(tasks), - prompt=seed_messages(src, pt["node"]), - source_trace_id=src.id, - resume_node=pt["node"], - resume_kind=pt["kind"], - snapshot_ref=snapshot_ref_of(src, pt["node"]), - original_task=src.task.model_dump(), - original_reward=src.reward, - ) - ) + tasks: list[ReplayTask] = [] + for src in _iter_traces(self.config.buffer_glob): + for pt in resume_points(src, kinds=kinds): + tasks.append( + ReplayTask( + idx=len(tasks), + prompt=seed_messages(src, pt["node"]), + source_trace_id=src.id, + resume_node=pt["node"], + resume_kind=pt["kind"], + snapshot_ref=snapshot_ref_of(src, pt["node"]), + original_task=src.task.model_dump(), + original_reward=src.reward, + ) + ) return tasks async def setup(self, task: ReplayTask, runtime: Runtime) -> None: - # Exec/sandbox replay: restore the sandbox to the resume point before the harness runs. - # Skeleton: snapshot capture isn't wired yet, so ref is None -> nothing to restore. + # Offline exec/sandbox replay: restore to the resume point before the harness runs. + # (Online restores inside ReplayHarness, where the sampled ref is known.) Skeleton: + # snapshot capture isn't wired yet, so refs are None -> nothing to restore. if task.snapshot_ref is not None: await runtime.restore(task.snapshot_ref) - @reward - def replay_reward(self, trace: Trace, task: ReplayTask) -> float: - # TODO: reuse the ORIGINAL env's verifier — reconstruct its taskset + the original Task - # from ``task.original_task`` and delegate (``inner.score(trace, runtime)``); for the - # judge-style mode, compare the model's verdict against ``task.original_reward``. - # Skeleton placeholder so the env is runnable end-to-end. - return 0.0 - - -__all__ = ["ReplayTaskset", "ReplayTasksetConfig", "ReplayTask"] + async def score(self, trace: Trace, runtime: Runtime) -> None: + # Reuse the ORIGINAL env's verifier: run its rewards/metrics over the replay trace with + # the original task swapped in. (Judge-style scoring would instead compare the model's + # verdict against ``original_reward`` here.) + if self._inner is None: + return + replay = trace.info.get("replay") # online: harness-stashed provenance + original_dump = replay["original_task"] if replay else getattr(trace.task, "original_task", {}) + if not original_dump: + return + replay_task = trace.task + trace.task = WireTask.model_validate(original_dump) + try: + await self._inner.score(trace, runtime) + finally: + trace.task = replay_task + + +# Bundle the online sampling harness so `--harness.id replay` resolves alongside the taskset. +from replay.harness import ReplayHarness, ReplayHarnessConfig # noqa: E402 + +__all__ = [ + "ReplayTaskset", + "ReplayTasksetConfig", + "ReplayTask", + "ReplayHarness", + "ReplayHarnessConfig", +] From 4a65a3d742a00050a871127cb10862294f174e58 Mon Sep 17 00:00:00 2001 From: mikasenghaas Date: Tue, 30 Jun 2026 00:22:42 +0000 Subject: [PATCH 4/9] feat(replay): add 'recheck' (try-again) resume mode (Option B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The 'check your work, fix if wrong' mode: take the full sampled rollout, append a user turn (config.followup), and re-roll — scored by the same original verifier as the rest. - Structural, tag-free: the recheck point is the rollout's final-answer leaf, so it works for ANY rollout (linear or compacting) and is identical across A/B (no tag dependency). - selector: recheck_points() + build_seed() (appends the follow-up user turn for recheck; plain prefix otherwise). resume_points() now yields recheck alongside compaction kinds. - taskset + harness: new `followup` config; `recheck` added to default `kinds`; both build seeds via build_seed so offline and online modes get the appended turn. Co-Authored-By: Claude Opus 4.8 (1M context) --- environments/replay/replay/harness.py | 12 ++++++++---- environments/replay/replay/selector.py | 26 +++++++++++++++++++++++--- environments/replay/replay/taskset.py | 13 +++++++++---- 3 files changed, 40 insertions(+), 11 deletions(-) diff --git a/environments/replay/replay/harness.py b/environments/replay/replay/harness.py index 2c196c76d8..2b508b22e3 100644 --- a/environments/replay/replay/harness.py +++ b/environments/replay/replay/harness.py @@ -28,14 +28,18 @@ from verifiers.v1.runtimes import ProgramResult, Runtime from verifiers.v1.trace import Trace, WireTrace -from replay.selector import resume_points, seed_messages, snapshot_ref_of +from replay.selector import build_seed, resume_points, snapshot_ref_of class ReplayHarnessConfig(DefaultHarnessConfig): buffer_glob: str = "" """Glob of stored-rollout JSONL files to sample from (the live, possibly growing, buffer).""" - kinds: list[str] = ["compaction_after", "compaction_before"] - """Which resume-point kinds to sample.""" + kinds: list[str] = ["recheck", "compaction_after", "compaction_before"] + """Which resume-point kinds to sample (``recheck`` = append a check-your-work turn).""" + followup: str = ( + "Check your work. If anything is wrong, fix it and give the corrected final answer." + ) + """The user turn appended for ``recheck`` points.""" class ReplayHarness(DefaultHarness): @@ -74,7 +78,7 @@ async def launch( } # Seed the default chat loop with the replay prefix (mirror DefaultHarness.launch). - seed = seed_messages(src, point["node"]) + seed = build_seed(src, point, self.config.followup) env = {**self.config.env} env["INITIAL_MESSAGES"] = json.dumps([message_to_wire(m) for m in seed]) args = [f"--base-url={endpoint}", f"--api-key={secret}", f"--model={ctx.model}"] diff --git a/environments/replay/replay/selector.py b/environments/replay/replay/selector.py index 266088e99b..15c667b19f 100644 --- a/environments/replay/replay/selector.py +++ b/environments/replay/replay/selector.py @@ -8,8 +8,9 @@ from __future__ import annotations +from verifiers.v1 import graph from verifiers.v1.trace import Trace -from verifiers.v1.types import Messages +from verifiers.v1.types import Messages, UserMessage def get_tag(trace: Trace, node_id: int) -> str | None: @@ -34,12 +35,31 @@ def seed_messages(trace: Trace, node_id: int) -> Messages: return path +def recheck_points(trace: Trace) -> list[dict]: + """The 'try again' point: the rollout's final-answer leaf (the last branch's leaf). Purely + structural — no producer tag needed — so it works for any rollout, compacting or linear.""" + leaves = graph.leaves(trace) + return [{"node": max(leaves), "kind": "recheck"}] if leaves else [] + + def resume_points(trace: Trace, *, kinds: set[str]) -> list[dict]: - """Compaction resume points whose tag is in ``kinds``. Each: ``node`` id and ``kind`` - ("compaction_before" | "compaction_after").""" + """Resume points whose ``kind`` is in ``kinds``. ``compaction_before``/``compaction_after`` + come from the harness tags; ``recheck`` is the structural final-answer leaf. Each: ``node`` + id and ``kind``.""" points: list[dict] = [] for nid in range(len(trace.nodes)): tag = get_tag(trace, nid) if tag in kinds: points.append({"node": nid, "kind": tag}) + if "recheck" in kinds: + points.extend(recheck_points(trace)) return points + + +def build_seed(trace: Trace, point: dict, followup: str) -> Messages: + """The seed conversation for a resume point: the ``root->node`` prefix, plus — for a + ``recheck`` point — an appended user turn asking the model to check and fix its work.""" + msgs = seed_messages(trace, point["node"]) + if point["kind"] == "recheck": + msgs = [*msgs, UserMessage(content=followup)] + return msgs diff --git a/environments/replay/replay/taskset.py b/environments/replay/replay/taskset.py index 22f63f7bee..bf68664594 100644 --- a/environments/replay/replay/taskset.py +++ b/environments/replay/replay/taskset.py @@ -35,7 +35,7 @@ from verifiers.v1.taskset import Taskset, TasksetConfig from verifiers.v1.trace import Trace, WireTrace -from replay.selector import resume_points, seed_messages, snapshot_ref_of +from replay.selector import build_seed, resume_points, snapshot_ref_of class ReplayTask(Task): @@ -54,8 +54,13 @@ class ReplayTasksetConfig(TasksetConfig): mode: Literal["offline", "online"] = "offline" buffer_glob: str = "" """Glob of stored-rollout JSONL files, e.g. ``.../rollouts/step_*/train_rollouts.jsonl``.""" - kinds: list[str] = ["compaction_after", "compaction_before"] - """Which resume-point kinds to replay from.""" + kinds: list[str] = ["recheck", "compaction_after", "compaction_before"] + """Which resume-point kinds to replay from (``recheck`` = append a check-your-work turn to + the full rollout; ``compaction_*`` = resume around a compaction).""" + followup: str = ( + "Check your work. If anything is wrong, fix it and give the corrected final answer." + ) + """The user turn appended for ``recheck`` points.""" pool_size: int = 1024 """Online mode: number of virtual task slots (num_tasks); the harness samples per rollout.""" inner: SerializeAsAny[TasksetConfig] = TasksetConfig() @@ -88,7 +93,7 @@ def load_tasks(self) -> list[ReplayTask]: tasks.append( ReplayTask( idx=len(tasks), - prompt=seed_messages(src, pt["node"]), + prompt=build_seed(src, pt, self.config.followup), source_trace_id=src.id, resume_node=pt["node"], resume_kind=pt["kind"], From 21c15458306da5fcf82e632bbe354f99349413ad Mon Sep 17 00:00:00 2001 From: mikasenghaas Date: Tue, 30 Jun 2026 00:29:44 +0000 Subject: [PATCH 5/9] refactor(replay): dedup config defaults + buffer reader; add judge mode (Option B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cleanup: - Shared DEFAULT_KINDS / DEFAULT_FOLLOWUP constants in selector.py, referenced by both the taskset and harness configs (was duplicated literals). - Buffer reading centralized as selector.iter_traces; taskset drops its glob/json copy. - resume_points folds recheck/judge final-leaf points in one place. Judge mode (config: add "judge" to kinds): - A judge replay point presents the rollout's transcript ("was this correct? yes/no") instead of continuing it (selector.judge_prompt / build_seed). - ReplayTaskset.score grades the model's verdict against the original rollout's reward (original_reward > judge_threshold) — a self-supervised correctness label; no inner verifier or snapshot needed. recheck/compaction_* still reuse the original verifier. Co-Authored-By: Claude Opus 4.8 (1M context) --- environments/replay/replay/harness.py | 16 ++-- environments/replay/replay/selector.py | 83 ++++++++++++----- environments/replay/replay/taskset.py | 124 +++++++++++++++---------- 3 files changed, 144 insertions(+), 79 deletions(-) diff --git a/environments/replay/replay/harness.py b/environments/replay/replay/harness.py index 2b508b22e3..54a3863a56 100644 --- a/environments/replay/replay/harness.py +++ b/environments/replay/replay/harness.py @@ -28,17 +28,21 @@ from verifiers.v1.runtimes import ProgramResult, Runtime from verifiers.v1.trace import Trace, WireTrace -from replay.selector import build_seed, resume_points, snapshot_ref_of +from replay.selector import ( + DEFAULT_FOLLOWUP, + DEFAULT_KINDS, + build_seed, + resume_points, + snapshot_ref_of, +) class ReplayHarnessConfig(DefaultHarnessConfig): buffer_glob: str = "" """Glob of stored-rollout JSONL files to sample from (the live, possibly growing, buffer).""" - kinds: list[str] = ["recheck", "compaction_after", "compaction_before"] - """Which resume-point kinds to sample (``recheck`` = append a check-your-work turn).""" - followup: str = ( - "Check your work. If anything is wrong, fix it and give the corrected final answer." - ) + kinds: list[str] = DEFAULT_KINDS + """Which replay kinds to sample (see ReplayTaskset; add ``"judge"`` to opt in).""" + followup: str = DEFAULT_FOLLOWUP """The user turn appended for ``recheck`` points.""" diff --git a/environments/replay/replay/selector.py b/environments/replay/replay/selector.py index 15c667b19f..d4c09ac8c0 100644 --- a/environments/replay/replay/selector.py +++ b/environments/replay/replay/selector.py @@ -1,17 +1,24 @@ -"""Consumer-side resume-point selection for the replay buffer (Option B: ``MessageNode`` fields). +"""Consumer-side buffer reading + resume-point selection (Option B: ``MessageNode`` fields). -Reads the compaction tags the producing harness stamped on ``MessageNode.kind`` and turns them -into resume points the ReplayTaskset can sample. ``seed_messages`` builds the replay prefix (the -root->node path); ``snapshot_ref_of`` returns the durable sandbox handle to restore (None until -per-turn snapshot capture is wired). +Turns the compaction tags the producing harness stamped on ``MessageNode.kind`` into resume +points the ReplayTaskset/ReplayHarness sample. ``build_seed`` produces the seed conversation per +mode; ``snapshot_ref_of`` returns the durable sandbox handle to restore (None until per-turn +snapshot capture is wired). Only ``get_tag``/``snapshot_ref_of`` differ from Option A — everything +else (buffer reading, resume-point logic, seed building) is shared. """ from __future__ import annotations +import glob +import json + from verifiers.v1 import graph -from verifiers.v1.trace import Trace +from verifiers.v1.trace import Trace, WireTrace from verifiers.v1.types import Messages, UserMessage +DEFAULT_KINDS = ["recheck", "compaction_after", "compaction_before"] +DEFAULT_FOLLOWUP = "Check your work. If anything is wrong, fix it and give the corrected final answer." + def get_tag(trace: Trace, node_id: int) -> str | None: """Read a node's replay tag (Option B: from the typed ``MessageNode.kind`` field).""" @@ -24,6 +31,16 @@ def snapshot_ref_of(trace: Trace, node_id: int) -> str | None: return trace.nodes[node_id].snapshot_ref +def iter_traces(buffer_glob: str): + """Yield each stored rollout (``WireTrace``) from the buffer glob, in file then line order.""" + for path in sorted(glob.glob(buffer_glob)): + with open(path) as f: + for line in f: + line = line.strip() + if line: + yield WireTrace.model_validate(json.loads(line)) + + def seed_messages(trace: Trace, node_id: int) -> Messages: """The replay prefix: messages along root->node_id, in order.""" path: Messages = [] @@ -35,30 +52,48 @@ def seed_messages(trace: Trace, node_id: int) -> Messages: return path -def recheck_points(trace: Trace) -> list[dict]: - """The 'try again' point: the rollout's final-answer leaf (the last branch's leaf). Purely - structural — no producer tag needed — so it works for any rollout, compacting or linear.""" - leaves = graph.leaves(trace) - return [{"node": max(leaves), "kind": "recheck"}] if leaves else [] - - def resume_points(trace: Trace, *, kinds: set[str]) -> list[dict]: """Resume points whose ``kind`` is in ``kinds``. ``compaction_before``/``compaction_after`` - come from the harness tags; ``recheck`` is the structural final-answer leaf. Each: ``node`` - id and ``kind``.""" - points: list[dict] = [] - for nid in range(len(trace.nodes)): - tag = get_tag(trace, nid) - if tag in kinds: - points.append({"node": nid, "kind": tag}) - if "recheck" in kinds: - points.extend(recheck_points(trace)) + come from the harness tags; ``recheck`` and ``judge`` are the structural final-answer leaf + (re-roll vs. judge-the-attempt). Each: ``node`` id and ``kind``.""" + points: list[dict] = [ + {"node": nid, "kind": get_tag(trace, nid)} + for nid in range(len(trace.nodes)) + if get_tag(trace, nid) in kinds + ] + leaves = graph.leaves(trace) + if leaves: + final = max(leaves) # the rollout's final-answer leaf + points += [{"node": final, "kind": k} for k in ("recheck", "judge") if k in kinds] return points +def render_transcript(trace: Trace, node_id: int) -> str: + """The conversation along root->node_id as plain text, for a judge prompt.""" + lines = [] + for m in seed_messages(trace, node_id): + content = m.content if isinstance(m.content, str) else (m.content or "") + lines.append(f"{m.role}: {content}") + return "\n".join(lines) + + +def judge_prompt(trace: Trace, node_id: int) -> str: + """A 'was this attempt correct?' prompt presenting the rollout's transcript.""" + return ( + "You are judging whether a previous attempt solved its task correctly.\n\n" + f"--- transcript ---\n{render_transcript(trace, node_id)}\n--- end transcript ---\n\n" + "Was the final answer correct? Reply with exactly 'yes' or 'no'." + ) + + def build_seed(trace: Trace, point: dict, followup: str) -> Messages: - """The seed conversation for a resume point: the ``root->node`` prefix, plus — for a - ``recheck`` point — an appended user turn asking the model to check and fix its work.""" + """The seed conversation for a resume point: + - ``judge`` -> a single user turn presenting the rollout to be graded; + - ``recheck`` -> the full rollout prefix + an appended check-your-work user turn; + - ``compaction_*`` -> the plain ``root->node`` prefix. + """ + if point["kind"] == "judge": + return [UserMessage(content=judge_prompt(trace, point["node"]))] msgs = seed_messages(trace, point["node"]) if point["kind"] == "recheck": msgs = [*msgs, UserMessage(content=followup)] diff --git a/environments/replay/replay/taskset.py b/environments/replay/replay/taskset.py index bf68664594..82854cd31e 100644 --- a/environments/replay/replay/taskset.py +++ b/environments/replay/replay/taskset.py @@ -1,30 +1,32 @@ -"""The replay-buffer taskset: resume old rollouts from tagged compaction points. +"""The replay-buffer taskset: turn old rollouts into new training tasks. -Two sourcing modes (``config.mode``): +Sourcing modes (``config.mode``): -- ``offline`` — ``load_tasks`` materializes one task per compaction resume point found in the - buffer at env-server start (snapshot of a prior run's, or this run's, rollouts). The replay - prefix (``root->node``) is the task ``prompt``; run with the default harness. +- ``offline`` — ``load_tasks`` materializes one task per resume point found in the buffer at + env-server start. The seed conversation is the task ``prompt``; run with the default harness. - ``online`` — ``load_tasks`` returns ``pool_size`` virtual slots and the bundled - :class:`ReplayHarness` samples a stored trace + resume point from the *live* buffer at rollout - time, so a growing buffer (this run's own rollouts) is picked up as it fills. - -Each compaction exposes two resume points: -- ``compaction_after`` -> prefix ends at the compaction message; the model continues solving. -- ``compaction_before`` -> prefix is the pre-compaction context; the model writes the compaction - itself, then continues. - -Scoring reuses the ORIGINAL env's verifier: ``config.inner`` names the taskset the rollouts came -from; ``score`` runs its ``@reward``/``@metric`` over the replay trace with the original task -swapped in (so e.g. a math verifier checks the replayed continuation's final answer against the -original ground truth). The original task + reward ride on the replay task (offline) or in -``trace.info["replay"]`` (online). + :class:`ReplayHarness` samples a stored trace + resume point from the *live* buffer per rollout, + so a growing buffer (this run's own rollouts) is picked up as it fills. + +Replay kinds (``config.kinds``): + +- ``recheck`` -> the full rollout + an appended "check your work" turn; re-roll. +- ``compaction_after`` -> resume from a compaction message; continue solving. +- ``compaction_before``-> resume before a compaction; the model writes the compaction, then continues. +- ``judge`` -> present the rollout's transcript and ask "was this correct?". + +Scoring: + +- ``recheck`` / ``compaction_*`` reuse the ORIGINAL env's verifier (``config.inner``): ``score`` + runs its rewards/metrics over the replay trace with the original task swapped in. +- ``judge`` compares the model's yes/no verdict against the original rollout's reward + (``original_reward > judge_threshold``) — a self-supervised correctness label, no inner verifier. + +The original task + reward ride on the replay task (offline) or in ``trace.info["replay"]`` (online). """ from __future__ import annotations -import glob -import json from typing import Literal from pydantic import SerializeAsAny @@ -33,14 +35,21 @@ from verifiers.v1.runtimes import Runtime from verifiers.v1.task import Task, WireTask from verifiers.v1.taskset import Taskset, TasksetConfig -from verifiers.v1.trace import Trace, WireTrace +from verifiers.v1.trace import Trace -from replay.selector import build_seed, resume_points, snapshot_ref_of +from replay.selector import ( + DEFAULT_FOLLOWUP, + DEFAULT_KINDS, + build_seed, + iter_traces, + resume_points, + snapshot_ref_of, +) class ReplayTask(Task): - """A replayed prefix plus the provenance needed to score it with the original verifier - (offline mode; online mode carries the same provenance in ``trace.info["replay"]``).""" + """A replayed seed plus the provenance needed to score it (offline mode; online mode carries + the same provenance in ``trace.info["replay"]``).""" source_trace_id: str = "" resume_node: int = -1 @@ -54,32 +63,34 @@ class ReplayTasksetConfig(TasksetConfig): mode: Literal["offline", "online"] = "offline" buffer_glob: str = "" """Glob of stored-rollout JSONL files, e.g. ``.../rollouts/step_*/train_rollouts.jsonl``.""" - kinds: list[str] = ["recheck", "compaction_after", "compaction_before"] - """Which resume-point kinds to replay from (``recheck`` = append a check-your-work turn to - the full rollout; ``compaction_*`` = resume around a compaction).""" - followup: str = ( - "Check your work. If anything is wrong, fix it and give the corrected final answer." - ) + kinds: list[str] = DEFAULT_KINDS + """Which replay kinds to source from (see module docstring); add ``"judge"`` to opt in.""" + followup: str = DEFAULT_FOLLOWUP """The user turn appended for ``recheck`` points.""" + judge_threshold: float = 0.5 + """``judge`` mode: the original rollout counts as correct when ``original_reward >`` this.""" pool_size: int = 1024 """Online mode: number of virtual task slots (num_tasks); the harness samples per rollout.""" inner: SerializeAsAny[TasksetConfig] = TasksetConfig() - """The ORIGINAL env's taskset (the verifier to reuse). Empty id => scoring is a no-op.""" + """The ORIGINAL env's taskset (the verifier to reuse). Empty id => no inner scoring.""" -def _iter_traces(buffer_glob: str): - for path in sorted(glob.glob(buffer_glob)): - with open(path) as f: - for line in f: - line = line.strip() - if line: - yield WireTrace.model_validate(json.loads(line)) +def _parse_verdict(trace: Trace) -> bool | None: + """The model's yes/no judgment from its last response (``judge`` mode). None if unclear.""" + msgs = trace.assistant_messages + text = (msgs[-1].content or "").strip().lower() if msgs else "" + if text.startswith("yes") or text.startswith("correct"): + return True + if text.startswith("no") or text.startswith("incorrect"): + return False + has_yes, has_no = "yes" in text, "no" in text + return has_yes if has_yes != has_no else None class ReplayTaskset(Taskset[ReplayTask, ReplayTasksetConfig]): def __init__(self, config: ReplayTasksetConfig) -> None: super().__init__(config) - # Reuse the original env's verifier when configured (skeleton: empty id => no scoring). + # Reuse the original env's verifier when configured (empty id => no inner scoring). self._inner: Taskset | None = load_taskset(config.inner) if config.inner.id else None def load_tasks(self) -> list[ReplayTask]: @@ -88,7 +99,7 @@ def load_tasks(self) -> list[ReplayTask]: return [ReplayTask(idx=i, prompt=None) for i in range(self.config.pool_size)] kinds = set(self.config.kinds) tasks: list[ReplayTask] = [] - for src in _iter_traces(self.config.buffer_glob): + for src in iter_traces(self.config.buffer_glob): for pt in resume_points(src, kinds=kinds): tasks.append( ReplayTask( @@ -106,20 +117,35 @@ def load_tasks(self) -> list[ReplayTask]: async def setup(self, task: ReplayTask, runtime: Runtime) -> None: # Offline exec/sandbox replay: restore to the resume point before the harness runs. - # (Online restores inside ReplayHarness, where the sampled ref is known.) Skeleton: + # (Online restores inside ReplayHarness; judge mode needs no sandbox.) Skeleton: # snapshot capture isn't wired yet, so refs are None -> nothing to restore. - if task.snapshot_ref is not None: + if task.snapshot_ref is not None and task.resume_kind != "judge": await runtime.restore(task.snapshot_ref) async def score(self, trace: Trace, runtime: Runtime) -> None: - # Reuse the ORIGINAL env's verifier: run its rewards/metrics over the replay trace with - # the original task swapped in. (Judge-style scoring would instead compare the model's - # verdict against ``original_reward`` here.) - if self._inner is None: - return replay = trace.info.get("replay") # online: harness-stashed provenance - original_dump = replay["original_task"] if replay else getattr(trace.task, "original_task", {}) - if not original_dump: + if replay: + kind, original_dump, original_reward = ( + replay["kind"], + replay["original_task"], + replay["original_reward"], + ) + else: + t = trace.task + kind = getattr(t, "resume_kind", "") + original_dump = getattr(t, "original_task", {}) + original_reward = getattr(t, "original_reward", 0.0) + + if kind == "judge": + # Grade the model's verdict against the original rollout's actual reward. + verdict = _parse_verdict(trace) + correct = original_reward > self.config.judge_threshold + trace.record_reward("judge_match", 1.0 if verdict == correct else 0.0) + return + + # recheck / compaction_*: reuse the ORIGINAL env's verifier over the replay trace, + # with the original task swapped in. + if self._inner is None or not original_dump: return replay_task = trace.task trace.task = WireTask.model_validate(original_dump) From 8a1ab7964009c80ca2cf1d85bf887f04033c48a9 Mon Sep 17 00:00:00 2001 From: mikasenghaas Date: Tue, 30 Jun 2026 00:30:18 +0000 Subject: [PATCH 6/9] fix(replay): online harness skips snapshot restore for judge mode (no sandbox needed) Co-Authored-By: Claude Opus 4.8 (1M context) --- environments/replay/replay/harness.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/environments/replay/replay/harness.py b/environments/replay/replay/harness.py index 54a3863a56..857bddd9d9 100644 --- a/environments/replay/replay/harness.py +++ b/environments/replay/replay/harness.py @@ -69,7 +69,7 @@ async def launch( src, point = sample ref = snapshot_ref_of(src, point["node"]) - if ref is not None: # exec/sandbox replay; skeleton refs are None -> skip + if ref is not None and point["kind"] != "judge": # judge needs no sandbox; skeleton refs are None await runtime.restore(ref) # Stash provenance so ReplayTaskset.score can reuse the original verifier. From ecddd735cbf7b39a88d533afdc21f5d0fde8dca1 Mon Sep 17 00:00:00 2001 From: mikasenghaas Date: Tue, 30 Jun 2026 00:41:29 +0000 Subject: [PATCH 7/9] refactor(replay): split into four separate tasksets, one per mode (Option B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Was a single ReplayTaskset with a `kinds` list; now each mode is its own selectable taskset (taskset id), so an env config picks exactly one task type — no mode mixing. Structure: one distribution (`environments/replay`) ships five top-level modules — - replay_common: shared base (BaseReplayTaskset/BaseReplayHarness parameterized by KIND, buffer sourcing offline+online, seeding, snapshot restore, scoring) + selector. - replay_recheck / replay_judge / replay_compaction_after / replay_compaction_before: thin modules each fixing KIND and bundling a harness (auto-selected via default_harness_id). recheck/compaction_* reuse the original verifier; judge grades the verdict against the original reward. The bundled harness handles both modes: offline (materialized task.prompt -> default chat loop) and online (sample this KIND from the live buffer). Dropped the now-unused DEFAULT_KINDS constant. Co-Authored-By: Claude Opus 4.8 (1M context) --- environments/replay/pyproject.toml | 12 +- environments/replay/replay/__init__.py | 18 -- environments/replay/replay/harness.py | 122 ---------- environments/replay/replay/taskset.py | 167 ------------- environments/replay/replay_common/__init__.py | 23 ++ environments/replay/replay_common/base.py | 219 ++++++++++++++++++ .../{replay => replay_common}/selector.py | 9 +- .../replay_compaction_after/__init__.py | 20 ++ .../replay_compaction_before/__init__.py | 21 ++ environments/replay/replay_judge/__init__.py | 21 ++ .../replay/replay_recheck/__init__.py | 20 ++ 11 files changed, 338 insertions(+), 314 deletions(-) delete mode 100644 environments/replay/replay/__init__.py delete mode 100644 environments/replay/replay/harness.py delete mode 100644 environments/replay/replay/taskset.py create mode 100644 environments/replay/replay_common/__init__.py create mode 100644 environments/replay/replay_common/base.py rename environments/replay/{replay => replay_common}/selector.py (89%) create mode 100644 environments/replay/replay_compaction_after/__init__.py create mode 100644 environments/replay/replay_compaction_before/__init__.py create mode 100644 environments/replay/replay_judge/__init__.py create mode 100644 environments/replay/replay_recheck/__init__.py diff --git a/environments/replay/pyproject.toml b/environments/replay/pyproject.toml index 7d409d74fb..dab2541a34 100644 --- a/environments/replay/pyproject.toml +++ b/environments/replay/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "replay" version = "0.1.0" -description = "replay — a replay-buffer taskset that resumes old rollouts from compaction points." +description = "replay — replay-buffer tasksets (recheck / judge / compaction_before / compaction_after)." requires-python = ">=3.10" dependencies = ["verifiers"] @@ -10,4 +10,12 @@ requires = ["hatchling"] build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] -packages = ["replay"] +# One distribution ships five top-level modules: the shared library + the four selectable +# tasksets (each a taskset id, e.g. `--taskset.id replay_recheck`). +packages = [ + "replay_common", + "replay_recheck", + "replay_judge", + "replay_compaction_after", + "replay_compaction_before", +] diff --git a/environments/replay/replay/__init__.py b/environments/replay/replay/__init__.py deleted file mode 100644 index 35b8b6974a..0000000000 --- a/environments/replay/replay/__init__.py +++ /dev/null @@ -1,18 +0,0 @@ -"""replay — resume old rollouts from tagged compaction points (a replay buffer). - -Offline (``ReplayTaskset`` materializes tasks from the buffer) or online (``ReplayHarness`` -samples the live buffer per rollout). Scoring reuses the original env's verifier. - -Option B: compaction tags + snapshot refs are read from typed ``MessageNode`` fields (see -``replay/selector.py``).""" - -from replay.harness import ReplayHarness, ReplayHarnessConfig -from replay.taskset import ReplayTask, ReplayTaskset, ReplayTasksetConfig - -__all__ = [ - "ReplayTaskset", - "ReplayTasksetConfig", - "ReplayTask", - "ReplayHarness", - "ReplayHarnessConfig", -] diff --git a/environments/replay/replay/harness.py b/environments/replay/replay/harness.py deleted file mode 100644 index 857bddd9d9..0000000000 --- a/environments/replay/replay/harness.py +++ /dev/null @@ -1,122 +0,0 @@ -"""ReplayHarness — rollout-time sampling for an online (growing) replay buffer. - -The offline path materializes tasks in ``ReplayTaskset.load_tasks``, which runs once at -env-server start and so can't see rollouts written later. This harness instead samples a stored -trace + resume point from the *live* buffer on each rollout, so this run's own rollouts are -replayed as the buffer fills. It restores the resume point's sandbox snapshot, seeds the default -chat loop with the replay prefix (``root->node``), and stashes provenance in ``trace.info`` so -``ReplayTaskset.score`` can reuse the original verifier. - -It reuses the default harness's program (a growing-message-list chat loop) and seeds it via the -same ``INITIAL_MESSAGES`` channel the default harness uses for a Messages prompt. -""" - -from __future__ import annotations - -import glob -import json -import random -from pathlib import Path - -from verifiers.v1.clients import RolloutContext -from verifiers.v1.dialects.chat import message_to_wire -from verifiers.v1.harnesses.default.harness import ( - PROGRAM_SOURCE, - DefaultHarness, - DefaultHarnessConfig, -) -from verifiers.v1.runtimes import ProgramResult, Runtime -from verifiers.v1.trace import Trace, WireTrace - -from replay.selector import ( - DEFAULT_FOLLOWUP, - DEFAULT_KINDS, - build_seed, - resume_points, - snapshot_ref_of, -) - - -class ReplayHarnessConfig(DefaultHarnessConfig): - buffer_glob: str = "" - """Glob of stored-rollout JSONL files to sample from (the live, possibly growing, buffer).""" - kinds: list[str] = DEFAULT_KINDS - """Which replay kinds to sample (see ReplayTaskset; add ``"judge"`` to opt in).""" - followup: str = DEFAULT_FOLLOWUP - """The user turn appended for ``recheck`` points.""" - - -class ReplayHarness(DefaultHarness): - """Subclasses the default harness (its chat-loop program) but seeds from a sampled buffer - rollout instead of ``task.prompt``.""" - - SUPPORTS_MESSAGE_PROMPT = True - - async def launch( - self, - ctx: RolloutContext, - trace: Trace, - runtime: Runtime, - endpoint: str, - secret: str, - mcp_urls: dict[str, str], - ) -> ProgramResult: - rng = random.Random(trace.id) # deterministic per rollout, varies across rollouts - sample = self._sample(rng) - if sample is None: # buffer empty (warmup) or no matching resume points yet - trace.stop("replay_buffer_empty") - return ProgramResult(exit_code=0, stdout="", stderr="") - src, point = sample - - ref = snapshot_ref_of(src, point["node"]) - if ref is not None and point["kind"] != "judge": # judge needs no sandbox; skeleton refs are None - await runtime.restore(ref) - - # Stash provenance so ReplayTaskset.score can reuse the original verifier. - trace.info["replay"] = { - "source_id": src.id, - "resume_node": point["node"], - "kind": point["kind"], - "original_task": src.task.model_dump(), - "original_reward": src.reward, - } - - # Seed the default chat loop with the replay prefix (mirror DefaultHarness.launch). - seed = build_seed(src, point, self.config.followup) - env = {**self.config.env} - env["INITIAL_MESSAGES"] = json.dumps([message_to_wire(m) for m in seed]) - args = [f"--base-url={endpoint}", f"--api-key={secret}", f"--model={ctx.model}"] - if mcp_urls: - args.append( - "--mcp-config=" - + json.dumps( - {"mcpServers": {name: {"url": url} for name, url in mcp_urls.items()}} - ) - ) - program = await runtime.prepare_uv_script(PROGRAM_SOURCE, self.config.env) - return await runtime.run_program([*program, *args], env) - - def _sample(self, rng: random.Random) -> tuple[Trace, dict] | None: - """Scan the live buffer in random order; return the first (trace, resume point) found. - Re-globs every rollout, so files written after env-server start are included.""" - kinds = set(self.config.kinds) - files = sorted(glob.glob(self.config.buffer_glob)) - rng.shuffle(files) - for path in files: - try: - lines = Path(path).read_text().splitlines() - except OSError: - continue - rng.shuffle(lines) - for line in lines: - line = line.strip() - if not line: - continue - src = WireTrace.model_validate(json.loads(line)) - points = resume_points(src, kinds=kinds) - if points: - return src, rng.choice(points) - return None - - -__all__ = ["ReplayHarness", "ReplayHarnessConfig"] diff --git a/environments/replay/replay/taskset.py b/environments/replay/replay/taskset.py deleted file mode 100644 index 82854cd31e..0000000000 --- a/environments/replay/replay/taskset.py +++ /dev/null @@ -1,167 +0,0 @@ -"""The replay-buffer taskset: turn old rollouts into new training tasks. - -Sourcing modes (``config.mode``): - -- ``offline`` — ``load_tasks`` materializes one task per resume point found in the buffer at - env-server start. The seed conversation is the task ``prompt``; run with the default harness. -- ``online`` — ``load_tasks`` returns ``pool_size`` virtual slots and the bundled - :class:`ReplayHarness` samples a stored trace + resume point from the *live* buffer per rollout, - so a growing buffer (this run's own rollouts) is picked up as it fills. - -Replay kinds (``config.kinds``): - -- ``recheck`` -> the full rollout + an appended "check your work" turn; re-roll. -- ``compaction_after`` -> resume from a compaction message; continue solving. -- ``compaction_before``-> resume before a compaction; the model writes the compaction, then continues. -- ``judge`` -> present the rollout's transcript and ask "was this correct?". - -Scoring: - -- ``recheck`` / ``compaction_*`` reuse the ORIGINAL env's verifier (``config.inner``): ``score`` - runs its rewards/metrics over the replay trace with the original task swapped in. -- ``judge`` compares the model's yes/no verdict against the original rollout's reward - (``original_reward > judge_threshold``) — a self-supervised correctness label, no inner verifier. - -The original task + reward ride on the replay task (offline) or in ``trace.info["replay"]`` (online). -""" - -from __future__ import annotations - -from typing import Literal - -from pydantic import SerializeAsAny - -from verifiers.v1.loaders import load_taskset -from verifiers.v1.runtimes import Runtime -from verifiers.v1.task import Task, WireTask -from verifiers.v1.taskset import Taskset, TasksetConfig -from verifiers.v1.trace import Trace - -from replay.selector import ( - DEFAULT_FOLLOWUP, - DEFAULT_KINDS, - build_seed, - iter_traces, - resume_points, - snapshot_ref_of, -) - - -class ReplayTask(Task): - """A replayed seed plus the provenance needed to score it (offline mode; online mode carries - the same provenance in ``trace.info["replay"]``).""" - - source_trace_id: str = "" - resume_node: int = -1 - resume_kind: str = "" - snapshot_ref: str | None = None - original_task: dict = {} - original_reward: float = 0.0 - - -class ReplayTasksetConfig(TasksetConfig): - mode: Literal["offline", "online"] = "offline" - buffer_glob: str = "" - """Glob of stored-rollout JSONL files, e.g. ``.../rollouts/step_*/train_rollouts.jsonl``.""" - kinds: list[str] = DEFAULT_KINDS - """Which replay kinds to source from (see module docstring); add ``"judge"`` to opt in.""" - followup: str = DEFAULT_FOLLOWUP - """The user turn appended for ``recheck`` points.""" - judge_threshold: float = 0.5 - """``judge`` mode: the original rollout counts as correct when ``original_reward >`` this.""" - pool_size: int = 1024 - """Online mode: number of virtual task slots (num_tasks); the harness samples per rollout.""" - inner: SerializeAsAny[TasksetConfig] = TasksetConfig() - """The ORIGINAL env's taskset (the verifier to reuse). Empty id => no inner scoring.""" - - -def _parse_verdict(trace: Trace) -> bool | None: - """The model's yes/no judgment from its last response (``judge`` mode). None if unclear.""" - msgs = trace.assistant_messages - text = (msgs[-1].content or "").strip().lower() if msgs else "" - if text.startswith("yes") or text.startswith("correct"): - return True - if text.startswith("no") or text.startswith("incorrect"): - return False - has_yes, has_no = "yes" in text, "no" in text - return has_yes if has_yes != has_no else None - - -class ReplayTaskset(Taskset[ReplayTask, ReplayTasksetConfig]): - def __init__(self, config: ReplayTasksetConfig) -> None: - super().__init__(config) - # Reuse the original env's verifier when configured (empty id => no inner scoring). - self._inner: Taskset | None = load_taskset(config.inner) if config.inner.id else None - - def load_tasks(self) -> list[ReplayTask]: - if self.config.mode == "online": - # Virtual slots; ReplayHarness samples the live buffer per rollout. - return [ReplayTask(idx=i, prompt=None) for i in range(self.config.pool_size)] - kinds = set(self.config.kinds) - tasks: list[ReplayTask] = [] - for src in iter_traces(self.config.buffer_glob): - for pt in resume_points(src, kinds=kinds): - tasks.append( - ReplayTask( - idx=len(tasks), - prompt=build_seed(src, pt, self.config.followup), - source_trace_id=src.id, - resume_node=pt["node"], - resume_kind=pt["kind"], - snapshot_ref=snapshot_ref_of(src, pt["node"]), - original_task=src.task.model_dump(), - original_reward=src.reward, - ) - ) - return tasks - - async def setup(self, task: ReplayTask, runtime: Runtime) -> None: - # Offline exec/sandbox replay: restore to the resume point before the harness runs. - # (Online restores inside ReplayHarness; judge mode needs no sandbox.) Skeleton: - # snapshot capture isn't wired yet, so refs are None -> nothing to restore. - if task.snapshot_ref is not None and task.resume_kind != "judge": - await runtime.restore(task.snapshot_ref) - - async def score(self, trace: Trace, runtime: Runtime) -> None: - replay = trace.info.get("replay") # online: harness-stashed provenance - if replay: - kind, original_dump, original_reward = ( - replay["kind"], - replay["original_task"], - replay["original_reward"], - ) - else: - t = trace.task - kind = getattr(t, "resume_kind", "") - original_dump = getattr(t, "original_task", {}) - original_reward = getattr(t, "original_reward", 0.0) - - if kind == "judge": - # Grade the model's verdict against the original rollout's actual reward. - verdict = _parse_verdict(trace) - correct = original_reward > self.config.judge_threshold - trace.record_reward("judge_match", 1.0 if verdict == correct else 0.0) - return - - # recheck / compaction_*: reuse the ORIGINAL env's verifier over the replay trace, - # with the original task swapped in. - if self._inner is None or not original_dump: - return - replay_task = trace.task - trace.task = WireTask.model_validate(original_dump) - try: - await self._inner.score(trace, runtime) - finally: - trace.task = replay_task - - -# Bundle the online sampling harness so `--harness.id replay` resolves alongside the taskset. -from replay.harness import ReplayHarness, ReplayHarnessConfig # noqa: E402 - -__all__ = [ - "ReplayTaskset", - "ReplayTasksetConfig", - "ReplayTask", - "ReplayHarness", - "ReplayHarnessConfig", -] diff --git a/environments/replay/replay_common/__init__.py b/environments/replay/replay_common/__init__.py new file mode 100644 index 0000000000..4ad4f0626d --- /dev/null +++ b/environments/replay/replay_common/__init__.py @@ -0,0 +1,23 @@ +"""replay_common — shared base for the replay-buffer tasksets. + +The four selectable tasksets live in sibling top-level modules, each fixing one ``KIND``: +``replay_recheck``, ``replay_judge``, ``replay_compaction_after``, ``replay_compaction_before``. +This module is the shared library (buffer sourcing, seeding, scoring); it is not itself a +selectable taskset. +""" + +from replay_common.base import ( + BaseReplayHarness, + BaseReplayTaskset, + ReplayConfig, + ReplayHarnessConfig, + ReplayTask, +) + +__all__ = [ + "BaseReplayTaskset", + "BaseReplayHarness", + "ReplayConfig", + "ReplayHarnessConfig", + "ReplayTask", +] diff --git a/environments/replay/replay_common/base.py b/environments/replay/replay_common/base.py new file mode 100644 index 0000000000..972600b75b --- /dev/null +++ b/environments/replay/replay_common/base.py @@ -0,0 +1,219 @@ +"""Shared base for the replay-buffer tasksets. + +Each concrete taskset/harness fixes one ``KIND`` (``recheck`` / ``judge`` / ``compaction_after`` +/ ``compaction_before``) and lives in its own selectable module (``replay_recheck``, ...). This +base holds everything they share: buffer sourcing (offline materialize + online sampling), +seed building, snapshot restore, and scoring. + +Sourcing (``config.mode``): +- ``offline`` — ``load_tasks`` materializes one task per resume point of this taskset's ``KIND``; + the bundled harness sees a non-empty ``task.prompt`` and just runs the default chat loop. +- ``online`` — ``load_tasks`` returns ``pool_size`` virtual slots and the harness samples a + ``KIND`` resume point from the *live* buffer per rollout. + +Scoring: +- ``recheck`` / ``compaction_*`` reuse the ORIGINAL env's verifier (``config.inner``). +- ``judge`` grades the model's verdict against the original rollout's reward (no inner verifier). +""" + +from __future__ import annotations + +import glob +import json +import random +from pathlib import Path +from typing import ClassVar, Literal + +from pydantic import SerializeAsAny + +from verifiers.v1.clients import RolloutContext +from verifiers.v1.dialects.chat import message_to_wire +from verifiers.v1.harnesses.default.harness import ( + PROGRAM_SOURCE, + DefaultHarness, + DefaultHarnessConfig, +) +from verifiers.v1.loaders import load_taskset +from verifiers.v1.runtimes import ProgramResult, Runtime +from verifiers.v1.task import Task, WireTask +from verifiers.v1.taskset import Taskset, TasksetConfig +from verifiers.v1.trace import Trace, WireTrace + +from replay_common.selector import ( + DEFAULT_FOLLOWUP, + build_seed, + iter_traces, + resume_points, + snapshot_ref_of, +) + + +class ReplayTask(Task): + """A replayed seed plus the provenance needed to score it (offline mode; online mode carries + the same provenance in ``trace.info["replay"]``).""" + + source_trace_id: str = "" + resume_node: int = -1 + snapshot_ref: str | None = None + original_task: dict = {} + original_reward: float = 0.0 + + +class ReplayConfig(TasksetConfig): + mode: Literal["offline", "online"] = "offline" + buffer_glob: str = "" + """Glob of stored-rollout JSONL files, e.g. ``.../rollouts/step_*/train_rollouts.jsonl``.""" + followup: str = DEFAULT_FOLLOWUP + """The user turn appended for the ``recheck`` taskset.""" + judge_threshold: float = 0.5 + """``judge`` taskset: the original rollout counts as correct when ``original_reward >`` this.""" + pool_size: int = 1024 + """Online mode: number of virtual task slots (num_tasks); the harness samples per rollout.""" + inner: SerializeAsAny[TasksetConfig] = TasksetConfig() + """The ORIGINAL env's taskset (the verifier to reuse). Empty id => no inner scoring.""" + + +class ReplayHarnessConfig(DefaultHarnessConfig): + buffer_glob: str = "" + """Online mode: the live buffer to sample from.""" + followup: str = DEFAULT_FOLLOWUP + + +def _parse_verdict(trace: Trace) -> bool | None: + """The model's yes/no judgment from its last response (``judge``). None if unclear.""" + msgs = trace.assistant_messages + text = (msgs[-1].content or "").strip().lower() if msgs else "" + if text.startswith("yes") or text.startswith("correct"): + return True + if text.startswith("no") or text.startswith("incorrect"): + return False + has_yes, has_no = "yes" in text, "no" in text + return has_yes if has_yes != has_no else None + + +class BaseReplayTaskset(Taskset[ReplayTask, ReplayConfig]): + """Concrete subclasses set ``KIND``; everything else is shared.""" + + KIND: ClassVar[str] = "" + + def __init__(self, config: ReplayConfig) -> None: + super().__init__(config) + # Reuse the original env's verifier when configured (empty id => no inner scoring). + self._inner: Taskset | None = load_taskset(config.inner) if config.inner.id else None + + def load_tasks(self) -> list[ReplayTask]: + if self.config.mode == "online": + return [ReplayTask(idx=i, prompt=None) for i in range(self.config.pool_size)] + tasks: list[ReplayTask] = [] + for src in iter_traces(self.config.buffer_glob): + for pt in resume_points(src, kinds={self.KIND}): + tasks.append( + ReplayTask( + idx=len(tasks), + prompt=build_seed(src, pt, self.config.followup), + source_trace_id=src.id, + resume_node=pt["node"], + snapshot_ref=snapshot_ref_of(src, pt["node"]), + original_task=src.task.model_dump(), + original_reward=src.reward, + ) + ) + return tasks + + async def setup(self, task: ReplayTask, runtime: Runtime) -> None: + # Offline exec/sandbox replay: restore to the resume point before the harness runs. + # (Online restores inside the harness; judge needs no sandbox.) Skeleton: refs are None. + if task.snapshot_ref is not None and self.KIND != "judge": + await runtime.restore(task.snapshot_ref) + + async def score(self, trace: Trace, runtime: Runtime) -> None: + replay = trace.info.get("replay") # online: harness-stashed provenance + original_dump = replay["original_task"] if replay else getattr(trace.task, "original_task", {}) + original_reward = replay["original_reward"] if replay else getattr(trace.task, "original_reward", 0.0) + + if self.KIND == "judge": + verdict = _parse_verdict(trace) + correct = original_reward > self.config.judge_threshold + trace.record_reward("judge_match", 1.0 if verdict == correct else 0.0) + return + + # recheck / compaction_*: reuse the ORIGINAL verifier with the original task swapped in. + if self._inner is None or not original_dump: + return + replay_task = trace.task + trace.task = WireTask.model_validate(original_dump) + try: + await self._inner.score(trace, runtime) + finally: + trace.task = replay_task + + +class BaseReplayHarness(DefaultHarness): + """Bundled with each taskset (auto-selected via ``default_harness_id``). Offline: a materialized + ``task.prompt`` is present, so defer to the default chat loop. Online: sample a ``KIND`` resume + point from the live buffer and seed the loop with it.""" + + KIND: ClassVar[str] = "" + SUPPORTS_MESSAGE_PROMPT = True + + async def launch( + self, + ctx: RolloutContext, + trace: Trace, + runtime: Runtime, + endpoint: str, + secret: str, + mcp_urls: dict[str, str], + ) -> ProgramResult: + if trace.task.prompt is not None: # offline: materialized seed -> default behavior + return await super().launch(ctx, trace, runtime, endpoint, secret, mcp_urls) + + rng = random.Random(trace.id) + sample = self._sample(rng) + if sample is None: # buffer empty (warmup) / no matching points yet + trace.stop("replay_buffer_empty") + return ProgramResult(exit_code=0, stdout="", stderr="") + src, point = sample + + ref = snapshot_ref_of(src, point["node"]) + if ref is not None and self.KIND != "judge": # judge needs no sandbox; skeleton refs are None + await runtime.restore(ref) + + trace.info["replay"] = { + "source_id": src.id, + "resume_node": point["node"], + "kind": point["kind"], + "original_task": src.task.model_dump(), + "original_reward": src.reward, + } + seed = build_seed(src, point, self.config.followup) + env = {**self.config.env} + env["INITIAL_MESSAGES"] = json.dumps([message_to_wire(m) for m in seed]) + args = [f"--base-url={endpoint}", f"--api-key={secret}", f"--model={ctx.model}"] + if mcp_urls: + args.append( + "--mcp-config=" + + json.dumps({"mcpServers": {n: {"url": u} for n, u in mcp_urls.items()}}) + ) + program = await runtime.prepare_uv_script(PROGRAM_SOURCE, self.config.env) + return await runtime.run_program([*program, *args], env) + + def _sample(self, rng: random.Random) -> tuple[Trace, dict] | None: + """Scan the live buffer in random order; return the first (trace, ``KIND`` point) found.""" + files = sorted(glob.glob(self.config.buffer_glob)) + rng.shuffle(files) + for path in files: + try: + lines = Path(path).read_text().splitlines() + except OSError: + continue + rng.shuffle(lines) + for line in lines: + line = line.strip() + if not line: + continue + src = WireTrace.model_validate(json.loads(line)) + points = resume_points(src, kinds={self.KIND}) + if points: + return src, rng.choice(points) + return None diff --git a/environments/replay/replay/selector.py b/environments/replay/replay_common/selector.py similarity index 89% rename from environments/replay/replay/selector.py rename to environments/replay/replay_common/selector.py index d4c09ac8c0..1bc9c0ce11 100644 --- a/environments/replay/replay/selector.py +++ b/environments/replay/replay_common/selector.py @@ -1,10 +1,10 @@ """Consumer-side buffer reading + resume-point selection (Option B: ``MessageNode`` fields). Turns the compaction tags the producing harness stamped on ``MessageNode.kind`` into resume -points the ReplayTaskset/ReplayHarness sample. ``build_seed`` produces the seed conversation per -mode; ``snapshot_ref_of`` returns the durable sandbox handle to restore (None until per-turn -snapshot capture is wired). Only ``get_tag``/``snapshot_ref_of`` differ from Option A — everything -else (buffer reading, resume-point logic, seed building) is shared. +points the tasksets/harnesses sample. ``build_seed`` produces the seed conversation per mode; +``snapshot_ref_of`` returns the durable sandbox handle to restore (None until per-turn snapshot +capture is wired). Only ``get_tag``/``snapshot_ref_of`` differ from Option A — everything else +(buffer reading, resume-point logic, seed building) is shared. """ from __future__ import annotations @@ -16,7 +16,6 @@ from verifiers.v1.trace import Trace, WireTrace from verifiers.v1.types import Messages, UserMessage -DEFAULT_KINDS = ["recheck", "compaction_after", "compaction_before"] DEFAULT_FOLLOWUP = "Check your work. If anything is wrong, fix it and give the corrected final answer." diff --git a/environments/replay/replay_compaction_after/__init__.py b/environments/replay/replay_compaction_after/__init__.py new file mode 100644 index 0000000000..f1e5f0a2ba --- /dev/null +++ b/environments/replay/replay_compaction_after/__init__.py @@ -0,0 +1,20 @@ +"""replay_compaction_after — resume from a compaction message. + +Seed the model with the post-compaction context (``[system, user(notes)]``) and let it continue +solving; score with the original env's verifier (``config.inner``). +""" + +from typing import ClassVar + +from replay_common.base import BaseReplayHarness, BaseReplayTaskset + + +class CompactionAfterTaskset(BaseReplayTaskset): + KIND: ClassVar[str] = "compaction_after" + + +class CompactionAfterHarness(BaseReplayHarness): + KIND: ClassVar[str] = "compaction_after" + + +__all__ = ["CompactionAfterTaskset", "CompactionAfterHarness"] diff --git a/environments/replay/replay_compaction_before/__init__.py b/environments/replay/replay_compaction_before/__init__.py new file mode 100644 index 0000000000..d29083c66f --- /dev/null +++ b/environments/replay/replay_compaction_before/__init__.py @@ -0,0 +1,21 @@ +"""replay_compaction_before — resume just before a compaction. + +Seed the model with the pre-compaction context (the branch leaf the compaction summarized) so its +continuation *writes* the compaction itself, then keeps solving; score with the original env's +verifier (``config.inner``). +""" + +from typing import ClassVar + +from replay_common.base import BaseReplayHarness, BaseReplayTaskset + + +class CompactionBeforeTaskset(BaseReplayTaskset): + KIND: ClassVar[str] = "compaction_before" + + +class CompactionBeforeHarness(BaseReplayHarness): + KIND: ClassVar[str] = "compaction_before" + + +__all__ = ["CompactionBeforeTaskset", "CompactionBeforeHarness"] diff --git a/environments/replay/replay_judge/__init__.py b/environments/replay/replay_judge/__init__.py new file mode 100644 index 0000000000..dca3c398bb --- /dev/null +++ b/environments/replay/replay_judge/__init__.py @@ -0,0 +1,21 @@ +"""replay_judge — the correctness-judge taskset. + +Present a sampled rollout's transcript and ask "was this correct? yes/no"; grade the model's +verdict against the rollout's actual reward (``original_reward > judge_threshold``). A +self-supervised correctness label — no inner verifier or sandbox needed. +""" + +from typing import ClassVar + +from replay_common.base import BaseReplayHarness, BaseReplayTaskset + + +class JudgeTaskset(BaseReplayTaskset): + KIND: ClassVar[str] = "judge" + + +class JudgeHarness(BaseReplayHarness): + KIND: ClassVar[str] = "judge" + + +__all__ = ["JudgeTaskset", "JudgeHarness"] diff --git a/environments/replay/replay_recheck/__init__.py b/environments/replay/replay_recheck/__init__.py new file mode 100644 index 0000000000..54913aff40 --- /dev/null +++ b/environments/replay/replay_recheck/__init__.py @@ -0,0 +1,20 @@ +"""replay_recheck — the "try again" taskset. + +Re-roll a sampled rollout after appending a "check your work, fix if wrong" user turn; score +the corrected attempt with the original env's verifier (``config.inner``). +""" + +from typing import ClassVar + +from replay_common.base import BaseReplayHarness, BaseReplayTaskset + + +class RecheckTaskset(BaseReplayTaskset): + KIND: ClassVar[str] = "recheck" + + +class RecheckHarness(BaseReplayHarness): + KIND: ClassVar[str] = "recheck" + + +__all__ = ["RecheckTaskset", "RecheckHarness"] From e3e6262428e03dbc690e0863b70baaba1b5bad88 Mon Sep 17 00:00:00 2001 From: mikasenghaas Date: Tue, 30 Jun 2026 22:57:55 +0000 Subject: [PATCH 8/9] fix(replay): correct harness config type; simplify resume_points (Option B) - Bug: BaseReplayHarness subclassed DefaultHarness without parameterizing Harness[ReplayHarnessConfig], so harness_config_type resolved to DefaultHarnessConfig (no buffer_glob/followup) -> self.config.buffer_glob would AttributeError in the online path. Now subclasses Harness[ReplayHarnessConfig] directly and seeds the default program itself (one unified offline/online seed path via _resolve_seed). - resume_points: single get_tag per node instead of two. Co-Authored-By: Claude Opus 4.8 (1M context) --- environments/replay/replay_common/base.py | 66 ++++++++++--------- environments/replay/replay_common/selector.py | 10 +-- 2 files changed, 40 insertions(+), 36 deletions(-) diff --git a/environments/replay/replay_common/base.py b/environments/replay/replay_common/base.py index 972600b75b..71bd1b46af 100644 --- a/environments/replay/replay_common/base.py +++ b/environments/replay/replay_common/base.py @@ -7,7 +7,7 @@ Sourcing (``config.mode``): - ``offline`` — ``load_tasks`` materializes one task per resume point of this taskset's ``KIND``; - the bundled harness sees a non-empty ``task.prompt`` and just runs the default chat loop. + the bundled harness seeds the chat loop from the task's prompt. - ``online`` — ``load_tasks`` returns ``pool_size`` virtual slots and the harness samples a ``KIND`` resume point from the *live* buffer per rollout. @@ -28,16 +28,14 @@ from verifiers.v1.clients import RolloutContext from verifiers.v1.dialects.chat import message_to_wire -from verifiers.v1.harnesses.default.harness import ( - PROGRAM_SOURCE, - DefaultHarness, - DefaultHarnessConfig, -) +from verifiers.v1.harness import Harness, HarnessConfig +from verifiers.v1.harnesses.default.harness import PROGRAM_SOURCE from verifiers.v1.loaders import load_taskset from verifiers.v1.runtimes import ProgramResult, Runtime from verifiers.v1.task import Task, WireTask from verifiers.v1.taskset import Taskset, TasksetConfig from verifiers.v1.trace import Trace, WireTrace +from verifiers.v1.types import Messages from replay_common.selector import ( DEFAULT_FOLLOWUP, @@ -73,7 +71,7 @@ class ReplayConfig(TasksetConfig): """The ORIGINAL env's taskset (the verifier to reuse). Empty id => no inner scoring.""" -class ReplayHarnessConfig(DefaultHarnessConfig): +class ReplayHarnessConfig(HarnessConfig): buffer_glob: str = "" """Online mode: the live buffer to sample from.""" followup: str = DEFAULT_FOLLOWUP @@ -148,13 +146,17 @@ async def score(self, trace: Trace, runtime: Runtime) -> None: trace.task = replay_task -class BaseReplayHarness(DefaultHarness): - """Bundled with each taskset (auto-selected via ``default_harness_id``). Offline: a materialized - ``task.prompt`` is present, so defer to the default chat loop. Online: sample a ``KIND`` resume - point from the live buffer and seed the loop with it.""" +class BaseReplayHarness(Harness[ReplayHarnessConfig]): + """Seeds the default chat-loop program (``PROGRAM_SOURCE``) from a replay prefix. Offline: the + task already carries the materialized seed. Online: sample a ``KIND`` resume point from the + live buffer. Concrete subclasses set ``KIND``.""" KIND: ClassVar[str] = "" SUPPORTS_MESSAGE_PROMPT = True + SUPPORTS_MCP = True + + async def setup(self, runtime: Runtime) -> None: + await runtime.prepare_uv_script(PROGRAM_SOURCE, self.config.env) async def launch( self, @@ -165,20 +167,32 @@ async def launch( secret: str, mcp_urls: dict[str, str], ) -> ProgramResult: - if trace.task.prompt is not None: # offline: materialized seed -> default behavior - return await super().launch(ctx, trace, runtime, endpoint, secret, mcp_urls) + seed = await self._resolve_seed(trace, runtime) + if seed is None: # warmup: live buffer empty / no matching points yet + return ProgramResult(exit_code=0, stdout="", stderr="") + env = {**self.config.env} + env["INITIAL_MESSAGES"] = json.dumps([message_to_wire(m) for m in seed]) + args = [f"--base-url={endpoint}", f"--api-key={secret}", f"--model={ctx.model}"] + if mcp_urls: + args.append( + "--mcp-config=" + + json.dumps({"mcpServers": {n: {"url": u} for n, u in mcp_urls.items()}}) + ) + program = await runtime.prepare_uv_script(PROGRAM_SOURCE, self.config.env) + return await runtime.run_program([*program, *args], env) - rng = random.Random(trace.id) - sample = self._sample(rng) - if sample is None: # buffer empty (warmup) / no matching points yet + async def _resolve_seed(self, trace: Trace, runtime: Runtime) -> Messages | None: + if trace.task.prompt is not None: # offline: task carries the materialized seed + return list(trace.task.prompt) + # online: sample a KIND resume point from the live buffer + sample = self._sample(random.Random(trace.id)) + if sample is None: trace.stop("replay_buffer_empty") - return ProgramResult(exit_code=0, stdout="", stderr="") + return None src, point = sample - ref = snapshot_ref_of(src, point["node"]) - if ref is not None and self.KIND != "judge": # judge needs no sandbox; skeleton refs are None + if ref is not None and self.KIND != "judge": # judge needs no sandbox; skeleton refs None await runtime.restore(ref) - trace.info["replay"] = { "source_id": src.id, "resume_node": point["node"], @@ -186,17 +200,7 @@ async def launch( "original_task": src.task.model_dump(), "original_reward": src.reward, } - seed = build_seed(src, point, self.config.followup) - env = {**self.config.env} - env["INITIAL_MESSAGES"] = json.dumps([message_to_wire(m) for m in seed]) - args = [f"--base-url={endpoint}", f"--api-key={secret}", f"--model={ctx.model}"] - if mcp_urls: - args.append( - "--mcp-config=" - + json.dumps({"mcpServers": {n: {"url": u} for n, u in mcp_urls.items()}}) - ) - program = await runtime.prepare_uv_script(PROGRAM_SOURCE, self.config.env) - return await runtime.run_program([*program, *args], env) + return build_seed(src, point, self.config.followup) def _sample(self, rng: random.Random) -> tuple[Trace, dict] | None: """Scan the live buffer in random order; return the first (trace, ``KIND`` point) found.""" diff --git a/environments/replay/replay_common/selector.py b/environments/replay/replay_common/selector.py index 1bc9c0ce11..99c12df7b8 100644 --- a/environments/replay/replay_common/selector.py +++ b/environments/replay/replay_common/selector.py @@ -55,11 +55,11 @@ def resume_points(trace: Trace, *, kinds: set[str]) -> list[dict]: """Resume points whose ``kind`` is in ``kinds``. ``compaction_before``/``compaction_after`` come from the harness tags; ``recheck`` and ``judge`` are the structural final-answer leaf (re-roll vs. judge-the-attempt). Each: ``node`` id and ``kind``.""" - points: list[dict] = [ - {"node": nid, "kind": get_tag(trace, nid)} - for nid in range(len(trace.nodes)) - if get_tag(trace, nid) in kinds - ] + points: list[dict] = [] + for nid in range(len(trace.nodes)): + tag = get_tag(trace, nid) + if tag in kinds: + points.append({"node": nid, "kind": tag}) leaves = graph.leaves(trace) if leaves: final = max(leaves) # the rollout's final-answer leaf From 7b79984feef803086368260af96766781d3c3188 Mon Sep 17 00:00:00 2001 From: mikasenghaas Date: Wed, 1 Jul 2026 00:26:04 +0000 Subject: [PATCH 9/9] fix(replay): online sampling keyed by task index for GRPO-consistent groups (Option B) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Was seeded by trace.id (unique per rollout), so the group_size rollouts of one group each sampled a different source rollout — GRPO's group-relative baseline would then compare across different problems. Seed by trace.task.idx (shared across a group) so all rollouts of a group replay the same source+point (N diverse continuations of one point). Freshness is preserved by the growing buffer (glob picks up new files, shifting each index's draw). Offline was already correct (a group = N rollouts of one materialized task). Co-Authored-By: Claude Opus 4.8 (1M context) --- environments/replay/replay_common/base.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/environments/replay/replay_common/base.py b/environments/replay/replay_common/base.py index 71bd1b46af..07eb20fd62 100644 --- a/environments/replay/replay_common/base.py +++ b/environments/replay/replay_common/base.py @@ -184,8 +184,12 @@ async def launch( async def _resolve_seed(self, trace: Trace, runtime: Runtime) -> Messages | None: if trace.task.prompt is not None: # offline: task carries the materialized seed return list(trace.task.prompt) - # online: sample a KIND resume point from the live buffer - sample = self._sample(random.Random(trace.id)) + # online: sample a KIND resume point from the live buffer. Seed the RNG with the task + # index (shared across a GRPO group), NOT the rollout id, so every rollout of a group + # replays the SAME source+point — N diverse continuations of one point, which is what the + # group-relative baseline needs. Freshness still comes from the buffer growing: glob picks + # up new files, which shifts each index's draw over time. + sample = self._sample(random.Random(trace.task.idx)) if sample is None: trace.stop("replay_buffer_empty") return None