feat: initial maze-ml GA visualizer

README.md

@@ -0,0 +1,37 @@
+# maze-ml
+
+100 neural-network agents learn to navigate a maze using a **Genetic Algorithm** — animated in real time with pygame.
+
+Inspired by the "AI learns to..." YouTube Shorts format.
+
+## Quick start
+
+```bash
+pip install -r requirements.txt
+python main.py
+```
+
+## Options
+
+| Flag | Default | Description |
+|------|---------|-------------|
+| `--fast` | off | Train without rendering, show result at end |
+| `--generations` | 200 | Number of GA generations |
+| `--maze-size` | 25 | Maze dimensions (NxN) |
+| `--seed` | random | Maze seed for reproducibility |
+| `--speed` | 60 | FPS cap for visual mode |
+
+## Architecture
+
+- `config.py` — all hyperparameters as a dataclass
+- `maze.py` — iterative DFS maze generation (bitmask cells)
+- `agent.py` — Agent + tiny NeuralNet (8→12→4, pure numpy)
+- `genetic.py` — tournament selection, uniform crossover, gaussian mutation
+- `visualizer.py` — pygame renderer with trail decay, fitness-ranked colors, HUD
+- `main.py` — training loop + argparse
+
+## How it works
+
+Each generation, 100 agents simultaneously traverse the maze controlled by small neural networks (8 inputs → 12 hidden → 4 outputs). Fitness rewards progress toward the goal and reaching it. The top 30% of agents are selected as parents for the next generation via tournament selection.
+
+After ~20–50 generations, agents learn to reliably navigate the maze.

agent.py

@@ -0,0 +1,78 @@
+import numpy as np
+from maze import get_walls, is_passable, NORTH, SOUTH, EAST, WEST
+from config import MazeConfig
+
+class NeuralNet:
+    WEIGHT_COUNT = 8*12 + 12 + 12*4 + 4  # = 160
+
+    def __init__(self, weights=None):
+        if weights is None:
+            weights = np.random.randn(self.WEIGHT_COUNT) * 0.5
+        self.weights = weights.astype(np.float32)
+
+    def forward(self, x: np.ndarray) -> int:
+        # x shape: (8,) float32
+        # Slices: W1=weights[0:96].reshape(8,12), b1=weights[96:108]
+        #         W2=weights[108:156].reshape(12,4), b2=weights[156:160]
+        # Forward: h = ReLU(x @ W1 + b1), logits = h @ W2 + b2
+        # Return argmax(logits) — deterministic, no sampling
+        W1 = self.weights[0:96].reshape(8, 12)
+        b1 = self.weights[96:108]
+        W2 = self.weights[108:156].reshape(12, 4)
+        b2 = self.weights[156:160]
+        h = np.maximum(0.0, x @ W1 + b1)
+        logits = h @ W2 + b2
+        return int(np.argmax(logits))
+
+
+class Agent:
+    def __init__(self, net: NeuralNet, start_row: int, start_col: int):
+        self.net = net
+        self.row = start_row
+        self.col = start_col
+        self.alive = True
+        self.reached_goal = False
+        self.steps = 0
+        self.fitness = 0.0
+        self.trail: list[tuple[int,int]] = []  # pixel coords
+        self.dist_traveled_toward = 0.0
+
+    def _pixel_center(self, cfg) -> tuple[int, int]:
+        return (self.col * cfg.cell_size + cfg.cell_size // 2,
+                self.row * cfg.cell_size + cfg.cell_size // 2)
+
+    def get_inputs(self, grid, goal_row, goal_col, cfg) -> np.ndarray:
+        wN, wS, wE, wW = get_walls(grid, self.row, self.col)
+        dx = (goal_col - self.col) / cfg.maze_cols
+        dy = (goal_row - self.row) / cfg.maze_rows
+        px = self.col / cfg.maze_cols
+        py = self.row / cfg.maze_rows
+        return np.array([wN, wS, wE, wW, dx, dy, px, py], dtype=np.float32)
+
+    def step(self, grid, goal_row, goal_col, cfg):
+        if self.reached_goal:
+            return
+        x = self.get_inputs(grid, goal_row, goal_col, cfg)
+        action = self.net.forward(x)
+        dr = [-1, 1, 0, 0][action]
+        dc = [0, 0, 1, -1][action]
+        prev_dist = abs(goal_row - self.row) + abs(goal_col - self.col)
+        if is_passable(grid, self.row, self.col, action):
+            self.row += dr
+            self.col += dc
+            new_dist = abs(goal_row - self.row) + abs(goal_col - self.col)
+            self.dist_traveled_toward += max(0.0, prev_dist - new_dist)
+        self.steps += 1
+        self.trail.append(self._pixel_center(cfg))
+        if self.row == goal_row and self.col == goal_col:
+            self.reached_goal = True
+
+    def compute_fitness(self, goal_row, goal_col, cfg):
+        dist = abs(goal_row - self.row) + abs(goal_col - self.col)
+        self.fitness = (
+            self.dist_traveled_toward
+            - dist * 0.5
+            + (cfg.goal_bonus if self.reached_goal else 0.0)
+            - self.steps * cfg.step_penalty
+        )
+        return self.fitness

config.py

@@ -0,0 +1,31 @@
+from dataclasses import dataclass
+
+@dataclass
+class MazeConfig:
+    # Maze geometry
+    maze_cols: int = 25
+    maze_rows: int = 25
+    cell_size: int = 24
+    wall_width: int = 2
+
+    # GA / Training
+    population: int = 100
+    max_steps: int = 300
+    elite_frac: float = 0.30
+    mutation_std: float = 0.05
+    mutation_rate: float = 0.20
+
+    # Neural net
+    n_inputs: int = 8
+    n_hidden: int = 12
+    n_outputs: int = 4
+
+    # Fitness
+    goal_bonus: float = 1000.0
+    step_penalty: float = 0.01
+
+    # Rendering
+    fps: int = 60
+    trail_decay: float = 0.88
+    window_title: str = "Maze ML — GA Visualizer"
+    hud_font_size: int = 18

genetic.py

@@ -0,0 +1,56 @@
+import numpy as np
+from agent import Agent, NeuralNet
+from config import MazeConfig
+
+
+def tournament_select(weights: np.ndarray, fitnesses: np.ndarray, n_parents: int, k: int = 5) -> np.ndarray:
+    n = len(fitnesses)
+    selected = np.empty((n_parents, weights.shape[1]), dtype=np.float32)
+    for i in range(n_parents):
+        indices = np.random.choice(n, size=k, replace=False)
+        winner = indices[np.argmax(fitnesses[indices])]
+        selected[i] = weights[winner]
+    return selected
+
+
+def uniform_crossover(p1: np.ndarray, p2: np.ndarray) -> np.ndarray:
+    mask = np.random.rand(len(p1)) < 0.5
+    child = np.where(mask, p1, p2)
+    return child.astype(np.float32)
+
+
+def mutate(weights: np.ndarray, std: float = 0.05, rate: float = 0.20) -> np.ndarray:
+    mask = np.random.rand(len(weights)) < rate
+    noise = np.random.randn(len(weights)).astype(np.float32) * std
+    return (weights + mask * noise).astype(np.float32)
+
+
+def next_generation(agents: list[Agent], cfg: MazeConfig) -> list[NeuralNet]:
+    fitnesses = np.array([a.fitness for a in agents], dtype=np.float32)
+    all_weights = np.stack([a.net.weights for a in agents], axis=0)
+
+    n_elite = 10
+    elite_indices = np.argsort(fitnesses)[::-1][:n_elite]
+    elite_weights = all_weights[elite_indices]
+
+    n_pool = max(1, int(len(agents) * cfg.elite_frac))
+    pool_indices = np.argsort(fitnesses)[::-1][:n_pool]
+    pool_weights = all_weights[pool_indices]
+    pool_fitnesses = fitnesses[pool_indices]
+
+    n_offspring = cfg.population - n_elite
+    parents = tournament_select(pool_weights, pool_fitnesses, n_parents=n_offspring * 2)
+
+    new_nets: list[NeuralNet] = []
+
+    for w in elite_weights:
+        new_nets.append(NeuralNet(w.copy()))
+
+    for i in range(n_offspring):
+        p1 = parents[i * 2]
+        p2 = parents[i * 2 + 1]
+        child = uniform_crossover(p1, p2)
+        child = mutate(child, std=cfg.mutation_std, rate=cfg.mutation_rate)
+        new_nets.append(NeuralNet(child))
+
+    return new_nets

main.py

@@ -0,0 +1,138 @@
+import argparse
+import sys
+import os
+from config import MazeConfig
+from maze import generate
+from agent import Agent, NeuralNet
+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) -> 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)
+        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=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")
+    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
+
+    grid = generate(cfg.maze_rows, cfg.maze_cols, seed=args.seed)
+
+    # Initialize population
+    nets = [NeuralNet() for _ in range(cfg.population)]
+
+    print(f"Starting training: {args.generations} generations, population={cfg.population}")
+
+    if args.fast:
+        # --- Headless fast mode: no pygame, no display ---
+        print("Fast mode: rendering disabled during training.")
+
+        for gen in range(args.generations):
+            agents = [Agent(net, 0, 0) for net in 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)
+
+        print("Training complete.")
+        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 = [Agent(net, 0, 0) for net in nets]
+
+            ok = run_generation_visual(agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, gen)
+            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
+        best_nets = nets
+        print("\nTraining complete. Showing final generation replay...")
+        replay_agents = [Agent(net, 0, 0) for net in best_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)
+
+        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()

maze.py

@@ -0,0 +1,108 @@
+import numpy as np
+import random
+
+NORTH, SOUTH, EAST, WEST = 0, 1, 2, 3
+_DX = [0, 0, 1, -1]   # col delta for N,S,E,W
+_DY = [-1, 1, 0, 0]   # row delta
+_OPP = [1, 0, 3, 2]   # opposite direction
+
+
+def generate(rows: int, cols: int, seed: int | None = None) -> np.ndarray:
+    rng = random.Random(seed)
+    grid = np.full((rows, cols), 0b1111, dtype=np.uint8)
+    visited = np.zeros((rows, cols), dtype=bool)
+
+    start_r, start_c = 0, 0
+    visited[start_r, start_c] = True
+    stack = [(start_r, start_c)]
+
+    while stack:
+        row, col = stack[-1]
+        directions = [NORTH, SOUTH, EAST, WEST]
+        rng.shuffle(directions)
+
+        moved = False
+        for d in directions:
+            nr = row + _DY[d]
+            nc = col + _DX[d]
+            if 0 <= nr < rows and 0 <= nc < cols and not visited[nr, nc]:
+                # Carve passage: clear wall bit on current cell
+                grid[row, col] &= (~(1 << d)) & 0xFF
+                # Clear opposite wall bit on neighbour
+                grid[nr, nc] &= (~(1 << _OPP[d])) & 0xFF
+                visited[nr, nc] = True
+                stack.append((nr, nc))
+                moved = True
+                break
+
+        if not moved:
+            stack.pop()
+
+    return grid
+
+
+def get_walls(grid: np.ndarray, row: int, col: int) -> tuple[bool, bool, bool, bool]:
+    cell = int(grid[row, col])
+    wall_n = float((cell >> NORTH) & 1)
+    wall_s = float((cell >> SOUTH) & 1)
+    wall_e = float((cell >> EAST) & 1)
+    wall_w = float((cell >> WEST) & 1)
+    return (wall_n, wall_s, wall_e, wall_w)
+
+
+def is_passable(grid: np.ndarray, row: int, col: int, direction: int) -> bool:
+    cell = int(grid[row, col])
+    return not bool((cell >> direction) & 1)
+
+
+def render_static(grid: np.ndarray, cfg) -> 'pygame.Surface':
+    import pygame
+
+    rows, cols = grid.shape
+    cell_size = getattr(cfg, 'cell_size', 32)
+    wall_thickness = getattr(cfg, 'wall_thickness', 2)
+
+    width = cols * cell_size
+    height = rows * cell_size
+
+    surface = pygame.Surface((width, height))
+    surface.fill((15, 15, 20))
+
+    wall_color = (40, 40, 50)
+    t = wall_thickness
+
+    for row in range(rows):
+        for col in range(cols):
+            x = col * cell_size
+            y = row * cell_size
+            cell = int(grid[row, col])
+
+            # North wall
+            if (cell >> NORTH) & 1:
+                pygame.draw.rect(surface, wall_color, (x, y, cell_size, t))
+
+            # South wall
+            if (cell >> SOUTH) & 1:
+                pygame.draw.rect(surface, wall_color, (x, y + cell_size - t, cell_size, t))
+
+            # East wall
+            if (cell >> EAST) & 1:
+                pygame.draw.rect(surface, wall_color, (x + cell_size - t, y, t, cell_size))
+
+            # West wall
+            if (cell >> WEST) & 1:
+                pygame.draw.rect(surface, wall_color, (x, y, t, cell_size))
+
+    # Start cell (0, 0): small cyan circle
+    start_cx = cell_size // 2
+    start_cy = cell_size // 2
+    start_radius = max(2, cell_size // 6)
+    pygame.draw.circle(surface, (0, 220, 220), (start_cx, start_cy), start_radius)
+
+    # Goal cell (rows-1, cols-1): green filled circle
+    goal_x = (cols - 1) * cell_size + cell_size // 2
+    goal_y = (rows - 1) * cell_size + cell_size // 2
+    goal_radius = max(3, cell_size // 4)
+    pygame.draw.circle(surface, (0, 200, 80), (goal_x, goal_y), goal_radius)
+
+    return surface

requirements.txt

@@ -0,0 +1,2 @@
+pygame>=2.5.0
+numpy>=1.26.0

visualizer.py

@@ -0,0 +1,130 @@
+import pygame
+import numpy as np
+from config import MazeConfig
+from maze import render_static
+
+class Visualizer:
+    def __init__(self, screen: pygame.Surface, grid: np.ndarray, cfg: MazeConfig):
+        self.screen = screen
+        self.grid = grid
+        self.cfg = cfg
+        self.maze_w = cfg.maze_cols * cfg.cell_size
+        self.maze_h = cfg.maze_rows * cfg.cell_size
+        self.hud_x = self.maze_w + 10
+        # Pre-render static maze
+        self.maze_surface = render_static(grid, cfg)
+        # 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))
+        # Font
+        pygame.font.init()
+        self.font = pygame.font.SysFont("monospace", cfg.hud_font_size)
+        self.font_sm = pygame.font.SysFont("monospace", cfg.hud_font_size - 4)
+        self.clock = pygame.time.Clock()
+        self._goal_flash = 0
+        self._first_goal = False
+
+    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)
+
+    def _draw_hud(self, gen: int, best_fitness: float, step: int, alive: int, reached: int):
+        # Dark sidebar background
+        sidebar_rect = pygame.Rect(self.maze_w, 0, 200, self.maze_h)
+        pygame.draw.rect(self.screen, (10, 10, 16), sidebar_rect)
+
+        lines = [
+            ("GEN", f"{gen:>4}"),
+            ("STEP", f"{step:>4}/{self.cfg.max_steps}"),
+            ("ALIVE", f"{alive:>4}/{self.cfg.population}"),
+            ("GOAL", f"{reached:>4}"),
+            ("BEST", f"{best_fitness:>8.1f}"),
+        ]
+        y = 20
+        accent = (0, 220, 200)
+        for label, value in lines:
+            lbl_surf = self.font_sm.render(label, True, (100, 100, 120))
+            val_surf = self.font.render(value, True, accent)
+            self.screen.blit(lbl_surf, (self.hud_x, y))
+            self.screen.blit(val_surf, (self.hud_x, y + 16))
+            y += 52
+
+        # Legend
+        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))
+        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))
+
+    def render(self, agents, gen: int, step: int):
+        # Sort by fitness for ranking colors
+        sorted_agents = sorted(agents, key=lambda a: a.fitness)
+
+        # Check first goal
+        reached = sum(1 for a in agents if a.reached_goal)
+        if reached > 0 and not self._first_goal:
+            self._first_goal = True
+            self._goal_flash = 45
+
+        # Blit static maze
+        self.screen.blit(self.maze_surface, (0, 0))
+
+        # Fade and draw trails
+        self._fade_trails()
+        for rank, agent in enumerate(sorted_agents):
+            if agent.trail:
+                color = self._rank_color(rank, len(sorted_agents))
+                px, py = agent.trail[-1]
+                pygame.draw.circle(self.trail_surf, (*color, 130), (px, py), 3)
+        self.screen.blit(self.trail_surf, (0, 0))
+
+        # Draw agents on top
+        for rank, agent in enumerate(sorted_agents):
+            color = self._rank_color(rank, len(sorted_agents))
+            px = agent.col * self.cfg.cell_size + self.cfg.cell_size // 2
+            py = agent.row * self.cfg.cell_size + self.cfg.cell_size // 2
+            radius = 5 if agent.reached_goal else 3
+            pygame.draw.circle(self.screen, color, (px, py), radius)
+            if agent.reached_goal:
+                pygame.draw.circle(self.screen, (255, 255, 255), (px, py), radius + 2, 1)
+
+        # Goal flash overlay
+        if self._goal_flash > 0:
+            gx = (self.cfg.maze_cols - 1) * self.cfg.cell_size + self.cfg.cell_size // 2
+            gy = (self.cfg.maze_rows - 1) * self.cfg.cell_size + self.cfg.cell_size // 2
+            alpha = int(200 * self._goal_flash / 45)
+            flash_surf = pygame.Surface((self.cfg.cell_size, self.cfg.cell_size), pygame.SRCALPHA)
+            flash_surf.fill((255, 220, 0, alpha))
+            self.screen.blit(flash_surf, (gx - self.cfg.cell_size//2, gy - self.cfg.cell_size//2))
+            self._goal_flash -= 1
+
+        # HUD
+        best_fit = max((a.fitness for a in agents), default=0.0)
+        alive = sum(1 for a in agents if not a.reached_goal)
+        self._draw_hud(gen, best_fit, step, alive, reached)
+
+        pygame.display.flip()
+
+    def reset_trails(self):
+        self.trail_surf.fill((0, 0, 0, 0))
+        self._first_goal = False
+
+    def new_maze(self, grid, cfg):
+        self.grid = grid
+        self.cfg = cfg
+        self.maze_surface = render_static(grid, cfg)
+        self.reset_trails()