| """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): |
| 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) |
| 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) |
|
|
