feat(orchestrator): hierarchical GRPO for proposer-solver envs

configs/debug/algo/hierarchical_grpo.toml

@@ -0,0 +1,57 @@
+max_steps = 20
+seq_len = 8192
+
+[model]
+name = "Qwen/Qwen3-4B"
+
+[wandb]
+project = "algorithms-debug"
+name = "debug-hierarchical-grpo"
+
+[orchestrator]
+batch_size = 32
+group_size = 4
+
+[orchestrator.algo]
+type = "hierarchical_grpo"
+episode_agents = ["solver"]
+
+[orchestrator.renderer]
+name = "prime-qwen3"
+
+[orchestrator.train.sampling]
+max_completion_tokens = 4096
+
+[[orchestrator.train.env]]
+name = "proposer-solver"
+env.taskset = { id = "proposer-solver-v1" }
+env.n = 4
+env.proposer.harness = { id = "null" }
+env.proposer.runtime = { type = "subprocess" }
+env.solver.harness = { id = "null" }
+env.solver.runtime = { type = "subprocess" }
+
+[orchestrator.eval]
+interval = 1
+num_examples = 6
+
+[orchestrator.eval.sampling]
+temperature = 1.0
+max_completion_tokens = 4096
+
+[[orchestrator.eval.env]]
+name = "proposer-solver"
+env.taskset = { id = "proposer-solver-v1" }
+env.n = 4
+env.proposer.harness = { id = "null" }
+env.proposer.runtime = { type = "subprocess" }
+env.solver.harness = { id = "null" }
+env.solver.runtime = { type = "subprocess" }
+
+[trainer.optim]
+lr = 3e-6
+
+[ckpt]
+
+[inference]
+gpu_memory_utilization = 0.5

docs/algorithms.md

@@ -17,6 +17,7 @@ This page covers the math and the configurable algorithmic components: the algor
   - [Custom Loss](#custom-loss)
 - [Advantage](#advantage)
   - [Default Advantage](#default-advantage)
+  - [Hierarchical Advantage (Proposer–Solver)](#hierarchical-advantage-proposersolver)
   - [Authoring an Algorithm](#authoring-an-algorithm)
   - [Reference Scoring](#reference-scoring)
 - [Filters](#filters)
@@ -67,6 +68,7 @@ type = "grpo"  # the default
 |---|---|---|---|
 | `grpo` | policy | `rl` on actions | Standard group-relative RL. |
 | `max_rl` | policy | `rl` on actions | MaxRL ([arXiv:2602.02710](https://arxiv.org/abs/2602.02710)): GRPO's centered reward normalized by the group **mean** instead of the standard deviation — the gradient is unbiased for the order-`group_size` truncation of the maximum-likelihood objective, upweighting hard examples like `1/p`. |
+| `hierarchical_grpo` | policy | `rl` on actions | Two-level GRPO for task-generating envs (proposer-solver): agents in `episode_agents` baseline within their own episode (sibling attempts at the same minted task), every other agent across the group (parallel attempts at the same source task). Needs `episode_agents`. See [Hierarchical Advantage](#hierarchical-advantage-proposersolver). |
 | `opd` | policy | `ref_kl` on actions | On-policy distillation ([Thinking Machines](https://thinkingmachines.ai/blog/on-policy-distillation/)): the policy samples, per-token reverse KL against a reference model as the gradient signal. Needs a `teacher`. |
 | `sft` | *(the teacher)* | `ce` on actions | Hard distillation: a frozen model generates rollouts, the policy trains with CE on its tokens. Needs a frozen `sampling.source` (the teacher it samples from). |
 | `opsd` | policy | `ref_kl` on actions | SDFT ([arXiv:2601.19897](https://arxiv.org/abs/2601.19897)): the model is its own reference, conditioned on an expert demonstration. The teacher *is* the live policy (the paper's setting, no extra deployment) — no model to configure. |
@@ -140,6 +142,7 @@ At runtime, each env's resolved config builds two objects: a `Sampler` (`prime_r
 | `grpo` | `GRPOAlgorithm` | `score_group`: group-norm credit (optional length penalty) |
 | `echo` | `EchoAlgorithm` | `score_rollout`: weighted ce on observation tokens; `score_group`: group-norm credit (inherited) |
 | `max_rl` | `MaxRLAlgorithm` | `score_group`: mean-normalized group credit |
+| `hierarchical_grpo` | `HierarchicalGRPOAlgorithm` | `score_group`: two-level peer-set credit (per-episode / per-group) |
 | `opd` | `OPDAlgorithm` | `score_rollout`: own-context prefill under the teacher |
 | `opsd` | `OPSDAlgorithm` | `score_rollout`: demo-conditioned prefill under the live policy |
 | `sft` | `SFTDistillAlgorithm` | `score_group`: group-norm credit (feeds filters) |
@@ -176,7 +179,7 @@ $$
 \mathcal{L} = \frac{\sum \mathcal{L}_{rl}}{N_{rl}} + \frac{\sum \mathcal{L}_{ce}}{N_{ce}} + \frac{\sum \mathcal{L}_{ref\_kl}}{N_{ref\_kl}}
 $$
 
-- `rl` — the configured RL loss (`[trainer.loss]`): DPPO + KL by default, or a [custom loss](#custom-loss). Fed by the group-relative algorithms (`grpo`, `max_rl`, and `echo`'s action tokens).
+- `rl` — the configured RL loss (`[trainer.loss]`): DPPO + KL by default, or a [custom loss](#custom-loss). Fed by the advantage-assigning algorithms (`grpo`, `max_rl`, `hierarchical_grpo`, and `echo`'s action tokens).
 - `ce` — masked NLL. Used for frozen-model tokens (`sft`) and env-observation tokens (`echo`).
 - `ref_kl` — the per-token reverse KL to a reference model ($\log \pi_{\text{ref}} - \log \pi$) as the policy-gradient signal, importance-ratio corrected with a one-sided trust region (`opd`, `opsd`). Requires `ref_logprobs` from a [reference scoring](#reference-scoring); the scoring model must be a vLLM server (it's the only one that exposes `prompt_logprobs`).
 
@@ -278,6 +281,7 @@ The per-token training signal is set by `algo.type` and the [algorithm](#the-alg
 |---|---|---|
 | `grpo` | `rl` | Group-norm: reward minus per-group baseline, optional length penalty. |
 | `max_rl` | `rl` | Mean-normalized group credit (maximum-likelihood RL). |
+| `hierarchical_grpo` | `rl` | Two-level peer-set credit for task-generating envs: episode-scoped agents center within their episode, the rest across the group. |
 | `echo` | `rl` + `ce` | Group-norm on action tokens, plus weighted CE on env-provided tokens selected by message role (each role's `alpha` is its ECHO λ), optionally narrowed by a user filter. |
 | `opd` | `ref_kl` | On-policy distillation: per-token reverse KL to a reference model (`teacher`, an inline frozen hosted model), evaluated in the trainer from shipped reference logprobs. No credit — rollouts keep `advantages = None` (advantage-based filters never fire) and ship no advantage stream; `group_size` only fans out sampling. |
 | `opsd` | `ref_kl` | SDFT: per-token reverse KL to a demo-conditioned reference. No credit — rollouts keep `advantages = None` (advantage-based filters never fire) and ship no advantage stream. |
@@ -299,6 +303,33 @@ type = "grpo"
 type = "linear"
 ```
 
+### Hierarchical Advantage (Proposer–Solver)
+
+Task-generating envs break the flat-group assumption twice over. In `proposer-solver-v1`, one episode is a proposer trace plus the n solver traces its minted problem fanned out to, and a group is `group_size` episodes of the same source task — the proposer asked `group_size` times to invent a problem. Pooling solver rewards across episodes baselines them against attempts at *different* problems; pooling proposer and solver traces mixes reward scales across agents (the proposer earns learnability, the solvers correctness).
+
+`hierarchical_grpo` computes GRPO baselines at two levels of that episode tree. Every trainable trace is mean-centered against its *peer set*:
+
+- agents listed in `episode_agents` (the solvers) against the same-agent traces of **their own episode** — sibling attempts at the same minted task;
+- every other agent (the proposer) against its same-agent traces **across the group** — parallel attempts at the same source task, each rewarded by what its minted task did to the solvers (`proposer-solver-v1`'s learnability, `4p(1−p)`, peaks when half the solvers crack the problem).
+
+Rewards never mix across agents or across minted tasks. A peer set of one centers to zero advantage (GRPO's singleton convention; the zero-advantage filter drops it). `episode_agents` is required — which agents are episode-scoped is env truth the algorithm must not guess:
+
+```toml
+[orchestrator.algo]
+type = "hierarchical_grpo"
+episode_agents = ["solver"]
+
+[[orchestrator.train.env]]
+name = "proposer-solver"
+env.taskset = { id = "proposer-solver-v1" }
+env.n = 4  # solvers per proposed problem
+env.proposer.harness = { id = "null" }
+env.proposer.runtime = { type = "subprocess" }
+env.solver.harness = { id = "null" }
+env.solver.runtime = { type = "subprocess" }
+```
+
+Here `group_size` is the number of proposals per source task and `env.n` the solvers per proposal, so one group carries `group_size × (1 + n)` traces. To train only one side, flip the env's own switches (`env.train_proposer` / `env.train_solver`) — the untrainable side's traces never reach the advantage computation.
 
 ### Authoring an Algorithm
 

docs/training.md

@@ -59,7 +59,7 @@ A condensed view of the knobs you'll most often tune. For trainer-side paralleli
 | `orchestrator.batch_size` | Tasks per trainer step. |
 | `orchestrator.group_size` | Rollouts generated per task. |
 | `orchestrator.max_off_policy_steps` | How many distinct policies may have contributed to one rollout before it's discarded (default 8). The main off-policy dial on long agentic rollouts — bump for throughput, lower for tighter on-policyness. Watch `errored_rollouts` and `mismatch_kl/all/mean` when tuning. |
-| `[orchestrator.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). |
+| `[orchestrator.algo]` | Training algorithm — its `type` names it (`grpo` default, `max_rl`, `hierarchical_grpo`, `opd`, `opsd`, `sft`, `echo`). See [Algorithms](#algorithms). |
 | `[[orchestrator.train.env]]` | Training environments. List multiple tables for multi-env training; weight them via `ratio`. See [Configuration § Environments](configuration.md#environments-orchestratortrainenv). |
 | `[[orchestrator.eval.env]]` + `orchestrator.eval.interval` | Eval environments and cadence (default every 100 steps). |
 
@@ -89,6 +89,7 @@ The RL entrypoint supports several training algorithms, switched via `[orchestra
 |---|---|---|
 | `grpo` (default) | None | Standard group-relative RL |
 | `max_rl` | None | [MaxRL](https://arxiv.org/abs/2602.02710): GRPO with mean-normalized advantages (maximum-likelihood RL) |
+| `hierarchical_grpo` | None | Two-level GRPO for task-generating envs (e.g. `proposer-solver-v1`): episode-scoped agents baseline within their episode, the rest across the group |
 | `opd` | Required, must be vLLM (needs `prompt_logprobs`) | [On-policy distillation](https://thinkingmachines.ai/blog/on-policy-distillation/): the policy generates rollouts, the trainer minimizes per-token reverse KL to a reference model |
 | `sft` | Required, any OpenAI-compatible endpoint | Hard-distill: a frozen model generates rollouts, the policy trains on its tokens |
 | `opsd` | None — the live policy is its own reference (no deployment) | [SDFT](https://arxiv.org/abs/2601.19897): the model is its own reference conditioned on expert demonstrations |

packages/prime-rl-configs/src/prime_rl/configs/algorithm.py

@@ -232,6 +232,29 @@ class MaxRLAlgoConfig(BaseAlgoConfig):
     action_loss_type: ClassVar[ActionLossType] = "rl"
 
 
+class HierarchicalGRPOAlgoConfig(BaseAlgoConfig):
+    type: Literal["hierarchical_grpo"] = "hierarchical_grpo"
+    """Hierarchical GRPO for task-generating envs (``proposer-solver-v1`` and
+    friends): GRPO baselines computed at two levels of the episode tree. A
+    group is ``group_size`` episodes of the same source task, and one episode
+    holds a proposer trace plus the n solver traces its minted task fanned out
+    to. Every trainable trace is mean-centered against its *peer set*: agents
+    listed in ``episode_agents`` (the solvers) against the same-agent traces of
+    their own episode — sibling attempts at the same minted task — and every
+    other agent (the proposer) against its same-agent traces across the group —
+    parallel attempts at the same source task. Rewards never mix across agents
+    or across minted tasks."""
+
+    action_loss_type: ClassVar[ActionLossType] = "rl"
+
+    episode_agents: list[str] = Field(min_length=1)
+    """Agents whose peer set is their own episode (``["solver"]`` for
+    ``proposer-solver-v1``): their reward compares only against sibling traces
+    of the same agent in the same episode, because each episode mints its own
+    task and rewards are not exchangeable across episodes. Required — which
+    agents are episode-scoped is env truth the algorithm must not guess."""
+
+
 class OPDAlgoConfig(BaseAlgoConfig):
     type: Literal["opd"] = "opd"
     """On-policy distillation: the per-token signal is the reverse KL to
@@ -305,7 +328,13 @@ def require_frozen_source(self):
 
 
 AlgoConfig: TypeAlias = Annotated[
-    GRPOAlgoConfig | EchoAlgoConfig | MaxRLAlgoConfig | OPDAlgoConfig | OPSDAlgoConfig | SFTAlgoConfig,
+    GRPOAlgoConfig
+    | EchoAlgoConfig
+    | MaxRLAlgoConfig
+    | HierarchicalGRPOAlgoConfig
+    | OPDAlgoConfig
+    | OPSDAlgoConfig
+    | SFTAlgoConfig,
     Field(discriminator="type"),
 ]
 """The training algorithm: sampling plus the per-token training signal (credit
@@ -314,6 +343,7 @@ def require_frozen_source(self):
 
 - ``grpo`` — policy group sampling, group-relative advantage, RL loss (the default).
 - ``max_rl`` — GRPO with mean-normalized advantages (maximum-likelihood RL).
+- ``hierarchical_grpo`` — two-level GRPO for task-generating envs: episode-scoped agents baseline within their episode, the rest across the group. Needs ``episode_agents``.
 - ``opd`` — on-policy distillation: policy samples, per-token reverse KL against a reference model. Needs ``teacher``.
 - ``opsd`` — SDFT: policy samples, demo-conditioned reverse KL against the live policy (the teacher is the policy itself).
 - ``sft`` — a frozen model samples, the policy trains with CE on its tokens. Needs a frozen ``sampling.source``.

skills/configs/SKILL.md

@@ -54,7 +54,7 @@ CLI: `--orchestrator.train.env.0.env.taskset.id reverse-text-v1`.
 
 **Discriminated unions** — set the `type` field to pick the variant (`[orchestrator.algo] type = "max_rl"`). Omit `type` to keep the default variant.
 
-**Algorithms** — `[orchestrator.algo] type = "grpo" | "max_rl" | "opd" | "opsd" | "sft" | "echo"` — the type names the algorithm (credit assignment + loss routing, fused), and each type's class defaults are its vetted setting; any other key you set is your own assembly (e.g. `[orchestrator.algo.roles.user] alpha = 0.1` for echo — setting any echo role replaces the whole role table). There is no preset layer, and no config hook that points at user code — a new algorithm is a named class in the repo (subclass `Algorithm`, register it). Per-env override: `[orchestrator.train.env.algo] type = "opd"` (the env assembles its own algorithm). prime-rl only hosts the trainable policy; frozen models are inline external endpoints on the algorithm, named where the model is used — `[orchestrator.algo.teacher]` for opd (the frozen model scored against), `[orchestrator.algo.sampling.source]` for sft (the model it samples from), each with `name` + `base_url`. There is no shared `teacher` slot. opsd declares no model — it self-distills against the live policy. See `docs/algorithms.md`.
+**Algorithms** — `[orchestrator.algo] type = "grpo" | "max_rl" | "hierarchical_grpo" | "opd" | "opsd" | "sft" | "echo"` — the type names the algorithm (credit assignment + loss routing, fused), and each type's class defaults are its vetted setting; any other key you set is your own assembly (e.g. `[orchestrator.algo.roles.user] alpha = 0.1` for echo — setting any echo role replaces the whole role table). There is no preset layer, and no config hook that points at user code — a new algorithm is a named class in the repo (subclass `Algorithm`, register it). Per-env override: `[orchestrator.train.env.algo] type = "opd"` (the env assembles its own algorithm). prime-rl only hosts the trainable policy; frozen models are inline external endpoints on the algorithm, named where the model is used — `[orchestrator.algo.teacher]` for opd (the frozen model scored against), `[orchestrator.algo.sampling.source]` for sft (the model it samples from), each with `name` + `base_url`. There is no shared `teacher` slot. opsd declares no model — it self-distills against the live policy. See `docs/algorithms.md`.
 
 **`BaseModel | None` fields** — bare flag enables defaults; nested override enables and sets:
 

src/prime_rl/orchestrator/algo/__init__.py

@@ -5,8 +5,9 @@
 turns the signal half into runtime objects (the sampling half is the env's
 :class:`~prime_rl.orchestrator.sampler.Sampler`):
 
-- one module per algorithm (``grpo``, ``echo``, ``max_rl``, ``opd``,
-  ``opsd``, ``sft``) — each named class owns its scoring hooks
+- one module per algorithm (``grpo``, ``echo``, ``max_rl``,
+  ``hierarchical_grpo``, ``opd``, ``opsd``, ``sft``) — each named class owns
+  its scoring hooks
   (``score_rollout`` / ``score_group``) and declares what it needs (loss
   component, a "teacher", ...). One instance per env, built by
   :func:`build_algorithm`. A new credit-assignment scheme is a new named class:
@@ -31,6 +32,7 @@
 from prime_rl.orchestrator.algo.base import Algorithm, connect_frozen_pool
 from prime_rl.orchestrator.algo.echo import EchoAlgorithm
 from prime_rl.orchestrator.algo.grpo import GRPOAlgorithm
+from prime_rl.orchestrator.algo.hierarchical_grpo import HierarchicalGRPOAlgorithm
 from prime_rl.orchestrator.algo.max_rl import MaxRLAlgorithm
 from prime_rl.orchestrator.algo.opd import OPDAlgorithm
 from prime_rl.orchestrator.algo.opsd import OPSDAlgorithm
@@ -48,6 +50,7 @@
     "grpo": GRPOAlgorithm,
     "echo": EchoAlgorithm,
     "max_rl": MaxRLAlgorithm,
+    "hierarchical_grpo": HierarchicalGRPOAlgorithm,
     "opd": OPDAlgorithm,
     "opsd": OPSDAlgorithm,
     "sft": SFTDistillAlgorithm,
@@ -69,6 +72,7 @@ def build_algorithm(config: AlgoConfig, policy_pool: InferencePool) -> Algorithm
     "Algorithm",
     "EchoAlgorithm",
     "GRPOAlgorithm",
+    "HierarchicalGRPOAlgorithm",
     "MaxRLAlgorithm",
     "OPDAlgorithm",
     "OPSDAlgorithm",

src/prime_rl/orchestrator/algo/hierarchical_grpo.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from prime_rl.configs.algorithm import HierarchicalGRPOAlgoConfig
+from prime_rl.orchestrator.algo.base import Algorithm
+
+if TYPE_CHECKING:
+    from prime_rl.orchestrator.types import Rollout
+    from prime_rl.utils.client import InferencePool
+
+
+class HierarchicalGRPOAlgorithm(Algorithm):
+    """Hierarchical GRPO for task-generating envs (proposer-solver-v1 and
+    friends): GRPO baselines computed at two levels of the episode tree.
+
+    A group is ``group_size`` episodes of the same source task — the proposer
+    asked ``group_size`` times to invent a problem — and each episode holds the
+    proposer's trace plus the n solver traces its minted task fanned out to.
+    Neither level is safe under plain GRPO: pooling solver rewards across
+    episodes baselines them against attempts at *different* problems, and
+    pooling proposer and solver traces mixes reward scales across agents.
+
+    Every trainable trace is instead mean-centered against its peer set:
+
+    - agents in ``episode_agents`` (the solvers) against the same-agent traces
+      of their own episode — sibling attempts at the same minted task;
+    - every other agent (the proposer) against its same-agent traces across
+      the group — parallel attempts at the same source task, rewarded by what
+      their minted tasks did to the solvers (e.g. learnability).
+
+    A peer set of one centers to zero advantage (GRPO's singleton convention;
+    the zero-advantage filter drops it)."""
+
+    def __init__(self, config: HierarchicalGRPOAlgoConfig, policy_pool: InferencePool):
+        super().__init__(config, policy_pool)
+        self.episode_agents = set(config.episode_agents)
+
+    async def score_group(self, group: list[Rollout]) -> None:
+        peers: dict[tuple[str | None, str | None], list[Rollout]] = defaultdict(list)
+        for rollout in group:
+            episode_scoped = rollout.agent_name in self.episode_agents
+            key = (rollout.agent_name, rollout.episode_id if episode_scoped else None)
+            peers[key].append(rollout)
+        for members in peers.values():
+            baseline = sum(rollout.reward for rollout in members) / len(members)
+            for rollout in members:
+                rollout.assign_advantages(rollout.reward - baseline)