139 lines
4.8 KiB
Python
139 lines
4.8 KiB
Python
import argparse
|
|
import sys
|
|
import os
|
|
from config import MazeConfig
|
|
from maze import generate
|
|
from agent import Agent, NeuralNet
|
|
from genetic import next_generation
|
|
|
|
|
|
def handle_events() -> bool:
|
|
"""Returns False if user wants to quit."""
|
|
import pygame
|
|
for event in pygame.event.get():
|
|
if event.type == pygame.QUIT:
|
|
return False
|
|
if event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
|
|
return False
|
|
return True
|
|
|
|
|
|
def run_generation_headless(agents, grid, goal_row, goal_col, cfg):
|
|
"""Run one generation without any rendering."""
|
|
for step in range(cfg.max_steps):
|
|
for agent in agents:
|
|
agent.step(grid, goal_row, goal_col, cfg)
|
|
|
|
|
|
def run_generation_visual(agents, grid, goal_row, goal_col, cfg, viz, clock, fps, gen) -> bool:
|
|
"""Run one generation with rendering. Returns False if user quit."""
|
|
import pygame
|
|
for step in range(cfg.max_steps):
|
|
for agent in agents:
|
|
agent.step(grid, goal_row, goal_col, cfg)
|
|
if not handle_events():
|
|
return False
|
|
viz.render(agents, gen, step)
|
|
clock.tick(fps)
|
|
return True
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(description="Maze ML — Genetic Algorithm Visualizer")
|
|
parser.add_argument("--fast", action="store_true", help="Train without rendering (headless)")
|
|
parser.add_argument("--generations", type=int, default=200)
|
|
parser.add_argument("--maze-size", type=int, default=25, help="Maze cols and rows (square)")
|
|
parser.add_argument("--seed", type=int, default=None)
|
|
parser.add_argument("--speed", type=int, default=60, help="FPS cap for visual mode")
|
|
args = parser.parse_args()
|
|
|
|
cfg = MazeConfig(maze_cols=args.maze_size, maze_rows=args.maze_size)
|
|
goal_row = cfg.maze_rows - 1
|
|
goal_col = cfg.maze_cols - 1
|
|
|
|
grid = generate(cfg.maze_rows, cfg.maze_cols, seed=args.seed)
|
|
|
|
# Initialize population
|
|
nets = [NeuralNet() for _ in range(cfg.population)]
|
|
|
|
print(f"Starting training: {args.generations} generations, population={cfg.population}")
|
|
|
|
if args.fast:
|
|
# --- Headless fast mode: no pygame, no display ---
|
|
print("Fast mode: rendering disabled during training.")
|
|
|
|
for gen in range(args.generations):
|
|
agents = [Agent(net, 0, 0) for net in nets]
|
|
run_generation_headless(agents, grid, goal_row, goal_col, cfg)
|
|
|
|
for agent in agents:
|
|
agent.compute_fitness(goal_row, goal_col, cfg)
|
|
|
|
best_fit = max(a.fitness for a in agents)
|
|
reached = sum(1 for a in agents if a.reached_goal)
|
|
print(f"Gen {gen+1:>4}/{args.generations} | Best fitness: {best_fit:>10.2f} | Reached goal: {reached}/{cfg.population}")
|
|
|
|
nets = next_generation(agents, cfg)
|
|
|
|
print("Training complete.")
|
|
sys.exit(0)
|
|
|
|
else:
|
|
# --- Visual mode: initialise pygame ---
|
|
import pygame
|
|
from visualizer import Visualizer
|
|
|
|
pygame.init()
|
|
win_w = cfg.maze_cols * cfg.cell_size + 200
|
|
win_h = cfg.maze_rows * cfg.cell_size
|
|
screen = pygame.display.set_mode((win_w, win_h))
|
|
pygame.display.set_caption(cfg.window_title)
|
|
clock = pygame.time.Clock()
|
|
|
|
viz = Visualizer(screen, grid, cfg)
|
|
|
|
for gen in range(args.generations):
|
|
agents = [Agent(net, 0, 0) for net in nets]
|
|
|
|
ok = run_generation_visual(agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, gen)
|
|
if not ok:
|
|
print("Quit by user.")
|
|
pygame.quit()
|
|
sys.exit(0)
|
|
|
|
for agent in agents:
|
|
agent.compute_fitness(goal_row, goal_col, cfg)
|
|
|
|
best_fit = max(a.fitness for a in agents)
|
|
reached = sum(1 for a in agents if a.reached_goal)
|
|
print(f"Gen {gen+1:>4}/{args.generations} | Best fitness: {best_fit:>10.2f} | Reached goal: {reached}/{cfg.population}")
|
|
|
|
nets = next_generation(agents, cfg)
|
|
viz.reset_trails()
|
|
|
|
# Final visual replay with best generation
|
|
best_nets = nets
|
|
print("\nTraining complete. Showing final generation replay...")
|
|
replay_agents = [Agent(net, 0, 0) for net in best_nets]
|
|
viz.reset_trails()
|
|
for step in range(cfg.max_steps):
|
|
for agent in replay_agents:
|
|
agent.step(grid, goal_row, goal_col, cfg)
|
|
agent.compute_fitness(goal_row, goal_col, cfg)
|
|
if not handle_events():
|
|
break
|
|
viz.render(replay_agents, args.generations, step)
|
|
clock.tick(args.speed)
|
|
|
|
print("Done. Press ESC or close window to exit.")
|
|
running = True
|
|
while running:
|
|
running = handle_events()
|
|
clock.tick(30)
|
|
|
|
pygame.quit()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|