feat: TOP-D algorithm (trust region policy distillation)

docs/algorithms.md

@@ -48,7 +48,7 @@ name = "Qwen/Qwen3-32B"
 base_url = ["http://localhost:8001/v1"]
 ```
 
-Model *roles* are algorithm-local vocabulary — each algorithm names its reference on the field where the model is actually used, and there is no shared `teacher` slot. `opd` declares a `teacher` field (the frozen model whose reverse KL the policy distills toward); `sft`'s teacher *is* its `sampling.source` (the frozen model it imitates); `opsd` self-distills against the live policy and names no model at all. No role exists outside the algorithm that declares it: the dispatcher, sink, and trainer branch on liveness alone, never on what an algorithm calls a model.
+Model *roles* are algorithm-local vocabulary — each algorithm names its reference on the field where the model is actually used, and there is no shared `teacher` slot. `opd` and `topd` declare a `teacher` field (the frozen model the policy distills toward); `sft`'s teacher *is* its `sampling.source` (the frozen model it imitates); `opsd` self-distills against the live policy and names no model at all. No role exists outside the algorithm that declares it: the dispatcher, sink, and trainer branch on liveness alone, never on what an algorithm calls a model.
 
 So for `opd` set `[orchestrator.algo.teacher]`; for `sft` set `[orchestrator.algo.sampling.source]`; `opsd` needs neither. `opd`'s teacher must be a frozen endpoint — it is typed `FrozenModelConfig`, so `"policy"` isn't representable (the KL would be identically zero); `opsd`'s teacher *is* the live policy by definition (self-distillation conditioned on a demonstration), so it exposes no reference to configure.
 
@@ -68,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`. |
 | `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`. |
+| `topd` | policy | `rl` on actions | Trust Region Policy Distillation ([arXiv:2607.04751](https://arxiv.org/abs/2607.04751)): on-policy distillation through a *proximal teacher* — the per-token reward `log(α·ρ + 1−α)` (`ρ` = teacher/sampler probability ratio) is floored at `log(1−α)` where OPD's `log ρ` diverges, then compiled to token-level group-normalized advantages on the `rl` loss (whose importance ratio + trust region are the paper's internal trust region iterations). Needs a `teacher`; `alpha` defaults to `0.2`. |
 | `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. |
 | `echo` | policy | `rl` on actions + weighted `ce` on observations | ECHO: standard GRPO plus a cross-entropy loss on env-provided tokens already present in the rollout, selected by message role (needs the renderer's role attribution). Defaults to tool-response bodies at `alpha = 0.1` (ECHO's λ); set `roles` to train other roles, each at its own weight. |
@@ -134,6 +135,7 @@ At runtime, each env's resolved config builds two objects: a `Sampler` (`prime_r
 | `echo` | `EchoAlgorithm` | `score_rollout`: weighted ce on observation tokens; `score_group`: group-norm credit (inherited) |
 | `max_rl` | `MaxRLAlgorithm` | `score_group`: mean-normalized group credit |
 | `opd` | `OPDAlgorithm` | `score_rollout`: own-context prefill under the teacher |
+| `topd` | `TOPDAlgorithm` | `score_rollout`: own-context prefill under the teacher; `score_group`: bounded rewards → token-level returns → group-norm credit |
 | `opsd` | `OPSDAlgorithm` | `score_rollout`: demo-conditioned prefill under the live policy |
 | `sft` | `SFTDistillAlgorithm` | `score_group`: group-norm credit (feeds filters) |
 
@@ -273,6 +275,7 @@ The per-token training signal is set by `algo.type` and the [algorithm](#the-alg
 | `max_rl` | `rl` | Mean-normalized group credit (maximum-likelihood RL). |
 | `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 (`model`, 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. |
+| `topd` | `rl` | TOP-D: the bounded distillation reward `log(α·ρ + 1−α)` per action token (`ρ` = teacher/sampler probability ratio), compiled on the orchestrator — per-token return = reward + mean future reward, z-normalized token-level across the group — and shipped as a real advantage stream (advantage-based filters apply). `group_size` is the normalization cohort. |
 | `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. |
 | `sft` | `ce` | Cross-entropy on the sampled tokens. Assigns no advantage — trains on every sampled token. |
 
@@ -318,9 +321,10 @@ Each per-token list must match the rollout's completion-token count exactly —
 
 ### Reference Scoring
 
-`OPDAlgorithm` / `OPSDAlgorithm` do their model I/O in `score_rollout`: as each rollout arrives they query a reference (the sample's own context for `opd`, the demo-conditioned context for `opsd`) and attach per-token reference logprobs to each sample. Rollouts are consumed serially by the orchestrator's main loop and each carries only a handful of samples, so the in-flight request count is naturally bounded — no explicit concurrency cap:
+`OPDAlgorithm` / `TOPDAlgorithm` / `OPSDAlgorithm` do their model I/O in `score_rollout`: as each rollout arrives they query a reference (the sample's own context for `opd`/`topd`, the demo-conditioned context for `opsd`) and attach per-token reference logprobs to each sample. Rollouts are consumed serially by the orchestrator's main loop and each carries only a handful of samples, so the in-flight request count is naturally bounded — no explicit concurrency cap:
 
 - `opd` — score each sample's own context under the `teacher` (a frozen [model reference](#model-references)) via prefill; fills `ref_logprobs` for the `ref_kl` loss component (on-policy distillation). The `teacher` is typed `FrozenModelConfig`, so `"policy"` isn't representable (the KL would be identically zero).
+- `topd` — the same own-context prefill under the `teacher`, but the logprobs never reach the trainer: `score_group` consumes them into token-level advantages (bounded reward → length-normalized future return → group z-normalization) and clears `ref_logprobs` — the signal ships as an ordinary advantage stream on the `rl` loss.
 - `opsd` — SDFT: prepend an expert demonstration as a leading system message (`template`, with a `{demonstration}` placeholder) and score the sample under that demo-conditioned context. The sample is scored verbatim (`hint_block + token_ids`, slicing the hint's logprobs back off), so the join is BPE-clean and it's robust to tool/multimodal prompts and any number of turns. The scoring reference *is* the live policy — self-distillation names no teacher. opsd builds its own renderer to tokenize the hint block: the tokenizer is always the live policy's (not configurable — there is no separate model), and only the `renderer` family is settable (defaults to `"auto"`, resolved from the policy tokenizer; set it to match a non-auto policy renderer). The demonstration is read from the example's `info[demo_key]`, falling back to a top-level rollout field of the same name (e.g. `answer`).
 
 ```toml
@@ -329,7 +333,7 @@ type = "opsd"
 demo_key = "demonstration"
 ```
 
-Scoring runs at arrival, *before* the pre-batch filters, so a rollout that is later filtered still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape (advantage-based filters never fire for opd/opsd anyway, since neither assigns an advantage).
+Scoring runs at arrival, *before* the pre-batch filters, so a rollout that is later filtered still cost its reference compute — accepted for the simpler one-rollout-at-a-time shape (advantage-based filters never fire for opd/opsd, since neither assigns an advantage; topd assigns one, so its filtered rollouts still cost teacher prefill).
 
 ## Filters
 

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

@@ -1,7 +1,8 @@
 """Algorithm abstraction: sampling and the per-token training signal.
 
 An algorithm is a named, self-contained config — a discriminated union keyed
-on ``type`` (``grpo``, ``max_rl``, ``opd``, ``opsd``, ``sft``, ``echo``).
+on ``type`` (``grpo``, ``max_rl``, ``opd``, ``topd``, ``opsd``, ``sft``,
+``echo``).
 The bundle *is* the algorithm: each variant carries
 its sampling component and its credit-assignment / loss-routing parameters,
 and its class defaults are the vetted setting — ``type = "opd"`` with a
@@ -251,6 +252,35 @@ class OPDAlgoConfig(BaseAlgoConfig):
     demo-conditioned self-teaching)."""
 
 
+class TOPDAlgoConfig(BaseAlgoConfig):
+    type: Literal["topd"] = "topd"
+    """Trust Region Policy Distillation (TOP-D, arXiv:2607.04751): on-policy
+    distillation through a *proximal teacher* — the probability-space
+    interpolation ``α·π_teacher + (1−α)·π_student`` — giving the bounded
+    per-token reward ``log(α·ρ + 1−α)`` (``ρ`` = teacher/sampler probability
+    ratio), floored at ``log(1−α)`` where standard OPD's ``log ρ`` diverges.
+    The signal is compiled on the orchestrator into token-level advantages
+    (per-token return = reward + mean future reward, z-normalized across the
+    group's tokens) consumed by the ``rl`` loss component, whose importance
+    ratio and trust region provide the paper's internal trust region
+    iterations. Needs a frozen ``teacher``; ``group_size`` is the
+    normalization cohort."""
+
+    action_loss_type: ClassVar[ActionLossType] = "rl"
+
+    teacher: FrozenModelConfig
+    """The teacher — an inline frozen hosted model (``name`` + ``base_url``)
+    the proximal teacher interpolates toward. Required, and necessarily a
+    frozen endpoint (as for ``opd``: the ratio against the policy itself
+    carries no signal)."""
+
+    alpha: float = Field(0.2, gt=0, lt=1)
+    """Proximal-teacher interpolation coefficient α ∈ (0, 1). Floors the
+    per-token reward at ``log(1−α)``; α → 1 recovers standard OPD's unbounded
+    log-ratio reward. The paper reports α ∈ {0.1, 0.2, 0.3} as uniformly
+    stable and trains with 0.1–0.2."""
+
+
 class OPSDAlgoConfig(BaseAlgoConfig):
     type: Literal["opsd"] = "opsd"
     """On-policy self-distillation (SDFT, https://arxiv.org/abs/2601.19897):
@@ -305,7 +335,7 @@ def require_frozen_source(self):
 
 
 AlgoConfig: TypeAlias = Annotated[
-    GRPOAlgoConfig | EchoAlgoConfig | MaxRLAlgoConfig | OPDAlgoConfig | OPSDAlgoConfig | SFTAlgoConfig,
+    GRPOAlgoConfig | EchoAlgoConfig | MaxRLAlgoConfig | OPDAlgoConfig | TOPDAlgoConfig | OPSDAlgoConfig | SFTAlgoConfig,
     Field(discriminator="type"),
 ]
 """The training algorithm: sampling plus the per-token training signal (credit
@@ -315,6 +345,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).
 - ``opd`` — on-policy distillation: policy samples, per-token reverse KL against a reference model. Needs ``teacher``.
+- ``topd`` — TOP-D: opd through a proximal teacher — bounded per-token reward compiled to token-level group-normalized advantages on the rl loss. 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``.
 - ``echo`` — GRPO on action tokens + weighted CE on tool-response observation tokens.

skills/configs/SKILL.md

@@ -51,7 +51,7 @@ CLI: `--env.0.id reverse-text --env.1.id math-env`.
 
 **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" | "opd" | "topd" | "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/topd (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

@@ -6,7 +6,7 @@
 :class:`~prime_rl.orchestrator.sampler.Sampler`):
 
 - one module per algorithm (``grpo``, ``echo``, ``max_rl``, ``opd``,
-  ``opsd``, ``sft``) — each named class owns its scoring hooks
+  ``topd``, ``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:
@@ -36,6 +36,7 @@
 from prime_rl.orchestrator.algo.opsd import OPSDAlgorithm
 from prime_rl.orchestrator.algo.routing import stamp_advantages, stamp_loss_routing
 from prime_rl.orchestrator.algo.sft import SFTDistillAlgorithm
+from prime_rl.orchestrator.algo.topd import TOPDAlgorithm
 from prime_rl.orchestrator.types import Rollout
 
 if TYPE_CHECKING:
@@ -49,6 +50,7 @@
     "echo": EchoAlgorithm,
     "max_rl": MaxRLAlgorithm,
     "opd": OPDAlgorithm,
+    "topd": TOPDAlgorithm,
     "opsd": OPSDAlgorithm,
     "sft": SFTDistillAlgorithm,
 }
@@ -74,6 +76,7 @@ def build_algorithm(config: AlgoConfig, policy_pool: InferencePool) -> Algorithm
     "OPSDAlgorithm",
     "Rollout",
     "SFTDistillAlgorithm",
+    "TOPDAlgorithm",
     "build_algorithm",
     "connect_frozen_pool",
     "stamp_advantages",