fix: genome GA converges via BFS fitness + two-rate prefix-aware mutation

- Replace Manhattan distance with BFS distances (actual maze path length)
- Switch from NN to direct sequence genome (what YouTube Shorts actually use)
- Mutation-only GA (crossover breaks positional maze paths)
- Two-rate mutation: low rate before best-step (preserve prefix), high after (explore tail)
- Auto-seed selection finds maze with short BFS path
- Default maze 15x15, cell_size=36, max_steps=150
- Typically converges in 5-15 generations

agent.py

@@ -1,78 +1,99 @@
 import numpy as np
-from maze import get_walls, is_passable, NORTH, SOUTH, EAST, WEST
+from maze import get_walls, is_passable, bfs_distances
 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)
+class Genome:
+    """Direct action-sequence genome: each gene is a move (0=N 1=S 2=E 3=W).
+    The GA evolves the sequence directly — no NN weights to get in the way.
+    This is the approach used in "AI learns maze" viral videos.
+    """
+    N_ACTIONS = 4
 
-    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))
+    def __init__(self, genes: np.ndarray | None = None, cfg: MazeConfig | None = None):
+        steps = cfg.max_steps if cfg is not None else 300
+        if genes is None:
+            genes = np.random.randint(0, self.N_ACTIONS, size=steps, dtype=np.uint8)
+        self.genes = genes
+
+    @property
+    def weights(self) -> np.ndarray:
+        """Alias so genetic.py can treat Genome like NeuralNet."""
+        return self.genes.astype(np.float32)
+
+    @classmethod
+    def from_weights(cls, w: np.ndarray) -> "Genome":
+        return cls(genes=np.clip(np.round(w), 0, 3).astype(np.uint8))
 
 
 class Agent:
-    def __init__(self, net: NeuralNet, start_row: int, start_col: int):
-        self.net = net
+    def __init__(self, genome: Genome, start_row: int, start_col: int):
+        self.net = genome          # kept as .net so visualizer/main.py don't need changes
+        self.genome = genome
         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
+        self.trail: list[tuple[int, int]] = []
+        self._visited: set[tuple[int, int]] = {(start_row, start_col)}
+        self._visit_counts: dict[tuple[int, int], int] = {(start_row, start_col): 1}
+        self.wall_hits = 0
+        self._min_bfs: int = 10000   # closest BFS approach to goal this episode
+        self._best_step: int = 0     # genome index when min_bfs was achieved
+        self._bfs: np.ndarray | None = None  # set externally before stepping
 
-    def _pixel_center(self, cfg) -> tuple[int, int]:
+    def _pixel_center(self, cfg: MazeConfig) -> 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:
+    def step(self, grid, goal_row: int, goal_col: int, cfg: MazeConfig):
+        if self.reached_goal or self.steps >= len(self.genome.genes):
             return
-        x = self.get_inputs(grid, goal_row, goal_col, cfg)
-        action = self.net.forward(x)
+        action = int(self.genome.genes[self.steps])
         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)
+        else:
+            self.wall_hits += 1
+
         self.steps += 1
+        cell = (self.row, self.col)
+        self._visited.add(cell)
+        self._visit_counts[cell] = self._visit_counts.get(cell, 0) + 1
         self.trail.append(self._pixel_center(cfg))
+
+        if self._bfs is not None:
+            bfs_d = int(self._bfs[self.row, self.col])
+            if bfs_d >= 0 and bfs_d < self._min_bfs:
+                self._min_bfs = bfs_d
+                self._best_step = self.steps  # record when we achieved best approach
+
         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)
+    def compute_fitness(self, goal_row: int, goal_col: int, cfg: MazeConfig) -> float:
+        # BFS distance from start (0,0) to goal = max_bfs
+        # best BFS approach: the closest the agent ever got (through actual maze corridors)
+        max_bfs = int(self._bfs[0, 0]) if self._bfs is not None else (goal_row + goal_col)
+        best_bfs = self._min_bfs if self._min_bfs < 10000 else max_bfs
+        unique_cells = len(self._visited)
+        total_visits = sum(self._visit_counts.values())
+        revisit_penalty = total_visits - unique_cells
+
         self.fitness = (
-            self.dist_traveled_toward
-            - dist * 0.5
+            (max_bfs - best_bfs) * 12.0             # real maze proximity (primary)
+            + unique_cells * 1.0                     # exploration
+            - revisit_penalty * 0.1                  # penalize looping
+            - self.wall_hits * 0.05
             + (cfg.goal_bonus if self.reached_goal else 0.0)
-            - self.steps * cfg.step_penalty
+            + ((cfg.max_steps - self.steps) * 5.0 if self.reached_goal else 0.0)
         )
         return self.fitness
+
+
+# Keep NeuralNet as alias so any external code importing it still works
+NeuralNet = Genome

config.py

@@ -3,21 +3,23 @@ from dataclasses import dataclass
 @dataclass
 class MazeConfig:
     # Maze geometry
-    maze_cols: int = 25
-    maze_rows: int = 25
-    cell_size: int = 24
+    maze_cols: int = 15
+    maze_rows: int = 15
+    cell_size: int = 36
     wall_width: int = 2
 
     # GA / Training
     population: int = 100
-    max_steps: int = 300
+    max_steps: int = 150
     elite_frac: float = 0.30
-    mutation_std: float = 0.05
-    mutation_rate: float = 0.20
+    n_elite: int = 3          # few elites → maintains population diversity
+    mutation_std: float = 0.08
+    mutation_rate: float = 0.35
+    temperature: float = 0.7  # softmax sampling temperature (lower = more greedy)
 
-    # Neural net
-    n_inputs: int = 8
-    n_hidden: int = 12
+    # Neural net — 14 inputs: 4 walls + 2 goal dir + 2 pos + 4 last-action + 1 blocked + 1 revisit
+    n_inputs: int = 14
+    n_hidden: int = 16
     n_outputs: int = 4
 
     # Fitness

genetic.py

@@ -1,56 +1,87 @@
 import numpy as np
-from agent import Agent, NeuralNet
+from agent import Agent, Genome
 from config import MazeConfig
 
+_AgentInfo = tuple  # (genes, best_step)
 
-def tournament_select(weights: np.ndarray, fitnesses: np.ndarray, n_parents: int, k: int = 5) -> np.ndarray:
+
+def tournament_select(genes_list: list[np.ndarray], fitnesses: np.ndarray,
+                      n_parents: int, k: int = 5) -> list[np.ndarray]:
     n = len(fitnesses)
-    selected = np.empty((n_parents, weights.shape[1]), dtype=np.float32)
-    for i in range(n_parents):
+    selected = []
+    for _ in range(n_parents):
         indices = np.random.choice(n, size=k, replace=False)
         winner = indices[np.argmax(fitnesses[indices])]
-        selected[i] = weights[winner]
+        selected.append(genes_list[winner].copy())
     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 single_point_crossover(g1: np.ndarray, g2: np.ndarray) -> np.ndarray:
+    point = np.random.randint(1, len(g1))
+    return np.concatenate([g1[:point], g2[point:]])
 
 
-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 mutate_seq(genes: np.ndarray, rate: float) -> np.ndarray:
+    """Replace each gene with a random action with probability rate."""
+    mask = np.random.rand(len(genes)) < rate
+    noise = np.random.randint(0, 4, size=len(genes), dtype=np.uint8)
+    result = genes.copy()
+    result[mask] = noise[mask]
+    return result
 
 
-def next_generation(agents: list[Agent], cfg: MazeConfig) -> list[NeuralNet]:
+def mutate_two_rate(genes: np.ndarray, pivot: int,
+                    low_rate: float = 0.005, high_rate: float = 0.30) -> np.ndarray:
+    """Prefix-aware mutation: low rate before pivot (preserve good path),
+    high rate after pivot (aggressively explore the stuck region)."""
+    result = genes.copy()
+    n = len(genes)
+    pivot = min(pivot, n)
+    # Prefix: preserve
+    mask_pre = np.random.rand(pivot) < low_rate
+    result[:pivot][mask_pre] = np.random.randint(0, 4, mask_pre.sum(), dtype=np.uint8)
+    # Tail: explore
+    tail = n - pivot
+    if tail > 0:
+        mask_tail = np.random.rand(tail) < high_rate
+        result[pivot:][mask_tail] = np.random.randint(0, 4, mask_tail.sum(), dtype=np.uint8)
+    return result
+
+
+def next_generation(agents: list[Agent], cfg: MazeConfig) -> list[Genome]:
     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)
+    genes_list = [a.genome.genes for a in agents]
 
-    n_elite = 10
-    elite_indices = np.argsort(fitnesses)[::-1][:n_elite]
-    elite_weights = all_weights[elite_indices]
+    # Elites: top n_elite survive unchanged
+    n_elite = cfg.n_elite
+    elite_idx = np.argsort(fitnesses)[::-1][:n_elite]
+    new_genomes: list[Genome] = [Genome(genes_list[i].copy()) for i in elite_idx]
 
-    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]
+    # Parents pool: top elite_frac
+    n_pool = max(2, int(len(agents) * cfg.elite_frac))
+    pool_idx = np.argsort(fitnesses)[::-1][:n_pool]
+    pool = [genes_list[i] for i in pool_idx]
+    pool_fits = fitnesses[pool_idx]
+
+    # Build pool with (genes, best_step) pairs
+    pool_agents = [agents[i] for i in pool_idx]
+    pool_genes = [a.genome.genes for a in pool_agents]
+    pool_pivots = [a._best_step for a in pool_agents]
 
     n_offspring = cfg.population - n_elite
-    parents = tournament_select(pool_weights, pool_fitnesses, n_parents=n_offspring * 2)
+    parent_indices = [
+        pool_idx[np.argmax(pool_fits[np.random.choice(len(pool_fits), 5, replace=False)])]
+        for _ in range(n_offspring)
+    ]
 
-    new_nets: list[NeuralNet] = []
+    for idx in parent_indices:
+        parent_agent = agents[idx]
+        child = mutate_two_rate(
+            parent_agent.genome.genes,
+            pivot=parent_agent._best_step,
+            low_rate=0.003,
+            high_rate=0.30,
+        )
+        new_genomes.append(Genome(child))
 
-    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
+    return new_genomes

main.py

@@ -2,8 +2,8 @@ import argparse
 import sys
 import os
 from config import MazeConfig
-from maze import generate
-from agent import Agent, NeuralNet
+from maze import generate, bfs_distances, find_short_seed
+from agent import Agent, Genome
 from genetic import next_generation
 
 
@@ -57,10 +57,26 @@ def main():
     goal_row = cfg.maze_rows - 1
     goal_col = cfg.maze_cols - 1
 
+    if args.seed is None:
+        print("Auto-selecting a maze with a short solution path...")
+        chosen_seed, grid, bfs = find_short_seed(cfg.maze_rows, cfg.maze_cols)
+        print(f"Using seed={chosen_seed}, BFS path={int(bfs[0,0])} steps")
+    else:
         grid = generate(cfg.maze_rows, cfg.maze_cols, seed=args.seed)
+        bfs = bfs_distances(grid, goal_row, goal_col)
+        path_len = int(bfs[0, 0])
+        print(f"Maze BFS path: {path_len} steps")
+        if path_len > cfg.maze_rows * cfg.maze_cols // 3:
+            print(f"  [hint: long path ({path_len} steps) -- try omitting --seed for auto-selection]")
+
+    def make_agents(genomes):
+        agents = [Agent(g, 0, 0) for g in genomes]
+        for a in agents:
+            a._bfs = bfs
+        return agents
 
     # Initialize population
-    nets = [NeuralNet() for _ in range(cfg.population)]
+    nets = [Genome(cfg=cfg) for _ in range(cfg.population)]
 
     print(f"Starting training: {args.generations} generations, population={cfg.population}")
 
@@ -69,7 +85,7 @@ def main():
         print("Fast mode: rendering disabled during training.")
 
         for gen in range(args.generations):
-            agents = [Agent(net, 0, 0) for net in nets]
+            agents = make_agents(nets)
             run_generation_headless(agents, grid, goal_row, goal_col, cfg)
 
             for agent in agents:
@@ -94,7 +110,8 @@ def main():
         clock = pygame.time.Clock()
         viz = Visualizer(screen, grid, cfg)
 
-        replay_agents = [Agent(net, 0, 0) for net in nets]
+        replay_agents = make_agents(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.")
@@ -120,7 +137,7 @@ def main():
         viz = Visualizer(screen, grid, cfg)
 
         for gen in range(args.generations):
-            agents = [Agent(net, 0, 0) for net in nets]
+            agents = make_agents(nets)
 
             ok = run_generation_visual(agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, gen, args.steps_per_frame)
             if not ok:
@@ -140,7 +157,8 @@ def main():
 
         # Final visual replay with best generation
         print("\nTraining complete. Showing final generation replay...")
-        replay_agents = [Agent(net, 0, 0) for net in nets]
+        replay_agents = make_agents(nets)
+
         viz.reset_trails()
         run_generation_visual(replay_agents, grid, goal_row, goal_col, cfg, viz, clock, args.speed, args.generations, args.steps_per_frame)
 

maze.py

@@ -1,5 +1,6 @@
 import numpy as np
 import random
+from collections import deque
 
 NORTH, SOUTH, EAST, WEST = 0, 1, 2, 3
 _DX = [0, 0, 1, -1]   # col delta for N,S,E,W
@@ -41,6 +42,44 @@ def generate(rows: int, cols: int, seed: int | None = None) -> np.ndarray:
     return grid
 
 
+def find_short_seed(rows: int, cols: int, max_ratio: float = 6.5, tries: int = 200) -> tuple[int, np.ndarray, np.ndarray]:
+    """Return (seed, grid, bfs_dist) with a short goal path.
+    Tries up to `tries` seeds, picks the one whose BFS path < rows*cols/max_ratio.
+    Falls back to the best found if none qualifies.
+    """
+    target = int(rows * cols / max_ratio)
+    best = None
+    for seed in range(tries):
+        grid = generate(rows, cols, seed=seed)
+        bfs = bfs_distances(grid, rows - 1, cols - 1)
+        path_len = int(bfs[0, 0])
+        if best is None or path_len < best[0]:
+            best = (path_len, seed, grid, bfs)
+        if path_len <= target:
+            return seed, grid, bfs
+    _, seed, grid, bfs = best
+    return seed, grid, bfs
+
+
+def bfs_distances(grid: np.ndarray, goal_row: int, goal_col: int) -> np.ndarray:
+    """BFS distance from every cell to (goal_row, goal_col) through passable walls.
+    Returns int array shape (rows, cols); unreachable cells = -1.
+    """
+    rows, cols = grid.shape
+    dist = np.full((rows, cols), -1, dtype=np.int32)
+    dist[goal_row, goal_col] = 0
+    q = deque([(goal_row, goal_col)])
+    while q:
+        r, c = q.popleft()
+        for d in range(4):
+            if is_passable(grid, r, c, d):
+                nr, nc = r + _DY[d], c + _DX[d]
+                if 0 <= nr < rows and 0 <= nc < cols and dist[nr, nc] == -1:
+                    dist[nr, nc] = dist[r, c] + 1
+                    q.append((nr, nc))
+    return dist
+
+
 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)