"""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() # Reconstruct a deterministic epoch order after a partial-epoch resume. 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 optimizer must belong to the current PPO stage. 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 # Teacher anchors never reuse PPO ratios or old advantages. 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: # Roll back an incomplete epoch/round, including optimizer and sampling RNG. 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()