router: seed EnergyRouter centers from data quantiles instead of a fixed linspace

The 2026-07-22 rollout benchmark's router_gating diagnostic showed the
10-expert EnergyRouter's default linspace(-2, 2, n_experts) init assumes a
roughly uniform z-normalized energy distribution, leaving experts heavily
overlapping instead of partitioning the range. Add an optional
centers_init kwarg (backward compatible, defaults to the old linspace) and
have giant train estimate it from a reservoir sample of the real energy
column, collected during the existing normalizer-fitting pass.

giant/data/transforms.py

@@ -228,6 +228,54 @@ class _WelfordAccumulator:
         return norm
 
 
+class _ReservoirSampler:
+    """Uniform random sample of a fixed capacity drawn from a data stream.
+
+    Algorithm R (Vitter 1985), vectorized per chunk so it stays cheap over
+    hundreds of millions of rows: use to get a representative subsample of
+    a column for a distribution estimate (e.g. quantiles) without
+    materializing the full column.
+
+        sampler = _ReservoirSampler(capacity=100_000)
+        for chunk in data:
+            sampler.update(chunk)
+        sample = sampler.sample
+    """
+
+    def __init__(self, capacity: int, seed: int = 0) -> None:
+        self.capacity = capacity
+        self.n_seen = 0
+        self._rng = np.random.default_rng(seed)
+        self._reservoir = np.empty(0, dtype=np.float64)
+
+    def update(self, values: np.ndarray) -> None:
+        values = np.asarray(values, dtype=np.float64).reshape(-1)
+        if values.size == 0:
+            return
+        n_before = self.n_seen
+        if n_before < self.capacity:
+            take = min(values.size, self.capacity - n_before)
+            self._reservoir = np.concatenate([self._reservoir, values[:take]])
+            values = values[take:]
+            n_before += take
+        self.n_seen = n_before + values.size
+        if values.size == 0 or self.capacity == 0:
+            return
+        # remaining elements are past the fill phase: element at 1-based
+        # stream position j replaces a uniformly random reservoir slot with
+        # probability capacity/j, which yields a uniform sample overall.
+        positions = n_before + np.arange(1, values.size + 1)
+        accept = self._rng.random(values.size) < (self.capacity / positions)
+        accept_idx = np.nonzero(accept)[0]
+        if accept_idx.size > 0:
+            slots = self._rng.integers(0, self.capacity, size=accept_idx.size)
+            self._reservoir[slots] = values[accept_idx]
+
+    @property
+    def sample(self) -> np.ndarray:
+        return self._reservoir.astype(np.float32)
+
+
 def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
     """World-frame unit vector pointing from pre_pos to post_pos.
 

giant/model/network.py

@@ -1,6 +1,7 @@
 import inspect
 import math
 import re
+from collections.abc import Sequence
 
 import torch
 import torch.nn as nn
@@ -600,10 +601,14 @@ class EnergyRouter(Router):
     """Soft turn-on gate over normalized pre-step log-energy.
 
     Reads `cond_cont[:, energy_idx]` (ignores cond_cat). Learnable (or
-    fixed) 1-D centers, initialized spread across [-2, 2] — roughly the
-    z-normalized energy range. `gate(e) = softmax_i(-(e - c_i)^2 / tau)`,
-    differentiable in e; as tau -> 0 this hardens to nearest-center
-    (Voronoi) selection, which is exactly what `top1` uses at eval.
+    fixed) 1-D centers. By default initialized spread evenly across
+    [-2, 2] — an assumed-uniform z-normalized energy range that may not
+    match the true (often skewed) distribution and can leave experts
+    overlapping instead of partitioning the range; pass `centers_init` to
+    seed them from data (e.g. energy quantiles) instead.
+    `gate(e) = softmax_i(-(e - c_i)^2 / tau)`, differentiable in e; as
+    tau -> 0 this hardens to nearest-center (Voronoi) selection, which is
+    exactly what `top1` uses at eval.
     """
 
     def __init__(
@@ -612,11 +617,20 @@ class EnergyRouter(Router):
         temperature: float = 0.5,
         learn_centers: bool = True,
         energy_idx: int = 3,
+        centers_init: Sequence[float] | None = None,
     ) -> None:
         super().__init__(n_experts)
         self.temperature = temperature
         self.energy_idx = energy_idx
-        centers = torch.linspace(-2.0, 2.0, n_experts)
+        if centers_init is None:
+            centers = torch.linspace(-2.0, 2.0, n_experts)
+        else:
+            if len(centers_init) != n_experts:
+                raise ValueError(
+                    f"centers_init has {len(centers_init)} values, "
+                    f"expected n_experts={n_experts}"
+                )
+            centers = torch.tensor(list(centers_init), dtype=torch.float32)
         if learn_centers:
             self.centers = nn.Parameter(centers)
         else:

giant/pipeline.py

@@ -20,7 +20,7 @@ from giant.data.loader import (
     build_index_maps_from_files,
     build_process_map_from_files,
 )
-from giant.data.transforms import build_features, _WelfordAccumulator
+from giant.data.transforms import build_features, _WelfordAccumulator, _ReservoirSampler
 from giant.data.dataset import make_event_split, StreamingStepsDataset
 from giant.model.network import build_models, build_critics
 from giant.train import train as run_training
@@ -83,6 +83,18 @@ def run_train_job(
     cond_acc = _WelfordAccumulator(COND_DIM)
     tgt_acc = _WelfordAccumulator(X_DIM)
     sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM)
+    # EnergyRouter's default center spread (linspace over [-2, 2]) assumes
+    # the z-normalized energy column is roughly uniform, which real energy
+    # spectra rarely are — collect a reservoir sample here (reusing this
+    # same pass, not a second scan) so centers can instead be seeded from
+    # actual data quantiles below.
+    energy_router_active = (
+        router_cfg.get("enabled") and router_cfg.get("type") == "energy"
+    )
+    energy_idx = router_cfg.get("energy_idx", 3)
+    energy_sampler = (
+        _ReservoirSampler(capacity=100_000) if energy_router_active else None
+    )
     for path in files:
         for chunk in iter_file_chunks(path):
             mask = np.isin(chunk["event_id"], events_arr)
@@ -99,6 +111,8 @@ def run_train_job(
             )
             cond_acc.update(cond_cont)
             tgt_acc.update(target_s1)
+            if energy_sampler is not None:
+                energy_sampler.update(cond_cont[:, energy_idx])
             sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
             sec_phys = sec_cont[:, :, 4:6][sec_valid]
             if len(sec_phys) > 0:
@@ -107,6 +121,18 @@ def run_train_job(
     tgt_norm = tgt_acc.to_normalizer()
     sec_phys_norm = sec_phys_acc.to_normalizer()
 
+    if energy_sampler is not None and energy_sampler.n_seen > 0:
+        assert cond_norm.mean is not None and cond_norm.std is not None
+        normalized_sample = (
+            energy_sampler.sample - cond_norm.mean[energy_idx]
+        ) / cond_norm.std[energy_idx]
+        quantiles = np.linspace(0.0, 1.0, router_cfg["n_experts"])
+        centers_init = np.quantile(normalized_sample, quantiles).astype(np.float32)
+        router_cfg["centers_init"] = centers_init.tolist()
+        echo(
+            f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
+        )
+
     train_ds = StreamingStepsDataset(
         files=files,
         split_events=train_events,

tests/test_router.py

@@ -97,6 +97,42 @@ def test_build_router_ignores_unrecognized_kwargs():
     assert router.temperature == 0.3
 
 
+def test_energy_router_default_centers_are_linspace():
+    router = EnergyRouter(n_experts=4)
+    torch.testing.assert_close(router.centers, torch.linspace(-2.0, 2.0, 4))
+
+
+def test_energy_router_centers_init_overrides_default():
+    centers_init = [-1.0, 0.0, 0.5, 3.0]
+    router = EnergyRouter(n_experts=4, centers_init=centers_init)
+    torch.testing.assert_close(router.centers, torch.tensor(centers_init))
+
+
+def test_energy_router_centers_init_wrong_length_raises():
+    try:
+        EnergyRouter(n_experts=4, centers_init=[0.0, 1.0])
+    except ValueError:
+        return
+    raise AssertionError("expected ValueError for centers_init length mismatch")
+
+
+def test_energy_router_centers_init_respects_learn_centers_flag():
+    learned = EnergyRouter(
+        n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True
+    )
+    fixed = EnergyRouter(
+        n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False
+    )
+    assert isinstance(learned.centers, torch.nn.Parameter)
+    assert not isinstance(fixed.centers, torch.nn.Parameter)
+
+
+def test_build_router_threads_centers_init_through_energy_router():
+    centers_init = [-1.5, -0.5, 0.5, 1.5]
+    router = build_router("energy", 4, centers_init=centers_init)
+    torch.testing.assert_close(router.centers, torch.tensor(centers_init))
+
+
 def test_build_router_unknown_type_raises():
     try:
         build_router("nonexistent", 4)