audit: enforce paired task order after neutral warmup

4 files changed, 18 insertions, 0 deletions

diff --git a/experiments/analyze_kp_innovation_confirmation.py b/experiments/analyze_kp_innovation_confirmation.py
index adac5dc..37b06b5 100644
--- a/experiments/analyze_kp_innovation_confirmation.py
+++ b/experiments/analyze_kp_innovation_confirmation.py
@@ -76,6 +76,12 @@ def main():
                 raise ValueError(
                     f"tracked-dirty MT-3 seed {seed} {condition}")
             source_commits.add(record["provenance"]["git_commit"])
+            if condition != "clean":
+                warmup = record.get("predictor_warmup", {})
+                if (warmup.get("instruction_present") is not False
+                        or warmup.get("task_loader_state_restored") is not True):
+                    raise ValueError(
+                        f"MT-3 seed {seed} {condition} warmup invariant failed")
             split = record["split"]
             if not (split["validation_examples"] == 0
                     and split["validation_index_sha256"] is None
diff --git a/experiments/analyze_kp_innovation_full.py b/experiments/analyze_kp_innovation_full.py
index 2b5394d..c01468b 100644
--- a/experiments/analyze_kp_innovation_full.py
+++ b/experiments/analyze_kp_innovation_full.py
@@ -70,6 +70,10 @@ def main():
         if record.get("calibration_metric_space") != (
                 "reciprocal_local_activity_products_with_mixed_apical_traffic"):
             raise ValueError(f"MT-2 {rule} metric-space drift")
+        warmup = record.get("predictor_warmup", {})
+        if (warmup.get("instruction_present") is not False
+                or warmup.get("task_loader_state_restored") is not True):
+            raise ValueError(f"MT-2 {rule} neutral-warmup invariant failed")
     if len(source_commits) != 1:
         raise ValueError("MT-2 conditions must share one source revision")
 
diff --git a/experiments/analyze_kp_innovation_short.py b/experiments/analyze_kp_innovation_short.py
index 5f21dea..2e7b9f4 100644
--- a/experiments/analyze_kp_innovation_short.py
+++ b/experiments/analyze_kp_innovation_short.py
@@ -64,6 +64,10 @@ def main():
         if record.get("calibration_metric_space") != (
                 "reciprocal_local_activity_products_with_mixed_apical_traffic"):
             raise ValueError(f"MT-1 {rule} metric-space drift")
+        warmup = record.get("predictor_warmup", {})
+        if (warmup.get("instruction_present") is not False
+                or warmup.get("task_loader_state_restored") is not True):
+            raise ValueError(f"MT-1 {rule} neutral-warmup invariant failed")
     if len(source_commits) != 1:
         raise ValueError("MT-1 conditions must share one source revision")
 
diff --git a/experiments/conv_run.py b/experiments/conv_run.py
index c25c48c..a1ebb6a 100644
--- a/experiments/conv_run.py
+++ b/experiments/conv_run.py
@@ -439,12 +439,16 @@ def run(args):
             counters["predictor_warmup_examples"] += x.shape[0]
         del forward
         train.g.set_state(loader_state)
+        loader_state_restored = torch.equal(train.g.get_state(), loader_state)
+        if not loader_state_restored:
+            raise AssertionError("predictor warmup did not restore loader state")
         log["predictor_warmup"] = {
             "steps": args.predictor_warmup_steps,
             "first_mse": predictor_metrics[0],
             "mean_mse": sum(predictor_metrics) / len(predictor_metrics),
             "last_mse": predictor_metrics[-1],
             "instruction_present": False,
+            "task_loader_state_restored": loader_state_restored,
         }
         if args.mode == "kp_traffic":
             audit_x, _ = traffic_calibration_batch