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

@@ -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
+    step = 0
-    for step in range(cfg.max_steps):
+    while step < cfg.max_steps:
-        for agent in agents:
+        # Advance simulation N steps per render frame
-            agent.step(grid, goal_row, goal_col, cfg)
+        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):
+        run_generation_visual(replay_agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, args.generations, args.steps_per_frame)
-            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

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)))
+        self.trail_surf.blit(self._fade_surf, (0, 0), special_flags=pygame.BLEND_RGBA_SUB)
-        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)
 
     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
+        self.screen.blit(self._gradient_bar, (self.hud_x, y))
-        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))
         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