feat(orchestrator)!: make the episode the unit the orchestrator passes around

[mikasenghaas]

Summary

Follow-up to #3165. The orchestrator passes episodes end to end — the env's own vf.Episode,
extended only with what prime-rl genuinely adds — instead of flattening them into loose traces at
the boundary and rebuilding the grouping downstream.

Depends on PrimeIntellect-ai/verifiers#2252; deps/verifiers is pinned to that branch and needs a
re-pin to main once it merges.

An episode that produced nothing needs no stand-in trace

A cancellation, or a task that raised before reaching the env, was reported as a fabricated
Rollout carrying a fake error. That trace has no real agent, so it took AgentInfo's default
name, and once #3165 keyed metrics by agent those failures landed under a phantom agent subtree —
in a proposer-solver env, cancellations inflated a seat nobody ran. (Reported by bugbot on #3165.)

verifiers already has the shape: run_episode records the reason on episode.errors and returns
the episode with ok false. The dispatcher records the cancellation on its own counters and emits
that same shape, so no stand-in exists and one vocabulary covers every cause.

The episode carries the dispatch

class Episode(vf.WireEpisode):   group_id      # the comparison group; verifiers has no notion of one
class Rollout(vf.Trace[DataT]):  env_name, group_id, samples, is_filtered, filter_results

Everything else has a place on the episode already: the env it ran (env.name) and the run it
belongs to (run, whose metadata says whether the run trains on it or measures itself with it, and
carries the step and the policy versions generation spanned). So kind, policy_version,
off_policy_steps and eval_step leave the trace — they describe the episode, which is the thing
that was dispatched.

Staleness is derived, not stored. It used to be written twice: the dispatcher counted weight
updates per in-flight episode, then the main loop discarded that and recomputed
(step - 1) - policy_version at ship. Now the dispatcher records the span generation covered and
every reading derives from it — TrainMetadata.off_policy_steps for what trains,
EvalMetadata.off_policy_steps (the drift) for what only measures.

Algorithms score episodes

score_group / finalize_group take the group's episodes, so an algorithm can compare within an
episode as well as across them — hierarchical_grpo keys its solver baselines off episode.id
instead of a foreign key copied onto every trace.

Credit moves onto MessageNode: Rollout.advantages becomes a derived read, assign_advantages
writes each node's trainable tokens, and stamp_advantages copies branch.advantages onto the
sample built from that branch rather than slicing one flat stream by offset. Unassigned stays
distinct from assigned-zero all the way to the trainer.

The sample monitors take episodes

log_samples / log_eval_samples take episodes across all five implementations, and the W&B table
gains agent and branch_idx columns — one row per branch, so a multi-agent episode reads as its
seats rather than as one blurred cell.

traces.jsonl is written by verifiers

One episode per line, through vf's own write_episode, so a prime-rl record reads exactly like one
its read_episodes produces. An episode that produced no traces is written too — its errors are
the record of why nothing came back.

Breaking

• Rollout no longer carries kind, policy_version, off_policy_steps, eval_step or
  episode_id; read them off the episode (run.metadata.type, run.metadata.step,
  run.metadata.policy, run.metadata.off_policy_steps, env.name).
• Rollout.advantages is read-only (derived from the nodes) and assign_advantages takes a scalar
  only — the full-length per-token list is gone. No shipped algorithm used it.
• Algorithm.score_group / finalize_group take list[Episode], not list[Rollout].
• Monitor.log_samples / log_eval_samples take episodes, not traces.
• InflightRollout → InflightEpisode (rollout_count → episodes_owed); RolloutDispatcher
  takes a run_id.
• traces.jsonl rows are episodes, not traces. verifiers' read_episodes reads both.
• {scope}/{subset}/<agent>/has_error/mean no longer counts cancellations or task failures — an
  episode nobody ran belongs to no seat. Those report under
  dispatcher/{cancelled,errored}/{train,eval}.

Verification

• uv run pytest tests/unit -m "not gpu" — 507 passed (CI's own selection).

• ruff check / ruff format --check at the pinned 0.13.0.

• A 6-step reverse-text run with eval at steps 2/4/6, no errors. The records:

  // step_6/train/all/traces.jsonl
  "run": {"type": "train", "id": "2f3d94e4...",
          "metadata": {"type": "train", "step": 6, "policy": {"start": 2, "end": 4}}}
  "env": {"id": "reverse-text-v1", "name": "reverse-text"}
  
  // step_6/eval/all/traces.jsonl — same run id, told apart by its metadata
  "run": {"type": "train", "id": "2f3d94e4...",
          "metadata": {"type": "eval", "step": 6, "policy": {"start": 4, "end": 4}}}

  That train episode spans v2 → v4: it outlived two weight updates mid-generation, which the
  stored counter could not express.

Follow-ups

• Splitting degeneracy detection (a measurement of every trace) from the drop policy (a decision).
  apply_filters stops at the first hit, so a rollout flagged as gibberish is never measured for
  repetition — a monitoring rate that depends on filter order. Left alone here on purpose.
• Rollout still carries env_name / group_id for the consumers that work in flat trace lists
  (the sink's pending batch, per-agent metric bucketing).
• The sinks, the ship path and the eval summary have no unit coverage; every bug found while
  building this reached a live run through a green suite.

🤖 Generated with Claude Code