| """Scripted Snake agents: the oracle used for supervision and the baselines. |
| |
| All agents implement `act(state: SnakeState) -> str`. |
| |
| Path search is *time-aware*: a body cell at index i (head=0) of a snake of |
| length L becomes free after L - i moves, so a path may pass where the tail |
| currently is if it arrives late enough. This makes BFS paths exact as long as |
| no food is eaten along the way (food is only ever the final cell). |
| """ |
| from __future__ import annotations |
|
|
| from collections import deque |
| from functools import lru_cache |
|
|
| import numpy as np |
|
|
| from .env import ACTIONS, DELTAS, SnakeState |
|
|
|
|
| @lru_cache(maxsize=16) |
| def _neighbors(h: int, w: int) -> tuple[tuple[tuple[int, str], ...], ...]: |
| out = [] |
| for cell in range(h * w): |
| r, c = divmod(cell, w) |
| nb = [] |
| for a in ACTIONS: |
| dr, dc = DELTAS[a] |
| nr, nc = r + dr, c + dc |
| if 0 <= nr < h and 0 <= nc < w: |
| nb.append((nr * w + nc, a)) |
| out.append(tuple(nb)) |
| return tuple(out) |
|
|
|
|
| def _flat(state: SnakeState) -> tuple[list[int], int]: |
| w = state.width |
| body = [r * w + c for r, c in state.body] |
| food = state.food[0] * w + state.food[1] if state.food is not None else -1 |
| return body, food |
|
|
|
|
| def _free_times(body: list[int], area: int) -> list[int]: |
| ft = [0] * area |
| L = len(body) |
| for i, cell in enumerate(body): |
| ft[cell] = L - i |
| return ft |
|
|
|
|
| def _bfs(nbrs, area: int, start: int, ft: list[int], goal: int = -1): |
| dist = [-1] * area |
| parent = [-1] * area |
| dist[start] = 0 |
| q = deque([start]) |
| while q: |
| u = q.popleft() |
| du = dist[u] + 1 |
| for v, _ in nbrs[u]: |
| if dist[v] < 0 and du >= ft[v]: |
| dist[v] = du |
| parent[v] = u |
| if v == goal: |
| return dist, parent |
| q.append(v) |
| return dist, parent |
|
|
|
|
| def _reachable_count(nbrs, area: int, start: int, ft: list[int]) -> int: |
| dist, _ = _bfs(nbrs, area, start, ft) |
| return sum(1 for d in dist if d > 0) |
|
|
|
|
| def _tail_reachable(nbrs, area: int, body: list[int]) -> tuple[bool, int]: |
| if len(body) < 2: |
| return True, 0 |
| ft = _free_times(body, area) |
| tail = body[-1] |
| dist, _ = _bfs(nbrs, area, body[0], ft, goal=tail) |
| return dist[tail] > 0, dist[tail] |
|
|
|
|
| def safe_moves(state: SnakeState) -> list[str]: |
| """Moves that do not die on the very next step.""" |
| body, _ = _flat(state) |
| nbrs = _neighbors(state.height, state.width) |
| ft = _free_times(body, state.height * state.width) |
| return [a for v, a in nbrs[body[0]] if ft[v] <= 1] |
|
|
|
|
| class Agent: |
| name = "agent" |
|
|
| def act(self, state: SnakeState) -> str: |
| raise NotImplementedError |
|
|
| def reset(self, seed: int | None = None) -> None: |
| pass |
|
|
|
|
| class RandomAgent(Agent): |
| name = "random" |
|
|
| def __init__(self, seed: int = 0): |
| self.rng = np.random.default_rng(seed) |
|
|
| def reset(self, seed=None): |
| if seed is not None: |
| self.rng = np.random.default_rng(seed) |
|
|
| def act(self, state): |
| return ACTIONS[int(self.rng.integers(4))] |
|
|
|
|
| class GreedyFoodAgent(Agent): |
| """Minimise Manhattan distance to food; avoid *immediately* fatal moves when possible.""" |
| name = "greedy" |
|
|
| def act(self, state): |
| safe = safe_moves(state) or list(ACTIONS) |
| hr, hc = state.body[0] |
| fr, fc = state.food if state.food else (hr, hc) |
| return min(safe, key=lambda a: abs(hr + DELTAS[a][0] - fr) + abs(hc + DELTAS[a][1] - fc)) |
|
|
|
|
| class BFSAgent(Agent): |
| """Shortest (time-aware) path to food, no look-ahead; else the safe move with most space.""" |
| name = "bfs" |
|
|
| def act(self, state): |
| h, w = state.height, state.width |
| area = h * w |
| nbrs = _neighbors(h, w) |
| body, food = _flat(state) |
| ft = _free_times(body, area) |
| if food >= 0: |
| path = _path_to(nbrs, area, body[0], ft, food) |
| if path: |
| return _direction(w, body[0], path[0]) |
| return _max_space_move(nbrs, area, body) or ACTIONS[0] |
|
|
|
|
| class SurvivalAgent(Agent): |
| """Flood-fill heuristic: maximise reachable space, break ties by distance to food.""" |
| name = "survival" |
|
|
| def act(self, state): |
| h, w = state.height, state.width |
| area = h * w |
| nbrs = _neighbors(h, w) |
| body, food = _flat(state) |
| ft = _free_times(body, area) |
| best, best_key = None, None |
| fr, fc = divmod(food, w) if food >= 0 else divmod(body[0], w) |
| for v, a in nbrs[body[0]]: |
| if ft[v] > 1: |
| continue |
| nb = [v] + (body if v == food else body[:-1]) |
| space = _reachable_count(nbrs, area, v, _free_times(nb, area)) |
| r, c = divmod(v, w) |
| key = (min(space, len(nb) + 1), -(abs(r - fr) + abs(c - fc))) |
| if best_key is None or key > best_key: |
| best, best_key = a, key |
| return best or ACTIONS[0] |
|
|
|
|
| class OracleAgent(Agent): |
| """Strong teacher: eat only if the snake can still reach its tail afterwards, |
| otherwise chase the tail along the longest safe route, otherwise maximise space. |
| |
| `analyze` also returns the *set* of equally good moves. Ties are broken by a fixed |
| priority (UP > DOWN > LEFT > RIGHT) instead of BFS queue order, so the hard label |
| is a learnable function of the state (and this rule also scored best: 95.0 vs |
| 91.5-92.1 for "close the larger/smaller axis gap first", 60 games 10x10). |
| """ |
| name = "oracle" |
|
|
| def act(self, state): |
| return self.analyze(state)[0] |
|
|
| def analyze(self, state) -> tuple[str, list[str], str]: |
| """-> (move, optimal_moves, mode) with mode in {eat, tail, space, dead}.""" |
| h, w = state.height, state.width |
| area = h * w |
| nbrs = _neighbors(h, w) |
| body, food = _flat(state) |
| L = len(body) |
| ft = _free_times(body, area) |
| head = body[0] |
| safe = [(v, a) for v, a in nbrs[head] if ft[v] <= 1] |
| if not safe: |
| return ACTIONS[0], list(ACTIONS), "dead" |
| if food >= 0: |
| dist, _ = _bfs(nbrs, area, head, ft, goal=food) |
| d = dist[food] |
| if d > 0: |
| good = [] |
| for v, a in safe: |
| if v == food: |
| virtual = [v] + body |
| else: |
| nb = [v] + body[:-1] |
| path = _path_to(nbrs, area, v, _free_times(nb, area), food) |
| if len(path) != d - 1: |
| continue |
| virtual = (path[::-1] + nb)[: L + 1] |
| if len(virtual) >= area or _tail_reachable(nbrs, area, virtual)[0]: |
| good.append(a) |
| if good: |
| return good[0], good, "eat" |
| |
| cands = [] |
| for v, a in safe: |
| nb = [v] + (body if v == food else body[:-1]) |
| ok, dt = _tail_reachable(nbrs, area, nb) |
| if ok: |
| cands.append((dt, a)) |
| if cands: |
| best = max(dt for dt, _ in cands) |
| good = [a for dt, a in cands if dt == best] |
| return good[0], good, "tail" |
| spaces = [] |
| for v, a in safe: |
| nb = [v] + (body if v == food else body[:-1]) |
| spaces.append((_reachable_count(nbrs, area, v, _free_times(nb, area)), a)) |
| best = max(sp for sp, _ in spaces) |
| good = [a for sp, a in spaces if sp == best] |
| return good[0], good, "space" |
|
|
|
|
| def _path_to(nbrs, area, start, ft, goal) -> list[int]: |
| dist, parent = _bfs(nbrs, area, start, ft, goal) |
| if dist[goal] <= 0: |
| return [] |
| path = [goal] |
| while parent[path[-1]] != start: |
| path.append(parent[path[-1]]) |
| return path[::-1] |
|
|
|
|
| def _direction(w: int, a: int, b: int) -> str: |
| ar, ac = divmod(a, w) |
| br, bc = divmod(b, w) |
| return {(-1, 0): "UP", (1, 0): "DOWN", (0, -1): "LEFT", (0, 1): "RIGHT"}[(br - ar, bc - ac)] |
|
|
|
|
| def _max_space_move(nbrs, area, body) -> str | None: |
| ft = _free_times(body, area) |
| best, best_space = None, -1 |
| for v, a in nbrs[body[0]]: |
| if ft[v] > 1: |
| continue |
| nb = [v] + body[:-1] |
| space = _reachable_count(nbrs, area, v, _free_times(nb, area)) |
| if space > best_space: |
| best, best_space = a, space |
| return best |
|
|
|
|
| AGENTS = {c.name: c for c in (RandomAgent, GreedyFoodAgent, BFSAgent, SurvivalAgent, OracleAgent)} |
|
|
|
|
| def make_agent(name: str, seed: int = 0) -> Agent: |
| cls = AGENTS[name] |
| return cls(seed) if cls is RandomAgent else cls() |
|
|
