| """QOJ residual training: teacher/offline warmup, fresh PPO and snapshot league.""" |
| import argparse |
| import json |
| import random |
| import sys |
| import time |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parents[1] |
| if __package__ in (None, ''): |
| sys.path.insert(0, str(ROOT)) |
|
|
| import torch |
|
|
| from qojbot.curriculum import append_jsonl, reconcile_jsonl, seed_splits |
| from qojbot.history import write_json |
| from QOJDoudizhuModel.dataset import cap_candidates, compact, iter_samples, prepare |
| from QOJDoudizhuModel.evaluate import evaluate |
| from QOJDoudizhuModel.learning import batches, collect, losses, update, warmup, rollout_kl |
| from QOJDoudizhuModel.network import ModelConfig |
| from QOJDoudizhuModel.expansion import expand_model |
| from QOJDoudizhuModel.policy import BASELINE_DIR, QOJPolicy |
| from QOJDoudizhuModel.storage import atomic_save, cpu_tree, digest, export, source_hashes, teacher_hashes |
|
|
|
|
| def make_parser(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument('--data', type=Path, default=ROOT / 'data/qoj-history') |
| p.add_argument('--output', type=Path, default=ROOT / 'QOJDoudizhuModel/runs/main') |
| p.add_argument('--teacher-dir', type=Path, default=BASELINE_DIR) |
| p.add_argument('--device', choices=('auto', 'cpu', 'mps', 'cuda'), default='mps') |
| p.add_argument('--width', type=int, default=128) |
| p.add_argument('--heads', type=int, default=4) |
| p.add_argument('--rank-layers', type=int, default=2) |
| p.add_argument('--history-layers', type=int, default=2) |
| p.add_argument('--init-checkpoint', type=Path, |
| help='continue learned weights in an expanded model; starts a new optimizer/run') |
| p.add_argument('--opponent-dir', type=Path, |
| help='target DouZero opponent for PPO and validation/test selection') |
| p.add_argument('--rebase-prior', type=Path, |
| help='switch frozen prior, resetting only residual actor heads and gates') |
| p.add_argument('--history-weight', type=float, default=.15) |
| p.add_argument('--distill-weight', type=float, default=0.) |
| p.add_argument('--offline-anchor-weight', type=float, default=.35) |
| p.add_argument('--max-ppo-kl', type=float, default=0., |
| help='rollback a whole round if exact rollout KL exceeds this limit; 0 disables') |
| p.add_argument('--rollback-tolerance', type=float, default=0., |
| help='rollback validation regressions beyond this score tolerance; 0 disables') |
| p.add_argument('--targeted-league', action='store_true') |
| p.add_argument('--teacher-anchor-states', type=int, default=2048) |
| p.add_argument('--checkpoint-interval', type=int, default=500, |
| help='commit offline progress every N minibatches') |
| p.add_argument('--warmup-games', type=int, default=32) |
| p.add_argument('--offline-epochs', type=int, default=2) |
| p.add_argument('--rounds', type=int, default=128) |
| p.add_argument('--games-per-round', type=int, default=12) |
| p.add_argument('--ppo-epochs', type=int, default=2) |
| p.add_argument('--batch-states', type=int, default=16) |
| p.add_argument('--max-batch-actions', type=int, default=4096) |
| p.add_argument('--candidate-cap', type=int, default=64) |
| p.add_argument('--max-history-games', type=int, default=0) |
| p.add_argument('--holdout', type=float, default=.1) |
| p.add_argument('--learning-rate', type=float, default=1e-4) |
| p.add_argument('--ppo-learning-rate', type=float, default=2e-5) |
| p.add_argument('--ppo-start', choices=('best', 'latest'), default='best', |
| help='start PPO from validation best or the latest offline representation') |
| p.add_argument('--anchor-weight', type=float, default=.02) |
| p.add_argument('--league-size', type=int, default=4) |
| p.add_argument('--validation-interval', type=int, default=16) |
| p.add_argument('--validation-deals', type=int, default=12) |
| p.add_argument('--test-deals', type=int, default=48) |
| p.add_argument('--seed', type=int, default=20261002) |
| p.add_argument('--validation-seed', type=int, default=30261002) |
| p.add_argument('--test-seed', type=int, default=40261002) |
| p.add_argument('--mode', choices=('fixed', 'bidding', 'mixed'), default='mixed') |
| p.add_argument('--resume', action='store_true') |
| p.add_argument('--skip-final-test', action='store_true', |
| help='smoke/continuing runs only; never claims final playing strength') |
| return p |
|
|
|
|
| def holdout_loss(policy, directory, manifest, split, args): |
| policy.model.eval() |
| total, n = 0., 0 |
| with torch.no_grad(): |
| for batch in batches(iter_samples(directory, manifest, split), args.batch_states, |
| args.max_batch_actions): |
| loss, _ = losses(policy.model, batch, policy.device, anchor_weight=args.offline_anchor_weight, |
| history_weight=args.history_weight, distill_weight=args.distill_weight) |
| total += float(loss.cpu()) * len(batch) |
| n += len(batch) |
| return {'samples': n, 'loss': total / n if n else None} |
|
|
|
|
| def train(args): |
| for key in ('width', 'heads', 'rank_layers', 'history_layers', 'checkpoint_interval', |
| 'games_per_round', 'ppo_epochs', 'batch_states', 'max_batch_actions', |
| 'validation_interval', 'validation_deals', 'test_deals', 'league_size', |
| 'teacher_anchor_states'): |
| if getattr(args, key) < 1: |
| raise ValueError(f'{key} must be positive') |
| if args.width % args.heads or args.warmup_games < 0 or args.offline_epochs < 0 or args.rounds < 0: |
| raise ValueError('width must be divisible by heads; counts must be nonnegative') |
| if not 0 <= args.holdout < .5 or (args.candidate_cap != 0 and args.candidate_cap < 2): |
| raise ValueError('invalid holdout/candidate cap') |
| if min(args.learning_rate, args.ppo_learning_rate) <= 0 or args.anchor_weight < 0: |
| raise ValueError('invalid optimizer settings') |
| if min(args.history_weight, args.distill_weight, args.offline_anchor_weight, |
| args.max_ppo_kl, args.rollback_tolerance) < 0: |
| raise ValueError('invalid supervision/safety settings') |
| count = args.warmup_games + args.rounds * args.games_per_round |
| seed_splits(args.seed, max(count, 1), args.validation_seed, args.validation_deals, |
| args.test_seed, args.test_deals) |
| output = args.output.resolve() |
| protected = [args.teacher_dir.resolve(), ROOT / 'models', ROOT / 'Douzero_Resnet', |
| ROOT / 'douzero-resnet-2.0', |
| args.data.resolve()] |
| if args.init_checkpoint: |
| path = args.init_checkpoint.resolve() |
| protected.append(path if path.is_dir() else path.parent) |
| if args.opponent_dir: |
| protected.append(args.opponent_dir.resolve()) |
| if args.rebase_prior: |
| protected.append(args.rebase_prior.resolve()) |
| if any(output == path or path in output.parents for path in protected): |
| raise ValueError('output must not overwrite input weights/data') |
| if args.resume and not (output / 'resume.pt').exists(): |
| raise ValueError('resume.pt missing') |
| if not args.resume and output.exists() and any(output.iterdir()): |
| raise ValueError('output is nonempty; use --resume or new output') |
| torch.set_num_threads(1) |
| torch.manual_seed(args.seed) |
| rng = random.Random(args.seed) |
| model_config = ModelConfig(width=args.width, heads=args.heads, |
| rank_layers=args.rank_layers, history_layers=args.history_layers) |
| policy = QOJPolicy(args.init_checkpoint, device=args.device, config=model_config, |
| teacher_dir=None if args.rebase_prior else args.teacher_dir) |
| initialization = None |
| if args.init_checkpoint: |
| path = args.init_checkpoint |
| path = path / 'best.pt' if path.is_dir() else path |
| previous = sum(p.numel() for p in policy.model.parameters()) |
| policy.model = expand_model(policy.model, model_config).eval() |
| initialization = {'checkpoint': str(path.resolve()), 'sha256': digest(path), |
| 'previous_parameters': previous} |
| if args.rebase_prior: |
| from qojbot.douzero import DouZeroPolicy |
| previous_teacher = teacher_hashes(policy.teacher) |
| policy.teacher = DouZeroPolicy(args.rebase_prior, device=str(policy.device), batch_size=256) |
| if policy.teacher.model_type not in ('lstm', 'resnet_best'): |
| raise ValueError('unsupported rebased prior') |
| for model in policy.teacher.models.values(): |
| model.requires_grad_(False) |
| with torch.no_grad(): |
| for head in policy.model.actors: |
| head.weight.zero_() |
| head.bias.zero_() |
| policy.model.prior_gate.zero_() |
| initialization = {**(initialization or {}), 'rebased_from_teacher': previous_teacher, |
| 'reset_parameters': 'actor heads and prior gates only'} |
| opponent = None |
| if args.opponent_dir: |
| from qojbot.douzero import DouZeroPolicy |
| opponent = (policy.teacher if args.opponent_dir.resolve() == next(iter(policy.teacher.paths.values())).parent |
| else DouZeroPolicy(args.opponent_dir, device=str(policy.device), batch_size=256)) |
| output.mkdir(parents=True, exist_ok=True) |
| config = {key: str(value.resolve()) if isinstance(value, Path) else value |
| for key, value in vars(args).items() if key not in ('resume', 'rounds', 'skip_final_test')} |
| config.update(teacher_sha256=teacher_hashes(policy.teacher), source_sha256=source_hashes()) |
| config.update(model_config=policy.model.config.dict(), initialization=initialization, |
| parameters=sum(p.numel() for p in policy.model.parameters()), |
| opponent_sha256=teacher_hashes(opponent) if opponent else None) |
| print(json.dumps({'model': config['model_config'], 'parameters': config['parameters'], |
| 'initialization': initialization, |
| 'opponent': str(args.opponent_dir or args.teacher_dir)}), flush=True) |
| if not args.resume: |
| export(policy, output / 'initial.pt', {'initialization': initialization}) |
| write_json(output / 'launch.json', config) |
| optimizer = torch.optim.AdamW(policy.model.parameters(), lr=args.learning_rate, weight_decay=1e-4) |
| progress = {'offline_epoch': 0, 'offline_batch': 0, 'round': 0, 'updates': 0, 'games': 0, |
| 'best_score': None, 'best_stage': 'baseline', 'stage': 'offline', |
| 'best_modes': {}, 'validation_rollbacks': 0, 'kl_rollbacks': 0} |
| metrics, league = [], [] |
| accepted = None |
| resume_path = output / 'resume.pt' |
| if args.resume: |
| saved = torch.load(resume_path, weights_only=True) |
| if saved['config'] != config: |
| raise ValueError('resume configuration/source/teacher mismatch') |
| policy.model.load_state_dict(saved['model']) |
| optimizer.load_state_dict(saved['optimizer']) |
| progress, metrics, league = saved['progress'], saved['metrics'], saved['league'] |
| accepted = saved.get('accepted') |
| rng.setstate(saved['random_state']) |
| torch.set_rng_state(saved['torch_rng']) |
| if args.rounds < progress['round']: |
| raise ValueError('rounds smaller than committed progress') |
| reconcile_jsonl(output / 'train_games.jsonl', progress['games']) |
| manifest = prepare(args.data, output / 'dataset', policy, args.seed, args.max_history_games, |
| candidate_cap=args.candidate_cap, holdout=args.holdout) |
| print(json.dumps({'dataset': manifest['counts'], 'valid_games': manifest['valid_games'], |
| 'rejected_games': len(manifest['rejected'])}), flush=True) |
| synthetic_path = output / 'warmup.pt' |
| if synthetic_path.exists(): |
| synthetic = torch.load(synthetic_path, weights_only=True) |
| else: |
| synthetic, records = warmup(policy, args.warmup_games, args.seed, rng) |
| synthetic = [compact(cap_candidates(sample, args.candidate_cap)) for sample in synthetic] |
| atomic_save(synthetic, synthetic_path) |
| write_json(output / 'warmup_games.json', records) |
| warmup_hash = digest(synthetic_path) |
| if args.resume and (saved['dataset_fingerprint'] != manifest['fingerprint'] or |
| saved['warmup_sha256'] != warmup_hash): |
| raise ValueError('resume dataset/warmup integrity mismatch') |
| anchor_rng = random.Random(f'{args.seed}:teacher-anchors') |
| anchor_pool = anchor_rng.sample(synthetic, min(args.teacher_anchor_states, len(synthetic))) |
| if not anchor_pool: |
| for sample in iter_samples(output / 'dataset', manifest, 'train'): |
| anchor_pool.append(sample) |
| if len(anchor_pool) == args.teacher_anchor_states: |
| break |
| modes = ('fixed', 'bidding') if args.mode == 'mixed' else (args.mode,) |
|
|
| def checkpoint(status): |
| atomic_save({'config': config, 'model': policy.model.state_dict(), |
| 'optimizer': optimizer.state_dict(), 'progress': progress, 'metrics': metrics[-200:], |
| 'league': league, 'random_state': rng.getstate(), 'torch_rng': torch.get_rng_state(), |
| 'accepted': accepted, |
| 'dataset_fingerprint': manifest['fingerprint'], 'warmup_sha256': warmup_hash}, |
| resume_path) |
| export(policy, output / 'latest.pt', progress) |
| report = {'status': status, 'device': str(policy.device), 'config': config, |
| 'progress': progress, 'dataset': {'fingerprint': manifest['fingerprint'], |
| 'counts': manifest['counts'], 'valid_games': manifest['valid_games'], |
| 'rejected_games': len(manifest['rejected'])}, 'metrics': metrics[-200:], |
| 'limitation': 'holdout imitation loss is not playing strength; best selected on paired validation'} |
| write_json(output / 'report.json', report) |
| return report |
|
|
| def validate(stage): |
| nonlocal accepted |
| result = evaluate(policy, args.validation_deals, args.validation_seed, modes, opponent) |
| write_json(output / f'validation-{stage}.json', result) |
| score = result['selection_score'] |
| mode_scores = {key: value['resnet']['mean_score'] for key, value in result['summary'].items()} |
| improves = progress['best_score'] is None or ( |
| score > progress['best_score'] and (not args.rollback_tolerance or all( |
| value >= progress['best_modes'].get(key, value) - args.rollback_tolerance |
| for key, value in mode_scores.items()))) |
| rolled_back = False |
| if improves: |
| progress.update(best_score=score, best_stage=stage) |
| progress['best_modes'] = mode_scores |
| export(policy, output / 'best.pt', progress) |
| accepted = {'model': cpu_tree(policy.model.state_dict()), |
| 'optimizer': cpu_tree(optimizer.state_dict())} |
| elif args.rollback_tolerance and accepted is not None and ( |
| score < progress['best_score'] - args.rollback_tolerance or any( |
| value < progress['best_modes'].get(key, value) - args.rollback_tolerance |
| for key, value in mode_scores.items())): |
| policy.model.load_state_dict(accepted['model']) |
| rates = [group['lr'] for group in optimizer.param_groups] |
| optimizer.load_state_dict(accepted['optimizer']) |
| for group, rate in zip(optimizer.param_groups, rates): |
| group['lr'] = max(rate / 2, 1e-6) |
| progress['validation_rollbacks'] += 1 |
| league.clear() |
| rolled_back = True |
| metrics.append({'stage': stage, 'validation_score': score, |
| 'best_score': progress['best_score'], 'rolled_back': rolled_back}) |
| print(json.dumps(metrics[-1]), flush=True) |
|
|
| if not resume_path.exists(): |
| validate('baseline') |
| checkpoint('running') |
| try: |
| while progress['offline_epoch'] < args.offline_epochs: |
| policy.model.train() |
| |
| epoch_rng = random.Random(f'{args.seed}:offline:{progress["offline_epoch"]}') |
| training_samples = iter_samples(output / 'dataset', manifest, 'train', epoch_rng) |
| shuffled = list(synthetic) |
| def epoch_batches(): |
| yield from batches(training_samples, args.batch_states, args.max_batch_actions) |
| epoch_rng.shuffle(shuffled) |
| yield from batches(shuffled, args.batch_states, args.max_batch_actions) |
| started = time.monotonic() |
| for batch_index, batch in enumerate(epoch_batches()): |
| if batch_index < progress.get('offline_batch', 0): |
| continue |
| stats = update(policy.model, optimizer, batch, policy.device, |
| anchor_weight=args.offline_anchor_weight, |
| history_weight=args.history_weight, distill_weight=args.distill_weight) |
| progress['updates'] += 1 |
| progress['offline_batch'] = batch_index + 1 |
| if progress['offline_batch'] % args.checkpoint_interval == 0: |
| metrics.append({'epoch': progress['offline_epoch'] + 1, |
| 'batch': progress['offline_batch'], 'offline': stats, |
| 'seconds': time.monotonic() - started}) |
| print(json.dumps(metrics[-1]), flush=True) |
| checkpoint('running') |
| progress['offline_epoch'] += 1 |
| progress['offline_batch'] = 0 |
| metrics.append({'epoch': progress['offline_epoch'], |
| 'offline_validation': holdout_loss(policy, output / 'dataset', manifest, |
| 'validation', args)}) |
| validate(f'offline-{progress["offline_epoch"]}') |
| checkpoint('running') |
| if progress['stage'] == 'offline': |
| if args.ppo_start == 'best': |
| selected = QOJPolicy(output / 'best.pt', device=str(policy.device), teacher=policy.teacher) |
| policy.model.load_state_dict(selected.model.state_dict()) |
| optimizer = torch.optim.AdamW(policy.model.parameters(), lr=args.ppo_learning_rate, |
| weight_decay=1e-4) |
| if args.rollback_tolerance: |
| |
| accepted = {'model': cpu_tree(policy.model.state_dict()), |
| 'optimizer': cpu_tree(optimizer.state_dict())} |
| progress['stage'] = 'ppo' |
| checkpoint('running') |
| while progress['round'] < args.rounds: |
| round_index = progress['round'] |
| started = time.monotonic() |
| before = ({'model': cpu_tree(policy.model.state_dict()), |
| 'optimizer': cpu_tree(optimizer.state_dict())} if args.max_ppo_kl else None) |
| samples, records = collect(policy, args.games_per_round, |
| round_index * args.games_per_round, |
| args.seed + args.warmup_games, rng, |
| modes[round_index % len(modes)], league, opponent, |
| targeted=args.targeted_league) |
| stop, update_stats = False, [] |
| stats = {'samples': 0, 'actions': 0, 'loss': None, 'approx_kl': 0.} |
| for epoch in range(args.ppo_epochs): |
| rng.shuffle(samples) |
| for batch in batches(samples, args.batch_states, args.max_batch_actions): |
| stats = update(policy.model, optimizer, batch, policy.device, |
| ppo=True, anchor_weight=args.anchor_weight, |
| history_weight=args.history_weight) |
| progress['updates'] += 1 |
| update_stats.append(stats) |
| if stats['approx_kl'] > .04: |
| stop = True |
| break |
| if stop: |
| break |
| |
| if anchor_pool: |
| anchor_batch = rng.sample(anchor_pool, min(args.batch_states, len(anchor_pool))) |
| update(policy.model, optimizer, anchor_batch, policy.device, |
| anchor_weight=args.offline_anchor_weight, |
| history_weight=args.history_weight, distill_weight=args.distill_weight) |
| progress['updates'] += 1 |
| exact_kl = (rollout_kl(policy, samples, args.batch_states, args.max_batch_actions) |
| if args.max_ppo_kl else None) |
| kl_rejected = exact_kl is not None and exact_kl > args.max_ppo_kl |
| if kl_rejected: |
| policy.model.load_state_dict(before['model']) |
| optimizer.load_state_dict(before['optimizer']) |
| for group in optimizer.param_groups: |
| group['lr'] = max(group['lr'] / 2, 1e-6) |
| progress['kl_rollbacks'] += 1 |
| progress['round'] += 1 |
| progress['games'] += len(records) |
| if update_stats: |
| total = sum(item['samples'] for item in update_stats) |
| stats = {key: sum(item[key] * item['samples'] for item in update_stats) / total |
| for key in update_stats[0] if key not in ('samples', 'actions')} |
| stats.update(samples=total, actions=sum(item['actions'] for item in update_stats)) |
| metrics.append({'round': progress['round'], 'ppo': stats, |
| 'samples': len(samples), 'early_kl_stop': stop, |
| 'exact_rollout_kl': exact_kl, 'kl_rejected': kl_rejected, |
| 'seconds': time.monotonic() - started}) |
| print(json.dumps(metrics[-1]), flush=True) |
| for record in records: |
| append_jsonl(output / 'train_games.jsonl', record) |
| if progress['round'] % args.validation_interval == 0 or progress['round'] == args.rounds: |
| validate(f'round-{progress["round"]}') |
| league.append(cpu_tree(policy.model.state_dict())) |
| league = league[-args.league_size:] |
| checkpoint('running') |
| report = checkpoint('testing' if not args.skip_final_test else 'complete') |
| if not args.skip_final_test: |
| selected = QOJPolicy(output / 'best.pt', device=str(policy.device), teacher=policy.teacher) |
| final = evaluate(selected, args.test_deals, args.test_seed, modes, opponent) |
| write_json(output / 'test.json', final) |
| report['test'] = {key: final[key] for key in ('games', 'selection_score', 'summary')} |
| report['offline_test'] = holdout_loss(selected, output / 'dataset', manifest, 'test', args) |
| report['status'] = 'complete' |
| report['test']['positive_score_supported'] = final['positive_score_supported'] |
| write_json(output / 'report.json', report) |
| return report |
| except KeyboardInterrupt: |
| |
| saved = torch.load(resume_path, weights_only=True) |
| policy.model.load_state_dict(saved['model']) |
| optimizer.load_state_dict(saved['optimizer']) |
| progress, metrics, league = saved['progress'], saved['metrics'], saved['league'] |
| accepted = saved.get('accepted') |
| rng.setstate(saved['random_state']) |
| torch.set_rng_state(saved['torch_rng']) |
| reconcile_jsonl(output / 'train_games.jsonl', progress['games']) |
| return checkpoint('interrupted') |
|
|
|
|
| def main(): |
| try: |
| report = train(make_parser().parse_args()) |
| except (ValueError, RuntimeError, OSError) as exc: |
| raise SystemExit(str(exc)) from exc |
| print(json.dumps({key: report[key] for key in ('status', 'device', 'progress')}, indent=2)) |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
