feat(orchestrator): RAE — role-conditioned advantage estimation for self-play envs (SPIRAL)

configs/debug/algo/rae.toml

@@ -0,0 +1,53 @@
+max_steps = 20
+seq_len = 4096
+
+[model]
+name = "Qwen/Qwen3-0.6B"
+
+[wandb]
+project = "algorithms-debug"
+name = "debug-rae"
+
+[orchestrator]
+batch_size = 32
+group_size = 1
+
+[orchestrator.algo]
+type = "rae"
+
+[orchestrator.renderer]
+name = "prime-qwen3"
+
+[orchestrator.train.sampling]
+max_completion_tokens = 512
+
+[[orchestrator.train.env]]
+name = "kuhn-poker"
+env.taskset = { id = "kuhn-poker-v1" }
+env.player0.harness = { id = "null" }
+env.player0.runtime = { type = "subprocess" }
+env.player1.harness = { id = "null" }
+env.player1.runtime = { type = "subprocess" }
+
+[orchestrator.eval]
+interval = 1
+num_examples = 16
+
+[orchestrator.eval.sampling]
+max_completion_tokens = 512
+
+[[orchestrator.eval.env]]
+name = "kuhn-poker"
+env.taskset = { id = "kuhn-poker-v1" }
+env.player0.harness = { id = "null" }
+env.player0.runtime = { type = "subprocess" }
+env.player1.harness = { id = "null" }
+env.player1.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)
+  - [Self-Play Advantage (RAE)](#self-play-advantage-rae)
   - [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`. |
+| `rae` | policy | `rl` on actions | RAE (SPIRAL, [arXiv:2506.24119](https://arxiv.org/abs/2506.24119)): reward minus a per-agent EMA baseline of that agent's own rewards — the estimator for multi-agent self-play envs, where the group mean would mix the agents' opposite reward scales. See [Self-Play Advantage](#self-play-advantage-rae). |
 | `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 |
+| `rae` | `RAEAlgorithm` | `score_group`: per-agent EMA-baseline credit |
 | `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`, `rae`, 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). |
+| `rae` | `rl` | Reward minus a per-agent EMA baseline (SPIRAL's role-conditioned advantage estimation) — for multi-agent self-play envs. |
 | `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,27 @@ type = "grpo"
 type = "linear"
 ```
 
+### Self-Play Advantage (RAE)
+
+Group-relative baselines assume the group is exchangeable attempts by one agent. A multi-agent self-play env breaks that: one episode yields one trace per agent, all trainable, and in a zero-sum game the rewards sum to ~0 whatever the policy does — the group mean carries no information, and centering against it converts any structural asymmetry (a first-mover edge) into permanent credit for one agent.
+
+`rae` implements SPIRAL's role-conditioned advantage estimation ([arXiv:2506.24119](https://arxiv.org/abs/2506.24119)): each agent keeps an exponential-moving-average baseline of its own rewards, and every trace's advantage is its reward minus its agent's baseline — measured against the *pre-update* baseline (the unbiased order), then folded in at `decay` (SPIRAL's α, default 0.95). The algorithm instance is per-env, so baselines are keyed per (env, agent) — the paper's per (game, role). Advantages are not normalized, and `group_size` is free (RAE needs no sibling rollouts; `group_size = 1` is fine). Baselines live in orchestrator memory and re-warm from 0 over ~`1/(1 − decay)` traces per agent after a restart.
+
+```toml
+[orchestrator.algo]
+type = "rae"
+decay = 0.95
+
+[[orchestrator.train.env]]
+name = "kuhn-poker"
+env.taskset = { id = "kuhn-poker-v1" }
+env.player0.harness = { id = "null" }
+env.player0.runtime = { type = "subprocess" }
+env.player1.harness = { id = "null" }
+env.player1.runtime = { type = "subprocess" }
+```
+
+Both of `kuhn-poker-v1`'s agents late-bind to the run's own model — shared-policy self-play against a continuously improving opponent. Pin one agent to a frozen endpoint (`env.player1.model = ...`) for asymmetric play; its traces are marked untrainable by the env and never reach the advantage computation. A single-agent env under `rae` degrades to REINFORCE with an EMA baseline.
 
 ### 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`, `rae`, `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) |
+| `rae` | None | [SPIRAL](https://arxiv.org/abs/2506.24119)'s role-conditioned advantage estimation: reward minus a per-agent EMA baseline, for multi-agent self-play envs (e.g. `kuhn-poker-v1`) |
 | `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,27 @@ class MaxRLAlgoConfig(BaseAlgoConfig):
     action_loss_type: ClassVar[ActionLossType] = "rl"
 
 
+class RAEAlgoConfig(BaseAlgoConfig):
+    type: Literal["rae"] = "rae"
+    """RAE — role-conditioned advantage estimation (SPIRAL,
+    https://arxiv.org/abs/2506.24119): scalar advantage = reward minus a
+    per-agent EMA baseline of that agent's own rewards, consumed by the ``rl``
+    loss component. The advantage estimator for multi-agent self-play envs
+    (``kuhn-poker-v1`` and friends): in a zero-sum game the group mean is ~0
+    whatever the policy does, so a group-relative baseline mixes the agents'
+    opposite reward scales — a structural first-mover edge would read as
+    permanent credit. Per-agent baselines measure each agent against its own
+    expected reward instead. Works with any ``group_size`` (including 1)."""
+
+    action_loss_type: ClassVar[ActionLossType] = "rl"
+
+    decay: float = Field(0.95, ge=0.0, lt=1.0)
+    """EMA decay of the per-agent baselines (SPIRAL's α): after a trace is
+    scored, its agent's baseline moves as ``baseline ← decay · baseline +
+    (1 − decay) · reward``. Baselines start at 0 and live in orchestrator
+    memory — a restart re-warms them over ~1/(1 − decay) traces per agent."""
+
+
 class OPDAlgoConfig(BaseAlgoConfig):
     type: Literal["opd"] = "opd"
     """On-policy distillation: the per-token signal is the reverse KL to
@@ -305,7 +326,7 @@ def require_frozen_source(self):
 
 
 AlgoConfig: TypeAlias = Annotated[
-    GRPOAlgoConfig | EchoAlgoConfig | MaxRLAlgoConfig | OPDAlgoConfig | OPSDAlgoConfig | SFTAlgoConfig,
+    GRPOAlgoConfig | EchoAlgoConfig | MaxRLAlgoConfig | RAEAlgoConfig | OPDAlgoConfig | OPSDAlgoConfig | SFTAlgoConfig,
     Field(discriminator="type"),
 ]
 """The training algorithm: sampling plus the per-token training signal (credit
@@ -314,6 +335,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).
+- ``rae`` — reward minus a per-agent EMA baseline (SPIRAL), for multi-agent self-play envs.
 - ``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" | "rae" | "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,7 +5,7 @@
 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``,
+- one module per algorithm (``grpo``, ``echo``, ``max_rl``, ``rae``, ``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
@@ -34,6 +34,7 @@
 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
+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
@@ -48,6 +49,7 @@
     "grpo": GRPOAlgorithm,
     "echo": EchoAlgorithm,
     "max_rl": MaxRLAlgorithm,
+    "rae": RAEAlgorithm,
     "opd": OPDAlgorithm,
     "opsd": OPSDAlgorithm,
     "sft": SFTDistillAlgorithm,
@@ -72,6 +74,7 @@ def build_algorithm(config: AlgoConfig, policy_pool: InferencePool) -> Algorithm
     "MaxRLAlgorithm",
     "OPDAlgorithm",
     "OPSDAlgorithm",
+    "RAEAlgorithm",
     "Rollout",
     "SFTDistillAlgorithm",
     "build_algorithm",

src/prime_rl/orchestrator/algo/rae.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from prime_rl.configs.algorithm import RAEAlgoConfig
+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 RAEAlgorithm(Algorithm):
+    """Role-conditioned advantage estimation (SPIRAL, arXiv:2506.24119):
+    credit = reward minus a per-agent EMA baseline of that agent's rewards.
+
+    The advantage estimator for multi-agent self-play envs (kuhn-poker-v1 and
+    friends): in a zero-sum game the group mean is ~0 whatever the policy does,
+    and centering all agents against it mixes their opposite reward scales — a
+    structural first-mover edge would read as permanent credit for one agent.
+    A per-agent baseline tracks each agent's own expected reward instead, so
+    credit measures improvement over that agent's usual result. The instance is
+    per-env, so baselines are keyed per (env, agent) — the paper's per
+    (game, role). A single-agent env degrades to REINFORCE with an EMA baseline.
+
+    Each trace is scored against the baseline *before* its reward updates it
+    (the reference implementation's unbiased order). Baselines start at 0 and
+    live in orchestrator memory: a restart re-warms them over ~1/(1 − decay)
+    traces per agent."""
+
+    def __init__(self, config: RAEAlgoConfig, policy_pool: InferencePool):
+        super().__init__(config, policy_pool)
+        self.decay = config.decay
+        self.baselines: dict[str | None, float] = defaultdict(float)
+
+    async def score_group(self, group: list[Rollout]) -> None:
+        for rollout in group:
+            baseline = self.baselines[rollout.agent_name]
+            rollout.assign_advantages(rollout.reward - baseline)
+            self.baselines[rollout.agent_name] = self.decay * baseline + (1.0 - self.decay) * rollout.reward