| """Headless, deterministic Snake. |
| |
| Cells are flattened indices `r * width + c`. The body is stored head-first in a |
| deque plus a bytearray occupancy grid, so a step is O(1) and thousands of games |
| run without any rendering. |
| """ |
| from __future__ import annotations |
|
|
| from collections import deque |
| from dataclasses import dataclass, field |
|
|
| import numpy as np |
|
|
| ACTIONS = ("UP", "DOWN", "LEFT", "RIGHT") |
| DELTAS = {"UP": (-1, 0), "DOWN": (1, 0), "LEFT": (0, -1), "RIGHT": (0, 1)} |
| OPPOSITE = {"UP": "DOWN", "DOWN": "UP", "LEFT": "RIGHT", "RIGHT": "LEFT"} |
|
|
|
|
| @dataclass |
| class RewardConfig: |
| food: float = 1.0 |
| death: float = -1.0 |
| step: float = 0.0 |
| win: float = 1.0 |
|
|
| @classmethod |
| def from_dict(cls, d: dict | None) -> "RewardConfig": |
| return cls(**(d or {})) |
|
|
|
|
| @dataclass |
| class SnakeConfig: |
| height: int = 10 |
| width: int = 10 |
| initial_length: int = 3 |
| |
| |
| reverse_action: str = "die" |
| starvation_steps: int | None = None |
| max_steps: int | None = None |
| reward: RewardConfig = field(default_factory=RewardConfig) |
|
|
| @classmethod |
| def from_dict(cls, d: dict | None) -> "SnakeConfig": |
| d = dict(d or {}) |
| board = d.pop("board_size", None) |
| if board is not None: |
| d["height"] = d["width"] = int(board) |
| d["reward"] = RewardConfig.from_dict(d.get("reward")) |
| return cls(**d) |
|
|
| @property |
| def starvation_limit(self) -> int: |
| return self.starvation_steps or 2 * self.height * self.width |
|
|
|
|
| @dataclass |
| class SnakeState: |
| """Immutable-ish snapshot, the `state` MiniJEV receives for Snake.""" |
| height: int |
| width: int |
| body: list[tuple[int, int]] |
| food: tuple[int, int] | None |
| direction: str |
| steps_since_food: int = 0 |
|
|
| def to_dict(self) -> dict: |
| return {"height": self.height, "width": self.width, "body": [list(p) for p in self.body], |
| "food": list(self.food) if self.food else None, "direction": self.direction, |
| "steps_since_food": self.steps_since_food} |
|
|
| @classmethod |
| def from_dict(cls, d: dict) -> "SnakeState": |
| return cls(d["height"], d["width"], [tuple(p) for p in d["body"]], |
| tuple(d["food"]) if d.get("food") else None, d["direction"], d.get("steps_since_food", 0)) |
|
|
|
|
| class SnakeEnv: |
| def __init__(self, config: SnakeConfig | None = None, seed: int | None = None): |
| self.cfg = config or SnakeConfig() |
| self.h, self.w = self.cfg.height, self.cfg.width |
| self.area = self.h * self.w |
| self.rng = np.random.default_rng(seed) |
| self.body: deque[int] = deque() |
| self.occ = bytearray(self.area) |
| self.reset() |
|
|
| |
| def seed(self, seed: int) -> None: |
| self.rng = np.random.default_rng(seed) |
|
|
| def reset(self, seed: int | None = None) -> SnakeState: |
| if seed is not None: |
| self.seed(seed) |
| self.body.clear() |
| self.occ = bytearray(self.area) |
| self.direction = ACTIONS[int(self.rng.integers(4))] |
| dr, dc = DELTAS[self.direction] |
| L = min(self.cfg.initial_length, min(self.h, self.w) // 2) |
| |
| r = int(self.rng.integers(L, self.h - L)) if self.h > 2 * L else self.h // 2 |
| c = int(self.rng.integers(L, self.w - L)) if self.w > 2 * L else self.w // 2 |
| for i in range(L): |
| cell = (r - dr * i) * self.w + (c - dc * i) |
| self.body.append(cell) |
| self.occ[cell] = 1 |
| self.score = 0 |
| self.steps = 0 |
| self.steps_since_food = 0 |
| self.done = False |
| self.end_reason: str | None = None |
| self._place_food() |
| return self.state() |
|
|
| def _place_food(self) -> None: |
| free = self.area - len(self.body) |
| if free <= 0: |
| self.food = -1 |
| return |
| |
| k = int(self.rng.integers(free)) |
| if free > self.area // 4: |
| while True: |
| cell = int(self.rng.integers(self.area)) |
| if not self.occ[cell]: |
| self.food = cell |
| return |
| for cell in range(self.area): |
| if not self.occ[cell]: |
| if k == 0: |
| self.food = cell |
| return |
| k -= 1 |
|
|
| def step(self, action: str | int) -> tuple[float, bool]: |
| """Returns (reward, done). Call `state()` for the observation.""" |
| if self.done: |
| raise RuntimeError("step() on finished game; call reset()") |
| if not isinstance(action, str): |
| action = ACTIONS[int(action)] |
| if self.cfg.reverse_action == "ignore" and len(self.body) > 1 and action == OPPOSITE[self.direction]: |
| action = self.direction |
| rw = self.cfg.reward |
| dr, dc = DELTAS[action] |
| head = self.body[0] |
| r, c = divmod(head, self.w) |
| nr, nc = r + dr, c + dc |
| self.direction = action |
| self.steps += 1 |
| self.steps_since_food += 1 |
| if not (0 <= nr < self.h and 0 <= nc < self.w): |
| return self._finish("wall", rw.death) |
| cell = nr * self.w + nc |
| eating = cell == self.food |
| tail = self.body[-1] |
| if self.occ[cell] and not (cell == tail and not eating): |
| return self._finish("body", rw.death) |
| if not eating: |
| self.body.pop() |
| self.occ[tail] = 0 |
| self.body.appendleft(cell) |
| self.occ[cell] = 1 |
| reward = rw.step |
| if eating: |
| self.score += 1 |
| self.steps_since_food = 0 |
| reward += rw.food |
| if len(self.body) == self.area: |
| return self._finish("win", reward + rw.win) |
| self._place_food() |
| if self.steps_since_food >= self.cfg.starvation_limit: |
| return self._finish("starvation", rw.death) |
| if self.cfg.max_steps and self.steps >= self.cfg.max_steps: |
| self.done, self.end_reason = True, "max_steps" |
| return reward, True |
| return reward, False |
|
|
| def _finish(self, reason: str, reward: float) -> tuple[float, bool]: |
| self.done, self.end_reason = True, reason |
| return reward, True |
|
|
| |
| def state(self) -> SnakeState: |
| w = self.w |
| return SnakeState(self.h, self.w, [divmod(b, w) for b in self.body], |
| divmod(self.food, w) if self.food >= 0 else None, |
| self.direction, self.steps_since_food) |
|
|
| def render_ascii(self) -> str: |
| grid = [["." for _ in range(self.w)] for _ in range(self.h)] |
| for i, b in enumerate(self.body): |
| r, c = divmod(b, self.w) |
| grid[r][c] = "H" if i == 0 else "o" |
| if self.food >= 0: |
| r, c = divmod(self.food, self.w) |
| grid[r][c] = "*" |
| return "\n".join("".join(row) for row in grid) |
|
|
|
|
| class VectorSnakeEnv: |
| """N independent games stepped together (auto-reset); used for batched |
| inference in RL/evaluation. Deterministic given `seed`.""" |
|
|
| def __init__(self, n: int, config: SnakeConfig | None = None, seed: int = 0): |
| self.envs = [SnakeEnv(config, seed=seed + i) for i in range(n)] |
| self.next_seed = seed + n |
| self.n = n |
|
|
| def states(self) -> list[SnakeState]: |
| return [e.state() for e in self.envs] |
|
|
| def step(self, actions) -> tuple[np.ndarray, np.ndarray, list[dict]]: |
| rewards = np.zeros(self.n, dtype=np.float32) |
| dones = np.zeros(self.n, dtype=bool) |
| infos: list[dict] = [{} for _ in range(self.n)] |
| for i, (env, a) in enumerate(zip(self.envs, actions)): |
| rewards[i], dones[i] = env.step(a) |
| if dones[i]: |
| infos[i] = {"score": env.score, "steps": env.steps, "end_reason": env.end_reason} |
| env.reset(seed=self.next_seed) |
| self.next_seed += 1 |
| return rewards, dones, infos |
|
|
