verifiers 0.3.1.dev57__tar.gz → 0.3.1.dev59__tar.gz

{verifiers-0.3.1.dev57 → verifiers-0.3.1.dev59}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: verifiers
-Version: 0.3.1.dev57
+Version: 0.3.1.dev59
 Summary: Verifiers: Environments for LLM Reinforcement Learning
 Project-URL: Homepage, https://github.com/primeintellect-ai/verifiers
 Project-URL: Documentation, https://github.com/primeintellect-ai/verifiers

{verifiers-0.3.1.dev57 → verifiers-0.3.1.dev59}/verifiers/v1/__init__.py

@@ -33,7 +33,19 @@ from verifiers.v1.configs.task import (
 from verifiers.v1.configs.taskset import TasksetConfig
 from verifiers.v1.env import Env
 from verifiers.v1.envs.single_agent import SingleAgentEnv, SingleAgentEnvConfig
-from verifiers.v1.episode import Episode, WireEpisode
+from verifiers.v1.episode import (
+    EnvInfo,
+    Episode,
+    EvalRunInfo,
+    EvalWorkInfo,
+    GroupInfo,
+    PolicySpan,
+    RunInfo,
+    TrainRunInfo,
+    TrainWorkInfo,
+    WireEpisode,
+    WorkInfo,
+)
 from verifiers.v1.errors import (
     EnvError,
     HarnessError,
@@ -79,18 +91,15 @@ from verifiers.v1.trace import (
     AgentSpan,
     Branch,
     Error,
-    EvalRunInfo,
     InterceptRecord,
     ModelCall,
     PolicyEvent,
     Reward,
-    RunInfo,
     TimeSpan,
     TimeSplit,
     Timing,
     Trace,
     TraceTask,
-    TrainRunInfo,
     VersionInfo,
     WireTrace,
 )
@@ -207,15 +216,21 @@ __all__ = [  # noqa: RUF022 - grouped by public API area
     "TraceTask",
     "WireTrace",
     "Reward",
+    "EnvInfo",
     "Episode",
     "WireEpisode",
+    "GroupInfo",
     "TRACE_VERSION",
     "AgentInfo",
     "RunInfo",
     "EvalRunInfo",
+    "EvalWorkInfo",
     "ModelCall",
     "PolicyEvent",
     "TrainRunInfo",
+    "TrainWorkInfo",
+    "WorkInfo",
+    "PolicySpan",
     "VersionInfo",
     "State",
     "StateT",

{verifiers-0.3.1.dev57 → verifiers-0.3.1.dev59}/verifiers/v1/cli/eval/resume.py

@@ -4,9 +4,9 @@
 it. `load` keeps the good saved rollouts and re-runs what's owed: missing rollouts
 (never written) and errored ones (dropped and redone).
 
-A saved rollout is matched to a selected task by content: `task_key` hashes the
-task's wire data. Tasks with identical data are interchangeable, a task whose data
-changed since the interrupted run re-runs, and nothing depends on `data.idx`.
+A saved episode is matched to a selected task by hashing `episode.task.data`.
+Tasks with identical data are interchangeable, a task whose data changed since the
+interrupted run re-runs, and nothing depends on `data.idx`.
 """
 
 import json
@@ -16,18 +16,17 @@ from pathlib import Path
 
 from pydantic_core import from_json
 
-from verifiers.v1.cli.output import TRACES_FILE, sniff_episode
-from verifiers.v1.episode import Episode, WireEpisode
+from verifiers.v1.cli.output import TRACES_FILE
+from verifiers.v1.episode import WireEpisode
 from verifiers.v1.task import task_key
-from verifiers.v1.trace import WireTrace
 
 
 def load(
     resume_dir: Path,
     selected_keys: list[str],
     num_rollouts: int,
-    complete: Callable[[Episode], bool] | None = None,
-) -> tuple[list[Episode], dict[str, int]]:
+    complete: Callable[[WireEpisode], bool] | None = None,
+) -> tuple[list[WireEpisode], dict[str, int]]:
     """Load the good saved rollouts and diff them against the run's target: returns
     (kept episodes, rollouts owed per task key). `selected_keys` is one key per
     selected task (duplicates allowed — a key selected k times is owed up to
@@ -40,13 +39,8 @@ def load(
         key: count * num_rollouts for key, count in Counter(selected_keys).items()
     }
 
-    def parse(row: dict) -> Episode:
-        if sniff_episode(row):
-            return WireEpisode.model_validate(row)
-        return Episode.of(WireTrace.model_validate(row))
-
     verdict = complete if complete is not None else (lambda episode: episode.ok)
-    good: dict[str, list[tuple[bytes, Episode]]] = defaultdict(list)
+    good: dict[str, list[tuple[bytes, WireEpisode]]] = defaultdict(list)
     if path.exists():
         with path.open("rb") as results:
             for line in results:
@@ -57,12 +51,9 @@ def load(
                         row = from_json(line)
                     except ValueError:
                         row = json.loads(line)
-                    # The task rides each trace; a traceless record (a failure
-                    # before any trace minted) has no task and is owed again.
-                    if sniff_episode(row):
-                        key = task_key(row["traces"][0]["task"]["data"])
-                    else:
-                        key = task_key(row["task"]["data"])
+                    if "traces" not in row:
+                        continue
+                    key = task_key(row["task"]["data"])
                 except (ValueError, KeyError, IndexError, TypeError):
                     # A torn final line (the run died mid-write) or a foreign shape
                     # is not a keepable rollout — it's owed again, never a crash.
@@ -70,7 +61,7 @@ def load(
                 if key not in targets or len(good[key]) >= targets[key]:
                     continue
                 try:
-                    episode = parse(row)
+                    episode = WireEpisode.model_validate(row)
                     if not verdict(episode):
                         continue
                 # A malformed row from any task/episode plugin is owed again.
@@ -80,7 +71,7 @@ def load(
                     (line if line.endswith(b"\n") else line + b"\n", episode)
                 )
     keep: list[bytes] = []
-    episodes: list[Episode] = []
+    episodes: list[WireEpisode] = []
     owed: dict[str, int] = {}
     for key, target in targets.items():
         rows = good.get(key, [])

{verifiers-0.3.1.dev57 → verifiers-0.3.1.dev59}/verifiers/v1/cli/eval/runner.py

@@ -4,6 +4,7 @@ import asyncio
 import contextlib
 import logging
 import time
+from typing import cast
 
 from verifiers.v1.cli.dashboard import dashboard
 from verifiers.v1.cli.eval import resume
@@ -16,8 +17,7 @@ from verifiers.v1.cli.resume import distribute
 from verifiers.v1.clients import ModelContext
 from verifiers.v1.configs.cli.eval import EvalConfig
 from verifiers.v1.env import Env, RunSlot
-from verifiers.v1.episode import Episode
-from verifiers.v1.trace import EvalRunInfo
+from verifiers.v1.episode import Episode, EvalRunInfo
 
 logger = logging.getLogger(__name__)
 
@@ -46,7 +46,13 @@ async def run_eval(env: Env, config: EvalConfig) -> list[Episode]:
     finished: list[Episode] = []
     if config.resume:
         keys = [task.hash for task in tasks]
-        finished, owed = resume.load(out, keys, config.num_rollouts, env.complete)
+        loaded, owed = resume.load(
+            out,
+            keys,
+            config.num_rollouts,
+            lambda episode: env.complete(cast(Episode, episode)),
+        )
+        finished = [cast(Episode, episode) for episode in loaded]
         if not owed:  # already complete - report it and exit successfully
             print(resume.nothing_to_resume_msg(out, len(tasks), config.num_rollouts))
             raise SystemExit(0)
@@ -72,8 +78,7 @@ async def run_eval(env: Env, config: EvalConfig) -> list[Episode]:
     write_lock = asyncio.Lock()
 
     async def on_complete(episode: Episode) -> None:
-        for trace in episode.traces:
-            trace.record_run(EvalRunInfo(id=config.run.id, name=config.run.name))
+        episode.record_run(EvalRunInfo(id=config.run.id, name=config.run.name))
         await append_episode(out, episode, write_lock)
 
     # Serving resources (shared tool servers, interception) come up once for the
@@ -170,7 +175,8 @@ async def run_eval_server(config: EvalConfig) -> list[Episode]:
         finished: list[Episode] = []
         if config.resume:
             keys = [task.hash for task in tasks]
-            finished, owed = resume.load(out, keys, config.num_rollouts)
+            loaded, owed = resume.load(out, keys, config.num_rollouts)
+            finished = [cast(Episode, episode) for episode in loaded]
             counts = distribute(keys, owed, config.num_rollouts)
             if not owed:  # already complete - report it and exit successfully
                 print(resume.nothing_to_resume_msg(out, len(plan), config.num_rollouts))
@@ -205,10 +211,9 @@ async def run_eval_server(config: EvalConfig) -> list[Episode]:
                     sampling=config.sampling,
                     **payload,
                 )
-            for trace in episode.traces:
-                trace.record_run(EvalRunInfo(id=config.run.id, name=config.run.name))
+            episode.record_run(EvalRunInfo(id=config.run.id, name=config.run.name))
             await append_episode(out, episode, write_lock)
-            return [episode]
+            return [cast(Episode, episode)]
 
         # Each rollout is its own `run` request, dispatched least-busy across workers.
         units = [run_unit(payload) for payload, n in plan for _ in range(n)]

{verifiers-0.3.1.dev57 → verifiers-0.3.1.dev59}/verifiers/v1/cli/output.py

@@ -6,8 +6,7 @@ whole episode owed on resume, and a failure before any trace minted still leaves
 its errors on disk. config.toml is the run's resolved config in the format the
 CLI reads (`@ config.toml`), so a run is re-runnable from its own output. Lines
 append as episodes complete, so results are durable mid-run. Files written
-before the episode atom (one bare trace per line) are still readable:
-`read_episodes` sniffs the line shape.
+by this surface contain episodes only.
 """
 
 import asyncio
@@ -19,8 +18,10 @@ from pathlib import Path
 from pydantic import BaseModel, TypeAdapter
 
 from verifiers.v1.configs.cli.eval import EvalConfig
-from verifiers.v1.episode import Episode, WireEpisode
-from verifiers.v1.trace import Trace
+from verifiers.v1.episode import EnvInfo, Episode, WireEpisode
+from verifiers.v1.state import StateT
+from verifiers.v1.task import DataT
+from verifiers.v1.trace import AgentConfigT, Trace
 from verifiers.v1.utils.aio import run_shielded
 
 TRACES_FILE = "traces.jsonl"
@@ -85,7 +86,9 @@ def save_config(
     )  # fresh; appended to as rollouts complete
 
 
-def write_episode(results_dir: Path, episode: Episode) -> None:
+def write_episode(
+    results_dir: Path, episode: Episode[DataT, StateT, AgentConfigT]
+) -> None:
     """Serialize and append one rollout episode in the worker thread."""
     # Preserve fields declared by typed Trace subclasses nested in the episode.
     data = type_adapter(type(episode)).dump_json(episode, exclude_none=True)
@@ -93,40 +96,29 @@ def write_episode(results_dir: Path, episode: Episode) -> None:
         f.write(data + b"\n")
 
 
-def sniff_episode(row: dict) -> bool:
-    """Whether a parsed traces.jsonl row is an `Episode` (vs a pre-episode
-    bare trace, recognizable by its message graph)."""
-    return "traces" in row and "nodes" not in row
-
-
-def read_episodes(results_dir: Path, trace_type: type) -> list[Episode]:
+def read_episodes(results_dir: Path, trace_type: type) -> list[WireEpisode]:
     """Load a run's saved rollouts from `traces.jsonl` with traces typed as
     `trace_type` (`Trace[WireTaskData, ...]` reads any taskset's file without
-    importing it). A pre-episode line (one bare trace) is wrapped as a single-trace
-    record, so both file generations read uniformly."""
+    importing it)."""
     trace_adapter = type_adapter(trace_type)
-    episodes: list[Episode] = []
+    episodes: list[WireEpisode] = []
     with (results_dir / TRACES_FILE).open(encoding="utf-8") as f:
         for line in f:
             if not line.strip():
                 continue
             row = json.loads(line)
-            if sniff_episode(row):
-                # Validate the shell wire-typed (unknown task fields preserved), then
-                # re-type the traces as the caller asked.
-                record = WireEpisode.model_validate({**row, "traces": []})
-                record.traces = [
-                    trace_adapter.validate_python(t) for t in row.get("traces") or []
-                ]
-                episodes.append(record)
-            else:
-                trace = trace_adapter.validate_python(row)
-                episodes.append(Episode.of(trace))
+            record = WireEpisode.model_validate({**row, "traces": []})
+            record.traces = [
+                trace_adapter.validate_python(trace) for trace in row["traces"]
+            ]
+            episodes.append(record)
     return episodes
 
 
 async def append_episode(
-    results_dir: Path, episode: Episode, lock: asyncio.Lock
+    results_dir: Path,
+    episode: Episode[DataT, StateT, AgentConfigT],
+    lock: asyncio.Lock,
 ) -> None:
     """Append one finished rollout episode without blocking the event loop. The run's
     shared lock preserves whole-line ordering, and awaiting the worker preserves
@@ -147,4 +139,10 @@ async def append_trace(
     """Append one finished trace as a single-agent rollout episode — the writers that
     complete trace-at-a-time (eval runners, gepa, replay) all go
     through here."""
-    await append_episode(results_dir, Episode.of(trace, env=env), lock)
+    episode = Episode(
+        env=EnvInfo(id=env),
+        task=trace.task,
+        traces=[trace],
+        ok=trace.ok,
+    )
+    await append_episode(results_dir, episode, lock)

{verifiers-0.3.1.dev57 → verifiers-0.3.1.dev59}/verifiers/v1/env.py

@@ -31,7 +31,7 @@ from verifiers.v1.interception import (
 from verifiers.v1.mcp import SharedToolServer, serve_shared
 from verifiers.v1.runtimes import SubprocessConfig, runtime_is_local
 from verifiers.v1.task import Task
-from verifiers.v1.trace import Error, Trace
+from verifiers.v1.trace import Error, Trace, TraceTask
 from verifiers.v1.utils.generic import concrete_type, deep_merge
 from verifiers.v1.utils.memory import trim_memory_periodically
 from verifiers.v1.utils.retries import run_episode_with_retry
@@ -60,7 +60,7 @@ class RunSlot:
     @classmethod
     def finished(cls, episode: Episode) -> "RunSlot":
         return cls(
-            task=Task(episode.traces[0].task.data),
+            task=Task(episode.task.data),
             traces=list(episode.traces),
             episode=episode,
             done=True,
@@ -252,7 +252,15 @@ class Env(ABC, Generic[ConfigT]):
         completed subset, its exception on the episode's `errors`. `on_trace` observes
         each agent-run's trace at mint; `on_discard` its abandonment (a per-agent
         retry mints a replacement)."""
-        episode = Episode(env=EnvInfo(id=self.config.env_id))
+        episode = Episode(
+            env=EnvInfo(id=self.config.env_id),
+            task=TraceTask(
+                type=type(task).__name__,
+                data=task.data,
+                key=task.key,
+                hash=task.hash,
+            ),
+        )
         agents = self._episode_agents(ctx, episode.traces, on_trace, on_discard)
         try:
             async with asyncio.timeout(self.config.timeout.episode):

{verifiers-0.3.1.dev57 → verifiers-0.3.1.dev59}/verifiers/v1/episode.py

@@ -1,14 +1,14 @@
 """The episode — one run's traces plus their shared standing, whole."""
 
 import uuid
-from typing import Any, Generic
+from typing import Annotated, Any, Generic, Literal, Self
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, model_validator
 
 from verifiers.v1.configs.agent import WireAgentConfig
 from verifiers.v1.state import State, StateT
 from verifiers.v1.task import DataT, WireTaskData
-from verifiers.v1.trace import EXCLUDE_FIELDS, AgentConfigT, Error, Trace
+from verifiers.v1.trace import EXCLUDE_FIELDS, AgentConfigT, Error, Trace, TraceTask
 from verifiers.v1.types import Usage
 
 
@@ -17,6 +17,70 @@ class EnvInfo(BaseModel):
 
     id: str = ""
     """`EnvConfig.env_id`, e.g. `agentic-judge+gsm8k`."""
+    name: str | None = None
+    """The name the consumer runs the env under (e.g. an orchestrator env key)."""
+
+
+class GroupInfo(BaseModel):
+    """The consumer-defined rollout group this episode belongs to."""
+
+    id: str
+
+
+class PolicySpan(BaseModel):
+    """Live policy versions spanned while generating an episode."""
+
+    start: int
+    end: int
+
+    @model_validator(mode="after")
+    def ordered(self) -> Self:
+        if self.end < self.start:
+            raise ValueError("policy span end must be at least its start")
+        return self
+
+    @property
+    def drift(self) -> int:
+        return self.end - self.start
+
+
+class TrainWorkInfo(BaseModel):
+    type: Literal["train"] = "train"
+    step: int
+    """The training step at which the episode was dispatched."""
+    policy: PolicySpan | None = None
+    """The live policy used for generation; None for frozen-policy work."""
+
+
+class EvalWorkInfo(BaseModel):
+    type: Literal["eval"] = "eval"
+    step: int
+    """The training step at which the episode was dispatched."""
+    policy: PolicySpan | None = None
+    """The live policy used for generation; None for frozen-policy work."""
+
+
+WorkInfo = Annotated[TrainWorkInfo | EvalWorkInfo, Field(discriminator="type")]
+
+
+class TrainRunInfo(BaseModel):
+    """A training run and the kind of work this episode performed for it."""
+
+    type: Literal["train"] = "train"
+    id: str
+    name: str | None = None
+    work: WorkInfo
+
+
+class EvalRunInfo(BaseModel):
+    """A standalone evaluation run."""
+
+    type: Literal["eval"] = "eval"
+    id: str
+    name: str | None = None
+
+
+RunInfo = Annotated[TrainRunInfo | EvalRunInfo, Field(discriminator="type")]
 
 
 class Episode(BaseModel, Generic[DataT, StateT, AgentConfigT]):
@@ -26,6 +90,12 @@ class Episode(BaseModel, Generic[DataT, StateT, AgentConfigT]):
 
     env: EnvInfo = Field(default_factory=EnvInfo)
     """The env that produced this episode."""
+    task: TraceTask[DataT]
+    """The task dispatched to the env, including when no trace was produced."""
+    group: GroupInfo | None = None
+    """Consumer-assigned rollout group, if this episode was grouped."""
+    run: RunInfo | None = None
+    """The run this episode belongs to, consumer-stamped."""
     ok: bool = False
     """Whether the episode completed successfully."""
     errors: list[Error] = Field(default_factory=list)
@@ -33,6 +103,10 @@ class Episode(BaseModel, Generic[DataT, StateT, AgentConfigT]):
     traces: list[Trace[DataT, StateT, AgentConfigT]] = Field(default_factory=list)
     """Every agent's trace, in completion order."""
 
+    def record_run(self, run: RunInfo) -> None:
+        """Record the run identity on the dispatched episode."""
+        self.run = run
+
     @property
     def last_error(self) -> Error | None:
         """The last episode-level error captured across attempts."""
@@ -80,11 +154,6 @@ class Episode(BaseModel, Generic[DataT, StateT, AgentConfigT]):
             exclude_none=True,
         )
 
-    @classmethod
-    def of(cls, trace: Trace, env: str = "") -> "Episode":
-        """The single-agent record: one trace as its own episode."""
-        return cls(env=EnvInfo(id=env), traces=[trace], ok=trace.ok)
-
 
 WireEpisode = Episode[WireTaskData, State, WireAgentConfig]
 """Record loader for consumers without the run's packages: unknown task fields

{verifiers-0.3.1.dev57 → verifiers-0.3.1.dev59}/verifiers/v1/graph.py

@@ -106,6 +106,11 @@ class MessageNode(BaseModel):
     consumer's RL algorithm assigns it, which is not the same as a credit of zero: a group whose
     rewards were all equal is assigned zeros and carries no gradient, while an unassigned node was
     never scored at all."""
+    reference_logprobs: list[float] | None = None
+    """Reference-model logprobs over the sampled tokens, in the same compact layout as
+    `logprobs`. None means no reference model scored this node."""
+    loss_weights: dict[str, list[float]] | None = None
+    """Named loss-weight streams aligned to `token_ids`, consumer-stamped."""
     multi_modal_data: SkipJsonSchema[MultiModalData | None] = None
     """The renderer items for the images this message's content introduces (pixel tensors,
     grids, hashes, placeholders) — the only carrier of the pixels from the env server to the

{verifiers-0.3.1.dev57 → verifiers-0.3.1.dev59}/verifiers/v1/runtimes/prime.py

@@ -36,9 +36,10 @@ from verifiers.v1.utils.prime import ensure_prime_auth
 
 logger = logging.getLogger(__name__)
 
-CONTAINER_LIFETIME = 24 * 60 * 60
-"""Fixed lifetime (seconds) of container sandboxes; only VM sandboxes support an
-unbounded lifetime."""
+IDLE_FALLBACK_LIFETIME = 30 * 24 * 60 * 60
+"""Lifetime (seconds) given to container sandboxes that set an idle timeout: the
+platform requires a finite lifetime there as a safety fallback should idle detection
+fail. Far above any real run, so effectively unbounded."""
 
 
 BASE_LABELS: list[str] = []
@@ -101,20 +102,6 @@ class PrimeConfig(NetworkPolicyConfig):
         )
         return self
 
-    @model_validator(mode="after")
-    def _validate_idle_timeout(self) -> "PrimeConfig":
-        if (
-            not self.vm
-            and self.idle_timeout is not None
-            and self.idle_timeout > CONTAINER_LIFETIME
-        ):
-            raise ValueError(
-                f"idle_timeout ({self.idle_timeout}s) must not exceed the "
-                f"{CONTAINER_LIFETIME}s ({CONTAINER_LIFETIME // 3600}h) container "
-                "sandbox lifetime"
-            )
-        return self
-
 
 class PrimeRuntimeInfo(PrimeConfig, BaseRuntimeInfo):
     image_cached: bool | None = None
@@ -178,8 +165,14 @@ class PrimeRuntime(Runtime):
             "memory_gb": self.config.memory,
             "disk_size_gb": self.config.disk,
             "gpu_count": gpu_count,
-            # -1 is prime's convention for no lifetime limit (VM-only)
-            "timeout_minutes": -1 if self.config.vm else CONTAINER_LIFETIME // 60,
+            # -1 is prime's convention for no lifetime limit; containers with an
+            # idle timeout must carry a finite lifetime as a safety fallback (which
+            # must exceed the idle timeout)
+            "timeout_minutes": (
+                -1
+                if self.config.vm or idle_minutes is None
+                else max(IDLE_FALLBACK_LIFETIME // 60, idle_minutes + 1)
+            ),
             "idle_timeout_minutes": idle_minutes,
             "gpu_type": gpu_type,
             "region": self.config.region,
@@ -281,17 +274,21 @@ class PrimeRuntime(Runtime):
         )
 
     async def run(self, argv: list[str], env: dict[str, str]) -> ProgramResult:
+        # Poll the job by hand: the SDK's run_background_job needs a finite deadline,
+        # but the sandbox has no lifetime limit — the rollout's stage timeouts bound
+        # this via cancellation instead.
         try:
-            result = await self._client.run_background_job(
+            job = await self._client.start_background_job(
                 self.info.id,
                 shlex.join(argv),
-                # the SDK poll needs a finite deadline: the container lifetime, or
-                # effectively unbounded for VM sandboxes
-                timeout=365 * 24 * 60 * 60 if self.config.vm else CONTAINER_LIFETIME,
                 working_dir=self.config.workdir,
                 env=self.process_env(env),
-                poll_interval=1,
             )
+            while True:
+                result = await self._client.get_background_job(self.info.id, job)
+                if result.completed:
+                    break
+                await asyncio.sleep(1)
         except (
             Exception
         ) as e:  # a sandbox/API failure is one rollout's problem, not the eval's

{verifiers-0.3.1.dev57 → verifiers-0.3.1.dev59}/verifiers/v1/trace.py

@@ -4,7 +4,7 @@ import time
 import traceback
 import uuid
 from collections.abc import Callable, Iterable, Iterator, Mapping
-from typing import TYPE_CHECKING, Annotated, Any, Generic, Literal
+from typing import TYPE_CHECKING, Any, Generic
 
 import numpy as np
 from pydantic import BaseModel, Field, PrivateAttr
@@ -141,28 +141,6 @@ class Reward(BaseModel):
         return self.score * self.weight
 
 
-class EvalRunInfo(BaseModel):
-    type: Literal["eval"] = "eval"
-
-    id: str
-    name: str | None = None
-    """Human-readable run name, consumer-stamped."""
-    step: int | None = None
-
-
-class TrainRunInfo(BaseModel):
-    type: Literal["train"] = "train"
-
-    id: str
-    name: str | None = None
-    """Human-readable run name, consumer-stamped."""
-    step: int | None = None
-
-
-RunInfo = Annotated[EvalRunInfo | TrainRunInfo, Field(discriminator="type")]
-"""The run a trace belongs to, discriminated on `type`."""
-
-
 class PolicyEvent(BaseModel):
     """A framework policy decision applied to a model call before inference."""
 
@@ -277,6 +255,34 @@ class Branch(BaseModel):
             return None
         return self.spread(lambda node: node.advantages)
 
+    @property
+    def reference_logprobs(self) -> list[float] | None:
+        """Reference-model logprobs aligned to `token_ids`, or None when unscored."""
+        if all(node.reference_logprobs is None for node in self.nodes):
+            return None
+        return self.spread(lambda node: node.reference_logprobs)
+
+    def loss_weights(self, name: str) -> list[float] | None:
+        """One named loss-weight stream aligned to `token_ids`, or None when absent."""
+        if all(
+            node.loss_weights is None or name not in node.loss_weights
+            for node in self.nodes
+        ):
+            return None
+        weights: list[float] = []
+        for node in self.nodes:
+            node_weights = (node.loss_weights or {}).get(name)
+            if node_weights is None:
+                weights.extend([0.0] * len(node.token_ids))
+                continue
+            if len(node_weights) != len(node.token_ids):
+                raise ValueError(
+                    f"loss weight stream {name!r} must align with node token_ids: "
+                    f"got {len(node_weights)}, expected {len(node.token_ids)}"
+                )
+            weights.extend(node_weights)
+        return weights
+
     @property
     def multi_modal_data(self) -> MultiModalData | None:
         """Node image data concatenated in token order for training; never persisted."""
@@ -361,9 +367,6 @@ class Trace(BaseModel, Generic[DataT, StateT, AgentConfigT]):
     """Unique ID for this trace, auto-generated."""
     verifiers: VersionInfo = Field(default_factory=_current_build)
     """The verifiers version that produced this trace."""
-    run: RunInfo | None = None
-    """The run this trace belongs to (eval or train), consumer-stamped."""
-
     task: TraceTask[DataT]
     """The task data that seeded this trace."""
     agent: AgentInfo[AgentConfigT]
@@ -564,12 +567,6 @@ class Trace(BaseModel, Generic[DataT, StateT, AgentConfigT]):
         if response.usage is not None:
             self.extra_usage.append(response.usage)
 
-    def record_run(self, run: RunInfo | None = None, **info: Any) -> None:
-        """Record the run identity (eval / train), and optional extra info."""
-        if run is not None:
-            self.run = run
-        self.info.update(info)
-
     def stop(self, condition: str) -> None:
         """Stop the trace, optionally with a stop condition."""
         self.is_completed = True