coorl-lost-cities/tests/lost_cities_jax/test_differential.py

from __future__ import annotations

import json
import os
import random

import jax
import jax.numpy as jnp
import numpy as np
import pytest

from lost_cities_jax import batched_legal_mask, batched_reset_from_order, batched_step, score
from lost_cities_jax.types import N_ACTIONS
from reference import lost_cities_ref as ref
from tests.lost_cities_jax.conftest import game_count
from tests.lost_cities_jax.helpers import shuffled_order

BATCHED_SCORE = jax.jit(jax.vmap(score))
LOW_WORD_MASK = (1 << 64) - 1


@pytest.mark.slow
def test_jax_matches_reference_random_legal_policy(pytestconfig, tmp_path):
    games = game_count(
        pytestconfig,
        env_name="LOST_CITIES_JAX_DIFF_GAMES",
        local=100,
        ci=100_000,
        full=1_000_000,
    )
    batch_size = int(os.environ.get("LOST_CITIES_JAX_DIFF_BATCH", "8192"))
    policy_rng = random.Random(20260704)

    for batch_start in range(0, games, batch_size):
        current_batch = min(batch_size, games - batch_start)
        deck_orders = [
            shuffled_order(10_000_000 + batch_start + idx) for idx in range(current_batch)
        ]
        jax_state = batched_reset_from_order(jnp.asarray(deck_orders, dtype=jnp.int8))
        ref_states = [ref.reset_from_order(deck_order) for deck_order in deck_orders]
        action_histories: list[list[int]] = [[] for _ in range(current_batch)]

        for step_idx in range(500):
            jax_masks = np.asarray(batched_legal_mask(jax_state), dtype=bool)
            jax_mask_low, jax_mask_high = _pack_masks_to_words(jax_masks)
            jax_done = np.asarray(jax_state.done, dtype=bool)
            actions = np.zeros((current_batch,), dtype=np.int32)
            ref_rewards: list[list[float]] = []
            ref_done_values: list[bool] = []
            active = 0

            for batch_idx, ref_state in enumerate(ref_states):
                game_idx = batch_start + batch_idx
                ref_bits = ref.legal_action_bits(ref_state)
                if int(jax_mask_low[batch_idx]) != (ref_bits & LOW_WORD_MASK) or int(
                    jax_mask_high[batch_idx]
                ) != (ref_bits >> 64):
                    _dump_failure(
                        tmp_path,
                        game_idx,
                        step_idx,
                        deck_orders[batch_idx],
                        action_histories[batch_idx],
                    )
                    ref_mask = ref.legal_action_mask(ref_state)
                    diff = [
                        idx
                        for idx, pair in enumerate(
                            zip(list(jax_masks[batch_idx]), ref_mask, strict=False)
                        )
                        if pair[0] != pair[1]
                    ]
                    pytest.fail(
                        f"legal mask mismatch game={game_idx} step={step_idx} diff={diff[:20]}"
                    )

                if bool(jax_done[batch_idx]) != ref_state.done:
                    _dump_failure(
                        tmp_path,
                        game_idx,
                        step_idx,
                        deck_orders[batch_idx],
                        action_histories[batch_idx],
                    )
                    pytest.fail(f"done mismatch game={game_idx} step={step_idx}")

                if ref_state.done:
                    ref_rewards.append([0.0, 0.0])
                    ref_done_values.append(True)
                    continue

                active += 1
                action = ref.nth_legal_action(ref_bits, policy_rng.randrange(ref_bits.bit_count()))
                actions[batch_idx] = action
                action_histories[batch_idx].append(action)
                ref_states[batch_idx], ref_reward, ref_done = ref.step_in_place(
                    ref_state, action, validate=False
                )
                ref_rewards.append(ref_reward)
                ref_done_values.append(ref_done)

            if active == 0:
                jax_scores = np.asarray(BATCHED_SCORE(jax_state), dtype=np.float32)
                for batch_idx, ref_state in enumerate(ref_states):
                    ref_score = np.asarray(ref.score(ref_state), dtype=np.float32)
                    if not np.array_equal(jax_scores[batch_idx], ref_score):
                        game_idx = batch_start + batch_idx
                        _dump_failure(
                            tmp_path,
                            game_idx,
                            step_idx,
                            deck_orders[batch_idx],
                            action_histories[batch_idx],
                        )
                        pytest.fail(
                            f"score mismatch game={game_idx} step={step_idx}: "
                            f"{jax_scores[batch_idx].tolist()} != {ref_score.tolist()}"
                        )
                break

            jax_state, jax_reward, jax_done_after = batched_step(jax_state, jnp.asarray(actions))
            jax_reward = np.asarray(jax_reward, dtype=np.float32)
            jax_done_after = np.asarray(jax_done_after, dtype=bool)
            for batch_idx, (ref_reward, ref_done) in enumerate(
                zip(ref_rewards, ref_done_values, strict=False)
            ):
                if bool(jax_done_after[batch_idx]) != ref_done or not np.array_equal(
                    jax_reward[batch_idx], np.asarray(ref_reward, dtype=np.float32)
                ):
                    game_idx = batch_start + batch_idx
                    _dump_failure(
                        tmp_path,
                        game_idx,
                        step_idx,
                        deck_orders[batch_idx],
                        action_histories[batch_idx],
                    )
                    pytest.fail(f"step result mismatch game={game_idx} step={step_idx}")
        else:
            for batch_idx, ref_state in enumerate(ref_states):
                if not ref_state.done:
                    game_idx = batch_start + batch_idx
                    _dump_failure(
                        tmp_path, game_idx, 500, deck_orders[batch_idx], action_histories[batch_idx]
                    )
                    pytest.fail(f"game did not finish within guard loop: game={game_idx}")


def _pack_masks_to_words(masks: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    packed = np.packbits(masks[:, :N_ACTIONS], axis=1, bitorder="little")
    low = np.ascontiguousarray(packed[:, :8]).view("<u8").reshape((-1,))
    high = np.ascontiguousarray(packed[:, 8:12]).view("<u4").reshape((-1,))
    return low, high


def _dump_failure(
    tmp_path,
    game_idx: int,
    step_idx: int,
    deck_order: list[int],
    actions: list[int],
) -> None:
    path = tmp_path / f"lost_cities_jax_diff_failure_{game_idx}_{step_idx}.json"
    path.write_text(
        json.dumps(
            {
                "game_idx": game_idx,
                "step_idx": step_idx,
                "deck_order": deck_order,
                "actions": actions,
            },
            indent=2,
        )
    )