tests/test_spec_ablations.py

"""Per-ablation behavioural spec tests (step 8).

One deterministic behavioural test per ablation mechanism of
the ablation spec §2, with expected values from the pinned
sources (SPG, Jaques 2017, Kirkpatrick 2017, Sun 2019, Hu 2021,
Yu 2020, Kim 2025) or derivations written in the docstrings. The
group-C trainable-set tests classify by reproduction: apply the
registry optimizer to all-ones gradients and read off which
parameters take a non-zero update. xfail(strict=True) marks canonical-vs-implemented
disagreements from the defect register or the step-8 findings list.

The minihack twin file carries the same mechanisms in its framework.
"""

from __future__ import annotations

import logging
import math
from pathlib import Path

import jax
import jax.numpy as jnp
import numpy as np
import optax
import pytest
import yaml
from tests.conftest import SEED

from experiments.rl_finetuning.ablations.losses import (
    LossContext,
    _ewc_penalty,
    make_loss_advantage_clip,
    make_loss_baseline,
    make_loss_bc_all,
    make_loss_bc_wins,
    make_loss_entropy_bonus,
    make_loss_ewc,
    make_loss_kl_penalty,
    make_loss_low_t,
    make_loss_normalized_adv,
    make_loss_t_curriculum,
    make_loss_t_curriculum_jit,
    make_loss_trust_region_kl,
)
from experiments.rl_finetuning.ablations.optimizers import (
    gradient_surgery,
    make_lora_params,
    make_optimizer_lora_only,
    merge_lora_into_base,
)
from experiments.rl_finetuning.ablations.registry import REGISTRY
from experiments.rl_finetuning.ablations.training import (
    _build_reward_model,
    _compute_advantages,
    _init_replay_buffer,
    _push_to_buffer,
    _reward_model_train_step,
)
from src.diffusion.schedules import SCHEDULE_MAP
from src.planners.model import build_model, init_params

V, H, OBS = 4, 4, 8
B = 4
TINY = {
    "D_MODEL": 16, "N_HEADS": 2, "N_LAYERS": 2, "D_FF": 16,
    "OBS_ENCODER_LAYERS": 1, "OBS_ENCODER_WIDTH": 16, "PLAN_HORIZON": H,
}
LINEAR = SCHEDULE_MAP["linear"]
COSINE = SCHEDULE_MAP["cosine"]


def _uniform_ctx(config=None, horizon=H, vocab=V):
    """LossContext with a uniform-logits stub model (params ignored)."""

    def apply_fn(params, obs, z, t, rng):
        return jnp.zeros((obs.shape[0], horizon, vocab))

    return LossContext(
        apply_fn=apply_fn, ref_params=None, schedule_fn=LINEAR[0],
        schedule_deriv_fn=LINEAR[1], num_actions=vocab,
        config=config or {},
    )


def _logit_ctx(ref_logits, config=None, horizon=H):
    """LossContext whose stub model broadcasts `params` as the logits.

    ref_params holds the reference distribution's logits, so
    KL(current || pretrained) has the closed form used in the tests.
    """

    def apply_fn(params, obs, z, t, rng):
        return jnp.broadcast_to(params, (obs.shape[0], horizon, params.shape[-1]))

    return LossContext(
        apply_fn=apply_fn, ref_params=jnp.asarray(ref_logits),
        schedule_fn=LINEAR[0], schedule_deriv_fn=LINEAR[1],
        num_actions=ref_logits.shape[-1], config=config or {},
    )


def _batch(b=B, horizon=H, obs_dim=OBS, vocab=V, key=1):
    k = jax.random.PRNGKey(key)
    acts = jax.random.randint(k, (b, horizon), 0, vocab)
    obs = jnp.zeros((b, obs_dim))
    return acts, obs, jnp.ones(b)


# ---------------------------------------------------------------------------
# baseline_rl and the advantage pipeline (spec-ablations §2 baseline row)
# ---------------------------------------------------------------------------


def test_compute_advantages_standard_branch_closed_form():
    """weight = clip(max(R,0)/(mean(max(R,0))+eps), 0.1, 5.0).

    Source: spec-ablations §2 baseline_rl effective params (SPG eq (5)
    positive branch with the pinned deviations). Derivation: returns
    [0,1,2,3] -> clipped mean 1.5 -> raw weights [0, 2/3, 4/3, 2] ->
    floor lifts the first to 0.1.
    """
    adv, mean, std = _compute_advantages(
        jnp.array([0.0, 1.0, 2.0, 3.0]), 0.1, 5.0, wins_only=False,
        win_thresh=0.5, use_running_stats=False, ema_decay=0.99,
        running_mean=jnp.zeros(()), running_std=jnp.ones(()),
    )
    assert np.allclose(np.asarray(adv), [0.1, 2 / 3, 4 / 3, 2.0], atol=1e-4)
    assert float(mean) == pytest.approx(1.5, abs=1e-6)
    assert float(std) == pytest.approx(math.sqrt(1.25), abs=1e-5)


def test_compute_advantages_running_stats_branch_closed_form():
    """running_stats: EMA of batch mean/std, adv = clip((w-mu)/sigma + 1,
    0.1, 5.0).

    Source: spec-ablations §2 running_stats row. Derivation with
    ema_decay d=0.5, prior mean 0 / std 1, batch [0,1,2,3]:
    new_mean = 0.5*0 + 0.5*1.5 = 0.75;
    new_std = 0.5*1 + 0.5*(sqrt(1.25)) = 1.0590;
    adv_i = clip((w_i - 0.75)/1.0590 + 1, 0.1, 5) =
    [0.2918, 1.2361, 2.1804, 3.1246].
    """
    adv, mean, std = _compute_advantages(
        jnp.array([0.0, 1.0, 2.0, 3.0]), 0.1, 5.0, wins_only=False,
        win_thresh=0.5, use_running_stats=True, ema_decay=0.5,
        running_mean=jnp.zeros(()), running_std=jnp.ones(()),
    )
    new_std = 0.5 * 1.0 + 0.5 * math.sqrt(1.25)
    expected = np.clip((np.array([0, 1, 2, 3.0]) - 0.75) / new_std + 1.0, 0.1, 5.0)
    assert np.allclose(np.asarray(adv), expected, atol=1e-4)
    assert float(mean) == pytest.approx(0.75, abs=1e-6)
    assert float(std) == pytest.approx(new_std, abs=1e-5)


def test_baseline_loss_is_linear_in_the_advantages():
    """The advantage weight is a per-sample multiplier on the ELBO
    (SPG eq (5) positive branch: A * L_ELBO), so scaling every
    advantage by 2 exactly doubles the loss under the same RNG."""
    ctx = _uniform_ctx()
    loss_fn = make_loss_baseline(ctx)
    acts, obs, valid = _batch()
    rng = jax.random.PRNGKey(SEED)
    adv = jnp.array([0.5, 1.0, 1.5, 2.0])
    l1 = float(loss_fn(None, acts, obs, valid, rng, adv))
    l2 = float(loss_fn(None, acts, obs, valid, rng, 2.0 * adv))
    assert l2 == pytest.approx(2.0 * l1, rel=1e-6)


# ---------------------------------------------------------------------------
# bc_wins (defect §8.5)
# ---------------------------------------------------------------------------


def test_bc_wins_averages_uniformly_over_winning_windows():
    """Canonical bc_wins ('Uniform ELBO on win windows', win = return >
    win_threshold, spec-ablations §2; was defect §8.5): a batch with no
    winning window carries no training signal (loss exactly 0), and an
    all-winning batch reduces to the plain uniform ELBO.

    The win masks are produced by the pipeline's own
    _compute_advantages(wins_only=True), exactly as in the training loop.
    """
    ctx = _uniform_ctx()
    loss_fn = make_loss_bc_wins(ctx)
    acts, obs, valid = _batch()
    rng = jax.random.PRNGKey(SEED)

    def mask(returns):
        m, _, _ = _compute_advantages(
            jnp.array(returns), 0.1, 5.0, wins_only=True, win_thresh=0.5,
            use_running_stats=False, ema_decay=0.99,
            running_mean=jnp.zeros(()), running_std=jnp.ones(()),
        )
        return m

    assert float(loss_fn(None, acts, obs, valid, rng, mask([0.0, 0.1, 0.2, 0.3]))) == 0.0
    all_wins = float(loss_fn(None, acts, obs, valid, rng, mask([1.0, 2.0, 3.0, 4.0])))
    uniform = float(make_loss_baseline(ctx)(None, acts, obs, valid, rng, jnp.ones(B)))
    assert all_wins == pytest.approx(uniform, abs=0.0)


# ---------------------------------------------------------------------------
# bc_all -- the unweighted-rollout arm
# ---------------------------------------------------------------------------


def test_bc_all_ignores_the_advantages_and_averages_over_the_whole_batch():
    """bc_all separates "trained on self-generated rollouts" from
    "weighted those rollouts by return": it must discard whatever weight
    vector the pipeline hands it and reduce to the plain uniform ELBO.

    Three weight vectors a return-weighted arm would score differently --
    all-ones, a spread of real advantages, and the explicit None the
    factory forwards -- must all give the identical loss, and that loss
    must equal the baseline's on uniform weights. Unlike bc_wins it keeps
    the losing windows, so an all-zero weight vector still carries signal.
    The minihack twin asserts the same four equalities.
    """
    ctx = _uniform_ctx()
    loss_fn = make_loss_bc_all(ctx)
    acts, obs, valid = _batch()
    rng = jax.random.PRNGKey(SEED)

    uniform = float(make_loss_baseline(ctx)(None, acts, obs, valid, rng, jnp.ones(B)))
    for adv in (jnp.ones(B), jnp.array([10.0, 0.0, 1.0, 1.1]), None, jnp.zeros(B)):
        got = float(loss_fn(None, acts, obs, valid, rng, adv))
        assert got == pytest.approx(uniform, abs=0.0), adv
    assert uniform != 0.0


# ---------------------------------------------------------------------------
# advantage_clip / normalized_adv (spec-ablations §2)
# ---------------------------------------------------------------------------


def test_advantage_clip_clips_the_weight_to_the_documented_band():
    """advantage_clip clips the return-weight itself to [1-eps, 1+eps]
    (project-specific; PPO's ratio clip rejected as source per
    the source index). With eps=0.2 the ablation loss equals the baseline
    loss on manually clipped advantages, under the same RNG."""
    ctx = _uniform_ctx({"ADV_CLIP_EPS": 0.2})
    acts, obs, valid = _batch()
    rng = jax.random.PRNGKey(SEED)
    adv = jnp.array([10.0, 0.0, 1.0, 1.1])
    got = float(make_loss_advantage_clip(ctx)(None, acts, obs, valid, rng, adv))
    want = float(
        make_loss_baseline(ctx)(None, acts, obs, valid, rng, jnp.clip(adv, 0.8, 1.2))
    )
    assert got == pytest.approx(want, abs=0.0)


def test_normalized_adv_standardises_over_the_batch():
    """normalized_adv applies (A - mean)/(std + 1e-8) over the batch
    (What Matters C67; spec-ablations §2). Equals the baseline loss on
    manually standardised advantages under the same RNG."""
    ctx = _uniform_ctx()
    acts, obs, valid = _batch()
    rng = jax.random.PRNGKey(SEED)
    adv = jnp.array([10.0, 0.0, 1.0, 1.1])
    norm = (adv - adv.mean()) / (adv.std() + 1e-8)
    got = float(make_loss_normalized_adv(ctx)(None, acts, obs, valid, rng, adv))
    want = float(make_loss_baseline(ctx)(None, acts, obs, valid, rng, norm))
    assert got == pytest.approx(want, abs=0.0)


# ---------------------------------------------------------------------------
# kl_penalty / trust_region_kl (Jaques 2017 family)
# ---------------------------------------------------------------------------

_P_LOGITS = jnp.log(jnp.array([0.7, 0.1, 0.1, 0.1]))
_Q_LOGITS = jnp.zeros(4)  # uniform reference
# KL(p || uniform) = sum p ln(4p) = 0.7 ln 2.8 + 0.3 ln 0.4
_KL_PQ = 0.7 * math.log(2.8) + 0.3 * math.log(0.4)
_P1_LOGITS = jnp.log(jnp.array([0.55, 0.15, 0.15, 0.15]))
_KL_P1Q = 0.55 * math.log(2.2) + 0.45 * math.log(0.6)

_KL_H = 512  # with H=512, P(a row draws no masked position) = E[alpha^H]
# ~ 1/513 per row under t~U(eps,1) linear, so all 4 rows almost surely
# contribute the constant per-position KL and the masked average is exact.


def test_kl_penalty_adds_coef_times_the_closed_form_kl():
    """kl_penalty adds kl_coef * KL(current || pretrained) on masked
    positions (Jaques 2017 eqs (2)-(4); spec-ablations §2 kl_penalty).

    Derivation: constant per-position logits make the masked-position
    KL equal KL(p||q) = 0.7 ln 2.8 + 0.3 ln 0.4 = 0.445846 in every row
    with at least one mask. Loss difference between kl_coef 0.3 and 0.1
    under identical RNG isolates 0.2 * KL.
    """
    acts, _, valid = _batch(horizon=_KL_H)
    obs = jnp.zeros((B, OBS))
    rng = jax.random.PRNGKey(SEED)
    losses = {}
    for coef in (0.1, 0.3):
        ctx = _logit_ctx(_Q_LOGITS, {"KL_COEF": coef}, horizon=_KL_H)
        losses[coef] = float(
            make_loss_kl_penalty(ctx)(_P_LOGITS, acts, obs, valid, rng, jnp.ones(B))
        )
    got_kl = (losses[0.3] - losses[0.1]) / 0.2
    assert got_kl == pytest.approx(_KL_PQ, rel=1e-3)


def test_kl_penalty_is_zero_when_current_equals_pretrained():
    """KL(p||p) = 0: with params == ref the penalty vanishes, so the
    kl_penalty loss equals its own RL term. The penalty shares the
    ELBO's forward pass, so the RL term is the plain core loss at the
    same RNG."""
    from experiments.rl_finetuning.ablations.losses import _core_loss

    ctx = _logit_ctx(_Q_LOGITS, {"KL_COEF": 0.1}, horizon=_KL_H)
    acts, _, valid = _batch(horizon=_KL_H)
    obs = jnp.zeros((B, OBS))
    rng = jax.random.PRNGKey(SEED)
    got = float(
        make_loss_kl_penalty(ctx)(_Q_LOGITS, acts, obs, valid, rng, jnp.ones(B))
    )
    want = float(_core_loss(ctx, _Q_LOGITS, rng, acts, obs, valid, jnp.ones(B)))
    assert got == pytest.approx(want, abs=1e-6)


def test_penalty_terms_share_the_elbo_forward_pass():
    """The KL and entropy penalties are evaluated at the ELBO's own
    noised sample, not at an independently drawn t (PARITY
    'Ablation-suite mechanics'; the minihack twin shares the pass).

    Counting model calls is the direct evidence: kl_penalty runs the
    current model once (shared with the ELBO) plus the frozen
    reference once, and entropy_bonus runs it once in total. An
    independent draw would add a second current-model forward.
    """
    acts, _, valid = _batch(horizon=_KL_H)
    obs = jnp.zeros((B, OBS))
    rng = jax.random.PRNGKey(SEED)
    seen: list[jax.Array] = []

    def counting_ctx(config):
        def apply_fn(params, obs_, z, t, key):
            seen.append(t)
            return jnp.broadcast_to(
                params, (obs_.shape[0], _KL_H, params.shape[-1])
            )

        return LossContext(
            apply_fn=apply_fn,
            ref_params=jnp.asarray(_Q_LOGITS),
            schedule_fn=LINEAR[0],
            schedule_deriv_fn=LINEAR[1],
            num_actions=_Q_LOGITS.shape[-1],
            config=config,
        )

    make_loss_kl_penalty(counting_ctx({"KL_COEF": 0.1}))(
        _P_LOGITS, acts, obs, valid, rng, jnp.ones(B)
    )
    assert len(seen) == 2, "kl_penalty must run current + reference once each"
    assert jnp.allclose(seen[0], seen[1]), "both forwards must use the ELBO's t"

    seen.clear()
    make_loss_entropy_bonus(counting_ctx({"ENTROPY_COEF": 0.01}))(
        _P_LOGITS, acts, obs, valid, rng, jnp.ones(B)
    )
    assert len(seen) == 1, "entropy_bonus must reuse the ELBO's logits"


def test_trust_region_barrier_is_zero_below_and_quadratic_above():
    """trust_region_kl adds a quadratic barrier c*max(KL-delta,0)^2
    (spec-ablations §2 trust_region row; delta=0.05).

    Below the threshold (params == ref, KL=0) the barrier is exactly 0.
    Above it, the barrier for two KL levels K0=0.445846 and K1=0.203781
    (closed forms as in the KL test) satisfies the quadratic ratio
    ((K0-delta)/(K1-delta))^2 = (0.395846/0.153781)^2 = 6.6262 - this
    pins the barrier's form without pinning the project-specific c.
    """
    from experiments.rl_finetuning.ablations.losses import _core_loss

    acts, _, valid = _batch(horizon=_KL_H)
    obs = jnp.zeros((B, OBS))
    rng = jax.random.PRNGKey(SEED)

    def barrier(cur_logits):
        ctx = _logit_ctx(_Q_LOGITS, {"TRUST_REGION_KL": 0.05}, horizon=_KL_H)
        total = float(
            make_loss_trust_region_kl(ctx)(cur_logits, acts, obs, valid, rng, jnp.ones(B))
        )
        rl = float(_core_loss(ctx, cur_logits, rng, acts, obs, valid, jnp.ones(B)))
        return total - rl

    assert barrier(_Q_LOGITS) == pytest.approx(0.0, abs=1e-6)
    b0, b1 = barrier(_P_LOGITS), barrier(_P1_LOGITS)
    assert b0 > 0 and b1 > 0
    want_ratio = ((_KL_PQ - 0.05) / (_KL_P1Q - 0.05)) ** 2
    assert b0 / b1 == pytest.approx(want_ratio, rel=2e-2)


# ---------------------------------------------------------------------------
# ewc (Kirkpatrick 2017 eq (3), lambda-reparameterised per the source index)
# ---------------------------------------------------------------------------


def test_ewc_penalty_closed_form_and_factory_scaling():
    """EWC adds lambda * sum_i F_i (theta_i - theta*_i)^2.

    Source: Kirkpatrick 2017 eq (3); the repo folds the paper's 1/2
    into lambda (documented reparameterisation, spec-ablations §2).
    Derivation: F={a:[1,2]}, theta={a:[3,5]}, theta*={a:[1,1]} ->
    penalty = 1*(2^2) + 2*(4^2) = 36; with ewc_lambda=100 the factory
    loss exceeds the same-RNG core loss by exactly 3600.
    """
    fisher = {"a": jnp.array([1.0, 2.0])}
    theta = {"a": jnp.array([3.0, 5.0])}
    ref = {"a": jnp.array([1.0, 1.0])}
    assert float(_ewc_penalty(fisher, theta, ref)) == pytest.approx(36.0)

    from experiments.rl_finetuning.ablations.losses import _core_loss

    def apply_fn(params, obs, z, t, rng):
        return jnp.zeros((obs.shape[0], H, V))

    ctx = LossContext(
        apply_fn=apply_fn, ref_params=ref, schedule_fn=LINEAR[0],
        schedule_deriv_fn=LINEAR[1], num_actions=V,
        config={"EWC_LAMBDA": 100.0},
    )
    acts, obs, valid = _batch()
    rng = jax.random.PRNGKey(SEED)
    got = float(make_loss_ewc(ctx, fisher)(theta, acts, obs, valid, rng, jnp.ones(B)))
    rl = float(_core_loss(ctx, theta, rng, acts, obs, valid, jnp.ones(B)))
    assert got - rl == pytest.approx(3600.0, rel=1e-5)


# ---------------------------------------------------------------------------
# entropy_bonus (standard tier, cf. Mnih 2016)
# ---------------------------------------------------------------------------


def test_entropy_bonus_subtracts_coef_times_the_closed_form_entropy():
    """entropy_bonus subtracts entropy_coef * H(p_theta) on masked
    positions (spec-ablations §2). Derivation: constant per-position
    p=[0.7,0.1,0.1,0.1] gives H = -(0.7 ln 0.7 + 0.3 ln 0.1) = 0.940448
    (the masked average is globally normalised, so it is exact whenever
    any position is masked). The coefficient difference 0.03-0.01
    isolates -0.02 * H.
    """
    acts, obs, valid = _batch()
    rng = jax.random.PRNGKey(SEED)
    losses = {}
    for coef in (0.01, 0.03):
        ctx = _logit_ctx(_Q_LOGITS, {"ENTROPY_COEF": coef})
        losses[coef] = float(
            make_loss_entropy_bonus(ctx)(_P_LOGITS, acts, obs, valid, rng, jnp.ones(B))
        )
    entropy = -(0.7 * math.log(0.7) + 0.3 * math.log(0.1))
    assert (losses[0.01] - losses[0.03]) / 0.02 == pytest.approx(entropy, rel=1e-4)


# ---------------------------------------------------------------------------
# low_t / t_curriculum (Kim 2025; spec-ablations §2)
# ---------------------------------------------------------------------------


def _recording_ctx(config, records):
    def apply_fn(params, obs, z, t, rng):
        records.append(np.asarray(t))
        return jnp.zeros((obs.shape[0], H, V))

    return LossContext(
        apply_fn=apply_fn, ref_params=None, schedule_fn=LINEAR[0],
        schedule_deriv_fn=LINEAR[1], num_actions=V, config=config,
    )


def test_low_t_restricts_sampling_to_the_low_noise_regime():
    """low_t trains only on t in [eps, t_max_low=0.2]
    (spec-ablations §2 low_t row)."""
    records: list[np.ndarray] = []
    ctx = _recording_ctx({"T_MAX_LOW": 0.2}, records)
    acts, obs, valid = _batch(b=64)
    make_loss_low_t(ctx)(None, acts, obs, jnp.ones(64), jax.random.PRNGKey(0), jnp.ones(64))
    t = np.concatenate(records)
    assert t.min() >= 1e-5 - 1e-9 and t.max() <= 0.2 + 1e-6


@pytest.mark.parametrize("variant", ["list", "jit"])
def test_t_curriculum_anneals_high_noise_to_low_noise(variant):
    """t_curriculum anneals the t window from [0.8, 1.0] to
    [eps, 0.2] linearly over 200 iterations, high-noise (easy) first
    (Kim 2025, simplified linear anneal per the source index; params
    t_start=0.8, t_end=0.2, steps=200, spec-ablations §1.6).

    Expected windows: iter 0 -> [0.8, 1.0]; iter 100 (frac 0.5) ->
    [0.4, 0.6]; iter >= 200 -> [eps, 0.2].
    """
    config = {"T_CURRICULUM_START": 0.8, "T_CURRICULUM_END": 0.2,
              "T_CURRICULUM_STEPS": 200}
    acts, obs, _ = _batch(b=64)
    for it, (lo, hi) in [(0, (0.8, 1.0)), (100, (0.4, 0.6)), (200, (1e-5, 0.2))]:
        records: list[np.ndarray] = []
        ctx = _recording_ctx(config, records)
        if variant == "list":
            fn = make_loss_t_curriculum(ctx, [it])
            fn(None, acts, obs, jnp.ones(64), jax.random.PRNGKey(0), jnp.ones(64))
        else:
            fn = make_loss_t_curriculum_jit(ctx)
            fn(None, acts, obs, jnp.ones(64), jax.random.PRNGKey(0), jnp.ones(64),
               jnp.array(it))
        t = np.concatenate(records)
        assert t.min() >= lo - 1e-6, (it, t.min(), lo)
        assert t.max() <= hi + 1e-6, (it, t.max(), hi)


# ---------------------------------------------------------------------------
# reward_filtering / action_diversity (spec-ablations §2, step-9 seams)
# ---------------------------------------------------------------------------


def test_reward_filter_keeps_strictly_above_the_percentile():
    """reward_filtering keeps windows with return STRICTLY above the
    batch percentile (spec-ablations §2, step-9 amendment: same
    boundary in both repos).

    Derivation: returns 1..8, 75th percentile (linear interpolation) =
    6.25 -> keep {7, 8}. All-equal returns: percentile == value, so
    strict > keeps nothing.
    """
    from experiments.rl_finetuning.ablations.training import reward_filter_mask

    keep = np.asarray(reward_filter_mask(jnp.arange(1.0, 9.0), 75))
    assert keep.tolist() == [False] * 6 + [True, True]
    assert not np.asarray(reward_filter_mask(jnp.full(5, 2.0), 75)).any()


def test_reward_filter_percentile_ignores_invalid_windows():
    """The threshold comes from the valid windows only; an invalid
    window has no return and must not move it (cross-implementation note 'Ablation-suite
    mechanics', filter row handling - the minihack twin removes the
    rows outright, which is the same rule).

    Derivation: valid returns 1..8 (75th percentile 6.25 -> keep
    {7, 8}); four invalid rows carrying 100.0 would raise the
    percentile to 82.75 and keep nothing.
    """
    from experiments.rl_finetuning.ablations.training import reward_filter_mask

    returns = jnp.concatenate([jnp.arange(1.0, 9.0), jnp.full(4, 100.0)])
    valid = jnp.array([True] * 8 + [False] * 4)

    keep = np.asarray(reward_filter_mask(returns, 75, valid))
    assert keep.tolist() == [False] * 6 + [True, True] + [False] * 4
    # Without the validity mask the same batch keeps nothing.
    assert not np.asarray(reward_filter_mask(returns, 75))[:8].any()


def test_an_empty_reward_filter_batch_warns(caplog):
    """A kept count of zero is logged; a non-empty one is silent.

    Tied returns leave nothing strictly above the percentile, so the
    keep-mask is all-False, the ELBO and its gradient are exactly zero, and
    the iteration is indistinguishable from a normal one in every logged
    quantity. The boundary stays strict, per spec-ablations §2; only the
    degenerate case becomes visible.
    """
    from experiments.rl_finetuning.ablations.training import (
        reward_filter_mask,
        warn_if_reward_filter_kept_nothing,
    )

    valid = jnp.ones(4, dtype=bool)

    with caplog.at_level(logging.WARNING):
        tied = reward_filter_mask(jnp.full(4, 2.0), 75, valid)
        warn_if_reward_filter_kept_nothing(tied, 75)
    assert "kept 0 of 4 windows" in caplog.text

    caplog.clear()
    with caplog.at_level(logging.WARNING):
        spread = reward_filter_mask(jnp.arange(1.0, 5.0), 75, valid)
        warn_if_reward_filter_kept_nothing(spread, 75)
    assert caplog.text == ""


def test_the_reward_filtering_branch_is_wired_to_the_warning():
    """The filter's only call site emits the warning.

    Source-anchored: the call sits inside the jitted scan body, which no
    unit test can reach without running a rollout. Without this the helper
    could be correct and never invoked -- the shape of the original defect.
    """
    import inspect

    from experiments.rl_finetuning.ablations import training

    src = inspect.getsource(training)
    filter_call = "mask = reward_filter_mask(flat_returns, reward_filter_pct, mask)"
    assert filter_call in src
    assert (
        f"{filter_call}\n                "
        "warn_if_reward_filter_kept_nothing(mask, reward_filter_pct)"
    ) in src


def test_action_diversity_discards_degenerate_plans():
    """action_diversity keeps only windows with more than one distinct
    action (spec-ablations §2)."""
    from experiments.rl_finetuning.ablations.training import action_diversity_mask

    acts = jnp.array([[1, 1, 1, 1], [1, 2, 1, 1], [0, 0, 0, 0]])
    assert np.asarray(action_diversity_mask(acts)).tolist() == [False, True, False]


# ---------------------------------------------------------------------------
# llrd (Sun 2019: eta_{k-1} = xi * eta_k, top-down from the head)
# ---------------------------------------------------------------------------


def test_llrd_learning_rates_decay_geometrically_from_the_head():
    """LLRD gives the head base_lr and each layer at depth d from the
    top base_lr * decay^d; the observation encoder sits below the
    lowest block (Sun 2019; spec-ablations §2 llrd row, decay 0.9).

    Method: one AdamW step on all-ones gradients with weight decay 0
    and the norm clip disabled; the first-step AdamW update magnitude
    is lr * ghat/(sqrt(vhat)+eps) = lr/(1+1e-5) uniformly, so update
    ratios equal LR ratios. With N_LAYERS=2:
    head(Dense_4)=base; block_1=base*0.9; block_0=base*0.81;
    everything else 0.9^3.
    """
    model = build_model(TINY, V)
    params = init_params(model, jax.random.PRNGKey(SEED), OBS, H)
    config = {**TINY, "LR": 1e-3, "LLRD_DECAY": 0.9, "N_LAYERS": 2,
              "WEIGHT_DECAY": 0.0, "MAX_GRAD_NORM": 1e9}
    tx = REGISTRY["llrd"].optimizer_factory(config, params)
    state = tx.init(params)
    updates, _ = tx.update(jax.tree.map(jnp.ones_like, params), state, params)
    flat = jax.tree_util.tree_flatten_with_path(updates)[0]

    def group_of(path_str: str) -> str:
        if "TransformerBlock_0" in path_str:
            return "block_0"
        if "TransformerBlock_1" in path_str:
            return "block_1"
        if "Dense_4" in path_str:
            return "head"
        return "obs_enc"

    mags: dict[str, set[float]] = {}
    for path, leaf in flat:
        p = "/".join(str(k.key) for k in path)
        mags.setdefault(group_of(p), set()).add(round(float(jnp.abs(leaf).max()), 10))

    base = max(mags["head"])
    expected = {"head": 1.0, "block_1": 0.9, "block_0": 0.81, "obs_enc": 0.9**3}
    for group, rel in expected.items():
        got = max(mags[group])
        assert got / base == pytest.approx(rel, rel=1e-4), (group, got / base)


# ---------------------------------------------------------------------------
# lora (Hu 2021 eq (3): h = W0 x + (alpha/r) B A x, B zero-initialised)
# ---------------------------------------------------------------------------


def test_lora_delta_is_zero_at_init_low_rank_and_scaled_by_alpha_over_r():
    """LoRA per Hu 2021 eq (3)/§4: B=0 at init (delta exactly zero),
    the delta is (alpha/r)*A@B with rank <= r, on attention kernels.

    The delta expectation is recomputed with NumPy from the A/B factors,
    independently of merge_lora_into_base.
    """
    rank, alpha = 8, 16.0
    model = build_model(TINY, V)
    params = init_params(model, jax.random.PRNGKey(SEED), OBS, H)
    lora = make_lora_params(params, rank, jax.random.PRNGKey(7))
    assert lora, "no attention kernels matched"
    for ab in lora.values():
        assert np.allclose(np.asarray(ab["B"]), 0.0)

    merged0 = merge_lora_into_base(params, lora, alpha, rank)
    for (pa, la), (_pb, lb) in zip(
        jax.tree_util.tree_flatten_with_path(params)[0],
        jax.tree_util.tree_flatten_with_path(merged0)[0],
        strict=True,
    ):
        assert np.array_equal(np.asarray(la), np.asarray(lb)), pa

    perturbed = {
        k: {"A": ab["A"], "B": jnp.ones_like(ab["B"])} for k, ab in lora.items()
    }
    merged = merge_lora_into_base(params, perturbed, alpha, rank)
    flat_base = {
        "/".join(str(k.key) for k in path): leaf
        for path, leaf in jax.tree_util.tree_flatten_with_path(params)[0]
    }
    flat_merged = {
        "/".join(str(k.key) for k in path): leaf
        for path, leaf in jax.tree_util.tree_flatten_with_path(merged)[0]
    }
    for path_str, ab in perturbed.items():
        a = np.asarray(ab["A"])
        bmat = np.asarray(ab["B"])
        want = (alpha / rank) * (a @ bmat)
        got = (
            np.asarray(flat_merged[path_str]) - np.asarray(flat_base[path_str])
        ).reshape(want.shape)
        assert np.allclose(got, want, atol=1e-5), path_str
        assert np.linalg.matrix_rank(got) <= rank


# ---------------------------------------------------------------------------
# gradient_surgery (Yu 2020 Alg 1, one-sided per the source index)
# ---------------------------------------------------------------------------


def test_pcgrad_projection_closed_form_and_one_sidedness():
    """PCGrad: if g_rl . g_bc < 0, g_rl <- g_rl - (g_rl.g_bc/|g_bc|^2) g_bc;
    otherwise unchanged (Yu 2020 Alg 1; one-sided variant per the source index).

    Derivation: g_rl=[1,0], g_bc=[-1,1]: dot=-1, |g_bc|^2=2 ->
    projected = [1,0] - (-1/2)[-1,1] = [0.5, 0.5], orthogonal to g_bc.
    Non-conflicting g_bc=[1,1] leaves g_rl untouched.
    """
    g_rl = {"w": jnp.array([1.0, 0.0])}
    out = gradient_surgery(g_rl, {"w": jnp.array([-1.0, 1.0])})
    assert np.allclose(np.asarray(out["w"]), [0.5, 0.5], atol=1e-6)
    assert float(out["w"] @ jnp.array([-1.0, 1.0])) == pytest.approx(0.0, abs=1e-6)
    out2 = gradient_surgery(g_rl, {"w": jnp.array([1.0, 1.0])})
    assert np.allclose(np.asarray(out2["w"]), [1.0, 0.0])


# ---------------------------------------------------------------------------
# mixed_replay (self-replay ring buffer; spec-ablations §2 + §5 res. 12)
# ---------------------------------------------------------------------------


def test_mixed_replay_ring_buffer_holds_the_runs_own_windows():
    """mixed_replay's buffer holds the run's own rollout windows
    (self-replay; the 'offline data' one-liner is recorded stale by
    traceability §5 res. 12) with ring-wrap FIFO semantics and a
    first-n push cap.

    Derivation: capacity 4, push 3 rows with returns [0,1,2] (write idx
    0..2), then 3 with [3,4,5] (indices 3,0,1) -> buffer returns
    [4,5,2,3]; count caps at 4. A push with n_new=2 takes only the
    first 2 rows of its batch.
    """
    buf = _init_replay_buffer(4, H, OBS)

    def rows(vals):
        n = len(vals)
        return (
            jnp.tile(jnp.arange(H, dtype=jnp.int32), (n, 1)),
            jnp.zeros((n, OBS)),
            jnp.ones(n, dtype=bool),
            jnp.array(vals, dtype=jnp.float32),
        )

    a, o, v, r = rows([0.0, 1.0, 2.0])
    buf = _push_to_buffer(buf, a, o, v, r, 3)
    a, o, v, r = rows([3.0, 4.0, 5.0])
    buf = _push_to_buffer(buf, a, o, v, r, 3)
    assert np.allclose(np.asarray(buf.returns), [4.0, 5.0, 2.0, 3.0])
    assert int(buf.count) == 4

    buf2 = _init_replay_buffer(4, H, OBS)
    a, o, v, r = rows([7.0, 8.0, 9.0])
    buf2 = _push_to_buffer(buf2, a, o, v, r, 2)
    assert int(buf2.count) == 2
    assert np.allclose(np.asarray(buf2.returns[:2]), [7.0, 8.0])


# ---------------------------------------------------------------------------
# reward_model (spec-ablations §2: MLP obs -> return, MSE)
# ---------------------------------------------------------------------------


def test_reward_model_learns_a_linear_return_map():
    """The reward model is an MLP regressor on (obs -> return) trained
    with MSE (spec-ablations §2 reward_model row: width 64, depth 2).
    50 gradient steps on a fixed linear target must cut the MSE by more
    than half (deterministic under the fixed seed).
    """
    _, rm_state = _build_reward_model(
        OBS, jax.random.PRNGKey(SEED), width=64, depth=2, lr=1e-3
    )
    k = jax.random.PRNGKey(3)
    obs = jax.random.normal(k, (64, OBS))
    targets = obs[:, 0] * 2.0 + 1.0
    _, loss0 = _reward_model_train_step(rm_state, obs, targets)
    state = rm_state
    for _ in range(50):
        state, loss = _reward_model_train_step(state, obs, targets)
    assert float(loss) < 0.5 * float(loss0)


# ---------------------------------------------------------------------------
# running_stats effective decay (defect §8.4) - step-7 closure method
# ---------------------------------------------------------------------------


def test_running_stats_uses_the_configured_decay():
    """The running_stats training closure must consume
    RUNNING_STATS_EMA_DECAY (spec-ablations §1.6: 0.99) for the
    advantage EMA and EMA_DECAY for the eval weights (was defect §8.4:
    one shadowed variable served both). Reuses the step-7 reproduction:
    build the shipped make_run_ablation closure with distinct sentinel
    decays and read the captured cells.
    """
    from tests.conftest import load_config

    from experiments.rl_finetuning.ablations.training import make_run_ablation
    from src.planners.model import make_apply_fns

    config = {
        **load_config("configs/defaults.yaml"),
        **load_config("experiments/rl_finetuning/configs/ablations_default.yaml"),
        **TINY, "NUM_ACTIONS": V, "MAX_ITER": 1, "NUM_ENVS": 2,
        "NUM_STEPS": 8, "BATCH_SIZE": 3, "EVAL_EVERY": 1, "EVAL_STEPS": 1,
        "EVAL_REPLAN": 1, "USE_WANDB": False, "SEED": 0,
        "RUNNING_STATS_EMA_DECAY": 0.111, "EMA_DECAY": 0.999,
    }
    model = build_model(config, V)
    params = init_params(model, jax.random.PRNGKey(0), OBS, H)
    apply_eval, apply_train = make_apply_fns(model)
    run = make_run_ablation(
        spec=REGISTRY["running_stats"], config=config, pretrained_params=params,
        apply_train=apply_train, apply_eval=apply_eval, env=None, env_params=None,
        schedule_fn=COSINE[0], schedule_deriv_fn=COSINE[1],
        num_actions=V, obs_dim=OBS,
    )
    cells = dict(zip(run.__code__.co_freevars, run.__closure__ or (), strict=False))
    assert "running_stats_ema_decay" in cells, "closure extraction failed"
    assert cells["running_stats_ema_decay"].cell_contents == pytest.approx(0.111)
    assert cells["eval_ema_decay"].cell_contents == pytest.approx(0.999)


# ---------------------------------------------------------------------------
# Group C: trainable-parameter sets (defects §8.1/§8.2 + step-8 finding)
# step-7 reproduction method: ones-grads through the registry optimizer
# ---------------------------------------------------------------------------


def _trainable_modules(name: str) -> frozenset[str]:
    model = build_model(TINY, V)
    params = init_params(model, jax.random.PRNGKey(SEED), OBS, H)
    tx = REGISTRY[name].optimizer_factory({**TINY, "WEIGHT_DECAY": 0.0}, params)
    state = tx.init(params)
    updates, _ = tx.update(jax.tree.map(jnp.ones_like, params), state, params)
    flat = jax.tree_util.tree_flatten_with_path(updates)[0]
    modules = set()
    for path, leaf in flat:
        p = "/".join(str(k.key) for k in path)
        if bool(jnp.any(leaf != 0)):
            modules.add(p.rsplit("/", 1)[0])
    return frozenset(modules)


def _all_modules() -> frozenset[str]:
    model = build_model(TINY, V)
    params = init_params(model, jax.random.PRNGKey(SEED), OBS, H)
    flat = jax.tree_util.tree_flatten_with_path(params)[0]
    return frozenset(
        "/".join(str(k.key) for k in path).rsplit("/", 1)[0] for path, _ in flat
    )


# Tiny-arch module map (OBS_ENCODER_LAYERS=1, N_LAYERS=2): Dense_0 obs
# encoder, Dense_1 obs projection, Dense_2/Dense_3 time embedding,
# Dense_4 action head, Embed_0 action embedding, LayerNorm_0/1 top level.
_HEAD = frozenset({"params/Dense_4"})


def test_frozen_backbone_trains_the_head_and_token_embeddings():
    """Canonical set (spec-ablations §2, step-9 amendment): the action
    head plus the token-interface embeddings (action embedding and the
    time-embedding MLP); the backbone (obs encoder incl. projection,
    transformer blocks, all LayerNorms) is frozen. Tiny-arch modules:
    Dense_2/Dense_3 (t-emb), Dense_4 (head), Embed_0."""
    assert _trainable_modules("frozen_backbone") == frozenset(
        {"params/Dense_2", "params/Dense_3", "params/Dense_4", "params/Embed_0"}
    )


def test_head_only_trains_only_the_final_projection():
    """Canonical set (spec-ablations §2, step-9 amendment): exactly the
    final action projection."""
    assert _trainable_modules("head_only") == _HEAD


def test_attention_only_trains_only_the_attention_projections():
    """Canonical set (spec-ablations §2, step-9 amendment): exactly the
    per-block attention projections Q/K/V/O; norms and head frozen."""
    expected = frozenset(
        m for m in _all_modules() if "MultiHeadDotProductAttention_" in m
    )
    assert _trainable_modules("attention_only") == expected


def test_ffn_only_trains_only_the_ffn_layers():
    """Canonical set (spec-ablations §2, step-9 amendment): exactly the
    two FFN Dense layers inside each TransformerBlock; norms and head
    frozen."""
    expected = frozenset(
        m
        for m in _all_modules()
        if "TransformerBlock_" in m and ("/Dense_0" in m or "/Dense_1" in m)
    )
    assert _trainable_modules("ffn_only") == expected


@pytest.mark.parametrize("top_n", [1, 2])
def test_layer_ablation_trains_only_the_top_blocks_and_head(top_n):
    """Canonical set (spec-ablations §2, step-9 amendment): all
    parameters of the top-n transformer blocks plus the action head.
    With N_LAYERS=2 the top-1 set is TransformerBlock_1 + head; top-2
    adds TransformerBlock_0.
    """
    kept = {f"params/TransformerBlock_{i}" for i in range(2 - top_n, 2)}
    expected = frozenset(
        m for m in _all_modules() if any(m.startswith(k) for k in kept)
    ) | _HEAD
    assert _trainable_modules(f"layer_ablation_top{top_n}") == expected


# ---------------------------------------------------------------------------
# The freeze family at the PRODUCTION architecture
# ---------------------------------------------------------------------------
# Every ablation whose optimiser trains a strict subset of the parameters.
_FREEZING_ABLATIONS = (
    "frozen_backbone",
    "head_only",
    "attention_only",
    "ffn_only",
    "layer_ablation_top1",
    "layer_ablation_top2",
    "layer_ablation_top3",
    "lora",
)
# The other two partitioning optimisers, asked the opposite question: LLRD
# splits the tree by depth and baseline_rl not at all, so neither may freeze
# anything. Same mechanism, inverted expectation.
_NON_FREEZING_ABLATIONS = ("llrd", "baseline_rl")

_ABLATIONS_DEFAULT = (
    Path(__file__).resolve().parents[1]
    / "experiments"
    / "rl_finetuning"
    / "configs"
    / "ablations_default.yaml"
)


@pytest.fixture(scope="module")
def prod():
    """``(config, params)`` for the model the suite actually trains.

    Read from the shipped ``ablations_default.yaml`` through the suite's own
    upper-casing convention, with the action count and observation width
    taken from the real environment, so the layout tracks production instead
    of a copy of it. At the pins that is 113 leaves and 9,334,289 parameters:
    ``N_LAYERS=6``, ``D_MODEL=384``, ``OBS_ENCODER_LAYERS=2``.

    The TINY arch the tests above use is a different layout, not a smaller
    one: the group-C selectors derive their Dense indices arithmetically from
    ``OBS_ENCODER_LAYERS``, so the head is ``Dense_4`` at TINY and ``Dense_5``
    here, and ``layer_ablation_top3`` is unreachable with two blocks.
    """
    from src.planners.env import make_env

    config = {
        k.upper(): v
        for k, v in yaml.safe_load(_ABLATIONS_DEFAULT.read_text()).items()
    }
    env, env_params = make_env(config, 1)
    config["NUM_ACTIONS"] = env.action_space(env_params).n
    obs_dim = env.observation_space(env_params).shape[0]

    model = build_model(config, config["NUM_ACTIONS"])
    params = init_params(
        model, jax.random.PRNGKey(SEED), obs_dim, config["PLAN_HORIZON"]
    )
    return config, params


def _prod_deltas(config: dict, params, name: str) -> dict[str, float]:
    """Per-leaf ``max|parameter delta|`` after one step of *name*'s optimiser.

    Built the way ``run_ablation`` builds it, LoRA branch included, and driven
    by a non-zero gradient on EVERY leaf, the nominally frozen ones included.
    The measurement is the applied parameter delta rather than the update, so
    AdamW's decoupled weight decay is inside it: at the shipped non-zero
    ``WEIGHT_DECAY`` a frozen leaf that reached the optimiser would move.
    """
    if name == "lora":
        lora = make_lora_params(
            params, config.get("LORA_RANK", 8), jax.random.PRNGKey(SEED)
        )
        tree = {"base": params, "lora": lora}
        tx = make_optimizer_lora_only(config, params, lora)
    else:
        tree = params
        tx = REGISTRY[name].optimizer_factory(config, params)

    grads = jax.tree.map(lambda p: jnp.full_like(p, 0.5), tree)
    updates, _ = tx.update(grads, tx.init(tree), tree)
    updated = optax.apply_updates(tree, updates)

    flat_new = jax.tree_util.tree_flatten_with_path(updated)[0]
    flat_old = jax.tree_util.tree_flatten_with_path(tree)[0]
    return {
        "/".join(str(k.key) for k in path): float(jnp.max(jnp.abs(new - old)))
        for (path, new), (_, old) in zip(flat_new, flat_old, strict=True)
    }


def _moved(deltas: dict[str, float]) -> frozenset[str]:
    return frozenset(name for name, d in deltas.items() if d != 0.0)


@pytest.mark.parametrize("name", _FREEZING_ABLATIONS)
def test_frozen_leaves_move_exactly_zero_at_the_production_architecture(
    name, prod
):
    """Every leaf either moves a full optimiser step or does not move at all.

    This is the standing guard for the FREEZE defect. ``optax.masked`` leaves
    NON-selected updates untouched rather than zeroing them, so masking in the
    trainable parameters handed every "frozen" leaf its raw clipped gradient
    as an update -- roughly SGD at lr 1.0, in the ascent direction. It survived
    four months and a green suite because nothing measured a parameter delta,
    and nothing measured one at the architecture the suite runs.

    Two properties, both needed:

    - no leakage: a leaf's delta is exactly 0.0 or at least a tenth of the
      learning rate. A partial freeze, a stale gradient or decoupled weight
      decay reaching a frozen leaf all land between those and fail here,
      where an order-of-magnitude tolerance would pass them.
    - the partition is real: at least one leaf frozen and at least one leaf
      trained. An optimiser that trains nothing is the failure the minihack
      twin raises on; here the whole-tree ``set_to_zero`` label makes it a
      valid optimiser that silently updates nothing.
    """
    config, params = prod
    deltas = _prod_deltas(config, params, name)
    moved = _moved(deltas)
    floor = config["LR"] / 10

    leaking = {n: d for n, d in deltas.items() if 0.0 < d < floor}
    assert not leaking, f"{name}: leaves moved by less than a full step: {leaking}"
    assert moved, f"{name} trained nothing"
    assert moved < frozenset(deltas), f"{name} froze nothing"


@pytest.mark.parametrize("name", _NON_FREEZING_ABLATIONS)
def test_the_non_freezing_optimisers_freeze_nothing(name, prod):
    """LLRD and baseline_rl must reach every leaf.

    LLRD is the other ``multi_transform`` user, so it is swept for the same
    defect class from the opposite side: its label function has to be total,
    because any leaf it fails to label would be dropped rather than merely
    slowed. Its slowest group is ``lr * 0.9 ** (N_LAYERS + 1)``, comfortably
    above the leakage floor, so "slow" and "frozen" stay distinguishable.
    """
    config, params = prod
    deltas = _prod_deltas(config, params, name)

    assert _moved(deltas) == frozenset(deltas), (
        f"{name} froze "
        f"{sorted(set(deltas) - _moved(deltas))}"
    )
    assert min(deltas.values()) >= config["LR"] / 10


# ---------------------------------------------------------------------------
# Group C: trainable sets at the PRODUCTION architecture
# ---------------------------------------------------------------------------
# The Dense index layout the expectations below are derived for. At
# OBS_ENCODER_LAYERS=2 the top-level Dense modules are: 0-1 obs encoder MLP,
# 2 obs projection, 3-4 time-embedding MLP, 5 action head -- one index higher
# throughout than the TINY layout the tests above assert, where the head is
# Dense_4. `_dense_index_map` computes exactly this arithmetically, so a test
# that only ever sees OBS_ENCODER_LAYERS=1 cannot tell a correct map from one
# that is off by the encoder depth.
_PROD_OBS_ENCODER_LAYERS = 2
_PROD_N_LAYERS = 6
_PROD_HEAD = frozenset({"params/Dense_5"})
# Head, plus the token-interface embeddings: the time-embedding MLP
# (Dense_3, Dense_4) and the action embedding.
_PROD_FROZEN_BACKBONE = _PROD_HEAD | frozenset(
    {"params/Dense_3", "params/Dense_4", "params/Embed_0"}
)


def _prod_modules(params) -> frozenset[str]:
    return frozenset(
        "/".join(str(k.key) for k in path).rsplit("/", 1)[0]
        for path, _ in jax.tree_util.tree_flatten_with_path(params)[0]
    )


def _prod_trainable_modules(config: dict, params, name: str) -> frozenset[str]:
    return frozenset(
        leaf.rsplit("/", 1)[0]
        for leaf, delta in _prod_deltas(config, params, name).items()
        if delta != 0.0
    )


def test_the_production_layout_is_the_one_the_group_c_expectations_assume(prod):
    """Pin the architecture the expectations below are derived for.

    Without this a config change would surface as an unreadable set
    mismatch in every group-C case at once, rather than as one failure
    saying the layout moved and the derivations need redoing.
    """
    config, params = prod
    assert config["OBS_ENCODER_LAYERS"] == _PROD_OBS_ENCODER_LAYERS
    assert config["N_LAYERS"] == _PROD_N_LAYERS
    assert len(jax.tree_util.tree_leaves(params)) == 113


def test_frozen_backbone_at_production_trains_the_head_and_token_embeddings(prod):
    """Canonical set (spec-ablations §2, step-9 amendment) at the real
    Dense layout: Dense_3/Dense_4 (t-emb), Dense_5 (head), Embed_0. The
    obs encoder including its projection, the six transformer blocks and
    all LayerNorms are frozen."""
    config, params = prod
    assert _prod_trainable_modules(config, params, "frozen_backbone") == (
        _PROD_FROZEN_BACKBONE
    )


def test_head_only_at_production_trains_only_the_final_projection(prod):
    """Canonical set: exactly the final action projection, which is
    Dense_5 here and Dense_4 at TINY -- the index defect §8.1/§8.2 turned
    on."""
    config, params = prod
    trainable = _prod_trainable_modules(config, params, "head_only")
    assert trainable == _PROD_HEAD
    assert trainable < _prod_trainable_modules(config, params, "frozen_backbone")


def test_attention_only_at_production_trains_only_the_attention_projections(prod):
    """Canonical set: exactly the per-block attention projections across
    all six blocks; the LayerNorms, the FFNs and the head are frozen."""
    config, params = prod
    expected = frozenset(
        m for m in _prod_modules(params) if "MultiHeadDotProductAttention_" in m
    )
    assert len(expected) == 4 * _PROD_N_LAYERS
    assert _prod_trainable_modules(config, params, "attention_only") == expected


def test_ffn_only_at_production_trains_only_the_ffn_layers(prod):
    """Canonical set: exactly the two FFN Dense layers inside each of the
    six blocks; norms and head frozen. The per-block Dense_0/Dense_1 must
    not be confused with the top-level Dense_0/Dense_1, which are the obs
    encoder -- the substring collision behind §8.1."""
    config, params = prod
    expected = frozenset(
        m
        for m in _prod_modules(params)
        if "TransformerBlock_" in m and ("/Dense_0" in m or "/Dense_1" in m)
    )
    assert len(expected) == 2 * _PROD_N_LAYERS
    assert _prod_trainable_modules(config, params, "ffn_only") == expected


@pytest.mark.parametrize("top_n", [1, 2, 3])
def test_layer_ablation_at_production_trains_only_the_top_blocks_and_head(
    top_n, prod
):
    """Canonical set: every parameter of the top-n transformer blocks plus
    the action head.

    `layer_ablation_top3` is only reachable here: the TINY arch has two
    blocks, so the top-3 arm was the one registry entry no test touched.
    With N_LAYERS=6 the top-1 set is TransformerBlock_5 + head, top-2 adds
    block 4 and top-3 block 3.
    """
    config, params = prod
    kept = {
        f"params/TransformerBlock_{i}"
        for i in range(_PROD_N_LAYERS - top_n, _PROD_N_LAYERS)
    }
    expected = frozenset(
        m for m in _prod_modules(params) if any(m.startswith(k) for k in kept)
    ) | _PROD_HEAD
    assert _prod_trainable_modules(
        config, params, f"layer_ablation_top{top_n}"
    ) == expected


# ---------------------------------------------------------------------------
# Suite loss estimator (cross-repo twin; NELBO per spec-method §3.1/§3.4)
# ---------------------------------------------------------------------------


def test_suite_loss_uses_the_nelbo_weight_and_per_token_normalisation():
    """The suite's per-sample loss is the NELBO estimator
    w(t) * sum_masked(CE) / H (spec-method §3.1/§3.4; spec-ablations §2
    baseline row: 'return-weighted ELBO').

    Derivation (cosine schedule, uniform logits, t ~ U(eps, 0.2) via the
    low_t factory, unit advantages): E[per-sample] = E_t[w(t) ln V
    (1-alpha)/1] = ln V * E_t[-alpha'(t)] = ln V * (1 - cos(0.1 pi))/0.2
    = 0.244715 ln V. Statistical: B=16384 rows; per-row second moment
    is ~0.62 (ln V)^2 (w ~ 2/t and Bin(H=8, 1-alpha) masking), giving
    sigma ~ 0.0061 ln V; bound 0.03 ln V ~ 4.9 sigma. The minihack twin
    asserts the same value and xfails: its suite loss drops w(t) and
    normalises by the realised masked count (step-8 finding).
    """
    b = 16384

    def apply_fn(params, obs, z, t, rng):
        return jnp.zeros((obs.shape[0], H, V))

    ctx = LossContext(
        apply_fn=apply_fn, ref_params=None, schedule_fn=COSINE[0],
        schedule_deriv_fn=COSINE[1], num_actions=V, config={"T_MAX_LOW": 0.2},
    )
    k = jax.random.PRNGKey(2)
    acts = jax.random.randint(k, (b, H), 0, V)
    obs = jnp.zeros((b, 1))
    loss = float(
        make_loss_low_t(ctx)(None, acts, obs, jnp.ones(b), jax.random.PRNGKey(0),
                             jnp.ones(b))
    )
    expected = math.log(V) * (1 - math.cos(0.1 * math.pi)) / 0.2
    assert abs(loss - expected) < 0.03 * math.log(V)