TrainModel.py

import os
import sys

# Force Pygame into headless mode for maximum training speed
os.environ["SDL_VIDEODRIVER"] = "dummy"

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from PushPullEnvironment import PushPullEnv

class BestSurvivalCallback(BaseCallback):
    def __init__(self, save_dir, milestone, verbose=0):
        super().__init__(verbose)
        self.save_dir = save_dir
        self.next_milestone = milestone
        os.makedirs(self.save_dir, exist_ok=True)

    def _on_step(self):
        dones = self.locals.get("dones")
        if dones is not None and dones[0]:
            for info in self.locals.get("infos", []):
                if "time_survived" in info:
                    survived = info["time_survived"]
                    if survived >= self.next_milestone:
                        save_path = os.path.join(self.save_dir, f"model_{survived:.2f}s")
                        self.model.save(save_path)
                        self.next_milestone = float(int(survived) + 1)
        return True

if __name__ == "__main__":
    os.makedirs("./saves", exist_ok=True)
    models = sorted([f for f in os.listdir("./saves") if f.endswith(".zip")])

    print("\n--- TRAINING MENU ---")
    print("0: Train from scratch")
    for i, m in enumerate(models, 1):
        print(f"{i}: Continue {m}")

    try:
        choice = int(input("\nEnter number: "))
    except ValueError:
        sys.exit("Invalid input. Exiting.")

    env = PushPullEnv() # render_mode=None by default

    if choice == 0:
        print("Starting fresh training...")
        model = PPO("MlpPolicy", env, verbose=1)
        callback = BestSurvivalCallback("./saves", 1.0)
    elif 1 <= choice <= len(models):
        model_name = models[choice - 1]
        model_path = os.path.join("./saves", model_name[:-4]) # Strip .zip for SB3
        print(f"Loading {model_name}...")
        model = PPO.load(model_path, env=env)
        
        # Extract milestone from filename
        try:
            survived_str = model_name.replace("model_", "").replace("s.zip", "")
            milestone = float(int(float(survived_str)) + 1)
        except:
            milestone = 1.0
            
        callback = BestSurvivalCallback("./saves", milestone)
    else:
        sys.exit("Invalid choice. Exiting.")

    try:
        model.learn(total_timesteps=2000000, callback=callback)
    except KeyboardInterrupt:
        print("Training interrupted. Saving...")
    finally:
        env.stop_video = True