From 9128325593311925167ad813f22c3c4b62f7ce82 Mon Sep 17 00:00:00 2001 From: hallerite Date: Fri, 24 Jul 2026 14:17:02 +0000 Subject: [PATCH] feat(orchestrator): hierarchical GRPO for proposer-solver envs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add the `hierarchical_grpo` algorithm: GRPO baselines at two levels of a task-generating env's episode tree. A group is group_size episodes of the same source task (the proposer asked group_size times); each episode holds the proposer's trace plus the n solver traces its minted task fanned out to. Agents in `episode_agents` (the solvers) mean-center against their own episode's same-agent traces — sibling attempts at the same minted task — and every other agent (the proposer) against its same-agent traces across the group. Rewards never mix across agents or minted tasks; singleton peer sets center to zero. configs/debug/algo/hierarchical_grpo.toml runs proposer-solver-v1 with Qwen3-4B (4 proposals per topic x 4 solvers per proposal, learnability reward on the proposer), both sides trainable. Co-Authored-By: Claude Fable 5 --- configs/debug/algo/hierarchical_grpo.toml | 57 +++++++++++++++ docs/algorithms.md | 33 ++++++++- docs/training.md | 3 +- .../src/prime_rl/configs/algorithm.py | 32 ++++++++- skills/configs/SKILL.md | 2 +- src/prime_rl/orchestrator/algo/__init__.py | 8 ++- .../orchestrator/algo/hierarchical_grpo.py | 49 +++++++++++++ tests/unit/orchestrator/test_advantage.py | 70 +++++++++++++++++++ 8 files changed, 248 insertions(+), 6 deletions(-) create mode 100644 configs/debug/algo/hierarchical_grpo.toml create mode 100644 src/prime_rl/orchestrator/algo/hierarchical_grpo.py diff --git a/configs/debug/algo/hierarchical_grpo.toml b/configs/debug/algo/hierarchical_grpo.toml new file mode 100644 index 0000000000..0518432c89 --- /dev/null +++ b/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 diff --git a/docs/algorithms.md b/docs/algorithms.md index 8088d28516..e43dc56a56 100644 --- a/docs/algorithms.md +++ b/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 diff --git a/docs/training.md b/docs/training.md index ce2bacae6d..3698b3a76a 100644 --- a/docs/training.md +++ b/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 | diff --git a/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py b/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py index b3282475af..41e7f8cb7c 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py +++ b/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``. diff --git a/skills/configs/SKILL.md b/skills/configs/SKILL.md index 8153564277..55d7689162 100644 --- a/skills/configs/SKILL.md +++ b/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: diff --git a/src/prime_rl/orchestrator/algo/__init__.py b/src/prime_rl/orchestrator/algo/__init__.py index 8d1baa60a3..01947ca9a3 100644 --- a/src/prime_rl/orchestrator/algo/__init__.py +++ b/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", diff --git a/src/prime_rl/orchestrator/algo/hierarchical_grpo.py b/src/prime_rl/orchestrator/algo/hierarchical_grpo.py new file mode 100644 index 0000000000..0debde3126 --- /dev/null +++ b/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) diff --git a/tests/unit/orchestrator/test_advantage.py b/tests/unit/orchestrator/test_advantage.py index 247184fee6..1092234d53 100644 --- a/tests/unit/orchestrator/test_advantage.py +++ b/tests/unit/orchestrator/test_advantage.py @@ -5,10 +5,12 @@ from prime_rl.configs.algorithm import ( GRPOAlgoConfig, + HierarchicalGRPOAlgoConfig, LinearLengthPenaltyConfig, MaxRLAlgoConfig, ) 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.trajectories import trace_to_samples from prime_rl.orchestrator.types import Rollout @@ -183,6 +185,74 @@ def test_max_rl_mean_normalized(): assert _max_rl(_make_group(rewards=[1.0, 1.0])) == pytest.approx([0.0, 0.0]) +# -------------------------------------------------------------------------- +# Hierarchical GRPO: two-level baselines for task-generating envs. +# -------------------------------------------------------------------------- + + +def _episode_rollout(reward: float, agent_name: str, episode_id: str) -> Rollout: + """A one-turn rollout stamped as one agent's trace of one episode — what a + task-generating env's traces look like by the time ``score_group`` sees them.""" + rollout = _build_rollout(reward, sampled_lengths=[2]) + rollout.agent = vf.AgentInfo(config=vf.AgentConfig(), name=agent_name) + rollout.episode_id = episode_id + return rollout + + +def _hier_grpo(group: list[Rollout], episode_agents: list[str]) -> list[float]: + """Drive ``HierarchicalGRPOAlgorithm.score_group`` and read back the scalars.""" + algo = HierarchicalGRPOAlgorithm(HierarchicalGRPOAlgoConfig(episode_agents=episode_agents), policy_pool=None) + asyncio.run(algo.score_group(group)) + return [_scalar(rollout) for rollout in group] + + +def test_hier_grpo_solvers_center_within_their_episode(): + """Solvers are baselined only against siblings attempting the same minted + task: an episode whose problem got 1-of-2 solved carries signal, while + pooling all four solver rewards (mean 0.25) would mis-credit both episodes.""" + group = [ + _episode_rollout(1.0, "solver", "ep1"), + _episode_rollout(0.0, "solver", "ep1"), + _episode_rollout(0.0, "solver", "ep2"), + _episode_rollout(0.0, "solver", "ep2"), + ] + advs = _hier_grpo(group, episode_agents=["solver"]) + assert advs[:2] == pytest.approx([0.5, -0.5]) + assert advs[2:] == pytest.approx([0.0, 0.0]) + + +def test_hier_grpo_proposers_center_across_the_group(): + """The proposer's peer set is its own traces across the group's episodes — + solver rewards never leak into its baseline, nor its into theirs.""" + group = [ + _episode_rollout(1.0, "proposer", "ep1"), + _episode_rollout(1.0, "solver", "ep1"), + _episode_rollout(0.0, "proposer", "ep2"), + _episode_rollout(1.0, "solver", "ep2"), + ] + advs = _hier_grpo(group, episode_agents=["solver"]) + assert advs[0] == pytest.approx(0.5) + assert advs[2] == pytest.approx(-0.5) + # each episode's lone solver is a singleton peer set: zero advantage + assert advs[1] == pytest.approx(0.0) + assert advs[3] == pytest.approx(0.0) + + +def test_hier_grpo_proposer_solver_shape(): + """The proposer-solver-v1 shape end to end: 2 proposals x 2 solvers, all + peer sets centered independently and each summing to zero.""" + group = [ + _episode_rollout(1.0, "proposer", "ep1"), + _episode_rollout(1.0, "solver", "ep1"), + _episode_rollout(0.0, "solver", "ep1"), + _episode_rollout(0.0, "proposer", "ep2"), + _episode_rollout(1.0, "solver", "ep2"), + _episode_rollout(1.0, "solver", "ep2"), + ] + advs = _hier_grpo(group, episode_agents=["solver"]) + assert advs == pytest.approx([0.5, 0.5, -0.5, -0.5, 0.0, 0.0]) + + # -------------------------------------------------------------------------- # GRPO linear length penalty: pass_rate-scaled penalty before the baseline. # --------------------------------------------------------------------------