diff --git a/.gitignore b/.gitignore index 67b5764..e72abde 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,8 @@ DerivedData/ *.xcodeproj .DS_Store +.venv/ +Tools/doom/sauerkraut/models/ +__pycache__/ +_vizdoom.ini +/*.mp4 diff --git a/Package.swift b/Package.swift index 1a75b92..0110320 100644 --- a/Package.swift +++ b/Package.swift @@ -54,6 +54,7 @@ let package = Package( .executableTarget(name: "ImageSortCheck", dependencies: ["ImageSort", "FluidUse"]), .executableTarget(name: "ImageSortDemo", dependencies: ["ImageSort"], exclude: ["README.md"]), .executableTarget(name: "KevCheck", dependencies: ["FluidUse"]), + .executableTarget(name: "GLiClassServe", dependencies: ["FluidUse"]), .executableTarget( name: "KevGuessWhoDemo", dependencies: ["FluidUse", "SortAnything"], exclude: ["README.md", "demo.sh"]), .testTarget( diff --git a/Sources/GLiClassServe/main.swift b/Sources/GLiClassServe/main.swift new file mode 100644 index 0000000..c43da33 --- /dev/null +++ b/Sources/GLiClassServe/main.swift @@ -0,0 +1,68 @@ +import FluidUse +import Foundation + +/// GLiClass decisions over stdin/stdout, one JSON object per line, for harnesses outside Swift +/// (Tools/doom). Prints `{"ready": true, ...}` once the model is loaded. +/// +/// in: {"text": "...", "labels": ["a", "b"], "prompt": "..."} +/// out: {"index": 0, "probabilities": [0.9, 0.1], "ms": 1.8} +/// +/// swift run -c release GLiClassServe [--precision lut8] [--model-dir ] +@main +struct GLiClassServe { + struct Request: Decodable { + let text: String + let labels: [String] + let prompt: String? + } + + struct Response: Encodable { + var index: Int? + var probabilities: [Float]? + var ms: Double? + var error: String? + } + + static func main() async throws { + let arguments = Array(CommandLine.arguments.dropFirst()) + func option(_ name: String) -> String? { + arguments.firstIndex(of: name).flatMap { arguments.indices.contains($0 + 1) ? arguments[$0 + 1] : nil } + } + let precision = option("--precision") ?? "lut8" + let configuration = GLiClassManager.Configuration(lengths: [128], precision: precision) + let started = Date() + let manager: GLiClassManager + if let directory = option("--model-dir") { + manager = try await GLiClassManager.load( + from: URL(fileURLWithPath: directory), configuration: configuration) + } else { + manager = try await GLiClassManager.load(configuration: configuration) + } + _ = try await manager.classify(text: "warm up", labels: ["yes", "no"]) + let loadSeconds = Date().timeIntervalSince(started) + emit("{\"ready\": true, \"precision\": \"\(precision)\", \"load_s\": \(loadSeconds)}") + + let decoder = JSONDecoder() + let encoder = JSONEncoder() + while let line = readLine() { + guard !line.isEmpty else { continue } + var response = Response() + do { + let request = try decoder.decode(Request.self, from: Data(line.utf8)) + let begin = DispatchTime.now().uptimeNanoseconds + let answer = try await manager.classify( + text: request.text, labels: request.labels, prompt: request.prompt) + response.ms = Double(DispatchTime.now().uptimeNanoseconds - begin) / 1e6 + response.index = answer.selectedIndex + response.probabilities = answer.probabilities + } catch { + response.error = error.localizedDescription + } + emit(String(decoding: try encoder.encode(response), as: UTF8.self)) + } + } + + private static func emit(_ line: String) { + FileHandle.standardOutput.write(Data((line + "\n").utf8)) + } +} diff --git a/Tools/doom/README.md b/Tools/doom/README.md new file mode 100644 index 0000000..05d6c79 --- /dev/null +++ b/Tools/doom/README.md @@ -0,0 +1,28 @@ +# GLiClass plays ViZDoom `defend_the_center` + +Headless check in the Jev style: the model never sees pixels. Code reads ViZDoom's labels buffer +and writes the state as text (health, ammo, each visible monster's bearing, distance, and whether it +is on the crosshair). GLiClass picks attack / turn left / turn right through `GLiClassServe`, and the +engine holds the action for 3 tics (about 12 decisions per second at 35 Hz). + +```bash +uv venv -p 3.12 .venv && uv pip install -p .venv vizdoom numpy +swift build -c release --product GLiClassServe +.venv/bin/python Tools/doom/defend_the_center.py --seeds 20 +``` + +Seeds 1–20, M5 Pro, GLiClass Edge Apps v2 L128: + +| Policy | Kills | Survived | Agrees with aimer | ms per call | +| --- | ---: | ---: | ---: | ---: | +| Hand-coded aimer | 14.80 | 26.9 s | — | — | +| Random | 1.00 | 9.0 s | 32% | — | +| GLiClass LUT8, bare labels | 0.00 | 8.3 s | 71% | 3.6 | +| GLiClass LUT8, consequence labels | 14.00 | 25.3 s | 91% | 4.0 (p95 ≤ 5.7) | +| GLiClass fp16, consequence labels | 13.75 | 24.6 s | 88% | 3.3 | + +With bare labels GLiClass turns the right way but never fires. With labels that state each action's +consequence ("attack: shoots the monster on the crosshair", "turn left: moves the crosshair off the +monster"), as the Tetris and 2048 shortlists do, it comes within one kill of the aimer (6 wins, 5 ties, +9 losses per seed). The consequence labels carry most of the decision, so this shows the loop is fast +enough for a 35 Hz shooter, not that GLiClass plays Doom better than 20 lines of code. diff --git a/Tools/doom/defend_the_center.py b/Tools/doom/defend_the_center.py new file mode 100644 index 0000000..ade8d06 --- /dev/null +++ b/Tools/doom/defend_the_center.py @@ -0,0 +1,217 @@ +"""Headless ViZDoom defend_the_center check: GLiClass vs a hand-coded aimer on the same seeds. + +The model never sees pixels. Code turns the labels buffer into text (health, ammo, each visible +monster's bearing and distance, whether it is on the crosshair); the policy picks attack / turn +left / turn right, and the engine holds that action for --tics tics. + + uv pip install vizdoom numpy + swift build -c release --product GLiClassServe + python Tools/doom/defend_the_center.py --policies aimer,random,gliclass,gliclass-described --seeds 20 +""" + +import argparse +import json +import math +import os +import random +import statistics +import subprocess +import sys +import time + +import vizdoom as vzd + +ACTIONS = ["attack", "turn left", "turn right"] +BUTTONS = {"attack": [0, 0, 1], "turn left": [1, 0, 0], "turn right": [0, 1, 0]} +PROMPT = ( + "You are playing Doom, standing in the middle of a round room. Monsters walk toward you. " + "Shoot a monster only when it is on the crosshair; otherwise turn toward the nearest one. " + "Choose the next action." +) +SCREEN_CENTER = 160 + + +class GLiClassProcess: + def __init__(self, binary, precision): + self.proc = subprocess.Popen( + [binary, "--precision", precision], stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True, bufsize=1 + ) + ready = json.loads(self.proc.stdout.readline()) + if not ready.get("ready"): + raise RuntimeError(f"GLiClassServe failed to start: {ready}") + self.load_s = ready["load_s"] + + def classify(self, text, labels): + self.proc.stdin.write(json.dumps({"text": text, "labels": labels, "prompt": PROMPT}) + "\n") + response = json.loads(self.proc.stdout.readline()) + if response.get("error"): + raise RuntimeError(response["error"]) + return response["index"], response["probabilities"], response["ms"] + + def close(self): + self.proc.stdin.close() + self.proc.wait() + + +def observe(state): + ammo, health, angle, px, py = state.game_variables[:5] + monsters = [] + for label in state.labels: + if label.object_name == "DoomPlayer": + continue + dx, dy = label.object_position_x - px, label.object_position_y - py + bearing = (math.degrees(math.atan2(dy, dx)) - angle + 180) % 360 - 180 # + = left + monsters.append( + { + "bearing": bearing, + "distance": math.hypot(dx, dy), + "on_crosshair": label.x <= SCREEN_CENTER <= label.x + label.width, + } + ) + monsters.sort(key=lambda m: m["distance"]) + return {"ammo": int(ammo), "health": int(health), "monsters": monsters} + + +def describe(obs): + parts = [f"Health {obs['health']}, ammo {obs['ammo']}."] + if not obs["monsters"]: + parts.append("No monsters in view.") + else: + parts.append(f"{len(obs['monsters'])} monsters in view.") + for i, m in enumerate(obs["monsters"][:3]): + side = "left" if m["bearing"] > 0 else "right" + aim = "on the crosshair" if m["on_crosshair"] else "not on the crosshair" + name = "Nearest" if i == 0 else "Next" + parts.append(f"{name}: {abs(m['bearing']):.0f} degrees {side}, {m['distance']:.0f} units away, {aim}.") + return " ".join(parts) + + +def described_labels(obs): + """Labels that state each action's consequence, like the Tetris and 2048 shortlists.""" + target = next((m for m in obs["monsters"] if m["on_crosshair"]), None) + if obs["ammo"] == 0: + attack = "attack: out of ammo, does nothing" + elif target: + attack = f"attack: shoots the monster on the crosshair {target['distance']:.0f} units away" + else: + attack = "attack: no monster on the crosshair, wastes a bullet" + nearest = obs["monsters"][0] if obs["monsters"] else None + if target: + left = "turn left: moves the crosshair off the monster" + right = "turn right: moves the crosshair off the monster" + elif nearest is None: + left, right = "turn left: search for monsters", "turn right: search for monsters" + elif nearest["bearing"] > 0: + left = f"turn left: toward the nearest monster {abs(nearest['bearing']):.0f} degrees left" + right = "turn right: away from the nearest monster" + else: + left = "turn left: away from the nearest monster" + right = f"turn right: toward the nearest monster {abs(nearest['bearing']):.0f} degrees right" + return [attack, left, right] + + +def aimer(obs): + if obs["ammo"] > 0 and any(m["on_crosshair"] for m in obs["monsters"]): + return "attack" + if obs["monsters"]: + return "turn left" if obs["monsters"][0]["bearing"] > 0 else "turn right" + return "turn left" + + +def make_game(): + game = vzd.DoomGame() + game.load_config(os.path.join(vzd.scenarios_path, "defend_the_center.cfg")) + game.set_window_visible(False) + game.set_labels_buffer_enabled(True) + for variable in ("ANGLE", "POSITION_X", "POSITION_Y", "KILLCOUNT"): + game.add_available_game_variable(getattr(vzd.GameVariable, variable)) + game.init() + return game + + +def run_episode(game, seed, policy, tics, model, rng): + game.set_seed(seed) + game.new_episode() + decisions, shots, wasted, disagreements, ms = 0, 0, 0, 0, [] + kills = 0 + while not game.is_episode_finished(): + state = game.get_state() + obs = observe(state) + kills = int(state.game_variables[5]) + if policy == "aimer": + action = aimer(obs) + elif policy == "random": + action = rng.choice(ACTIONS) + else: + labels = described_labels(obs) if policy == "gliclass-described" else ACTIONS + index, _, elapsed = model.classify(describe(obs), labels) + action = ACTIONS[index] + ms.append(elapsed) + if action != aimer(obs): + disagreements += 1 + if action == "attack" and obs["ammo"] > 0: + shots += 1 + if not any(m["on_crosshair"] for m in obs["monsters"]): + wasted += 1 + decisions += 1 + game.make_action(BUTTONS[action], tics) + tics_alive = game.get_episode_time() + died = game.is_player_dead() + return { + "seed": seed, + "kills": kills, + "tics": tics_alive, + "died": died, + "decisions": decisions, + "shots": shots, + "wasted_shots": wasted, + "aimer_disagreement": disagreements / max(decisions, 1), + "ms_median": statistics.median(ms) if ms else None, + "ms_p95": sorted(ms)[int(0.95 * (len(ms) - 1))] if ms else None, + } + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--policies", default="aimer,random,gliclass,gliclass-described") + parser.add_argument("--seeds", type=int, default=20) + parser.add_argument("--first-seed", type=int, default=1) + parser.add_argument("--tics", type=int, default=3, help="tics each decision is held (35 tics = 1 s)") + parser.add_argument("--precision", default="lut8") + parser.add_argument("--binary", default=".build/release/GLiClassServe") + parser.add_argument("--out", default=None) + args = parser.parse_args() + + policies = args.policies.split(",") + model = GLiClassProcess(args.binary, args.precision) if any(p.startswith("gliclass") for p in policies) else None + if model: + print(f"GLiClass {args.precision} loaded in {model.load_s:.1f} s", file=sys.stderr) + game = make_game() + results = {} + for policy in policies: + rng = random.Random(0) + episodes = [] + started = time.time() + for seed in range(args.first_seed, args.first_seed + args.seeds): + episodes.append(run_episode(game, seed, policy, args.tics, model, rng)) + results[policy] = episodes + kills = [e["kills"] for e in episodes] + line = ( + f"{policy:20s} kills {statistics.mean(kills):5.2f} (min {min(kills)}, max {max(kills)})" + f" survived {statistics.mean(e['tics'] for e in episodes) / 35:5.1f} s" + f" wasted {sum(e['wasted_shots'] for e in episodes) / max(sum(e['shots'] for e in episodes), 1):5.1%}" + f" vs aimer {statistics.mean(e['aimer_disagreement'] for e in episodes):5.1%}" + ) + if episodes[0]["ms_median"] is not None: + line += f" {statistics.median(e['ms_median'] for e in episodes):.2f} ms/call" + print(line + f" ({time.time() - started:.0f} s wall)") + game.close() + if model: + model.close() + if args.out: + with open(args.out, "w") as f: + json.dump({"tics": args.tics, "precision": args.precision, "results": results}, f, indent=1) + + +if __name__ == "__main__": + main() diff --git a/Tools/doom/sauerkraut/README.md b/Tools/doom/sauerkraut/README.md new file mode 100644 index 0000000..fe73fb6 --- /dev/null +++ b/Tools/doom/sauerkraut/README.md @@ -0,0 +1,64 @@ +# SauerkrautLM-Doom-MultiVec on Core ML + +[SauerkrautLM-Doom-MultiVec-1.3M](https://huggingface.co/VAGOsolutions/SauerkrautLM-Doom-MultiVec-1.3M) +(Apache 2.0, VAGO solutions) converted to Core ML and playing ViZDoom `defend_the_center`. The Core ML +models are on Hugging Face at +[FluidInference/sauerkrautlm-doom-coreml](https://huggingface.co/FluidInference/sauerkrautlm-doom-coreml); +the demo downloads them on first run. + +```bash +Tools/doom/sauerkraut/demo.sh # game window + asitop + decision log +Tools/doom/sauerkraut/demo.sh --record doom.mp4 --episodes 1 # also writes an mp4 of the window +.venv/bin/python Tools/doom/sauerkraut/play.py --check 100 # headless score check, seeds 10000-10099 +``` + +`demo.sh` also opens one Ghostty window split into two rows with tmux: `sudo asitop` on top (type +your password there) and `tail -f /tmp/doom-demo.log` below, one colored line per decision with the +action, the four probabilities, and the Core ML call time. Without Ghostty or tmux it opens two Terminal.app windows. The game window is resizable and waits on each episode's first frame until you press space +(`--autostart` skips that). Keys: space start / pause, n next episode, q quit. + +## What the model reads + +1026 tokens per frame: a character per cell of a 40×25 grid (plus newlines, `[CLS]`, `[SEP]`) and a +learned embedding of each cell's depth, 16 bins. Upstream's ASCII path overflows `uint8` and writes +`@` for every cell, in training and at inference, so the depth bins carry all of the information. The +demo panel draws those bins. It is not reading pixels or game state. + +## Results, M5 Pro, seeds 10000–10099, 4 tics per decision, 2100-tic episodes + +| Model | Mean kills | Mean survival | Full 60 s | ms per decision | +| --- | ---: | ---: | ---: | ---: | +| PyTorch (upstream) | 20.42 (sd 5.32) | 50.5 s | 31 / 100 | GPU (MPS) 10.5 fp32, 8.4 fp16; CPU 57.7 (1 thread) | +| Core ML fp32, GPU | 20.42 (sd 5.32), identical on 100/100 seeds | 50.5 s | 31 / 100 | 2.5 | +| Core ML fp16, 1026 tokens, GPU (demo default) | 20.54 (sd 5.27) | 50.6 s | 35 / 100 | 1.2 | + +Speeds are median ms per decision on the same frame, idle machine, called from Python. On the GPU, +Core ML is about 4× faster than PyTorch MPS at fp32 and about 7× at fp16. + +For reference on the same seeds: our hand-coded aimer 13.05 kills and GLiClass with consequence +labels 11.98 (both in `Tools/doom/defend_the_center.py`, three buttons, 320×240), random 1.26. An +independent 1000-episode evaluation of the PyTorch model +([tiny-doom-defender](https://huggingface.co/spaces/anakin87/tiny-doom-defender)) reports 20.38. + +The Neural Engine is slower than the GPU for this model (7.8–8.4 ms); `--units CPU_AND_NE` to compare. + +The window renders every game tic, and ViZDoom's rendering consumes game randomness, so a seed does +not replay the headless path exactly. The play is equally strong: seeds 10000–10029 average 19.70 +kills / 50.5 s rendered against 20.37 / 50.7 s headless. Rendered seeds 10016 and 10005 reach 25 kills +and survive the full 60 s. + +## Conversion + +`convert.py` re-implements the ModernBERT forward with static masks (Hugging Face's mask construction +does not trace through coremltools) and exports fp32/fp16 with an attention mask at 1100 tokens. +`convert_ane.py` exports the demo model: fixed 1026 tokens (no frame has padding) and one-hot matmul +embeddings. Both check parity against upstream PyTorch on real frames and need the upstream package: + +```bash +uv venv -p 3.12 convenv && uv pip install -p convenv vizdoom==1.3.0 torch==2.7.* transformers==4.56.2 \ + coremltools git+https://github.com/VAGOsolutions/SauerkrautLM-Doom-MultiVec +cd Tools/doom/sauerkraut && ../../../convenv/bin/python convert.py && ../../../convenv/bin/python convert_ane.py +# then pass one with --model path/to/SauerkrautDoom_L1026_fp16.mlpackage +``` + +`sauer_eval.py` is the PyTorch/Core ML episode evaluator used for the table. diff --git a/Tools/doom/sauerkraut/convert.py b/Tools/doom/sauerkraut/convert.py new file mode 100644 index 0000000..7339fd4 --- /dev/null +++ b/Tools/doom/sauerkraut/convert.py @@ -0,0 +1,82 @@ +"""Core ML export of SauerkrautLM-Doom-MultiVec-1.3M: static-mask re-implementation of the traced forward.""" +import os, sys, numpy as np, torch, torch.nn as nn, torch.nn.functional as F, coremltools as ct +from huggingface_hub import snapshot_download +sys.path.insert(0, os.path.dirname(__file__)) +from sauer_eval import load_torch, make_game, Encoder, MODEL_ID, MAX_TOKENS, FRAME_SKIP + +NEG = -1e4 + +def rope(theta, head_dim, n): + inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + f = torch.outer(torch.arange(n, dtype=torch.float32), inv) + e = torch.cat((f, f), -1) + return e.cos()[None, None], e.sin()[None, None] + +def rot_half(x, h): + return torch.cat((-x[..., h:], x[..., :h]), -1) + +class Static(nn.Module): + def __init__(self, clf, n=MAX_TOKENS): + super().__init__() + enc = clf.encoder; c = enc.config + self.enc, self.clf = enc, clf + self.h, self.nh = c.hidden_size, c.num_attention_heads + self.hd = self.h // self.nh + self.local = [i % c.global_attn_every_n_layers != 0 for i in range(c.num_hidden_layers)] + gc, gs = rope(c.global_rope_theta, self.hd, n); lc, ls = rope(c.local_rope_theta, self.hd, n) + self.register_buffer("gcos", gc); self.register_buffer("gsin", gs) + self.register_buffer("lcos", lc); self.register_buffer("lsin", ls) + idx = torch.arange(n) + win = ((idx[None] - idx[:, None]).abs() > c.local_attention // 2).float() * NEG + self.register_buffer("win", win[None, None]) + + def forward(self, ids, mask, depth): + e = self.enc + x = e.embeddings.tok_embeddings(ids.long()) + e.depth_embedding(depth.long()) + x = e.embeddings.norm(x) + pad = ((1.0 - mask.float()) * NEG)[:, None, None, :] + local_bias = pad + self.win + for i, layer in enumerate(e.layers): + a = layer.attn + qkv = a.Wqkv(layer.attn_norm(x)).view(1, -1, 3, self.nh, self.hd) + q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) + cos, sin = (self.lcos, self.lsin) if self.local[i] else (self.gcos, self.gsin) + q = q * cos + rot_half(q, self.hd // 2) * sin; k = k * cos + rot_half(k, self.hd // 2) * sin + w = torch.matmul(q, k.transpose(2, 3)) * (self.hd ** -0.5) + (local_bias if self.local[i] else pad) + o = torch.matmul(torch.softmax(w, -1), v).transpose(1, 2).reshape(1, -1, self.h) + x = x + a.Wo(o) + x = x + layer.mlp(layer.mlp_norm(x)) + x = e.final_norm(x) + s = self.clf.attn_weight(x).squeeze(-1) + (1.0 - mask.float()) * NEG + pooled = (x * torch.softmax(s, 1).unsqueeze(-1)).sum(1) + return self.clf.classifier(pooled) + +if __name__ == "__main__": + model_dir = snapshot_download(MODEL_ID) + clf = load_torch(model_dir) + static = Static(clf).eval() + g = make_game(); enc = Encoder(model_dir); frames = [] + rng = np.random.default_rng(0) + for seed in (1, 2): + g.set_seed(seed); g.new_episode() + while not g.is_episode_finished(): + s = g.get_state(); frames.append(enc(s.screen_buffer, s.depth_buffer)) + g.make_action([int(x) for x in rng.integers(0, 2, 4)], FRAME_SKIP) + g.close() + frames = frames[::3][:80] + T = lambda f: [torch.from_numpy(x) for x in f] + with torch.no_grad(): + ref = torch.stack([clf(*[t.long() for t in T(f)[:2]], depth_ids=T(f)[2].long())["logits"][0] for f in frames]).numpy() + mine = torch.stack([static(*T(f))[0] for f in frames]).numpy() + print(f"static vs HF torch: max|dlogit| {np.abs(mine-ref).max():.2e}, lengths {sorted(set(int(f[1].sum()) for f in frames))}") + with torch.no_grad(): + traced = torch.jit.trace(static, tuple(T(frames[0]))) + inputs = [ct.TensorType(name=n, shape=(1, MAX_TOKENS), dtype=np.int32) for n in ("input_ids", "attention_mask", "depth_ids")] + for prec, name in ((ct.precision.FLOAT32, "fp32"), (ct.precision.FLOAT16, "fp16")): + ml = ct.convert(traced, inputs=inputs, outputs=[ct.TensorType(name="logits")], compute_precision=prec, + minimum_deployment_target=ct.target.macOS14, convert_to="mlprogram") + path = f"SauerkrautDoom_{name}.mlpackage"; ml.save(path) + for units in ("CPU_ONLY", "CPU_AND_GPU", "CPU_AND_NE"): + mm = ct.models.MLModel(path, compute_units=getattr(ct.ComputeUnit, units)) + out = np.stack([mm.predict({"input_ids": f[0], "attention_mask": f[1], "depth_ids": f[2]})["logits"][0] for f in frames]) + print(f"{name} {units}: max|dlogit| {np.abs(out-ref).max():.2e}, argmax agree {(out.argmax(1)==ref.argmax(1)).mean():.1%} / {len(frames)}") diff --git a/Tools/doom/sauerkraut/convert_ane.py b/Tools/doom/sauerkraut/convert_ane.py new file mode 100644 index 0000000..283da28 --- /dev/null +++ b/Tools/doom/sauerkraut/convert_ane.py @@ -0,0 +1,54 @@ +"""ANE variant: fixed 1026 tokens (no padding exists: 40x25 ASCII + 24 newlines + CLS/SEP), one-hot matmul embeddings.""" +import numpy as np, torch, coremltools as ct, time +from convert import * +N = 1026 + +class Ane(Static): + def __init__(self, clf): + super().__init__(clf, n=N) + e = self.enc + self.register_buffer("tok_table", e.embeddings.tok_embeddings.hash_embeddings.weight.detach().clone()) + self.register_buffer("depth_table", e.depth_embedding.depth_emb.weight.detach().clone()) + def forward(self, ids, depth): + e = self.enc + oh = (ids.float().unsqueeze(-1) == torch.arange(self.tok_table.shape[0]).float()).float() + od = (depth.float().unsqueeze(-1) == torch.arange(self.depth_table.shape[0]).float()).float() + he = e.embeddings.tok_embeddings + x = he.norm(he.projection(oh @ self.tok_table)) + od @ self.depth_table + x = e.embeddings.norm(x) + for i, layer in enumerate(e.layers): + a = layer.attn + qkv = a.Wqkv(layer.attn_norm(x)).view(1, -1, 3, self.nh, self.hd) + q, k, v = qkv.permute(2, 0, 3, 1, 4).unbind(0) + cos, sin = (self.lcos, self.lsin) if self.local[i] else (self.gcos, self.gsin) + q = q * cos + rot_half(q, self.hd // 2) * sin; k = k * cos + rot_half(k, self.hd // 2) * sin + w = torch.matmul(q, k.transpose(2, 3)) * (self.hd ** -0.5) + if self.local[i]: w = w + self.win + o = torch.matmul(torch.softmax(w, -1), v).transpose(1, 2).reshape(1, -1, self.h) + x = x + a.Wo(o) + x = x + layer.mlp(layer.mlp_norm(x)) + x = e.final_norm(x) + pooled = (x * torch.softmax(self.clf.attn_weight(x).squeeze(-1), 1).unsqueeze(-1)).sum(1) + return self.clf.classifier(pooled) + +clf = load_torch(snapshot_download(MODEL_ID)); m = Ane(clf).eval() +d = np.load("frames.npz"); ids, mask, dep = d["ids"], d["mask"], d["depth"] +assert (mask.sum(1) == N).all() +with torch.no_grad(): + ref = np.stack([clf(torch.from_numpy(ids[i:i+1]).long(), torch.from_numpy(mask[i:i+1]).long(), depth_ids=torch.from_numpy(dep[i:i+1]).long())["logits"][0].numpy() for i in range(len(ids))]) + mine = np.stack([m(torch.from_numpy(ids[i:i+1, :N]), torch.from_numpy(dep[i:i+1, :N]))[0].numpy() for i in range(len(ids))]) + print(f"ane-torch vs HF: {np.abs(mine-ref).max():.2e}") + tr = torch.jit.trace(m, (torch.from_numpy(ids[:1, :N]), torch.from_numpy(dep[:1, :N]))) +ml = ct.convert(tr, inputs=[ct.TensorType(name=n, shape=(1, N), dtype=np.int32) for n in ("input_ids", "depth_ids")], + outputs=[ct.TensorType(name="logits")], compute_precision=ct.precision.FLOAT16, + minimum_deployment_target=ct.target.macOS14, convert_to="mlprogram") +ml.save("SauerkrautDoom_L1026_fp16.mlpackage") +for u in ("CPU_ONLY", "CPU_AND_GPU", "CPU_AND_NE"): + mm = ct.models.MLModel("SauerkrautDoom_L1026_fp16.mlpackage", compute_units=getattr(ct.ComputeUnit, u)) + out = np.stack([mm.predict({"input_ids": ids[i:i+1, :N], "depth_ids": dep[i:i+1, :N]})["logits"][0] for i in range(len(ids))]) + f = {"input_ids": ids[:1, :N], "depth_ids": dep[:1, :N]} + for _ in range(20): mm.predict(f) + t = [] + for _ in range(200): + s = time.perf_counter(); mm.predict(f); t.append((time.perf_counter() - s) * 1000) + print(f"L1026 fp16 {u:12s} max|d| {np.abs(out-ref).max():.2e} agree {(out.argmax(1)==ref.argmax(1)).mean():.1%} median {np.median(t):.2f} ms p95 {np.percentile(t,95):.2f}") diff --git a/Tools/doom/sauerkraut/demo.sh b/Tools/doom/sauerkraut/demo.sh new file mode 100755 index 0000000..187938a --- /dev/null +++ b/Tools/doom/sauerkraut/demo.sh @@ -0,0 +1,11 @@ +#!/bin/zsh +# One-command SauerkrautLM-Doom demo: game window + asitop + decision log terminals. +# Tools/doom/sauerkraut/demo.sh # live, seeds from 10016 (25 kills / full 60 s) +# Tools/doom/sauerkraut/demo.sh --record doom.mp4 --episodes 1 +set -e +cd "$(dirname "$0")/../../.." +if [ ! -x .venv/bin/python ]; then + uv venv -q -p 3.12 .venv + uv pip install -q -p .venv vizdoom==1.3.0 numpy coremltools pygame imageio imageio-ffmpeg huggingface_hub +fi +exec .venv/bin/python Tools/doom/sauerkraut/play.py --seed 10016 "$@" diff --git a/Tools/doom/sauerkraut/play.py b/Tools/doom/sauerkraut/play.py new file mode 100644 index 0000000..2f70010 --- /dev/null +++ b/Tools/doom/sauerkraut/play.py @@ -0,0 +1,456 @@ +"""SauerkrautLM-Doom-MultiVec 1.3M on Core ML, playing ViZDoom defend_the_center. + +Split-screen demo: the game on the left; on the right the 40x25 depth grid the model reads, its four +action probabilities, and live ms per decision. No PyTorch at runtime. + +The model's input is 1026 tokens: a character per cell plus a learned embedding of that cell's depth +(16 bins). Upstream's ASCII path overflows uint8 and emits "@" for every cell, in training and here, +so the depth bins carry all the information; the panel draws them. + + .venv/bin/python Tools/doom/sauerkraut/play.py # window, real-time + .venv/bin/python Tools/doom/sauerkraut/play.py --record doom.mp4 --episodes 3 + .venv/bin/python Tools/doom/sauerkraut/play.py --check 10 # headless, no rendering + +Window mode also opens one Ghostty window split into two rows (tmux): `sudo asitop` (GPU/ANE/CPU +power) on top and `tail -f` of the decision log (DOOM_DEMO_LOG, default /tmp/doom-demo.log) below. +Pass --no-terminals to skip it. + +The window renders every game tic (35 fps). ViZDoom's rendering consumes game randomness, so a +seed does not replay the headless benchmark path exactly (`--check` does); the play is equally +strong (seeds 10000-10029: 19.70 kills / 50.5 s rendered vs 20.37 / 50.7 s headless). + +The window waits on each episode's first frame until you press space (--autostart to skip), and +can be resized freely; the layout scales to fit. + +Keys: space start / pause, n next episode, q quit. +""" + +import argparse +import json +import os +import statistics +import subprocess +import time + +import numpy as np +import vizdoom + +REPO_ID = "FluidInference/sauerkrautlm-doom-coreml" # Core ML port of VAGOsolutions/SauerkrautLM-Doom-MultiVec-1.3M +MODEL_NAME = "SauerkrautDoom_L1026_fp16.mlpackage" + +ROWS, COLS, DEPTH_BINS = 25, 40, 16 +NO_DEPTH = DEPTH_BINS +SEQ = 1 + ROWS * COLS + (ROWS - 1) + 1 # [CLS] + chars + newlines + [SEP] +CHARS = " .:-=+*#%@" +FRAME_SKIP, TIMEOUT = 4, 2100 +NAMES = ["shoot", "move_forward", "turn_left", "turn_right"] +BUTTONS = {"shoot": [1, 0, 0, 0], "move_forward": [0, 1, 0, 0], "turn_left": [0, 0, 1, 0], "turn_right": [0, 0, 0, 1]} +LOG_PATH = os.environ.get("DOOM_DEMO_LOG", "/tmp/doom-demo.log") +PYTORCH_GPU_MS = 8.4 # same model in PyTorch on the GPU (MPS, fp16), same M5 Pro and frame + + +# ---- observation: the upstream AsciiConverter.convert_with_depth + character tokenizer ---- +def downscale(img, h, w): + bh, bw = img.shape[0] // h, img.shape[1] // w + return img[: bh * h, : bw * w].reshape(h, bh, w, bw).mean(axis=(1, 3)) + + +def observe(depth, vocab): + coarse = downscale(depth.astype(np.uint8), ROWS, COLS).astype(np.uint8) + peak = coarse.max() + bright = 255 - (coarse * 255 / peak).astype(np.uint8) if peak > 0 else np.zeros_like(coarse) + levels = np.minimum((bright.astype(np.int64) * len(CHARS)) // 256, len(CHARS) - 1) + rows = ["".join(CHARS[v] for v in row) for row in levels] # upstream uint8 overflow: always "@" + + d = downscale(depth.astype(np.float32), ROWS, COLS).astype(np.uint8).astype(np.float32) # upstream truncates + lo, hi = d.min(), d.max() + norm = (d - lo) / (hi - lo) if hi > lo else np.zeros_like(d) + bins = np.clip((norm * DEPTH_BINS).astype(int), 0, DEPTH_BINS - 1) + + ids, depth_ids = [vocab["[CLS]"]], [NO_DEPTH] + for y in range(ROWS): + for x in range(COLS): + ids.append(vocab[rows[y][x]]) + depth_ids.append(int(bins[y, x])) + if y < ROWS - 1: + ids.append(vocab["\n"]) + depth_ids.append(NO_DEPTH) + ids.append(vocab["[SEP]"]) + depth_ids.append(NO_DEPTH) + return bins, np.array([ids], dtype=np.int32), np.array([depth_ids], dtype=np.int32) + + +def choose(probs): + """Upstream composite rule: add a shot when nearly as likely, add a turn/move runner-up.""" + ranked = np.argsort(probs)[::-1] + best, second = NAMES[ranked[0]], NAMES[ranked[1]] + buttons, label = list(BUTTONS[best]), [best] + + def press(name): + return [max(a, b) for a, b in zip(buttons, BUTTONS[name])] + + if best != "shoot" and probs[0] > probs[ranked[0]] * 0.75: + buttons = press("shoot") + label.append("shoot") + kind = {"move_forward": "move", "turn_left": "turn", "turn_right": "turn"} + if probs[ranked[1]] > 0.15 and second != "shoot" and kind.get(best) and kind.get(second): + if kind[best] != kind[second]: + buttons = press(second) + label.append(second) + return buttons, " + ".join(label) + + +class DecisionLog: + """ANSI-colored decision lines for a `tail -f` terminal next to asitop.""" + + def __init__(self, path, units): + self.file = open(path, "w", buffering=1) + self.device = {"CPU_AND_GPU": "GPU", "CPU_AND_NE": "Neural Engine", "CPU_ONLY": "CPU"}.get(units, units) + + def episode(self, seed): + self.file.write(f"\033[1;36m▶ episode · seed {seed}\033[0m\n") + + def decision(self, elapsed, probs, action, ms): + short = {"shoot": "shoot", "move_forward": "fwd", "turn_left": "left", "turn_right": "right"} + act = " + ".join(short[a] for a in action.split(" + ")) + scores = " ".join(f"{short[n]} {p:.2f}" for n, p in zip(NAMES, probs)) + self.file.write( + f"{elapsed:5.1f}s \033[33m→ {act:<18}\033[0m {scores} \033[1;31m{ms:4.1f} ms {self.device}\033[0m\n" + ) + + def result(self, kills, outcome, survived): + self.file.write(f"\033[32m■ {kills} kills · {outcome} {survived:.1f} s\033[0m\n") + + +def open_terminals(log_path): + """One Ghostty window, two rows (tmux): `sudo asitop` on top, where the password is typed, and the + decision log below. Falls back to Terminal.app when Ghostty or tmux is missing.""" + asitop = next( + (p for p in (os.path.expanduser("~/.local/bin/asitop"), "/opt/homebrew/bin/asitop", "/usr/local/bin/asitop") + if os.access(p, os.X_OK)), + None, + ) + if asitop is None: + print("asitop not found: `uv tool install asitop` to show power usage") + tail = f"tail -n 0 -f {log_path}" + tmux = next((p for p in ("/opt/homebrew/bin/tmux", "/usr/local/bin/tmux") if os.access(p, os.X_OK)), None) + ghostty = "/Applications/Ghostty.app" + if tmux and asitop and os.path.isdir(ghostty): + launcher = os.path.join(os.path.dirname(log_path), "doom-demo-terminal.sh") + with open(launcher, "w") as f: + f.write( + f"{tmux} kill-session -t doom-demo 2>/dev/null\n" + f"exec {tmux} new-session -s doom-demo 'sudo {asitop}' \\; split-window -v '{tail}' \\; " + "select-pane -t 0\n" + ) + subprocess.run( + [ + "open", "-na", ghostty, "--args", "--font-size=11", "--window-position-x=690", + "--window-position-y=30", "--window-width=110", "--window-height=52", + f"--command=/bin/zsh {launcher}", + ], + check=False, + ) + return + for command in ([f"sudo {asitop}"] if asitop else []) + [f"clear; {tail}"]: + script = f'tell application "Terminal"\nactivate\ndo script "{command}"\nend tell' + subprocess.run(["osascript", "-e", script], check=False, capture_output=True) + + +def make_game(): + game = vizdoom.DoomGame() + game.load_config(os.path.join(vizdoom.scenarios_path, "defend_the_center.cfg")) + game.set_window_visible(False) + game.set_screen_resolution(vizdoom.ScreenResolution.RES_640X480) + game.set_screen_format(vizdoom.ScreenFormat.RGB24) + game.set_render_hud(True) + game.set_depth_buffer_enabled(True) + game.clear_available_buttons() + for button in ("ATTACK", "MOVE_FORWARD", "TURN_LEFT", "TURN_RIGHT"): + game.add_available_button(getattr(vizdoom.Button, button)) + for variable in ("HEALTH", "AMMO2", "KILLCOUNT"): + game.add_available_game_variable(getattr(vizdoom.GameVariable, variable)) + game.set_episode_timeout(TIMEOUT) + game.set_mode(vizdoom.Mode.PLAYER) + game.init() + return game + + +class Policy: + def __init__(self, path, units): + import coremltools as ct + from huggingface_hub import snapshot_download + + # Real files, not the hub cache's symlinks: Core ML can't compile a symlinked .mlpackage. + repo = snapshot_download( + REPO_ID, allow_patterns=[f"{MODEL_NAME}/*", "tokenizer.json", "config.json"], + local_dir=os.path.expanduser("~/Library/Caches/FluidUse/sauerkrautlm-doom-coreml"), + ) + path = path or os.path.join(repo, MODEL_NAME) + with open(os.path.join(repo, "tokenizer.json")) as f: + self.vocab = json.load(f)["model"]["vocab"] + started = time.perf_counter() + self.model = ct.models.MLModel(path, compute_units=getattr(ct.ComputeUnit, units)) + self.load_s = time.perf_counter() - started + self.units = units + + def __call__(self, depth): + bins, ids, depth_ids = observe(depth, self.vocab) + assert ids.shape[1] == SEQ + started = time.perf_counter() + logits = self.model.predict({"input_ids": ids, "depth_ids": depth_ids})["logits"][0].astype(np.float64) + ms = (time.perf_counter() - started) * 1000 + probs = np.exp(logits - logits.max()) + return bins, probs / probs.sum(), ms + + +def check(policy, first_seed, episodes): + game, kills, tics, ms = make_game(), [], [], [] + for seed in range(first_seed, first_seed + episodes): + game.set_seed(seed) + game.new_episode() + while not game.is_episode_finished(): + _, probs, elapsed = policy(game.get_state().depth_buffer) + ms.append(elapsed) + game.make_action(choose(probs)[0], FRAME_SKIP) + kills.append(int(game.get_game_variable(vizdoom.GameVariable.KILLCOUNT))) + tics.append(game.get_episode_time()) + print(f"seed {seed}: {kills[-1]} kills, {tics[-1] / 35:.1f} s") + game.close() + print( + f"{episodes} episodes: {statistics.mean(kills):.2f} kills, {statistics.mean(tics) / 35:.1f} s, " + f"{statistics.median(ms):.2f} ms median per decision" + ) + + +class Viewer: + """Two rows so a terminal fits beside it: the game on top; below, the depth grid the model reads + and its action probabilities, current action, and ms per decision.""" + + W, H = 680, 815 # layout size; the window is resizable and the layout scales to fit + GAME = (20, 64, 640, 480) # x, y, w, h + + def __init__(self, record, headless): + if headless: + os.environ["SDL_VIDEODRIVER"] = "dummy" + os.environ.setdefault("SDL_VIDEO_WINDOW_POS", "0,30") + import pygame + + self.pg = pygame + pygame.init() + self.window = pygame.display.set_mode((self.W, self.H), pygame.RESIZABLE) + self.screen = pygame.Surface((self.W, self.H)) + pygame.display.set_caption("SauerkrautLM-Doom on Core ML") + mono = pygame.font.match_font("menlo,monaco,couriernew") + sans = pygame.font.match_font("helveticaneue,helvetica,arial") + self.mono = pygame.font.Font(mono, 16) + self.big = pygame.font.Font(sans, 24) + self.label = pygame.font.Font(sans, 17) + self.small = pygame.font.Font(sans, 14) + self.writer = None + if record: + import imageio + + self.writer = imageio.get_writer(record, fps=35, codec="libx264", quality=7, macro_block_size=8) + + def text(self, font, s, pos, color=(230, 230, 230)): + self.screen.blit(font.render(s, True, color), pos) + + def draw(self, frame, hud): + pg, s = self.pg, self.screen + grey = (150, 155, 165) + s.fill((16, 17, 20)) + self.text(self.big, "A 1.3M-parameter model plays Doom", (20, 8)) + self.text(self.small, "from a 40×25 depth grid · SauerkrautLM-Doom-MultiVec · Core ML on the Mac GPU", (21, 38), grey) + x, y, w, h = self.GAME + surface = pg.surfarray.make_surface(frame.swapaxes(0, 1)) + s.blit(pg.transform.smoothscale(surface, (w, h)), (x, y)) + self.text( + self.label, + f"kills {hud['kills']} health {hud['health']} ammo {hud['ammo']} " + f"{hud['elapsed']:4.1f} s / 60 s seed {hud['seed']}", + (20, y + h + 8), + ) + + top = y + h + 38 + self.text(self.small, "What the model reads: depth, near = bright", (20, top), grey) + cell_w, cell_h = 7, 8 + for gy, row in enumerate(hud["bins"]): + for gx, b in enumerate(row): + shade = int(235 - b * 13) + pg.draw.rect( + s, (shade // 3, shade, shade // 2), (20 + gx * cell_w, top + 22 + gy * cell_h, cell_w - 1, cell_h - 1) + ) + + px = 310 + self.text(self.small, "Action probabilities", (px, top), grey) + for i, name in enumerate(NAMES): + yy = top + 26 + i * 28 + p = hud["probs"][i] + self.text(self.small, name.replace("_", " "), (px, yy)) + pg.draw.rect(s, (45, 48, 55), (px + 105, yy + 2, 210, 14), border_radius=3) + pg.draw.rect(s, (90, 170, 250), (px + 105, yy + 2, int(210 * p), 14), border_radius=3) + self.text(self.small, f"{p:.2f}", (px + 322, yy)) + self.text(self.small, f"pressing: {hud['action']}", (px, top + 142), (250, 200, 90)) + self.text(self.mono, f"{hud['ms']:.1f} ms", (px, top + 178), (90, 170, 250)) + self.text(self.small, "per decision (Core ML, GPU)", (px + 90, top + 180), grey) + self.text(self.mono, f"{PYTORCH_GPU_MS:.1f} ms", (px, top + 204), grey) + self.text(self.small, "same model, PyTorch GPU (MPS)", (px + 90, top + 206), grey) + if hud.get("banner"): + b = self.big.render(hud["banner"], True, (255, 255, 255)) + pg.draw.rect(s, (0, 0, 0), (x, y + h // 2 - 28, w, 56)) + s.blit(b, (x + (w - b.get_width()) // 2, y + h // 2 - b.get_height() // 2)) + self.present() + if self.writer: + self.writer.append_data(pg.surfarray.array3d(s).swapaxes(0, 1)) + + def present(self): + """Scale the layout into the current window size, keeping its aspect ratio.""" + win = self.pg.display.get_surface() + ww, wh = win.get_size() + scale = min(ww / self.W, wh / self.H) + size = (max(1, int(self.W * scale)), max(1, int(self.H * scale))) + win.fill((16, 17, 20)) + win.blit(self.pg.transform.smoothscale(self.screen, size), ((ww - size[0]) // 2, (wh - size[1]) // 2)) + self.pg.display.flip() + + def wait_for_start(self, frame, hud): + """Hold on the first frame until space (start) or q (quit).""" + while True: + self.draw(frame, hud) + for event in self.pg.event.get(): + if event.type == self.pg.QUIT: + return False + if event.type == self.pg.KEYDOWN: + if event.key == self.pg.K_q: + return False + if event.key == self.pg.K_SPACE: + return True + time.sleep(1 / 30) + + def events(self): + for event in self.pg.event.get(): + if event.type == self.pg.QUIT: + return "quit" + if event.type == self.pg.KEYDOWN: + return {self.pg.K_q: "quit", self.pg.K_SPACE: "pause", self.pg.K_n: "next"}.get(event.key) + return None + + def close(self): + if self.writer: + self.writer.close() + self.pg.quit() + + +def play(policy, args): + log = DecisionLog(LOG_PATH, policy.units) + if not args.headless and not args.no_terminals: + open_terminals(LOG_PATH) + time.sleep(1.0) + viewer = Viewer(args.record, args.headless) + game = make_game() + realtime = not (args.record and args.headless) + args.autostart = args.autostart or args.headless + ms_window = [] + seed = args.seed + try: + for _ in range(args.episodes): + game.set_seed(seed) + game.new_episode() + state = game.get_state() + paused, quit_requested, skipped = False, False, False + if not args.autostart: + idle = { + "bins": np.full((ROWS, COLS), DEPTH_BINS - 1), "probs": np.full(4, 0.25), "action": "—", + "ms": 0.0, "kills": 0, "health": 100, "ammo": 26, "elapsed": 0.0, "seed": seed, + "banner": "Press space to start", + } + if not viewer.wait_for_start(state.screen_buffer, idle): + break + log.episode(seed) + bins, probs = np.full((ROWS, COLS), DEPTH_BINS - 1), np.full(4, 0.25) + while not game.is_episode_finished(): + command = viewer.events() + if command == "quit": + quit_requested = True + break + if command == "next": + skipped = True + break + if command == "pause": + paused = not paused + if paused: + time.sleep(0.03) + continue + bins, probs, ms = policy(state.depth_buffer) + ms_window = (ms_window + [ms])[-30:] + buttons, action = choose(probs) + log.decision(game.get_episode_time() / 35, probs, action, ms) + for _tic in range(FRAME_SKIP): # render every tic so video runs at 35 fps + tic_started = time.perf_counter() + game.make_action(buttons, 1) + if game.is_episode_finished(): + break + state = game.get_state() + health, ammo, kills = ( + int(game.get_game_variable(getattr(vizdoom.GameVariable, v))) + for v in ("HEALTH", "AMMO2", "KILLCOUNT") + ) + viewer.draw( + state.screen_buffer, + { + "bins": bins, "probs": probs, "action": action, "ms": statistics.median(ms_window), + "kills": kills, "health": health, "ammo": ammo, + "elapsed": game.get_episode_time() / 35, "seed": seed, + }, + ) + if realtime: + time.sleep(max(0.0, 1 / 35 - (time.perf_counter() - tic_started))) + if quit_requested: + break + kills = int(game.get_game_variable(vizdoom.GameVariable.KILLCOUNT)) + survived = game.get_episode_time() / 35 + outcome = "skipped at" if skipped else "died at" if game.is_player_dead() else "survived" + last = game.get_state() or state + for _ in range(70): # 2 s end card + viewer.draw( + last.screen_buffer if last is not None else np.zeros((480, 640, 3), np.uint8), + { + "bins": bins, "probs": probs, "action": "—", "ms": statistics.median(ms_window) if ms_window else 0.0, + "kills": kills, "health": 0 if game.is_player_dead() else "—", "ammo": "—", + "elapsed": survived, "seed": seed, + "banner": f"{kills} kills · {outcome} {survived:.1f} s", + }, + ) + if realtime: + time.sleep(1 / 35) + log.result(kills, outcome, survived) + print(f"seed {seed}: {kills} kills, {outcome} {survived:.1f} s") + seed += 1 + finally: + game.close() + viewer.close() + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model", default=None, help=f"local .mlpackage; default downloads {MODEL_NAME} from {REPO_ID}") + parser.add_argument("--units", default="CPU_AND_GPU", help="CPU_ONLY, CPU_AND_GPU, CPU_AND_NE, ALL") + parser.add_argument("--seed", type=int, default=10000) + parser.add_argument("--episodes", type=int, default=100) + parser.add_argument("--record", default=None, help="write an mp4 of the window") + parser.add_argument("--headless", action="store_true", help="no window; with --record, renders as fast as possible") + parser.add_argument("--autostart", action="store_true", help="start each episode without waiting for space") + parser.add_argument("--no-terminals", action="store_true", help="don't open the asitop and log Terminal windows") + parser.add_argument("--check", type=int, default=0, help="headless score check over N seeds, no rendering") + args = parser.parse_args() + policy = Policy(args.model, args.units) + print(f"Loaded {os.path.basename(args.model or MODEL_NAME)} ({args.units}) in {policy.load_s:.1f} s") + if args.check: + check(policy, args.seed, args.check) + else: + play(policy, args) + + +if __name__ == "__main__": + main() diff --git a/Tools/doom/sauerkraut/sauer_eval.py b/Tools/doom/sauerkraut/sauer_eval.py new file mode 100644 index 0000000..84cf681 --- /dev/null +++ b/Tools/doom/sauerkraut/sauer_eval.py @@ -0,0 +1,99 @@ +"""Evaluate SauerkrautLM-Doom-MultiVec (PyTorch or Core ML) on defend_the_center, Space-identical setup.""" +import argparse, json, os, statistics, time +import numpy as np, torch, vizdoom +from huggingface_hub import snapshot_download +from transformers import AutoTokenizer +from doom_multivec.ascii.converter import AsciiConverter +from doom_multivec.model.classifier import DoomMultiVecClassifier + +MODEL_ID = "VAGOsolutions/SauerkrautLM-Doom-MultiVec-1.3M" +DEPTH_BINS, NO_DEPTH, MAX_TOKENS, FRAME_SKIP, TIMEOUT = 16, 16, 1100, 4, 2100 +NAMES = ["shoot", "move_forward", "turn_left", "turn_right"] +BUTTONS = {"shoot": [1,0,0,0], "move_forward": [0,1,0,0], "turn_left": [0,0,1,0], "turn_right": [0,0,0,1]} + +def load_torch(model_dir): + w = torch.load(os.path.join(model_dir, "model.pt"), map_location="cpu", weights_only=False) + m = DoomMultiVecClassifier(model_dir, pool_mode="attention", num_actions=4) + m.load_state_dict(w); m.eval(); return m + +def make_game(): + g = vizdoom.DoomGame() + g.load_config(os.path.join(vizdoom.scenarios_path, "defend_the_center.cfg")) + g.set_window_visible(False) + g.set_screen_resolution(vizdoom.ScreenResolution.RES_640X480) + g.set_screen_format(vizdoom.ScreenFormat.RGB24) + g.set_render_hud(True); g.set_depth_buffer_enabled(True) + g.clear_available_buttons() + for b in ("ATTACK", "MOVE_FORWARD", "TURN_LEFT", "TURN_RIGHT"): + g.add_available_button(getattr(vizdoom.Button, b)) + for v in ("HEALTH", "AMMO2", "KILLCOUNT"): + g.add_available_game_variable(getattr(vizdoom.GameVariable, v)) + g.set_episode_timeout(TIMEOUT); g.set_mode(vizdoom.Mode.PLAYER); g.init(); return g + +class Encoder: + def __init__(self, model_dir): + self.tok = AutoTokenizer.from_pretrained(model_dir) + self.conv = AsciiConverter(width=40, height=25) + def __call__(self, screen, depth): + text, bins = self.conv.convert_with_depth(screen, depth, num_bins=DEPTH_BINS) + enc = self.tok(text, return_tensors="np", max_length=MAX_TOKENS, padding="max_length", truncation=True) + ids = enc["input_ids"].astype(np.int32); mask = enc["attention_mask"].astype(np.int32) + d = np.full((1, MAX_TOKENS), NO_DEPTH, dtype=np.int32) + n = min(len(bins), MAX_TOKENS - 2) + d[0, 1:1+n] = np.asarray(bins[:n]) + return ids, mask, d + +def composite(probs): + ranked = np.argsort(probs)[::-1] + best, second = NAMES[ranked[0]], NAMES[ranked[1]] + buttons = list(BUTTONS[best]) + press = lambda n: [max(a, b) for a, b in zip(buttons, BUTTONS[n])] + if best != "shoot" and probs[0] > probs[ranked[0]] * 0.75: buttons = press("shoot") + kind = {"move_forward": "move", "turn_left": "turn", "turn_right": "turn"} + if probs[ranked[1]] > 0.15 and second != "shoot" and kind.get(best) and kind.get(second) and kind[best] != kind[second]: + buttons = press(second) + return buttons + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--backend", default="torch") # torch | coreml + ap.add_argument("--mlpackage", default=None) + ap.add_argument("--units", default="CPU_AND_NE") + ap.add_argument("--seeds", default="10000:10100") + ap.add_argument("--out", default=None) + a = ap.parse_args() + model_dir = snapshot_download(MODEL_ID) + enc = Encoder(model_dir) + if a.backend == "torch": + torch.set_num_threads(int(os.environ.get("TORCH_THREADS", "1"))) + m = load_torch(model_dir) + def logits(ids, mask, d): + with torch.inference_mode(): + return m(torch.from_numpy(ids).long(), torch.from_numpy(mask).long(), depth_ids=torch.from_numpy(d).long())["logits"][0].numpy() + else: + import coremltools as ct + ml = ct.models.MLModel(a.mlpackage, compute_units=getattr(ct.ComputeUnit, a.units)) + names = {i.name for i in ml.get_spec().description.input} + def logits(ids, mask, d): + if "attention_mask" not in names: # L1026 variant: every frame is exactly 1026 tokens + assert mask.sum() == 1026 + return ml.predict({"input_ids": ids[:, :1026], "depth_ids": d[:, :1026]})["logits"][0] + return ml.predict({"input_ids": ids, "attention_mask": mask, "depth_ids": d})["logits"][0] + lo, hi = map(int, a.seeds.split(":")) + g = make_game(); kills_all, ms_all, tics_all, died_all = [], [], [], [] + t0 = time.time() + for seed in range(lo, hi): + g.set_seed(seed); g.new_episode() + while not g.is_episode_finished(): + s = g.get_state() + ids, mask, d = enc(s.screen_buffer, s.depth_buffer) + t = time.perf_counter(); lg = np.asarray(logits(ids, mask, d), dtype=np.float64); ms_all.append((time.perf_counter()-t)*1000) + p = np.exp(lg - lg.max()); p /= p.sum() + g.make_action(composite(p), FRAME_SKIP) + kills_all.append(int(g.get_game_variable(vizdoom.GameVariable.KILLCOUNT))); tics_all.append(g.get_episode_time()); died_all.append(g.is_player_dead()) + g.close() + print(f"{a.backend}{'/'+a.units if a.backend=='coreml' else ''}: {len(kills_all)} eps, mean kills {statistics.mean(kills_all):.2f} " + f"(sd {statistics.pstdev(kills_all):.2f}), model ms median {statistics.median(ms_all):.2f} p95 {np.percentile(ms_all,95):.2f}, wall {time.time()-t0:.0f}s, survived {statistics.mean(tics_all)/35:.1f}s avg, {sum(not d for d in died_all)} full episodes") + if a.out: json.dump({"tics": tics_all, "died": died_all, "kills": kills_all, "ms_median": statistics.median(ms_all)}, open(a.out, "w")) + +if __name__ == "__main__": main()