feat(orchestrator)!: align multi-agent types

docs/algorithms.md

@@ -166,12 +166,12 @@ At runtime, each env's resolved config builds two objects: a `Sampler` (`prime_r
 | `opsd` | `OPSDAlgorithm` | `score_rollout`: demo-conditioned prefill under the live policy |
 | `sft` | `SFTDistillAlgorithm` | `score_group`: group-norm credit (feeds filters) |
 
-Each class owns its hooks outright — reading one top to bottom reads the algorithm, and everything on the class is an override point. The two hooks are one scope-and-timing ladder — the wider scope is unlocked by a later barrier, so the two axes coincide. Each is handed the `Rollout` directly — the env's typed trace (`reward`, `nodes`, `num_turns`, ...) with `samples` attached, plus `assign_advantages` to write credit:
+Each class owns its hooks outright — reading one top to bottom reads the algorithm, and everything on the class is an override point. The two hooks are one scope-and-timing ladder — the wider scope is unlocked by a later barrier, so the two axes coincide. Each is handed the env's own data directly — a `Rollout` (the typed trace: `reward`, `nodes`, `num_turns`, ... with `samples` attached) on arrival, the group's `Episode`s at group time — plus `assign_advantages` to write credit:
 
-- `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`, scalar broadcast or per-token), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the pre-batch filters, so it pays compute on rollouts that may then be filtered out.
-- `score_group(group)` — the cohort, **before filtering** (filters read the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `Rollout`.
+- `async score_rollout(rollout)` — one rollout, **on arrival** (as it's tokenized, before its group is complete): rollout-local credit (`rollout.assign_advantages(...)`), observation ce weights, **or** model I/O — query a reference pool (e.g. `self.teacher_pool`, connected in `setup()` via `self.connect(...)`, or the live `self.policy_pool` for opsd) and attach per-token results (e.g. teacher logprobs) with bounded concurrency. No siblings. `echo` weights observation tokens here, identifying env-provided observation nodes by their non-sampled status and source step role attribution, applying the optional user filter, and writing the `ce_weights` stream. Model I/O runs *before* the pre-batch filters, so it pays compute on rollouts that may then be filtered out.
+- `score_group(group)` — the cohort, **before filtering** (filters read the streams), synchronous: group-relative credit (GRPO/MaxRL baselines). `group` is a list of `TrainEpisode`, so an algorithm can compare within an episode as well as across them (`hierarchical_grpo` does); `group_rollouts(group)` flattens it for the ones that only compare across.
 
-The pipeline drives the hooks through two non-virtual methods it never looks inside: `algorithm.finalize_rollout(rollout)` per arrival (rollout-local scoring + reference I/O) and `algorithm.finalize_group(rollouts)` per group (scoring + wire stamping; after this the records are frozen — groups die at stamping). Sample construction (interleaving) is pure pipeline — observation-token provenance is available through structural attribution (`node.sampled`, `node.is_content`) for any algorithm that trains on env-provided tokens.
+The pipeline drives the hooks through two non-virtual methods it never looks inside: `algorithm.finalize_rollout(rollout)` per arrival (rollout-local scoring + reference I/O) and `algorithm.finalize_group(episodes)` per group (scoring + wire stamping; after this the records are frozen — groups die at stamping). Sample construction (interleaving) is pure pipeline — observation-token provenance is available through structural attribution (`node.sampled`, `node.is_content`) for any algorithm that trains on env-provided tokens.
 
 Class-level declarations state what the algorithm needs: which loss component its action tokens feed (`action_loss_type`). Every class is constructed with its algorithm config plus the one host-owned resource it can't rebuild — the live policy pool (`self.policy_pool`). Everything else an algorithm needs it builds from its own config in `setup()`: `opd` connects its frozen `teacher`; `opsd` builds the renderer for its demonstration hint (tokenizer is always the live policy's — self-distillation has no separate model). The pipeline only ever calls the two `finalize_*` methods — writing your own algorithm is subclassing `Algorithm` and overriding the hooks its signal needs (see [Authoring an Algorithm](#authoring-an-algorithm)). Shared math (efficiency shaping, prefill alignment) lives as plain functions in `prime_rl.orchestrator.algo.advantage`.
 
@@ -420,26 +420,28 @@ Both of `kuhn-poker-v1`'s agents late-bind to the run's own model — shared-pol
 
 ### Authoring an Algorithm
 
-There is no config hook that points at user code — a new credit-assignment scheme is a new named algorithm in the repo. Subclass `Algorithm`, assign credit in the scoring hook whose timing fits your signal, and register the class. The hook receives the group's `Rollout`s (each the env's typed `verifiers.Trace` — turns, tool calls, metadata in `info` — with `samples` attached) and writes credit via `assign_advantages`:
+There is no config hook that points at user code — a new credit-assignment scheme is a new named algorithm in the repo. Subclass `Algorithm`, assign credit in the scoring hook whose timing fits your signal, and register the class. The hook receives the group's `TrainEpisode`s, each holding the env's typed `verifiers.Trace`s — turns, tool calls, metadata in `info` — with `samples` attached, and writes credit via `assign_advantages`:
 
 ```python
 # src/prime_rl/orchestrator/algo/my_algo.py
 import torch
 
 from prime_rl.orchestrator.algo.base import Algorithm
+from prime_rl.orchestrator.types import group_rollouts
 
 
 class MyAlgorithm(Algorithm):
     async def score_group(self, group):
-        rewards = torch.tensor([rollout.reward for rollout in group], dtype=torch.float32)
+        rollouts = group_rollouts(group)
+        rewards = torch.tensor([rollout.reward for rollout in rollouts], dtype=torch.float32)
         advantages = ...  # one value per rollout
-        for rollout, advantage in zip(group, advantages.tolist(), strict=True):
+        for rollout, advantage in zip(rollouts, advantages.tolist(), strict=True):
             rollout.assign_advantages(advantage)
 ```
 
-Add a typed `MyAlgoConfig` to `prime_rl.configs.algorithm` and its discriminated union, then register `"my_algo": MyAlgorithm` in `ALGORITHM_CLASSES`. Pick the hook by *when* your signal is ready: `score_rollout` for per-arrival credit or credit that needs a model call (it's `async`), `score_group` for group-relative credit (GRPO/MaxRL). `assign_advantages` takes a scalar (broadcast over the rollout's trainable tokens — the common case) or a full-length per-token list aligned to the concatenated sample token_ids (process rewards, step-level credit; `0.0` off-mask).
+Add a typed `MyAlgoConfig` to `prime_rl.configs.algorithm` and its discriminated union, then register `"my_algo": MyAlgorithm` in `ALGORITHM_CLASSES`. Pick the hook by *when* your signal is ready: `score_rollout` for per-arrival credit or credit that needs a model call (it's `async`), `score_group` for group-relative credit (GRPO/MaxRL). `assign_advantages` takes a scalar and broadcasts it over the rollout's trainable tokens, writing it onto the graph's nodes so branches sharing a node cannot disagree about its credit.
 
-Each per-token list must match the rollout's completion-token count exactly — validated loudly when the view writes it. Advantage-based filters and metrics derive from the streams (the zero-advantage filter checks for all-zero streams; logged distributions use per-rollout means). Signals that depend on the live policy's weights (like OPD's reverse KL) cannot be precomputed here; those are reference-scoring algorithms, evaluated in the trainer.
+Advantage-based filters and metrics derive from the streams (the zero-advantage filter checks for all-zero streams; logged distributions use per-rollout means). Signals that depend on the live policy's weights (like OPD's reverse KL) cannot be precomputed here; those are reference-scoring algorithms, evaluated in the trainer.
 
 ### Reference Scoring
 

pyproject.toml

@@ -33,7 +33,7 @@ dependencies = [
     "setproctitle>=1.3.0",
     "uvloop>=0.21.0",
     "torchtitan",
-    "verifiers[harbor]>=0.2.2.dev76",
+    "verifiers[harbor]>=0.2.2.dev81",
     "renderers",
     "dion",
     "tilelang>=0.1.8",

src/prime_rl/orchestrator/algo/__init__.py

@@ -17,10 +17,9 @@
   ``finalize_rollout`` / ``finalize_group`` methods the pipeline drives.
   Advantages are per-token everywhere they are stored or shipped — there is no
   scalar advantage in the pipeline. An algorithm assigns credit in its scoring
-  hook via ``Rollout.assign_advantages``: a scalar that is *broadcast* over the
-  rollout's completion tokens (uniform credit, the common case), or an explicit
-  full-length-N per-token list aligned to the concatenated sample token_ids
-  (0.0 off-mask).
+  hook via ``TrainRollout.assign_advantages``, which broadcasts one value over
+  the rollout's trainable tokens and writes it onto the graph's nodes, where
+  the tokens themselves live.
 - ``routing`` — wire-field stamping: per-token component weight streams
   (rl / ce / ref_kl) and the per-token advantage stream.
 """
@@ -39,7 +38,7 @@
 from prime_rl.orchestrator.algo.rae import RAEAlgorithm
 from prime_rl.orchestrator.algo.routing import stamp_advantages, stamp_loss_routing
 from prime_rl.orchestrator.algo.sft import SFTDistillAlgorithm
-from prime_rl.orchestrator.types import Rollout
+from prime_rl.orchestrator.types import TrainRollout
 
 if TYPE_CHECKING:
     from prime_rl.configs.algorithm import AlgoConfig
@@ -79,7 +78,7 @@ def build_algorithm(config: AlgoConfig, policy_pool: InferencePool) -> Algorithm
     "OPDAlgorithm",
     "OPSDAlgorithm",
     "RAEAlgorithm",
-    "Rollout",
+    "TrainRollout",
     "SFTDistillAlgorithm",
     "build_algorithm",
     "connect_frozen_pool",

src/prime_rl/orchestrator/algo/base.py

@@ -19,7 +19,7 @@
   I/O against another model — an inference pool the algorithm connected in
   ``setup()`` (a frozen teacher) or the live policy (opsd's self-distillation),
   queried with bounded concurrency. No siblings.
-- ``score_group(group)`` — the cohort, on group completion, *before* filtering
+- ``score_group(group)`` — the cohort of episodes, on group completion, *before* filtering
   (filters read the streams): group-relative credit (GRPO/MaxRL baselines).
 
 How rollouts are *produced* is not the algorithm's concern: that is the env's
@@ -44,12 +44,13 @@
 
 from prime_rl.configs.algorithm import ActionLossType, AlgoConfig, FrozenModelConfig
 from prime_rl.orchestrator.algo.routing import stamp_advantages, stamp_loss_routing
+from prime_rl.orchestrator.types import group_rollouts
 from prime_rl.utils.logger import get_logger
 
 if TYPE_CHECKING:
     from renderers import RendererConfig
 
-    from prime_rl.orchestrator.types import Rollout
+    from prime_rl.orchestrator.types import Episode, TrainRollout
     from prime_rl.utils.client import InferencePool
 
 
@@ -91,9 +92,10 @@ class Algorithm:
       (``action_loss_type``);
     - lifecycle — :meth:`setup` connects client pools to the frozen models
       the algorithm declares, resolving each reference via :meth:`connect`;
-    - the two scoring hooks, each ``async`` and given the :class:`Rollout`
-      directly — read the trace, write credit via
-      :meth:`Rollout.assign_advantages`. They are
+    - the two scoring hooks, each ``async`` and given the env's own data
+      directly — a :class:`TrainRollout` on arrival, the group's
+      :class:`Episode`\ s at group time — so a hook reads the trace and
+      writes credit via :meth:`TrainRollout.assign_advantages`. They are
       async so either stage may do I/O — e.g. a process-reward model or a
       teacher at arrival, or a judge at group time whose signal a pre-batch
       filter then reads; a hook that only does advantage math simply never
@@ -103,9 +105,9 @@ class Algorithm:
         observation ce weights, or per-token results from a model the algorithm
         connected in :meth:`setup` (e.g. teacher reference logprobs). Default:
         nothing.
-      - :meth:`score_group` — the cohort, *before* filtering (filters read the
-        streams): group-relative credit. Default: nothing — rollouts keep
-        ``advantages=None``, so advantage-based filters skip them.
+      - :meth:`score_group` — the cohort of episodes, *before* filtering
+        (filters read the streams): group-relative credit. Default: nothing —
+        rollouts keep ``advantages=None``, so advantage-based filters skip them.
 
     Model I/O lives in :meth:`score_rollout`: it runs at arrival, *before* the
     pre-batch filters, so it pays compute on rollouts that may then be filtered
@@ -137,29 +139,31 @@ async def connect(self, reference: FrozenModelConfig) -> InferencePool:
         self.connected_pools.append(pool)
         return pool
 
-    async def score_rollout(self, rollout: Rollout) -> None:
+    async def score_rollout(self, rollout: TrainRollout) -> None:
         """Arrival phase, one rollout, before its group is complete: write
         rollout-local credit (``rollout.assign_advantages``), observation ce
         weights (echo), or per-token results from a model — an inference pool
         connected in :meth:`setup`, or the live policy (opsd). No siblings, no
         group stats."""
 
-    async def score_group(self, group: list[Rollout]) -> None:
+    async def score_group(self, group: list[Episode]) -> None:
         """Group phase, the finalized cohort, before filtering: write
-        group-relative credit."""
+        group-relative credit. The cohort arrives as episodes, so an algorithm
+        can compare within one episode as well as across them; ``group_rollouts``
+        flattens it for the algorithms that only compare across."""
 
-    async def finalize_rollout(self, rollout: Rollout) -> None:
+    async def finalize_rollout(self, rollout: TrainRollout) -> None:
         """Arrival phase (non-virtual): rollout-local scoring as each rollout is
         tokenized."""
         if rollout.samples:
             await self.score_rollout(rollout)
 
-    async def finalize_group(self, rollouts: list[Rollout]) -> None:
+    async def finalize_group(self, episodes: list[Episode]) -> None:
         """Group phase (non-virtual): group-relative scoring, then stamp each
         sample's wire fields (the advantage stream + loss routing). After this
         the records are frozen — groups die at stamping."""
-        await self.score_group(rollouts)
-        for rollout in rollouts:
+        await self.score_group(episodes)
+        for rollout in group_rollouts(episodes):
             stamp_advantages(rollout)
             for sample in rollout.samples:
                 stamp_loss_routing(sample, self.action_loss_type)

src/prime_rl/orchestrator/algo/echo.py

@@ -11,7 +11,7 @@
 if TYPE_CHECKING:
     import verifiers.v1 as vf
 
-    from prime_rl.orchestrator.types import Rollout
+    from prime_rl.orchestrator.types import TrainRollout
     from prime_rl.utils.client import InferencePool
 
 
@@ -34,12 +34,12 @@ def __init__(self, config: EchoAlgoConfig, policy_pool: InferencePool):
         if config.filter is not None:
             self.filter_fn = partial(import_object(config.filter.import_path), **config.filter.kwargs)
 
-    async def score_rollout(self, rollout: Rollout) -> None:
+    async def score_rollout(self, rollout: TrainRollout) -> None:
         # Observation weighting is rollout-local; the group-relative GRPO
         # baseline is inherited unchanged as ``score_group``.
         self._weight_observations(rollout)
 
-    def _weight_observations(self, rollout: Rollout) -> None:
+    def _weight_observations(self, rollout: TrainRollout) -> None:
         """Write each sample's ``ce_weights`` stream over the env-provided
         observation tokens of later turns. Provenance is structural under v1:
         within a branch, the non-sampled nodes that follow the first model

src/prime_rl/orchestrator/algo/grpo.py

@@ -6,9 +6,10 @@
 
 from prime_rl.configs.algorithm import GRPOAlgoConfig
 from prime_rl.orchestrator.algo.base import Algorithm
+from prime_rl.orchestrator.types import group_rollouts
 
 if TYPE_CHECKING:
-    from prime_rl.orchestrator.types import Rollout
+    from prime_rl.orchestrator.types import Episode
     from prime_rl.utils.client import InferencePool
 
 
@@ -21,15 +22,16 @@ def __init__(self, config: GRPOAlgoConfig, policy_pool: InferencePool):
         super().__init__(config, policy_pool)
         self.length_penalty = config.length_penalty
 
-    async def score_group(self, group: list[Rollout]) -> None:
-        rewards = torch.tensor([rollout.reward for rollout in group], dtype=torch.float32)
+    async def score_group(self, group: list[Episode]) -> None:
+        rollouts = group_rollouts(group)
+        rewards = torch.tensor([rollout.reward for rollout in rollouts], dtype=torch.float32)
         length_penalty = self.length_penalty
         if length_penalty is None:
             advantages = rewards - rewards.mean()
         else:
-            output = torch.tensor([rollout.num_output_tokens for rollout in group], dtype=rewards.dtype)
-            total = torch.tensor([rollout.num_total_tokens for rollout in group], dtype=rewards.dtype)
-            turns = torch.tensor([rollout.num_turns for rollout in group], dtype=rewards.dtype)
+            output = torch.tensor([rollout.num_output_tokens for rollout in rollouts], dtype=rewards.dtype)
+            total = torch.tensor([rollout.num_total_tokens for rollout in rollouts], dtype=rewards.dtype)
+            turns = torch.tensor([rollout.num_turns for rollout in rollouts], dtype=rewards.dtype)
             input = total - output
             penalty_frac = (
                 length_penalty.num_output_tokens_weight * (output / output.max().clamp(min=1))
@@ -39,5 +41,5 @@ async def score_group(self, group: list[Rollout]) -> None:
             penalty = rewards.mean() * penalty_frac
             shaped_rewards = rewards - penalty
             advantages = shaped_rewards - shaped_rewards.mean()
-        for rollout, advantage in zip(group, advantages.tolist(), strict=True):
+        for rollout, advantage in zip(rollouts, advantages.tolist(), strict=True):
             rollout.assign_advantages(advantage)