| 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() |
|
|
