335 lines
12 KiB
Python
335 lines
12 KiB
Python
"""Render human-readable JAX PPO Lost Cities transcripts."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
from collections import Counter
|
|
from pathlib import Path
|
|
from typing import NamedTuple
|
|
|
|
import jax
|
|
import jax.numpy as jnp
|
|
import numpy as np
|
|
|
|
from lost_cities_jax.engine import (
|
|
board_score,
|
|
current_hand_sorted,
|
|
decode_action,
|
|
reset_from_order,
|
|
step,
|
|
)
|
|
from lost_cities_jax.opponents import policy_by_name
|
|
from lost_cities_jax.ppo import checkpoint_policy, load_config, make_shuffle_bank
|
|
from lost_cities_jax.types import (
|
|
CARDS_PER_COLOR,
|
|
DISCARD,
|
|
DRAW_DECK,
|
|
LOC_P0_BOARD,
|
|
MAX_STEPS,
|
|
N_CARDS,
|
|
N_COLORS,
|
|
PLAY,
|
|
State,
|
|
)
|
|
|
|
COLOR_NAMES = ["red", "green", "white", "blue", "yellow"]
|
|
STEP_JIT = jax.jit(step)
|
|
|
|
|
|
class PolicySpec(NamedTuple):
|
|
label: str
|
|
fn: object
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--config", required=True)
|
|
parser.add_argument("--checkpoint", required=True)
|
|
parser.add_argument("--opponent", default="heuristic_cautious")
|
|
parser.add_argument("--shuffle-bank-seed", type=int, default=20260704)
|
|
parser.add_argument("--pairs", type=int, default=10)
|
|
parser.add_argument("--output", required=True)
|
|
args = parser.parse_args()
|
|
|
|
cfg = load_config(args.config)
|
|
agent = PolicySpec("gate3_checkpoint", jax.jit(checkpoint_policy(cfg, args.checkpoint)))
|
|
opponent = PolicySpec(args.opponent, jax.jit(policy_by_name(args.opponent)))
|
|
orders = make_shuffle_bank(args.shuffle_bank_seed, args.pairs)
|
|
|
|
sections = []
|
|
summaries = []
|
|
for idx, order in enumerate(orders):
|
|
sections.append(f"\n\n## Pair {idx:02d} / learner seat 0\n")
|
|
text, summary = render_game(order, [agent, opponent], agent_seat=0, game_id=idx * 2)
|
|
sections.append(text)
|
|
summaries.append(summary)
|
|
|
|
sections.append(f"\n\n## Pair {idx:02d} / learner seat 1\n")
|
|
text, summary = render_game(order, [opponent, agent], agent_seat=1, game_id=idx * 2 + 1)
|
|
sections.append(text)
|
|
summaries.append(summary)
|
|
|
|
output = render_summary(summaries, args) + "".join(sections)
|
|
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
|
|
Path(args.output).write_text(output, encoding="utf-8")
|
|
|
|
|
|
def render_game(
|
|
deck_order: list[int], policies: list[PolicySpec], *, agent_seat: int, game_id: int
|
|
) -> tuple[str, dict]:
|
|
state = reset_from_order(jnp.asarray(deck_order, dtype=jnp.int8))
|
|
lines = [f"Game {game_id}: P0={policies[0].label}, P1={policies[1].label}"]
|
|
handshake_events = []
|
|
cautious_openings = []
|
|
cautious_useful_discards = []
|
|
|
|
while not bool(state.done):
|
|
player = int(state.to_move)
|
|
policy = policies[player]
|
|
key = jax.random.PRNGKey(game_id * 1000 + int(state.step_count))
|
|
hand = [int(x) for x in current_hand_sorted(state, player)]
|
|
action = int(policy.fn(state, jnp.asarray(player, dtype=jnp.int32), key))
|
|
hand_slot, place_type, draw_source = [int(x) for x in decode_action(jnp.asarray(action))]
|
|
card = hand[hand_slot]
|
|
drawn = peek_drawn_card(state, place_type, draw_source)
|
|
before_col_len = int(state.col_len[player, card_color(card)])
|
|
before_hand = [card_label(x) for x in hand if x < N_CARDS]
|
|
before_scores = [int(x) for x in board_score(state)]
|
|
|
|
if policy.label == "heuristic_cautious":
|
|
if place_type == PLAY and before_col_len == 0:
|
|
cautious_openings.append(
|
|
{
|
|
"ply": int(state.step_count),
|
|
"seat": player,
|
|
"card": card_label(card),
|
|
"rank": card_rank(card),
|
|
}
|
|
)
|
|
if place_type == DISCARD and can_play_card(state, 1 - player, card):
|
|
cautious_useful_discards.append(
|
|
{
|
|
"ply": int(state.step_count),
|
|
"seat": player,
|
|
"card": card_label(card),
|
|
"opponent": 1 - player,
|
|
}
|
|
)
|
|
|
|
if place_type == PLAY and is_handshake(card):
|
|
handshake_events.append(
|
|
{
|
|
"ply": int(state.step_count),
|
|
"seat": player,
|
|
"policy": policy.label,
|
|
"card": card_label(card),
|
|
"hand": before_hand,
|
|
}
|
|
)
|
|
|
|
next_state, _, _ = STEP_JIT(state, jnp.asarray(action, dtype=jnp.int32))
|
|
place_text = "PLAY" if place_type == PLAY else "DISCARD"
|
|
target = "expedition" if place_type == PLAY else "discard"
|
|
draw_text = "deck" if draw_source == DRAW_DECK else f"{COLOR_NAMES[draw_source - 1]} pile"
|
|
lines.append(
|
|
f"Ply {int(state.step_count):03d} P{player} {policy.label}: "
|
|
f"{place_text} {card_label(card)} to {target}; "
|
|
f"DRAW {draw_text} -> {card_label(drawn)}; "
|
|
f"score_before={before_scores}"
|
|
)
|
|
lines.extend(render_board(next_state))
|
|
state = next_state
|
|
|
|
final_scores = [int(x) for x in board_score(state)]
|
|
lines.append(f"Final score: P0={final_scores[0]} P1={final_scores[1]}")
|
|
lines.extend(render_final_breakdown(state))
|
|
agent_score = final_scores[agent_seat]
|
|
opponent_score = final_scores[1 - agent_seat]
|
|
summary = {
|
|
"game_id": game_id,
|
|
"agent_seat": agent_seat,
|
|
"agent_score_diff": agent_score - opponent_score,
|
|
"agent_opened_colors": int(jnp.sum(state.col_len[agent_seat] > 0)),
|
|
"length": int(state.step_count),
|
|
"max_steps": int(state.step_count) >= MAX_STEPS,
|
|
"handshake_events": handshake_events,
|
|
"cautious_openings": cautious_openings,
|
|
"cautious_useful_discards": cautious_useful_discards,
|
|
}
|
|
return "\n".join(lines) + "\n", summary
|
|
|
|
|
|
def render_summary(summaries: list[dict], args: argparse.Namespace) -> str:
|
|
diffs = np.asarray([item["agent_score_diff"] for item in summaries], dtype=np.float64)
|
|
opened = np.asarray([item["agent_opened_colors"] for item in summaries], dtype=np.float64)
|
|
lengths = np.asarray([item["length"] for item in summaries], dtype=np.float64)
|
|
handshakes = [event for item in summaries for event in item["handshake_events"]]
|
|
cautious_openings = [event for item in summaries for event in item["cautious_openings"]]
|
|
useful_discards = [event for item in summaries for event in item["cautious_useful_discards"]]
|
|
opening_ranks = Counter(str(event["rank"]) for event in cautious_openings)
|
|
low_openings = [
|
|
event for event in cautious_openings if event["rank"] not in ("HS",) and event["rank"] < 7
|
|
]
|
|
|
|
lines = [
|
|
"# Gate-3 Transcript Dump",
|
|
"",
|
|
f"config: `{args.config}`",
|
|
f"checkpoint: `{args.checkpoint}`",
|
|
f"opponent: `{args.opponent}`",
|
|
f"duplicate pairs: `{args.pairs}`",
|
|
"",
|
|
"## Summary",
|
|
"",
|
|
f"games: {len(summaries)}",
|
|
f"agent_score_diff_mean: {float(np.mean(diffs)):.3f}",
|
|
f"agent_score_diff_min_max: {int(np.min(diffs))} / {int(np.max(diffs))}",
|
|
f"agent_opened_colors_mean: {float(np.mean(opened)):.3f}",
|
|
f"game_length_mean: {float(np.mean(lengths)):.3f}",
|
|
f"game_length_min_p50_p95_max: {int(np.min(lengths))} / "
|
|
f"{float(np.quantile(lengths, 0.50)):.1f} / "
|
|
f"{float(np.quantile(lengths, 0.95)):.1f} / {int(np.max(lengths))}",
|
|
f"max_steps_rate: {float(np.mean([item['max_steps'] for item in summaries])):.3f}",
|
|
f"handshake_play_events: {len(handshakes)}",
|
|
f"cautious_openings: {len(cautious_openings)}",
|
|
f"cautious_opening_rank_counts: {json.dumps(dict(sorted(opening_ranks.items())))}",
|
|
f"cautious_low_openings_lt7: {len(low_openings)}",
|
|
f"cautious_discards_immediately_playable_by_opponent: {len(useful_discards)}",
|
|
"",
|
|
"## Handshake Play Contexts",
|
|
"",
|
|
]
|
|
if handshakes:
|
|
for event in handshakes[:80]:
|
|
lines.append(
|
|
f"- game_event ply={event['ply']} seat=P{event['seat']} "
|
|
f"policy={event['policy']} card={event['card']} hand={event['hand']}"
|
|
)
|
|
if len(handshakes) > 80:
|
|
lines.append(f"- ... {len(handshakes) - 80} more")
|
|
else:
|
|
lines.append("- none")
|
|
|
|
lines.extend(["", "## Cautious Opening Audit", ""])
|
|
if cautious_openings:
|
|
for event in cautious_openings[:80]:
|
|
lines.append(
|
|
f"- ply={event['ply']} seat=P{event['seat']} card={event['card']} "
|
|
f"rank={event['rank']}"
|
|
)
|
|
if len(cautious_openings) > 80:
|
|
lines.append(f"- ... {len(cautious_openings) - 80} more")
|
|
else:
|
|
lines.append("- none")
|
|
|
|
lines.extend(["", "## Cautious Useful Discard Audit", ""])
|
|
if useful_discards:
|
|
for event in useful_discards[:80]:
|
|
lines.append(
|
|
f"- ply={event['ply']} seat=P{event['seat']} card={event['card']} "
|
|
f"opponent=P{event['opponent']}"
|
|
)
|
|
if len(useful_discards) > 80:
|
|
lines.append(f"- ... {len(useful_discards) - 80} more")
|
|
else:
|
|
lines.append("- none")
|
|
lines.append("")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def render_board(state: State) -> list[str]:
|
|
scores = [int(x) for x in board_score(state)]
|
|
lines = [f" Board scores: P0={scores[0]} P1={scores[1]}"]
|
|
for player in range(2):
|
|
parts = []
|
|
for color in range(N_COLORS):
|
|
cards = board_cards(state, player, color)
|
|
score = color_score(cards)
|
|
parts.append(f"{COLOR_NAMES[color]}={format_cards(cards)}({score:+d})")
|
|
lines.append(f" P{player}: " + " | ".join(parts))
|
|
return lines
|
|
|
|
|
|
def render_final_breakdown(state: State) -> list[str]:
|
|
lines = ["Final color breakdown:"]
|
|
for player in range(2):
|
|
parts = []
|
|
for color in range(N_COLORS):
|
|
cards = board_cards(state, player, color)
|
|
parts.append(f"{COLOR_NAMES[color]} {format_cards(cards)} => {color_score(cards):+d}")
|
|
lines.append(f" P{player}: " + "; ".join(parts))
|
|
return lines
|
|
|
|
|
|
def board_cards(state: State, player: int, color: int) -> list[int]:
|
|
loc = np.asarray(state.card_loc)
|
|
board_loc = LOC_P0_BOARD + player
|
|
start = color * CARDS_PER_COLOR
|
|
cards = [card for card in range(start, start + CARDS_PER_COLOR) if int(loc[card]) == board_loc]
|
|
return sorted(
|
|
cards, key=lambda card: (0 if is_handshake(card) else 1, card_rank_value(card), card)
|
|
)
|
|
|
|
|
|
def color_score(cards: list[int]) -> int:
|
|
if not cards:
|
|
return 0
|
|
handshakes = sum(1 for card in cards if is_handshake(card))
|
|
ranks = [card_rank_value(card) for card in cards if not is_handshake(card)]
|
|
score = (sum(ranks) - 20) * (1 + handshakes)
|
|
if len(cards) >= 8:
|
|
score += 20
|
|
return score
|
|
|
|
|
|
def peek_drawn_card(state: State, place_type: int, draw_source: int) -> int:
|
|
del place_type
|
|
if draw_source == DRAW_DECK:
|
|
return int(state.deck_order[int(state.draw_ptr)])
|
|
color = draw_source - 1
|
|
length = int(state.pile_len[color])
|
|
return int(state.pile[color, length - 1])
|
|
|
|
|
|
def can_play_card(state: State, player: int, card: int) -> bool:
|
|
top = int(state.col_top[player, card_color(card)])
|
|
if is_handshake(card):
|
|
return top == 0
|
|
return card_rank_value(card) > top
|
|
|
|
|
|
def format_cards(cards: list[int]) -> str:
|
|
if not cards:
|
|
return "[]"
|
|
return "[" + ",".join(card_label(card) for card in cards) + "]"
|
|
|
|
|
|
def card_label(card: int) -> str:
|
|
if card < 0 or card >= N_CARDS:
|
|
return "none"
|
|
color = COLOR_NAMES[card_color(card)][0].upper()
|
|
if is_handshake(card):
|
|
return f"{color}HS{card % CARDS_PER_COLOR + 1}"
|
|
return f"{color}{card_rank_value(card)}"
|
|
|
|
|
|
def card_color(card: int) -> int:
|
|
return card // CARDS_PER_COLOR
|
|
|
|
|
|
def is_handshake(card: int) -> bool:
|
|
return card % CARDS_PER_COLOR < 3
|
|
|
|
|
|
def card_rank(card: int) -> int | str:
|
|
return "HS" if is_handshake(card) else card_rank_value(card)
|
|
|
|
|
|
def card_rank_value(card: int) -> int:
|
|
return card % CARDS_PER_COLOR - 1
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|