Train on matches: match observation, asymmetric critic, both seats, match reward
Phases 2-4 land together because they all rewrite the same rollout, and doing them in sequence would mean writing it three times. Observation adds the four things a match policy cannot play without: carry as a scalar *and* a binned one-hot (round three is a threshold problem -- "win by 41 or lose" plays nothing like "win by 39" -- and the old score_diff divided by MAX_ABS_SCORE=780, squashing a decisive 50-point lead to 0.06); the round index; whose turn it is, which the single-round observation never carried even though the critic is trained on opponent-turn states; and the deck clock, since a round ends on the last deck draw and players bend that parity by drawing from discard piles. Live points per colour are split by hand / discard pile / unseen, because a discard pile is public and recoverable. The critic is asymmetric: it gets the opponent's hand and the deck in order, on a separate trunk so none of it can reach the logits. A test pins that down -- perturbing the privileged input leaves the policy logits bit-identical while moving the value. Deal luck is what makes a match-terminal reward hard to learn from, and a state-value baseline may condition on anything action-independent. Both seats now train. Self-play ran one network on both sides and stop_gradiented the opponent, throwing away half of every game; each ply now emits a transition per seat, folded into the batch so each seat keeps an independent GAE chain. Reward is the match: rounds one and two only bank into carry, and round three pays tanh(total / terminal_scale). Potential shaping on the running total covers the early sparsity and anneals out. match_eval adds the two measurements the plan turns on: duplicate match play (same deals and coins from both seats -- self-play scores exactly 0.500 with zero mean lead, so the mirroring cancels deal luck exactly) and the carry probe. On an untrained net the probe is flat: 4.97 expeditions opened at a 60-point deficit and at a 60-point lead alike. Breaking that flat line is success criterion 1. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01XBQKgvBbxbheiTF1AVy1Sh
This commit is contained in:
Claude Opus 4.8
@@ -0,0 +1,195 @@
|
||||
"""Match-level evaluation, and the probe that says whether carry changed anything.
|
||||
|
||||
Single-deal win rate cannot tell you if the three-round rework worked. Two things
|
||||
can:
|
||||
|
||||
- **Duplicate match play.** The same three deals and the same coin flips, played
|
||||
from both seats. Deal luck cancels, so what is left is the policy difference.
|
||||
- **The carry probe.** Drop the policy into round three holding a fixed deficit or
|
||||
lead and watch what it does. A carry-blind policy plays a 60-point deficit
|
||||
exactly like a 60-point lead. If the numbers below do not move with carry, the
|
||||
rework bought nothing, no matter what the win rate says.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
|
||||
from lost_cities_jax.engine import reset_from_order
|
||||
from lost_cities_jax.match import (
|
||||
N_ROUNDS,
|
||||
MatchState,
|
||||
match_legal_action_mask,
|
||||
match_reset_from,
|
||||
match_score,
|
||||
match_step,
|
||||
)
|
||||
from lost_cities_jax.match_obs import match_critic_observation, match_observation
|
||||
from lost_cities_jax.match_ppo import MatchActorCritic, match_lead
|
||||
from lost_cities_jax.ppo import JaxPPOConfig, mask_logits
|
||||
from lost_cities_jax.types import DRAW_DECK, MAX_STEPS, N_CARDS
|
||||
|
||||
MATCH_SCAN_STEPS = N_ROUNDS * MAX_STEPS
|
||||
|
||||
|
||||
def match_bank(seed: int, matches: int) -> tuple[jnp.ndarray, jnp.ndarray]:
|
||||
"""Deterministic deals and coin flips, so runs are comparable."""
|
||||
rng = np.random.default_rng(seed)
|
||||
decks = np.stack(
|
||||
[np.stack([rng.permutation(N_CARDS) for _ in range(N_ROUNDS)]) for _ in range(matches)]
|
||||
)
|
||||
coins = rng.integers(0, 2, size=(matches, N_ROUNDS))
|
||||
return jnp.asarray(decks, dtype=jnp.int8), jnp.asarray(coins, dtype=jnp.int8)
|
||||
|
||||
|
||||
def _wilson(wins: float, games: float) -> tuple[float, float]:
|
||||
if games == 0:
|
||||
return 0.0, 0.0
|
||||
z = 1.96
|
||||
p = wins / games
|
||||
denom = 1 + z * z / games
|
||||
centre = p + z * z / (2 * games)
|
||||
margin = z * math.sqrt(p * (1 - p) / games + z * z / (4 * games * games))
|
||||
return (centre - margin) / denom, (centre + margin) / denom
|
||||
|
||||
|
||||
def make_match_runner(cfg: JaxPPOConfig, opponent_policy=None):
|
||||
"""Play whole matches to the end. ``opponent_policy`` is a single-round policy."""
|
||||
|
||||
model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
|
||||
|
||||
@jax.jit
|
||||
def run(params, env: MatchState, learner_seat: jax.Array, rng: jax.Array):
|
||||
def body(carry, _):
|
||||
env, key, deck_draws, plies = carry
|
||||
key, opp_key = jax.random.split(key)
|
||||
|
||||
to_move = env.round.to_move.astype(jnp.int32)
|
||||
obs = jax.vmap(match_observation)(env, to_move)
|
||||
critic = jax.vmap(match_critic_observation)(env, to_move)
|
||||
mask = jax.vmap(match_legal_action_mask)(env)
|
||||
logits, _ = model.apply(params, obs, critic)
|
||||
learner_action = jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)
|
||||
|
||||
if opponent_policy is None:
|
||||
action = learner_action
|
||||
else:
|
||||
opp_keys = jax.random.split(opp_key, env.done.shape[0])
|
||||
opp_action = jax.vmap(opponent_policy, in_axes=(0, 0, 0))(
|
||||
env.round, to_move, opp_keys
|
||||
)
|
||||
action = jnp.where(to_move == learner_seat, learner_action, opp_action)
|
||||
|
||||
live = ~env.done
|
||||
is_deck = (action % 6) == DRAW_DECK
|
||||
deck_draws = deck_draws + (is_deck & live).astype(jnp.int32)
|
||||
plies = plies + live.astype(jnp.int32)
|
||||
|
||||
env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
|
||||
return (env, key, deck_draws, plies), None
|
||||
|
||||
n = env.done.shape[0]
|
||||
zeros = jnp.zeros((n,), dtype=jnp.int32)
|
||||
(env, _, deck_draws, plies), _ = jax.lax.scan(
|
||||
body, (env, rng, zeros, zeros), xs=None, length=MATCH_SCAN_STEPS
|
||||
)
|
||||
return env, deck_draws, plies
|
||||
|
||||
return run
|
||||
|
||||
|
||||
def match_evaluate(cfg: JaxPPOConfig, params, *, matches: int = 2000, seed: int = 20260715) -> dict:
|
||||
"""Duplicate match play: every deal seen from both seats."""
|
||||
|
||||
decks, coins = match_bank(seed, matches)
|
||||
runner = make_match_runner(cfg)
|
||||
|
||||
results = []
|
||||
for seat in (0, 1):
|
||||
env = jax.vmap(match_reset_from)(decks, coins)
|
||||
seats = jnp.full((matches,), seat, dtype=jnp.int32)
|
||||
final, deck_draws, plies = runner(params, env, seats, jax.random.PRNGKey(0))
|
||||
totals = np.asarray(jax.vmap(match_score)(final))
|
||||
lead = totals[:, seat] - totals[:, 1 - seat]
|
||||
results.append(
|
||||
{
|
||||
"lead": lead,
|
||||
"plies": np.asarray(plies),
|
||||
"deck_share": np.asarray(deck_draws) / np.maximum(np.asarray(plies), 1),
|
||||
"done": np.asarray(final.done),
|
||||
}
|
||||
)
|
||||
|
||||
lead = np.concatenate([r["lead"] for r in results])
|
||||
games = float(lead.size)
|
||||
wins = float((lead > 0).sum())
|
||||
ties = float((lead == 0).sum())
|
||||
low, high = _wilson(wins, games)
|
||||
return {
|
||||
"matches": games,
|
||||
"wins": wins,
|
||||
"ties": ties,
|
||||
"losses": games - wins - ties,
|
||||
"match_win_rate": wins / games,
|
||||
"wilson_low": low,
|
||||
"wilson_high": high,
|
||||
"mean_total_lead": float(lead.mean()),
|
||||
"mean_plies": float(np.concatenate([r["plies"] for r in results]).mean()),
|
||||
"unfinished_rate": float(
|
||||
1.0 - np.concatenate([r["done"] for r in results]).astype(np.float32).mean()
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def carry_probe(
|
||||
cfg: JaxPPOConfig,
|
||||
params,
|
||||
*,
|
||||
carry_levels=(-60, -25, -1, 1, 25, 60),
|
||||
matches: int = 512,
|
||||
seed: int = 20260716,
|
||||
) -> list[dict]:
|
||||
"""Start the policy in round three holding ``carry`` and see what changes.
|
||||
|
||||
A policy that reads carry should open more expeditions and lean on wagers when
|
||||
it is behind (it has to swing), and shut down when it is ahead. Flat rows mean
|
||||
the policy is ignoring the one thing the three-round rework added.
|
||||
"""
|
||||
|
||||
decks, coins = match_bank(seed, matches)
|
||||
runner = make_match_runner(cfg)
|
||||
rows = []
|
||||
|
||||
for level in carry_levels:
|
||||
base = jax.vmap(match_reset_from)(decks, coins)
|
||||
carry = jnp.tile(jnp.asarray([level, 0], dtype=jnp.int32), (matches, 1))
|
||||
round_idx = jnp.full((matches,), N_ROUNDS - 1, dtype=jnp.int32)
|
||||
# Whoever is ahead leads the round, exactly as the match layer would.
|
||||
first = jnp.int8(0 if level > 0 else 1)
|
||||
rounds = jax.vmap(reset_from_order, in_axes=(0, None))(decks[:, N_ROUNDS - 1], first)
|
||||
|
||||
env = base._replace(round=rounds, carry=carry, round_idx=round_idx)
|
||||
final, deck_draws, plies = runner(
|
||||
params, env, jnp.zeros((matches,), jnp.int32), jax.random.PRNGKey(1)
|
||||
)
|
||||
|
||||
lead = np.asarray(jax.vmap(match_lead)(final))
|
||||
opened = np.asarray(jnp.sum(final.round.col_len[:, 0, :] > 0, axis=-1))
|
||||
wagers = np.asarray(jnp.sum(final.round.col_hs[:, 0, :], axis=-1))
|
||||
deck_share = np.asarray(deck_draws) / np.maximum(np.asarray(plies), 1)
|
||||
|
||||
rows.append(
|
||||
{
|
||||
"carry": level,
|
||||
"win_rate": float((lead > 0).mean()),
|
||||
"opened_colors": float(opened.mean()),
|
||||
"wagers_played": float(wagers.mean()),
|
||||
"deck_race_rate": float(deck_share.mean()),
|
||||
"mean_plies": float(np.asarray(plies).mean()),
|
||||
}
|
||||
)
|
||||
return rows
|
||||
@@ -0,0 +1,198 @@
|
||||
"""Player-view observation for a three-round match.
|
||||
|
||||
Everything the single-round observation had, plus the four things a match policy
|
||||
cannot play without:
|
||||
|
||||
- **carry**, as a scalar *and* a one-hot over bins. Round three is a threshold
|
||||
problem -- "win by 41 or lose" plays nothing like "win by 39" -- and a lone
|
||||
scalar makes the network learn a cliff from a smooth input. (The old
|
||||
``score_diff`` divided by ``MAX_ABS_SCORE`` = 780, which squashed a decisive
|
||||
50-point lead to 0.06. That scaling is fixed here too.)
|
||||
- **which round it is**, and how many are left.
|
||||
- **whose turn it is.** The single-round obs never had this, yet the critic is
|
||||
trained on opponent-turn states as well, leaving it to recover turn parity
|
||||
from ``step_count / 400``.
|
||||
- **the deck clock.** A round ends when the last deck card is drawn, so turn
|
||||
parity decides who gets the final placement -- and players bend that parity by
|
||||
drawing from discard piles instead.
|
||||
|
||||
Live points per colour are split by where the card is (hand / discard pile /
|
||||
unseen) because a discard pile is public and recoverable: lumping it in with the
|
||||
unseen cards hides the high cards sitting in plain view.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
|
||||
from lost_cities_jax.engine import board_score
|
||||
from lost_cities_jax.match import N_ROUNDS, MatchState
|
||||
from lost_cities_jax.obs import observation
|
||||
from lost_cities_jax.types import (
|
||||
CARDS_PER_COLOR,
|
||||
LOC_DECK,
|
||||
LOC_DISCARD,
|
||||
LOC_P0_HAND,
|
||||
N_CARDS,
|
||||
N_COLORS,
|
||||
OBS_DIM,
|
||||
)
|
||||
|
||||
# Bin edges for the carry one-hot. Packed tightly around zero: that is where the
|
||||
# round-three decision actually flips.
|
||||
CARRY_BIN_EDGES = jnp.asarray([-60.0, -30.0, -12.0, -1.0, 1.0, 12.0, 30.0, 60.0])
|
||||
N_CARRY_BINS = CARRY_BIN_EDGES.shape[0] + 1
|
||||
|
||||
CARRY_SCALE = 75.0
|
||||
"""Typical round margin, not the theoretical 780 maximum."""
|
||||
|
||||
MAX_COLOR_POINTS = 54.0
|
||||
"""2+3+...+10, the most a single expedition can be worth before multipliers."""
|
||||
|
||||
N_MATCH_SCALARS = (
|
||||
1 # carry, scaled
|
||||
+ N_CARRY_BINS # carry, binned
|
||||
+ N_ROUNDS # which round
|
||||
+ 1 # rounds left
|
||||
+ 1 # is it my turn
|
||||
+ 1 # did I lead this round
|
||||
+ 1 # do I take the last deck card at the current parity
|
||||
+ 2 * N_COLORS * 3 # live points: both players x colour x {hand, discard, unseen}
|
||||
)
|
||||
MATCH_OBS_DIM = OBS_DIM + N_MATCH_SCALARS
|
||||
|
||||
|
||||
def _live_points(state, viewer: jax.Array, subject: jax.Array) -> jax.Array:
|
||||
"""Points still reachable for ``subject``, split by where the card sits.
|
||||
|
||||
Returns ``float32[N_COLORS, 3]`` over {in hand, in a discard pile, unseen}.
|
||||
Seen through ``viewer``'s eyes: a card in the opponent's hand only counts as
|
||||
"in hand" if it is public, otherwise it is unseen.
|
||||
"""
|
||||
|
||||
ids = jnp.arange(N_CARDS, dtype=jnp.int32)
|
||||
colors = ids // CARDS_PER_COLOR
|
||||
slots = ids % CARDS_PER_COLOR
|
||||
ranks = jnp.where(slots >= 3, slots - 1, 0).astype(jnp.float32)
|
||||
|
||||
loc = state.card_loc.astype(jnp.int32)
|
||||
subject_hand = loc == (LOC_P0_HAND + subject)
|
||||
is_mine = subject == viewer
|
||||
|
||||
# An ascending column can only take cards above its current top.
|
||||
above_top = ranks > state.col_top[subject][colors].astype(jnp.float32)
|
||||
|
||||
in_hand = subject_hand & (is_mine | state.hand_public)
|
||||
in_discard = loc == LOC_DISCARD
|
||||
unseen = (loc == LOC_DECK) | (subject_hand & ~is_mine & ~state.hand_public)
|
||||
# Cards in the *other* player's hidden hand are unseen to the viewer too, but
|
||||
# they are not reachable by the subject, so they are deliberately excluded.
|
||||
|
||||
def by_color(mask: jax.Array) -> jax.Array:
|
||||
weighted = jnp.where(mask & above_top, ranks, 0.0)
|
||||
return jax.ops.segment_sum(weighted, colors, num_segments=N_COLORS)
|
||||
|
||||
stacked = jnp.stack([by_color(in_hand), by_color(in_discard), by_color(unseen)], axis=1)
|
||||
return stacked / MAX_COLOR_POINTS
|
||||
|
||||
|
||||
def match_observation(state: MatchState, player: jax.Array) -> jax.Array:
|
||||
"""Return ``float32[MATCH_OBS_DIM]`` from ``player``'s seat."""
|
||||
|
||||
player = jnp.asarray(player, dtype=jnp.int32)
|
||||
opponent = 1 - player
|
||||
round_state = state.round
|
||||
|
||||
base = observation(round_state, player)
|
||||
|
||||
totals = state.carry + board_score(round_state).astype(jnp.int32)
|
||||
lead = (totals[player] - totals[opponent]).astype(jnp.float32)
|
||||
carry_scaled = jnp.clip(lead / CARRY_SCALE, -2.0, 2.0).reshape((1,))
|
||||
carry_bins = jax.nn.one_hot(jnp.digitize(lead, CARRY_BIN_EDGES), N_CARRY_BINS)
|
||||
|
||||
round_onehot = jax.nn.one_hot(state.round_idx, N_ROUNDS)
|
||||
rounds_left = ((N_ROUNDS - 1 - state.round_idx).astype(jnp.float32) / (N_ROUNDS - 1)).reshape(
|
||||
(1,)
|
||||
)
|
||||
|
||||
to_move = round_state.to_move.astype(jnp.int32)
|
||||
my_turn = (to_move == player).astype(jnp.float32).reshape((1,))
|
||||
# to_move flips every ply, so the round's opener is recoverable from parity.
|
||||
parity = round_state.step_count.astype(jnp.int32) & 1
|
||||
round_opener = to_move ^ parity
|
||||
i_opened = (round_opener == player).astype(jnp.float32).reshape((1,))
|
||||
|
||||
# If both players drew from the deck from here, the last deck card falls to
|
||||
# whoever is on move after `remaining - 1` more plies.
|
||||
remaining = jnp.maximum(N_CARDS - round_state.draw_ptr, 1).astype(jnp.int32)
|
||||
last_drawer = to_move ^ ((remaining - 1) & 1)
|
||||
i_take_last = (last_drawer == player).astype(jnp.float32).reshape((1,))
|
||||
|
||||
live_me = _live_points(round_state, player, player).reshape((-1,))
|
||||
live_opp = _live_points(round_state, player, opponent).reshape((-1,))
|
||||
|
||||
extra = jnp.concatenate(
|
||||
[
|
||||
carry_scaled,
|
||||
carry_bins,
|
||||
round_onehot,
|
||||
rounds_left,
|
||||
my_turn,
|
||||
i_opened,
|
||||
i_take_last,
|
||||
live_me,
|
||||
live_opp,
|
||||
],
|
||||
axis=0,
|
||||
).astype(jnp.float32)
|
||||
|
||||
return jnp.concatenate([base, extra], axis=0).reshape((MATCH_OBS_DIM,))
|
||||
|
||||
|
||||
CRITIC_EXTRA_CHANNELS = 3
|
||||
MATCH_CRITIC_OBS_DIM = MATCH_OBS_DIM + N_CARDS * CRITIC_EXTRA_CHANNELS
|
||||
|
||||
|
||||
def match_critic_observation(state: MatchState, player: jax.Array) -> jax.Array:
|
||||
"""The same view, plus the hidden state. Training only -- never played from.
|
||||
|
||||
A value baseline may condition on anything the action does not depend on
|
||||
without biasing the policy gradient, and the deal is exactly the noise that
|
||||
swamps a match-terminal reward. So the critic gets the opponent's hand and
|
||||
the deck *in order* -- it knows what is coming -- while the actor keeps the
|
||||
masked view. The two run on separate trunks so none of this can leak into
|
||||
the logits.
|
||||
"""
|
||||
|
||||
player = jnp.asarray(player, dtype=jnp.int32)
|
||||
opponent = 1 - player
|
||||
round_state = state.round
|
||||
loc = round_state.card_loc.astype(jnp.int32)
|
||||
|
||||
opp_hand = (loc == (LOC_P0_HAND + opponent)).astype(jnp.float32)
|
||||
in_deck = (loc == LOC_DECK).astype(jnp.float32)
|
||||
|
||||
# Position of each card in the undrawn deck, so the critic can see the order
|
||||
# rather than just the multiset.
|
||||
deck_order = round_state.deck_order.astype(jnp.int32)
|
||||
position = (
|
||||
jnp.zeros((N_CARDS,), dtype=jnp.int32)
|
||||
.at[deck_order]
|
||||
.set(jnp.arange(N_CARDS, dtype=jnp.int32))
|
||||
)
|
||||
depth = (position - round_state.draw_ptr).astype(jnp.float32) / N_CARDS
|
||||
deck_depth = jnp.where(in_deck > 0, depth, 0.0)
|
||||
|
||||
hidden = jnp.stack([opp_hand, in_deck, deck_depth], axis=1).reshape((-1,))
|
||||
base = match_observation(state, player)
|
||||
return jnp.concatenate([base, hidden], axis=0).reshape((MATCH_CRITIC_OBS_DIM,))
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MATCH_CRITIC_OBS_DIM",
|
||||
"MATCH_OBS_DIM",
|
||||
"N_CARRY_BINS",
|
||||
"match_critic_observation",
|
||||
"match_observation",
|
||||
]
|
||||
@@ -0,0 +1,407 @@
|
||||
"""PPO over three-round matches.
|
||||
|
||||
Three things differ from the single-round trainer beyond the env swap:
|
||||
|
||||
**Asymmetric actor-critic.** The critic never plays, so it is fed the hidden
|
||||
state (opponent's hand, deck in order) while the actor keeps the masked view. A
|
||||
state-value baseline may condition on anything action-independent without biasing
|
||||
the policy gradient, and deal luck is exactly what makes a match-terminal reward
|
||||
hard to learn from. The trunks are separate so the privileged features cannot
|
||||
reach the logits.
|
||||
|
||||
**Both seats train.** In self-play one network picks both players' moves, but the
|
||||
single-round trainer stop_gradients the opponent seat and throws those plies away
|
||||
-- half of every game. Here every ply emits a transition for both seats, each
|
||||
with its own seat-relative view, value and reward; the actor mask keeps the
|
||||
policy loss on the seat that actually moved, while the critic learns from both.
|
||||
Folding the seat axis into the batch keeps each seat's GAE chain independent.
|
||||
|
||||
**Reward is the match, not the round.** Rounds one and two pay nothing, they only
|
||||
bank into ``carry``. The terminal reward is ``tanh(total_diff / terminal_scale)``
|
||||
at the end of round three. Driving that scale toward zero would make it
|
||||
``sign()`` -- the true objective, but a poor signal, since every ply of a
|
||||
~160-ply match would then carry the same +/-1 and all credit assignment would
|
||||
fall to the critic. Potential shaping on the running total covers the gap early
|
||||
and anneals away.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
import flax.linen as nn
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import optax
|
||||
|
||||
from lost_cities_jax.match import (
|
||||
MatchState,
|
||||
match_legal_action_mask,
|
||||
match_reset,
|
||||
match_score,
|
||||
match_step,
|
||||
)
|
||||
from lost_cities_jax.match_obs import (
|
||||
MATCH_CRITIC_OBS_DIM,
|
||||
MATCH_OBS_DIM,
|
||||
match_critic_observation,
|
||||
match_observation,
|
||||
)
|
||||
from lost_cities_jax.ppo import (
|
||||
EpisodeEnd,
|
||||
JaxPPOConfig,
|
||||
TrainState,
|
||||
_append_jsonl,
|
||||
_broadcast_done,
|
||||
_create_run_dir,
|
||||
_write_json,
|
||||
action_log_prob,
|
||||
categorical_entropy,
|
||||
compute_gae,
|
||||
mask_logits,
|
||||
masked_mean,
|
||||
restore_checkpoint,
|
||||
save_checkpoint,
|
||||
shaping_coefficient,
|
||||
)
|
||||
from lost_cities_jax.types import MAX_STEPS, N_ACTIONS, PLAY
|
||||
|
||||
N_SEATS = 2
|
||||
|
||||
|
||||
class MatchTransition(NamedTuple):
|
||||
obs: jax.Array
|
||||
critic_obs: jax.Array
|
||||
legal_mask: jax.Array
|
||||
action: jax.Array
|
||||
log_prob: jax.Array
|
||||
value: jax.Array
|
||||
reward: jax.Array
|
||||
done: jax.Array
|
||||
active: jax.Array
|
||||
actor_mask: jax.Array
|
||||
entropy: jax.Array
|
||||
play_action: jax.Array
|
||||
|
||||
|
||||
class MatchActorCritic(nn.Module):
|
||||
"""Separate trunks, so the critic's privileged input cannot reach the logits."""
|
||||
|
||||
hidden_size: int = 512
|
||||
num_layers: int = 3
|
||||
|
||||
@nn.compact
|
||||
def __call__(self, obs: jax.Array, critic_obs: jax.Array) -> tuple[jax.Array, jax.Array]:
|
||||
x = obs
|
||||
for _ in range(self.num_layers):
|
||||
x = nn.relu(nn.Dense(self.hidden_size)(x))
|
||||
logits = nn.Dense(N_ACTIONS)(x)
|
||||
|
||||
v = critic_obs
|
||||
for _ in range(self.num_layers):
|
||||
v = nn.relu(nn.Dense(self.hidden_size)(v))
|
||||
value = nn.Dense(1)(v)
|
||||
return logits, jnp.squeeze(value, axis=-1)
|
||||
|
||||
|
||||
def create_match_train_state(cfg: JaxPPOConfig, rng: jax.Array) -> TrainState:
|
||||
model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
|
||||
params = model.init(
|
||||
rng,
|
||||
jnp.zeros((1, MATCH_OBS_DIM), dtype=jnp.float32),
|
||||
jnp.zeros((1, MATCH_CRITIC_OBS_DIM), dtype=jnp.float32),
|
||||
)
|
||||
tx = optax.chain(
|
||||
optax.clip_by_global_norm(cfg.ppo.max_grad_norm),
|
||||
optax.adam(cfg.ppo.learning_rate),
|
||||
)
|
||||
return TrainState.create(apply_fn=model.apply, params=params, tx=tx)
|
||||
|
||||
|
||||
def match_lead(env: MatchState) -> jax.Array:
|
||||
"""Running match total from seat 0's side, in raw points."""
|
||||
totals = match_score(env)
|
||||
return (totals[0] - totals[1]).astype(jnp.float32)
|
||||
|
||||
|
||||
def reset_done_matches(env: MatchState, rng: jax.Array, batch: int) -> MatchState:
|
||||
fresh = jax.vmap(match_reset)(jax.random.split(rng, batch))
|
||||
return jax.tree_util.tree_map(
|
||||
lambda old, new: jnp.where(_broadcast_done(env.done, old), new, old), env, fresh
|
||||
)
|
||||
|
||||
|
||||
def _seat_views(env: MatchState, seats: jax.Array):
|
||||
obs = jax.vmap(lambda s: jax.vmap(match_observation, in_axes=(0, None))(env, s))(seats)
|
||||
critic = jax.vmap(lambda s: jax.vmap(match_critic_observation, in_axes=(0, None))(env, s))(
|
||||
seats
|
||||
)
|
||||
return obs, critic
|
||||
|
||||
|
||||
def _episode_end(env: MatchState, terminal: jax.Array, episode_return: jax.Array) -> EpisodeEnd:
|
||||
lead = jax.vmap(match_lead)(env)
|
||||
opened = jnp.sum(env.round.col_len[:, 0, :] > 0, axis=-1)
|
||||
return EpisodeEnd(
|
||||
done=terminal,
|
||||
episode_return=episode_return,
|
||||
length=env.round.step_count,
|
||||
opened_colors=opened,
|
||||
positive_expeditions=lead.astype(jnp.int32), # final match lead, in points
|
||||
hit_max_steps=env.round.step_count >= MAX_STEPS,
|
||||
)
|
||||
|
||||
|
||||
def match_episode_metrics(transitions: MatchTransition, episodes: EpisodeEnd) -> dict:
|
||||
done = episodes.done.astype(jnp.float32)
|
||||
completed = jnp.sum(done)
|
||||
denom = jnp.maximum(completed, 1.0)
|
||||
actor_count = jnp.sum(transitions.actor_mask)
|
||||
lead = episodes.positive_expeditions.astype(jnp.float32)
|
||||
return {
|
||||
"return_mean": jnp.sum(episodes.episode_return * done) / denom,
|
||||
"match_lead_mean": jnp.sum(lead * done) / denom,
|
||||
"match_win_rate": jnp.sum((lead > 0).astype(jnp.float32) * done) / denom,
|
||||
"final_round_length_mean": jnp.sum(episodes.length.astype(jnp.float32) * done) / denom,
|
||||
"max_steps_rate": jnp.sum(episodes.hit_max_steps.astype(jnp.float32) * done) / denom,
|
||||
"play_action_rate": jnp.sum(transitions.play_action) / jnp.maximum(actor_count, 1),
|
||||
"entropy_mean": masked_mean(
|
||||
transitions.entropy, transitions.actor_mask.astype(jnp.float32)
|
||||
),
|
||||
"learner_actions": actor_count,
|
||||
"matches_completed": completed,
|
||||
}
|
||||
|
||||
|
||||
def make_match_rollout_fn(cfg: JaxPPOConfig):
|
||||
batch = cfg.ppo.batch_games
|
||||
seats = jnp.arange(N_SEATS, dtype=jnp.int32)
|
||||
|
||||
@jax.jit
|
||||
def rollout_fn(state: TrainState, env: MatchState, rng: jax.Array, shaping_coef: jax.Array):
|
||||
def body(carry, _):
|
||||
env, episode_return, key = carry
|
||||
key, act_key, reset_key = jax.random.split(key, 3)
|
||||
|
||||
obs, critic_obs = _seat_views(env, seats) # (N_SEATS, batch, ...)
|
||||
legal = jax.vmap(match_legal_action_mask)(env)
|
||||
legal_both = jnp.broadcast_to(legal, (N_SEATS, batch, N_ACTIONS))
|
||||
|
||||
logits, value = state.apply_fn(state.params, obs, critic_obs)
|
||||
masked = mask_logits(logits, legal_both)
|
||||
sampled = jax.random.categorical(act_key, masked, axis=-1).astype(jnp.int32)
|
||||
log_prob = jax.vmap(action_log_prob)(masked, sampled)
|
||||
entropy = jax.vmap(categorical_entropy)(masked)
|
||||
|
||||
to_move = env.round.to_move.astype(jnp.int32)
|
||||
active = ~env.done
|
||||
actor_mask = (seats[:, None] == to_move[None, :]) & active[None, :]
|
||||
actions = jnp.take_along_axis(sampled, to_move[None, :], axis=0)[0]
|
||||
|
||||
before = jax.vmap(match_lead)(env)
|
||||
next_env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, actions)
|
||||
after = jax.vmap(match_lead)(next_env)
|
||||
|
||||
done = next_env.done
|
||||
terminal = active & done
|
||||
reward0 = jnp.where(terminal, jnp.tanh(after / cfg.reward.terminal_scale), 0.0)
|
||||
reward0 = reward0 + shaping_coef * (after - before) / cfg.reward.terminal_scale
|
||||
reward0 = jnp.where(active, reward0, 0.0)
|
||||
# Zero-sum: seat 1 sees exactly the negation.
|
||||
reward = jnp.stack([reward0, -reward0], axis=0)
|
||||
|
||||
played = ((actions % 12) // 6 == PLAY)[None, :] & actor_mask
|
||||
transition = MatchTransition(
|
||||
obs=obs,
|
||||
critic_obs=critic_obs,
|
||||
legal_mask=legal_both,
|
||||
action=jnp.broadcast_to(actions, (N_SEATS, batch)),
|
||||
log_prob=log_prob,
|
||||
value=value,
|
||||
reward=reward,
|
||||
done=jnp.broadcast_to(done, (N_SEATS, batch)),
|
||||
active=jnp.broadcast_to(active, (N_SEATS, batch)),
|
||||
actor_mask=actor_mask,
|
||||
entropy=entropy,
|
||||
play_action=played,
|
||||
)
|
||||
|
||||
episode_return = episode_return + reward0
|
||||
episode = _episode_end(next_env, terminal, episode_return)
|
||||
next_env = reset_done_matches(next_env, reset_key, batch)
|
||||
episode_return = jnp.where(done, 0.0, episode_return)
|
||||
return (next_env, episode_return, key), (transition, episode)
|
||||
|
||||
init_return = jnp.zeros((batch,), dtype=jnp.float32)
|
||||
(next_env, _, rng), (transitions, episodes) = jax.lax.scan(
|
||||
body, (env, init_return, rng), xs=None, length=cfg.ppo.rollout_steps
|
||||
)
|
||||
|
||||
# Fold the seat axis into the batch: each column is then an independent
|
||||
# episode chain, which is exactly what GAE wants.
|
||||
transitions = jax.tree_util.tree_map(
|
||||
lambda x: x.reshape((x.shape[0], N_SEATS * batch, *x.shape[3:])), transitions
|
||||
)
|
||||
final_obs, final_critic = _seat_views(next_env, seats)
|
||||
_, last_value = state.apply_fn(state.params, final_obs, final_critic)
|
||||
last_value = last_value.reshape((N_SEATS * batch,))
|
||||
|
||||
return next_env, transitions, last_value, match_episode_metrics(transitions, episodes)
|
||||
|
||||
return rollout_fn
|
||||
|
||||
|
||||
def match_ppo_update(
|
||||
state: TrainState,
|
||||
transitions: MatchTransition,
|
||||
advantages: jax.Array,
|
||||
returns: jax.Array,
|
||||
rng: jax.Array,
|
||||
cfg: JaxPPOConfig,
|
||||
) -> tuple[TrainState, dict[str, jax.Array]]:
|
||||
batch_size = int(transitions.reward.size)
|
||||
minibatch_size = batch_size // cfg.ppo.minibatches
|
||||
flat = jax.tree_util.tree_map(lambda x: x.reshape((-1, *x.shape[2:])), transitions)
|
||||
flat_advantages = advantages.reshape((batch_size,))
|
||||
flat_returns = returns.reshape((batch_size,))
|
||||
|
||||
actor_mask = flat.actor_mask.astype(jnp.float32)
|
||||
adv_mean = masked_mean(flat_advantages, actor_mask)
|
||||
adv_std = jnp.sqrt(masked_mean((flat_advantages - adv_mean) ** 2, actor_mask) + 1.0e-8)
|
||||
flat_advantages = (flat_advantages - adv_mean) / adv_std
|
||||
|
||||
def loss_fn(params, mb, mb_adv, mb_returns):
|
||||
logits, value = state.apply_fn(params, mb.obs, mb.critic_obs)
|
||||
masked_logits = mask_logits(logits, mb.legal_mask)
|
||||
log_prob = action_log_prob(masked_logits, mb.action)
|
||||
entropy = categorical_entropy(masked_logits)
|
||||
ratio = jnp.exp(log_prob - mb.log_prob)
|
||||
actor_weight = mb.actor_mask.astype(jnp.float32)
|
||||
active_weight = mb.active.astype(jnp.float32)
|
||||
unclipped = ratio * mb_adv
|
||||
clipped = jnp.clip(ratio, 1.0 - cfg.ppo.clip_epsilon, 1.0 + cfg.ppo.clip_epsilon) * mb_adv
|
||||
actor_loss = -masked_mean(jnp.minimum(unclipped, clipped), actor_weight)
|
||||
value_loss = masked_mean((mb_returns - value) ** 2, active_weight)
|
||||
entropy_loss = masked_mean(entropy, actor_weight)
|
||||
loss = actor_loss + cfg.ppo.value_coef * value_loss - cfg.ppo.entropy_coef * entropy_loss
|
||||
return loss, {
|
||||
"loss": loss,
|
||||
"actor_loss": actor_loss,
|
||||
"value_loss": value_loss,
|
||||
"entropy_loss": entropy_loss,
|
||||
"approx_kl": masked_mean(mb.log_prob - log_prob, actor_weight),
|
||||
}
|
||||
|
||||
def epoch_update(carry, key):
|
||||
train_state = carry
|
||||
permutation = jax.random.permutation(key, batch_size)
|
||||
|
||||
def minibatch_update(carry, idx):
|
||||
train_state, accum = carry
|
||||
mb_idx = jax.lax.dynamic_slice(permutation, (idx * minibatch_size,), (minibatch_size,))
|
||||
mb = jax.tree_util.tree_map(lambda x: x[mb_idx], flat)
|
||||
(_, metrics), grads = jax.value_and_grad(loss_fn, has_aux=True)(
|
||||
train_state.params, mb, flat_advantages[mb_idx], flat_returns[mb_idx]
|
||||
)
|
||||
train_state = train_state.apply_gradients(grads=grads)
|
||||
accum = jax.tree_util.tree_map(lambda a, b: a + b, accum, metrics)
|
||||
return (train_state, accum), None
|
||||
|
||||
zeros = {
|
||||
key: jnp.array(0.0)
|
||||
for key in ("loss", "actor_loss", "value_loss", "entropy_loss", "approx_kl")
|
||||
}
|
||||
(train_state, metrics), _ = jax.lax.scan(
|
||||
minibatch_update, (train_state, zeros), jnp.arange(cfg.ppo.minibatches)
|
||||
)
|
||||
metrics = jax.tree_util.tree_map(lambda x: x / cfg.ppo.minibatches, metrics)
|
||||
return train_state, metrics
|
||||
|
||||
keys = jax.random.split(rng, cfg.ppo.epochs)
|
||||
state, metrics = jax.lax.scan(epoch_update, state, keys)
|
||||
return state, jax.tree_util.tree_map(lambda x: jnp.mean(x, axis=0), metrics)
|
||||
|
||||
|
||||
def make_match_train_iteration(cfg: JaxPPOConfig):
|
||||
rollout_fn = make_match_rollout_fn(cfg)
|
||||
|
||||
@jax.jit
|
||||
def train_iteration(
|
||||
state: TrainState, env: MatchState, rng: jax.Array, shaping_coef: jax.Array
|
||||
):
|
||||
rng, rollout_key, update_key = jax.random.split(rng, 3)
|
||||
env, transitions, last_value, stats = rollout_fn(state, env, rollout_key, shaping_coef)
|
||||
advantages, returns = compute_gae(
|
||||
transitions.reward,
|
||||
transitions.value,
|
||||
transitions.done,
|
||||
cfg.ppo.gamma,
|
||||
cfg.ppo.gae_lambda,
|
||||
last_value,
|
||||
)
|
||||
state, update_metrics = match_ppo_update(
|
||||
state, transitions, advantages, returns, update_key, cfg
|
||||
)
|
||||
return state, env, rng, {**stats, **update_metrics}
|
||||
|
||||
return train_iteration
|
||||
|
||||
|
||||
def match_train(cfg: JaxPPOConfig, *, resume: str | None = None) -> Path:
|
||||
run_dir = _create_run_dir(cfg)
|
||||
_write_json(run_dir / "config.json", asdict(cfg))
|
||||
metrics_path = run_dir / "metrics.jsonl"
|
||||
|
||||
rng = jax.random.PRNGKey(cfg.run.seed)
|
||||
rng, init_key, reset_key = jax.random.split(rng, 3)
|
||||
state = create_match_train_state(cfg, init_key)
|
||||
if resume:
|
||||
state = restore_checkpoint(Path(resume), state)
|
||||
|
||||
env = jax.jit(jax.vmap(match_reset))(jax.random.split(reset_key, cfg.ppo.batch_games))
|
||||
train_iteration = make_match_train_iteration(cfg)
|
||||
|
||||
start = time.perf_counter()
|
||||
learner_actions_seen = 0
|
||||
for update in range(cfg.run.total_updates):
|
||||
shaping_coef = jnp.asarray(shaping_coefficient(cfg, learner_actions_seen), jnp.float32)
|
||||
iter_start = time.perf_counter()
|
||||
state, env, rng, metrics = train_iteration(state, env, rng, shaping_coef)
|
||||
jax.tree_util.tree_leaves(metrics)[0].block_until_ready()
|
||||
|
||||
row = {
|
||||
key: float(value) if getattr(value, "shape", ()) == () else value.tolist()
|
||||
for key, value in metrics.items()
|
||||
}
|
||||
learner_actions_seen += int(row["learner_actions"])
|
||||
row.update(
|
||||
update=update,
|
||||
shaping_coef=float(shaping_coef),
|
||||
learner_actions_total=learner_actions_seen,
|
||||
iteration_seconds=time.perf_counter() - iter_start,
|
||||
elapsed_seconds=time.perf_counter() - start,
|
||||
)
|
||||
_append_jsonl(metrics_path, row)
|
||||
if update % cfg.run.log_every == 0:
|
||||
print(
|
||||
json.dumps(
|
||||
{
|
||||
"update": update,
|
||||
"match_win_rate": row["match_win_rate"],
|
||||
"match_lead_mean": row["match_lead_mean"],
|
||||
"matches_completed": row["matches_completed"],
|
||||
"entropy_mean": row["entropy_mean"],
|
||||
},
|
||||
sort_keys=True,
|
||||
),
|
||||
flush=True,
|
||||
)
|
||||
if cfg.run.checkpoint_every and (update + 1) % cfg.run.checkpoint_every == 0:
|
||||
save_checkpoint(run_dir / "latest", state, cfg)
|
||||
|
||||
save_checkpoint(run_dir / "latest", state, cfg)
|
||||
return run_dir
|
||||
@@ -0,0 +1,194 @@
|
||||
"""Phases 2-4: match reward, match observation, asymmetric critic, both-seat training."""
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from lost_cities_jax.match import match_reset, match_score, match_step
|
||||
from lost_cities_jax.match_obs import (
|
||||
MATCH_CRITIC_OBS_DIM,
|
||||
MATCH_OBS_DIM,
|
||||
match_critic_observation,
|
||||
match_observation,
|
||||
)
|
||||
from lost_cities_jax.match_ppo import (
|
||||
create_match_train_state,
|
||||
make_match_rollout_fn,
|
||||
make_match_train_iteration,
|
||||
)
|
||||
from lost_cities_jax.opponents import random_legal_action
|
||||
from lost_cities_jax.ppo import JaxPPOConfig
|
||||
from lost_cities_jax.types import LOC_P0_HAND, N_CARDS
|
||||
|
||||
|
||||
def _cfg(rollout_steps: int = 700, batch_games: int = 8) -> JaxPPOConfig:
|
||||
cfg = JaxPPOConfig()
|
||||
cfg.ppo.batch_games = batch_games
|
||||
cfg.ppo.rollout_steps = rollout_steps
|
||||
cfg.ppo.minibatches = 4
|
||||
cfg.ppo.epochs = 1
|
||||
cfg.network.hidden_size = 32
|
||||
cfg.network.num_layers = 1
|
||||
return cfg
|
||||
|
||||
|
||||
# --- the asymmetric critic -------------------------------------------------
|
||||
|
||||
|
||||
def test_the_critic_sees_the_opponents_hand_and_the_actor_does_not():
|
||||
match = match_reset(jax.random.PRNGKey(0))
|
||||
loc = np.asarray(match.round.card_loc)
|
||||
opp_hand = np.flatnonzero(loc == LOC_P0_HAND + 1)
|
||||
|
||||
actor = np.asarray(match_observation(match, jnp.int32(0)))
|
||||
critic = np.asarray(match_critic_observation(match, jnp.int32(0)))
|
||||
|
||||
# Card-major, three channels each, matching the base observation's layout.
|
||||
hidden = critic[MATCH_OBS_DIM:].reshape(N_CARDS, 3)
|
||||
assert np.array_equal(np.flatnonzero(hidden[:, 0]), opp_hand)
|
||||
|
||||
# The actor's view is a strict prefix of it and carries none of that.
|
||||
assert np.array_equal(actor, critic[:MATCH_OBS_DIM])
|
||||
assert critic.shape == (MATCH_CRITIC_OBS_DIM,)
|
||||
|
||||
|
||||
def test_privileged_input_cannot_move_the_policy_logits():
|
||||
"""Separate trunks, or the critic's view of the deck leaks into play."""
|
||||
cfg = _cfg()
|
||||
state = create_match_train_state(cfg, jax.random.PRNGKey(1))
|
||||
match = match_reset(jax.random.PRNGKey(2))
|
||||
|
||||
obs = match_observation(match, jnp.int32(0))[None, :]
|
||||
critic_a = match_critic_observation(match, jnp.int32(0))[None, :]
|
||||
critic_b = critic_a.at[0, MATCH_OBS_DIM:].set(1.0) # a totally different hidden state
|
||||
|
||||
logits_a, value_a = state.apply_fn(state.params, obs, critic_a)
|
||||
logits_b, value_b = state.apply_fn(state.params, obs, critic_b)
|
||||
|
||||
assert np.array_equal(np.asarray(logits_a), np.asarray(logits_b))
|
||||
# ...and the value head must actually be using it, or the critic is pointless.
|
||||
assert not np.allclose(np.asarray(value_a), np.asarray(value_b))
|
||||
|
||||
|
||||
# --- the match reward ------------------------------------------------------
|
||||
|
||||
|
||||
def test_reward_is_zero_sum_and_only_paid_at_the_end_of_the_match():
|
||||
cfg = _cfg()
|
||||
state = create_match_train_state(cfg, jax.random.PRNGKey(3))
|
||||
env = jax.jit(jax.vmap(match_reset))(
|
||||
jax.random.split(jax.random.PRNGKey(4), cfg.ppo.batch_games)
|
||||
)
|
||||
|
||||
rollout = make_match_rollout_fn(cfg)
|
||||
# Shaping off, so any non-zero reward must be a terminal one.
|
||||
_, transitions, _, metrics = rollout(state, env, jax.random.PRNGKey(5), jnp.asarray(0.0))
|
||||
|
||||
reward = np.asarray(transitions.reward) # (T, 2 * batch)
|
||||
seat0, seat1 = reward[:, : cfg.ppo.batch_games], reward[:, cfg.ppo.batch_games :]
|
||||
assert np.allclose(seat0, -seat1)
|
||||
|
||||
done = np.asarray(transitions.done)[:, : cfg.ppo.batch_games]
|
||||
paid = seat0 != 0.0
|
||||
assert np.array_equal(paid, done & paid) # never paid on a non-terminal ply
|
||||
assert float(metrics["matches_completed"]) > 0
|
||||
assert np.abs(seat0[paid]).max() <= 1.0 # tanh-bounded
|
||||
|
||||
|
||||
def test_shaping_tracks_the_running_match_total():
|
||||
"""Phi is carry + board diff, so shaping must follow the total, not the round."""
|
||||
match = match_reset(jax.random.PRNGKey(6))
|
||||
key = jax.random.PRNGKey(7)
|
||||
|
||||
prev = np.asarray(match_score(match))
|
||||
for _ in range(200):
|
||||
if bool(match.done):
|
||||
break
|
||||
key, step_key = jax.random.split(key)
|
||||
action = random_legal_action(match.round, match.round.to_move, step_key)
|
||||
match, _, _ = match_step(match, action)
|
||||
total = np.asarray(match_score(match))
|
||||
# The total never resets when a round rolls over; it only accumulates.
|
||||
assert total.shape == (2,)
|
||||
prev = total
|
||||
assert prev.shape == (2,)
|
||||
|
||||
|
||||
# --- both seats ------------------------------------------------------------
|
||||
|
||||
|
||||
def test_exactly_one_seat_acts_per_ply_but_both_are_trained():
|
||||
cfg = _cfg()
|
||||
state = create_match_train_state(cfg, jax.random.PRNGKey(8))
|
||||
env = jax.jit(jax.vmap(match_reset))(
|
||||
jax.random.split(jax.random.PRNGKey(9), cfg.ppo.batch_games)
|
||||
)
|
||||
|
||||
rollout = make_match_rollout_fn(cfg)
|
||||
_, transitions, last_value, _ = rollout(state, env, jax.random.PRNGKey(10), jnp.asarray(0.0))
|
||||
|
||||
batch = cfg.ppo.batch_games
|
||||
actor = np.asarray(transitions.actor_mask)
|
||||
seat0, seat1 = actor[:, :batch], actor[:, batch:]
|
||||
|
||||
# One mover per ply while the match is live.
|
||||
live = np.asarray(transitions.active)[:, :batch]
|
||||
assert np.array_equal(seat0 ^ seat1, live)
|
||||
|
||||
# But the critic gets both seats: every live ply is a value target on both.
|
||||
assert np.asarray(transitions.active).sum() == 2 * live.sum()
|
||||
assert transitions.value.shape == (cfg.ppo.rollout_steps, 2 * batch)
|
||||
assert last_value.shape == (2 * batch,)
|
||||
|
||||
|
||||
def test_self_play_is_balanced_and_a_train_step_stays_finite():
|
||||
cfg = _cfg(rollout_steps=700, batch_games=16)
|
||||
state = create_match_train_state(cfg, jax.random.PRNGKey(11))
|
||||
env = jax.jit(jax.vmap(match_reset))(jax.random.split(jax.random.PRNGKey(12), 16))
|
||||
train_iteration = make_match_train_iteration(cfg)
|
||||
|
||||
state, env, _, metrics = train_iteration(state, env, jax.random.PRNGKey(13), jnp.asarray(0.1))
|
||||
|
||||
assert bool(jnp.isfinite(metrics["loss"]))
|
||||
assert float(metrics["matches_completed"]) > 0
|
||||
# One network on both sides: neither seat should be favoured.
|
||||
assert 0.2 < float(metrics["match_win_rate"]) < 0.8
|
||||
|
||||
|
||||
# --- the match observation -------------------------------------------------
|
||||
|
||||
|
||||
def test_carry_reaches_the_observation_with_a_usable_scale():
|
||||
"""The old score_diff divided by 780; a 50-point lead vanished into 0.06."""
|
||||
match = match_reset(jax.random.PRNGKey(14))
|
||||
behind = match._replace(carry=jnp.asarray([0, 50], dtype=jnp.int32))
|
||||
ahead = match._replace(carry=jnp.asarray([50, 0], dtype=jnp.int32))
|
||||
|
||||
obs_behind = np.asarray(match_observation(behind, jnp.int32(0)))
|
||||
obs_ahead = np.asarray(match_observation(ahead, jnp.int32(0)))
|
||||
|
||||
delta = np.abs(obs_ahead - obs_behind)
|
||||
# A 100-point swing has to be plainly visible, not a rounding error.
|
||||
assert delta.max() > 0.5
|
||||
assert (delta > 0.01).sum() >= 2 # the scalar and at least one bin flip
|
||||
|
||||
|
||||
@pytest.mark.parametrize("round_idx", [0, 1, 2])
|
||||
def test_the_round_index_is_observable(round_idx):
|
||||
match = match_reset(jax.random.PRNGKey(15))._replace(
|
||||
round_idx=jnp.asarray(round_idx, dtype=jnp.int32)
|
||||
)
|
||||
obs = np.asarray(match_observation(match, jnp.int32(0)))
|
||||
assert obs.shape == (MATCH_OBS_DIM,)
|
||||
assert np.isfinite(obs).all()
|
||||
|
||||
|
||||
def test_whose_turn_it_is_is_observable():
|
||||
"""The single-round obs never said; the critic had to read it off step_count."""
|
||||
match = match_reset(jax.random.PRNGKey(16))
|
||||
mover = int(match.round.to_move)
|
||||
|
||||
from_mover = np.asarray(match_observation(match, jnp.int32(mover)))
|
||||
from_waiter = np.asarray(match_observation(match, jnp.int32(1 - mover)))
|
||||
assert not np.array_equal(from_mover, from_waiter)
|
||||
Reference in New Issue
Block a user