RunModel.py

import os
import sys
import pygame
from stable_baselines3 import PPO
from PushPullEnvironment import PushPullEnv

if __name__ == "__main__":
    if not os.path.exists("./saves"):
        sys.exit("No './saves' directory found. Train a model first!")

    # Grabs anything starting with "model_" - no zip stuff here
    models = sorted([f for f in os.listdir("./saves") if f.startswith("model_")])
    
    if not models:
        sys.exit("No models found in './saves'. Train a model first!")

    print("\n--- INFERENCE MENU ---")
    for i, m in enumerate(models, 1):
        print(f"{i}: {m}")

    try:
        choice = int(input("\nEnter the number of the model to load: "))
        if not (1 <= choice <= len(models)):
            raise ValueError
    except ValueError:
        sys.exit("Invalid choice. Exiting.")

    model_name = models[choice - 1]
    
    # Path to the folder where TrainModel.py saves the models
    model_path = os.path.join("./saves", model_name)
        
    print(f"Loading {model_name}...")

    # Initialize environment with human rendering enabled
    env = PushPullEnv(render_mode="human")
    
    try:
        model = PPO.load(model_path)
    except Exception as e:
        env.stop_video = True
        sys.exit(f"Error loading model: {e}")

    obs, _ = env.reset()
    try:
        while True:
            action, _ = model.predict(obs, deterministic=False)
            obs, _, terminated, truncated, _ = env.step(action)
            
            env.render()
            
            # Restart if the AI loses
            if terminated or truncated:
                pygame.time.wait(1000)
                obs, _ = env.reset()
    finally:
        env.stop_video = True
        pygame.quit()