65 lines
2.2 KiB
Python
65 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
"""Generate deterministic JAX states for TypeScript engine parity tests."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import jax.numpy as jnp
|
|
import numpy as np
|
|
|
|
from lost_cities_jax.engine import legal_action_mask, reset_from_order, step
|
|
from lost_cities_jax.obs import observation
|
|
|
|
OUTPUT = Path(__file__).resolve().parents[1] / "web" / "src" / "game" / "parity-fixture.json"
|
|
|
|
|
|
def state_json(state) -> dict:
|
|
return {
|
|
"deckOrder": np.asarray(state.deck_order).astype(int).tolist(),
|
|
"drawPtr": int(state.draw_ptr),
|
|
"cardLoc": np.asarray(state.card_loc).astype(int).tolist(),
|
|
"handPublic": np.asarray(state.hand_public).astype(bool).tolist(),
|
|
"colTop": np.asarray(state.col_top).astype(int).tolist(),
|
|
"colHandshakes": np.asarray(state.col_hs).astype(int).tolist(),
|
|
"colLength": np.asarray(state.col_len).astype(int).tolist(),
|
|
"piles": [
|
|
np.asarray(state.pile[color, : int(state.pile_len[color])]).astype(int).tolist()
|
|
for color in range(5)
|
|
],
|
|
"toMove": int(state.to_move),
|
|
"stepCount": int(state.step_count),
|
|
"done": bool(state.done),
|
|
}
|
|
|
|
|
|
def main() -> None:
|
|
rng = np.random.default_rng(20260713)
|
|
order = rng.permutation(60).astype(np.int8)
|
|
state = reset_from_order(jnp.asarray(order))
|
|
rows = []
|
|
for index in range(24):
|
|
mask = np.asarray(legal_action_mask(state), dtype=bool)
|
|
rows.append(
|
|
{
|
|
"index": index,
|
|
"state": state_json(state),
|
|
"legalMask": mask.tolist(),
|
|
"observationP0": np.asarray(observation(state, jnp.int32(0))).tolist(),
|
|
"observationP1": np.asarray(observation(state, jnp.int32(1))).tolist(),
|
|
}
|
|
)
|
|
legal = np.flatnonzero(mask)
|
|
action = int(legal[(index * 17 + 3) % len(legal)])
|
|
rows[-1]["action"] = action
|
|
state, _, _ = step(state, jnp.int32(action))
|
|
if bool(state.done):
|
|
break
|
|
OUTPUT.write_text(json.dumps({"format": "jax-web-parity-v1", "rows": rows}) + "\n")
|
|
print(f"wrote {OUTPUT} ({len(rows)} states)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|