From af745473a3c12d8fbd17df9807e5ccb9ba00dac9 Mon Sep 17 00:00:00 2001 From: hallerite Date: Tue, 7 Jul 2026 05:16:11 +0000 Subject: [PATCH] feat: TOP-D algorithm (trust region policy distillation) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Implements Trust Region Policy Distillation (arXiv:2607.04751) as a new algorithm type "topd": on-policy distillation through a proximal teacher. The teacher prefill-scores each sample at arrival (opd's path), then group finalization compiles the bounded per-token reward log(alpha*rho + 1-alpha) into token-level advantages (length-normalized future return, group z-normalization) consumed by the standard rl loss — whose importance ratio and trust region provide the paper's internal trust region iterations. Co-Authored-By: Claude Fable 5 --- docs/algorithms.md | 10 +- .../src/prime_rl/configs/algorithm.py | 35 ++++++- skills/configs/SKILL.md | 2 +- src/prime_rl/orchestrator/algo/__init__.py | 5 +- src/prime_rl/orchestrator/algo/topd.py | 94 +++++++++++++++++++ tests/unit/orchestrator/test_algorithms.py | 48 ++++++++-- 6 files changed, 180 insertions(+), 14 deletions(-) create mode 100644 src/prime_rl/orchestrator/algo/topd.py diff --git a/docs/algorithms.md b/docs/algorithms.md index 4cdb83a08d..3abc5e329f 100644 --- a/docs/algorithms.md +++ b/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 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..af2511fb62 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/algorithm.py +++ b/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. diff --git a/skills/configs/SKILL.md b/skills/configs/SKILL.md index 7eebe87f86..9213f16c82 100644 --- a/skills/configs/SKILL.md +++ b/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: diff --git a/src/prime_rl/orchestrator/algo/__init__.py b/src/prime_rl/orchestrator/algo/__init__.py index 8d1baa60a3..df29d98c7c 100644 --- a/src/prime_rl/orchestrator/algo/__init__.py +++ b/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", diff --git a/src/prime_rl/orchestrator/algo/topd.py b/src/prime_rl/orchestrator/algo/topd.py new file mode 100644 index 0000000000..1d8563495a --- /dev/null +++ b/src/prime_rl/orchestrator/algo/topd.py @@ -0,0 +1,94 @@ +from __future__ import annotations + +import asyncio +import math +from typing import TYPE_CHECKING + +import torch + +from prime_rl.configs.algorithm import TOPDAlgoConfig +from prime_rl.orchestrator.algo.base import Algorithm +from prime_rl.utils.client import StaticInferencePool + +if TYPE_CHECKING: + from prime_rl.orchestrator.types import Rollout + from prime_rl.transport import TrainingSample + from prime_rl.utils.client import InferencePool + + +class TOPDAlgorithm(Algorithm): + """Trust Region Policy Distillation (TOP-D, arXiv:2607.04751). Needs a + teacher, like opd — but where opd's unbounded log-ratio reward diverges + when the teacher rejects a sampled token, TOP-D distills toward a *proximal + teacher*, the probability-space interpolation ``α·π_teacher + + (1−α)·π_sampler``. That construction never needs to be materialized: the + per-token reward reduces to ``r̃ = log(α·ρ + 1−α)`` with ``ρ`` the + teacher/sampler probability ratio, floored at ``log(1−α)``. + + The signal compiles to ordinary rl credit rather than a trainer-side KL: + each sample's trainable tokens get a return ``R̃_k = r̃_k + mean(r̃_{k+1:})`` + (the paper's length-normalized future return, Eq. 8), z-normalized + *token-level* across the whole group, and shipped as a per-token advantage + stream on the ``rl`` loss — whose importance ratio against the sampling + logprobs and trust-region masking are exactly the paper's internal trust + region iterations. Advantage-based filters apply as they would for grpo.""" + + def __init__(self, config: TOPDAlgoConfig, policy_pool: InferencePool): + super().__init__(config, policy_pool) + self.teacher = config.teacher + self.alpha = config.alpha + self.teacher_pool: StaticInferencePool | None = None # static teacher endpoint, connected in setup() + + async def setup(self) -> None: + pool = await self.connect(self.teacher) + if not isinstance(pool, StaticInferencePool): + raise TypeError("topd teacher must be a static endpoint — prefill scoring needs fixed endpoints") + self.teacher_pool = pool + + async def score_rollout(self, rollout: Rollout) -> None: + pool = self.teacher_pool + assert pool is not None, "teacher pool not connected — Algorithm.setup() must run first" + + async def score_sample(sample: TrainingSample) -> None: + sample.ref_logprobs = await pool.score(list(sample.token_ids)) + + await asyncio.gather(*(score_sample(sample) for sample in rollout.samples)) + + def _sample_returns(self, sample: TrainingSample) -> torch.Tensor: + """Token-level returns over the sample's trainable tokens: the TOP-D + reward plus the mean of the later trainable tokens' rewards — the + length-normalized future return of the paper's Eq. 8.""" + assert sample.ref_logprobs is not None, "sample not teacher-scored — score_rollout must run first" + mask = torch.tensor(sample.mask, dtype=torch.bool) + teacher_logprobs = torch.tensor(sample.ref_logprobs, dtype=torch.float32)[mask] + sampler_logprobs = torch.tensor(sample.logprobs, dtype=torch.float32)[mask] + # r̃ = log(α·ρ + 1−α), via logaddexp so large log-ratios never overflow + log_ratio = teacher_logprobs - sampler_logprobs + rewards = torch.logaddexp(log_ratio + math.log(self.alpha), torch.full_like(log_ratio, math.log1p(-self.alpha))) + future_sum = rewards.flip(0).cumsum(0).flip(0) - rewards + num_future = torch.arange(len(rewards) - 1, -1, -1, dtype=torch.float32) + return rewards + future_sum / num_future.clamp(min=1.0) + + async def score_group(self, group: list[Rollout]) -> None: + returns: list[list[torch.Tensor]] = [] + for rollout in group: + rollout_returns = [] + for sample in rollout.samples: + rollout_returns.append(self._sample_returns(sample)) + sample.ref_logprobs = None # consumed — the rl loss reads advantages, not reference logprobs + returns.append(rollout_returns) + + flat = [r for rollout_returns in returns for r in rollout_returns] + all_returns = torch.cat(flat) if flat else torch.empty(0) + if all_returns.numel() == 0: + return + mean = all_returns.mean() + std = all_returns.std() if all_returns.numel() > 1 else all_returns.new_zeros(()) + std = std.clamp(min=1e-6) # a group with no spread carries all-zero advantages, like a uniform grpo group + + for rollout, rollout_returns in zip(group, returns, strict=True): + stream: list[float] = [] + for sample, sample_returns in zip(rollout.samples, rollout_returns, strict=True): + advantages = iter(((sample_returns - mean) / std).tolist()) + stream.extend(next(advantages) if trainable else 0.0 for trainable in sample.mask) + rollout.assign_advantages(stream) diff --git a/tests/unit/orchestrator/test_algorithms.py b/tests/unit/orchestrator/test_algorithms.py index 1eae928a46..51c24d7bb3 100644 --- a/tests/unit/orchestrator/test_algorithms.py +++ b/tests/unit/orchestrator/test_algorithms.py @@ -8,7 +8,7 @@ from verifiers.v1.types import AssistantMessage, ToolMessage, UserMessage from prime_rl.configs.algorithm import AlgoConfig, FrozenModelConfig -from prime_rl.orchestrator.algo import EchoAlgorithm, stamp_advantages, stamp_loss_routing +from prime_rl.orchestrator.algo import EchoAlgorithm, TOPDAlgorithm, stamp_advantages, stamp_loss_routing from prime_rl.orchestrator.trajectories import trace_to_samples from prime_rl.orchestrator.types import Rollout from prime_rl.transport.types import TrainingSample @@ -38,6 +38,7 @@ def _ref_kind(ref): ("grpo", {}, "policy", "rl"), ("max_rl", {}, "policy", "rl"), ("opd", {"teacher": FROZEN}, "policy", "ref_kl"), + ("topd", {"teacher": FROZEN}, "policy", "rl"), ("sft", {"sampling": {"source": FROZEN}}, "frozen", "ce"), ("opsd", {}, "policy", "ref_kl"), ("echo", {}, "policy", "rl"), @@ -68,14 +69,15 @@ def test_echo_roles_require_at_least_one(): _build(type="echo", roles={}) -def test_opd_teacher_must_be_a_frozen_endpoint(): - # opd needs a teacher, and it must be frozen: a missing teacher is a - # structural error, and "policy" can't even be set — opd.teacher is typed - # FrozenModelConfig (the KL against the policy itself would be zero). +@pytest.mark.parametrize("algorithm_type", ["opd", "topd"]) +def test_distillation_teacher_must_be_a_frozen_endpoint(algorithm_type): + # opd/topd need a teacher, and it must be frozen: a missing teacher is a + # structural error, and "policy" can't even be set — the teacher is typed + # FrozenModelConfig (the signal against the policy itself would be zero). with pytest.raises(ValueError, match="Field required"): - _build(type="opd") + _build(type=algorithm_type) with pytest.raises(ValueError, match="FrozenModelConfig"): - _build(type="opd", teacher="policy") + _build(type=algorithm_type, teacher="policy") def test_sft_requires_teacher(): @@ -207,6 +209,38 @@ def test_assign_advantages_list_rejects_misaligned(): rollout.assign_advantages([0.5]) +# -------------------------------------------------------------------------- +# TOP-D: bounded distillation reward compiled to token-level group credit. +# -------------------------------------------------------------------------- + + +def _topd_sample(logprobs: list[float], ref_logprobs: list[float]) -> TrainingSample: + # One prompt token (mask False, filler logprobs), the rest trainable. + return TrainingSample( + token_ids=list(range(len(logprobs))), + mask=[False] + [True] * (len(logprobs) - 1), + logprobs=logprobs, + temperatures=[], + env_name="test-env", + ref_logprobs=ref_logprobs, + ) + + +def test_topd_token_level_group_advantages(): + # Hand-computed with the default alpha = 0.2: per trainable token, reward + # r = log(0.2 * exp(teacher - sampler) + 0.8); return = r + mean of the + # later trainable tokens' rewards; advantages z-normalize the returns + # across the whole group's tokens (token-level, unbiased std). + algo = TOPDAlgorithm(_build(type="topd", teacher=FROZEN), MagicMock()) + a = _make_rollout([_topd_sample([0.0, -0.5, -1.0, -0.2], [0.0, -0.3, -2.0, -0.2])]) + b = _make_rollout([_topd_sample([0.0, -0.4, -0.1], [0.0, -0.1, -0.6])]) + asyncio.run(algo.score_group([a, b])) + assert a.advantages == pytest.approx([0.0, 0.4772, -1.4909, 0.9075], abs=1e-3) + assert b.advantages == pytest.approx([0.0, 0.6532, -0.5469], abs=1e-3) + # Teacher logprobs are consumed into advantages, never shipped on the wire. + assert a.samples[0].ref_logprobs is None and b.samples[0].ref_logprobs is None + + # -------------------------------------------------------------------------- # Echo: weighted CE on env-provided observation tokens of later turns. #