File size: 8,402 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 | """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 # small survival bonus; keep tiny, starvation limit prevents stalling exploits
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
# "die": reversing into the neck is a collision (every choice has a distinct outcome)
# "ignore": reversal keeps the current heading (classic arcade behaviour)
reverse_action: str = "die"
starvation_steps: int | None = None # None -> 2 * area
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]] # head first
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()
# ------------------------------------------------------------------ core
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)
# place head so the body (trailing behind it) fits on the board
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-th free cell, found by scanning; cheap versus rejection sampling on full boards
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
# --------------------------------------------------------------- views
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
|