From 20d9e99d5f2b6c25ec9aae867fe0c53117fd98f1 Mon Sep 17 00:00:00 2001 From: kerboul Date: Sat, 13 Jun 2026 15:26:02 +0200 Subject: [PATCH] perf: pre-alloc fade surface, pre-render gradient bar, add --steps-per-frame fix: fast mode now opens pygame replay window after headless training --- main.py | 55 ++++++++++++++++++++++++++++++++++----------------- visualizer.py | 30 +++++++++++++++++----------- 2 files changed, 56 insertions(+), 29 deletions(-) diff --git a/main.py b/main.py index 9d8204b..cc1199b 100644 --- a/main.py +++ b/main.py @@ -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 diff --git a/visualizer.py b/visualizer.py index d4beec7..228a6bc 100644 --- a/visualizer.py +++ b/visualizer.py @@ -16,6 +16,12 @@ class Visualizer: # Trail surface with per-pixel alpha self.trail_surf = pygame.Surface((self.maze_w, self.maze_h), pygame.SRCALPHA) self.trail_surf.fill((0, 0, 0, 0)) + # Pre-allocate fade surface — reused every frame (avoids per-frame Surface creation) + self._fade_surf = pygame.Surface((self.maze_w, self.maze_h), pygame.SRCALPHA) + subtract = max(1, int(255 * (1.0 - cfg.trail_decay))) + self._fade_surf.fill((0, 0, 0, subtract)) + # Pre-render static gradient bar for HUD + self._gradient_bar = self._make_gradient_bar(170) # Font pygame.font.init() self.font = pygame.font.SysFont("monospace", cfg.hud_font_size) @@ -24,15 +30,21 @@ class Visualizer: self._goal_flash = 0 self._first_goal = False + def _make_gradient_bar(self, width: int) -> pygame.Surface: + surf = pygame.Surface((width, 10)) + for i in range(width): + t = i / width + surf.set_at((i, 0), (int(255 * (1 - t)), int(255 * t), 30)) + for y in range(1, 10): + surf.set_at((i, y), (int(255 * (1 - t)), int(255 * t), 30)) + return surf + def _rank_color(self, rank: int, total: int) -> tuple[int, int, int]: t = rank / max(1, total - 1) return (int(255 * t), int(255 * (1.0 - t)), 30) def _fade_trails(self): - subtract = max(1, int(255 * (1.0 - self.cfg.trail_decay))) - fade = pygame.Surface((self.maze_w, self.maze_h), pygame.SRCALPHA) - fade.fill((0, 0, 0, subtract)) - self.trail_surf.blit(fade, (0, 0), special_flags=pygame.BLEND_RGBA_SUB) + self.trail_surf.blit(self._fade_surf, (0, 0), special_flags=pygame.BLEND_RGBA_SUB) def _draw_hud(self, gen: int, best_fitness: float, step: int, alive: int, reached: int): # Dark sidebar background @@ -55,20 +67,16 @@ class Visualizer: self.screen.blit(val_surf, (self.hud_x, y + 16)) y += 52 - # Legend + # Legend — pre-rendered gradient bar (no per-frame line drawing) y += 20 legend_label = self.font_sm.render("RANK", True, (80, 80, 100)) self.screen.blit(legend_label, (self.hud_x, y)) y += 18 - bar_w = 170 - for i in range(bar_w): - t = i / bar_w - r, g = int(255 * t), int(255 * (1 - t)) - pygame.draw.line(self.screen, (r, g, 30), (self.hud_x + bar_w - i, y), (self.hud_x + bar_w - i, y + 10)) + self.screen.blit(self._gradient_bar, (self.hud_x, y)) best_lbl = self.font_sm.render("best", True, (0, 255, 30)) wrst_lbl = self.font_sm.render("worst", True, (255, 0, 30)) self.screen.blit(best_lbl, (self.hud_x, y + 14)) - self.screen.blit(wrst_lbl, (self.hud_x + bar_w - 35, y + 14)) + self.screen.blit(wrst_lbl, (self.hud_x + 120, y + 14)) def render(self, agents, gen: int, step: int): # Sort by fitness for ranking colors