import json import tempfile import unittest from pathlib import Path import torch from qojbot import rules as R from qojbot.adapter import observe, Strategy from qojbot.game import deal from qojbot.history import replay_game from QOJDoudizhuModel.dataset import belief_labels, cap_candidates, offline_samples, prepare from QOJDoudizhuModel.features import collate, encode_bid, encode_public, legal_bids from QOJDoudizhuModel.learning import Actor, batches, collect, losses, terminal_targets, update from QOJDoudizhuModel.network import ModelConfig, QOJNet, distribution_target from QOJDoudizhuModel.policy import QOJPolicy from QOJDoudizhuModel.storage import export ROOT = Path(__file__).resolve().parents[2] class ModelTests(unittest.TestCase): @classmethod def setUpClass(cls): torch.set_num_threads(1) cls.record = json.loads((ROOT / 'reference/game-8570.json').read_text()) cls.policy = QOJPolicy(config=ModelConfig(width=32, rank_layers=1, history_layers=1)) cls.state, cls.obs, cls.move, _ = replay_game(cls.record)[0] def test_public_features_ignore_privileged_fields_and_handle_empty_history(self): moves = R.legal_moves(self.obs.hand, self.obs.last) a = encode_public(self.obs, self.state, moves) state = {**self.state, 'hands': [[1], [2], [3]], 'result': {'deltas': [999]}} b = encode_public(self.obs, state, moves) for key in ('ranks', 'scalars', 'history', 'mask', 'actions'): torch.testing.assert_close(a[key], b[key]) empty = dict(a, history=torch.zeros_like(a['history']), mask=torch.zeros_like(a['mask'])) output = self.policy.model(collate([empty], 'cpu')) self.assertTrue(torch.isfinite(output['correction']).all()) def test_prior_zero_initialization_exactly_matches_baseline(self): for state, obs, move, _ in replay_game(self.record)[:8]: actual = self.policy.choose(obs, state) expected = self.policy.teacher.choose(obs, state) self.assertEqual(actual, expected) def test_action_chunking_and_state_batching_are_invariant(self): _, sample = self.policy.prepare(self.obs, self.state) batch = collate([sample, sample], 'cpu') with torch.no_grad(): a = self.policy.model(batch, 1) b = self.policy.model(batch, 512) for key in ('correction', 'score_logits', 'win_logits'): torch.testing.assert_close(a[key], b[key], atol=2e-6, rtol=2e-5) def test_same_cards_different_interpretations_are_distinct(self): hand = R.parse('333444555666') moves = R.legal_moves(hand) codes = {} for move in moves: codes.setdefault(move.code, []).append(move) duplicate = next(group for group in codes.values() if len(group) > 1) from QOJDoudizhuModel.features import action_features features = action_features(hand, duplicate) self.assertFalse(torch.equal(features[0], features[1])) def test_bidding_masks_and_forced_bid(self): state = {'seat': 1, 'turn': 1, 'hand': list(range(17)), 'bid': 2, 'log': []} self.assertEqual(legal_bids(state), [0, 3]) self.assertEqual(legal_bids(dict(state, must_bid=True)), [3]) bids, sample = self.policy.prepare_bid(state) self.assertIn(self.policy.choose_bid(state), bids) self.assertEqual(sample['role'], 3) self.assertTrue(torch.isfinite(self.policy.model(collate([sample], 'cpu'))['value']).all()) with self.assertRaises(ValueError): encode_bid(dict(state, turn=2), bids) def test_signed_distribution_normalized_and_negative_scores_representable(self): support = torch.linspace(-5, 5, 51) targets = distribution_target(torch.tensor([-96., -3., 0., 3., 96.]), support) torch.testing.assert_close(targets.sum(1), torch.ones(5)) self.assertTrue((targets @ support)[:2].lt(0).all()) rewards, scaled = terminal_targets([{'reward': -12}, {'reward': 6}], 'cpu') torch.testing.assert_close(scaled, torch.tensor([-2., 1.])) def test_belief_respects_known_bottom_and_unseen_rank_totals(self): samples = [] for state, obs, _, _ in replay_game(self.record): if obs.seat != obs.landlord: _, sample = self.policy.prepare(obs, state) samples.append(sample) if len(samples) == 3: break output = self.policy.model(collate(samples, 'cpu')) for index, sample in enumerate(samples): unseen = sample['ranks'][:, 1] * 4 known = sample['ranks'][:, 2] * 4 torch.testing.assert_close(output['belief_counts'][index].sum(-1), unseen) self.assertTrue((output['belief_counts'][index, :, 0] >= known - 1e-6).all()) def test_offline_samples_real_scores_and_beliefs(self): samples = list(offline_samples(self.record, self.policy, include_auto=True)) self.assertEqual(len(samples), len(replay_game(self.record))) for sample, (_, obs, _, _) in zip(samples, replay_game(self.record)): self.assertEqual(sample['reward'], self.record['result']['deltas'][obs.seat]) torch.testing.assert_close(sample['belief_target'].sum(1), torch.tensor(R.counts(obs.unseen), dtype=torch.float32)) capped = cap_candidates(samples[0], 2) self.assertEqual(len(capped['actions']), min(2, len(samples[0]['actions']))) def test_update_and_export_roundtrip(self): policy = QOJPolicy(config=ModelConfig(width=32, rank_layers=1, history_layers=1), teacher=self.policy.teacher) samples = list(offline_samples(self.record, policy, include_auto=True))[:2] optimizer = torch.optim.AdamW(policy.model.parameters(), lr=1e-3) stats = update(policy.model, optimizer, samples, 'cpu') self.assertGreater(stats['actions'], 0) with tempfile.TemporaryDirectory() as directory: path = Path(directory) / 'best.pt' export(policy, path) loaded = QOJPolicy(path, teacher=policy.teacher) self.assertEqual(policy.choose(self.obs, self.state), loaded.choose(self.obs, self.state)) strategy = Strategy(policy='qoj', checkpoint_dir=path) action, extra = strategy.decide(self.state) self.assertNotIn('fallback', strategy.diagnostics) self.assertIn(action, ('play', 'pass')) def test_history_human_action_and_outcome_do_not_drive_strong_distillation(self): samples = list(offline_samples(self.record, self.policy, include_auto=True))[:3] modified = [dict(s, selected=(s['selected']+1) % len(s['actions']), reward=-1000 * s['reward']) for s in samples] a, stats = losses(self.policy.model, samples, 'cpu', history_weight=0, distill_weight=.1) b, _ = losses(self.policy.model, modified, 'cpu', history_weight=0, distill_weight=.1) torch.testing.assert_close(a, b, atol=0, rtol=0) self.assertEqual(stats['value'], 0.) def test_resnet_prior_roundtrip_and_exact_initial_moves(self): from qojbot.douzero import DouZeroPolicy, DEFAULT_CHECKPOINT_DIR teacher = DouZeroPolicy(DEFAULT_CHECKPOINT_DIR) policy = QOJPolicy(teacher=teacher, config=ModelConfig(width=32, rank_layers=1, history_layers=1)) for state, obs, _, _ in replay_game(self.record)[:6]: self.assertEqual(policy.choose(obs, state), teacher.choose(obs, state)) with tempfile.TemporaryDirectory() as directory: path = Path(directory) / 'best.pt' export(policy, path) loaded = QOJPolicy(path) self.assertEqual(loaded.teacher.model_type, 'resnet_best') self.assertEqual(loaded.choose(self.obs, self.state), teacher.choose(self.obs, self.state)) def test_exact_rollout_kl_detects_policy_change(self): from QOJDoudizhuModel.learning import rollout_kl samples = list(offline_samples(self.record, self.policy, include_auto=True))[:3] for sample in samples: logits, _ = self.policy.predict(sample) sample['old_log_probs'] = logits.log_softmax(0) self.assertLess(rollout_kl(self.policy, samples), 1e-6) changed = [dict(s, old_log_probs=torch.zeros_like(s['base']).log_softmax(0)) for s in samples] self.assertGreater(rollout_kl(self.policy, changed), .01) def test_ppo_ratio_is_one_before_update_and_bids_collected(self): import random policy = QOJPolicy(config=ModelConfig(width=32, rank_layers=1, history_layers=1), teacher=self.policy.teacher) samples, records = collect(policy, 1, 4, 1000, random.Random(4), mode='bidding') self.assertTrue(any(sample['role'] == 3 for sample in samples)) self.assertEqual(sum(records[0]['deltas']), 0) _, stats = losses(policy.model, samples[:8], 'cpu', ppo=True) self.assertAlmostEqual(stats['approx_kl'], 0., places=5) self.assertTrue(all('belief_target' in sample for sample in samples if sample['role'] != 3)) def test_batched_losses_match_individual_states_with_padding(self): samples = list(offline_samples(self.record, self.policy, include_auto=True))[:4] samples = [cap_candidates(sample, cap) for sample, cap in zip(samples, (2, 3, 5, 8))] samples[0] = dict(samples[0], actions=samples[0]['actions'][:1], base=samples[0]['base'][:1], selected=0) for index, sample in enumerate(samples): logits, _ = self.policy.predict(sample) sample.update(old_logp=float(logits.log_softmax(0)[sample['selected']]) + .01 * index, advantage=(-1) ** index * .7) for ppo in (False, True): whole, _ = losses(self.policy.model, samples, 'cpu', ppo=ppo) individual = torch.stack([losses(self.policy.model, [sample], 'cpu', ppo=ppo)[0] for sample in samples]).mean() torch.testing.assert_close(whole, individual, atol=2e-7, rtol=2e-6) params = tuple(self.policy.model.parameters()) a = torch.autograd.grad(whole, params, allow_unused=True) b = torch.autograd.grad(individual, params, allow_unused=True) for actual, expected in zip(a, b): if actual is not None: self.assertTrue(torch.isfinite(actual).all()) torch.testing.assert_close(actual, expected, atol=3e-7, rtol=5e-5) def test_sharded_snapshot_groups_and_integrity(self): with tempfile.TemporaryDirectory() as directory: data, cache = Path(directory) / 'data', Path(directory) / 'cache' data.mkdir() for index in range(2): record = dict(self.record, id=index + 1, match={'id': index + 1}) (data / f'game-{index + 1}.json').write_text(json.dumps(record)) manifest = prepare(data, cache, self.policy, candidate_cap=8) self.assertEqual(len({s['split'] for s in manifest['assignments'].values()}), 1) self.assertEqual(len({s['group'] for s in manifest['assignments'].values()}), 1) self.assertEqual(prepare(data, cache, self.policy, candidate_cap=8)['fingerprint'], manifest['fingerprint']) (cache / manifest['shards'][0]['file']).write_bytes(b'bad') with self.assertRaisesRegex(ValueError, 'integrity'): prepare(data, cache, self.policy, candidate_cap=8) if __name__ == '__main__': unittest.main()