| """Snake demo: a background thread plays the game, the browser renders snapshots. |
| |
| Stdlib only (http.server). The simulation runs independently of rendering, so |
| "max" speed is limited by the agent, not by the browser. |
| |
| minijev snake demo --model runs/.../final.onnx |
| minijev snake play # you play with the arrow keys |
| """ |
| from __future__ import annotations |
|
|
| import json |
| import threading |
| import time |
| import webbrowser |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| from pathlib import Path |
|
|
| from ..agents import AGENTS, make_agent |
| from ..env import ACTIONS, OPPOSITE, SnakeConfig, SnakeEnv |
|
|
| STATIC = Path(__file__).parent / "static" |
| SPEEDS = {"realtime": 8.0, "fast": 60.0, "max": None} |
| QUESTION = "Choose the next move" |
|
|
|
|
| class Session: |
| def __init__(self, model: str | None, board_size: int, agent: str): |
| self.lock = threading.Lock() |
| self.decider = None |
| if model: |
| from minijev.inference.runtime import MiniJEV |
| self.decider = MiniJEV.load(model, threads=1) |
| self.model_path = model |
| self.board_size = board_size |
| self.agent_name = agent if (agent != "minijev" or self.decider) else "oracle" |
| self.agent = None if self.agent_name in ("minijev", "human") else make_agent(self.agent_name) |
| self.speed = "realtime" |
| self.seed = 1 |
| self.env = SnakeEnv(SnakeConfig.from_dict({"board_size": board_size}), seed=self.seed) |
| self.human_dir = self.env.direction |
| self.last = {"choice": None, "scores": {}, "confidence": None, "latency_ms": None} |
| self.lat_avg = None |
| self.games: list[int] = [] |
| self.best = 0 |
| self.last_end = None |
| self.steps_total = 0 |
| self.rate_t, self.rate_n, self.sim_rate = time.perf_counter(), 0, 0.0 |
| self.pause_until = 0.0 |
| self.running = True |
|
|
| |
| def control(self, msg: dict) -> None: |
| with self.lock: |
| if msg.get("speed") in SPEEDS: |
| self.speed = msg["speed"] |
| if "agent" in msg: |
| a = msg["agent"] |
| if a == "minijev" and not self.decider: |
| pass |
| elif a in AGENTS or a in ("minijev", "human"): |
| self.agent_name = a |
| self.agent = None if a in ("minijev", "human") else make_agent(a) |
| self._reset(clear_stats=True) |
| if msg.get("key") in ACTIONS: |
| k = msg["key"] |
| if not (len(self.env.body) > 1 and k == OPPOSITE[self.env.direction]): |
| self.human_dir = k |
| if msg.get("reset"): |
| self._reset() |
| if "board_size" in msg: |
| self.board_size = int(msg["board_size"]) |
| self._reset(clear_stats=True) |
|
|
| def _reset(self, clear_stats: bool = False): |
| self.seed += 1 |
| self.env = SnakeEnv(SnakeConfig.from_dict({"board_size": self.board_size}), seed=self.seed) |
| self.human_dir = self.env.direction |
| if clear_stats: |
| self.games, self.best = [], 0 |
|
|
| |
| def decide(self): |
| s = self.env.state() |
| t = time.perf_counter() |
| if self.agent_name == "minijev": |
| task = self.decider.task |
| d = self.decider.decide(s, QUESTION, list(task.default_choices)) |
| choice, conf = task.action_of(d.choice), d.confidence |
| scores = {task.action_of(k): v for k, v in d.scores.items()} |
| elif self.agent_name == "human": |
| choice, scores, conf = self.human_dir, {a: float(a == self.human_dir) for a in ACTIONS}, None |
| else: |
| choice = self.agent.act(s) |
| scores, conf = {a: float(a == choice) for a in ACTIONS}, None |
| lat = (time.perf_counter() - t) * 1000 |
| self.lat_avg = lat if self.lat_avg is None else 0.9 * self.lat_avg + 0.1 * lat |
| self.last = {"choice": choice, "scores": scores, "confidence": conf, "latency_ms": lat} |
| return choice |
|
|
| def run(self): |
| while self.running: |
| now = time.perf_counter() |
| if now < self.pause_until: |
| time.sleep(0.02) |
| continue |
| with self.lock: |
| choice = self.decide() |
| self.env.step(choice) |
| self.steps_total += 1 |
| self.rate_n += 1 |
| if self.env.done: |
| self.games.append(self.env.score) |
| self.best = max(self.best, self.env.score) |
| self.last_end = {"score": self.env.score, "reason": self.env.end_reason, "steps": self.env.steps} |
| self._reset() |
| if self.speed != "max": |
| self.pause_until = time.perf_counter() + 0.8 |
| if now - self.rate_t > 0.5: |
| self.sim_rate = self.rate_n / (now - self.rate_t) |
| self.rate_t, self.rate_n = now, 0 |
| speed = SPEEDS[self.speed] if self.agent_name != "human" else min(SPEEDS[self.speed] or 12, 12) |
| if speed: |
| time.sleep(max(0.0, 1.0 / speed - (time.perf_counter() - now))) |
| else: |
| time.sleep(0) |
|
|
| def snapshot(self) -> dict: |
| with self.lock: |
| s = self.env.state() |
| g = self.games |
| return {"height": s.height, "width": s.width, "body": s.body, "food": s.food, "direction": s.direction, |
| "score": self.env.score, "step": self.env.steps, "agent": self.agent_name, "speed": self.speed, |
| "decision": self.last, "latency_avg_ms": self.lat_avg, "sim_steps_per_sec": self.sim_rate, |
| "games_played": len(g), "avg_score": sum(g) / len(g) if g else None, "best_score": self.best, |
| "last_game": self.last_end, "has_model": self.decider is not None, |
| "model": Path(self.model_path).name if self.model_path else None, |
| "agents": ["minijev"] * bool(self.decider) + list(AGENTS) + ["human"]} |
|
|
|
|
| def serve(model: str | None = None, board_size: int = 10, port: int = 8765, agent: str = "minijev", |
| open_browser: bool = True): |
| session = Session(model, board_size, agent) |
| threading.Thread(target=session.run, daemon=True).start() |
|
|
| class Handler(BaseHTTPRequestHandler): |
| def log_message(self, *a): |
| pass |
|
|
| def _send(self, body: bytes, ctype: str, code: int = 200): |
| self.send_response(code) |
| self.send_header("Content-Type", ctype) |
| self.send_header("Cache-Control", "no-store") |
| self.send_header("Content-Length", str(len(body))) |
| self.end_headers() |
| self.wfile.write(body) |
|
|
| def do_GET(self): |
| if self.path.startswith("/api/state"): |
| self._send(json.dumps(session.snapshot()).encode(), "application/json") |
| elif self.path in ("/", "/index.html"): |
| self._send((STATIC / "index.html").read_bytes(), "text/html; charset=utf-8") |
| else: |
| self._send(b"not found", "text/plain", 404) |
|
|
| def do_POST(self): |
| if self.path.startswith("/api/control"): |
| n = int(self.headers.get("Content-Length", 0)) |
| session.control(json.loads(self.rfile.read(n) or b"{}")) |
| self._send(b"{}", "application/json") |
| else: |
| self._send(b"not found", "text/plain", 404) |
|
|
| httpd = ThreadingHTTPServer(("127.0.0.1", port), Handler) |
| url = f"http://127.0.0.1:{port}/" |
| print(f"MiniJEV Snake demo on {url} (agent={session.agent_name}, model={model}) Ctrl-C to stop") |
| if open_browser: |
| threading.Timer(0.5, lambda: webbrowser.open(url)).start() |
| try: |
| httpd.serve_forever() |
| except KeyboardInterrupt: |
| pass |
| finally: |
| session.running = False |
| httpd.server_close() |
|
|
