Upload train_all_phases.py

train_all_phases.py

@@ -399,7 +399,7 @@ class CurriculumEnv(gym.Env):
                     f"Stage {self.CURRICULUM_STAGES[self.stage_idx]} complete: "
                     f"win_rate={win_rate:.2%}, avg_reward={avg_reward:.1f}. "
                     f"Advancing to {self.CURRICULUM_STAGES[self.stage_idx + 1]}",
-                    "INFO",
+                    trackio.AlertLevel.INFO,
                 )
                 self.stage_idx += 1
                 self.stage_episodes = 0
@@ -508,8 +508,8 @@ class TrackioLoggingCallback(BaseCallback):
         self._last_mean_reward = 0.0
 
     def _on_training_start(self):
-        trackio.init_run(project=self.project, run_name=self.run_name)
-        trackio.alert("Training Started", f"{self.run_name} training began.", "INFO")
+        trackio.init(project=self.project, name=self.run_name)
+        trackio.alert("Training Started", f"{self.run_name} training began.", trackio.AlertLevel.INFO)
 
     def _on_step(self):
         if self.n_calls % self.log_interval == 0:
@@ -535,13 +535,13 @@ class TrackioLoggingCallback(BaseCallback):
 
                 if mean_r < -5.0 and self.num_timesteps > 50_000:
                     trackio.alert("Low Reward Warning",
-                        f"mean_reward={mean_r:.2f} at step {self.num_timesteps} -- may be camping.", "WARN")
+                        f"mean_reward={mean_r:.2f} at step {self.num_timesteps} -- may be camping.", trackio.AlertLevel.WARN)
         return True
 
     def _on_training_end(self):
         trackio.alert("Training Complete",
             f"Finished at {self.num_timesteps}. Final mean reward: {self._last_mean_reward:.2f}",
-            "INFO")
+            trackio.AlertLevel.INFO)
         trackio.finish()