File size: 6,014 Bytes
9ede8c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """
LoRA self-play fine-tune of an ALREADY-TRAINED checkpoint's value head --
this is the step that actually produced the deployed checkpoint, not a
fresh training run. Starts from the real, working checkpoint (encoder/
dynamics/decoder frozen, matching `train.py`'s own self-play stage, which
never touches them either) and fine-tunes the value head under self-play
through a LoRA-constrained delta rather than an unconstrained further Adam
update -- directly tests whether constraining HOW MUCH the value head can
move (not just how much data/how many rounds) helps it absorb self-play
signal without regressing calibration.
`opponent_strong_epsilon` (the curriculum variant): mixes the stronger
1-ply-deeper heuristic into training ALONGSIDE the self-play snapshots and
the original weak heuristic -- without this, self-play rounds only ever
train against [self-play snapshots, weak heuristic], and never the harder
fixed opponent at all. This is what produced this project's best-confirmed
result (see the whitepaper's results table).
rank=4/alpha=4.0: found to work well for a value head elsewhere in this
project's development on this same architecture; not independently
re-tuned for ConnectX specifically.
"""
import copy
import random
import torch
from connectx.env import ConnectXEnv
from scripts.train import make_selfplay_pool_opponent_fn, random_baseline_win_rate
from connectx.verifier import train_mc_value_onpolicy
from connectx.lora import LoRALinear, apply_lora
from connectx.search import load_checkpoint, evaluate
def _merge_and_unwrap(module):
"""Merge every LoRALinear's delta into its frozen base weight, then
replace the wrapper with the plain (now-merged) nn.Linear -- so the
saved checkpoint is an ORDINARY WorldModel state_dict, loadable by
every existing caller with zero LoRA-awareness needed downstream."""
for name, child in list(module.named_children()):
if isinstance(child, LoRALinear):
child.merge_into_base()
setattr(module, name, child.linear)
else:
_merge_and_unwrap(child)
def main(ckpt_path="checkpoints/connectx_checkpoint.pt", seed=0, rank=4, alpha=4.0,
selfplay_rounds=5, selfplay_epsilon=0.4, selfplay_mc_rounds_per_iter=15,
mc_problems_per_round=400, selfplay_pool_size=5,
opponent_epsilon=0.15, opponent_strong_epsilon=0.0,
save_path="checkpoints/connectx_checkpoint_lora_selfplay.pt"):
raw_ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
model, normalizer = load_checkpoint(ckpt_path)
rng = random.Random(seed)
torch.manual_seed(seed)
random.seed(seed) # env.py's opponent stochasticity reads the GLOBAL
# random module, not this local rng -- see train.py's own note
env = ConnectXEnv(width=raw_ckpt["board_width"], height=raw_ckpt["board_height"],
win_len=raw_ckpt["win_len"])
max_steps = (env.width * env.height) // 2 + 2
lora_params = apply_lora(model.value, rank=rank, alpha=alpha)
n_lora = sum(p.numel() for p in lora_params)
n_frozen = sum(p.numel() for p in model.value.parameters()) - n_lora
print(f"LoRA rank={rank} alpha={alpha} on model.value: "
f"{n_lora} trainable params, {n_frozen} frozen (base value head)")
frozen_pool = []
for sp_round in range(selfplay_rounds):
frozen_model = copy.deepcopy(model).eval()
for p in frozen_model.parameters():
p.requires_grad_(False)
frozen_pool.append(frozen_model)
if len(frozen_pool) > selfplay_pool_size:
frozen_pool.pop(0)
print(f"\nSelf-play LoRA fine-tune round {sp_round + 1}/{selfplay_rounds} "
f"(opponent_selfplay_epsilon={selfplay_epsilon}, pool_size={len(frozen_pool)})...")
opponent_env = ConnectXEnv(width=env.width, height=env.height, win_len=env.win_len)
selfplay_fn = make_selfplay_pool_opponent_fn(frozen_pool, normalizer, opponent_env)
selfplay_train_env = ConnectXEnv(width=env.width, height=env.height, win_len=env.win_len,
opponent_epsilon=opponent_epsilon,
opponent_strong_epsilon=opponent_strong_epsilon,
opponent_selfplay_epsilon=selfplay_epsilon,
opponent_policy_fn=selfplay_fn)
# Adam(model.value.parameters()) inside train_mc_value_onpolicy
# naturally trains ONLY the LoRA deltas here: the wrapped base
# linears have requires_grad=False (set by apply_lora), so their
# .grad stays None and Adam's step() skips them.
train_mc_value_onpolicy(selfplay_train_env, model, normalizer, rng,
n_rounds=selfplay_mc_rounds_per_iter,
n_problems_per_round=mc_problems_per_round,
max_steps=max_steps, unsolved_penalty=max_steps)
_merge_and_unwrap(model.value)
torch.save({
"model_state": model.state_dict(),
"norm_mean": normalizer.mean.cpu(),
"norm_std": normalizer.std.cpu(),
"state_dim": raw_ckpt["state_dim"],
"num_actions": raw_ckpt["num_actions"],
"latent_dim": raw_ckpt["latent_dim"],
"hidden_dim": raw_ckpt["hidden_dim"],
"board_width": env.width,
"board_height": env.height,
"win_len": env.win_len,
}, save_path)
print(f"\nSaved LoRA-self-play-fine-tuned checkpoint to {save_path}")
print("\n" + "=" * 20 + " EVALUATION (no oracle -- vs. random-legal-play baseline only) " + "=" * 20)
eval_rng = random.Random(999)
problems = [env.random_problem(eval_rng) for _ in range(150)]
random_baseline_win_rate(env, problems, max_steps, random.Random(1000))
evaluate(env, model, normalizer, problems, depth=1, beam_width=8, max_total_steps=max_steps,
label="Baseline A (model, depth=1)")
if __name__ == "__main__":
main()
|