Give CriticModel a registry-built trunk and StageModel base (gitea #57)

CriticModel was the one stage-shaped class left out of the trunk-registry
(gitea #33), block-conditioning-registry (gitea #34), and StageModel-base
(gitea #39) refactors: it hand-rolled a plain ResBlock stack, so a
routed/FiLM/AdaLN trunk was available to every generative stage model except
the critic competing against them under WGAN-GP.

CriticModel now subclasses StageModel (reusing its cond_enc construction, and
a stage-2 context-fusion helper factored out of Stage2OneShot onto the base)
and builds its body via build_trunk (output width 1) instead of a bespoke
ResBlock loop, so trunk.type/trunk.block_conditioning now affect the critic
too. Each stage's critic inherits its own generator's trunk config rather
than a new critic_trunk config key, mirroring the existing
critic_hidden_dim/critic_n_res_blocks "0 = inherit from generator" pattern.
Router mixing (MoE) for the critic stays out of scope. Since CriticModel is
training-only and never persisted for inference, and WGAN-GP is still
unbenchmarked, its state_dict shape has no back-compat burden.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

tests/test_network.py

@@ -8,7 +8,10 @@ from giant.model.network import (
     HISTORY_REGISTRY,
     AttentionHistory,
     ConditionEncoder,
+    CriticModel,
+    FilmResBlock,
     HistoryEncoder,
+    LinearTrunk,
     MarkovHistory,
     NoHistory,
     SinusoidalEmbedding,
@@ -944,7 +947,7 @@ def test_build_critics_particle_type_n_classes_overrides_conditioning_emb_dim():
     assert wider_critic is not None
     # k_max=3 slots, each CONT_SLOT_DIM + n_classes wide under wgan folding —
     # widening n_classes alone (emb_dim stays 4) must widen the critic input.
-    assert wider_critic.input_proj.in_features > default_n_classes_critic.input_proj.in_features
+    assert wider_critic.trunk.input_proj.in_features > default_n_classes_critic.trunk.input_proj.in_features
 
 
 # ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─
@@ -1021,12 +1024,12 @@ def test_build_critics_omitted_particle_type_matches_default_config():
     cfg["stage2_model"]["generator"] = "wgan"
     onehot_critic = build_critics(cfg)["stage2"]
     assert onehot_critic is not None
-    onehot_in_dim = onehot_critic.input_proj.in_features
+    onehot_in_dim = onehot_critic.trunk.input_proj.in_features
 
     cfg["stage2_model"]["particle_type"] = {"target": "physical"}
     physical_critic = build_critics(cfg)["stage2"]
     assert physical_critic is not None
-    physical_in_dim = physical_critic.input_proj.in_features
+    physical_in_dim = physical_critic.trunk.input_proj.in_features
 
     # onehot's per-slot type width is emb_dim classes vs. physical's fixed
     # (log-mass, charge) pair — different unless emb_dim happens to be 2, so
@@ -1046,15 +1049,15 @@ def test_build_critics_stage1_critic_hidden_dim_and_n_res_blocks_override_genera
 
     inherited = build_critics(cfg)["stage1"]
     assert inherited is not None
-    assert inherited.input_proj.out_features == 8
-    assert len(inherited.blocks) == 1
+    assert inherited.trunk.input_proj.out_features == 8
+    assert len(inherited.trunk.blocks) == 1
 
     cfg["stage1_model"]["wgan"]["critic_hidden_dim"] = 16
     cfg["stage1_model"]["wgan"]["critic_n_res_blocks"] = 3
     overridden = build_critics(cfg)["stage1"]
     assert overridden is not None
-    assert overridden.input_proj.out_features == 16
-    assert len(overridden.blocks) == 3
+    assert overridden.trunk.input_proj.out_features == 16
+    assert len(overridden.trunk.blocks) == 3
 
 
 def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_generator_size():
@@ -1065,15 +1068,15 @@ def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_genera
 
     inherited = build_critics(cfg)["stage2"]
     assert inherited is not None
-    assert inherited.input_proj.out_features == 8
-    assert len(inherited.blocks) == 1
+    assert inherited.trunk.input_proj.out_features == 8
+    assert len(inherited.trunk.blocks) == 1
 
     cfg["stage2_model"]["wgan"]["critic_hidden_dim"] = 16
     cfg["stage2_model"]["wgan"]["critic_n_res_blocks"] = 3
     overridden = build_critics(cfg)["stage2"]
     assert overridden is not None
-    assert overridden.input_proj.out_features == 16
-    assert len(overridden.blocks) == 3
+    assert overridden.trunk.input_proj.out_features == 16
+    assert len(overridden.trunk.blocks) == 3
 
 
 # ── StageModel base (gitea #39): Stage1Model/Stage2OneShot/Stage2Autoregressive
@@ -1234,3 +1237,34 @@ def test_stagemodel_time_emb_matches_objective_needs_time(cls, generator):
     assert model.generator_kind == generator
     assert model.noise_dim == 8
     assert (model.time_emb is not None) == build_objective(generator).needs_time
+
+
+# ── CriticModel uses the trunk/block registries + StageModel base (gitea #57) ─
+
+
+def test_critic_model_is_stagemodel_subclass():
+    assert issubclass(CriticModel, StageModel)
+
+
+@pytest.mark.parametrize("stage", ["stage1", "stage2"])
+def test_build_critics_threads_trunk_type_from_generator_config(stage):
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage1_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg[f"{stage}_model"]["trunk"] = {"type": "linear"}
+
+    critic = build_critics(cfg)[stage]
+    assert critic is not None
+    assert isinstance(critic.trunk, LinearTrunk)
+
+
+@pytest.mark.parametrize("stage", ["stage1", "stage2"])
+def test_build_critics_threads_block_conditioning_from_generator_config(stage):
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage1_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg[f"{stage}_model"]["trunk"] = {"block_conditioning": "film"}
+
+    critic = build_critics(cfg)[stage]
+    assert critic is not None
+    assert all(isinstance(block, FilmResBlock) for block in critic.trunk.blocks)

tests/test_wgan.py

@@ -2,7 +2,7 @@ import torch
 
 from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM, K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
-from giant.model.network import CriticModel, Stage1Model, Stage2OneShot
+from giant.model.network import CriticModel, LinearTrunk, Stage1Model, Stage2OneShot
 from giant.model.wgan import critic_loss, generator_loss, gradient_penalty
 from giant.sample import sample_secondaries_wgan, sample_wgan
 
@@ -115,6 +115,33 @@ def test_critic_output_shape():
     assert out.shape == (B,)
 
 
+def test_critic_model_honours_trunk_type_and_block_conditioning():
+    """gitea #57: CriticModel routes its body through build_trunk/build_block
+    like every generator stage model, instead of hand-rolling a plain
+    ResBlock stack."""
+    B = 8
+    critic = CriticModel(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        in_dim=X_DIM,
+        hidden_dim=32,
+        n_res_blocks=2,
+        stage="stage1",
+        trunk_type="linear",
+        block_conditioning="adaln",
+    )
+    assert isinstance(critic.trunk, LinearTrunk)
+    cond_cont, cond_cat = _cond(B)
+    real = torch.randn(B, X_DIM)
+    fake = torch.randn(B, X_DIM)
+    loss = critic_loss(lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0)
+    loss.backward()
+    for name, p in critic.named_parameters():
+        assert p.grad is not None, f"no grad for {name}"
+
+
 def test_sample_wgan_shape():
     B = 6
     model = _small_generator()