kerboul/maze-ml
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

perf: pre-alloc fade surface, pre-render gradient bar, add --steps-per-frame

Browse Source 
fix: fast mode now opens pygame replay window after headless training
This commit is contained in:
Ethan Puyaubreau
2026-06-13 15:26:02 +02:00
parent e0a0f5cd7f
commit 20d9e99d5f
2 changed files with 56 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

55
main.py
View File

@@ -25,12 +25,17 @@ def run_generation_headless(agents, grid, goal_row, goal_col, cfg):
agent.step(grid, goal_row, goal_col, cfg)
def run_generation_visual(agents, grid, goal_row, goal_col, cfg, viz, clock, fps, gen) -> bool:
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."""
import pygame
for step in range(cfg.max_steps):
for agent in agents:
agent.step(grid, goal_row, goal_col, cfg)
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)
@@ -45,6 +50,7 @@ def main():
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")
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)
@@ -59,7 +65,7 @@ def main():
print(f"Starting training: {args.generations} generations, population={cfg.population}")
if args.fast:
# --- Headless fast mode: no pygame, no display ---
# --- Headless fast mode: train without display ---
print("Fast mode: rendering disabled during training.")
for gen in range(args.generations):
@@ -75,7 +81,28 @@ def main():
nets = next_generation(agents, cfg)
print("Training complete.")
# 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 = [Agent(net, 0, 0) for net in 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:
@@ -95,7 +122,7 @@ def main():
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)
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()
@@ -112,18 +139,10 @@ def main():
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]
replay_agents = [Agent(net, 0, 0) for net in 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)
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