research/08-sdpo-grpo-integration.md

# SDPO ⊕ Dr. GRPO: wiring the on-policy KL-at-error-turns channel into a live RL loop

> **Design date:** 2026-05-28.
> **Scope:** A concrete, implementable design for adding the SDPO auxiliary
> loss channel (on-policy KL at error turns, teacher = same weights conditioned
> on a hint) as a **second loss head** on a live **Dr. GRPO** update step. Targets
> the two integration substrates already in this repo: the **PRIME-RL parity
> recipe** (`recipes/prime_rl/composer_loss.py`) and the **TRL `GRPOTrainer`
> subclass** (`trainer/composer_trainer.py`). Recommends the TRL subclass as the
> host and gives a ~70-LoC `ComposerGRPOTrainer` sketch.
> **Method:** Lead with local-file analysis of `loss.py`, `composer_loss.py`,
> `composer_trainer.py`, `data_collator.py`, plus `research/07` (HintGenerator)
> and `research/10` (the Dr. GRPO target). One bounded TRL API lookup
> (`mcp_exa_get_code_context_exa` on `huggingface/trl@main`) to confirm the
> `GRPOTrainer` loss-override surface; the DeepWiki follow-up timed out, so the
> version-robust guard in §4 documents both the `_compute_loss(self, model,
> inputs)` internal hook (what this repo already overrides) and the public
> `compute_loss(self, model, inputs, return_outputs=False,
> num_items_in_batch=None)` HF-Trainer wrapper.

---

## TL;DR

SDPO is **not** the GRPO-KL-to-reference term and must not be folded into it. It
is a **separate distillation head**: a generalized-JSD between the student's
on-policy logits and the **same model's** logits when its context has a hint
spliced in at the error turn, masked to the post-hint recovery tokens. The
integration is therefore "compute the Dr. GRPO loss as usual, then **add
`beta_sdpo · JSD_error_turns`** before `.backward()`."

- **Host = the TRL `GRPOTrainer` subclass.** It already exists
  (`ComposerReplicationTrainer`), already overrides the loss with exactly this
  `grpo + alpha*sdpo + beta*replay` shape, and — decisively — it has **full
  logits** in `_compute_loss`. The PRIME-RL recipe **cannot** host SDPO today:
  its `LossInputs` exposes per-token **log-probs only, not full vocabulary
  logits**, and `composer_loss.py` correctly raises `NotImplementedError` when
  `alpha_sdpo>0`. SDPO needs the full distribution; PRIME-RL is blocked until
  upstream exposes logits.
- **Attach point:** inside the Dr. GRPO update step, after the policy-gradient +
  k1-KL loss is computed on the minibatch, run one **student forward (grad)** +
  one **teacher forward (`no_grad`, hint-spliced context)**, take
  `generalized_jsd_loss` masked to `sdpo_loss_mask`, scale by `beta_sdpo`, and
  add. Single-epoch Dr. GRPO makes this clean: the teacher forward happens on
  **the same minibatch being updated**, so the KL is genuinely on-policy.
- **Dr. GRPO specifics are preserved untouched:** SDPO touches neither the
  advantage estimator (no std-norm, no length-standardization) nor the GRPO
  **k1** (`−log r`) KL-to-ref. It is purely additive.
- **CPU-testable:** a 1–2 rollout Dr. GRPO step on Qwen2.5-0.5B with the SDPO
  channel on, mirroring the existing `examples/sdpo_real_trace_train_smoke`.

---

## 1. The two in-repo substrates, and why TRL is the host

### 1.1 Substrate A — TRL `GRPOTrainer` subclass (`trainer/composer_trainer.py`)

Already in the repo and already the right shape. `ComposerReplicationTrainer`
subclasses `trl.GRPOTrainer` and overrides:

```python
def _compute_loss(self, model, inputs) -> torch.Tensor:
    grpo_loss  = super()._compute_loss(model, inputs)         # channel 1
    sdpo_kl    = self._compute_sdpo_loss(model, inputs)       # channel 2
    replay_dpo = self._compute_trace_replay_loss(model, inputs)
    return grpo_loss + self.alpha_sdpo*sdpo_kl + self.beta_replay*replay_dpo
```

`_compute_sdpo_loss` (lines 133–178) already does the **student forward (grad) +
teacher forward (`no_grad`) over `ctx_teacher_input_ids`**, the
`student_logits.shape == teacher_logits.shape` gate, and
`generalized_jsd_loss(..., labels=inputs["sdpo_loss_mask"], beta, temperature,
token_clip, reduction="batchmean")`. This is the SDPO channel, intact. **It has
full logits** — the prerequisite PRIME-RL lacks.

**Decisive property:** TRL hands the subclass `model` and `inputs` and lets it
return any scalar; full `.logits` are available for both the student and the
hint-conditioned teacher forward. SDPO is a drop-in.

### 1.2 Substrate B — PRIME-RL parity recipe (`recipes/prime_rl/composer_loss.py`)

PRIME-RL's `CustomLossConfig` takes an importable `loss_fn(inputs:
LossInputs)` called **once per sample** on **1-D `(seq,)` tensors**. Channel 1
(DPPO + k1-style KL on the importance ratio) is **byte-for-byte parity-verified**
against upstream `default_loss_fn` and is an excellent Dr.-GRPO-adjacent PG loss.

But SDPO is **deferred by construction**:

```python
# composer_loss.py, lines 257-268
teacher_lp = getattr(inputs, "teacher_logprobs", None)
if alpha_sdpo > 0:
    raise NotImplementedError(
        "SDPO channel in the PRIME-RL recipe is deferred. PRIME-RL v0.5 "
        "exposes (seq,) log-probs through LossInputs but not full vocabulary "
        "logits, and SDPO/OPSD requires the full distribution. ...")
```

`generalized_jsd_loss` calls `log_softmax(dim=-1)` over the vocab axis. With
only a `(seq,)` log-prob vector there is **no vocab axis** — softmax of a
1-element slice is identically 1.0 and `log` is 0, i.e. a mathematically
degenerate, silently-zero channel (the Wave-13 finding the docstring cites). So
SDPO in PRIME-RL is blocked **until upstream exposes per-token full logits**, not
a thing we can paper over.

### 1.3 Recommendation

**Host the SDPO aux channel in the TRL `GRPOTrainer` subclass.** Rationale:

1. **Logits available** — the one hard requirement SDPO has and PRIME-RL lacks.
2. **The override already exists** with the exact additive shape; we
   re-point channel 1 at Dr. GRPO and tighten the teacher forward (§4).
3. **Single-process, CPU-runnable** — matches the existing smoke harness, so the
   SDPO-on Dr.-GRPO step is testable today (§6) without PRIME-RL's 3-actor mesh.
4. PRIME-RL stays the **scale/parity** path for channel-1-only runs; SDPO lands
   there for free the moment `LossInputs.teacher_logits` (full distribution)
   exists upstream — the adapter is otherwise ready.

> One caveat to fix while we're here: the current `ComposerReplicationTrainer`
> channel 1 is *vanilla* GRPO (`super()._compute_loss`). The Composer target is
> **Dr. GRPO** (`research/10`): length-standardization removed, **no std-dev
> advantage normalization**, **k1** (`−log r`) KL, Adam, single-epoch. §3 + §4
> pin those into the subclass; SDPO rides on top unchanged.

---

## 2. The exact attach point + data flow

SDPO attaches **inside one Dr. GRPO update step, after the PG+KL loss is formed,
before backward**. It is one extra additive scalar. Concretely, per minibatch:

```
            ┌─────────────────────── one Dr. GRPO update step (single-epoch) ──────────────────────┐
rollout ──▶ │  Channel 1 (Dr. GRPO):                                                                │
trajectory  │    advantages = (R - group_mean)          # NO /std, NO length-standardization        │
(group of K)│    logπ_new = model(input_ids).logprobs   # the on-policy student forward (grad)      │
            │    log_r     = logπ_new - logπ_old         # log importance ratio (old = rollout-time) │
            │    pg   = -(advantages * exp(log_r))[resp_mask]                                        │
            │    kl   = (-log_r)[resp_mask]              # k1 estimator, NOT k3                       │
            │    L_drgrpo = (pg + beta_kl * kl).sum()                                                │
            │                                                                                        │
            │  Channel 2 (SDPO) — SAME minibatch, reuses the student forward where possible:         │
            │    error sites  ◀── reuse ingestion structural `tool_error` (research/07 §5)           │
            │       │            (turn.get("tool_error") is not None; single source of truth)        │
            │       ▼                                                                                 │
            │    HintGenerator.generate(ErrorContext)  ──▶ hint text  (research/07 §6, layered)      │
            │       │                                                                                 │
            │       ▼  data_collator splices hint at the error turn:                                 │
            │    ctx_teacher_input_ids  (hint system-msg + recovery turn, chat-template aligned)     │
            │    input_ids              (placeholder-of-equal-token-length so shapes match)          │
            │    sdpo_loss_mask         (1 on post-hint recovery tokens only)                         │
            │       │                                                                                 │
            │       ▼                                                                                 │
            │    student_logits = model(input_ids).logits                 # grad                      │
            │    with no_grad: teacher_logits = model(ctx_teacher_input_ids).logits   # stop-grad     │
            │    L_sdpo = generalized_jsd_loss(student, teacher,                                      │
            │                labels=sdpo_loss_mask, beta=jsd_beta,                                    │
            │                temperature=1.0, token_clip=0.05)           # masked to error turn       │
            │                                                                                        │
            │  total = L_drgrpo + beta_sdpo * L_sdpo      ──▶ .backward()  ──▶ Adam.step()            │
            └────────────────────────────────────────────────────────────────────────────────────────┘
```

Key flow facts:

- **Error-site detection is not re-invented.** The ingestion layer already sets
  `turn["tool_error"] = <error_kind>` (structural `is_error:true` flag first,
  string-tag fallback), and the collator's `_is_error_turn` keys on exactly that
  (`research/07` §5). The trainer **consumes** the collator's
  `ctx_teacher_input_ids` / `sdpo_loss_mask`; it does not detect errors itself.
- **HintGenerator is called at collation time**, not in the loss. Per
  `research/07` §6.1, the generator's only job is to produce the text spliced
  into the teacher context; the collator's `_build_hint_injected_trace` does the
  splice and the equal-length student alignment
  (`_build_aligned_student_for_sdpo`). The trainer sees finished tensors.
- **The teacher forward is on the live weights**, hint-conditioned, `no_grad`.
  It is *not* a separate model and *not* a re-rollout (`research/07` §1.3). One
  extra forward per SDPO minibatch.
- **The JSD is masked to the error turn** via `sdpo_loss_mask` (post-hint
  recovery tokens only), so SDPO supervises *exactly* the turn the hint targets,
  leaving the rest of the trajectory to channel 1.

---

## 3. Reconciling with Dr. GRPO specifics

`research/10` pins the algorithm. SDPO must coexist without perturbing any of it:

| Dr. GRPO property (`research/10` §2) | Where it lives | SDPO interaction |
|---|---|---|
| **No std-dev advantage normalization** | advantage estimator | **None.** SDPO never touches advantages. Keep `A = R - group_mean` (no `/std`). |
| **Length-standardization term removed** | PG reduction | **None.** SDPO is a separate head; do not re-introduce a `1/|y|` factor via SDPO's reduction either (use `batchmean` over masked error-turn tokens, which is SDPO's own normalization, independent of trajectory length). |
| **k1 KL = `−log r`** (NOT k3) | GRPO KL-to-ref term | **Distinct from SDPO.** The GRPO k1 KL regularizes the policy toward the *reference/old* policy on all response tokens. SDPO's JSD pulls the policy toward the *hint-conditioned self-teacher* on error-turn tokens. Two different targets, two different token sets, two different weights (`beta_kl` vs `beta_sdpo`). Never merge them. |
| **Single-epoch (a prompt is never trained twice)** | outer loop | **This is what makes SDPO clean.** The teacher forward happens on the *same minibatch being updated this step* — the student logits and the hint-conditioned teacher logits are both from the current weights on the current rollout, so the distilled KL is genuinely **on-policy** (SDPO's defining property). No stale-teacher / replay-buffer drift to reconcile. |
| **Adam, full-parameter, async rollouts** | optimizer / infra | **None.** SDPO adds gradient only through the student forward; Adam consumes the summed gradient transparently. Async/off-policy weight sync (PipelineRL-style) affects channel 1's `logπ_old`; SDPO's teacher is the *current* weights so it is unaffected. |

**The one thing to get right:** SDPO's JSD is **SEPARATE** from the GRPO
KL-to-ref. In the loss expression `total = L_drgrpo + beta_sdpo*L_sdpo`, the
`L_drgrpo` already *contains* its own `beta_kl * k1_kl`. Do not let `beta_sdpo`
masquerade as a KL coefficient or vice-versa; they are logged separately
(`loss/grpo_kl` vs `loss/sdpo_jsd`).

---

## 4. Implementation handles — `ComposerGRPOTrainer(GRPOTrainer)`

A focused subclass that (a) forces channel 1 into the Dr. GRPO regime and (b)
adds the SDPO head. This refines the existing `ComposerReplicationTrainer`; the
SDPO method is lifted almost verbatim from `composer_trainer.py:_compute_sdpo_loss`
(it is already correct), and the Dr. GRPO config is pinned via `GRPOConfig`.

### 4.1 The loss-override surface (version-robust)

The repo already overrides `_compute_loss(self, model, inputs)` — the internal
per-step loss hook TRL's `GRPOTrainer` exposes, and what this subclass keeps
using. Recent TRL wraps that in the HF `Trainer.compute_loss(self, model,
inputs, return_outputs=False, num_items_in_batch=None)`. To be robust to either
surface, override **`_compute_loss`** (present across the versions this repo
targets) and additionally provide a thin `compute_loss` shim that delegates, so
the subclass works whether TRL calls the internal or the public method:

```python
def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
    loss = self._compute_loss(model, inputs)            # our composed loss
    return (loss, None) if return_outputs else loss
```

If a future TRL drops `_compute_loss`, move the channel-1 call to
`super().compute_loss(model, inputs, return_outputs=True,
num_items_in_batch=num_items_in_batch)[0]` inside `_compute_loss` — the SDPO
add-on is unaffected.

### 4.2 The sketch (~70 LoC)

```python
# composer_replication/trainer/composer_grpo_trainer.py
from __future__ import annotations
from typing import Any
import logging, torch

try:
    from trl import GRPOTrainer, GRPOConfig            # noqa: F401
    _TRL = True
except ImportError:                                    # doc/test import without TRL
    GRPOTrainer = object; _TRL = False

from composer_replication.opsd import generalized_jsd_loss

logger = logging.getLogger(__name__)


def make_dr_grpo_config(**overrides: Any) -> "GRPOConfig":
    """Dr. GRPO regime (research/10 §2): no std-norm, no length-standardization,
    k1 KL, single-epoch, Adam. We pin what GRPOConfig exposes and assert the
    rest. TRL flag names drift across versions, so set defensively + log."""
    cfg_kwargs = dict(
        num_iterations=1,            # single-epoch: a prompt is never re-trained
        scale_rewards=False,         # << NO std-dev advantage normalization (Dr. GRPO)
        loss_type="dr_grpo",         # TRL's Dr. GRPO loss_type: drops length-standardization;
                                     #   if absent in your TRL, fall back to "grpo" and
                                     #   override the reduction (see assert below).
        optim="adamw_torch",         # Adam(W); Composer 2 uses Adam for RL
        beta=0.0,                    # GRPO KL-to-ref coeff; set >0 to enable the k1 term
    )
    cfg_kwargs.update(overrides)
    return GRPOConfig(**cfg_kwargs)


class ComposerGRPOTrainer(GRPOTrainer):  # type: ignore[misc,valid-type]
    """Dr. GRPO + SDPO (on-policy KL at error turns). SDPO is an ADDITIVE head;
    it never touches advantages or the GRPO-KL-to-ref term."""

    def __init__(self, *args: Any, beta_sdpo: float = 0.0, sdpo_jsd_beta: float = 0.5,
                 sdpo_temperature: float = 1.0, sdpo_token_clip: float | None = 0.05,
                 sdpo_warmup_steps: int = 0, beta_sdpo_max: float | None = None,
                 **kwargs: Any):
        if not _TRL:
            raise ImportError("ComposerGRPOTrainer requires TRL: pip install -e .[train]")
        super().__init__(*args, **kwargs)
        self.beta_sdpo = beta_sdpo
        self.beta_sdpo_max = beta_sdpo_max if beta_sdpo_max is not None else beta_sdpo
        self.sdpo_warmup_steps = sdpo_warmup_steps
        self.sdpo_jsd_beta = sdpo_jsd_beta
        self.sdpo_temperature = sdpo_temperature
        self.sdpo_token_clip = sdpo_token_clip
        # Dr. GRPO sanity pins (loud, not silent): if the TRL version ignored a
        # flag, surface it rather than train vanilla GRPO by accident.
        if getattr(self.args, "scale_rewards", True):
            logger.warning("Dr. GRPO requires scale_rewards=False (no std-norm); "
                           "GRPOConfig.scale_rewards=%s — advantages may be std-normalized.",
                           getattr(self.args, "scale_rewards", None))

    def _beta_sdpo_now(self) -> float:
        """Linear warmup so SDPO doesn't swamp the early policy gradient (§5)."""
        step = getattr(getattr(self, "state", None), "global_step", 0) or 0
        if self.sdpo_warmup_steps <= 0:
            return self.beta_sdpo
        frac = min(1.0, step / float(self.sdpo_warmup_steps))
        return self.beta_sdpo + frac * (self.beta_sdpo_max - self.beta_sdpo)

    def _compute_loss(self, model, inputs):
        drgrpo = super()._compute_loss(model, inputs)          # channel 1 (Dr. GRPO, k1 KL)
        sdpo   = self._compute_sdpo_loss(model, inputs)        # channel 2 (additive)
        beta   = self._beta_sdpo_now()
        total  = drgrpo + beta * sdpo
        if self.state.global_step % getattr(self.args, "logging_steps", 50) == 0:
            self.log({"loss/grpo": float(drgrpo.detach()),
                      "loss/sdpo_jsd": float(sdpo.detach()),
                      "loss/beta_sdpo": beta, "loss/total": float(total.detach())})
        return total

    def _compute_sdpo_loss(self, model, inputs):
        if (self._beta_sdpo_now() == 0.0
                or "ctx_teacher_input_ids" not in inputs
                or inputs["ctx_teacher_input_ids"].numel() == 0):
            return torch.zeros((), device=next(model.parameters()).device, requires_grad=True)
        student = model(input_ids=inputs["input_ids"]).logits           # grad
        with torch.no_grad():
            teacher = model(input_ids=inputs["ctx_teacher_input_ids"]).logits   # stop-grad
        if student.shape != teacher.shape:                              # collator alignment guard
            logger.warning("SDPO shape mismatch student=%s teacher=%s; skipping step.",
                           student.shape, teacher.shape)
            return torch.zeros((), device=student.device, requires_grad=True)
        return generalized_jsd_loss(student_logits=student, teacher_logits=teacher,
                                    labels=inputs.get("sdpo_loss_mask"),
                                    beta=self.sdpo_jsd_beta, temperature=self.sdpo_temperature,
                                    token_clip=self.sdpo_token_clip, reduction="batchmean")

    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        loss = self._compute_loss(model, inputs)
        return (loss, None) if return_outputs else loss
```

### 4.3 How error-turn batches reach the trainer

**Reuse `ComposerDataCollator` verbatim** — it already emits
`ctx_teacher_input_ids` + `sdpo_loss_mask` and (critically) the
**equal-length student** via `_build_aligned_student_for_sdpo` (the placeholder
trick that keeps `student_logits.shape == teacher_logits.shape` so the JSD gate
passes; the Gemini-W19 alias bug is already handled there). Wiring:

```python
gen = default_layered(judge_client=small_model).as_collator_hook()   # research/07 §6
collator = ComposerDataCollator(tokenizer=tok,
             config=CollatorConfig(hint_generator=gen, enable_sdpo=True,
                                   enable_replay_dpo=False))
trainer = ComposerGRPOTrainer(model=model, args=make_dr_grpo_config(...),
             train_dataset=ds, data_collator=collator,
             beta_sdpo=0.1, sdpo_warmup_steps=50, sdpo_token_clip=0.05,
             reward_funcs=[my_rlvr_reward])
```

> **GRPO-rollout vs collator note.** TRL's `GRPOTrainer` generates rollouts
> internally and forms its own `inputs` (prompt + completions + advantages). For
> SDPO the error sites come from the *rollout trajectory itself* (tool errors in
> the just-generated completions), so the SDPO tensors must be built **from the
> live rollout**, not from a static dataset. Two equivalent integration modes:
> (1) **post-rollout hook** — override `_generate_and_score_completions` (or the
> rollout collation step) to run the structural `tool_error` detector +
> `HintGenerator` + `ComposerDataCollator._build_sdpo_fields` on the generated
> completions and stash `ctx_teacher_input_ids`/`sdpo_loss_mask` into `inputs`;
> (2) **offline-trace mode** (what the smoke uses) — feed pre-ingested
> error-bearing traces through the collator as the dataset, exercising the exact
> loss path on CPU. Mode (2) is the test; mode (1) is production. The
> `_compute_sdpo_loss` body is identical for both — it only reads the two SDPO
> keys.

---

## 5. Weighting, scheduling, and guardrails

So SDPO informs without swamping the policy gradient:

1. **Scale.** Start `beta_sdpo = 0.1` (the library default `alpha_sdpo`), not the
   `1.0` the smoke uses (the smoke over-weights deliberately to *prove the path
   fires*). The Dr. GRPO PG loss is a `sum()` over response tokens; SDPO is a
   `batchmean` JSD over error-turn tokens — different magnitudes. **Normalize
   first:** log `loss/grpo` and `loss/sdpo_jsd` separately for the first ~50
   steps and pick `beta_sdpo` so `beta_sdpo·sdpo_jsd ≈ 0.1–0.3 × |grpo|` at
   steady state. Do **not** assume `0.1` is calibrated across reductions.
2. **Warmup.** Linear `beta_sdpo` warmup over `sdpo_warmup_steps` (50–200) via
   `_beta_sdpo_now()`. Early in training the policy is far from any sensible
   distribution; a strong distillation pull then fights exploration. Let Dr. GRPO
   establish a reward signal, then ramp SDPO in.
3. **Per-token JSD clip = 0.05** (`sdpo_token_clip`, the OPSD `--jsd_token_clip`
   default, `research/07` §1.1/§7). Prevents a few high-divergence **stylistic**
   tokens at the error turn from dominating the distillation gradient — exactly
   what it exists for.
4. **Mask discipline.** SDPO supervises **only** `sdpo_loss_mask` tokens
   (post-hint recovery). If the mask is all-ignore (empty-recovery error site,
   ~67% of real Claude traces under `strip_thinking`), the collator already drops
   the row (`data_collator.py` L308) — the channel silently no-ops rather than
   emitting a degenerate ~ln(2) signal.

**KL-explosion / teacher-student-drift guardrails:**

- **SDPO drift is bounded-by-construction.** Teacher = same weights + hint,
  stop-grad. A *wrong* hint produces a noisier target at one masked turn, not a
  corrupted reward (`research/07` §1.3). There is no replay buffer and no
  separate teacher to drift apart — single-epoch keeps teacher and student on the
  same weights.
- **Watch `loss/sdpo_jsd` for collapse-to-zero or blow-up.** A *good* hint should
  *raise* divergence at the hinted turn (it shifts mass toward the fix); a
  persistently ~0 JSD means the hint isn't moving the teacher (prune that hint
  source, `research/07` §7 item 7). A diverging JSD means the clip is too loose or
  `beta_sdpo` too high — cap `beta_sdpo` and/or lower `token_clip`.
- **Guard the GRPO k1 KL independently.** Dr. GRPO's own `beta` (KL-to-ref) is
  the explosion guard for the *policy*; keep it at its tuned value. SDPO's `beta_sdpo`
  must not be conflated with it (§3). If total loss NaNs, bisect by zeroing
  `beta_sdpo` — if it persists, the bug is in channel 1, not SDPO.
- **Shape-gate is a hard stop, logged.** If collator alignment regresses,
  `_compute_sdpo_loss` skips the step with a warning rather than training on
  aliased pad tokens (the silent-degenerate failure mode).

---

## 6. CPU-testable vs GPU-only, and the smoke plan

### What is CPU-testable
- **The whole SDPO loss path** — student forward + hint-conditioned teacher
  forward + masked JSD + `.backward()` + `Adam.step()` — on **Qwen2.5-0.5B** with
  1–2 error-bearing rollouts. This is *exactly* what
  `examples/sdpo_real_trace_train_smoke/run.py` already proves for the free
  `compose_loss` composer; the new test wraps it in the Dr. GRPO step.
- **The additive composition** `total = drgrpo + beta_sdpo·sdpo` and the warmup
  schedule (assert `beta_sdpo` ramps, assert `loss/sdpo_jsd>0` on ≥1 step, assert
  a watched param moves).
- **Dr. GRPO config pins** — assert `scale_rewards=False`, `num_iterations=1`,
  k1-KL path selected (unit-level, no GPU).

### What is GPU-only
- **Real TRL `GRPOTrainer` rollout generation** (vLLM/transformers generation at
  batch size + group size K) — too slow on CPU for a live step; this is the
  production "mode (1)" path in §4.3.
- **Async weight sync / off-policy control**, MoE router replay, multi-region
  infra (`research/10` §4) — all out of scope for the SDPO channel test.
- **Convergence / quality** (does SDPO actually improve error-recovery) — needs a
  real RL run.

### Minimal smoke plan (`examples/sdpo_drgrpo_step_smoke/run.py`)
Analogous to the existing SDPO smoke; gates:

1. Build a Dr. GRPO minibatch from 1–2 ingested **error-bearing** Qwen traces via
   `ComposerDataCollator` (reuse `_discover_error_sessions` + the layered
   `HintGenerator`); assert `sdpo_loss_mask` has ≥1 in-loss position.
2. Construct a **synthetic Dr. GRPO channel-1 loss** standing in for
   `super()._compute_loss` (advantages = `R - group_mean`, **no /std**; k1 KL
   `−log r`; `sum()` reduction; no length-standardization) so the test runs
   **without** spinning up TRL's full rollout machinery on CPU — mirrors how the
   existing smoke uses LM-CE as the GRPO stub. Optionally also run a real
   `GRPOTrainer._compute_loss` path under a `@pytest.mark.gpu` guard.
3. `total = drgrpo_stub + beta_sdpo · _compute_sdpo_loss(...)`; `.backward()`;
   `Adam.step()`.
4. **Gates (exit 0 = PASS):** (a) all losses finite across steps; (b)
   `loss/sdpo_jsd > 0` on ≥1 step (SDPO fired — shape-gate passed, hint
   contributed real signal); (c) a watched parameter moved; (d) `beta_sdpo`
   warmup increases monotonically; (e) zeroing `beta_sdpo` reproduces the pure
   Dr. GRPO stub loss bit-for-bit (proves SDPO is purely additive). Exit 2 = SKIP
   (no error-bearing sessions / no chat-template model), matching the existing
   smoke's contract.

This is ~$0, CPU, single-process, and closes the one unproven edge: **a live
Dr. GRPO update step with the SDPO channel on**, end-to-end on a real HF model.

---

## 7. Citations

- **In-repo (authoritative substrate):** `composer_replication/loss.py`
  (`compose_loss` 3-channel composer + `generalized_jsd_loss` call);
  `recipes/prime_rl/composer_loss.py` (PRIME-RL adapter; SDPO `NotImplementedError`
  at L257-268; parity-verified channel 1); `recipes/prime_rl/prime_rl_recipe.md`
  (LossInputs shape, log-probs-not-logits limitation);
  `trainer/composer_trainer.py` (`ComposerReplicationTrainer._compute_loss` and
  `_compute_sdpo_loss` — the existing, correct SDPO head);
  `trainer/data_collator.py` (`ctx_teacher_input_ids` + `sdpo_loss_mask` +
  `_build_aligned_student_for_sdpo` equal-length alignment; hint-AND-recovery
  gate L308); `examples/sdpo_real_trace_train_smoke/run.py` (the proven CPU
  forward+backward+step harness this design's smoke extends).
- **`research/10-composer2-techreport-mining.md`** — the Dr. GRPO target:
  length-standardization removed, no std-dev advantage normalization, **k1**
  (`−log r`) KL not k3, Adam, single-epoch (a prompt never trained twice).
  arXiv:2603.24477 §4.1.
- **`research/07-sdpo-hint-generator.md`** — `HintGenerator` Protocol + layered
  composite, error-site detection alignment with ingestion `tool_error`, OPSD
  `--jsd_token_clip` stabilizer, the "wrong hint is bounded-bad" property.
- **SDPO** — arXiv:2601.20802 (on-policy self-distillation; teacher = same model
  conditioned on feedback, student stop-grad-free / teacher stop-grad, per-token
  KL on the student trajectory). **OPSD** — arXiv:2601.18734 (privileged-info
  teacher, generalized-JSD, token clip).
- **TRL** — `huggingface/trl@main` `trl/trainer/grpo_trainer.py`
  (`GRPOTrainer` loss-override surface; confirmed via
  `mcp_exa_get_code_context_exa`). The `_compute_loss(self, model, inputs)`
  internal hook is what this repo already overrides; the public
  `compute_loss(self, model, inputs, return_outputs=False,
  num_items_in_batch=None)` HF-Trainer wrapper is shimmed in §4.1 for
  version-robustness. (A confirmatory DeepWiki lookup timed out; the §4.1 guard
  is written to work under either surface.)
```