feat(orchestrator)!: align multi-agent types

[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. Three changes, each one falling out of the last.

1. 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. Prime-rl's own outcomes are minted the same way, so one vocabulary
covers every cause. envs.py stops unwrapping the episode, metrics.py stops rebuilding it, and
vf's own by_agent / num_turns / num_*_tokens replace the hand-rolled versions.

traces.jsonl now stores one episode per line — what verifiers writes and what its
read_episodes expects; prime-rl was writing the legacy bare-trace form. A traceless episode gets a
row too, with the reason on errors.

2. prime-rl stops extending the episode and the trace

Everything the orchestrator adds to a dispatch now has a place on verifiers' own types
(PrimeIntellect-ai/verifiers#2252): the env it ran (env.name) and the run it belongs to (run —
an id, plus a metadata saying whether the run trains on the episode or measures itself with it,
and carrying the step and policy span that go with that). The comparison group is prime-rl's alone —
verifiers has no notion of one — so it rides in episode.info, read through one accessor.

Episode = vf.WireEpisode      # alias — nothing added
Rollout = vf.Trace            # alias — nothing added

class TrainRollout(vf.Trace[DataT]):
    env_name, samples, is_filtered, filter_results

TrainEpisode, EvalEpisode, Episode.KIND and the Rollout base are gone; what were methods are
free functions over the vf types. TrainRollout is the one extension left, because trainer-bound
state has nowhere else to live — samples in particular cannot go in trace.info, which
serializes.

The main loop routes on run.metadata and fills in a train episode's step, the one fact the
dispatcher can't know: it is the batch window collecting when the episode lands, not when it was
dispatched. An eval's is known at dispatch, so EvalMetadata.step is required.

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. The dispatcher now records the policy span generation
covered, and every reading derives from it: the cancel check subtracts against the live version, the
in-flight gauges do the same, and off_policy_steps is a property of each metadata: versions behind
what is training for an episode trained on, and what drifted under it for an eval, which is the only
sense in which an eval can be off-policy. A frozen sampler records no span at all, so its staleness
reads None instead of a zero that looks fresh.

InflightRollout becomes InflightEpisode, and stamp() is the single place a dispatch's facts
become an episode's.

3. Algorithms score episodes; credit lives on the graph's nodes

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. group_rollouts(group) flattens for GRPO / MaxRL /
RAE, which only compare across.

Per-agent metrics are built from episodes too (narrow, keyed by vf.Episode.by_agent), so
solve_rates and pass_at_k bucket on the episode's group — and Rollout.group_id is gone. The
metrics stay flat over traces; the episode is only what they are built from.

Credit moves onto MessageNode (PrimeIntellect-ai/verifiers#2245, merged and pinned here).
TrainRollout.advantages becomes a derived read, assign_advantages(value) 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 back apart by offset. Unassigned stays distinct from
assigned-zero all the way to the trainer.

Breaking

Types

• Episode and Rollout are aliases for vf.WireEpisode / vf.Trace. TrainEpisode,
  EvalEpisode and Episode.KIND are gone; read the path off run.metadata.type, the step off
  run.metadata.step, the env off env.name, the group off info["group_id"], staleness off
  run.metadata.off_policy_steps.
• What were Episode methods are free functions in orchestrator.types: narrow, rollouts_of,
  group_id_of, env_name_of, run_of.
• Rollout no longer carries kind, policy_version, off_policy_steps, eval_step, group_id
  or episode_id. env_name, samples, is_filtered and filter_results are on TrainRollout.
• Monitor.log_samples / log_eval_samples take episodes, not traces. The wandb sample table gains
  agent and branch_idx columns, one branch per row.
• TrainRollout.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.
• InflightRollout → InflightEpisode (rollout_count → episodes_owed); stamp() takes a
  run_id, and RolloutDispatcher a matching ctor argument.

Algorithms

• Algorithm.score_group / finalize_group take list[Episode], not list[TrainRollout]. A custom
  algorithm comparing across the whole cohort wraps its body in group_rollouts(group).
• TraceMetrics(...) and pass_at_k(...) take episodes.

Metrics / on-disk

• 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 — 481 passed. Two exclusions on this box, both pre-existing and
  unrelated: tests/unit/train/models (GPU crash) and test_qwen3_vl_e2e (stale fixture, fix: token_id-formatted logprob tokens in the qwen3-vl fake engine #3161).
• Four live reverse-text runs on 2 GPUs:

  • credit — 12 steps, reward 0.17 → 0.75, Trainable 128/128 throughout, so node-assigned
    credit reaches the trainer.

  • episode collapse — train+eval, evals at steps 2/4/6; traces.jsonl rows carry
    run={'type': 'train', 'step': <batch window>} / {'type': 'eval', 'step': <eval epoch>} with
    the dispatch facts in info.

  • metrics rework — solved_*, avg@k and the per-agent subtree land unchanged on both
    subsets.

  • no subclasses — a 6-step run with eval. Both kinds of episode carry the same run id and are
    told apart by their metadata, and everything prime-rl used to keep on a subclass is on the
    record:

    // step_6/train/all/traces.jsonl
    "run":  {"type": "train", "id": "2bfeba20...",
             "metadata": {"type": "train", "step": 6, "policy": {"start": 3, "end": 3}}},
    "env":  {"id": "reverse-text-v1", "name": "reverse-text"},
    "info": {"group_id": "ec69996d-..."}
    
    // step_6/eval/all/traces.jsonl — same run id, different metadata
    "run":  {"type": "train", "id": "2bfeba20...",
             "metadata": {"type": "eval", "step": 6, "policy": {"start": 4, "end": 4}}}

Two latent bugs surfaced on the way and are fixed here: both sinks read a group's env and step off
group[0] — a trace — which IndexErrors when a whole group is cancelled and produces none. No
test covered that path.

Follow-ups

• Depends on feat(v1): run metadata on episode verifiers#2252; deps/verifiers is pinned to that branch and needs a
  re-pin to main once it merges.
• Up to date with main as of feat!: standalone env servers #3162 (standalone env servers), verified with a live run.
• TrainRollout is the remaining extension. Retiring it means moving samples / is_filtered /
  filter_results into a sink-owned table keyed by trace id, which metrics would then have to
  carry — a subclass traded for threading, so I would keep it.
• Filters are untouched here on purpose. Splitting degeneracy detection (a measurement of every
  trace) from the drop policy (a decision) is a change of its own — it fixes apply_filters stopping
  at the first hit, which makes a monitoring rate depend on filter order — and it goes in a
  follow-up rather than riding along with the episode work.
• The sinks, the ship path and the eval summary have no unit coverage. Four runtime
  AttributeErrors during this work reached a live run through a green suite, each one a read of a
  field the subclass used to have. Scripted sweeps now check every episode and run attribute read
  against the vf surface, but that is a stopgap for real tests.
• dropped/all/<name>/rate is only emitted in steps where that reason fired, so the per-reason
  series is sparse. Pre-existing under the old key.

🤖 Generated with Claude Code

[mikasenghaas]

Closing in favour of a minimal rewrite — this branch accumulated several changes (filters, the trace-type collapse, group-in-info, accessor helpers) that belong in their own PRs. Superseded by a fresh PR scoped to the episode work alone.