Merge pull request 'Add an Objective registry for the flow/ddpm/wgan generator choice (gitea #32)' (#51) from fix/issue-32 into master

Reviewed-on: #51

giant/model/builders.py

@@ -15,6 +15,7 @@ from giant.model.models import (
     resolve_type_n_classes,
     stage2_trunk_sec_dim,
 )
+from giant.model.objectives import build_objective
 from giant.model.routers import Router, _build_router_from_cfg
 
 # ---------------------------------------------------------------------------
@@ -64,10 +65,11 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
         if s1_spec.router.enabled:
             stage1_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
         generator = s1_spec.generator
+        objective = build_objective(generator)
         # wgan has no time_dim concept (no diffusion/flow time variable) —
         # matches the pre-dataclass .get("time_dim", 64) fallback, which
         # always hit its default for a wgan sub-block too.
-        time_dim = getattr(s1_spec, generator).time_dim if generator != "wgan" else 64
+        time_dim = getattr(s1_spec, generator).time_dim if objective.needs_time else 64
         n_sec_owner = s2_spec.n_sec.owner
         n_sec_head_k_max = s2_spec.k_max if n_sec_owner == "stage1" else None
         result["stage1"] = Stage1Model(
@@ -99,9 +101,10 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
             else:
                 stage2_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
         generator = s2_spec.generator
+        objective = build_objective(generator)
         # wgan has no time_dim concept — see the matching comment in stage 1
         # above.
-        time_dim = getattr(s2_spec, generator).time_dim if generator != "wgan" else 64
+        time_dim = getattr(s2_spec, generator).time_dim if objective.needs_time else 64
         n_sec_owner = s2_spec.n_sec.owner
         k_max = s2_spec.k_max
         particle_type_cfg = s2_spec.particle_type.to_dict()
@@ -180,7 +183,7 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
 
     result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
 
-    if s1_spec.active and s1_spec.generator == "wgan":
+    if s1_spec.active and build_objective(s1_spec.generator).is_adversarial:
         result["stage1"] = CriticModel(
             pdg_vocab=pdg_vocab,
             mat_vocab=mat_vocab,
@@ -194,11 +197,14 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
             stage="stage1",
         )
 
-    if s2_spec.active and s2_spec.generator == "wgan":
+    if s2_spec.active and build_objective(s2_spec.generator).is_adversarial:
         k_max = s2_spec.k_max
         particle_type_cfg = s2_spec.particle_type.to_dict()
         in_dim = stage2_trunk_sec_dim(
-            particle_type_cfg, "wgan", k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
+            particle_type_cfg,
+            s2_spec.generator,
+            k_max,
+            resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]),
         )
         result["stage2"] = CriticModel(
             pdg_vocab=pdg_vocab,

giant/model/models.py

@@ -8,6 +8,7 @@ from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SE
 from giant.model.encoders import ConditionEncoder
 from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory
 from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding
+from giant.model.objectives import build_objective
 from giant.model.routers import Router
 from giant.model.trunks import build_trunk
 
@@ -49,18 +50,18 @@ def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, em
     `k_max * SEC_SLOT_DIM`, the type slice folded into the same
     flow-matched/WGAN vector as the continuous stick/dir slots.
 
-    `target` in `("onehot", "embedding")`: under `generator == "wgan"` the
-    type slice is still folded in (adversarial for onehot via ST-Gumbel,
-    already-continuous for embedding), just `emb_dim` wide instead of
-    `PARTICLE_PHYS_DIM` wide: `k_max * (CONT_SLOT_DIM + emb_dim)`. Under
-    `generator in ("flow", "ddpm")` the type slice isn't part of this vector
-    at all — it's `Stage2OneShot.type_head`'s job instead — so the trunk
-    only covers `k_max * CONT_SLOT_DIM`.
+    `target` in `("onehot", "embedding")`: under an objective with
+    `folds_type_slice` (currently just wgan) the type slice is still folded
+    in (adversarial for onehot via ST-Gumbel, already-continuous for
+    embedding), just `emb_dim` wide instead of `PARTICLE_PHYS_DIM` wide:
+    `k_max * (CONT_SLOT_DIM + emb_dim)`. Otherwise (flow/ddpm) the type slice
+    isn't part of this vector at all — it's `Stage2OneShot.type_head`'s job
+    instead — so the trunk only covers `k_max * CONT_SLOT_DIM`.
     """
     target = particle_type_cfg.get("target", "physical")
     if target == "physical":
         return k_max * SEC_SLOT_DIM
-    if generator == "wgan":
+    if build_objective(generator).folds_type_slice:
         return k_max * (CONT_SLOT_DIM + emb_dim)
     return k_max * CONT_SLOT_DIM
 
@@ -105,10 +106,11 @@ class Stage1Model(nn.Module):
             if cond_enc is not None
             else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
         )
-        has_time = generator in ("flow", "ddpm")
+        objective = build_objective(generator)
+        has_time = objective.needs_time
         self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
         merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
-        in_dim = noise_dim if generator == "wgan" else x_dim
+        in_dim = objective.trunk_in_dim(x_dim, noise_dim)
         self.trunk = build_trunk(
             router, trunk_type, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout, block_conditioning
         )
@@ -159,11 +161,12 @@ class Stage2OneShot(nn.Module):
     the trunk's own
     flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by
     the caller via `stage2_trunk_sec_dim` — already reflects this). Under
-    `"onehot"`/`"embedding"` with `generator in ("flow", "ddpm")`, the type
+    `"onehot"`/`"embedding"` with an objective (`giant.model.objectives`) that
+    doesn't fold the type slice (flow/ddpm), the type
     slice is predicted by a separate `type_head` instead (same shape pattern
     as `n_sec_head`) — `sec_dim` then covers only the continuous
     stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors.
-    Under `generator == "wgan"` the type slice stays folded into `sec_dim`
+    Under a folding objective (wgan) the type slice stays folded into `sec_dim`
     (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is
     unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation.
 
@@ -213,10 +216,11 @@ class Stage2OneShot(nn.Module):
             nn.Linear(cond_out_dim + context_dim, cond_out_dim),
             nn.SiLU(),
         )
-        has_time = generator in ("flow", "ddpm")
+        objective = build_objective(generator)
+        has_time = objective.needs_time
         self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
         merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
-        in_dim = noise_dim if generator == "wgan" else sec_dim
+        in_dim = objective.trunk_in_dim(sec_dim, noise_dim)
         self.trunk = build_trunk(
             router,
             trunk_type,
@@ -237,7 +241,7 @@ class Stage2OneShot(nn.Module):
             )
         self.type_head = None
         target = self.particle_type_cfg.get("target", "physical")
-        if target != "physical" and generator in ("flow", "ddpm"):
+        if target != "physical" and not objective.folds_type_slice:
             emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
             self.type_head = nn.Sequential(
                 nn.Linear(cond_out_dim, hidden_dim // 2),
@@ -391,11 +395,12 @@ class Stage2Autoregressive(nn.Module):
             nn.SiLU(),
         )
 
-        has_time = generator in ("flow", "ddpm")
+        objective = build_objective(generator)
+        has_time = objective.needs_time
         self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
         merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
         token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim)
-        in_dim = noise_dim if generator == "wgan" else token_dim
+        in_dim = objective.trunk_in_dim(token_dim, noise_dim)
         self.trunk = build_trunk(
             router,
             trunk_type,
@@ -417,7 +422,7 @@ class Stage2Autoregressive(nn.Module):
             )
         self.type_head = None
         target = self.particle_type_cfg.get("target", "physical")
-        if target != "physical" and generator in ("flow", "ddpm"):
+        if target != "physical" and not objective.folds_type_slice:
             self.type_head = nn.Sequential(
                 nn.Linear(cond_out_dim, hidden_dim // 2),
                 nn.SiLU(),

giant/model/network.py

@@ -31,6 +31,15 @@ from giant.model.models import (
     stage2_trunk_sec_dim,
     stage2_type_dim,
 )
+from giant.model.objectives import (
+    OBJECTIVE_REGISTRY,
+    DdpmObjective,
+    FlowObjective,
+    Objective,
+    WganObjective,
+    build_objective,
+    register_objective,
+)
 from giant.model.routers import (
     ROUTER_REGISTRY,
     ComposedRouter,
@@ -64,11 +73,15 @@ __all__ = [
     "ConditionEncoder",
     "ContextAdapter",
     "CriticModel",
+    "DdpmObjective",
     "EnergyRouter",
     "ExpertTrunk",
     "FilmResBlock",
+    "FlowObjective",
     "HistoryEncoder",
     "MarkovHistory",
+    "OBJECTIVE_REGISTRY",
+    "Objective",
     "PdgRouter",
     "ProcessRouter",
     "ROUTER_REGISTRY",
@@ -81,6 +94,7 @@ __all__ = [
     "Stage2OneShot",
     "TRUNK_REGISTRY",
     "Trunk",
+    "WganObjective",
     "_CausalAttnBlock",
     "_build_router_from_cfg",
     "_check_router_conditioning_compat",
@@ -93,11 +107,13 @@ __all__ = [
     "build_critics",
     "build_expert_body",
     "build_models",
+    "build_objective",
     "build_router",
     "build_trunk",
     "cat_col_layout",
     "migrate_legacy_state_dict",
     "register_block",
+    "register_objective",
     "register_router",
     "register_trunk",
     "resolve_type_n_classes",

giant/model/objectives.py

@@ -0,0 +1,202 @@
+"""Generative objectives (flow/ddpm/wgan): `Objective` base + registry,
+mirroring `giant.model.routers`'s `Router` pattern (gitea #32). Each objective
+answers, in one place, the handful of questions every stage model/sampler/
+trainer used to re-derive independently from a bare `generator` string: does
+this stage need a time embedding, is it adversarial, does it fold the
+secondary type slice into its own trunk output, what does the trunk take as
+input, which stage-1/stage-2 loss does it train against.
+
+Self-contained (no dependency on `giant.model.models`, unlike `Router` which
+`giant.model.trunks` depends on) — `Objective` never needs to construct a
+stage model or critic itself, only describe one. This also sidesteps a
+`models.py` <-> `objectives.py` import cycle, since `models.py` calls
+`build_objective`.
+"""
+
+import inspect
+
+import torch
+
+from giant.model.schedule import (
+    CosineSchedule,
+    flow_matching_loss,
+    flow_matching_loss_secondary,
+    flow_matching_loss_secondary_ar,
+)
+
+# ---------------------------------------------------------------------------
+# Objective contract
+# ---------------------------------------------------------------------------
+
+
+class Objective:
+    """Contract for a pluggable generative objective. Not an `nn.Module` —
+    unlike `Router`, no objective owns learnable parameters, so a plain
+    strategy object is the honest fit.
+
+    `needs_time`/`is_adversarial`/`folds_type_slice`/`supports_stage2_decoder`
+    are set by each concrete subclass (no defaults here — a new objective
+    should have to state all four, not silently inherit one that happens to
+    be wrong for it). See `FlowObjective`/`DdpmObjective`/`WganObjective`.
+    """
+
+    needs_time: bool
+    is_adversarial: bool
+    folds_type_slice: bool
+    supports_stage2_decoder: bool = True
+
+    def trunk_in_dim(self, out_dim: int, noise_dim: int) -> int:
+        """Width of the trunk's own input — `out_dim` (denoising/flow-matching
+        a same-shape vector) for every non-adversarial objective;
+        `WganObjective` overrides to `noise_dim` (a single-pass noise-to-output
+        generator)."""
+        return out_dim
+
+    def build_schedule(self, n_steps: int, device: torch.device) -> CosineSchedule | None:
+        """Objective-owned auxiliary state a stage trainer must build once
+        and hold onto (device-placed) across its training loop. `None` for
+        every objective except `DdpmObjective` (its noise schedule)."""
+        return None
+
+    def stage1_loss(
+        self,
+        model: torch.nn.Module,
+        x1: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        *,
+        schedule: object | None = None,
+    ) -> torch.Tensor:
+        """Stage-1 training loss. Only implemented by non-adversarial
+        objectives — `WganObjective` is unused here, `WGANStageTrainer` has
+        its own G/D step instead."""
+        raise NotImplementedError(f"{type(self).__name__} has no stage1_loss")
+
+    def stage2_loss(
+        self,
+        model: torch.nn.Module,
+        x1_s2: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_ctx: torch.Tensor,
+        sec_mask: torch.Tensor,
+        *,
+        type_dim: int | None,
+        ar_inputs: dict[str, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        """Stage-2 secondary-decoder training loss, one-shot or
+        autoregressive depending on whether `ar_inputs` is given. Same
+        adversarial caveat as `stage1_loss`."""
+        raise NotImplementedError(f"{type(self).__name__} has no stage2_loss")
+
+
+OBJECTIVE_REGISTRY: dict[str, type[Objective]] = {}
+
+
+def register_objective(name: str):
+    def decorator(cls: type[Objective]) -> type[Objective]:
+        OBJECTIVE_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_objective(name: str, **kwargs) -> Objective:
+    """Factory: look up an `Objective` subclass by name (a `generator`
+    config value) from the registry.
+
+    Every registered objective is fed the same kwargs; kwargs not declared by
+    that type's constructor are silently dropped, so per-type hyperparameters
+    (e.g. `DdpmObjective`'s `n_steps`) can coexist in one call without
+    special-casing — same convention as `giant.model.routers.build_router`.
+    """
+    if name not in OBJECTIVE_REGISTRY:
+        raise ValueError(f"unknown generator/objective {name!r}; available: {sorted(OBJECTIVE_REGISTRY)}")
+    cls = OBJECTIVE_REGISTRY[name]
+    accepted = set(inspect.signature(cls.__init__).parameters) - {"self"}
+    filtered = {k: v for k, v in kwargs.items() if k in accepted}
+    return cls(**filtered)
+
+
+# ---------------------------------------------------------------------------
+# Concrete objectives
+# ---------------------------------------------------------------------------
+
+
+@register_objective("flow")
+class FlowObjective(Objective):
+    """Conditional flow matching (Lipman et al. 2022) — the primary
+    objective. ~10 ODE steps at inference (`giant.sample.sample_flow`)."""
+
+    needs_time = True
+    is_adversarial = False
+    folds_type_slice = False
+
+    def stage1_loss(self, model, x1, cond_cont, cond_cat, *, schedule=None) -> torch.Tensor:
+        return flow_matching_loss(model, x1, cond_cont, cond_cat)
+
+    def stage2_loss(
+        self,
+        model,
+        x1_s2,
+        cond_cont,
+        cond_cat,
+        stage1_ctx,
+        sec_mask,
+        *,
+        type_dim=None,
+        ar_inputs=None,
+    ) -> torch.Tensor:
+        if ar_inputs is not None:
+            return flow_matching_loss_secondary_ar(
+                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=type_dim,
+            )
+        return flow_matching_loss_secondary(model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=type_dim)
+
+
+@register_objective("ddpm")
+class DdpmObjective(Objective):
+    """Full DDPM ancestral sampling (Nichol & Dhariwal 2021 cosine schedule)
+    — the throwaway baseline. Stage-1 only: no `Stage2*` class has ever been
+    trained with `generator="ddpm"` in practice, so there's no stage-2 ddpm
+    loss to dispatch to (matches `FlowDDPMStageTrainer`'s pre-existing
+    stage-2 guard)."""
+
+    needs_time = True
+    is_adversarial = False
+    folds_type_slice = False
+    supports_stage2_decoder = False
+
+    def __init__(self, n_steps: int = 1000) -> None:
+        self.n_steps = n_steps
+
+    def build_schedule(self, n_steps: int, device: torch.device) -> CosineSchedule:
+        return CosineSchedule(T=n_steps).to(device)
+
+    def stage1_loss(self, model, x1, cond_cont, cond_cat, *, schedule=None) -> torch.Tensor:
+        assert schedule is not None, "DdpmObjective.stage1_loss needs a schedule (see build_schedule)"
+        return schedule.loss(model, x1, cond_cont, cond_cat)
+
+
+@register_objective("wgan")
+class WganObjective(Objective):
+    """WGAN-GP (Gulrajani et al. 2017) — single forward pass instead of an
+    ODE loop. `stage1_loss`/`stage2_loss` are unused: `WGANStageTrainer` owns
+    its own dual generator/critic step instead of a single scalar loss."""
+
+    needs_time = False
+    is_adversarial = True
+    folds_type_slice = True
+
+    def trunk_in_dim(self, out_dim: int, noise_dim: int) -> int:
+        return noise_dim

giant/sample.py

@@ -2,7 +2,7 @@ import torch
 import torch.nn.functional as F
 
 from giant.constants import CONT_SLOT_DIM, X_DIM
-from giant.model.network import Stage2Autoregressive, stage2_trunk_sec_dim
+from giant.model.network import DdpmObjective, Stage2Autoregressive, build_objective, stage2_trunk_sec_dim
 from giant.model.schedule import CosineSchedule
 
 
@@ -132,7 +132,7 @@ def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int:
 
 def _type_folded(sec_decoder: torch.nn.Module) -> bool:
     target = sec_decoder.particle_type_cfg.get("target", "physical")
-    return target == "physical" or sec_decoder.generator_kind == "wgan"
+    return target == "physical" or build_objective(sec_decoder.generator_kind).folds_type_slice
 
 
 def _decode_stage2_flat(
@@ -281,7 +281,7 @@ def sample_secondaries_ar(
     device = cond_cont.device
     k_max = sec_decoder.k_max
     type_dim = sec_decoder.type_dim
-    generator = sec_decoder.generator_kind
+    objective = build_objective(sec_decoder.generator_kind)
     target = sec_decoder.particle_type_cfg.get("target", "physical")
     type_folded = _type_folded(sec_decoder)
     token_dim = CONT_SLOT_DIM + type_dim if type_folded else CONT_SLOT_DIM
@@ -301,7 +301,7 @@ def sample_secondaries_ar(
         slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
         hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
 
-        if generator == "wgan":
+        if objective.is_adversarial:
             z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
             token = sec_decoder(
                 z,
@@ -383,10 +383,10 @@ def sample_stage1(
     ddpm_steps: int = 1000,
 ) -> tuple[torch.Tensor, torch.Tensor | None]:
     """Dispatches on `stage1_model.generator_kind`."""
-    kind = stage1_model.generator_kind
-    if kind == "wgan":
+    objective = build_objective(stage1_model.generator_kind)
+    if objective.is_adversarial:
         return sample_wgan(stage1_model, cond_cont, cond_cat)
-    if kind == "ddpm":
+    if isinstance(objective, DdpmObjective):
         schedule = CosineSchedule(T=ddpm_steps).to(cond_cont.device)
         return sample_ddpm(stage1_model, cond_cont, cond_cat, schedule)
     return sample_flow(stage1_model, cond_cont, cond_cat, steps=steps)
@@ -409,7 +409,7 @@ def sample_stage2(
     """
     if isinstance(sec_decoder, Stage2Autoregressive):
         return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
-    if sec_decoder.generator_kind == "wgan":
+    if build_objective(sec_decoder.generator_kind).is_adversarial:
         return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
     return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
 

giant/training/stage2_inputs.py

@@ -12,6 +12,7 @@ import torch
 import torch.nn.functional as F
 
 from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
+from giant.model.objectives import build_objective
 from giant.sample import sample_secondaries_ar
 
 
@@ -69,19 +70,20 @@ def _assemble_stage2_ar_target(
 
     - `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit,
       dir, log_mass, charge) as-is.
-    - `target` in `("onehot", "embedding")` + `generator in ("flow", "ddpm")`:
-      just the continuous stick/dir slots — the type slice isn't part of
-      this tensor at all (`type_head` handles it separately).
-    - `target` in `("onehot", "embedding")` + `generator == "wgan"`: stick/dir
-      slots concatenated with the per-slot type representation (a one-hot of
-      the true class, relaxed on the *generated* side only, by the caller;
-      or the conditioning's own detached embedding-table row).
+    - `target` in `("onehot", "embedding")` + an objective that doesn't fold
+      the type slice (flow/ddpm): just the continuous stick/dir slots — the
+      type slice isn't part of this tensor at all (`type_head` handles it
+      separately).
+    - `target` in `("onehot", "embedding")` + a folding objective (wgan):
+      stick/dir slots concatenated with the per-slot type representation (a
+      one-hot of the true class, relaxed on the *generated* side only, by the
+      caller; or the conditioning's own detached embedding-table row).
     """
     target = particle_type_cfg.get("target", "physical")
     if target == "physical":
         return sec_cont
     cont = sec_cont[..., :CONT_SLOT_DIM]
-    if generator != "wgan":
+    if not build_objective(generator).folds_type_slice:
         return cont
     type_repr = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
     return torch.cat([cont, type_repr], dim=-1)

giant/training/trainers.py

@@ -25,13 +25,7 @@ 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, resolve_type_n_classes, stage2_type_dim
-from giant.model.schedule import (
-    CosineSchedule,
-    flow_matching_loss,
-    flow_matching_loss_secondary,
-    flow_matching_loss_secondary_ar,
-)
+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 (
@@ -450,19 +444,22 @@ 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",):
+        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.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).
+        # 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.get("target", "physical") == "physical" else 0
 
         self.params = list(self.model.parameters())
@@ -472,7 +469,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             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.ddpm_schedule = self.objective.build_schedule(spec.ddpm_n_steps, device)
 
         self.train_metrics = [
             train_metric(key)
@@ -504,26 +501,9 @@ class FlowDDPMStageTrainer(StageTrainer):
 
     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(
+            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,
@@ -531,6 +511,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             stage1_ctx,
             sec_mask,
             type_dim=self._flow_type_dim,
+            ar_inputs=ar_inputs,
         )
 
     def _type_loss(
@@ -774,7 +755,7 @@ class WGANStageTrainer(StageTrainer):
         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
+            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,
@@ -787,7 +768,7 @@ class WGANStageTrainer(StageTrainer):
                 ar["slot_idx"],
             ).reshape(B, -1)
         else:
-            real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=True) * mask
+            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)
 
@@ -957,7 +938,7 @@ def build_stage_trainers(
         if model is None:
             continue
         spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1))
-        if spec.generator == "wgan":
+        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)"

tests/test_objectives.py

@@ -0,0 +1,250 @@
+"""Tests for `giant/model/objectives.py` — the generator/objective registry
+(gitea #32) that replaced bare `generator in ("flow", "ddpm", "wgan")`
+string checks scattered across models.py/sample.py/builders.py/
+stage2_inputs.py/trainers.py."""
+
+import pytest
+import torch
+
+from giant.constants import COND_DIM, CONT_SLOT_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
+from giant.model.network import (
+    OBJECTIVE_REGISTRY,
+    DdpmObjective,
+    FlowObjective,
+    Stage1Model,
+    Stage2Autoregressive,
+    Stage2OneShot,
+    WganObjective,
+    build_objective,
+)
+from giant.model.schedule import CosineSchedule, flow_matching_loss, flow_matching_loss_secondary
+
+_PHYS_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
+
+
+def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
+    cond_cont = torch.randn(B, COND_DIM)
+    cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
+    return cond_cont, cond_cat
+
+
+# ── registry ─────────────────────────────────────────────────────────────
+
+
+def test_registry_has_exactly_the_three_known_objectives():
+    assert set(OBJECTIVE_REGISTRY) == {"flow", "ddpm", "wgan"}
+
+
+def test_build_objective_returns_correct_concrete_type():
+    assert isinstance(build_objective("flow"), FlowObjective)
+    assert isinstance(build_objective("ddpm"), DdpmObjective)
+    assert isinstance(build_objective("wgan"), WganObjective)
+
+
+def test_build_objective_unknown_name_raises():
+    with pytest.raises(ValueError, match="unknown generator/objective"):
+        build_objective("bogus")
+
+
+def test_build_objective_filters_kwargs_by_signature():
+    # FlowObjective takes no constructor args — n_steps (a DdpmObjective-only
+    # kwarg) must be silently dropped, not raise a TypeError.
+    build_objective("flow", n_steps=500)
+    ddpm = build_objective("ddpm", n_steps=250)
+    assert isinstance(ddpm, DdpmObjective)
+    assert ddpm.n_steps == 250
+
+
+# ── flags ────────────────────────────────────────────────────────────────
+
+
+def test_flow_objective_flags():
+    obj = build_objective("flow")
+    assert obj.needs_time is True
+    assert obj.is_adversarial is False
+    assert obj.folds_type_slice is False
+    assert obj.supports_stage2_decoder is True
+
+
+def test_ddpm_objective_flags():
+    obj = build_objective("ddpm")
+    assert obj.needs_time is True
+    assert obj.is_adversarial is False
+    assert obj.folds_type_slice is False
+    assert obj.supports_stage2_decoder is False
+
+
+def test_wgan_objective_flags():
+    obj = build_objective("wgan")
+    assert obj.needs_time is False
+    assert obj.is_adversarial is True
+    assert obj.folds_type_slice is True
+    assert obj.supports_stage2_decoder is True
+
+
+# ── trunk_in_dim ─────────────────────────────────────────────────────────
+
+
+def test_trunk_in_dim_flow_and_ddpm_pass_through_out_dim():
+    assert build_objective("flow").trunk_in_dim(out_dim=9, noise_dim=8) == 9
+    assert build_objective("ddpm").trunk_in_dim(out_dim=9, noise_dim=8) == 9
+
+
+def test_trunk_in_dim_wgan_uses_noise_dim():
+    assert build_objective("wgan").trunk_in_dim(out_dim=9, noise_dim=8) == 8
+
+
+# ── ddpm schedule ────────────────────────────────────────────────────────
+
+
+def test_ddpm_build_schedule_has_requested_length():
+    schedule = build_objective("ddpm").build_schedule(n_steps=17, device=torch.device("cpu"))
+    assert isinstance(schedule, CosineSchedule)
+    assert schedule.T == 17
+
+
+def test_flow_and_wgan_build_schedule_is_none():
+    assert build_objective("flow").build_schedule(100, torch.device("cpu")) is None
+    assert build_objective("wgan").build_schedule(100, torch.device("cpu")) is None
+
+
+# ── stage1_loss parity ──────────────────────────────────────────────────
+
+
+def test_flow_objective_stage1_loss_matches_direct_call():
+    torch.manual_seed(0)
+    model = Stage1Model(
+        pdg_vocab=3, mat_vocab=2, particle_cfg=_PHYS_CFG, material_cfg=_PHYS_CFG, hidden_dim=16, n_res_blocks=1
+    )
+    cond_cont, cond_cat = _cond(4)
+    x1 = torch.randn(4, X_DIM)
+
+    torch.manual_seed(1)
+    expected = flow_matching_loss(model, x1, cond_cont, cond_cat)
+    torch.manual_seed(1)
+    actual = build_objective("flow").stage1_loss(model, x1, cond_cont, cond_cat)
+    assert torch.allclose(actual, expected)
+
+
+def test_ddpm_objective_stage1_loss_matches_direct_call():
+    torch.manual_seed(0)
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="ddpm",
+    )
+    cond_cont, cond_cat = _cond(4)
+    x1 = torch.randn(4, X_DIM)
+    objective = build_objective("ddpm", n_steps=50)
+    schedule = objective.build_schedule(50, torch.device("cpu"))
+    assert isinstance(schedule, CosineSchedule)
+
+    torch.manual_seed(1)
+    expected = schedule.loss(model, x1, cond_cont, cond_cat)
+    torch.manual_seed(1)
+    actual = objective.stage1_loss(model, x1, cond_cont, cond_cat, schedule=schedule)
+    assert torch.allclose(actual, expected)
+
+
+def test_ddpm_objective_stage1_loss_requires_a_schedule():
+    model = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="ddpm",
+    )
+    cond_cont, cond_cat = _cond(4)
+    with pytest.raises(AssertionError):
+        build_objective("ddpm").stage1_loss(model, torch.randn(4, X_DIM), cond_cont, cond_cat, schedule=None)
+
+
+# ── stage2_loss dispatch ─────────────────────────────────────────────────
+
+
+def test_flow_objective_stage2_loss_one_shot_matches_direct_call():
+    torch.manual_seed(0)
+    B, k_max = 4, 5
+    sec_dim = k_max * SEC_SLOT_DIM
+    model = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="flow",
+        sec_dim=sec_dim,
+        k_max=k_max,
+    )
+    cond_cont, cond_cat = _cond(B)
+    stage1_ctx = torch.randn(B, X_DIM)
+    x1_s2 = torch.randn(B, sec_dim)
+    sec_mask = torch.ones(B, k_max, dtype=torch.bool)
+
+    torch.manual_seed(1)
+    expected = flow_matching_loss_secondary(model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None)
+    torch.manual_seed(1)
+    actual = build_objective("flow").stage2_loss(
+        model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None, ar_inputs=None
+    )
+    assert torch.allclose(actual, expected)
+
+
+def test_flow_objective_stage2_loss_dispatches_to_ar_when_ar_inputs_given():
+    torch.manual_seed(0)
+    B, k_max = 4, 5
+    model = Stage2Autoregressive(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=_PHYS_CFG,
+        material_cfg=_PHYS_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        generator="flow",
+        k_max=k_max,
+    )
+    cond_cont, cond_cat = _cond(B)
+    stage1_ctx = torch.randn(B, X_DIM)
+    token_dim = CONT_SLOT_DIM + PARTICLE_PHYS_DIM
+    x1_s2 = torch.randn(B, k_max, token_dim)
+    sec_mask = torch.ones(B, k_max, dtype=torch.bool)
+    ar_inputs = {
+        "history_feat": torch.randn(B, k_max, token_dim),
+        "has_prev": torch.ones(B, k_max, dtype=torch.bool),
+        "remaining_frac": torch.rand(B, k_max),
+        "slot_idx": torch.linspace(0, 1, k_max).unsqueeze(0).expand(B, -1),
+    }
+
+    loss = build_objective("flow").stage2_loss(
+        model, x1_s2, cond_cont, cond_cat, stage1_ctx, sec_mask, type_dim=None, ar_inputs=ar_inputs
+    )
+    assert loss.dim() == 0
+    assert torch.isfinite(loss)
+
+
+def test_ddpm_objective_stage2_loss_not_implemented():
+    dummy_model = torch.nn.Module()
+    dummy = torch.zeros(1)
+    with pytest.raises(NotImplementedError):
+        build_objective("ddpm").stage2_loss(
+            dummy_model, dummy, dummy, dummy, dummy, torch.ones(1, 1, dtype=torch.bool), type_dim=None
+        )
+
+
+def test_wgan_objective_has_no_loss_methods():
+    dummy_model = torch.nn.Module()
+    dummy = torch.zeros(1)
+    objective = build_objective("wgan")
+    with pytest.raises(NotImplementedError):
+        objective.stage1_loss(dummy_model, dummy, dummy, dummy)
+    with pytest.raises(NotImplementedError):
+        objective.stage2_loss(
+            dummy_model, dummy, dummy, dummy, dummy, torch.ones(1, 1, dtype=torch.bool), type_dim=None
+        )