Update env/environment.py

env/environment.py

@@ -64,8 +64,6 @@ class SQLDebuggerEnvironment:
             raise ValueError(f"Failed to load task: {str(e)}")
 
         # ── Initialize DatabaseSimulator ──────────────────────────
-        # Always clear milestone state on episode reset.
-        self._milestones_earned = set()
         # Only initialize for Round 2 scenarios (have 'tables' key)
         if "tables" in task and "slow_queries" in task:
             self._db_sim         = DatabaseSimulator(task)
@@ -74,6 +72,7 @@ class SQLDebuggerEnvironment:
                 # Round 1 task — no DB simulator needed
             self._db_sim         = None
             self._baseline_score = 0.0
+            self._milestones_earned = set()
 
         # ── Reset episode state ───────────────────────────────────
         self._current_task = task
@@ -362,4 +361,4 @@ class SQLDebuggerEnvironment:
 #  SINGLETON INSTANCE (used by FastAPI)
 # ─────────────────────────────────────────────
 
-environment = SQLDebuggerEnvironment()
+environment = SQLDebuggerEnvironment()