Files
coorl-lost-cities/src/lost_cities_jax/human_play.py
T

618 lines
20 KiB
Python

"""Terminal human-play client and log summarizer for JAX PPO checkpoints."""
from __future__ import annotations
import json
import math
import time
import uuid
from collections import defaultdict
from collections.abc import Callable
from dataclasses import dataclass
from pathlib import Path
from typing import Any
import jax
import jax.numpy as jnp
import numpy as np
from lost_cities_jax.engine import (
board_score,
current_hand_sorted,
decode_action,
legal_action_mask,
reset_from_order,
step,
)
from lost_cities_jax.obs import observation
from lost_cities_jax.ppo import (
ActorCritic,
JaxPPOConfig,
create_train_state,
load_config,
mask_logits,
restore_checkpoint,
)
from lost_cities_jax.types import (
CARDS_PER_COLOR,
DISCARD,
DRAW_DECK,
LOC_P0_BOARD,
MAX_STEPS,
N_ACTIONS,
N_CARDS,
N_COLORS,
PLAY,
State,
)
COLOR_NAMES = ("R", "G", "W", "B", "Y")
COLOR_WORDS = {
"R": 0,
"RED": 0,
"G": 1,
"GREEN": 1,
"W": 2,
"WHITE": 2,
"B": 3,
"BLUE": 3,
"Y": 4,
"YELLOW": 4,
}
DEFAULT_HUMAN_LOG_DIR = Path("/mnt/2tbhdd/coolrl-lost-cities-artifacts/human-play")
DEFAULT_HUMAN_BANK_STATE = DEFAULT_HUMAN_LOG_DIR / "shuffle_bank_state.json"
DEFAULT_HUMAN_BANK_SEED = 20260706
@dataclass(frozen=True)
class PolicyEval:
action: int
top3: list[dict[str, Any]]
value: float
def play_cli(
*,
checkpoint: str | Path,
config: str | Path | None = None,
seat: int = 0,
seed: int = DEFAULT_HUMAN_BANK_SEED,
duplicate: bool = False,
log_dir: str | Path = DEFAULT_HUMAN_LOG_DIR,
bank_state: str | Path = DEFAULT_HUMAN_BANK_STATE,
input_fn: Callable[[str], str] = input,
output_fn: Callable[[str], None] = print,
) -> list[dict[str, Any]]:
cfg = load_config(infer_config_path(checkpoint) if config is None else config)
agent = load_agent(cfg, checkpoint)
log_dir = Path(log_dir)
deck_index, deck_order = reserve_deck_order(seed, Path(bank_state))
session_id = time.strftime("%Y%m%d-%H%M%S") + "-" + uuid.uuid4().hex[:8]
seats = session_human_seats(seat, duplicate)
records = []
set_diff = 0.0
for game_index, human_seat in enumerate(seats, start=1):
output_fn("")
output_fn(f"Game {game_index}/{len(seats)} - human seat P{human_seat}")
record = play_one_game(
cfg=cfg,
agent=agent,
deck_order=deck_order,
human_seat=human_seat,
session_id=session_id,
duplicate_set_id=session_id if duplicate else None,
duplicate_game_index=game_index if duplicate else None,
deck_seed=seed,
deck_index=deck_index,
input_fn=input_fn,
output_fn=output_fn,
)
records.append(record)
set_diff += float(record["human_score_diff"])
append_human_log(log_dir, record)
if duplicate:
output_fn("")
output_fn(f"Duplicate set total human score diff: {set_diff:+.0f}")
return records
def summarize_logs(log_dir: str | Path = DEFAULT_HUMAN_LOG_DIR) -> dict[str, Any]:
records = read_human_logs(log_dir)
games = len(records)
duplicate_sets: dict[str, list[dict[str, Any]]] = defaultdict(list)
singletons = []
for record in records:
set_id = record.get("duplicate_set_id")
if set_id:
duplicate_sets[set_id].append(record)
else:
singletons.append(record)
set_diffs = [
sum(float(game["human_score_diff"]) for game in group)
for group in duplicate_sets.values()
if len(group) == 2
]
game_diffs = [float(record["human_score_diff"]) for record in records]
ai_opened = [float(record["ai_opened_colors"]) for record in records]
lengths = [float(record["game_length"]) for record in records]
swings = collect_value_swings(records)[:10]
return {
"games": games,
"duplicate_sets": len(set_diffs),
"single_games": len(singletons),
"human_game_score_diff_mean": mean(game_diffs),
"human_game_score_diff_ci95": ci95(game_diffs),
"human_duplicate_set_diff_mean": mean(set_diffs),
"human_duplicate_set_diff_ci95": ci95(set_diffs),
"ai_opened_colors_mean": mean(ai_opened),
"game_length_mean": mean(lengths),
"value_swings_top10": swings,
}
def session_human_seats(seat: int, duplicate: bool) -> list[int]:
if seat not in {0, 1}:
raise ValueError("seat must be 0 or 1")
return [seat, 1 - seat] if duplicate else [seat]
def print_summary(summary: dict[str, Any], output_fn: Callable[[str], None] = print) -> None:
output_fn(json.dumps(summary, indent=2, sort_keys=True, ensure_ascii=False))
def play_one_game(
*,
cfg: JaxPPOConfig,
agent: tuple[Any, Any],
deck_order: list[int],
human_seat: int,
session_id: str,
duplicate_set_id: str | None,
duplicate_game_index: int | None,
deck_seed: int,
deck_index: int,
input_fn: Callable[[str], str],
output_fn: Callable[[str], None],
human_action_fn: Callable[[State, int], int] | None = None,
) -> dict[str, Any]:
params, model = agent
state = reset_from_order(jnp.asarray(deck_order, dtype=jnp.int8))
actions = []
move_logs = []
while not bool(state.done) and int(state.step_count) < MAX_STEPS:
player = int(state.to_move)
policy_eval = evaluate_agent_policy(cfg, params, model, state, player)
output_fn(render_public_state(state, human_seat))
if player == human_seat:
if human_action_fn is None:
action = prompt_human_action(state, human_seat, input_fn, output_fn)
else:
action = int(human_action_fn(state, human_seat))
actor = "human"
else:
action = policy_eval.action
actor = "ai"
output_fn(f"AI: {describe_action(state, action)}")
before = state
state, _, _ = step(state, jnp.asarray(action, dtype=jnp.int32))
actions.append(int(action))
move_logs.append(
{
"ply": int(before.step_count),
"actor": actor,
"player": player,
"action": int(action),
"action_text": describe_action(before, action),
"ai_policy_player": player,
"ai_top3": policy_eval.top3,
"ai_value": policy_eval.value,
}
)
output_fn(render_public_state(state, human_seat))
output_fn(render_score_breakdown(state))
scores = np.asarray(board_score(state), dtype=np.float32)
ai_seat = 1 - human_seat
comment = input_fn("Post-game comment (optional, Enter to skip): ").strip()
return {
"schema": "lost-cities-jax-human-play-v1",
"session_id": session_id,
"duplicate_set_id": duplicate_set_id,
"duplicate_game_index": duplicate_game_index,
"deck_seed": deck_seed,
"deck_index": deck_index,
"human_seat": human_seat,
"ai_seat": ai_seat,
"deck_order": [int(card) for card in deck_order],
"actions": actions,
"moves": move_logs,
"final_scores": {"p0": float(scores[0]), "p1": float(scores[1])},
"human_score_diff": float(scores[human_seat] - scores[ai_seat]),
"ai_score_diff": float(scores[ai_seat] - scores[human_seat]),
"score_breakdown": score_breakdown(state),
"game_length": int(state.step_count),
"ai_opened_colors": int(np.asarray(state.col_len)[ai_seat].astype(bool).sum()),
"human_opened_colors": int(np.asarray(state.col_len)[human_seat].astype(bool).sum()),
"max_steps": bool(int(state.step_count) >= MAX_STEPS),
"human_comment": comment,
"created_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"),
}
def load_agent(cfg: JaxPPOConfig, checkpoint: str | Path) -> tuple[Any, ActorCritic]:
state = create_train_state(cfg, jax.random.PRNGKey(0))
state = restore_checkpoint(Path(checkpoint), state)
model = ActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
return state.params, model
def infer_config_path(checkpoint: str | Path) -> Path:
path = Path(checkpoint).resolve()
candidates = []
for parent in [path, *path.parents]:
candidates.extend([parent / "main_ppo_config.json", parent / "config.json"])
for candidate in candidates:
if candidate.exists():
return candidate
msg = f"could not infer config for checkpoint; pass --config explicitly: {checkpoint}"
raise FileNotFoundError(msg)
def evaluate_agent_policy(
cfg: JaxPPOConfig,
params: Any,
model: ActorCritic,
state: State,
player: int,
) -> PolicyEval:
obs = observation(state, jnp.asarray(player, dtype=jnp.int32))[None, :]
legal = legal_action_mask(state)[None, :]
logits, value = model.apply(params, obs)
masked = mask_logits(logits, legal)
probs = np.asarray(jax.nn.softmax(masked, axis=-1)[0], dtype=np.float64)
legal_np = np.asarray(legal[0], dtype=bool)
order = np.argsort(-probs)
top = []
for action in order:
if not legal_np[action]:
continue
top.append(
{
"action": int(action),
"prob": float(probs[action]),
"text": describe_action(state, int(action)),
}
)
if len(top) == 3:
break
return PolicyEval(action=int(order[0]), top3=top, value=float(np.asarray(value)[0]))
def prompt_human_action(
state: State,
player: int,
input_fn: Callable[[str], str],
output_fn: Callable[[str], None],
) -> int:
while True:
raw = input_fn("Your move (e.g. play R7 draw deck): ").strip()
try:
action = parse_human_action(raw, state, player)
except ValueError as exc:
output_fn(f"Invalid input: {exc}")
continue
mask = np.asarray(legal_action_mask(state), dtype=bool)
if action < 0 or action >= N_ACTIONS or not mask[action]:
output_fn(f"Illegal move: {illegal_reason(state, action)}")
continue
return action
def parse_human_action(raw: str, state: State, player: int) -> int:
parts = raw.strip().split()
if len(parts) != 4 or parts[2].lower() != "draw":
msg = "expected '<play|discard> <card> draw <deck|R|G|W|B|Y>'"
raise ValueError(msg)
place_word, card_word, _, draw_word = parts
place = parse_place(place_word)
card_id = parse_card_token(card_word)
hand = [
int(card) for card in np.asarray(current_hand_sorted(state, player)) if int(card) < N_CARDS
]
matching = [idx for idx, card in enumerate(hand) if same_card_face(card, card_id)]
if not matching:
msg = f"card is not in your hand: {card_word}"
raise ValueError(msg)
hand_slot = matching[0]
draw_source = parse_draw_source(draw_word)
return encode_action(hand_slot, place, draw_source)
def parse_place(value: str) -> int:
lowered = value.lower()
if lowered == "play":
return PLAY
if lowered == "discard":
return DISCARD
msg = "placement must be 'play' or 'discard'"
raise ValueError(msg)
def parse_draw_source(value: str) -> int:
upper = value.upper()
if upper == "DECK":
return DRAW_DECK
if upper in COLOR_WORDS:
return COLOR_WORDS[upper] + 1
msg = "draw source must be deck or one of R/G/W/B/Y"
raise ValueError(msg)
def parse_card_token(value: str) -> int:
token = value.strip().upper()
if len(token) < 2:
raise ValueError("card token is too short")
color = COLOR_WORDS.get(token[0])
if color is None:
raise ValueError("card color must be R/G/W/B/Y")
suffix = token[1:]
if suffix in {"H", "HS", "W", "WAGER"}:
slot = 0
else:
try:
rank = int(suffix)
except ValueError as exc:
raise ValueError("rank must be 2..10 or HS") from exc
if rank < 2 or rank > 10:
raise ValueError("rank must be 2..10")
slot = rank + 1
return color * CARDS_PER_COLOR + slot
def same_card_face(left: int, right: int) -> bool:
return left // CARDS_PER_COLOR == right // CARDS_PER_COLOR and card_label(left) == card_label(
right
)
def encode_action(hand_slot: int, place_type: int, draw_source: int) -> int:
return hand_slot * 12 + place_type * 6 + draw_source
def illegal_reason(state: State, action: int) -> str:
if action < 0 or action >= N_ACTIONS:
return "action id outside action space"
hand_slot, place_type, draw_source = [int(x) for x in decode_action(jnp.asarray(action))]
hand = [int(card) for card in np.asarray(current_hand_sorted(state)) if int(card) < N_CARDS]
if hand_slot >= len(hand):
return "selected hand slot is empty"
card = hand[hand_slot]
color = card // CARDS_PER_COLOR
rank = card_rank(card)
player = int(state.to_move)
if place_type == PLAY:
top = int(np.asarray(state.col_top)[player, color])
if rank == 0 and top > 0:
return "handshake cannot be played after a number card"
if rank > 0 and rank <= top:
return f"number card must be above current top rank {top}"
if draw_source != DRAW_DECK:
pile_color = draw_source - 1
pile_len = int(np.asarray(state.pile_len)[pile_color])
if pile_len <= 0:
return "discard pile is empty"
if place_type == DISCARD and pile_color == color:
return "cannot draw the card you just discarded"
return "move is not legal under the current mask"
def render_public_state(state: State, player: int) -> str:
scores = np.asarray(board_score(state), dtype=np.float32)
opponent = 1 - player
lines = [
"",
f"Ply {int(state.step_count)} | to move P{int(state.to_move)} | deck left {N_CARDS - int(state.draw_ptr)}",
f"Board score: P{player} {scores[player]:+.0f} / P{opponent} {scores[opponent]:+.0f} "
f"(you {scores[player] - scores[opponent]:+.0f})",
f"Your hand: {' '.join(card_name(card) for card in np.asarray(current_hand_sorted(state, player)) if int(card) < N_CARDS)}",
"Boards:",
*render_boards(state),
"Discards:",
*render_piles(state),
]
return "\n".join(lines)
def render_boards(state: State) -> list[str]:
loc = np.asarray(state.card_loc)
lines = []
for player in range(2):
color_parts = []
for color in range(N_COLORS):
cards = [
card
for card in range(N_CARDS)
if loc[card] == LOC_P0_BOARD + player and card // CARDS_PER_COLOR == color
]
cards.sort(key=lambda card: (card_rank(card) > 0, card_rank(card), card))
color_parts.append(f"{COLOR_NAMES[color]}:[{' '.join(card_name(c) for c in cards)}]")
lines.append(f" P{player} " + " ".join(color_parts))
return lines
def render_piles(state: State) -> list[str]:
pile = np.asarray(state.pile)
pile_len = np.asarray(state.pile_len)
lines = []
for color in range(N_COLORS):
cards = [int(card) for card in pile[color, : int(pile_len[color])]]
lines.append(f" {COLOR_NAMES[color]}: {' '.join(card_name(card) for card in cards)}")
return lines
def render_score_breakdown(state: State) -> str:
breakdown = score_breakdown(state)
lines = ["Final scoring:"]
for player in range(2):
parts = []
for color in COLOR_NAMES:
item = breakdown[f"p{player}"][color]
parts.append(f"{color}:{item['score']:+.0f}")
lines.append(f" P{player} {' '.join(parts)} total {breakdown[f'p{player}']['total']:+.0f}")
return "\n".join(lines)
def score_breakdown(state: State) -> dict[str, Any]:
loc = np.asarray(state.card_loc)
col_len = np.asarray(state.col_len)
col_hs = np.asarray(state.col_hs)
data: dict[str, Any] = {}
for player in range(2):
total = 0
player_data: dict[str, Any] = {}
for color in range(N_COLORS):
ranks = [
card_rank(card)
for card in range(N_CARDS)
if loc[card] == LOC_P0_BOARD + player
and card // CARDS_PER_COLOR == color
and card_rank(card) > 0
]
length = int(col_len[player, color])
hs = int(col_hs[player, color])
rank_sum = sum(ranks)
score = 0 if length == 0 else (rank_sum - 20) * (1 + hs) + (20 if length >= 8 else 0)
total += score
player_data[COLOR_NAMES[color]] = {
"rank_sum": rank_sum,
"handshakes": hs,
"length": length,
"score": float(score),
}
player_data["total"] = float(total)
data[f"p{player}"] = player_data
return data
def describe_action(state: State, action: int) -> str:
hand_slot, place_type, draw_source = [int(x) for x in decode_action(jnp.asarray(action))]
player = int(state.to_move)
hand = [
int(card) for card in np.asarray(current_hand_sorted(state, player)) if int(card) < N_CARDS
]
card = hand[hand_slot] if hand_slot < len(hand) else -1
place = "play" if place_type == PLAY else "discard"
draw = "deck" if draw_source == DRAW_DECK else COLOR_NAMES[draw_source - 1]
return f"{place} {card_name(card)} draw {draw}"
def card_name(card: int) -> str:
if card < 0 or card >= N_CARDS:
return "?"
return COLOR_NAMES[card // CARDS_PER_COLOR] + card_label(card)
def card_label(card: int) -> str:
rank = card_rank(card)
return "HS" if rank == 0 else str(rank)
def card_rank(card: int) -> int:
slot = card % CARDS_PER_COLOR
return 0 if slot < 3 else slot - 1
def reserve_deck_order(seed: int, bank_state: Path) -> tuple[int, list[int]]:
bank_state.parent.mkdir(parents=True, exist_ok=True)
if bank_state.exists():
state = json.loads(bank_state.read_text(encoding="utf-8"))
next_index = int(state.get(str(seed), 0))
else:
state = {}
next_index = 0
state[str(seed)] = next_index + 1
bank_state.write_text(json.dumps(state, indent=2, sort_keys=True) + "\n", encoding="utf-8")
return next_index, human_deck_order(seed, next_index)
def human_deck_order(seed: int, index: int) -> list[int]:
rng = np.random.default_rng(seed + index)
return rng.permutation(N_CARDS).astype("int8").astype(int).tolist()
def append_human_log(log_dir: Path, record: dict[str, Any]) -> Path:
log_dir.mkdir(parents=True, exist_ok=True)
path = log_dir / "games.jsonl"
with path.open("a", encoding="utf-8") as handle:
handle.write(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n")
return path
def read_human_logs(log_dir: str | Path) -> list[dict[str, Any]]:
path = Path(log_dir) / "games.jsonl"
if not path.exists():
return []
return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line]
def collect_value_swings(records: list[dict[str, Any]]) -> list[dict[str, Any]]:
swings = []
for record in records:
previous = None
for move in record.get("moves", []):
value = float(move.get("ai_value", 0.0))
if previous is not None:
swings.append(
{
"session_id": record.get("session_id"),
"duplicate_game_index": record.get("duplicate_game_index"),
"ply": move.get("ply"),
"delta": value - previous,
"abs_delta": abs(value - previous),
"action_text": move.get("action_text"),
}
)
previous = value
swings.sort(key=lambda row: row["abs_delta"], reverse=True)
return swings
def mean(values: list[float]) -> float | None:
if not values:
return None
return float(np.mean(np.asarray(values, dtype=np.float64)))
def ci95(values: list[float]) -> list[float | None]:
if not values:
return [None, None]
arr = np.asarray(values, dtype=np.float64)
if arr.size <= 1:
value = float(arr[0])
return [value, value]
se = float(np.std(arr, ddof=1) / math.sqrt(arr.size))
avg = float(np.mean(arr))
return [avg - 1.959963984540054 * se, avg + 1.959963984540054 * se]
__all__ = [
"append_human_log",
"card_name",
"describe_action",
"human_deck_order",
"infer_config_path",
"parse_human_action",
"play_cli",
"play_one_game",
"read_human_logs",
"render_public_state",
"reserve_deck_order",
"session_human_seats",
"summarize_logs",
]