diff --git a/configs/debug/algo/rae.toml b/configs/debug/algo/rae.toml new file mode 100644 index 0000000000..e601277aa9 --- /dev/null +++ b/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 diff --git a/docs/algorithms.md b/docs/algorithms.md index 8088d28516..48a599f251 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) + - [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 diff --git a/docs/training.md b/docs/training.md index ce2bacae6d..12be2dc58f 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`, `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 | 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..1549c2db43 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,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``. diff --git a/skills/configs/SKILL.md b/skills/configs/SKILL.md index 8153564277..887340cfb3 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" | "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: diff --git a/src/prime_rl/orchestrator/algo/__init__.py b/src/prime_rl/orchestrator/algo/__init__.py index 8d1baa60a3..cc5ddd60ed 100644 --- a/src/prime_rl/orchestrator/algo/__init__.py +++ b/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", diff --git a/src/prime_rl/orchestrator/algo/rae.py b/src/prime_rl/orchestrator/algo/rae.py new file mode 100644 index 0000000000..c93b4d242f --- /dev/null +++ b/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