coorl-lost-cities/scripts/analyze_first_open_followup.py

#!/usr/bin/env python
"""Inspect model's post-forced-open behavior over the next K policy_player turns.

Diagnoses *why* forced-open continuation values are poor under self-play
rollouts: does the model play followup cards of the opened color, discard
them, or open another color instead?
"""

from __future__ import annotations

import argparse
import json
import sys
import time
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any

import numpy as np
import torch
from coolrl_lost_cities.games.classic.game import GameState, LostCitiesConfig

from coolrl_lost_cities.games.classic.bots import build_bot
from coolrl_lost_cities.games.classic.deep_cfr.config import config_from_dict
from coolrl_lost_cities.games.classic.deep_cfr.networks import DeepCFRMLP

_SCRIPT_DIR = Path(__file__).resolve().parent
if str(_SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(_SCRIPT_DIR))

from analyze_first_open_counterfactual import (  # noqa: E402
    AdvantagePolicy,
    _classify_action,
)


@dataclass
class FollowupBucket:
    candidates: int = 0
    same_color_plays: list[int] = field(default_factory=list)
    same_color_discards: list[int] = field(default_factory=list)
    other_open_new: list[int] = field(default_factory=list)
    other_play_existing: list[int] = field(default_factory=list)
    other_discard: list[int] = field(default_factory=list)
    draw_deck: list[int] = field(default_factory=list)
    draw_pile: list[int] = field(default_factory=list)
    same_color_held_at_force: list[int] = field(default_factory=list)
    same_color_played_in_window: list[int] = field(default_factory=list)
    same_color_discarded_in_window: list[int] = field(default_factory=list)
    final_score_diff: list[float] = field(default_factory=list)
    same_color_held_terminal: list[int] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        def stats(values: list[float] | list[int]) -> dict[str, float]:
            if not values:
                return {"mean": 0.0, "median": 0.0}
            return {
                "mean": float(np.mean(values)),
                "median": float(np.median(values)),
            }

        return {
            "candidates": self.candidates,
            "same_color_plays": stats(self.same_color_plays),
            "same_color_discards": stats(self.same_color_discards),
            "other_open_new": stats(self.other_open_new),
            "other_play_existing": stats(self.other_play_existing),
            "other_discard": stats(self.other_discard),
            "draw_deck": stats(self.draw_deck),
            "draw_pile": stats(self.draw_pile),
            "same_color_held_at_force": stats(self.same_color_held_at_force),
            "same_color_played_in_window": stats(self.same_color_played_in_window),
            "same_color_discarded_in_window": stats(self.same_color_discarded_in_window),
            "same_color_held_terminal": stats(self.same_color_held_terminal),
            "final_score_diff": stats(self.final_score_diff),
        }


def _load_checkpoint(
    checkpoint: Path, device: torch.device
) -> tuple[Any, LostCitiesConfig, AdvantagePolicy, int]:
    payload = torch.load(checkpoint, map_location="cpu")
    cfg = config_from_dict(payload["config"])
    game_config = LostCitiesConfig(**payload["game_config"])
    action_size = int(payload["action_size"])
    networks = [
        DeepCFRMLP.from_config(int(payload["input_dim"]), action_size, cfg.network).to(device)
        for _ in range(2)
    ]
    for network, state_dict in zip(networks, payload["advantage_networks"], strict=True):
        network.load_state_dict(state_dict)
        network.eval()
    policy = AdvantagePolicy(
        networks,
        device=device,
        encoding=cfg.encoding,
        epsilon=cfg.traversal.regret_matching_epsilon,
        fallback=cfg.regret_matching.all_negative_fallback,
    )
    return cfg, game_config, policy, int(payload.get("iteration", -1))


def _count_color_cards_in_hand(state: GameState, player: int, color: int) -> int:
    return sum(
        1 for card in state.hand_slots(player) if card is not None and int(card.color) == color
    )


def _label_followup_action(
    state: GameState, unified_action: int, player: int, forced_color: int
) -> str:
    """Categorise a followup action relative to the previously forced-open color."""
    card_action_size = state.config.hand_size * 2
    if unified_action == card_action_size:
        return "draw_deck"
    if unified_action > card_action_size:
        return "draw_pile"
    if unified_action % 2 == 1:
        slot = unified_action // 2
        card = state.hand_slots(player)[slot]
        if card is None:
            return "discard_invalid"
        return "discard_same" if int(card.color) == forced_color else "discard_other"
    slot = unified_action // 2
    card = state.hand_slots(player)[slot]
    if card is None:
        return "play_invalid"
    color = int(card.color)
    is_open_action = not state.expeditions[player][color]
    if color == forced_color:
        return "play_same"
    return "open_other" if is_open_action else "play_other_existing"


def _force_and_observe(
    base_state: GameState,
    *,
    forced_action: int,
    forced_color: int,
    bucket: FollowupBucket,
    policy: AdvantagePolicy,
    opponent_policy: Any,
    policy_player: int,
    window: int,
    max_steps: int,
) -> None:
    rollout = base_state.clone()
    same_at_force = _count_color_cards_in_hand(rollout, policy_player, forced_color)
    rollout.apply_action(rollout.from_unified_action(forced_action))
    counts: Counter[str] = Counter()
    policy_turns_seen = 0
    same_color_plays = 0
    same_color_discards = 0
    steps = 0
    while not rollout.terminal and steps < max_steps and policy_turns_seen < window:
        current = int(rollout.current_player)
        if current != policy_player:
            rollout.apply_action(opponent_policy.act(rollout))
            steps += 1
            continue
        unified = policy.select_unified(rollout)
        label = _label_followup_action(rollout, unified, policy_player, forced_color)
        counts[label] += 1
        if label == "play_same":
            same_color_plays += 1
        elif label == "discard_same":
            same_color_discards += 1
        rollout.apply_action(rollout.from_unified_action(unified))
        steps += 1
        policy_turns_seen += 1
    while not rollout.terminal and steps < max_steps:
        current = int(rollout.current_player)
        if current == policy_player:
            unified = policy.select_unified(rollout)
            rollout.apply_action(rollout.from_unified_action(unified))
        else:
            rollout.apply_action(opponent_policy.act(rollout))
        steps += 1

    bucket.candidates += 1
    bucket.same_color_plays.append(counts.get("play_same", 0))
    bucket.same_color_discards.append(counts.get("discard_same", 0))
    bucket.other_open_new.append(counts.get("open_other", 0))
    bucket.other_play_existing.append(counts.get("play_other_existing", 0))
    bucket.other_discard.append(counts.get("discard_other", 0))
    bucket.draw_deck.append(counts.get("draw_deck", 0))
    bucket.draw_pile.append(counts.get("draw_pile", 0))
    bucket.same_color_held_at_force.append(same_at_force)
    bucket.same_color_played_in_window.append(same_color_plays)
    bucket.same_color_discarded_in_window.append(same_color_discards)
    bucket.same_color_held_terminal.append(
        _count_color_cards_in_hand(rollout, policy_player, forced_color)
    )
    bucket.final_score_diff.append(float(rollout.score_diff(policy_player)))


def analyze_checkpoint(
    checkpoint: Path,
    *,
    games: int,
    seed: int,
    opponent: str,
    device: torch.device,
    max_steps: int,
    max_candidates: int,
    window: int,
) -> dict[str, Any]:
    _cfg, game_config, policy, iteration = _load_checkpoint(checkpoint, device)
    buckets: dict[str, FollowupBucket] = defaultdict(FollowupBucket)
    candidates_evaluated = 0
    started = time.perf_counter()

    for game_index in range(games):
        if candidates_evaluated >= max_candidates:
            break
        game_seed = seed + game_index
        swap = game_index % 2 == 1
        policy_player = 1 if swap else 0
        opponent_policy = build_bot(opponent, seed=game_seed * 2 + (1 - policy_player))
        state = GameState.new_game(game_config, seed=game_seed)
        for _step in range(max_steps):
            if state.terminal or candidates_evaluated >= max_candidates:
                break
            current = int(state.current_player)
            if current != policy_player:
                state.apply_action(opponent_policy.act(state))
                continue
            legal_actions, _policy_probs, _advantages = policy.distribution(state)
            labels = {
                int(action): _classify_action(state, int(action), current)
                for action in legal_actions
            }
            open_actions = [
                int(action) for action, label in labels.items() if label.startswith("open_")
            ]
            if open_actions:
                for open_action in open_actions:
                    if candidates_evaluated >= max_candidates:
                        break
                    color = int(state.hand_slots(current)[open_action // 2].color)
                    bucket = buckets[labels[open_action]]
                    _force_and_observe(
                        state,
                        forced_action=open_action,
                        forced_color=color,
                        bucket=bucket,
                        policy=policy,
                        opponent_policy=opponent_policy,
                        policy_player=policy_player,
                        window=window,
                        max_steps=max_steps,
                    )
                    candidates_evaluated += 1
            unified = policy.select_unified(state)
            state.apply_action(state.from_unified_action(unified))

    return {
        "checkpoint": str(checkpoint),
        "iteration": iteration,
        "opponent": opponent,
        "games": games,
        "seed": seed,
        "device": str(device),
        "max_candidates": max_candidates,
        "window": window,
        "elapsed_seconds": time.perf_counter() - started,
        "candidates_evaluated": candidates_evaluated,
        "buckets": {key: bucket.to_dict() for key, bucket in sorted(buckets.items())},
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("checkpoints", nargs="+", type=Path)
    parser.add_argument("--opponent", default="heuristic_cautious")
    parser.add_argument("--games", type=int, default=100)
    parser.add_argument("--seed", type=int, default=232_000)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--max-steps", type=int, default=10_000)
    parser.add_argument("--max-candidates", type=int, default=400)
    parser.add_argument(
        "--window",
        type=int,
        default=3,
        help="Number of policy_player turns to observe after the forced open.",
    )
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    device = torch.device(args.device)
    args.output.parent.mkdir(parents=True, exist_ok=True)
    rows = []
    for checkpoint in args.checkpoints:
        row = analyze_checkpoint(
            checkpoint,
            games=args.games,
            seed=args.seed,
            opponent=args.opponent,
            device=device,
            max_steps=args.max_steps,
            max_candidates=args.max_candidates,
            window=args.window,
        )
        rows.append(row)
        print(json.dumps(row, sort_keys=True))
    args.output.write_text("\n".join(json.dumps(row, sort_keys=True) for row in rows) + "\n")
    print(f"wrote {args.output}")


if __name__ == "__main__":
    main()