coorl-lost-cities/scripts/export_match_onnx.py

#!/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()