183 lines
6.7 KiB
Python
183 lines
6.7 KiB
Python
import argparse
|
||
import sys
|
||
import os
|
||
from config import MazeConfig
|
||
from maze import generate, bfs_distances, find_short_seed
|
||
from agent import Agent, Genome
|
||
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, steps_per_frame) -> bool:
|
||
"""Run one generation with rendering. Returns False if user quit."""
|
||
step = 0
|
||
while step < cfg.max_steps:
|
||
# Advance simulation N steps per render frame
|
||
for _ in range(steps_per_frame):
|
||
if step >= cfg.max_steps:
|
||
break
|
||
for agent in agents:
|
||
agent.step(grid, goal_row, goal_col, cfg)
|
||
step += 1
|
||
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=15, 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")
|
||
parser.add_argument("--steps-per-frame", type=int, default=3, help="Simulation steps rendered per frame (higher = faster)")
|
||
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
|
||
|
||
if args.seed is None:
|
||
print("Auto-selecting a maze with a short solution path...")
|
||
chosen_seed, grid, bfs = find_short_seed(cfg.maze_rows, cfg.maze_cols)
|
||
print(f"Using seed={chosen_seed}, BFS path={int(bfs[0,0])} steps")
|
||
else:
|
||
grid = generate(cfg.maze_rows, cfg.maze_cols, seed=args.seed)
|
||
bfs = bfs_distances(grid, goal_row, goal_col)
|
||
path_len = int(bfs[0, 0])
|
||
print(f"Maze BFS path: {path_len} steps")
|
||
if path_len > cfg.maze_rows * cfg.maze_cols // 3:
|
||
print(f" [hint: long path ({path_len} steps) -- try omitting --seed for auto-selection]")
|
||
|
||
# Genome must be at least 3× the BFS path or GA will never converge
|
||
bfs_path = int(bfs[0, 0])
|
||
needed = bfs_path * 3
|
||
if needed > cfg.max_steps:
|
||
cfg.max_steps = needed
|
||
print(f"Genome length adapted: max_steps={cfg.max_steps} (BFS={bfs_path}×3)")
|
||
|
||
def make_agents(genomes):
|
||
agents = [Agent(g, 0, 0) for g in genomes]
|
||
for a in agents:
|
||
a._bfs = bfs
|
||
return agents
|
||
|
||
# Initialize population
|
||
nets = [Genome(cfg=cfg) for _ in range(cfg.population)]
|
||
|
||
print(f"Starting training: {args.generations} generations, population={cfg.population}")
|
||
|
||
if args.fast:
|
||
# --- Headless fast mode: train without display ---
|
||
print("Fast mode: rendering disabled during training.")
|
||
|
||
for gen in range(args.generations):
|
||
agents = make_agents(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)
|
||
|
||
# Init pygame now for the visual replay
|
||
print("\nTraining complete. Opening replay window...")
|
||
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 + " — Replay")
|
||
clock = pygame.time.Clock()
|
||
viz = Visualizer(screen, grid, cfg)
|
||
|
||
replay_agents = make_agents(nets)
|
||
|
||
ok = run_generation_visual(replay_agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, args.generations, args.steps_per_frame)
|
||
if ok:
|
||
print("Done. Press ESC or close window to exit.")
|
||
running = True
|
||
while running:
|
||
running = handle_events()
|
||
clock.tick(30)
|
||
pygame.quit()
|
||
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 = make_agents(nets)
|
||
|
||
ok = run_generation_visual(agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, gen, args.steps_per_frame)
|
||
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
|
||
print("\nTraining complete. Showing final generation replay...")
|
||
replay_agents = make_agents(nets)
|
||
|
||
viz.reset_trails()
|
||
run_generation_visual(replay_agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, args.generations, args.steps_per_frame)
|
||
|
||
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()
|