Add on-device web client

web/src/model/policy.ts

@@ -0,0 +1,100 @@
+import * as ort from "onnxruntime-web/all";
+
+import { currentHandSorted, legalActionMask } from "../game/engine";
+import { observation } from "../game/observation";
+import { DISCARD, DRAW_DECK, PLAY, type GameState } from "../game/types";
+import { cardColor, cardRank, isHandshake } from "../game/cards";
+
+export type ExecutionProvider = "webgpu" | "wasm" | "heuristic";
+
+export interface Policy {
+  readonly provider: ExecutionProvider;
+  action(state: GameState): Promise<number>;
+}
+
+class OnnxPolicy implements Policy {
+  constructor(
+    private readonly session: ort.InferenceSession,
+    readonly provider: ExecutionProvider,
+  ) {}
+
+  async action(state: GameState): Promise<number> {
+    const obs = observation(state, state.toMove);
+    const result = await this.session.run({ obs: new ort.Tensor("float32", obs, [1, obs.length]) });
+    const logits = result.logits?.data;
+    if (!logits) throw new Error("ONNX model did not return a logits output");
+    const legal = legalActionMask(state);
+    let bestAction = -1;
+    let bestLogit = Number.NEGATIVE_INFINITY;
+    for (let action = 0; action < legal.length; action += 1) {
+      const logit = Number(logits[action]);
+      if (legal[action] && logit > bestLogit) {
+        bestAction = action;
+        bestLogit = logit;
+      }
+    }
+    if (bestAction < 0) throw new Error("model received a state without legal actions");
+    return bestAction;
+  }
+}
+
+class HeuristicPolicy implements Policy {
+  readonly provider = "heuristic" as const;
+
+  async action(state: GameState): Promise<number> {
+    const legal = legalActionMask(state);
+    const hand = currentHandSorted(state);
+    let bestAction = -1;
+    let bestScore = Number.NEGATIVE_INFINITY;
+    for (let action = 0; action < legal.length; action += 1) {
+      if (!legal[action]) continue;
+      const handSlot = Math.floor(action / 12);
+      const place = Math.floor((action % 12) / 6);
+      const draw = action % 6;
+      const card = hand[handSlot];
+      let score = place === PLAY ? cardRank(card) + (isHandshake(card) ? 7 : 0) : -cardRank(card);
+      if (draw !== DRAW_DECK) score += 2;
+      if (place === DISCARD && draw > 0 && draw - 1 === cardColor(card)) score -= 100;
+      if (score > bestScore) {
+        bestScore = score;
+        bestAction = action;
+      }
+    }
+    if (bestAction < 0) throw new Error("state has no legal actions");
+    return bestAction;
+  }
+}
+
+async function createSession(provider: "webgpu" | "wasm"): Promise<ort.InferenceSession> {
+  return ort.InferenceSession.create("/models/jax-ppo.onnx", {
+    executionProviders: [provider],
+    graphOptimizationLevel: "all",
+  });
+}
+
+let policyPromise: Promise<{ policy: Policy; warning?: string }> | undefined;
+
+async function initializePolicy(): Promise<{ policy: Policy; warning?: string }> {
+  const canUseWebGpu = "gpu" in navigator;
+  if (canUseWebGpu) {
+    try {
+      return { policy: new OnnxPolicy(await createSession("webgpu"), "webgpu") };
+    } catch (error) {
+      console.warn("WebGPU model initialization failed; trying WASM", error);
+    }
+  }
+  try {
+    return { policy: new OnnxPolicy(await createSession("wasm"), "wasm") };
+  } catch (error) {
+    console.warn("ONNX model initialization failed; using heuristic policy", error);
+    return {
+      policy: new HeuristicPolicy(),
+      warning: "ONNX 모델을 찾지 못해 로컬 휴리스틱으로 플레이합니다.",
+    };
+  }
+}
+
+export function loadPolicy(): Promise<{ policy: Policy; warning?: string }> {
+  policyPromise ??= initializePolicy();
+  return policyPromise;
+}