From 9699c1cf26057af5cac139fe0d05dbbbec3332a7 Mon Sep 17 00:00:00 2001
From: thatguy11325 <148832074+thatguy11325@users.noreply.github.com>
Date: Sat, 2 Nov 2024 16:41:54 -0400
Subject: [PATCH] sweeps are a bit safer

---
 pokemonred_puffer/sweep.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/pokemonred_puffer/sweep.py b/pokemonred_puffer/sweep.py
index 62814d8..f534015 100644
--- a/pokemonred_puffer/sweep.py
+++ b/pokemonred_puffer/sweep.py
@@ -202,6 +202,9 @@ def launch_sweep(
                             if (
                                 "environment/stats/required_count" in summary_metrics
                                 and "performance/uptime" in summary_metrics
+                                # Only count agents that have run more than 1M steps
+                                and "Overview/agent_steps" in summary_metrics
+                                and summary_metrics["overview/agent_steps"] > 1e6
                             ):
                                 obs_in = ObservationInParam(
                                     input={
@@ -242,7 +245,10 @@ def _fn():
             }
         )
         agent_config = update_base_config(base_config, agent_config)
-        train.train(config=agent_config, debug=debug, track=True)
+        try:
+            train.train(config=agent_config, debug=debug, track=True)
+        except Exception as e:
+            print(f"Exception in training: {e!r}")
 
     for _ in range(99999):
         # Manually reset the env id counter between runs