exclude is_los_nan from metrics

Author: sozforex

pufferlib/pufferl.py

@@ -187,12 +187,13 @@ def _train_worker(args):
 
     backend.close(pufferl)
 
-def _downsample_logs(all_logs, n):
+def _downsample_logs(all_logs, n, exclude_keys=()):
     if not all_logs:
         raise ValueError('Cannot downsample empty logs')
 
     expected_keys = set(all_logs[0])
-    metrics = {k: [[]] for k in all_logs[0]}
+    exclude_keys = set(exclude_keys)
+    metrics = {k: [[]] for k in all_logs[0] if k not in exclude_keys}
     logged_timesteps = all_logs[-1]['agent_steps']
     next_bin = logged_timesteps / (n - 1) if n > 1 else np.inf
     for idx, log in enumerate(all_logs):
@@ -206,7 +207,8 @@ def _downsample_logs(all_logs, n):
             )
 
         for k, v in log.items():
-            metrics[k][-1].append(v)
+            if k in metrics:
+                metrics[k][-1].append(v)
 
         if log['agent_steps'] < next_bin:
             continue
@@ -365,7 +367,12 @@ def _train(env_name, args, result_queue=None, verbose=False, sweep_early_stop=No
     # This version has the training perf logs and eval env logs
     all_logs.append(flat_logs)
 
-    metrics = _downsample_logs(all_logs, args['sweep']['downsample'])
+    exclude_keys = ()
+    if sweep_early_stop is not None:
+        exclude_keys = pufferlib.sweep.SWEEP_NON_METRIC_LOG_KEYS
+
+    metrics = _downsample_logs(
+        all_logs, args['sweep']['downsample'], exclude_keys=exclude_keys)
 
     # Match-mode: single observation at final-training cost. Protein's curve
     # fit collapses to one point — we only trust the match winrate, not any

pufferlib/sweep.py

@@ -29,6 +29,7 @@
     'is_loss_nan': False,
     'early_stop_threshold': EARLY_STOP_THRESHOLD_FLOOR,
 }
+SWEEP_NON_METRIC_LOG_KEYS = frozenset(('is_loss_nan',))
 
 def apply_early_stop_log_defaults(logs):
     logs.update(SWEEP_EARLY_STOP_LOG_DEFAULTS)