Ship borealis to the browser, and add the classic three-round mode
The web client now plays borealis (data/models.json), trained on the three-round match and taking the 501-dim match view rather than a bare round. Only the actor trunk is exported -- the critic exists to grade moves in training and never plays, so the graph physically cannot leak the opponent's hand or the deck, which beats promising not to call it. The TypeScript match layer and observation mirror match.py and match_obs.py. They have to agree to the bit: a mismatch throws nowhere, the ONNX policy just consumes a wrong vector and plays worse for reasons nobody can see. So the port is not trusted -- generate_match_parity_fixture.py emits 282 positions from real JAX play (mid-round, both seats, past a roll-over, with a live carry) and the TS output is checked against them to float32 round-off. Match mode is a menu toggle. A seed fixes all three deals and the coin flips, so a match stays a pure function of it. One-deal mode is unchanged from the player's side; borealis simply sees it as round one at a carry of zero, a position it has seen a great many times. Two bugs found by driving the built app in a browser, both silent: - The result card totalled the round, not the match. It read "-11 : 3" while the match stood at -96 : 66 -- it would have named the wrong winner. It now headlines the match total and breaks the round out beneath it. - Game records were being rejected. The client's schema went to v2 (it now records which model played; the old records stored the on-screen label, which stops identifying anything once there are two models) while serve_web_with_logs.py still only accepted v1, so every record would have 400'd into a console warning. v1 stays accepted -- the 111 existing games are altair. npm test and tsc were green through both. Hence web/.claude/skills/verify, which records the recipe and the selectors so the next session drives the app instead of re-deriving how. .gitignore excluded the new model, which would have shipped a 404: the deploy builds straight from the repo. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XBQKgvBbxbheiTF1AVy1Sh
This commit is contained in:
@@ -0,0 +1,138 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Export a match policy's actor trunk to ONNX for the browser.
|
||||
|
||||
Only the actor ships. The critic exists to grade moves during training and never
|
||||
plays, so its trunk -- and the privileged view of the opponent's hand and the deck
|
||||
that feeds it -- is dropped here rather than shipped and then not used. That also
|
||||
means the exported graph physically cannot leak hidden state, which is a stronger
|
||||
guarantee than promising not to call it.
|
||||
|
||||
MatchActorCritic lays the actor out as Dense_0..Dense_{num_layers} exactly as the
|
||||
single-round model does, so the graph construction is the same; only the input
|
||||
width and the checkpoint loader differ.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import hashlib
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
|
||||
from lost_cities_jax.match_obs import MATCH_CRITIC_OBS_DIM, MATCH_OBS_DIM
|
||||
from lost_cities_jax.match_ppo import Ablation, MatchActorCritic, create_match_train_state
|
||||
from lost_cities_jax.ppo import load_config, restore_checkpoint
|
||||
from lost_cities_jax.types import N_ACTIONS
|
||||
|
||||
|
||||
def build_argparser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--checkpoint", type=Path, required=True)
|
||||
parser.add_argument("--config", type=Path, default=Path("configs/jax_ppo/match-selfplay.yaml"))
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument("--codename", required=True, help="see data/models.json")
|
||||
return parser
|
||||
|
||||
|
||||
def export_model(checkpoint: Path, config: Path, output: Path, codename: str) -> None:
|
||||
try:
|
||||
import onnx
|
||||
from onnx import TensorProto, helper, numpy_helper
|
||||
from onnx.reference import ReferenceEvaluator
|
||||
except ImportError as exc:
|
||||
raise SystemExit("onnx is required; run with `uv run --with onnx ...`") from exc
|
||||
|
||||
cfg = load_config(config)
|
||||
state = restore_checkpoint(
|
||||
checkpoint, create_match_train_state(cfg, jax.random.PRNGKey(0), Ablation())
|
||||
)
|
||||
params = state.params
|
||||
dense = params["params"]
|
||||
|
||||
nodes = []
|
||||
initializers = []
|
||||
previous = "obs"
|
||||
for index in range(cfg.network.num_layers):
|
||||
layer = dense[f"Dense_{index}"]
|
||||
weight, bias = f"dense_{index}.weight", f"dense_{index}.bias"
|
||||
initializers.extend(
|
||||
[
|
||||
numpy_helper.from_array(np.asarray(layer["kernel"], dtype=np.float32), weight),
|
||||
numpy_helper.from_array(np.asarray(layer["bias"], dtype=np.float32), bias),
|
||||
]
|
||||
)
|
||||
nodes.append(helper.make_node("Gemm", [previous, weight, bias], [f"dense_{index}.linear"]))
|
||||
nodes.append(helper.make_node("Relu", [f"dense_{index}.linear"], [f"dense_{index}.relu"]))
|
||||
previous = f"dense_{index}.relu"
|
||||
|
||||
actor = dense[f"Dense_{cfg.network.num_layers}"]
|
||||
initializers.extend(
|
||||
[
|
||||
numpy_helper.from_array(np.asarray(actor["kernel"], dtype=np.float32), "actor.weight"),
|
||||
numpy_helper.from_array(np.asarray(actor["bias"], dtype=np.float32), "actor.bias"),
|
||||
]
|
||||
)
|
||||
nodes.append(helper.make_node("Gemm", [previous, "actor.weight", "actor.bias"], ["logits"]))
|
||||
|
||||
graph = helper.make_graph(
|
||||
nodes,
|
||||
f"coolrl-lost-cities-match-actor-{codename}",
|
||||
[helper.make_tensor_value_info("obs", TensorProto.FLOAT, [None, MATCH_OBS_DIM])],
|
||||
[helper.make_tensor_value_info("logits", TensorProto.FLOAT, [None, N_ACTIONS])],
|
||||
initializer=initializers,
|
||||
)
|
||||
model = helper.make_model(
|
||||
graph, producer_name="coolrl-lost-cities", opset_imports=[helper.make_opsetid("", 17)]
|
||||
)
|
||||
model.ir_version = 8
|
||||
onnx.checker.check_model(model)
|
||||
|
||||
# The exported graph has to agree with the trained one, not merely load.
|
||||
rng = np.random.default_rng(20260715)
|
||||
sample = rng.normal(size=(8, MATCH_OBS_DIM)).astype(np.float32)
|
||||
critic_stub = jnp.zeros((8, MATCH_CRITIC_OBS_DIM), dtype=jnp.float32)
|
||||
flax_model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
|
||||
expected, _ = flax_model.apply(params, jnp.asarray(sample), critic_stub)
|
||||
actual = ReferenceEvaluator(model).run(None, {"obs": sample})[0]
|
||||
np.testing.assert_allclose(actual, np.asarray(expected), rtol=2e-5, atol=2e-5)
|
||||
np.testing.assert_array_equal(
|
||||
np.argmax(actual, axis=1), np.argmax(np.asarray(expected), axis=1)
|
||||
)
|
||||
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
onnx.save(model, output)
|
||||
model_bytes = output.read_bytes()
|
||||
manifest = {
|
||||
"format": "coolrl-lost-cities-match-onnx-v1",
|
||||
"codename": codename,
|
||||
"model_file": output.name,
|
||||
"model_size_bytes": len(model_bytes),
|
||||
"model_sha256": hashlib.sha256(model_bytes).hexdigest(),
|
||||
"source_checkpoint": str(checkpoint),
|
||||
"source_config": config.name,
|
||||
"observation_size": MATCH_OBS_DIM,
|
||||
"action_size": N_ACTIONS,
|
||||
"hidden_size": cfg.network.hidden_size,
|
||||
"num_layers": cfg.network.num_layers,
|
||||
"dtype": "float32",
|
||||
"validation_max_abs_error": float(np.max(np.abs(actual - np.asarray(expected)))),
|
||||
}
|
||||
output.with_suffix(".json").write_text(json.dumps(manifest, indent=2) + "\n")
|
||||
print(f"exported {codename} -> {output} ({output.stat().st_size:,} bytes)")
|
||||
print(
|
||||
f" sha256 {manifest['model_sha256'][:12]} obs {MATCH_OBS_DIM} max err "
|
||||
f"{manifest['validation_max_abs_error']:.2e}"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = build_argparser().parse_args()
|
||||
export_model(args.checkpoint, args.config, args.output, args.codename)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,97 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Generate match states and their observations so TypeScript can be checked against JAX.
|
||||
|
||||
The two observation builders must agree to the bit. A mismatch does not throw --
|
||||
the ONNX policy consumes a wrong vector quite happily and plays worse for reasons
|
||||
nobody can see. So the port is not trusted; it is checked.
|
||||
|
||||
States are drawn from real random play so the fixture covers the awkward parts:
|
||||
mid-round, both seats to move, past a round roll-over, with a non-zero carry.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import jax
|
||||
import numpy as np
|
||||
|
||||
from lost_cities_jax.match import MatchState, match_reset_from, match_step
|
||||
from lost_cities_jax.match_obs import MATCH_OBS_DIM, match_observation
|
||||
from lost_cities_jax.opponents import random_legal_action
|
||||
from lost_cities_jax.types import N_CARDS
|
||||
|
||||
OUTPUT = Path(__file__).resolve().parents[1] / "web" / "src" / "game" / "match-parity-fixture.json"
|
||||
N_ROUNDS = 3
|
||||
|
||||
|
||||
def match_json(match: MatchState) -> dict:
|
||||
round_state = match.round
|
||||
return {
|
||||
"round": {
|
||||
"deckOrder": np.asarray(round_state.deck_order).astype(int).tolist(),
|
||||
"drawPtr": int(round_state.draw_ptr),
|
||||
"cardLoc": np.asarray(round_state.card_loc).astype(int).tolist(),
|
||||
"handPublic": np.asarray(round_state.hand_public).astype(bool).tolist(),
|
||||
"colTop": np.asarray(round_state.col_top).astype(int).tolist(),
|
||||
"colHandshakes": np.asarray(round_state.col_hs).astype(int).tolist(),
|
||||
"colLength": np.asarray(round_state.col_len).astype(int).tolist(),
|
||||
"piles": [
|
||||
np.asarray(round_state.pile[color, : int(round_state.pile_len[color])])
|
||||
.astype(int)
|
||||
.tolist()
|
||||
for color in range(5)
|
||||
],
|
||||
"toMove": int(round_state.to_move),
|
||||
"stepCount": int(round_state.step_count),
|
||||
"done": bool(round_state.done),
|
||||
},
|
||||
"deckOrders": np.asarray(match.deck_orders).astype(int).tolist(),
|
||||
"coinFlips": np.asarray(match.coin_flips).astype(int).tolist(),
|
||||
"roundIdx": int(match.round_idx),
|
||||
"carry": np.asarray(match.carry).astype(int).tolist(),
|
||||
"done": bool(match.done),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
rng = np.random.default_rng(20260715)
|
||||
key = jax.random.PRNGKey(7)
|
||||
rows = []
|
||||
|
||||
for match_index in range(6):
|
||||
decks = np.stack([rng.permutation(N_CARDS) for _ in range(N_ROUNDS)])
|
||||
coins = rng.integers(0, 2, size=(N_ROUNDS,))
|
||||
match = match_reset_from(decks.astype(np.int8), coins.astype(np.int8))
|
||||
|
||||
# Sample the opening position and then every 17th ply, which lands in all
|
||||
# three rounds and on both seats without hand-picking anything.
|
||||
ply = 0
|
||||
while not bool(match.done) and ply < 400:
|
||||
if ply % 17 == 0 or ply == 0:
|
||||
for player in (0, 1):
|
||||
obs = np.asarray(match_observation(match, player), dtype=np.float64)
|
||||
assert obs.shape == (MATCH_OBS_DIM,)
|
||||
rows.append(
|
||||
{
|
||||
"match": match_json(match),
|
||||
"player": player,
|
||||
"observation": [round(float(v), 7) for v in obs],
|
||||
}
|
||||
)
|
||||
key, step_key = jax.random.split(key)
|
||||
action = int(random_legal_action(match.round, match.round.to_move, step_key))
|
||||
match, _, _ = match_step(match, action)
|
||||
ply += 1
|
||||
del match_index
|
||||
|
||||
OUTPUT.write_text(
|
||||
json.dumps({"format": "jax-web-match-parity-v1", "obsDim": MATCH_OBS_DIM, "rows": rows})
|
||||
+ "\n"
|
||||
)
|
||||
print(f"wrote {len(rows)} rows -> {OUTPUT}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,86 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Serve the production web client and append completed games to JSONL."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from http import HTTPStatus
|
||||
from http.server import SimpleHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
MAX_RECORD_BYTES = 6_000_000
|
||||
|
||||
# v2 adds the opponent's identity (codename + hash) and the match layer -- three
|
||||
# deals, the coin flips, and which round each move belongs to. v1 is still
|
||||
# accepted because 111 games were recorded under it; they are all altair, which
|
||||
# data/models.json records since v1 has nowhere to say so.
|
||||
SUPPORTED_FORMATS = frozenset({"lost-cities-web-game-v1", "lost-cities-web-game-v2"})
|
||||
|
||||
|
||||
def parse_record(body: bytes) -> dict[str, Any]:
|
||||
if len(body) > MAX_RECORD_BYTES:
|
||||
raise ValueError("record is too large")
|
||||
value = json.loads(body)
|
||||
if not isinstance(value, dict) or value.get("format") not in SUPPORTED_FORMATS:
|
||||
raise ValueError("unsupported game record")
|
||||
if not isinstance(value.get("gameId"), str) or not isinstance(value.get("moves"), list):
|
||||
raise ValueError("invalid game record")
|
||||
return value
|
||||
|
||||
|
||||
def make_handler(dist: Path, output: Path):
|
||||
seen_ids: set[str] = set()
|
||||
if output.exists():
|
||||
for line in output.read_text(encoding="utf-8").splitlines():
|
||||
try:
|
||||
game_id = json.loads(line).get("gameId")
|
||||
if isinstance(game_id, str):
|
||||
seen_ids.add(game_id)
|
||||
except (json.JSONDecodeError, AttributeError):
|
||||
continue
|
||||
|
||||
class Handler(SimpleHTTPRequestHandler):
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
super().__init__(*args, directory=str(dist), **kwargs)
|
||||
|
||||
def do_POST(self) -> None: # noqa: N802
|
||||
if self.path != "/api/game-records":
|
||||
self.send_error(HTTPStatus.NOT_FOUND)
|
||||
return
|
||||
try:
|
||||
length = int(self.headers.get("content-length", "0"))
|
||||
record = parse_record(self.rfile.read(length))
|
||||
except (ValueError, json.JSONDecodeError) as error:
|
||||
self.send_error(HTTPStatus.BAD_REQUEST, str(error))
|
||||
return
|
||||
game_id = record["gameId"]
|
||||
if game_id not in seen_ids:
|
||||
output.parent.mkdir(parents=True, exist_ok=True)
|
||||
with output.open("a", encoding="utf-8") as stream:
|
||||
stream.write(
|
||||
json.dumps(record, ensure_ascii=False, separators=(",", ":")) + "\n"
|
||||
)
|
||||
seen_ids.add(game_id)
|
||||
self.send_response(HTTPStatus.NO_CONTENT)
|
||||
self.end_headers()
|
||||
|
||||
return Handler
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--host", default="0.0.0.0")
|
||||
parser.add_argument("--port", type=int, default=5173)
|
||||
parser.add_argument("--dist", type=Path, default=Path("web/dist"))
|
||||
parser.add_argument("--output", type=Path, default=Path("data/human-play/game-records.jsonl"))
|
||||
args = parser.parse_args()
|
||||
server = ThreadingHTTPServer((args.host, args.port), make_handler(args.dist, args.output))
|
||||
print(f"Serving {args.dist} on http://{args.host}:{args.port}", flush=True)
|
||||
print(f"Writing game records to {args.output}", flush=True)
|
||||
server.serve_forever()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user