File size: 8,744 Bytes
026f63e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
"""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:  # pragma: no cover - interface
        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"
        # tail chasing: keep an escape route, prefer the move farthest from the tail
        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()