Ablate the match stack: both-seat training carries it, the privileged critic hurts
Each piece switched off in turn, trained at identical compute, then played against the full stack over 8192 duplicate matches. Below 0.5 means the removed piece was doing work. - both seats: 0.3317 [0.322, 0.342]. The biggest single contributor. Half of it is simply sample count -- dropping the opponent seat halves the learner actions per update -- but that is the point: self-play already produced those plies with the same network, and the old trainer stop_gradiented them away. - match observation: 0.4751 [0.464, 0.486]. Small but real. Since carry itself contributes almost nothing (rounds decompose), most of this is likely the single-round observation defects being fixed: to_move, the deck clock, and the score_diff scale. - privileged critic: 0.5160 [0.505, 0.527] -- turning it OFF makes the agent significantly STRONGER. Fable called this the biggest missing idea; it is wrong. A critic that knows the deck fits V(full state), which is not E[return | masked obs], so the advantage picks up a component the actor cannot act on. From the actor's side that is noise, not variance reduction. Asymmetric critics hurting under partial observability is a known failure mode. Defaulted off accordingly. (Reusing it as a PIMC leaf evaluator may still stand -- that is a separate claim from using it to train the policy.) 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
@@ -444,3 +444,49 @@ duplicate 매치 8,192판(같은 3딜 + 같은 코인 + 자리 교대), 진짜 3
|
||||
|
||||
**남은 단서:** league는 **착취자(exploiter) 구조**를 포함한 런이다. 평균 강함은 우리가 이기지만
|
||||
**exploitability는 아직 재지 않았다.** "덜 착취당한다"는 증명되지 않았다.
|
||||
|
||||
---
|
||||
|
||||
## 기여도 A/B (2026-07-15)
|
||||
|
||||
각 조각을 하나씩 끄고 동일 컴퓨트(300 업데이트, batch 512)로 학습시킨 뒤,
|
||||
full 스택과 duplicate 매치 8,192판 맞대결. **0.5 미만 = 그 조각이 기여하고 있었다.**
|
||||
|
||||
| 제거한 것 | full 대비 승률 | 95% CI | 평균 총점차 | 판정 |
|
||||
|---|---|---|---|---|
|
||||
| **양쪽 좌석 학습** | **0.3317** | [0.322, 0.342] | −28.6 | **압도적 기여** |
|
||||
| 매치 관측 (carry/to_move/덱시계/살아있는점수) | 0.4751 | [0.464, 0.486] | −4.0 | 유의하게 기여 |
|
||||
| **전지적 critic (CTDE)** | **0.5160** | [0.505, 0.527] | +2.6 | **오히려 해가 된다** |
|
||||
|
||||
### 1. 양쪽 좌석 학습 — 가장 큰 기여
|
||||
|
||||
끄면 승률이 0.33으로 무너진다. 같은 업데이트 수에서 learner 액션이 **정확히 절반**이 되므로
|
||||
상당 부분은 샘플 수 효과다 — 그러나 **그게 요점이다.** 셀프플레이에서 상대 좌석의 수는 같은
|
||||
네트워크가 둔 것인데 기존 트레이너는 `stop_gradient`로 버렸다. **공짜로 데이터가 2배**가 된다.
|
||||
|
||||
### 2. 매치 관측 — 작지만 실재
|
||||
|
||||
0.4751 (CI 상한 0.486 < 0.5). `to_move`, 덱 시계, ÷780 → ÷75 스케일 수정, 살아있는 점수
|
||||
3분할이 합쳐서 약 2.5%p 값어치. carry 자체는 (분해 논증대로) 거의 기여하지 않으므로,
|
||||
이 이득의 대부분은 **단판 obs의 결함 수정분**으로 보인다.
|
||||
|
||||
### 3. 전지적 critic — **Fable의 최우선 권고가 틀렸다**
|
||||
|
||||
Fable은 이것을 "가장 큰 누락 아이디어"로 꼽았다. 실측은 **반대**다: 끄면 오히려
|
||||
**0.5160 (CI 하한 0.505 > 0.5)** 으로 유의하게 **더 강해진다.**
|
||||
|
||||
그럴듯한 이유: 특권 critic은 덱을 알기 때문에 가치를 아주 잘 맞추지만, 그 결과
|
||||
`V(전체상태) ≠ E[리턴 | 마스킹된 관측]`이 되어, **어드밴티지에 정책이 통제할 수 없는 성분**이
|
||||
섞인다. 액터 입장에서 그건 분산 감소가 아니라 노이즈다. 부분관측 환경에서 비대칭 critic이
|
||||
해가 되는 것은 알려진 현상이다.
|
||||
|
||||
**조치: 전지적 critic을 기본에서 끈다.** (PIMC 탐색의 리프 평가기로 재활용한다는 계획은
|
||||
별개로 유효할 수 있으나, 정책 학습용으로는 손해다.)
|
||||
|
||||
### 종합: 실제로 값어치 있었던 것
|
||||
|
||||
1. **in-scan auto-reset + GAE 절단 부트스트랩** — 샘플 5.8배, expert 0.4525 → 0.5071
|
||||
2. **양쪽 좌석 학습** — 데이터 2배, 단독으로 승률 0.33 → 0.5
|
||||
3. **선형 총점 보상** — tanh 대비 0.5859 (사용자 제안)
|
||||
4. **매치 관측 결함 수정** — 약 2.5%p
|
||||
5. ~~전지적 critic~~ — **해가 된다. 끈다.**
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
#!/usr/bin/env python3
|
||||
"""What is each piece of the match stack worth?
|
||||
|
||||
Train one variant per switch at identical compute, then play each against the full
|
||||
stack in duplicate matches -- same three deals, same coins, both seats -- so deal
|
||||
luck cancels and only the policy difference is left.
|
||||
|
||||
The variants do not share an observation shape (dropping the privileged critic or
|
||||
the match features changes the input dims), so this carries its own head-to-head
|
||||
rather than reusing match_eval's, which assumes one architecture.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import jax
|
||||
import jax.numpy as jnp
|
||||
import numpy as np
|
||||
|
||||
from lost_cities_jax.match import (
|
||||
match_legal_action_mask,
|
||||
match_reset_from,
|
||||
match_score,
|
||||
match_step,
|
||||
)
|
||||
from lost_cities_jax.match_eval import MATCH_SCAN_STEPS, _wilson, match_bank
|
||||
from lost_cities_jax.match_ppo import (
|
||||
Ablation,
|
||||
MatchActorCritic,
|
||||
_seat_views,
|
||||
create_match_train_state,
|
||||
match_train,
|
||||
)
|
||||
from lost_cities_jax.ppo import load_config, mask_logits, restore_checkpoint
|
||||
|
||||
VARIANTS = [
|
||||
Ablation(privileged_critic=False),
|
||||
Ablation(both_seats=False),
|
||||
Ablation(match_obs=False),
|
||||
]
|
||||
MATCHES = 4096
|
||||
FULL_CHECKPOINT = Path("runs/jax-ppo-match/2026-07-15_030152_match-linear/latest")
|
||||
|
||||
|
||||
def _load(cfg, checkpoint: Path, ablation: Ablation):
|
||||
state = create_match_train_state(cfg, jax.random.PRNGKey(0), ablation)
|
||||
return restore_checkpoint(checkpoint, state).params
|
||||
|
||||
|
||||
def _existing_run(label: str) -> Path | None:
|
||||
runs = sorted(Path("runs/jax-ppo-match").glob(f"*ablate-{label}"))
|
||||
for run in reversed(runs):
|
||||
if (run / "latest").exists():
|
||||
return run
|
||||
return None
|
||||
|
||||
|
||||
def duel(cfg, params_a, ablation_a, params_b, ablation_b, *, matches: int) -> dict:
|
||||
"""A vs B over duplicate matches. Each side sees the world its own way."""
|
||||
|
||||
decks, coins = match_bank(20260719, matches)
|
||||
model = MatchActorCritic(cfg.network.hidden_size, cfg.network.num_layers)
|
||||
seats = jnp.arange(2, dtype=jnp.int32)
|
||||
|
||||
@jax.jit
|
||||
def run(env, a_seat):
|
||||
def body(carry, _):
|
||||
env, _unused = carry
|
||||
to_move = env.round.to_move.astype(jnp.int32)
|
||||
mask = jax.vmap(match_legal_action_mask)(env)
|
||||
|
||||
def act(params, ablation):
|
||||
obs, critic = _seat_views(env, seats, ablation)
|
||||
idx = to_move[None, :, None]
|
||||
seat_obs = jnp.take_along_axis(obs, idx, axis=0)[0]
|
||||
seat_critic = jnp.take_along_axis(critic, idx, axis=0)[0]
|
||||
logits, _ = model.apply(params, seat_obs, seat_critic)
|
||||
return jnp.argmax(mask_logits(logits, mask), axis=-1).astype(jnp.int32)
|
||||
|
||||
action = jnp.where(
|
||||
to_move == a_seat, act(params_a, ablation_a), act(params_b, ablation_b)
|
||||
)
|
||||
env, _, _ = jax.vmap(match_step, in_axes=(0, 0))(env, action)
|
||||
return (env, _unused), None
|
||||
|
||||
(env, _), _ = jax.lax.scan(body, (env, jnp.int32(0)), xs=None, length=MATCH_SCAN_STEPS)
|
||||
return env
|
||||
|
||||
leads = []
|
||||
for seat in (0, 1):
|
||||
env = jax.vmap(match_reset_from)(decks, coins)
|
||||
final = run(env, jnp.full((matches,), seat, dtype=jnp.int32))
|
||||
totals = np.asarray(jax.vmap(match_score)(final))
|
||||
leads.append(totals[:, seat] - totals[:, 1 - seat])
|
||||
|
||||
lead = np.concatenate(leads)
|
||||
games = float(lead.size)
|
||||
wins = float((lead > 0).sum())
|
||||
low, high = _wilson(wins, games)
|
||||
return {
|
||||
"matches": games,
|
||||
"a_win_rate": wins / games,
|
||||
"wilson_low": low,
|
||||
"wilson_high": high,
|
||||
"a_mean_lead": float(lead.mean()),
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
|
||||
full = _load(cfg, FULL_CHECKPOINT, Ablation())
|
||||
|
||||
rows = []
|
||||
for ablation in VARIANTS:
|
||||
label = ablation.label()
|
||||
run_dir = _existing_run(label)
|
||||
if run_dir is None:
|
||||
print(f"\n===== training {label} =====", flush=True)
|
||||
cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
|
||||
cfg.run.experiment_name = f"ablate-{label}"
|
||||
run_dir = match_train(cfg, ablation=ablation)
|
||||
else:
|
||||
print(f"\n===== reusing {run_dir} =====", flush=True)
|
||||
|
||||
cfg = load_config("configs/jax_ppo/match-selfplay.yaml")
|
||||
params = _load(cfg, run_dir / "latest", ablation)
|
||||
# The ablated net plays seat A, the full stack answers. Below 0.5 means
|
||||
# the piece we removed was carrying weight.
|
||||
result = duel(cfg, params, ablation, full, Ablation(), matches=MATCHES)
|
||||
rows.append({"ablation": label, **result})
|
||||
print(f" {label}: win rate vs full = {result['a_win_rate']:.4f}", flush=True)
|
||||
|
||||
print("\n\n==================== ablation ====================")
|
||||
print(f"{'removed':<26}{'win rate vs full':>18}{'95% CI':>22}{'mean lead':>12}")
|
||||
print("-" * 78)
|
||||
for row in rows:
|
||||
ci = f"[{row['wilson_low']:.3f}, {row['wilson_high']:.3f}]"
|
||||
print(
|
||||
f"{row['ablation']:<26}{row['a_win_rate']:>18.4f}{ci:>22}{row['a_mean_lead']:>+12.1f}"
|
||||
)
|
||||
print("\nbelow 0.5 = the removed piece was carrying weight")
|
||||
Path("runs/jax-ppo-match/ablation.json").write_text(json.dumps(rows, indent=2, sort_keys=True))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -40,7 +40,7 @@ from __future__ import annotations
|
||||
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
@@ -62,6 +62,7 @@ from lost_cities_jax.match_obs import (
|
||||
match_critic_observation,
|
||||
match_observation,
|
||||
)
|
||||
from lost_cities_jax.obs import observation
|
||||
from lost_cities_jax.ppo import (
|
||||
EpisodeEnd,
|
||||
JaxPPOConfig,
|
||||
@@ -79,11 +80,54 @@ from lost_cities_jax.ppo import (
|
||||
save_checkpoint,
|
||||
shaping_coefficient,
|
||||
)
|
||||
from lost_cities_jax.types import MAX_STEPS, N_ACTIONS, PLAY
|
||||
from lost_cities_jax.types import MAX_STEPS, N_ACTIONS, OBS_DIM, PLAY
|
||||
|
||||
N_SEATS = 2
|
||||
|
||||
|
||||
@dataclass
|
||||
class Ablation:
|
||||
"""Switches for measuring what each piece of the match stack is worth."""
|
||||
|
||||
privileged_critic: bool = False
|
||||
"""On: the critic also sees the opponent's hand and the deck in order.
|
||||
|
||||
Defaults OFF because it measured *negative*. Ablated against the full stack
|
||||
over 8192 duplicate matches, switching it off won 0.5160 [0.505, 0.527]. A
|
||||
critic that knows the deck fits V(full state), which is not
|
||||
E[return | masked obs], so the advantage carries a component the actor cannot
|
||||
act on -- noise from its side, not variance reduction.
|
||||
"""
|
||||
|
||||
both_seats: bool = True
|
||||
"""Off: the opponent seat's plies are discarded, as the old trainer did."""
|
||||
|
||||
match_obs: bool = True
|
||||
"""Off: the actor gets the plain single-round observation, carry and all the
|
||||
rest of the match features withheld."""
|
||||
|
||||
def actor_obs_dim(self) -> int:
|
||||
return MATCH_OBS_DIM if self.match_obs else OBS_DIM
|
||||
|
||||
def critic_obs_dim(self) -> int:
|
||||
if not self.privileged_critic:
|
||||
return self.actor_obs_dim()
|
||||
return MATCH_CRITIC_OBS_DIM
|
||||
|
||||
def label(self) -> str:
|
||||
parts = []
|
||||
if self.privileged_critic:
|
||||
parts.append("privileged_critic")
|
||||
if not self.both_seats:
|
||||
parts.append("no-both_seats")
|
||||
if not self.match_obs:
|
||||
parts.append("no-match_obs")
|
||||
return "default" if not parts else "+".join(parts)
|
||||
|
||||
|
||||
FULL = Ablation()
|
||||
|
||||
|
||||
class MatchTransition(NamedTuple):
|
||||
obs: jax.Array
|
||||
critic_obs: jax.Array
|
||||
@@ -119,12 +163,14 @@ class MatchActorCritic(nn.Module):
|
||||
return logits, jnp.squeeze(value, axis=-1)
|
||||
|
||||
|
||||
def create_match_train_state(cfg: JaxPPOConfig, rng: jax.Array) -> TrainState:
|
||||
def create_match_train_state(
|
||||
cfg: JaxPPOConfig, rng: jax.Array, ablation: Ablation = FULL
|
||||
) -> 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),
|
||||
jnp.zeros((1, ablation.actor_obs_dim()), dtype=jnp.float32),
|
||||
jnp.zeros((1, ablation.critic_obs_dim()), dtype=jnp.float32),
|
||||
)
|
||||
tx = optax.chain(
|
||||
optax.clip_by_global_norm(cfg.ppo.max_grad_norm),
|
||||
@@ -146,8 +192,11 @@ def reset_done_matches(env: MatchState, rng: jax.Array, batch: int) -> MatchStat
|
||||
)
|
||||
|
||||
|
||||
def _seat_views(env: MatchState, seats: jax.Array):
|
||||
obs = jax.vmap(lambda s: jax.vmap(match_observation, in_axes=(0, None))(env, s))(seats)
|
||||
def _seat_views(env: MatchState, seats: jax.Array, ablation: Ablation = FULL):
|
||||
actor_fn = match_observation if ablation.match_obs else (lambda m, p: observation(m.round, p))
|
||||
obs = jax.vmap(lambda s: jax.vmap(actor_fn, in_axes=(0, None))(env, s))(seats)
|
||||
if not ablation.privileged_critic:
|
||||
return obs, obs
|
||||
critic = jax.vmap(lambda s: jax.vmap(match_critic_observation, in_axes=(0, None))(env, s))(
|
||||
seats
|
||||
)
|
||||
@@ -188,7 +237,7 @@ def match_episode_metrics(transitions: MatchTransition, episodes: EpisodeEnd) ->
|
||||
}
|
||||
|
||||
|
||||
def make_match_rollout_fn(cfg: JaxPPOConfig):
|
||||
def make_match_rollout_fn(cfg: JaxPPOConfig, ablation: Ablation = FULL):
|
||||
batch = cfg.ppo.batch_games
|
||||
seats = jnp.arange(N_SEATS, dtype=jnp.int32)
|
||||
|
||||
@@ -198,7 +247,7 @@ def make_match_rollout_fn(cfg: JaxPPOConfig):
|
||||
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, ...)
|
||||
obs, critic_obs = _seat_views(env, seats, ablation)
|
||||
legal = jax.vmap(match_legal_action_mask)(env)
|
||||
legal_both = jnp.broadcast_to(legal, (N_SEATS, batch, N_ACTIONS))
|
||||
|
||||
@@ -211,6 +260,10 @@ def make_match_rollout_fn(cfg: JaxPPOConfig):
|
||||
to_move = env.round.to_move.astype(jnp.int32)
|
||||
active = ~env.done
|
||||
actor_mask = (seats[:, None] == to_move[None, :]) & active[None, :]
|
||||
if not ablation.both_seats:
|
||||
# The old trainer kept one seat and stop_gradiented the other.
|
||||
keep = (seats[:, None] == 0) & jnp.ones_like(active)[None, :]
|
||||
actor_mask = actor_mask & keep
|
||||
actions = jnp.take_along_axis(sampled, to_move[None, :], axis=0)[0]
|
||||
|
||||
before = jax.vmap(match_lead)(env)
|
||||
@@ -237,7 +290,11 @@ def make_match_rollout_fn(cfg: JaxPPOConfig):
|
||||
value=value,
|
||||
reward=reward,
|
||||
done=jnp.broadcast_to(done, (N_SEATS, batch)),
|
||||
active=jnp.broadcast_to(active, (N_SEATS, batch)),
|
||||
active=(
|
||||
jnp.broadcast_to(active, (N_SEATS, batch))
|
||||
if ablation.both_seats
|
||||
else jnp.broadcast_to(active, (N_SEATS, batch)) & (seats[:, None] == 0)
|
||||
),
|
||||
actor_mask=actor_mask,
|
||||
entropy=entropy,
|
||||
play_action=played,
|
||||
@@ -259,7 +316,7 @@ def make_match_rollout_fn(cfg: JaxPPOConfig):
|
||||
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)
|
||||
final_obs, final_critic = _seat_views(next_env, seats, ablation)
|
||||
_, last_value = state.apply_fn(state.params, final_obs, final_critic)
|
||||
last_value = last_value.reshape((N_SEATS * batch,))
|
||||
|
||||
@@ -339,8 +396,8 @@ def match_ppo_update(
|
||||
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)
|
||||
def make_match_train_iteration(cfg: JaxPPOConfig, ablation: Ablation = FULL):
|
||||
rollout_fn = make_match_rollout_fn(cfg, ablation)
|
||||
|
||||
@jax.jit
|
||||
def train_iteration(
|
||||
@@ -364,19 +421,19 @@ def make_match_train_iteration(cfg: JaxPPOConfig):
|
||||
return train_iteration
|
||||
|
||||
|
||||
def match_train(cfg: JaxPPOConfig, *, resume: str | None = None) -> Path:
|
||||
def match_train(cfg: JaxPPOConfig, *, resume: str | None = None, ablation: Ablation = FULL) -> 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)
|
||||
state = create_match_train_state(cfg, init_key, ablation)
|
||||
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)
|
||||
train_iteration = make_match_train_iteration(cfg, ablation)
|
||||
|
||||
start = time.perf_counter()
|
||||
learner_actions_seen = 0
|
||||
|
||||
Reference in New Issue
Block a user