Add JAX PPO ladder verification pass

This commit is contained in:
2026-07-05 00:17:49 +09:00
parent bb52ef9ec1
commit 7fe0bfdcfe
6 changed files with 14243 additions and 3 deletions
+40
View File
@@ -0,0 +1,40 @@
run:
experiment_name: jax-ppo-gate3-exploiter
seed: 20260705
learner_seat: 0
total_updates: 250
log_every: 1
checkpoint_every: 10
artifact_root: /mnt/2tbhdd/coolrl-lost-cities-artifacts/verification
opponent:
name: gate3_checkpoint_frozen
network:
hidden_size: 512
num_layers: 3
ppo:
batch_games: 8192
rollout_steps: 400
gamma: 1.0
gae_lambda: 0.95
clip_epsilon: 0.2
entropy_coef: 0.01
value_coef: 0.5
max_grad_norm: 0.5
learning_rate: 0.0003
epochs: 4
minibatches: 128
reward:
terminal_scale: 50.0
potential_shaping_initial: 1.0
potential_shaping_final: 0.0
potential_shaping_anneal_steps: 5000000
evaluation:
games: 10000
duplicate: true
shuffle_bank_seed: 20260704
batch_games: 8192
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,191 @@
# Verification Pass - 2026-07-04
**Status:** FAIL - stop before league self-play.
**Implementation start commit:** `bb52ef9`.
**Primary checkpoint under audit:** `/mnt/2tbhdd/coolrl-lost-cities-artifacts/jax-ppo-static-opponents/2026-07-04_230150_jax-ppo-cautious/latest`.
**Verification artifacts:** `/mnt/2tbhdd/coolrl-lost-cities-artifacts/verification/2026-07-04_phase1/`.
## Decision
Do not start snapshot-pool league self-play from the current ladder artifacts.
The duplicate/Wilson protocol is reproducible, and the obvious seat-perspective
bug was not found, but the cautious anchor is not a trustworthy permanent Elo
anchor: it stalls heavily in mirror play and gives the trained policy a very
clean win while the trained policy opens almost all colors.
## Transcript Audit
Transcript file:
```text
docs/reports/verification-pass-2026-07-04-transcripts.txt
```
Rendered protocol: 10 fixed shuffle-bank deck orders, duplicate seat-swapped,
for 20 total games of gate-3 checkpoint vs `heuristic_cautious`.
Summary from the transcript dump:
| Metric | Value |
| --- | ---: |
| Games | 20 |
| Agent mean score diff | +144.500 |
| Agent score diff min/max | -93 / +230 |
| Agent opened colors mean | 5.000 |
| Mean game length | 159.200 |
| Length min / p50 / p95 / max | 63 / 109 / 400 / 400 |
| Max-steps rate | 0.150 |
| Handshake play events | 86 |
| Cautious openings | 78 |
| Cautious low openings below rank 7 | 3 |
| Cautious discards immediately playable by opponent | 162 |
Interpretation: the sample confirms the concern behind this pass. The agent is
not winning by a restrained 2-3 color expert pattern; it opens all 5 colors on
average and still wins by a wide margin. The cautious policy is also leaking
many immediately useful discard tops.
## Protocol Recheck
The original ladder evaluation was audited in code. The `eval` path uses:
- explicit shuffle banks generated from `evaluation.shuffle_bank_seed`;
- duplicate seat-swapped loops over learner seat 0 and learner seat 1;
- Wilson confidence intervals from aggregate wins.
The protocol itself is not the main failure. The failure is that the original
gate summary omitted `max_steps_rate` and score-diff distribution canaries.
Recomputed gate-3 duplicate evaluation with distribution canaries:
```text
/mnt/2tbhdd/coolrl-lost-cities-artifacts/verification/2026-07-04_phase1/gate3_reverification_distribution.json
```
| Metric | Value |
| --- | ---: |
| Games | 20,000 |
| Win rate | 0.95960 |
| Wilson 95% | [0.95678, 0.96224] |
| Mean score diff | +142.9178 |
| Score diff p05 / p50 / p95 | +21 / +152 / +227 |
| Opened colors/game | 4.98965 |
| Positive expeditions/game | 4.09470 |
| Play action rate | 0.27639 |
| Mean game length | 185.3179 |
| Game length p50 / p95 / max | 115 / 400 / 400 |
| Max-steps rate | 0.24070 |
The 24.07% forced-end rate is a high-severity canary failure for using this
gate as a clean league baseline.
## Perspective And Mirror Tests
New tests added:
- hand-authored heuristic behavior checks for weak unopened hands, strong
balanced opens, and P1 own-board perspective;
- static-policy duplicate mirror checks for all three heuristic policies;
- gate-3 checkpoint duplicate self-mirror check, skipped on CPU-only runs and
executed under CUDA JAX locally.
Commands run:
```bash
uv run pytest -q tests/lost_cities_jax/test_ppo_stack.py
uv run --with 'jax[cuda12]' pytest -q \
tests/lost_cities_jax/test_ppo_stack.py::test_gate3_checkpoint_duplicate_self_mirror_score_diff_is_zero
```
Results:
- CPU focused tests: 11 passed, 1 skipped.
- CUDA checkpoint self-mirror test: 1 passed.
Mirror evaluation artifacts:
```text
/mnt/2tbhdd/coolrl-lost-cities-artifacts/verification/2026-07-04_phase1/mirror_discard_only.json
/mnt/2tbhdd/coolrl-lost-cities-artifacts/verification/2026-07-04_phase1/mirror_heuristic_balanced.json
/mnt/2tbhdd/coolrl-lost-cities-artifacts/verification/2026-07-04_phase1/mirror_heuristic_cautious.json
/mnt/2tbhdd/coolrl-lost-cities-artifacts/verification/2026-07-04_phase1/mirror_gate3_checkpoint.json
```
| Mirror policy | Mean diff | Win/loss symmetry | Max-steps rate | Play action rate | Mean length |
| --- | ---: | --- | ---: | ---: | ---: |
| `discard_only` | 0.0 | exact | 0.0000 | 0.0000 | 44.0000 |
| `heuristic_balanced` | 0.0 | exact | 0.4683 | 0.0422 | 231.1210 |
| `heuristic_cautious` | 0.0 | exact | 0.8956 | 0.0053 | 368.8435 |
| gate-3 checkpoint | 0.0 | exact | 0.0000 | 0.7839 | 45.4292 |
Interpretation: the broad P0/P1 perspective bug is unlikely. Duplicate mirror
score differences cancel exactly. The serious issue is heuristic quality:
`heuristic_cautious` mirror play almost never opens expeditions and reaches
forced termination in 89.56% of games.
## Exploiter Baseline
Frozen opponent: gate-3 checkpoint. New exploiter: random init PPO, same 8192 x
400 rollout shape, 250 updates, shaping annealed by the standard schedule.
Training run:
```text
/mnt/2tbhdd/coolrl-lost-cities-artifacts/verification/2026-07-04_235948_jax-ppo-gate3-exploiter/
```
Final train-row canaries:
| Metric | Value |
| --- | ---: |
| return_mean | +0.24118 |
| play_action_rate | 0.71362 |
| opened_colors_mean | 4.91016 |
| positive_expeditions_mean | 2.07922 |
| game_length_mean | 49.12524 |
| max_steps_rate | 0.00000 |
Duplicate exploiter evaluation:
```text
/mnt/2tbhdd/coolrl-lost-cities-artifacts/verification/2026-07-04_phase1/exploiter_vs_gate3_duplicate.json
```
| Metric | Value |
| --- | ---: |
| Games | 20,000 |
| Exploiter win rate | 0.67115 |
| Wilson 95% | [0.66461, 0.67763] |
| Mean score diff | +24.1371 |
| Score diff p05 / p50 / p95 | -50 / +21 / +110 |
| Opened colors/game | 4.89710 |
| Positive expeditions/game | 2.12865 |
| Play action rate | 0.74612 |
| Mean game length | 48.0017 |
| Max-steps rate | 0.00000 |
This is now the measured exploitability baseline. It is not by itself a stop
condition, but it reinforces that the current gate-3 checkpoint is not a
finished robust policy.
## Findings
| Severity | Finding | Evidence | Consequence |
| --- | --- | --- | --- |
| High | `heuristic_cautious` is a stalling anchor, not a strong cautious baseline. | Mirror max-steps rate 89.56%, play action rate 0.0053. | Do not use it as a permanent Elo anchor without fixing/replacing it. |
| High | Gate-3 result is contaminated by forced-end games and over-opening. | Gate-3 max-steps rate 24.07%, opened colors/game 4.98965. | The ladder PASS remains reproducible, but it should not become a league baseline. |
| Medium | Cautious discards many immediately playable cards to the opponent. | Transcript sample: 162 such discards in 20 games. | The agent may be exploiting discard leakage rather than learning robust play. |
| Medium | A same-budget PPO exploiter beats the gate-3 checkpoint. | Exploiter duplicate win rate 67.115%, mean diff +24.1371. | League success should reduce this number, but Phase 2 should wait for anchor repair. |
| Low | No broad seat perspective bug found. | Static and checkpoint duplicate mirrors have exact zero mean diff. | Seat symmetry is not the likely explanation for the ladder result. |
## Stop Condition
Phase 1 did not pass cleanly. Per the work order, Phase 2 is not started.
Recommended next work:
1. Replace or repair `heuristic_balanced` and `heuristic_cautious` so mirror
play has near-zero forced-end rate and realistic play/open rates.
2. Add `max_steps_rate` and game-length quantiles to every gate report.
3. Rerun the static ladder gates with the repaired anchors before initializing
a snapshot-pool league.
+334
View File
@@ -0,0 +1,334 @@
"""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()
+351 -3
View File
@@ -147,6 +147,16 @@ class EvalBatch(NamedTuple):
games: jax.Array games: jax.Array
class MatchBatch(NamedTuple):
score_diff: jax.Array
game_length: jax.Array
play_action_sum: jax.Array
learner_action_sum: jax.Array
opened_colors: jax.Array
positive_expeditions: jax.Array
max_steps: jax.Array
def load_config(path: str | Path | None = None, overrides: dict | None = None) -> JaxPPOConfig: def load_config(path: str | Path | None = None, overrides: dict | None = None) -> JaxPPOConfig:
data: dict = {} data: dict = {}
if path is not None: if path is not None:
@@ -179,7 +189,12 @@ def create_train_state(cfg: JaxPPOConfig, rng: jax.Array) -> TrainState:
return TrainState.create(apply_fn=model.apply, params=params, tx=tx) return TrainState.create(apply_fn=model.apply, params=params, tx=tx)
def train(cfg: JaxPPOConfig, *, resume: str | None = None) -> Path: def train(
cfg: JaxPPOConfig,
*,
resume: str | None = None,
opponent_policy=None,
) -> Path:
run_dir = _create_run_dir(cfg) run_dir = _create_run_dir(cfg)
_write_json(run_dir / "config.json", asdict(cfg)) _write_json(run_dir / "config.json", asdict(cfg))
metrics_path = run_dir / "metrics.jsonl" metrics_path = run_dir / "metrics.jsonl"
@@ -192,7 +207,8 @@ def train(cfg: JaxPPOConfig, *, resume: str | None = None) -> Path:
if resume: if resume:
state = restore_checkpoint(Path(resume), state) state = restore_checkpoint(Path(resume), state)
opponent_policy = policy_by_name(cfg.opponent.name) if opponent_policy is None:
opponent_policy = policy_by_name(cfg.opponent.name)
train_iteration = make_train_iteration(cfg, opponent_policy) train_iteration = make_train_iteration(cfg, opponent_policy)
start = time.perf_counter() start = time.perf_counter()
@@ -584,6 +600,246 @@ def make_eval_batch_fn(cfg: JaxPPOConfig, opponent_policy):
return eval_batch return eval_batch
def checkpoint_policy(cfg: JaxPPOConfig, checkpoint: str | Path):
state = create_train_state(cfg, jax.random.PRNGKey(0))
state = restore_checkpoint(Path(checkpoint), state)
return checkpoint_policy_from_params(cfg, state.params)
def checkpoint_policy_from_params(cfg: JaxPPOConfig, params):
model = ActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
def policy(state: State, player: jax.Array, rng: jax.Array) -> jax.Array:
del rng
obs = observation(state, player)[None, :]
legal = legal_action_mask(state)[None, :]
logits, _ = model.apply(params, obs)
return jnp.argmax(mask_logits(logits, legal), axis=-1)[0].astype(jnp.int32)
return policy
def train_against_checkpoint(
cfg: JaxPPOConfig,
opponent_cfg: JaxPPOConfig,
opponent_checkpoint: str | Path,
*,
resume: str | None = None,
) -> Path:
policy = checkpoint_policy(opponent_cfg, opponent_checkpoint)
return train(cfg, resume=resume, opponent_policy=policy)
def evaluate_checkpoint_match(
learner_cfg: JaxPPOConfig,
learner_checkpoint: str | Path,
opponent_cfg: JaxPPOConfig,
opponent_checkpoint: str | Path,
*,
games: int | None = None,
duplicate: bool | None = None,
output: str | Path | None = None,
) -> dict:
learner_policy = checkpoint_policy(learner_cfg, learner_checkpoint)
opponent_policy = checkpoint_policy(opponent_cfg, opponent_checkpoint)
games = games or learner_cfg.evaluation.games
duplicate = learner_cfg.evaluation.duplicate if duplicate is None else duplicate
return evaluate_policy_match(
learner_policy,
opponent_policy,
games=games,
duplicate=duplicate,
shuffle_bank_seed=learner_cfg.evaluation.shuffle_bank_seed,
batch_games=learner_cfg.evaluation.batch_games,
output=output,
)
def evaluate_checkpoint_vs_static(
learner_cfg: JaxPPOConfig,
learner_checkpoint: str | Path,
opponent_name: str,
*,
games: int | None = None,
duplicate: bool | None = None,
output: str | Path | None = None,
) -> dict:
learner_policy = checkpoint_policy(learner_cfg, learner_checkpoint)
opponent_policy = policy_by_name(opponent_name)
games = games or learner_cfg.evaluation.games
duplicate = learner_cfg.evaluation.duplicate if duplicate is None else duplicate
result = evaluate_policy_match(
learner_policy,
opponent_policy,
games=games,
duplicate=duplicate,
shuffle_bank_seed=learner_cfg.evaluation.shuffle_bank_seed,
batch_games=learner_cfg.evaluation.batch_games,
output=output,
)
result["opponent"] = opponent_name
return result
def evaluate_static_mirror(
policy_name: str,
*,
games: int,
duplicate: bool = True,
shuffle_bank_seed: int = 20260704,
batch_games: int = 8192,
output: str | Path | None = None,
) -> dict:
policy = policy_by_name(policy_name)
return evaluate_policy_match(
policy,
policy,
games=games,
duplicate=duplicate,
shuffle_bank_seed=shuffle_bank_seed,
batch_games=batch_games,
output=output,
)
def evaluate_policy_match(
learner_policy,
opponent_policy,
*,
games: int,
duplicate: bool,
shuffle_bank_seed: int,
batch_games: int,
output: str | Path | None = None,
) -> dict:
eval_batch = make_policy_match_batch_fn(learner_policy, opponent_policy)
diffs = []
lengths = []
play_actions = 0.0
learner_actions = 0.0
opened_sum = 0.0
positive_sum = 0.0
max_steps_sum = 0.0
for seat in [0, 1] if duplicate else [0]:
for start in range(0, games, batch_games):
count = min(batch_games, games - start)
orders = make_shuffle_bank(shuffle_bank_seed + start, count)
batch = eval_batch(jnp.asarray(orders, dtype=jnp.int8), seat)
batch = jax.tree_util.tree_map(lambda x: x.block_until_ready(), batch)
diffs.append(jax.device_get(batch["score_diff"]))
lengths.append(jax.device_get(batch["game_length"]))
play_actions += float(jnp.sum(batch["play_action_sum"]))
learner_actions += float(jnp.sum(batch["learner_action_sum"]))
opened_sum += float(jnp.sum(batch["opened_colors"]))
positive_sum += float(jnp.sum(batch["positive_expeditions"]))
max_steps_sum += float(jnp.sum(batch["max_steps"]))
import numpy as np
diff_arr = np.concatenate(diffs).astype(np.float64)
length_arr = np.concatenate(lengths).astype(np.float64)
total_games = int(diff_arr.shape[0])
wins = int(np.sum(diff_arr > 0))
losses = int(np.sum(diff_arr < 0))
ties = int(np.sum(diff_arr == 0))
wilson_low, wilson_high = wilson_interval(wins, total_games)
quantiles = np.quantile(diff_arr, [0.0, 0.05, 0.25, 0.5, 0.75, 0.95, 1.0])
hist_counts, hist_edges = np.histogram(
diff_arr,
bins=[-400, -300, -200, -100, -50, 0, 50, 100, 200, 300, 400, 600],
)
result = {
"duplicate": duplicate,
"games": total_games,
"wins": float(wins),
"losses": float(losses),
"ties": float(ties),
"win_rate": wins / total_games,
"wilson_low": wilson_low,
"wilson_high": wilson_high,
"mean_score_diff": float(np.mean(diff_arr)),
"score_diff_sum": float(np.sum(diff_arr)),
"score_diff_quantiles": {
"min": float(quantiles[0]),
"p05": float(quantiles[1]),
"p25": float(quantiles[2]),
"p50": float(quantiles[3]),
"p75": float(quantiles[4]),
"p95": float(quantiles[5]),
"max": float(quantiles[6]),
},
"score_diff_histogram": {
"edges": hist_edges.astype(float).tolist(),
"counts": hist_counts.astype(int).tolist(),
},
"mean_game_length": float(np.mean(length_arr)),
"game_length_quantiles": {
"min": float(np.min(length_arr)),
"p50": float(np.quantile(length_arr, 0.50)),
"p95": float(np.quantile(length_arr, 0.95)),
"max": float(np.max(length_arr)),
},
"max_steps_rate": max_steps_sum / total_games,
"play_action_rate": play_actions / max(learner_actions, 1.0),
"opened_colors_per_game": opened_sum / total_games,
"positive_expeditions_per_game": positive_sum / total_games,
}
if output:
_write_json(Path(output), result)
return result
def make_policy_match_batch_fn(learner_policy, opponent_policy):
@jax.jit
def eval_batch(deck_orders: jax.Array, learner_seat: int):
learner = jnp.asarray(learner_seat, dtype=jnp.int32)
opponent = 1 - learner
env_state = jax.vmap(reset_from_order)(deck_orders)
keys = jax.random.split(jax.random.PRNGKey(0), deck_orders.shape[0])
def body(carry, _):
env, key = carry
key, learner_key, opponent_key = jax.random.split(key, 3)
learner_keys = jax.random.split(learner_key, deck_orders.shape[0])
opponent_keys = jax.random.split(opponent_key, deck_orders.shape[0])
learner_actions = jax.vmap(learner_policy, in_axes=(0, None, 0))(
env, learner, learner_keys
)
opponent_actions = jax.vmap(opponent_policy, in_axes=(0, None, 0))(
env, opponent, opponent_keys
)
learner_turn = env.to_move.astype(jnp.int32) == learner
active = ~env.done
actions = jnp.where(learner_turn, learner_actions, opponent_actions)
next_env, _, _ = jax.vmap(step, in_axes=(0, 0))(env, actions)
place_type = (actions % 12) // 6
return (next_env, key), (
(place_type == PLAY) & learner_turn & active,
learner_turn & active,
)
(final_env, _), (play_actions, learner_actions) = jax.lax.scan(
body, (env_state, keys[0]), xs=None, length=MAX_STEPS
)
scores = jax.vmap(board_score)(final_env)
diff = scores[:, learner] - scores[:, opponent]
color_scores = jax.vmap(color_scores_for_player, in_axes=(0, None))(final_env, learner)
opened = jnp.sum(final_env.col_len[:, learner, :] > 0, axis=-1)
positive = jnp.sum(color_scores > 0, axis=-1)
return MatchBatch(
score_diff=diff,
game_length=final_env.step_count,
play_action_sum=jnp.sum(play_actions, axis=0),
learner_action_sum=jnp.sum(learner_actions, axis=0),
opened_colors=opened,
positive_expeditions=positive,
max_steps=final_env.step_count >= MAX_STEPS,
)._asdict()
return eval_batch
def rollout_metrics( def rollout_metrics(
transitions: Transition, final_env: State, learner: jax.Array transitions: Transition, final_env: State, learner: jax.Array
) -> dict[str, jax.Array]: ) -> dict[str, jax.Array]:
@@ -716,6 +972,13 @@ def cli_main(argv: list[str] | None = None) -> None:
train_parser.add_argument("--resume") train_parser.add_argument("--resume")
train_parser.add_argument("--set", action="append", default=[]) train_parser.add_argument("--set", action="append", default=[])
exploiter_parser = sub.add_parser("train-exploiter")
exploiter_parser.add_argument("--config", required=True)
exploiter_parser.add_argument("--opponent-config", required=True)
exploiter_parser.add_argument("--opponent-checkpoint", required=True)
exploiter_parser.add_argument("--resume")
exploiter_parser.add_argument("--set", action="append", default=[])
eval_parser = sub.add_parser("eval") eval_parser = sub.add_parser("eval")
eval_parser.add_argument("--config", required=True) eval_parser.add_argument("--config", required=True)
eval_parser.add_argument("--checkpoint", required=True) eval_parser.add_argument("--checkpoint", required=True)
@@ -725,17 +988,61 @@ def cli_main(argv: list[str] | None = None) -> None:
eval_parser.add_argument("--output") eval_parser.add_argument("--output")
eval_parser.add_argument("--set", action="append", default=[]) eval_parser.add_argument("--set", action="append", default=[])
checkpoint_eval_parser = sub.add_parser("eval-checkpoint")
checkpoint_eval_parser.add_argument("--config", required=True)
checkpoint_eval_parser.add_argument("--checkpoint", required=True)
checkpoint_eval_parser.add_argument("--opponent-config", required=True)
checkpoint_eval_parser.add_argument("--opponent-checkpoint", required=True)
checkpoint_eval_parser.add_argument("--games", type=int)
checkpoint_eval_parser.add_argument("--duplicate", action="store_true")
checkpoint_eval_parser.add_argument("--output")
checkpoint_eval_parser.add_argument("--set", action="append", default=[])
checkpoint_static_eval_parser = sub.add_parser("eval-checkpoint-static")
checkpoint_static_eval_parser.add_argument("--config", required=True)
checkpoint_static_eval_parser.add_argument("--checkpoint", required=True)
checkpoint_static_eval_parser.add_argument("--opponent", required=True)
checkpoint_static_eval_parser.add_argument("--games", type=int)
checkpoint_static_eval_parser.add_argument("--duplicate", action="store_true")
checkpoint_static_eval_parser.add_argument("--output")
checkpoint_static_eval_parser.add_argument("--set", action="append", default=[])
static_mirror_parser = sub.add_parser("eval-static-mirror")
static_mirror_parser.add_argument("--policy", required=True)
static_mirror_parser.add_argument("--games", type=int, default=10_000)
static_mirror_parser.add_argument("--batch-games", type=int, default=8192)
static_mirror_parser.add_argument("--shuffle-bank-seed", type=int, default=20260704)
static_mirror_parser.add_argument("--duplicate", action="store_true")
static_mirror_parser.add_argument("--output")
args = parser.parse_args(argv) args = parser.parse_args(argv)
cfg = load_config(args.config, overrides=parse_overrides(args.set)) cfg = (
load_config(args.config, overrides=parse_overrides(args.set))
if hasattr(args, "config")
else None
)
if args.command == "rollout-smoke": if args.command == "rollout-smoke":
assert cfg is not None
row = random_rollout(cfg) row = random_rollout(cfg)
print(json.dumps(row, indent=2, sort_keys=True)) print(json.dumps(row, indent=2, sort_keys=True))
if args.output: if args.output:
_write_json(Path(args.output), row) _write_json(Path(args.output), row)
elif args.command == "train": elif args.command == "train":
assert cfg is not None
run_dir = train(cfg, resume=args.resume) run_dir = train(cfg, resume=args.resume)
print(run_dir) print(run_dir)
elif args.command == "train-exploiter":
assert cfg is not None
opponent_cfg = load_config(args.opponent_config)
run_dir = train_against_checkpoint(
cfg,
opponent_cfg,
args.opponent_checkpoint,
resume=args.resume,
)
print(run_dir)
elif args.command == "eval": elif args.command == "eval":
assert cfg is not None
result = evaluate( result = evaluate(
cfg, cfg,
args.checkpoint, args.checkpoint,
@@ -745,6 +1052,40 @@ def cli_main(argv: list[str] | None = None) -> None:
output=args.output, output=args.output,
) )
print(json.dumps(result, indent=2, sort_keys=True)) print(json.dumps(result, indent=2, sort_keys=True))
elif args.command == "eval-checkpoint":
assert cfg is not None
opponent_cfg = load_config(args.opponent_config)
result = evaluate_checkpoint_match(
cfg,
args.checkpoint,
opponent_cfg,
args.opponent_checkpoint,
games=args.games,
duplicate=args.duplicate,
output=args.output,
)
print(json.dumps(result, indent=2, sort_keys=True))
elif args.command == "eval-checkpoint-static":
assert cfg is not None
result = evaluate_checkpoint_vs_static(
cfg,
args.checkpoint,
args.opponent,
games=args.games,
duplicate=args.duplicate,
output=args.output,
)
print(json.dumps(result, indent=2, sort_keys=True))
elif args.command == "eval-static-mirror":
result = evaluate_static_mirror(
args.policy,
games=args.games,
duplicate=args.duplicate,
shuffle_bank_seed=args.shuffle_bank_seed,
batch_games=args.batch_games,
output=args.output,
)
print(json.dumps(result, indent=2, sort_keys=True))
def _create_run_dir(cfg: JaxPPOConfig) -> Path: def _create_run_dir(cfg: JaxPPOConfig) -> Path:
@@ -817,10 +1158,17 @@ __all__ = [
"ActorCritic", "ActorCritic",
"JaxPPOConfig", "JaxPPOConfig",
"TrainState", "TrainState",
"checkpoint_policy",
"checkpoint_policy_from_params",
"cli_main", "cli_main",
"create_train_state", "create_train_state",
"evaluate", "evaluate",
"evaluate_checkpoint_match",
"evaluate_checkpoint_vs_static",
"evaluate_policy_match",
"evaluate_static_mirror",
"load_config", "load_config",
"random_rollout", "random_rollout",
"train", "train",
"train_against_checkpoint",
] ]
+161
View File
@@ -1,10 +1,14 @@
from __future__ import annotations from __future__ import annotations
from pathlib import Path
import jax import jax
import jax.numpy as jnp import jax.numpy as jnp
import numpy as np import numpy as np
import pytest
from lost_cities_jax import legal_action_mask, reset from lost_cities_jax import legal_action_mask, reset
from lost_cities_jax.engine import current_hand_sorted, decode_action, reset_from_order
from lost_cities_jax.opponents import ( from lost_cities_jax.opponents import (
discard_only_action, discard_only_action,
heuristic_balanced_action, heuristic_balanced_action,
@@ -18,12 +22,29 @@ from lost_cities_jax.ppo import (
RunConfig, RunConfig,
create_train_state, create_train_state,
evaluate, evaluate,
evaluate_checkpoint_match,
evaluate_static_mirror,
load_config, load_config,
make_train_iteration, make_train_iteration,
policy_by_name, policy_by_name,
random_rollout, random_rollout,
train, train,
) )
from lost_cities_jax.types import (
CARDS_PER_COLOR,
HAND_SIZE,
LOC_DECK,
LOC_P0_BOARD,
LOC_P0_HAND,
N_CARDS,
N_COLORS,
PLAY,
)
GATE3_CHECKPOINT = (
"/mnt/2tbhdd/coolrl-lost-cities-artifacts/jax-ppo-static-opponents/"
"2026-07-04_230150_jax-ppo-cautious/latest"
)
def tiny_config(tmp_path) -> JaxPPOConfig: def tiny_config(tmp_path) -> JaxPPOConfig:
@@ -49,6 +70,84 @@ def test_static_opponents_return_legal_actions():
assert mask[action] assert mask[action]
def test_cautious_discards_weak_unopened_hand():
low_cards = [_rank_card(0, rank) for rank in range(2, 7)]
state = _manual_state(p0_hand=low_cards, to_move=0)
action = int(heuristic_cautious_action(state, jnp.int32(0), jax.random.PRNGKey(1)))
_, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
assert place_type != PLAY
def test_balanced_opens_strong_unopened_rank():
strong = _rank_card(0, 10)
state = _manual_state(p0_hand=[strong, _rank_card(1, 2), _rank_card(2, 3)], to_move=0)
action = int(heuristic_balanced_action(state, jnp.int32(0), jax.random.PRNGKey(1)))
hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
hand = [int(x) for x in current_hand_sorted(state, 0)]
assert place_type == PLAY
assert hand[hand_slot] == strong
def test_cautious_uses_own_board_when_playing_as_p1():
existing = _rank_card(0, 5)
playable = _rank_card(0, 6)
state = _manual_state(
p1_hand=[
playable,
_rank_card(1, 2),
_rank_card(1, 3),
_rank_card(2, 2),
_rank_card(2, 3),
_rank_card(3, 2),
_rank_card(3, 3),
_rank_card(4, 2),
],
p1_board=[existing],
to_move=1,
)
action = int(heuristic_cautious_action(state, jnp.int32(1), jax.random.PRNGKey(1)))
hand_slot, place_type, _ = [int(x) for x in decode_action(jnp.asarray(action))]
hand = [int(x) for x in current_hand_sorted(state, 1)]
assert place_type == PLAY
assert hand[hand_slot] == playable
@pytest.mark.parametrize(
"policy_name", ["discard_only", "heuristic_balanced", "heuristic_cautious"]
)
def test_static_policy_duplicate_mirror_score_diff_is_zero(policy_name):
result = evaluate_static_mirror(
policy_name,
games=16,
duplicate=True,
shuffle_bank_seed=20260704,
batch_games=16,
)
assert result["games"] == 32
assert result["mean_score_diff"] == 0.0
assert result["wins"] == result["losses"]
def test_gate3_checkpoint_duplicate_self_mirror_score_diff_is_zero():
checkpoint = Path(GATE3_CHECKPOINT)
if not checkpoint.exists():
pytest.skip(f"gate-3 checkpoint not available: {checkpoint}")
if not any(device.platform == "gpu" for device in jax.local_devices()):
pytest.skip("gate-3 checkpoint was saved with CUDA sharding; run under CUDA JAX")
cfg = load_config("configs/jax_ppo/cautious.yaml")
result = evaluate_checkpoint_match(
cfg,
checkpoint,
cfg,
checkpoint,
games=4,
duplicate=True,
)
assert result["games"] == 8
assert result["mean_score_diff"] == 0.0
assert result["wins"] == result["losses"]
def test_random_rollout_smoke(tmp_path): def test_random_rollout_smoke(tmp_path):
row = random_rollout(tiny_config(tmp_path)) row = random_rollout(tiny_config(tmp_path))
assert row["env_steps"] == 8 * 16 assert row["env_steps"] == 8 * 16
@@ -84,3 +183,65 @@ def test_load_config_file():
cfg = load_config("configs/jax_ppo/smoke.yaml") cfg = load_config("configs/jax_ppo/smoke.yaml")
assert cfg.opponent.name == "discard_only" assert cfg.opponent.name == "discard_only"
assert cfg.ppo.batch_games == 64 assert cfg.ppo.batch_games == 64
def _rank_card(color: int, rank: int) -> int:
return color * CARDS_PER_COLOR + rank + 1
def _manual_state(
*,
p0_hand: list[int] | None = None,
p1_hand: list[int] | None = None,
p0_board: list[int] | None = None,
p1_board: list[int] | None = None,
to_move: int = 0,
):
p0_hand = _fill_hand(p0_hand or [], set((p0_board or []) + (p1_board or [])))
p1_hand = _fill_hand(p1_hand or [], set(p0_hand + (p0_board or []) + (p1_board or [])))
rest = [card for card in range(N_CARDS) if card not in set(p0_hand + p1_hand)]
deck_order = jnp.asarray(p0_hand + p1_hand + rest, dtype=jnp.int8)
state = reset_from_order(deck_order)
card_loc = jnp.full((N_CARDS,), LOC_DECK, dtype=jnp.int8)
card_loc = card_loc.at[jnp.asarray(p0_hand)].set(LOC_P0_HAND)
card_loc = card_loc.at[jnp.asarray(p1_hand)].set(LOC_P0_HAND + 1)
if p0_board:
card_loc = card_loc.at[jnp.asarray(p0_board)].set(LOC_P0_BOARD)
if p1_board:
card_loc = card_loc.at[jnp.asarray(p1_board)].set(LOC_P0_BOARD + 1)
col_top, col_hs, col_len = _columns(p0_board or [], p1_board or [])
return state._replace(
card_loc=card_loc,
col_top=col_top,
col_hs=col_hs,
col_len=col_len,
to_move=jnp.asarray(to_move, dtype=jnp.int8),
)
def _fill_hand(cards: list[int], reserved: set[int]) -> list[int]:
result = list(cards)
blocked = set(result) | reserved
for card in range(N_CARDS):
if len(result) == HAND_SIZE:
return sorted(result)
if card not in blocked:
result.append(card)
blocked.add(card)
raise AssertionError("could not fill hand")
def _columns(p0_board: list[int], p1_board: list[int]):
col_top = np.zeros((2, N_COLORS), dtype=np.int8)
col_hs = np.zeros((2, N_COLORS), dtype=np.int8)
col_len = np.zeros((2, N_COLORS), dtype=np.int8)
for player, cards in enumerate([p0_board, p1_board]):
for card in cards:
color = card // CARDS_PER_COLOR
slot = card % CARDS_PER_COLOR
col_len[player, color] += 1
if slot < 3:
col_hs[player, color] += 1
else:
col_top[player, color] = max(col_top[player, color], slot - 1)
return jnp.asarray(col_top), jnp.asarray(col_hs), jnp.asarray(col_len)