eval: deterministic validation (seeded prompts + flow sampling); log-only grad-norm MAD

egomimic/algo/hpt.py

@@ -23,6 +23,21 @@
 # it has exactly one consumer, below).
 STD_SCALE = 0.02
 
+# Deterministic evaluation. Every RNG used while the algo is in eval mode is
+# seeded from ``EVAL_BASE_SEED + rank * EVAL_RANK_STRIDE + eval_step``, where
+# the eval step counts ``forward_eval`` calls (batches) since the last
+# ``reset_eval_step()`` (called from ``ModelWrapper.on_validation_start``).
+# Two validation passes over the same batches therefore replay the same prompts
+# and the same sampling noise, while different ranks and different batches
+# within a pass still see different draws.
+EVAL_BASE_SEED = 0
+EVAL_RANK_STRIDE = 1_000_003
+# Offset of the val-loss RNG from the sampling generator's seed. It must not be
+# a small number: the eval step advances by 1 per forward_eval call, and a
+# torch.Generator seeded with S draws the same stream as a global manual_seed
+# of S, so an offset of 1 gave batch i's val loss the noise batch i+1 sampled.
+EVAL_LOSS_STRIDE = 7_919_369
+
 
 class HPTModel(nn.Module):
     """
@@ -659,7 +674,7 @@ def compute_loss(self, batch):
         total_loss = action_loss + shared_action_loss + auxiliary_action_loss
         return total_loss
 
-    def forward(self, domain, data):
+    def forward(self, domain, data, generator=None):
         """
         Forward pass of the HPTModel to compute actions.
 
@@ -669,6 +684,10 @@ def forward(self, domain, data):
             The domain corresponding to the input data.
         data : dict
             Dictionary containing input data for various modalities.
+        generator : torch.Generator, optional
+            RNG handed to the denoising heads' sampler so that evaluation is
+            reproducible (see ``HPT.forward_eval``). ``None`` (the default)
+            keeps the global RNG, so every other caller is unchanged.
 
         Returns
         -------
@@ -678,19 +697,23 @@ def forward(self, domain, data):
         features, block_outputs = self.forward_features(domain, data)
         action = {}
 
+        # Only the denoising (diffusion / flow matching) heads sample noise and
+        # accept a generator; the plain MLP heads take the features alone.
+        head_kwargs = {"generator": generator} if self.diffusion else {}
+
         if self.diffusion:
             features = (features, domain)
 
         if domain in self.heads:
-            action[domain] = self.heads[domain](features)
+            action[domain] = self.heads[domain](features, **head_kwargs)
 
         if self.shared_action:
-            action["shared"] = self.heads["shared"](features)
+            action["shared"] = self.heads["shared"](features, **head_kwargs)
 
         if domain in self.auxiliary_ac_keys:
             for key in self.auxiliary_ac_keys[domain]:
                 if f"{domain}_{key}" in self.heads:
-                    action[key] = self.heads[f"{domain}_{key}"](features)
+                    action[key] = self.heads[f"{domain}_{key}"](features, **head_kwargs)
 
         return action
 
@@ -855,6 +878,8 @@ def __init__(
         model = HPTModel(**trunk)
         model.auxiliary_ac_keys = self.auxiliary_ac_keys
 
+        self.eval_base_seed = int(kwargs.get("eval_base_seed", EVAL_BASE_SEED))
+        self._eval_step = 0
         self.multitask = kwargs.get("multitask", False)
         self.device = kwargs.get(
             "device", torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -956,16 +981,76 @@ def __init__(
 
         self.training_step = 0
 
+    # ------------------------------------------------------------------
+    # Deterministic evaluation
+    # ------------------------------------------------------------------
+    def reset_eval_step(self) -> None:
+        """Start a new validation pass: the next ``forward_eval`` call is step 0.
+
+        Called from ``ModelWrapper.on_validation_start``. Two passes over the
+        same loader (same order, same number of batches) then draw the same
+        prompts and the same sampling noise.
+        """
+        self._eval_step = 0
+
+    @staticmethod
+    def _eval_rank() -> int:
+        """Distributed rank, or 0 outside a process group."""
+        if torch.distributed.is_available() and torch.distributed.is_initialized():
+            return torch.distributed.get_rank()
+        return 0
+
+    def _eval_batch_seed(self) -> int:
+        """Seed for the current eval batch: base + rank stride + eval step."""
+        return (
+            self.eval_base_seed + self._eval_rank() * EVAL_RANK_STRIDE + self._eval_step
+        )
+
+    def _eval_device(self) -> torch.device:
+        return (
+            torch.device(self.device)
+            if self.device is not None
+            else torch.device("cpu")
+        )
+
+    def _eval_generator(self) -> torch.Generator:
+        """Seeded generator on the device the sampling noise is drawn on."""
+        device = self._eval_device()
+        generator = torch.Generator(device=device)
+        generator.manual_seed(self._eval_batch_seed())
+        return generator
+
+    def _eval_prompt_rng(self) -> random.Random:
+        """Seeded Python RNG for annotation choice during eval."""
+        return random.Random(self._eval_batch_seed())
+
+    def _eval_fork_devices(self) -> list:
+        """Devices ``torch.random.fork_rng`` must save/restore (CPU is implicit)."""
+        device = self._eval_device()
+        if device.type != "cuda":
+            return []
+        return [
+            device.index if device.index is not None else torch.cuda.current_device()
+        ]
+
     def _build_prompts(self, _batch, batch_size: int) -> list[str]:
         """Sample one annotation per batch item, falling back to default_prompt
         on empty / missing annotations. Mirrors the Pi algo flow.
+
+        In eval the choice comes from a per-batch seeded RNG regardless of
+        ``annotation_sampling_mode``: that keeps prompt variety (unlike always
+        taking ``sample[0]``) while making two validation passes over the same
+        batches produce identical prompt lists. Training is unchanged.
         """
         if self.annotation_key is None or self.annotation_key not in _batch:
             return [self.default_prompt] * batch_size
+        eval_rng = None if self.nets.training else self._eval_prompt_rng()
         prompts = []
         for sample in _batch[self.annotation_key]:
             if not sample:
                 prompts.append(self.default_prompt)
+            elif eval_rng is not None:
+                prompts.append(sample[eval_rng.randint(0, len(sample) - 1)])
             elif self.annotation_sampling_mode == "random":
                 prompts.append(sample[random.randint(0, len(sample) - 1)])
             else:  # "first"
@@ -1101,6 +1186,12 @@ def forward_eval(self, batch):
             }
         """
         unnorm_preds = {}
+        # One seeded generator per batch: the flow-matching noise below and
+        # the prompts picked in ``process_batch_for_training`` share this batch's
+        # seed, so replaying the pass replays every draw.
+        generator = self._eval_generator()
+        batch_seed = self._eval_batch_seed()
+        fork_devices = self._eval_fork_devices()
         for embodiment_id, _batch in batch.items():
             embodiment_name = get_embodiment(embodiment_id).lower()
             cam_keys = self.camera_keys[embodiment_id]
@@ -1119,16 +1210,24 @@ def forward_eval(self, batch):
             # so keep a fresh copy for the forward() call below.
             forward_data = self._clone_batch(hpt_batch["data"])
 
-            # BC val loss — same call as forward_training.
-            if self.freeze_repr:
-                val_loss = self.nets["policy"].compute_loss_depth(
-                    hpt_batch, depth=self.freeze_depth
-                )
-            else:
-                val_loss = self.nets["policy"].compute_loss(hpt_batch)
+            # BC val loss — same call as forward_training. The training loss
+            # draws its own noise / timestep from the global RNG deep inside the
+            # head (``DenoisingPolicy.predict``), which no generator kwarg
+            # reaches; fork and seed the global RNG around the call instead so
+            # the reported val loss is reproducible too.
+            with torch.random.fork_rng(devices=fork_devices):
+                torch.manual_seed(batch_seed + EVAL_LOSS_STRIDE)
+                if self.freeze_repr:
+                    val_loss = self.nets["policy"].compute_loss_depth(
+                        hpt_batch, depth=self.freeze_depth
+                    )
+                else:
+                    val_loss = self.nets["policy"].compute_loss(hpt_batch)
             unnorm_preds[f"{embodiment_name}_loss"] = val_loss
 
-            actions = self.nets["policy"].forward(hpt_batch["domain"], forward_data)
+            actions = self.nets["policy"].forward(
+                hpt_batch["domain"], forward_data, generator=generator
+            )
             predictions = OrderedDict()
 
             for key in actions:
@@ -1153,6 +1252,10 @@ def forward_eval(self, batch):
             for key in unnorm_actions:
                 unnorm_preds[f"{embodiment_name}_{key}"] = unnorm_actions[key]
 
+        # Advance once per forward_eval call (after the batch, so the prompts
+        # built for this batch and the noise drawn for it share a seed).
+        # ``reset_eval_step`` rewinds it at on_validation_start.
+        self._eval_step += 1
         return unnorm_preds
 
     @override

egomimic/hydra_configs/trainer/default.yaml

@@ -17,4 +17,11 @@ check_val_every_n_epoch: 100
 
 # set True to to ensure deterministic results
 # makes training slower but gives more reproducibility than just setting seeds
-deterministic: False
+deterministic: False
+
+# Fixed gradient clipping, applied by Lightning in configure_gradient_clipping.
+# The MAD spike detector in ModelWrapper.on_after_backward is log-only and never
+# touches gradients; this is the only thing that clips. null = no clipping (the
+# value is picked from the grad-norm curve of a log-only baseline run).
+gradient_clip_val: null
+gradient_clip_algorithm: norm

egomimic/models/denoising_policy.py

@@ -80,9 +80,16 @@ def sample_action(self, global_cond, embodiment_name, generator=None):
         )
         return self.inference(noise, global_cond, generator)
 
-    def forward(self, global_cond):
+    def forward(self, global_cond, generator=None):
+        """Sample actions from the conditioning.
+
+        ``generator`` is threaded down to ``sample_action`` so the eval path can
+        make sampling reproducible (see ``HPT.forward_eval``). It defaults to
+        ``None``, i.e. the global RNG, so training and every other caller are
+        unchanged.
+        """
         cond, embodiment = global_cond
-        return self.sample_action(cond, embodiment)
+        return self.sample_action(cond, embodiment, generator=generator)
 
     def predict(self, actions, global_cond) -> Tuple[torch.Tensor, torch.Tensor]:
         """

egomimic/pl_utils/pl_model.py

@@ -184,38 +184,23 @@ def on_after_backward(self):
                 info["policy_grad_norms_mad_threshold"] = threshold
                 grad_norm_flagged = grad_norm_val > threshold
                 info["policy_grad_norms_mad_flag"] = float(grad_norm_flagged)
-                if grad_norm_flagged:
-                    torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm=median)
-                    if self.trainer.is_global_zero:
-                        print(
-                            "[GRAD_NORM_SPIKE] "
-                            f"step={self.global_step} "
-                            f"grad_norm={grad_norm_val:.4f} "
-                            f"median={median:.4f} "
-                            f"mad={mad:.4f} "
-                            f"threshold={threshold:.4f}",
-                            flush=True,
-                        )
-
-        if not grad_norm_flagged:
-            self.grad_norm_history.append(grad_norm_val)
+                if grad_norm_flagged and self.trainer.is_global_zero:
+                    print(
+                        "[GRAD_NORM_SPIKE] "
+                        f"step={self.global_step} "
+                        f"grad_norm={grad_norm_val:.4f} "
+                        f"median={median:.4f} "
+                        f"mad={mad:.4f} "
+                        f"threshold={threshold:.4f}",
+                        flush=True,
+                    )
+
+        # Log-only detector: every step enters the rolling window; clipping is
+        # trainer.gradient_clip_val.
+        self.grad_norm_history.append(grad_norm_val)
         for k, v in info.items():
             self.log("Train/" + k, v, on_step=False, on_epoch=True, sync_dist=True)
 
-    def on_before_optimizer_step(self, optimizer):
-        if not self.enable_grad_norm:
-            return
-        grad_norm = torch.nn.utils.clip_grad_norm_(
-            self.parameters(), max_norm=float("inf")
-        )
-        self.log(
-            "Train/policy_grad_norms_clipped",
-            float(grad_norm),
-            on_step=False,
-            on_epoch=True,
-            sync_dist=True,
-        )
-
     def _val_heads(self) -> dict:
         return {
             "valid": self.evaluator,
@@ -224,6 +209,14 @@ def _val_heads(self) -> dict:
         }
 
     def on_validation_start(self):
+        # Deterministic eval: rewind the algo's per-batch RNG step so this
+        # pass replays the same prompts and sampling noise as any other pass
+        # over the same batches (egomimic/algo/hpt.py). Algos without the hook
+        # (e.g. PI) are unaffected.
+        reset_eval_rng = getattr(self.model, "reset_eval_step", None)
+        if callable(reset_eval_rng):
+            reset_eval_rng()
+
         heads = [h for h in self._val_heads().values() if h is not None]
         if not heads:
             return