perf: pre-alloc fade surface, pre-render gradient bar, add --steps-per-frame
fix: fast mode now opens pygame replay window after headless training
This commit is contained in:
55
main.py
55
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
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user