giant/giant/training/trainers.py

"""Per-stage trainers: optimizer(s), EMA, LR schedule, and the per-batch step.

`StageSpec` resolves one stage's slice of the config once, so the two
concrete trainers share a single constructor shape instead of ~24 keyword
arguments each, and `StageTrainer` carries every piece that used to be
copy-pasted between them (cosine warmup, EMA, checkpoint state, LR resume,
train/eval toggling).

Each trainer also *declares* the metrics it emits, as `MetricSpec` lists —
that declaration is the single source of truth for `metrics.csv` and W&B
columns (see `giant.training.metrics`) — and exposes the three small hooks
(`batch_loss`, `summary`, `val_objective`) that let the epoch loop treat
adversarial and non-adversarial stages identically.
"""

import copy
import math
from dataclasses import dataclass, field

import torch
import torch.nn.functional as F
import torch.optim as optim

from giant.constants import CONT_SLOT_DIM
from giant.model.network import Router, stage2_type_dim
from giant.model.schedule import (
    CosineSchedule,
    flow_matching_loss,
    flow_matching_loss_secondary,
    flow_matching_loss_secondary_ar,
)
from giant.model.wgan import generator_loss, gradient_penalty
from giant.training.metrics import MetricSpec, stage_metric, train_metric, val_metric
from giant.training.stage2_inputs import (
    _assemble_stage2_ar_inputs_scheduled,
    _assemble_stage2_ar_target,
    _gumbel_tau,
    _relax_onehot_type_slice,
    _stage2_tf_prob,
)


@torch.no_grad()
def _update_ema(
    ema_model: torch.nn.Module, model: torch.nn.Module, decay: float
) -> None:
    for ema_p, p in zip(ema_model.parameters(), model.parameters()):
        ema_p.mul_(decay).add_(p, alpha=1 - decay)


def _stage_router(model: torch.nn.Module) -> Router | None:
    """A stage model's Router, if its trunk is routed — else None.

    Post-step-2 refactor the router lives at `model.trunk.router`
    (`giant.model.network.RoutedTrunk`), not `model.router` directly.
    """
    trunk = getattr(model, "trunk", None)
    return getattr(trunk, "router", None)


def _cosine_warmup_lambda(warmup_steps: int, total_steps: int):
    """Linear warmup for `warmup_steps`, then cosine decay to zero over the
    remainder — the LR schedule both trainers use, in their own step units
    (optimizer steps for flow/ddpm, generator steps for WGAN)."""

    def _lr_lambda(step: int) -> float:
        if warmup_steps > 0 and step < warmup_steps:
            return (step + 1) / warmup_steps
        t = step - warmup_steps
        T = max(total_steps - warmup_steps, 1)
        return 0.5 * (1.0 + math.cos(math.pi * min(t, T) / T))

    return _lr_lambda


def _batch_to_device(batch: tuple, device: torch.device) -> tuple:
    return tuple(t.to(device) for t in batch)


@dataclass(frozen=True)
class StageSpec:
    """One stage's resolved training configuration.

    Built once by `StageSpec.from_config`, which is the only place that reads
    the `cfg` dict — so a new config key means one new field and one new read,
    not another argument threaded through two constructors.
    """

    name: str
    is_stage2: bool
    generator: str
    decoder: str = "one_shot"

    # loss weights
    lambda_weight: float = 1.0
    n_sec_lambda: float = 0.1

    # particle-type target (stage 2 only)
    particle_type: dict = field(default_factory=lambda: {"target": "physical"})
    particle_type_emb_dim: int = 16

    # optimization
    lr: float = 3e-4
    weight_decay: float = 0.01
    ema_decay: float = 0.9999
    warmup_epochs: int = 0
    epochs: int = 1
    steps_per_epoch: int = 1

    # routing auxiliaries
    lambda_balance: float = 0.0
    lambda_proc: float = 0.0
    lambda_entropy: float = 0.0
    gumbel_tau_start: float = 1.0
    gumbel_tau_end: float = 0.1

    # autoregressive stage 2
    teacher_forcing: str = "always"
    tf_p_start: float = 1.0
    tf_p_end: float = 1.0
    ar_sample_steps: int = 10

    # generator-specific
    ddpm_n_steps: int = 1000
    n_critic: int = 5
    gp_weight: float = 10.0
    critic_lr: float = 0.0
    type_gumbel_tau_start: float = 1.0
    type_gumbel_tau_end: float = 0.1

    @classmethod
    def from_config(
        cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int
    ) -> "StageSpec":
        t = cfg["train"]
        stage_cfg = cfg[f"{name}_model"]
        router_cfg = stage_cfg.get("router") or {}
        wgan_cfg = stage_cfg.get("wgan") or {}
        ar_cfg = (stage_cfg.get("autoregressive") or {}) if is_stage2 else {}
        return cls(
            name=name,
            is_stage2=is_stage2,
            generator=stage_cfg["generator"],
            decoder=stage_cfg.get("decoder", "one_shot") if is_stage2 else "one_shot",
            lambda_weight=stage_cfg.get("lambda", 1.0),
            n_sec_lambda=cfg["stage2_model"].get("n_sec", {}).get("lambda", 0.1),
            particle_type=cfg["stage2_model"].get("particle_type")
            or {"target": "physical"},
            particle_type_emb_dim=cfg["conditioning"]["particle"]["emb_dim"],
            # train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge
            # (giant/config.py), so they read directly; the field defaults
            # below exist only for tests that construct StageSpec by hand.
            lr=t["lr"],
            weight_decay=t["weight_decay"],
            ema_decay=t["ema_decay"],
            warmup_epochs=t["warmup_epochs"],
            epochs=t["epochs"],
            steps_per_epoch=max(steps_per_epoch, 1),
            lambda_balance=router_cfg.get("lambda_balance", 0.0),
            lambda_proc=router_cfg.get("lambda_proc", 0.0),
            lambda_entropy=router_cfg.get("lambda_entropy", 0.0),
            gumbel_tau_start=router_cfg.get("gumbel_tau_start", 1.0),
            gumbel_tau_end=router_cfg.get("gumbel_tau_end", 0.1),
            teacher_forcing=ar_cfg.get("teacher_forcing", "always"),
            tf_p_start=ar_cfg.get("tf_p_start", 1.0),
            tf_p_end=ar_cfg.get("tf_p_end", 1.0),
            # AR self-sampling under scheduled/never teacher forcing reuses
            # train.validate_steps as its flow-matching ODE step count — no
            # dedicated config key for this (docs/v0.3.0-design.md §3.3 lists
            # tf_p_start/tf_p_end/attn_n_heads/attn_n_layers only).
            ar_sample_steps=t["validate_steps"],
            ddpm_n_steps=stage_cfg.get("ddpm", {}).get("n_steps", 1000),
            n_critic=wgan_cfg.get("n_critic", 5),
            gp_weight=wgan_cfg.get("gp_weight", 10.0),
            critic_lr=wgan_cfg.get("critic_lr", 0.0),
            type_gumbel_tau_start=wgan_cfg.get("gumbel_tau_start", 1.0),
            type_gumbel_tau_end=wgan_cfg.get("gumbel_tau_end", 0.1),
        )


class StageTrainer:
    """One active stage's optimizer(s), EMA, and per-batch step.

    Reads only the shared batch tuple `(cond_cont, cond_cat, x1_s1, n_sec,
    sec_cont, proc_idx, sec_type_idx)` — stage 2 always conditions on the
    ground-truth `x1_s1` (`stage2_model.stage1_context = "truth"`,
    stage-level teacher forcing; `"sampled"` is not implemented — see
    docs/v0.3.0-design.md §3.3), so stage trainers never need each other's
    output at train time. This means "stage-2-only training is a cheap
    ablation, not new plumbing" (design doc §7) falls out for free: a trainer
    only exists for active stages, and inactive stages are simply never
    constructed.

    Grad-norm clipping is per-stage here — v0.2's single shared optimizer
    clipped both stages' gradients jointly; splitting per stage is a small,
    disclosed behavior change. It doesn't affect Adam's per-parameter update
    math itself (no cross-parameter coupling), only the clip threshold's
    scope.
    """

    #: Metrics this trainer emits, declared once — `giant.training.metrics`
    #: derives every CSV/W&B column from these. Instance attributes rather
    #: than class constants because some are conditional on the stage's own
    #: configuration (see `WGANStageTrainer.__init__`).
    train_metrics: list[MetricSpec]
    val_metrics: list[MetricSpec]
    stage_metrics: list[MetricSpec]

    #: False for adversarial stages, which have no monotone per-batch
    #: validation loss worth averaging (see `val_objective`).
    supports_val_loss: bool = True

    #: Built by the subclass (the optimizer flavour differs) and wired to the
    #: schedule via `_init_lr_schedule`.
    optimizer: optim.Optimizer
    lr_sched: optim.lr_scheduler.LambdaLR
    total_steps: int

    def __init__(
        self,
        spec: StageSpec,
        model: torch.nn.Module,
        device: torch.device,
        extra_modules: tuple[torch.nn.Module, ...] = (),
    ) -> None:
        self.spec = spec
        self.name = spec.name
        self.is_stage2 = spec.is_stage2
        self.generator = spec.generator
        self.decoder = spec.decoder
        self.device = device
        self.model = model.to(device)
        self.router = _stage_router(self.model)
        self._modules = (self.model, *extra_modules)

        self.particle_type_cfg = dict(spec.particle_type or {"target": "physical"})
        self.particle_type_emb_dim = spec.particle_type_emb_dim
        self.ema_decay = spec.ema_decay

        self.ema_model: torch.nn.Module | None = None
        if spec.ema_decay > 0:
            self.ema_model = copy.deepcopy(self.model).eval()
            for p in self.ema_model.parameters():
                p.requires_grad_(False)

    # --- schedule -------------------------------------------------------

    def _init_lr_schedule(
        self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int
    ) -> None:
        self._lr_lambda = _cosine_warmup_lambda(warmup_steps, total_steps)
        self.total_steps = total_steps
        self.lr_sched = optim.lr_scheduler.LambdaLR(optimizer, self._lr_lambda)

    # --- per-batch (subclass responsibility) ----------------------------

    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
        raise NotImplementedError

    def val_loss(self, batch: tuple, device: torch.device) -> dict:
        raise NotImplementedError

    # --- reporting hooks ------------------------------------------------

    def batch_loss(self, stats: dict) -> float:
        """The single number this stage contributes to the progress bar's
        smoothed loss."""
        raise NotImplementedError

    def summary(self, means: dict) -> str:
        """This stage's fragment of the end-of-epoch console line."""
        raise NotImplementedError

    def val_objective(
        self, train_means: dict, val_means: dict, marginal_kl: float
    ) -> float:
        """This stage's contribution to the best-checkpoint selection score."""
        raise NotImplementedError

    # --- mode / state ---------------------------------------------------

    def sampling_model(self) -> torch.nn.Module:
        return self.ema_model if self.ema_model is not None else self.model

    def train_mode(self) -> None:
        for module in self._modules:
            module.train()

    def eval_mode(self) -> None:
        for module in self._modules:
            module.eval()

    # --- stage-2 secondary assembly (shared by both trainer subclasses) ---

    def _ar_inputs(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_ctx: torch.Tensor,
        sec_cont: torch.Tensor,
        sec_type_idx: torch.Tensor,
        n_sec: torch.Tensor,
        epoch: int | None,
    ) -> dict[str, torch.Tensor]:
        """Per-token AR conditioning for this stage's secondary decoder.

        `epoch=None` means full teacher forcing (`p_tf=1.0`) regardless of
        `spec.teacher_forcing` — the val-loss convention, kept in this one
        place so both trainer subclasses honor it identically.
        """
        p_tf = (
            1.0
            if epoch is None
            else _stage2_tf_prob(
                self.spec.teacher_forcing,
                self.spec.tf_p_start,
                self.spec.tf_p_end,
                epoch,
                self.spec.epochs,
            )
        )
        return _assemble_stage2_ar_inputs_scheduled(
            self.model,
            cond_cont,
            cond_cat,
            stage1_ctx,
            sec_cont,
            sec_type_idx,
            n_sec,
            self.particle_type_cfg,
            self.model.cond_enc,
            self.particle_type_emb_dim,
            p_tf,
            self.spec.ar_sample_steps,
        )

    def _sec_target(
        self,
        sec_cont: torch.Tensor,
        sec_type_idx: torch.Tensor,
        generator: str,
        *,
        flatten: bool,
    ) -> torch.Tensor:
        """Ground-truth stage-2 target for this stage's secondary decoder,
        per the (particle-type target, generator) width rules in
        `_assemble_stage2_ar_target`. `flatten=True` gives `Stage2OneShot`'s
        flattened `(B, K*token_dim)` form (the old `_real`); `flatten=False`
        gives `Stage2Autoregressive`'s per-token `(B, K, token_dim)` form (the
        old `_ar_target`) — the two are the same tensor modulo `.flatten(1)`,
        so the width rules live in one place (`stage2_inputs.py`)."""
        target = _assemble_stage2_ar_target(
            sec_cont,
            sec_type_idx,
            self.particle_type_cfg,
            generator,
            self.model.cond_enc,
            self.particle_type_emb_dim,
        )
        return target.flatten(1) if flatten else target

    @staticmethod
    def _sec_mask(
        n_sec: torch.Tensor, k_max: int, device: torch.device
    ) -> torch.Tensor:
        """`(B, K_MAX)` bool prefix mask: slot k is valid iff `k < n_sec`."""
        return torch.arange(k_max, device=device).unsqueeze(0) < n_sec.unsqueeze(1)

    def _n_sec_loss(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_ctx: torch.Tensor,
        n_sec: torch.Tensor,
        device: torch.device,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """`(l_nsec, nsec_acc)` for this stage's multiplicity classifier —
        zeros when the stage owns no `n_sec_head` (stage 1 now that n_sec
        defaults to stage 2, or any stage without the head). Owns the only
        stage1-vs-stage2 `predict_n_sec` signature split, shared by the
        non-adversarial and WGAN trainers.

        Gated on `n_sec_head is None`, not on `n_sec.mode`: a future
        `mode="stop_token"` model (design doc §11.2, currently rejected in
        `validate_config`) carries no head and would train its EOS signal in
        the generator/AR loss path instead, so this correctly stays zero.
        """
        if self.model.n_sec_head is None:
            zero = torch.zeros((), device=device)
            return zero, zero
        logits = (
            self.model.predict_n_sec(cond_cont, cond_cat, stage1_ctx)
            if self.is_stage2
            else self.model.predict_n_sec(cond_cont, cond_cat)
        )
        l_nsec = F.cross_entropy(logits, n_sec)
        nsec_acc = (logits.argmax(dim=-1) == n_sec).float().mean()
        return l_nsec, nsec_acc

    @staticmethod
    def _step_optimizer(
        optimizer: optim.Optimizer, loss: torch.Tensor, params: list
    ) -> float:
        """`zero_grad -> backward -> clip_grad_norm_(1.0) -> step`, returning
        the pre-clip grad norm. The one place the grad-clip constant lives."""
        optimizer.zero_grad()
        loss.backward()
        grad_norm = torch.nn.utils.clip_grad_norm_(params, 1.0)
        optimizer.step()
        return grad_norm.item()

    def _extra_state(self) -> dict:
        """Subclass state beyond model/optimizer/lr_sched/EMA."""
        return {}

    def _load_extra_state(self, sd: dict) -> None:
        return None

    def state_dict(self) -> dict:
        sd = {
            "model": self.model.state_dict(),
            "optimizer": self.optimizer.state_dict(),
            "lr_sched": self.lr_sched.state_dict(),
        }
        if self.ema_model is not None:
            sd["model_ema"] = self.ema_model.state_dict()
        sd.update(self._extra_state())
        return sd

    def load_state_dict(self, sd: dict) -> None:
        self.model.load_state_dict(sd["model"])
        self.optimizer.load_state_dict(sd["optimizer"])
        self.lr_sched.load_state_dict(sd["lr_sched"])
        if self.ema_model is not None:
            self.ema_model.load_state_dict(sd.get("model_ema", sd["model"]))
        self._load_extra_state(sd)

    def _resume_extra_lr(self, lr: float) -> None:
        return None

    def resume_lr(self, lr: float) -> None:
        """Restore the configured `lr`'s authority after `load_state_dict`
        restored the checkpoint's own base LR."""
        self.lr_sched.base_lrs = [lr for _ in self.lr_sched.base_lrs]
        resumed_lr = lr * self._lr_lambda(self.lr_sched.last_epoch)
        for group in self.optimizer.param_groups:
            group["lr"] = resumed_lr
        self._resume_extra_lr(lr)


class FlowDDPMStageTrainer(StageTrainer):
    """flow or ddpm generator for a single stage."""

    def __init__(
        self, spec: StageSpec, model: torch.nn.Module, device: torch.device
    ) -> None:
        if spec.is_stage2 and spec.generator not in ("flow",):
            raise NotImplementedError(
                f"stage2_model.generator={spec.generator!r} is accepted by the "
                "schema but not implemented in v0.3.0 for stage 2 (only "
                "'flow' and 'wgan' have a stage-2 secondary-decoder loss — "
                "see docs/v0.3.0-design.md §11.2)"
            )
        super().__init__(spec, model, device)
        self.particle_type_lambda = self.particle_type_cfg.get("lambda", 1.0)
        # Width of the type slice actually folded into x1_s2 by _sec_target,
        # under this trainer's generator (flow/ddpm only — see the
        # NotImplementedError above): "physical" keeps it folded in
        # (PARTICLE_PHYS_DIM wide, unchanged from v0.2); "onehot"/"embedding"
        # pull it out into model.type_head instead (0 here).
        self._flow_type_dim = (
            None
            if self.particle_type_cfg.get("target", "physical") == "physical"
            else 0
        )

        self.params = list(self.model.parameters())
        self.optimizer = optim.AdamW(
            self.params, lr=spec.lr, weight_decay=spec.weight_decay
        )
        self._init_lr_schedule(
            self.optimizer,
            warmup_steps=spec.warmup_epochs * spec.steps_per_epoch,
            total_steps=max(spec.epochs * spec.steps_per_epoch, 1),
        )
        self.ddpm_schedule = (
            CosineSchedule(T=spec.ddpm_n_steps).to(device)
            if spec.generator == "ddpm"
            else None
        )

        self.train_metrics = [
            train_metric(key)
            for key in (
                "loss",
                "loss_gen",
                "loss_nsec",
                "loss_balance",
                "loss_proc",
                "loss_entropy",
                "nsec_acc",
                "loss_type",
                "type_acc",
                "grad_norm",
            )
        ]
        self.val_metrics = [
            val_metric(key)
            for key in (
                "loss",
                "loss_gen",
                "loss_nsec",
                "nsec_acc",
                "loss_type",
                "type_acc",
            )
        ]
        self.stage_metrics = [stage_metric("lr")]

    def _generator_loss(
        self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None
    ):
        if not self.is_stage2:
            if self.generator == "flow":
                return flow_matching_loss(self.model, x1_s1, cond_cont, cond_cat)
            assert self.ddpm_schedule is not None
            return self.ddpm_schedule.loss(self.model, x1_s1, cond_cont, cond_cat)
        if self.decoder == "autoregressive":
            assert ar_inputs is not None
            return flow_matching_loss_secondary_ar(
                self.model,
                x1_s2,
                cond_cont,
                cond_cat,
                stage1_ctx,
                ar_inputs["history_feat"],
                ar_inputs["has_prev"],
                ar_inputs["remaining_frac"],
                ar_inputs["slot_idx"],
                sec_mask,
                type_dim=self._flow_type_dim,
            )
        return flow_matching_loss_secondary(
            self.model,
            x1_s2,
            cond_cont,
            cond_cat,
            stage1_ctx,
            sec_mask,
            type_dim=self._flow_type_dim,
        )

    def _type_loss(
        self,
        cond_cont,
        cond_cat,
        stage1_ctx,
        sec_type_idx,
        sec_mask,
        device,
        ar_inputs=None,
    ):
        """CE (`target="onehot"`) or MSE (`target="embedding"`) loss for the
        stage-2 model's `type_head` — the non-adversarial counterpart to
        WGANStageTrainer's ST-Gumbel-into-the-critic path (decision 2/5).
        Zero when this stage has no `type_head` (stage 1, or
        `particle_type.target = "physical"`)."""
        l_type = torch.zeros((), device=device)
        type_acc = torch.zeros((), device=device)
        type_head = getattr(self.model, "type_head", None)
        if not self.is_stage2 or type_head is None:
            return l_type, type_acc
        if self.decoder == "autoregressive":
            assert ar_inputs is not None
            type_out = self.model.predict_type(
                cond_cont,
                cond_cat,
                stage1_ctx,
                ar_inputs["history_feat"],
                ar_inputs["has_prev"],
                ar_inputs["remaining_frac"],
                ar_inputs["slot_idx"],
            )
        else:
            type_out = self.model.predict_type(cond_cont, cond_cat, stage1_ctx)
        mask = sec_mask.float()
        denom = mask.sum().clamp(min=1)
        if self.particle_type_cfg.get("target") == "onehot":
            ce = F.cross_entropy(
                type_out.transpose(1, 2), sec_type_idx, reduction="none"
            )
            l_type = (ce * mask).sum() / denom
            type_acc = (
                (type_out.argmax(-1) == sec_type_idx).float() * mask
            ).sum() / denom
        else:  # "embedding"
            target_vec = self.model.cond_enc.pdg_emb(sec_type_idx).detach()
            se = ((type_out - target_vec) ** 2).mean(-1)
            l_type = (se * mask).sum() / denom
        return l_type, type_acc

    def _compute(
        self, batch: tuple, device: torch.device, epoch: int | None = None
    ) -> dict:
        """`epoch=None` (the `val_loss` path) always uses full teacher
        forcing (`p_tf=1.0`) regardless of `spec.teacher_forcing` — validation
        should stay a stable, non-stochastic ground-truth comparison; only
        the training `step` path schedules `p_tf` by epoch."""
        (
            cond_cont,
            cond_cat,
            x1_s1,
            n_sec,
            sec_cont,
            proc_idx,
            sec_type_idx,
        ) = _batch_to_device(batch, device)
        sec_mask = self._sec_mask(n_sec, sec_cont.size(1), device)
        stage1_ctx = x1_s1.detach()

        x1_s2 = None
        ar_inputs = None
        if self.is_stage2 and self.decoder == "autoregressive":
            ar_inputs = self._ar_inputs(
                cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch
            )
            x1_s2 = self._sec_target(
                sec_cont, sec_type_idx, self.generator, flatten=False
            )
        elif self.is_stage2:
            x1_s2 = self._sec_target(
                sec_cont, sec_type_idx, self.generator, flatten=True
            )

        l_gen = self._generator_loss(
            cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs
        )
        l_nsec, nsec_acc = self._n_sec_loss(
            cond_cont, cond_cat, stage1_ctx, n_sec, device
        )

        l_type, type_acc = self._type_loss(
            cond_cont,
            cond_cat,
            stage1_ctx,
            sec_type_idx,
            sec_mask,
            device,
            ar_inputs=ar_inputs,
        )

        l_balance = l_proc = l_entropy = torch.zeros((), device=device)
        if self.router is not None:
            l_balance = self.router.balance_loss(cond_cont, cond_cat)
            l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx)
            l_entropy = self.router.entropy_loss(cond_cont, cond_cat)

        total = (
            self.spec.lambda_weight * l_gen
            + self.spec.n_sec_lambda * l_nsec
            + self.particle_type_lambda * l_type
        )
        if self.spec.lambda_balance > 0:
            total = total + self.spec.lambda_balance * l_balance
        if self.spec.lambda_proc > 0:
            total = total + self.spec.lambda_proc * l_proc
        if self.spec.lambda_entropy > 0:
            total = total + self.spec.lambda_entropy * l_entropy

        return {
            "loss": total,
            "loss_gen": l_gen,
            "loss_nsec": l_nsec,
            "loss_type": l_type,
            "type_acc": type_acc,
            "loss_balance": l_balance,
            "loss_proc": l_proc,
            "loss_entropy": l_entropy,
            "nsec_acc": nsec_acc,
        }

    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
        if self.router is not None:
            self.router.gumbel_tau = _gumbel_tau(
                global_step,
                self.total_steps,
                self.spec.gumbel_tau_start,
                self.spec.gumbel_tau_end,
            )
        epoch = global_step // self.spec.steps_per_epoch
        out = self._compute(batch, device, epoch=epoch)
        grad_norm = self._step_optimizer(self.optimizer, out["loss"], self.params)
        self.lr_sched.step()
        if self.ema_model is not None:
            _update_ema(self.ema_model, self.model, self.ema_decay)
        stats = {key: value.item() for key, value in out.items()}
        stats["grad_norm"] = grad_norm
        stats["lr"] = self.optimizer.param_groups[0]["lr"]
        return stats

    @torch.no_grad()
    def val_loss(self, batch: tuple, device: torch.device) -> dict:
        return {
            key: value.item() for key, value in self._compute(batch, device).items()
        }

    # --- reporting ------------------------------------------------------

    def batch_loss(self, stats: dict) -> float:
        return stats["loss"]

    def summary(self, means: dict) -> str:
        return f"{self.name}[loss={means.get('loss', 0.0):.3f}]"

    def val_objective(
        self, train_means: dict, val_means: dict, marginal_kl: float
    ) -> float:
        return val_means.get("loss", 0.0)


class WGANStageTrainer(StageTrainer):
    """WGAN-GP generator+critic for a single stage (see giant/model/wgan.py).

    Ports `_wgan_train_step` to operate on one stage instead of two fused
    together — the critic updates every batch; every `n_critic`-th batch
    additionally updates the generator (`did_g_step`). The (non-adversarial)
    n_sec classifier, when this stage's model owns it, updates every batch
    regardless — folded into whichever generator optimizer step happens this
    batch, same precedent as v0.2.
    """

    supports_val_loss = False

    def __init__(
        self,
        spec: StageSpec,
        model: torch.nn.Module,
        critic: torch.nn.Module,
        device: torch.device,
    ) -> None:
        self.critic = critic.to(device)
        super().__init__(spec, model, device, extra_modules=(self.critic,))
        self.n_critic = max(spec.n_critic, 1)
        self.gp_weight = spec.gp_weight
        self.critic_lr = spec.critic_lr

        self.g_params = list(self.model.parameters())
        self.d_params = list(self.critic.parameters())
        # WGAN-GP recipe (Gulrajani et al. 2017): Adam, beta1=0, no weight decay.
        self.optimizer = optim.Adam(self.g_params, lr=spec.lr, betas=(0.0, 0.9))
        self.optimizer_d = optim.Adam(
            self.d_params,
            lr=spec.critic_lr if spec.critic_lr > 0 else spec.lr,
            betas=(0.0, 0.9),
        )

        # Generator steps fire every n_critic-th batch, so warmup/decay must
        # be counted in those units, matching v0.2.
        gen_steps_per_epoch = max(spec.steps_per_epoch // self.n_critic, 1)
        self._init_lr_schedule(
            self.optimizer,
            warmup_steps=spec.warmup_epochs * gen_steps_per_epoch,
            total_steps=max(spec.epochs * gen_steps_per_epoch, 1),
        )

        train_keys = [
            "d_loss",
            "g_loss",
            "wasserstein",
            "gp_loss",
            "loss_nsec",
            "nsec_acc",
            "grad_norm_d",
            "grad_norm_g",
        ]
        if self.is_stage2 and self.particle_type_cfg.get("target") == "onehot":
            # §11.4 differentiability instrumentation — only meaningful when
            # the type slice is a straight-through Gumbel relaxation.
            train_keys += ["grad_norm_type_slice", "grad_norm_cont_slice"]
        self.train_metrics = [train_metric(key) for key in train_keys]
        self.val_metrics = []
        self.stage_metrics = [stage_metric("lr"), stage_metric("critic_lr")]

    def _stage2_real_and_fake(self, batch_tensors, stage1_ctx, global_step, device):
        """Build `(real, fake_raw, mask, critic_fn)` for stage 2, covering
        both decoders and all three particle-type targets. `fake_raw` still
        needs the caller's straight-through relaxation under
        `particle_type.target = "onehot"`, and neither tensor is masked-and-
        multiplied on the fake side yet."""
        cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors
        B = cond_cont.size(0)
        type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_emb_dim)
        slot_width = CONT_SLOT_DIM + type_dim
        k_max = sec_cont.size(1)

        sec_mask = self._sec_mask(n_sec, k_max, device)
        mask = sec_mask.unsqueeze(-1).expand(-1, -1, slot_width).reshape(B, -1).float()

        def critic_fn(x):
            return self.critic(x, cond_cont, cond_cat, stage1_ctx)

        if self.decoder == "autoregressive":
            epoch = global_step // self.spec.steps_per_epoch
            ar = self._ar_inputs(
                cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch
            )
            real = (
                self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(
                    B, -1
                )
                * mask
            )
            z = torch.randn(B, k_max, self.model.noise_dim, device=device)
            fake_raw = self.model(
                z,
                cond_cont,
                cond_cat,
                stage1_ctx,
                ar["history_feat"],
                ar["has_prev"],
                ar["remaining_frac"],
                ar["slot_idx"],
            ).reshape(B, -1)
        else:
            real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=True) * mask
            z = torch.randn(B, self.model.noise_dim, device=device)
            fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx)

        return real, fake_raw, mask, critic_fn

    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
        (
            cond_cont,
            cond_cat,
            x1_s1,
            n_sec,
            sec_cont,
            _proc_idx,
            sec_type_idx,
        ) = _batch_to_device(batch, device)
        B = cond_cont.size(0)
        stage1_ctx = x1_s1.detach()
        grad_probe: dict[str, float] = {}

        if not self.is_stage2:
            real = x1_s1

            def critic_fn(x):
                return self.critic(x, cond_cont, cond_cat)

            z = torch.randn(B, self.model.noise_dim, device=device)
            fake = self.model(z, cond_cont, cond_cat)
            mask = None
        else:
            real, fake_raw, mask, critic_fn = self._stage2_real_and_fake(
                (cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
                stage1_ctx,
                global_step,
                device,
            )
            if self.particle_type_cfg.get("target", "physical") == "onehot":
                # Straight-through Gumbel-softmax relaxation of the type
                # slice only (decision 5) — the critic must see a hard
                # one-hot forward (matching what "real" data looks like)
                # while gradient still flows smoothly to the generator.
                # grad_probe captures the §11.4 gradient-magnitude
                # instrumentation — see _relax_onehot_type_slice's docstring.
                tau = _gumbel_tau(
                    global_step,
                    self.total_steps,
                    self.spec.type_gumbel_tau_start,
                    self.spec.type_gumbel_tau_end,
                )
                fake_raw = _relax_onehot_type_slice(
                    fake_raw,
                    sec_cont.size(1),
                    CONT_SLOT_DIM,
                    stage2_type_dim(self.particle_type_cfg, self.particle_type_emb_dim),
                    tau,
                    grad_probe=grad_probe,
                )
            fake = fake_raw * mask

        # --- critic step (every batch) ---
        fake_detached = fake.detach()
        real_score = critic_fn(real)
        fake_score = critic_fn(fake_detached)
        gp = gradient_penalty(critic_fn, real, fake_detached, mask=mask)
        d_loss = fake_score.mean() - real_score.mean() + self.gp_weight * gp
        wasserstein = (real_score.mean() - fake_score.mean()).detach()

        grad_norm_d = self._step_optimizer(self.optimizer_d, d_loss, self.d_params)

        # --- generator (+ n_sec) step ---
        did_g_step = global_step % self.n_critic == 0
        l_nsec, nsec_acc = self._n_sec_loss(
            cond_cont, cond_cat, stage1_ctx, n_sec, device
        )

        # On a non-generator-step batch with no n_sec_head on this stage
        # (n_sec now defaults to stage 2, decision 1), there's nothing for
        # the generator optimizer to do this batch — g_loss would otherwise
        # be a graph-less zero tensor, which .backward() rejects outright.
        skip_g_step = not did_g_step and self.model.n_sec_head is None
        if did_g_step:
            g_loss_adv = generator_loss(critic_fn, fake)
            g_loss = (
                self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
            )
        else:
            g_loss_adv = torch.zeros((), device=device)
            g_loss = self.spec.n_sec_lambda * l_nsec
        if skip_g_step:
            grad_norm_g = 0.0
        else:
            grad_norm_g = self._step_optimizer(self.optimizer, g_loss, self.g_params)

        if did_g_step:
            self.lr_sched.step()
            if self.ema_model is not None:
                _update_ema(self.ema_model, self.model, self.ema_decay)

        return {
            "d_loss": d_loss.item(),
            "g_loss": g_loss_adv.item(),
            "wasserstein": wasserstein.item(),
            "gp_loss": gp.item(),
            "loss_nsec": l_nsec.item(),
            "nsec_acc": nsec_acc.item(),
            "did_g_step": did_g_step,
            "grad_norm": grad_norm_d + grad_norm_g,
            "grad_norm_d": grad_norm_d,
            "grad_norm_g": grad_norm_g,
            "grad_norm_type_slice": grad_probe.get("type", 0.0),
            "grad_norm_cont_slice": grad_probe.get("cont", 0.0),
            "lr": self.optimizer.param_groups[0]["lr"],
            "critic_lr": self.optimizer_d.param_groups[0]["lr"],
        }

    # --- reporting ------------------------------------------------------

    def batch_loss(self, stats: dict) -> float:
        return stats["d_loss"] + stats["g_loss"]

    def summary(self, means: dict) -> str:
        return (
            f"{self.name}[d={means.get('d_loss', 0.0):.3f} "
            f"g={means.get('g_loss', 0.0):.3f}]"
        )

    def val_objective(
        self, train_means: dict, val_means: dict, marginal_kl: float
    ) -> float:
        """No monotone per-batch WGAN loss fit for averaging, so
        best-checkpoint selection uses the real marginal-KL signal when
        `validate_marginals` produced one, and falls back to this epoch's own
        Wasserstein-distance magnitude otherwise.

        Behavior change vs. every run up to v0.3.0: the pre-refactor code
        meant to do exactly this, but its guard
        (`{n: kl for n in wgan_names if n not in val_loss_per_stage}`) could
        never fire — `val_loss_per_stage` was pre-seeded with `0.0` for every
        stage, so a WGAN stage contributed a flat `0.0` and the marginal KL
        was recorded in the metrics row without ever influencing `best.pt`.
        Runs from before this commit therefore selected their best checkpoint
        on the non-adversarial stages alone."""
        if math.isfinite(marginal_kl):
            return marginal_kl
        return abs(train_means.get("wasserstein", 0.0))

    # --- state ----------------------------------------------------------

    def _extra_state(self) -> dict:
        return {
            "critic": self.critic.state_dict(),
            "optimizer_d": self.optimizer_d.state_dict(),
        }

    def _load_extra_state(self, sd: dict) -> None:
        self.critic.load_state_dict(sd["critic"])
        self.optimizer_d.load_state_dict(sd["optimizer_d"])

    def _resume_extra_lr(self, lr: float) -> None:
        resumed_critic_lr = self.critic_lr if self.critic_lr > 0 else lr
        for group in self.optimizer_d.param_groups:
            group["lr"] = resumed_critic_lr


def build_stage_trainers(
    cfg: dict,
    models: dict[str, torch.nn.Module | None],
    critics: dict[str, torch.nn.Module | None],
    device: torch.device,
    total_train_batches: int,
) -> dict[str, StageTrainer]:
    """One trainer per active stage — `models[name] is None` means that stage
    is `active = false` and is simply never constructed."""
    trainers: dict[str, StageTrainer] = {}
    for name, is_stage2 in (("stage1", False), ("stage2", True)):
        model = models.get(name)
        if model is None:
            continue
        spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1))
        if spec.generator == "wgan":
            critic = critics.get(name)
            assert critic is not None, (
                f"{name}_model.generator='wgan' requires a critic (see "
                "giant.model.network.build_critics)"
            )
            trainers[name] = WGANStageTrainer(spec, model, critic, device)
        else:
            trainers[name] = FlowDDPMStageTrainer(spec, model, device)
    return trainers