"""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 # ------------------------------------------------------------- control 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 # ----------------------------------------------------------------- loop 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) # yield to the HTTP threads 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()