GATES: How Consensus Gating Fixed the Broken Promise of Self-Distillation in Language Models

# ─────────────────────────────────────────────────────────────────────────────
# GATES: Gated Asymmetric Trajectory Self-Distillation
# Stein, Huang, Goldstein · University of Maryland, College Park
# arXiv:2602.20574v1 [cs.LG] · February 2026
#
# Full implementation: consensus gate, off-policy + on-policy distillation,
# KL regularization, document-leakage filtering, evaluation suite
# ─────────────────────────────────────────────────────────────────────────────

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np
import random
import re
from dataclasses import dataclass, field
from typing import List, Tuple, Optional, Dict
from collections import Counter


# ─── Section 1: Configuration ─────────────────────────────────────────────────

@dataclass
class GATESConfig:
    """
    Hyperparameters matching the GATES paper (Appendix A.3).
    """
    # Model
    model_name: str = "Qwen/Qwen3-4B-Base"
    embed_dim: int = 2048          # hidden dim (simplified for demo)

    # Training
    batch_size: int = 32            # questions per gradient step
    k_tutor: int = 8               # tutor rollouts per question
    k_student: int = 8             # student rollouts per question
    train_temp: float = 0.5        # sampling temperature for rollouts
    max_completion_tokens: int = 512
    num_epochs: int = 1            # single epoch over training split

    # Consensus gate
    consensus_min: int = 4         # minimum agreeing tutor answers (4/8)
    consensus_k: int = 8

    # Advantage clipping
    advantage_clip: float = 5.0    # a in clip(log πT - log πS, [-a, a])

    # Loss weights
    lambda_off: float = 1.0
    lambda_on: float = 0.1
    lambda_cons: float = 0.0       # disabled by default (§4.3)
    lambda_kl: float = 0.02

    # Optimizer
    lr: float = 1e-6
    weight_decay: float = 0.01
    grad_clip: float = 1.0

    # Evaluation
    eval_temp: float = 0.6         # maj@8 sampling temperature
    eval_samples: int = 8
    eval_max_tokens: int = 1024

    # Document leakage keywords (Appendix A.4)
    leakage_keywords: List[str] = field(default_factory=lambda: [
        "document", "passage", "text", "according to",
        "the article", "as stated", "mentioned in"
    ])

    random_seed: int = 42


# ─── Section 2: Data Structures ───────────────────────────────────────────────

@dataclass
class TrainingExample:
    """
    A document–question pair used during GATES training.
    No verified answer is stored; supervision is derived entirely from tutor consensus.
    """
    document: str
    question: str
    example_id: str = ""


@dataclass
class Rollout:
    """A single sampled completion from tutor or student."""
    text: str                          # full completion text
    token_ids: List[int]              # token ids of completion
    log_probs: torch.Tensor            # per-token log-probs under sampling model
    extracted_answer: Optional[str]   # last \boxed{...} expression or None
    has_boxed: bool = False           # whether a parsable boxed answer exists


@dataclass
class ConsensusResult:
    """Output of the two-level consensus gate for one training question."""
    question_gate: bool              # g_i: True if question passes consensus
    consensus_answer: Optional[str]  # a* — modal answer, or None
    eligible_tutor_idx: List[int]   # indices of rollouts with e_{i,j}=1
    agreement_count: int


# ─── Section 3: Answer Extraction ────────────────────────────────────────────

def extract_boxed_answer(text: str) -> Optional[str]:
    """
    Extract the LAST \\boxed{...} expression from a completion.
    Paper uses math-verify library; we provide a regex approximation.
    """
    pattern = r'\\boxed\{([^}]*)\}'
    matches = re.findall(pattern, text)
    return matches[-1].strip() if matches else None


def answers_equivalent(a: str, b: str) -> bool:
    """
    Check symbolic equivalence of two math answers.
    In production: use math-verify [OpenAI, 2024] for robust equivalence.
    Here: normalized string comparison as a simple approximation.
    """
    def normalize(s):
        s = s.strip().lower()
        s = re.sub(r'\s+', '', s)     # remove whitespace
        s = re.sub(r'[,\.]', '', s)    # remove punctuation
        return s
    return normalize(a) == normalize(b)


# ─── Section 4: Document Leakage Guardrails ───────────────────────────────────

def check_document_leakage(text: str, keywords: List[str]) -> bool:
    """
    Returns True if the text appears to explicitly reference a source document.
    Appendix A.4: we filter questions containing these references from distillation.
    """
    text_lower = text.lower()
    return any(kw in text_lower for kw in keywords)


# ─── Section 5: Two-Level Consensus Gate ─────────────────────────────────────

def consensus_gate(
        rollouts: List[Rollout],
        cfg: GATESConfig,
        leakage_check_enabled: bool = True
) -> ConsensusResult:
    """
    Two-level consensus gate (Section 3.2, Figure 2).

    Gate 1 (question-level): g_i = 1 if ≥ τ/k rollouts agree on same answer.
    Gate 2 (rollout-level): e_{i,j} = 1 if rollout j matches consensus answer a*
                             AND passes document-leakage guardrails.

    Returns ConsensusResult with question_gate, consensus_answer, eligible indices.
    """
    # Extract answers from all rollouts
    answers = []
    for r in rollouts:
        ans = r.extracted_answer
        answers.append(ans if (ans is not None and r.has_boxed) else None)

    # ── Gate 1: find modal answer ──────────────────────────────────────────────
    valid_answers = [a for a in answers if a is not None]
    if not valid_answers:
        return ConsensusResult(
            question_gate=False, consensus_answer=None,
            eligible_tutor_idx=[], agreement_count=0)

    # Cluster equivalent answers and count
    clusters: Dict[str, int] = {}
    representative: Dict[str, str] = {}
    for ans in valid_answers:
        placed = False
        for rep in list(clusters.keys()):
            if answers_equivalent(ans, rep):
                clusters[rep] += 1
                placed = True
                break
        if not placed:
            clusters[ans] = 1
            representative[ans] = ans

    modal_key = max(clusters, key=lambda k: clusters[k])
    agreement_count = clusters[modal_key]
    question_gate = agreement_count >= cfg.consensus_min
    consensus_answer = modal_key if question_gate else None

    if not question_gate:
        return ConsensusResult(
            question_gate=False, consensus_answer=None,
            eligible_tutor_idx=[], agreement_count=agreement_count)

    # ── Gate 2: rollout-level eligibility ─────────────────────────────────────
    eligible = []
    for j, (rollout, ans) in enumerate(zip(rollouts, answers)):
        if ans is None:
            continue
        if not answers_equivalent(ans, consensus_answer):
            continue
        if leakage_check_enabled and check_document_leakage(
                rollout.text, cfg.leakage_keywords):
            continue   # exclude document-dependent trajectories
        eligible.append(j)

    return ConsensusResult(
        question_gate=len(eligible) > 0,
        consensus_answer=consensus_answer,
        eligible_tutor_idx=eligible,
        agreement_count=agreement_count)


# ─── Section 6: Off-Policy Distillation Loss ─────────────────────────────────

def off_policy_loss(
        model: nn.Module,
        eligible_tutor_rollouts: List[Rollout],
        question_tokens: torch.Tensor,          # (seq_len,) prompt token ids
        device: torch.device
) -> Optional[torch.Tensor]:
    """
    Off-policy NLL distillation loss — Formula (1).

    For each eligible tutor trajectory, compute the student's log-probability
    of predicting each tutor token conditioned on the question prompt (no doc).
    Loss is the negative average per-token log-prob across all eligible rollouts.

    Returns None if no eligible rollouts to avoid zero-denominator.
    """
    if not eligible_tutor_rollouts:
        return None

    total_log_prob = torch.tensor(0.0, device=device, requires_grad=True)
    total_tokens = 0

    for rollout in eligible_tutor_rollouts:
        tutor_tokens = torch.tensor(rollout.token_ids, device=device)
        L = len(tutor_tokens)
        if L == 0:
            continue

        # Build full sequence: [question_tokens | tutor_tokens]
        full_seq = torch.cat([question_tokens, tutor_tokens]).unsqueeze(0)
        prompt_len = len(question_tokens)

        # Forward pass through student (no document in context)
        with torch.enable_grad():
            logits = model(full_seq).logits  # (1, total_len, vocab_size)

        # Completion region: [prompt_len-1 : prompt_len+L-1]
        completion_logits = logits[0, prompt_len - 1: prompt_len + L - 1]
        student_log_probs = F.log_softmax(completion_logits, dim=-1)

        # Per-token NLL on tutor tokens: ℓ_{i,j,t}(θ) = log π_θ(y^T_t | y^T_{
        token_log_probs = student_log_probs[
            torch.arange(L, device=device), tutor_tokens]

        total_log_prob = total_log_prob + token_log_probs.sum()
        total_tokens += L

    if total_tokens == 0:
        return None

    return -total_log_prob / total_tokens   # Formula (1)


# ─── Section 7: On-Policy Distillation Loss ───────────────────────────────────

def compute_per_token_advantage(
        tutor_model: nn.Module,
        student_rollout: Rollout,
        question_tokens: torch.Tensor,
        document_tokens: torch.Tensor,
        cfg: GATESConfig,
        device: torch.device
) -> torch.Tensor:
    """
    Compute per-token advantage A_t — Formula (2).

    A_t = clip(log π_T(y^S_t | y^S_{
    student_tokens = torch.tensor(student_rollout.token_ids, device=device)
    L = len(student_tokens)
    if L == 0:
        return torch.zeros(0, device=device)

    # Student log-probs (no document): π_θ(y^S_t | y^S_{
    q_seq = torch.cat([question_tokens, student_tokens]).unsqueeze(0)
    plen_q = len(question_tokens)
    with torch.no_grad():
        logits_s = tutor_model(q_seq).logits[0, plen_q - 1: plen_q + L - 1]
    log_pi_s = F.log_softmax(logits_s, dim=-1)[
        torch.arange(L, device=device), student_tokens]

    # Tutor log-probs (with document): π_T(y^S_t | y^S_{
    dq_seq = torch.cat([document_tokens, question_tokens, student_tokens]).unsqueeze(0)
    plen_dq = len(document_tokens) + len(question_tokens)
    with torch.no_grad():
        logits_t = tutor_model(dq_seq).logits[0, plen_dq - 1: plen_dq + L - 1]
    log_pi_t = F.log_softmax(logits_t, dim=-1)[
        torch.arange(L, device=device), student_tokens]

    # Advantage: clip(log π_T − log π_S, [−a, a]) — gradient does not flow through
    advantage = (log_pi_t - log_pi_s).clamp(-cfg.advantage_clip, cfg.advantage_clip)
    return advantage.detach()


def on_policy_loss(
        model: nn.Module,
        student_rollouts: List[Rollout],
        advantages: List[torch.Tensor],
        question_tokens: torch.Tensor,
        gate_passed: bool,
        device: torch.device
) -> Optional[torch.Tensor]:
    """
    On-policy distillation loss — Formula (3).

    ℒ_on(θ) = −(1/Σ g_i L^S_{i,j}) Σ g_i · A_{i,j,t} · ℓ^S_{i,j,t}(θ)

    Only applied when question passes the consensus gate (gate_passed=True).
    """
    if not gate_passed or not student_rollouts:
        return None

    total_loss = torch.tensor(0.0, device=device, requires_grad=True)
    total_tokens = 0
    plen = len(question_tokens)

    for rollout, adv in zip(student_rollouts, advantages):
        student_tokens = torch.tensor(rollout.token_ids, device=device)
        L = len(student_tokens)
        if L == 0 or len(adv) != L:
            continue

        full_seq = torch.cat([question_tokens, student_tokens]).unsqueeze(0)
        with torch.enable_grad():
            logits = model(full_seq).logits[0, plen - 1: plen + L - 1]

        log_pi = F.log_softmax(logits, dim=-1)[
            torch.arange(L, device=device), student_tokens]

        # Advantage-weighted NLL: Σ_t A_{i,j,t} · ℓ^S_{i,j,t}(θ)
        weighted = (adv.to(device) * log_pi).sum()
        total_loss = total_loss + weighted
        total_tokens += L

    if total_tokens == 0:
        return None

    return -total_loss / total_tokens   # Formula (3)


# ─── Section 8: KL Regularization ────────────────────────────────────────────

def kl_regularization_loss(
        model: nn.Module,
        ref_model: nn.Module,
        student_rollouts: List[Rollout],
        question_tokens: torch.Tensor,
        gate_passed: bool,
        device: torch.device
) -> Optional[torch.Tensor]:
    """
    KL regularization toward frozen reference policy — Formula (6/7).
    Prevents catastrophic drift from pretrained base during self-distillation.

    LKL(θ) ≈ β · E[KL(π_θ(·|q) || π_ref(·|q))] — averaged over student rollouts
    """
    if not gate_passed or not student_rollouts:
        return None

    total_kl = torch.tensor(0.0, device=device, requires_grad=True)
    total_tokens = 0
    plen = len(question_tokens)

    for rollout in student_rollouts:
        student_tokens = torch.tensor(rollout.token_ids, device=device)
        L = len(student_tokens)
        if L == 0:
            continue

        full_seq = torch.cat([question_tokens, student_tokens]).unsqueeze(0)
        with torch.enable_grad():
            logits_model = model(full_seq).logits[0, plen - 1: plen + L - 1]
        with torch.no_grad():
            logits_ref = ref_model(full_seq).logits[0, plen - 1: plen + L - 1]

        log_pi = F.log_softmax(logits_model, dim=-1)
        log_ref = F.log_softmax(logits_ref, dim=-1)

        # token-level KL: Σ_vocab π(v) (log π(v) - log π_ref(v))
        kl = F.kl_div(log_ref, log_pi.exp(), reduction='sum')
        total_kl = total_kl + kl
        total_tokens += L

    if total_tokens == 0:
        return None

    return total_kl / total_tokens


# ─── Section 9: GATES Trainer ─────────────────────────────────────────────────

class GATESTrainer:
    """
    Full GATES training loop combining:
        - Tutor sampling under document context
        - Two-level consensus gating
        - Off-policy trajectory distillation (primary signal)
        - On-policy advantage-weighted distillation (secondary signal)
        - KL regularization toward frozen reference
        - Document-leakage filtering
        - Joint gradient update over all active losses

    Implements Algorithm 1 (training procedure) from the GATES paper.
    """

    def __init__(self, model: nn.Module, ref_model: nn.Module,
                 cfg: GATESConfig, device: torch.device):
        self.model = model
        self.ref_model = ref_model
        self.ref_model.eval()
        for p in self.ref_model.parameters():
            p.requires_grad = False

        self.cfg = cfg
        self.device = device
        self.optimizer = optim.AdamW(
            model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)

    def _sample_rollout(self, prompt_tokens: torch.Tensor,
                        max_new: int, greedy: bool = False) -> Rollout:
        """
        Sample one completion from the model.
        In production: use vLLM or HF generate() with proper sampling parameters.
        This stub demonstrates the interface.
        """
        with torch.no_grad():
            generated = self.model.generate(
                prompt_tokens.unsqueeze(0).to(self.device),
                max_new_tokens=max_new,
                temperature=0.0 if greedy else self.cfg.train_temp,
                do_sample=not greedy,
                pad_token_id=0
            )
        completion_ids = generated[0, len(prompt_tokens):].tolist()
        text = ""  # would decode here in production
        extracted = extract_boxed_answer(text)
        return Rollout(
            text=text,
            token_ids=completion_ids,
            log_probs=torch.zeros(len(completion_ids)),
            extracted_answer=extracted,
            has_boxed=extracted is not None
        )

    def train_step(self, examples: List[TrainingExample],
                   tokenizer) -> dict:
        """
        One gradient step over a batch of document–question pairs.

        For each example:
            1. Sample k tutor rollouts with document in context
            2. Apply consensus gate
            3. If gate passes: sample k student rollouts, compute all losses
            4. Accumulate gradients; step optimizer after full batch
        """
        self.model.train()
        batch_losses = []
        batch_stats = {'gated': 0, 'skipped': 0, 'total': len(examples)}

        self.optimizer.zero_grad()

        for ex in examples:
            # ── Tokenize prompts (prompt tokens only; loss on completions) ─────
            tutor_prompt = (
                f"Document:\n{ex.document}\n\nQuestion:\n{ex.question}\n"
                "Solve step by step. Put FINAL answer in \\boxed{}.\n"
                "Do NOT mention the document. Solution:\n"
            )
            student_prompt = (
                f"Question:\n{ex.question}\n"
                "Solve step by step. Put FINAL answer in \\boxed{}.\nSolution:\n"
            )
            doc_q_tokens = torch.tensor(
                tokenizer.encode(tutor_prompt), device=self.device)
            q_tokens = torch.tensor(
                tokenizer.encode(student_prompt), device=self.device)
            doc_tokens = doc_q_tokens[:len(doc_q_tokens) - len(q_tokens)]

            # ── Step 1: Sample k tutor rollouts ───────────────────────────────
            tutor_rollouts = [
                self._sample_rollout(doc_q_tokens, self.cfg.max_completion_tokens)
                for _ in range(self.cfg.k_tutor)
            ]

            # ── Step 2: Consensus gate ─────────────────────────────────────────
            gate_result = consensus_gate(tutor_rollouts, self.cfg)

            if not gate_result.question_gate:
                batch_stats['skipped'] += 1
                continue   # zero loss for this question

            batch_stats['gated'] += 1
            eligible_rollouts = [tutor_rollouts[j]
                                 for j in gate_result.eligible_tutor_idx]

            # ── Step 3a: Off-policy distillation loss ─────────────────────────
            l_off = off_policy_loss(
                self.model, eligible_rollouts, q_tokens, self.device)

            # ── Step 3b: Sample k student rollouts ────────────────────────────
            student_rollouts = [
                self._sample_rollout(q_tokens, self.cfg.max_completion_tokens)
                for _ in range(self.cfg.k_student)
            ]

            # ── Step 3c: Per-token advantages for on-policy loss ──────────────
            advantages = [
                compute_per_token_advantage(
                    self.model, s, q_tokens, doc_tokens, self.cfg, self.device)
                for s in student_rollouts
            ]

            # ── Step 3d: On-policy distillation loss ──────────────────────────
            l_on = on_policy_loss(
                self.model, student_rollouts, advantages,
                q_tokens, gate_result.question_gate, self.device)

            # ── Step 3e: KL regularization ────────────────────────────────────
            l_kl = kl_regularization_loss(
                self.model, self.ref_model, student_rollouts,
                q_tokens, gate_result.question_gate, self.device)

            # ── Step 3f: Composite loss — Formula (7) ─────────────────────────
            loss = torch.tensor(0.0, device=self.device, requires_grad=True)
            if l_off is not None:
                loss = loss + self.cfg.lambda_off * l_off
            if l_on is not None:
                loss = loss + self.cfg.lambda_on * l_on
            if l_kl is not None:
                loss = loss + self.cfg.lambda_kl * l_kl

            if loss.requires_grad:
                loss.backward()
                batch_losses.append(loss.item())

        # ── Gradient update after full batch ──────────────────────────────────
        if batch_losses:
            nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.grad_clip)
            self.optimizer.step()

        return {
            'mean_loss': sum(batch_losses) / max(len(batch_losses), 1),
            'gated_fraction': batch_stats['gated'] / max(batch_stats['total'], 1),
            **batch_stats
        }

    def train(self, train_examples: List[TrainingExample],
              tokenizer, log_every: int = 10):
        """Full training loop — one epoch over the training split."""
        random.shuffle(train_examples)
        steps = range(0, len(train_examples), self.cfg.batch_size)
        for step_idx, start in enumerate(steps):
            batch = train_examples[start: start + self.cfg.batch_size]
            stats = self.train_step(batch, tokenizer)
            if step_idx % log_every == 0:
                print(
                    f"Step {step_idx:4d} | Loss: {stats['mean_loss']:6.4f} | "
                    f"Gated: {stats['gated_fraction']:.2%} of questions"
                )


# ─── Section 10: Evaluation Suite ────────────────────────────────────────────

def evaluate_student(
        model: nn.Module,
        examples: List[TrainingExample],
        oracle_answers: List[str],
        tokenizer,
        cfg: GATESConfig,
        mode: str = "maj8"
) -> dict:
    """
    Evaluate document-free student accuracy under three decoding strategies:
        - greedy:  single sample, temperature=0
        - maj8:    majority voting over 8 samples (paper primary metric)
        - pass8:   pass@8, True if any of 8 samples is correct

    oracle_answers: ground-truth answers used ONLY for evaluation (not training).
    """
    model.eval()
    correct_greedy = 0
    correct_maj = 0
    correct_pass = 0
    total = len(examples)

    with torch.no_grad():
        for ex, oracle in zip(examples, oracle_answers):
            student_prompt = (
                f"Question:\n{ex.question}\n"
                "Solve step by step. Put FINAL answer in \\boxed{}.\nSolution:\n"
            )
            q_tokens = torch.tensor(
                tokenizer.encode(student_prompt), device="cpu")

            # ── Greedy sample ─────────────────────────────────────────────────
            greedy_out = model.generate(
                q_tokens.unsqueeze(0),
                max_new_tokens=cfg.eval_max_tokens,
                temperature=0.0, do_sample=False, pad_token_id=0)
            greedy_text = tokenizer.decode(greedy_out[0, len(q_tokens):])
            greedy_ans = extract_boxed_answer(greedy_text)
            if greedy_ans and answers_equivalent(greedy_ans, oracle):
                correct_greedy += 1

            # ── 8 sampled completions for maj@8 and pass@8 ────────────────────
            sampled_answers = []
            for _ in range(cfg.eval_samples):
                out = model.generate(
                    q_tokens.unsqueeze(0),
                    max_new_tokens=cfg.eval_max_tokens,
                    temperature=cfg.eval_temp, do_sample=True, pad_token_id=0)
                text = tokenizer.decode(out[0, len(q_tokens):])
                ans = extract_boxed_answer(text)
                if ans:
                    sampled_answers.append(ans)

            # maj@8: majority vote
            if sampled_answers:
                # cluster answers by equivalence
                clusters: Dict[str, int] = {}
                for a in sampled_answers:
                    placed = False
                    for rep in clusters:
                        if answers_equivalent(a, rep):
                            clusters[rep] += 1; placed = True; break
                    if not placed:
                        clusters[a] = 1
                majority = max(clusters, key=clusters.__getitem__)
                if answers_equivalent(majority, oracle):
                    correct_maj += 1

                # pass@8: any correct
                if any(answers_equivalent(a, oracle) for a in sampled_answers):
                    correct_pass += 1

    return {
        'greedy_acc': correct_greedy / max(total, 1),
        'maj8_acc': correct_maj / max(total, 1),
        'pass8_acc': correct_pass / max(total, 1),
        'total': total,
    }


# ─── Section 11: Main Entry Point ────────────────────────────────────────────

if __name__ == "__main__":
    print("=" * 72)
    print("GATES: Gated Asymmetric Trajectory Self-Distillation")
    print("Stein, Huang, Goldstein · University of Maryland · arXiv:2602.20574v1")
    print("=" * 72)

    cfg = GATESConfig()
    torch.manual_seed(cfg.random_seed)
    random.seed(cfg.random_seed)

    print("\n[1] Configuration:")
    print(f"    Tutor rollouts:    k={cfg.k_tutor}")
    print(f"    Consensus min:     {cfg.consensus_min}/{cfg.consensus_k}")
    print(f"    Loss weights:      λ_off={cfg.lambda_off}, λ_on={cfg.lambda_on}, λ_KL={cfg.lambda_kl}")
    print(f"    Oracle reward:     λ_cons={cfg.lambda_cons} (disabled by default)")
    print(f"    Advantage clip:    ±{cfg.advantage_clip}")

    print("\n[2] Demonstrating consensus gate logic...")
    # Simulate 8 tutor rollouts: 6 agree on "7", 2 on "21"
    simulated_rollouts = []
    for i in range(6):
        simulated_rollouts.append(Rollout(
            text=f"42 = 2 × 3 × 7; largest prime is 7. \\boxed{{7}}",
            token_ids=[100, 101, 102],
            log_probs=torch.zeros(3),
            extracted_answer="7", has_boxed=True
        ))
    for i in range(2):
        simulated_rollouts.append(Rollout(
            text=f"42 / 2 = 21; largest factor is 21. \\boxed{{21}}",
            token_ids=[200, 201],
            log_probs=torch.zeros(2),
            extracted_answer="21", has_boxed=True
        ))

    result = consensus_gate(simulated_rollouts, cfg)
    print(f"    Question gate:     {result.question_gate}")
    print(f"    Consensus answer:  {result.consensus_answer}")
    print(f"    Agreement count:   {result.agreement_count}/{cfg.consensus_k}")
    print(f"    Eligible rollouts: {len(result.eligible_tutor_idx)} of {cfg.k_tutor}")

    print("\n[3] Demonstrating answer extraction...")
    test_completions = [
        "Step 1: factor 42 = 2 × 3 × 7. The primes are 2, 3, 7. \\boxed{7}",
        "Computing... final answer is \\boxed{42} no wait \\boxed{7}",
        "I'm not sure. Perhaps \\boxed{21}? No, \\boxed{7}.",
        "The answer is probably 7."  # no boxed
    ]
    for c in test_completions:
        ans = extract_boxed_answer(c)
        print(f"    '{c[:50]}...' → {ans}")

    print("\n[4] Demonstrating leakage filter...")
    clean = "42 = 2 × 3 × 7, so the largest prime factor is 7."
    leaky = "According to the document, 42 = 2 × 3 × 7, so the answer is 7."
    print(f"    Clean trajectory leak detected: {check_document_leakage(clean, cfg.leakage_keywords)}")
    print(f"    Leaky trajectory leak detected: {check_document_leakage(leaky, cfg.leakage_keywords)}")

    print("\n[5] Demonstrating answer equivalence...")
    pairs = [("7", "7"), ("7", " 7 "), ("7", "21"), ("1/2", "1/2")]
    for a, b in pairs:
        print(f"    '{a}' == '{b}': {answers_equivalent(a, b)}")

    print("\n" + "=" * 72)
    print("GATES implementation components:")
    print("  GATESConfig              — full hyperparameter config (matches Appendix A.3)")
    print("  TrainingExample/Rollout  — data structures for doc-question pairs")
    print("  extract_boxed_answer     — last \\boxed{} extraction from completions")
    print("  answers_equivalent       — symbolic answer comparison")
    print("  check_document_leakage   — guardrail for doc-dependent trajectories")
    print("  consensus_gate           — 2-level gate: question-level + rollout-level")
    print("  off_policy_loss          — NLL on eligible tutor trajectories (Formula 1)")
    print("  compute_per_token_advantage — clipped log πT − log πS (Formula 2)")
    print("  on_policy_loss           — advantage-weighted student NLL (Formula 3)")
    print("  kl_regularization_loss   — KL to frozen reference (Formula 6)")
    print("  GATESTrainer.train_step  — full per-batch training step (Algorithm 1)")
    print("  evaluate_student         — greedy / maj@8 / pass@8 accuracy")
    print("=" * 72)