From fa50f22cd373997881887c067883370cc441b3e8 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Sun, 13 Sep 2026 14:21:23 -0400 Subject: [PATCH 1/2] configs: declare per-embodiment dims once per HPT model config stem input_dim, infer_ac_dims, act_dim and enc-dec output_dim now interpolate robomimic_model.dims.{.proprio,.action,action_width}. Resolved values are unchanged (pinned by tests/unit/test_model_dims.py). Pi configs untouched. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01MWtR1x5tjREwoxQZAnUGrh Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01BbQNNaMAbTNMypegmJkv4s From fba621898e2909c637a59c1cadd3637acb900d01 Mon Sep 17 00:00:00 2001 From: RyanPCo Date: Tue, 15 Sep 2026 14:41:09 -0400 Subject: [PATCH 2/2] eval: deterministic validation (seeded prompts + flow sampling); log-only grad-norm MAD Validating one checkpoint twice now gives identical prompts, actions and metrics. * HPT gets a per-pass RNG seed, base + rank * stride + pass counter, where the counter advances once per forward_eval call and is rewound by ModelWrapper.on_validation_start. forward_eval builds a torch.Generator from it and threads it down HPTModel.forward -> DenoisingPolicy.forward -> sample_action, and forks/seeds the global RNG around the BC val-loss call (DenoisingPolicy.predict draws its own noise and timestep, out of reach of any generator kwarg). Every generator kwarg defaults to None, so training and all other callers are untouched. The val-loss seed is offset by EVAL_LOSS_STRIDE rather than by 1: a torch.Generator seeded with S draws the same stream as a global manual_seed(S), so +1 would have handed batch i's val loss exactly the noise batch i+1 samples with. eval_base_seed is a real __init__ kwarg, so a second eval pass can be made to draw different noise on purpose. * _build_prompts picks annotations with that same seeded RNG whenever the algo is in eval mode, regardless of annotation_sampling_mode, so eval keeps prompt variety instead of collapsing to sample[0]. Train mode is unchanged. * The MAD grad-norm spike detector in on_after_backward is now log-only: it still logs the raw norm, the threshold, the flag and the [GRAD_NORM_SPIKE] line, but never rescales gradients, and every step (flagged or not) goes into the rolling median/MAD window. Fixed clipping moves to Lightning: trainer/default.yaml gains gradient_clip_val (null for now) and gradient_clip_algorithm: norm. The on_before_optimizer_step hook goes with it: Precision._after_closure calls that hook BEFORE _clip_gradients, so "policy_grad_norms_clipped" was the pre-clip norm under another name -- bit-identical to policy_grad_norms_raw, and no more informative once gradient_clip_val is set. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01VBmhpVy7ZcMtsBJWh38cJp --- egomimic/algo/hpt.py | 127 +++++- egomimic/hydra_configs/trainer/default.yaml | 9 +- egomimic/models/denoising_policy.py | 11 +- egomimic/pl_utils/pl_model.py | 51 +-- tests/unit/test_eval_determinism.py | 436 ++++++++++++++++++++ 5 files changed, 590 insertions(+), 44 deletions(-) create mode 100644 tests/unit/test_eval_determinism.py diff --git a/egomimic/algo/hpt.py b/egomimic/algo/hpt.py index 979a536b79..04ec415c3a 100644 --- a/egomimic/algo/hpt.py +++ b/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 diff --git a/egomimic/hydra_configs/trainer/default.yaml b/egomimic/hydra_configs/trainer/default.yaml index 62f147c27d..e75492a53f 100644 --- a/egomimic/hydra_configs/trainer/default.yaml +++ b/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 \ No newline at end of file +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 diff --git a/egomimic/models/denoising_policy.py b/egomimic/models/denoising_policy.py index 645a8c44e4..b9180e0d37 100644 --- a/egomimic/models/denoising_policy.py +++ b/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]: """ diff --git a/egomimic/pl_utils/pl_model.py b/egomimic/pl_utils/pl_model.py index ca6868deb0..4c73a04c17 100644 --- a/egomimic/pl_utils/pl_model.py +++ b/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 diff --git a/tests/unit/test_eval_determinism.py b/tests/unit/test_eval_determinism.py new file mode 100644 index 0000000000..c9a5c7cbdf --- /dev/null +++ b/tests/unit/test_eval_determinism.py @@ -0,0 +1,436 @@ +"""Deterministic evaluation. + +Validating the same checkpoint twice must give the same prompts, the same +sampled actions and the same metrics. Three pieces make that true: + +* ``HPT._build_prompts`` picks annotations with a per-pass seeded RNG whenever + the algo is in eval mode (``annotation_sampling_mode`` only steers training); +* ``HPT.forward_eval`` builds a ``torch.Generator`` seeded from + ``EVAL_BASE_SEED + rank * EVAL_RANK_STRIDE + eval_step`` and threads it + down to the denoising head's sampler, and forks/seeds the global RNG around + the BC val-loss call; ``ModelWrapper.on_validation_start`` rewinds the + step, so pass N and pass N+1 replay the same draws. + +Plus the gradient-clipping change that rides along in this branch: the MAD +spike detector in ``ModelWrapper.on_after_backward`` is log-only (fixed +clipping is Lightning's ``trainer.gradient_clip_val``). + +Everything here is CPU-only and builds the smallest object that still runs the +real code path: the real ``FMPolicy`` Euler sampler, the real +``HPTModel.forward`` head dispatch and the real ``HPT.forward_eval``, with only +the stems/trunk (``forward_features``) and the norm stats stubbed out. +""" + +from __future__ import annotations + +import random +import types +from collections import deque + +import hydra +import pytest +import torch +import torch.nn as nn + +from egomimic.algo.hpt import EVAL_BASE_SEED, HPT, HPTModel +from egomimic.models.fm_policy import FMPolicy +from egomimic.pl_utils.pl_model import ModelWrapper + +DOMAIN = "eva_bimanual" +EMB_ID = 6 # EMBODIMENT.EVA_BIMANUAL; get_embodiment(6).lower() == DOMAIN +AC_KEY = "actions_cartesian" +ANNOTATION_KEY = "annotations" +BATCH = 4 +HORIZON = 3 +ACTION_DIM = 4 +FEAT_DIM = 5 + + +# --------------------------------------------------------------------------- +# Fixtures: a minimal but real HPT +# --------------------------------------------------------------------------- +class _TinyVelocity(nn.Module): + """Stands in for ConditionalUnet1D: v(x_t, t, cond), same call signature.""" + + def __init__(self, dim: int): + super().__init__() + self.lin = nn.Linear(dim, dim) + + def forward(self, x_t, t, global_cond): + return self.lin(x_t) * t.reshape(-1, 1, 1) + + +class _StubNormStats: + def unnormalize(self, predictions, embodiment_id): + return dict(predictions) + + +def _make_algo(annotation_key: str | None = ANNOTATION_KEY) -> HPT: + """A real HPT wired to a real FMPolicy head, without stems/trunk/dataset. + + ``HPT.__init__`` needs norm stats, configs and a dataset; every attribute + ``forward_eval`` / ``_build_prompts`` actually read is set here instead. + """ + torch.manual_seed(0) + head = FMPolicy( + model=_TinyVelocity(ACTION_DIM), + action_horizon=HORIZON, + infer_ac_dims={DOMAIN: ACTION_DIM}, + num_inference_steps=2, + ) + model = HPTModel.__new__(HPTModel) + nn.Module.__init__(model) + model.diffusion = True + model.shared_action = False + model.auxiliary_ac_keys = {} + model.heads = nn.ModuleDict({DOMAIN: head}) + model.device = torch.device("cpu") + features = ( + torch.arange(BATCH * FEAT_DIM, dtype=torch.float32).reshape(BATCH, FEAT_DIM) + / 10.0 + ) + model.forward_features = lambda domain, data: (features, None) + + algo = HPT.__new__(HPT) + algo.nets = nn.ModuleDict({"policy": model}) + algo.device = torch.device("cpu") + algo.norm_stats = _StubNormStats() + algo.camera_keys = {EMB_ID: []} + algo.proprio_keys = {EMB_ID: []} + algo.lang_keys = {EMB_ID: []} + algo.ac_keys = {EMB_ID: AC_KEY} + algo.auxiliary_ac_keys = {} + algo.shared_ac_key = None + algo.is_6dof = False + algo.freeze_repr = False + algo.freeze_depth = 8 + algo.eval_base_seed = EVAL_BASE_SEED + algo._eval_step = 0 + algo.annotation_key = annotation_key + algo.annotation_modality = "annotation" + algo.annotation_sampling_mode = "random" + algo.default_prompt = "do the task" + algo.train_image_augs = None + algo.eval_image_augs = None + algo.nets.eval() + return algo + + +def _make_batch() -> dict: + torch.manual_seed(1) + return { + EMB_ID: { + AC_KEY: torch.randn(BATCH, HORIZON, ACTION_DIM), + "pad_mask": torch.ones(BATCH, HORIZON, 1), + "embodiment": torch.tensor([EMB_ID]), + } + } + + +def _annotation_batch(n_options: int = 4) -> dict: + return { + ANNOTATION_KEY: [ + [f"sample{i}_option{j}" for j in range(n_options)] for i in range(BATCH) + ] + } + + +def _prompt_pass(algo: HPT, raw_batch: dict, n_batches: int = 3) -> list[list[str]]: + """One validation pass: N batches, step advancing as forward_eval does.""" + algo.reset_eval_step() + out = [] + for _ in range(n_batches): + out.append(algo._build_prompts(raw_batch, BATCH)) + algo._eval_step += 1 # what forward_eval does after each batch + return out + + +# --------------------------------------------------------------------------- +# 1. Prompts +# --------------------------------------------------------------------------- +def test_eval_prompts_are_reproducible_across_passes(): + algo = _make_algo() + raw = _annotation_batch() + first = _prompt_pass(algo, raw) + second = _prompt_pass(algo, raw) + assert first == second + + +def test_eval_prompts_keep_variety_and_are_not_always_first(): + algo = _make_algo() + raw = _annotation_batch() + prompts = [p for batch in _prompt_pass(algo, raw, n_batches=5) for p in batch] + firsts = [options[0] for options in raw[ANNOTATION_KEY]] + assert any( + p not in firsts for p in prompts + ), f"eval prompts collapsed to sample[0]: {prompts}" + assert len(set(prompts)) > 1 + + +def test_eval_prompts_ignore_sampling_mode(): + """`annotation_sampling_mode` is an algo-level knob; eval overrides it.""" + algo = _make_algo() + raw = _annotation_batch() + algo.annotation_sampling_mode = "random" + as_random = _prompt_pass(algo, raw) + algo.annotation_sampling_mode = "first" + as_first = _prompt_pass(algo, raw) + assert as_random == as_first + assert as_first[0] != [options[0] for options in raw[ANNOTATION_KEY]] + + +def test_eval_prompts_fall_back_to_default_on_empty(): + algo = _make_algo() + raw = {ANNOTATION_KEY: [[], ["a", "b"], [], ["c"]]} + prompts = algo._build_prompts(raw, 4) + assert prompts[0] == algo.default_prompt + assert prompts[2] == algo.default_prompt + assert prompts[1] in {"a", "b"} and prompts[3] == "c" + + +def test_missing_annotation_key_uses_default_prompt(): + algo = _make_algo(annotation_key=None) + assert algo._build_prompts({}, 3) == [algo.default_prompt] * 3 + + +@pytest.mark.parametrize("mode", ["random", "first"]) +def test_train_prompt_behaviour_is_unchanged(mode): + """In train mode the parent's logic still applies verbatim.""" + algo = _make_algo() + algo.annotation_sampling_mode = mode + algo.nets.train() + raw = _annotation_batch() + + random.seed(12345) + got = algo._build_prompts(raw, BATCH) + + random.seed(12345) + expected = [] + for options in raw[ANNOTATION_KEY]: + if mode == "random": + expected.append(options[random.randint(0, len(options) - 1)]) + else: + expected.append(options[0]) + assert got == expected + + +# --------------------------------------------------------------------------- +# 2. Seeded flow sampling +# --------------------------------------------------------------------------- +def _eval_pass(algo: HPT, batch: dict, n_batches: int = 2) -> list[dict]: + algo.reset_eval_step() + return [algo.forward_eval(batch) for _ in range(n_batches)] + + +def _actions(preds: list[dict]) -> list[torch.Tensor]: + return [p[f"{DOMAIN}_{AC_KEY}"] for p in preds] + + +def test_two_validation_passes_give_bit_identical_actions(): + algo = _make_algo() + batch = _make_batch() + first = _eval_pass(algo, batch) + second = _eval_pass(algo, batch) + for a, b in zip(_actions(first), _actions(second)): + assert torch.equal(a, b) + + +def test_val_loss_is_identical_across_passes(): + algo = _make_algo() + batch = _make_batch() + first = _eval_pass(algo, batch) + second = _eval_pass(algo, batch) + for a, b in zip(first, second): + assert torch.equal(a[f"{DOMAIN}_loss"], b[f"{DOMAIN}_loss"]) + + +def test_batches_within_a_pass_use_different_noise(): + """Without the reset the step has moved on, so the draws differ.""" + algo = _make_algo() + batch = _make_batch() + algo.reset_eval_step() + first = algo.forward_eval(batch) + second = algo.forward_eval(batch) # no reset: eval step is now 1 + assert not torch.equal(first[f"{DOMAIN}_{AC_KEY}"], second[f"{DOMAIN}_{AC_KEY}"]) + + +def test_forward_eval_does_not_disturb_the_global_rng(): + """Eval must not shift the training RNG stream (fork_rng + generator).""" + algo = _make_algo() + batch = _make_batch() + torch.manual_seed(7) + before = torch.randn(3) + torch.manual_seed(7) + algo.reset_eval_step() + algo.forward_eval(batch) + after = torch.randn(3) + assert torch.equal(before, after) + + +def test_batch_seed_depends_on_rank_and_step(monkeypatch): + algo = _make_algo() + algo.reset_eval_step() + monkeypatch.setattr(HPT, "_eval_rank", staticmethod(lambda: 0)) + rank0 = algo._eval_batch_seed() + monkeypatch.setattr(HPT, "_eval_rank", staticmethod(lambda: 3)) + rank3 = algo._eval_batch_seed() + assert rank0 != rank3 + algo._eval_step += 1 + assert algo._eval_batch_seed() == rank3 + 1 + + +def test_val_loss_noise_does_not_collide_with_the_next_batch(monkeypatch): + """A torch.Generator seeded with S draws the same stream as a global + manual_seed(S), so a +1 offset would give batch i's val loss the very noise + batch i+1 samples with.""" + from egomimic.algo.hpt import EVAL_LOSS_STRIDE + + assert EVAL_LOSS_STRIDE > 1000 + algo = _make_algo() + algo.reset_eval_step() + seeds = set() + for _ in range(4): + pass_seed = algo._eval_batch_seed() + seeds.add(pass_seed) + seeds.add(pass_seed + EVAL_LOSS_STRIDE) + algo._eval_step += 1 + assert len(seeds) == 8, "a val-loss seed repeats a sampling seed" + + +def test_eval_base_seed_is_a_real_knob(): + """A second eval pass can be made to draw DIFFERENT noise on purpose.""" + algo = _make_algo() + assert algo.eval_base_seed == EVAL_BASE_SEED + algo.reset_eval_step() + default = algo._eval_batch_seed() + algo.eval_base_seed = 12345 + assert algo._eval_batch_seed() == default + 12345 + + +def test_generator_is_threaded_to_the_head(): + """HPTModel.forward -> DenoisingPolicy.forward -> sample_action(generator).""" + algo = _make_algo() + model = algo.nets["policy"] + data = {"action": torch.zeros(BATCH, HORIZON, ACTION_DIM)} + gen = torch.Generator(device="cpu") + + gen.manual_seed(42) + a = model.forward(DOMAIN, data, generator=gen)[DOMAIN] + gen.manual_seed(42) + b = model.forward(DOMAIN, data, generator=gen)[DOMAIN] + gen.manual_seed(43) + c = model.forward(DOMAIN, data, generator=gen)[DOMAIN] + assert torch.equal(a, b) + assert not torch.equal(a, c) + + +def test_on_validation_start_resets_the_eval_step(): + algo = _make_algo() + algo.reset_eval_step() + algo._eval_step = 17 + wrapper = types.SimpleNamespace( + model=algo, + device=torch.device("cpu"), + _val_heads=lambda: {"valid": None}, + ) + ModelWrapper.on_validation_start(wrapper) + assert algo._eval_step == 0 + + +# --------------------------------------------------------------------------- +# 3. Log-only MAD gradient detector +# --------------------------------------------------------------------------- +class _GradHookStub: + """Just enough LightningModule for ModelWrapper.on_after_backward.""" + + def __init__(self, param: nn.Parameter, history): + self.enable_grad_norm = True + self.grad_norm_mad_scale = ModelWrapper.grad_norm_mad_scale + self.grad_norm_mad_min_count = ModelWrapper.grad_norm_mad_min_count + self.grad_norm_mad_window = ModelWrapper.grad_norm_mad_window + self.grad_norm_history = deque(history, maxlen=self.grad_norm_mad_window) + self._param = param + self.global_step = 4242 + self.trainer = types.SimpleNamespace(is_global_zero=True) + self.logged = {} + + def parameters(self): + return [self._param] + + def log(self, name, value, **kwargs): + self.logged[name] = value + + +def _spiking_stub(spike_norm: float = 50.0): + param = nn.Parameter(torch.zeros(4)) + grad = torch.full((4,), spike_norm / 2.0) # ||grad||_2 == spike_norm + param.grad = grad + history = [ + 0.9 if i % 2 else 1.1 for i in range(ModelWrapper.grad_norm_mad_min_count) + ] + return _GradHookStub(param, history), param + + +def test_on_after_backward_never_modifies_gradients(): + stub, param = _spiking_stub() + before = param.grad.clone() + ModelWrapper.on_after_backward(stub) + assert torch.equal(param.grad, before), "MAD detector rescaled the gradients" + + +def test_on_after_backward_still_flags_and_logs_the_spike(): + stub, _ = _spiking_stub() + ModelWrapper.on_after_backward(stub) + logged = stub.logged + assert logged["Train/policy_grad_norms_mad_flag"] == 1.0 + assert logged["Train/policy_grad_norms_raw"] == pytest.approx(50.0, rel=1e-5) + assert "Train/policy_grad_norms_mad_threshold" in logged + + +def test_flagged_steps_are_appended_to_the_history(): + """Log-only detector: every step goes into the rolling median/MAD window.""" + stub, _ = _spiking_stub() + n_before = len(stub.grad_norm_history) + ModelWrapper.on_after_backward(stub) + assert len(stub.grad_norm_history) == n_before + 1 + assert stub.grad_norm_history[-1] == pytest.approx(50.0, rel=1e-5) + + +# --------------------------------------------------------------------------- +# 4. Trainer config knob +# --------------------------------------------------------------------------- +def test_trainer_default_exposes_gradient_clip_val(compose_resolve): + cfg = compose_resolve("train_zarr_cartesian", []) + assert "gradient_clip_val" in cfg.trainer + assert cfg.trainer.gradient_clip_val is None # no clipping until chosen + assert cfg.trainer.gradient_clip_algorithm == "norm" + + +@pytest.mark.parametrize("trainer_cfg", ["default", "ddp", "ddp_pi", "debug"]) +def test_every_trainer_config_keeps_the_knob(trainer_cfg, compose_resolve): + cfg = compose_resolve("train_zarr_cartesian", [f"trainer={trainer_cfg}"]) + assert cfg.trainer.gradient_clip_val is None + assert cfg.trainer.gradient_clip_algorithm == "norm" + + +def test_gradient_clip_override_reaches_the_trainer(tmp_path, compose_resolve): + cfg = compose_resolve( + "train_zarr_cartesian", + [ + "trainer=default", + "trainer.gradient_clip_val=1.0", + "trainer.accelerator=cpu", + "trainer.precision=32", + "trainer.max_epochs=1", + "trainer.min_epochs=1", + f"paths.output_dir={tmp_path}", + ], + ) + trainer = hydra.utils.instantiate(cfg.trainer, logger=False, callbacks=None) + assert trainer.gradient_clip_val == 1.0 + assert trainer.gradient_clip_algorithm == "norm" + + +def test_lightning_module_does_not_override_gradient_clipping(): + """Nothing in ModelWrapper may intercept Lightning's clipping hook.""" + assert "configure_gradient_clipping" not in vars(ModelWrapper)