File size: 8,101 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 | """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()
|