experiment: add one-sided residual stability margin

1 files changed, 6 insertions, 1 deletions

diff --git a/experiments/diagnose_kp_traffic_nonfinite.py b/experiments/diagnose_kp_traffic_nonfinite.py
index 99a4e99..7b7389f 100644
--- a/experiments/diagnose_kp_traffic_nonfinite.py
+++ b/experiments/diagnose_kp_traffic_nonfinite.py
@@ -90,6 +90,7 @@ def main():
                         choices=("nlms20", "closed_form"),
                         default="nlms20")
     parser.add_argument("--predictor_every", type=int, default=16)
+    parser.add_argument("--stability_margin", type=float, default=0.0)
     parser.add_argument("--device", default="cuda")
     parser.add_argument("--max_steps", type=int, default=352)
     parser.add_argument("--out", required=True)
@@ -98,6 +99,8 @@ def main():
         raise ValueError("max_steps must be positive")
     if args.predictor_every < 0:
         raise ValueError("predictor_every must be nonnegative")
+    if args.stability_margin < 0:
+        raise ValueError("stability_margin must be nonnegative")
 
     torch.manual_seed(0)
     if str(args.device).startswith("cuda"):
@@ -137,7 +140,8 @@ def main():
     warmup_mse = []
     if args.predictor_mode == "closed_form":
         closed_form_fit = net.predictor_closed_form_fit(
-            calibration_forward["hiddens"])
+            calibration_forward["hiddens"],
+            stability_margin=args.stability_margin)
         warmup_mse.append(closed_form_fit["mse"])
     else:
         iterator = iter(train)
@@ -228,6 +232,7 @@ def main():
         "rule": args.rule,
         "predictor_mode": args.predictor_mode,
         "predictor_every": args.predictor_every,
+        "stability_margin": args.stability_margin,
         "max_steps": args.max_steps,
         "split": split,
         "traffic_calibration": calibration,