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
from typing import NamedTuple

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

from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig
from giant.constants import CONT_SLOT_DIM
from giant.data.dataset import StepBatch
from giant.model.network import Router, build_objective, resolve_type_n_classes, stage2_type_dim
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,
    _stop_target_and_mask,
)


@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: StepBatch, device: torch.device) -> StepBatch:
    return type(batch)(*(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
    n_sec_mode: str = "head"

    # particle-type target (stage 2 only)
    particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
    particle_type_n_classes: 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 = TrainConfig.from_dict(cfg["train"])
        # n_sec/particle_type/decoder/autoregressive/wgan's gumbel_tau_* are
        # stage-2-only concepts, always read off s2_spec (guarded by
        # is_stage2 where the stage-1 StageSpec needs a different value) —
        # historically n_sec/particle_type were read from stage2_model
        # unconditionally even for the stage-1 StageSpec, preserved here for
        # behavioral parity. stage_spec covers the fields both stage configs
        # share structurally (generator, lambda, router, ddpm, and wgan's
        # base fields — Stage2ModelConfig's sub-configs all subclass
        # stage 1's).
        s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
        stage_spec = s2_spec if is_stage2 else Stage1ModelConfig.from_dict(cfg["stage1_model"])
        return cls(
            name=name,
            is_stage2=is_stage2,
            generator=stage_spec.generator,
            decoder=s2_spec.decoder if is_stage2 else "one_shot",
            lambda_weight=stage_spec.lambda_weight,
            n_sec_lambda=s2_spec.n_sec.lambda_weight,
            n_sec_mode=s2_spec.n_sec.mode,
            particle_type=s2_spec.particle_type,
            particle_type_n_classes=resolve_type_n_classes(
                s2_spec.particle_type, cfg["conditioning"]["particle"]["emb_dim"]
            ),
            # train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge
            # (giant/config.py), so TrainConfig.from_dict never has to fall
            # back to a literal here; 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=stage_spec.router.lambda_balance,
            lambda_proc=stage_spec.router.lambda_proc,
            lambda_entropy=stage_spec.router.lambda_entropy,
            gumbel_tau_start=stage_spec.router.gumbel_tau_start,
            gumbel_tau_end=stage_spec.router.gumbel_tau_end,
            teacher_forcing=s2_spec.autoregressive.teacher_forcing if is_stage2 else cls.teacher_forcing,
            tf_p_start=s2_spec.autoregressive.tf_p_start if is_stage2 else cls.tf_p_start,
            tf_p_end=s2_spec.autoregressive.tf_p_end if is_stage2 else cls.tf_p_end,
            # 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 (the autoregressive config lists
            # tf_p_start/tf_p_end/attn_n_heads/attn_n_layers only).
            ar_sample_steps=t.validate_steps,
            ddpm_n_steps=stage_spec.ddpm.n_steps,
            n_critic=stage_spec.wgan.n_critic,
            gp_weight=stage_spec.wgan.gp_weight,
            critic_lr=stage_spec.wgan.critic_lr,
            type_gumbel_tau_start=s2_spec.wgan.gumbel_tau_start if is_stage2 else cls.type_gumbel_tau_start,
            type_gumbel_tau_end=s2_spec.wgan.gumbel_tau_end if is_stage2 else cls.type_gumbel_tau_end,
        )


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

    Reads only the shared `StepBatch` (`giant.data.dataset`) — stage 2 always
    conditions on the ground-truth `x1_s1` (`stage2_model.stage1_context =
    "truth"`, stage-level teacher forcing; `"sampled"` is not implemented),
    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" 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 = spec.particle_type
        self.particle_type_n_classes = spec.particle_type_n_classes
        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: StepBatch, device: torch.device, global_step: int) -> dict:
        raise NotImplementedError

    def val_loss(self, batch: StepBatch, 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_n_classes,
            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_n_classes,
        )
        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 `mode =
        "stop_token"` model carries no head at all (see `_stop_loss` for its
        EOS signal instead), so this correctly stays zero for it.
        """
        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

    def _stop_loss(
        self,
        cond_cont: torch.Tensor,
        cond_cat: torch.Tensor,
        stage1_ctx: torch.Tensor,
        n_sec: torch.Tensor,
        device: torch.device,
        ar_inputs: dict[str, torch.Tensor] | None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """`(l_stop, stop_acc)` for `n_sec.mode = "stop_token"`'s per-slot EOS
        head (`Stage2Autoregressive.predict_stop`) — zeros when this stage
        owns no `stop_head` (every other `n_sec.mode`), the same gating
        convention `_n_sec_loss` uses for `n_sec_head`. The two heads are
        mutually exclusive (`giant.model.builders`), so exactly one of
        `_n_sec_loss`/`_stop_loss` is ever non-zero for a given stage.

        Masked BCE against `_stop_target_and_mask`'s per-slot target — one
        slot wider than `sec_mask` (the stop slot itself, `k == n_sec`, needs
        supervision even though it holds no real secondary)."""
        stop_head = getattr(self.model, "stop_head", None)
        if stop_head is None:
            zero = torch.zeros((), device=device)
            return zero, zero
        assert ar_inputs is not None
        logits = self.model.predict_stop(
            cond_cont,
            cond_cat,
            stage1_ctx,
            ar_inputs["history_feat"],
            ar_inputs["has_prev"],
            ar_inputs["remaining_frac"],
            ar_inputs["slot_idx"],
        )
        target, mask = _stop_target_and_mask(n_sec, logits.size(1), device)
        mask_f = mask.float()
        denom = mask_f.sum().clamp(min=1)
        bce = F.binary_cross_entropy_with_logits(logits, target, reduction="none")
        l_stop = (bce * mask_f).sum() / denom
        stop_acc = (((logits >= 0).float() == target).float() * mask_f).sum() / denom
        return l_stop, stop_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:
        objective = build_objective(spec.generator, n_steps=spec.ddpm_n_steps)
        if spec.is_stage2 and not objective.supports_stage2_decoder:
            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)"
            )
        super().__init__(spec, model, device)
        self.objective = objective
        self.particle_type_lambda = self.particle_type_cfg.lambda_weight
        # Width of the type slice actually folded into x1_s2 by _sec_target,
        # under this trainer's objective (flow/ddpm only — see the
        # NotImplementedError above, neither folds the type slice): "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.target == "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 = self.objective.build_schedule(spec.ddpm_n_steps, device)

        self.train_metrics = [
            train_metric(key)
            for key in (
                "loss",
                "loss_gen",
                "loss_nsec",
                "loss_stop",
                "loss_balance",
                "loss_proc",
                "loss_entropy",
                "nsec_acc",
                "stop_acc",
                "loss_type",
                "type_acc",
                "grad_norm",
            )
        ]
        self.val_metrics = [
            val_metric(key)
            for key in (
                "loss",
                "loss_gen",
                "loss_nsec",
                "nsec_acc",
                "loss_stop",
                "stop_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:
            return self.objective.stage1_loss(self.model, x1_s1, cond_cont, cond_cat, schedule=self.ddpm_schedule)
        assert self.decoder != "autoregressive" or ar_inputs is not None
        return self.objective.stage2_loss(
            self.model,
            x1_s2,
            cond_cont,
            cond_cat,
            stage1_ctx,
            sec_mask,
            type_dim=self._flow_type_dim,
            ar_inputs=ar_inputs,
        )

    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.
        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.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: StepBatch, 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_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs)

        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:
            if self.spec.lambda_balance > 0:
                l_balance = self.router.balance_loss(cond_cont, cond_cat)
            if self.spec.lambda_proc > 0:
                l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx)
            if self.spec.lambda_entropy > 0:
                l_entropy = self.router.entropy_loss(cond_cont, cond_cat)

        total = (
            self.spec.lambda_weight * l_gen
            + self.spec.n_sec_lambda * (l_nsec + l_stop)
            + 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_stop": l_stop,
            "loss_type": l_type,
            "type_acc": type_acc,
            "loss_balance": l_balance,
            "loss_proc": l_proc,
            "loss_entropy": l_entropy,
            "nsec_acc": nsec_acc,
            "stop_acc": stop_acc,
        }

    def step(self, batch: StepBatch, 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: StepBatch, 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 _Stage2RealFakeBatch(NamedTuple):
    """Subset of `StepBatch` that `_stage2_real_and_fake` needs."""

    cond_cont: torch.Tensor
    cond_cat: torch.Tensor
    n_sec: torch.Tensor
    sec_cont: torch.Tensor
    sec_type_idx: torch.Tensor


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",
            "loss_stop",
            "stop_acc",
            "grad_norm_d",
            "grad_norm_g",
        ]
        if self.is_stage2 and self.particle_type_cfg.target == "onehot":
            # 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: _Stage2RealFakeBatch, stage1_ctx, global_step, device):
        """Build `(real, fake_raw, mask, critic_fn, ar_inputs)` 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. `ar_inputs` is `None` under
        `decoder = "one_shot"`; under `"autoregressive"` it's the same dict
        `_ar_inputs` built to condition `self.model` above — returned so the
        caller's `_stop_loss` reuses it instead of paying for a second
        (possibly self-sampling) `_ar_inputs` call."""
        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_n_classes)
        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)

        ar_inputs = None
        if self.decoder == "autoregressive":
            epoch = global_step // self.spec.steps_per_epoch
            ar_inputs = 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, self.generator, 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_inputs["history_feat"],
                ar_inputs["has_prev"],
                ar_inputs["remaining_frac"],
                ar_inputs["slot_idx"],
            ).reshape(B, -1)
        else:
            real = self._sec_target(sec_cont, sec_type_idx, self.generator, 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, ar_inputs

    def step(self, batch: StepBatch, 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] = {}

        ar_inputs = None
        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, ar_inputs = self._stage2_real_and_fake(
                _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
                stage1_ctx,
                global_step,
                device,
            )
            if self.particle_type_cfg.target == "onehot":
                # Straight-through Gumbel-softmax relaxation of the type
                # slice only — 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 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_n_classes),
                    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)
        l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs)

        # On a non-generator-step batch with no n_sec_head/stop_head on this
        # stage (n_sec now defaults to stage 2), 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 and self.model.stop_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 + l_stop)
        else:
            g_loss_adv = torch.zeros((), device=device)
            g_loss = self.spec.n_sec_lambda * (l_nsec + l_stop)
        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(),
            "loss_stop": l_stop.item(),
            "stop_acc": stop_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} 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 build_objective(spec.generator).is_adversarial:
            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