File size: 4,801 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 | """Snake as a MiniJEV task: SnakeState -> multi-channel board image.
Channels (board padded by a 1-cell wall border; smaller boards in a batch are
padded with more wall, which is semantically exact):
0 wall 1 head 2 body (excl. head) 3 tail 4 food
5 time-until-free: (L - i) / area for body segment i (head=0)
6 row offset from the head (r - r_head) / max(H, W)
7 col offset from the head (c - c_head) / max(H, W)
Channel 5 lets the model reason about which cells will open up as the snake
moves - the same information the oracle's time-aware search uses. Channels 6-7
make every cell egocentric ("the food is 3 up, 2 left"), which removes the
long learning plateau of discovering relative geometry from absolute positions.
"""
from __future__ import annotations
from typing import Sequence
import numpy as np
from minijev.tasks import Task, ragged_pack, ragged_take
from .env import ACTIONS, SnakeState
N_CHANNELS = 8
DIRS = {a: i for i, a in enumerate(ACTIONS)}
SEMANTIC_CHOICES = {"UP": "move up", "DOWN": "move down", "LEFT": "move left", "RIGHT": "move right"}
class SnakeTask(Task):
"""task.choice_style: "semantic" -> choices "move up/down/left/right" (v2 models);
"legacy" (default, what pre-v2 checkpoints were trained with) -> "UP/DOWN/LEFT/RIGHT"."""
name = "snake"
state_encoder = "grid"
default_question = "Choose the next move"
@property
def choice_alias(self) -> dict[str, str]:
return SEMANTIC_CHOICES if self.cfg.get("choice_style", "legacy") == "semantic" else {}
@property
def default_choices(self) -> tuple[str, ...]:
return tuple(self.model_label(a) for a in ACTIONS)
def encoder_class(self):
from .encoder import SnakeStateEncoder
return SnakeStateEncoder
def encoder_config(self) -> dict:
return {"in_channels": N_CHANNELS}
def states_to_arrays(self, states: Sequence[SnakeState]) -> dict[str, np.ndarray]:
bodies = [[r * s.width + c for r, c in s.body] for s in states]
flat, offsets = ragged_pack(bodies, np.int16)
return {
"body": flat,
"body_offsets": offsets,
"food": np.array([s.food[0] * s.width + s.food[1] if s.food else -1 for s in states], np.int16),
"height": np.array([s.height for s in states], np.uint8),
"width": np.array([s.width for s in states], np.uint8),
"direction": np.array([DIRS[s.direction] for s in states], np.uint8),
}
def example_states(self, n=4, seed=0):
from .agents import OracleAgent
from .env import SnakeConfig, SnakeEnv
states, oracle = [], OracleAgent()
env = SnakeEnv(SnakeConfig.from_dict({"board_size": self.cfg.get("board_size", 10)}), seed=seed)
while len(states) < n:
states.append(env.state())
for _ in range(7): # spread samples along the game
env.step(oracle.act(env.state()))
if env.done:
env.reset()
return states
def encode_arrays(self, arrays, idx):
return {"state_grid": encode_grid(arrays, np.asarray(idx))}
def encode_grid(arrays: dict[str, np.ndarray], idx: np.ndarray) -> np.ndarray:
B = len(idx)
hs = arrays["height"][idx].astype(np.int64)
ws = arrays["width"][idx].astype(np.int64)
H, W = int(hs.max()) + 2, int(ws.max()) + 2
x = np.zeros((B, N_CHANNELS, H, W), dtype=np.float32)
x[:, 0] = 1.0
for h, w in set(zip(hs.tolist(), ws.tolist())):
sel = np.nonzero((hs == h) & (ws == w))[0]
x[sel, 0, 1:h + 1, 1:w + 1] = 0.0
cells, lengths = ragged_take(arrays["body"], arrays["body_offsets"], idx)
cells = cells.astype(np.int64)
sample = np.repeat(np.arange(B), lengths)
pos = np.arange(len(cells)) - np.repeat(np.cumsum(lengths) - lengths, lengths) # index from head
Ls = np.repeat(lengths, lengths)
w_s = ws[sample]
r, c = cells // w_s + 1, cells % w_s + 1
head = pos == 0
x[sample[head], 1, r[head], c[head]] = 1.0
x[sample[~head], 2, r[~head], c[~head]] = 1.0
tail = pos == Ls - 1
x[sample[tail], 3, r[tail], c[tail]] = 1.0
x[sample, 5, r, c] = (Ls - pos) / (hs[sample] * w_s)
hr, hc = np.zeros(B, np.int64), np.zeros(B, np.int64)
hr[sample[head]], hc[sample[head]] = r[head], c[head]
scale = np.maximum(hs, ws).astype(np.float32)[:, None, None]
x[:, 6] = (np.arange(H)[None, :, None] - hr[:, None, None]) / scale
x[:, 7] = (np.arange(W)[None, None, :] - hc[:, None, None]) / scale
food = arrays["food"][idx].astype(np.int64)
has = np.nonzero(food >= 0)[0]
x[has, 4, food[has] // ws[has] + 1, food[has] % ws[has] + 1] = 1.0
return x
def make_task(cfg: dict) -> SnakeTask:
return SnakeTask(cfg)
|