first prototype

wizard_ai/tests/test_ai.py

@@ -0,0 +1,80 @@
+import json
+import random
+
+import pytest
+
+from wizard_game import GameOptions, WizardGame
+
+from wizard_ai.agents import HeuristicAgent, ModelAgent, RandomAgent, make_agent
+from wizard_ai.model import FEATURE_SIZE, ActionSpace, UniformPolicy, encode_observation, load_policy
+from wizard_ai.training import SelfPlayTrainer, TrainingConfig, WizardEnv, evaluate
+
+
+@pytest.mark.parametrize("num_players", [3, 4, 5, 6])
+def test_encoding_has_fixed_size_and_masks_match(num_players):
+    game = WizardGame(num_players, GameOptions(hidden_bids=True), seed=num_players)
+    rng = random.Random(0)
+    while not game.is_finished:
+        obs = game.observation(game.current_player)
+        assert len(encode_observation(obs)) == FEATURE_SIZE
+        mask = ActionSpace.mask(obs)
+        legal = {ActionSpace.index(a) for a in obs.legal_actions}
+        assert {i for i, allowed in enumerate(mask) if allowed} == legal
+        game.apply(obs.player, rng.choice(obs.legal_actions))
+
+
+def test_action_space_round_trip():
+    for index in range(ActionSpace.SIZE):
+        assert ActionSpace.index(ActionSpace.action(index)) == index
+
+
+@pytest.mark.parametrize("agent", [HeuristicAgent(), RandomAgent(1), ModelAgent(UniformPolicy(1))])
+def test_agents_play_legal_games(agent):
+    game = WizardGame(4, GameOptions(not_equal=True), seed=3)
+    while not game.is_finished:
+        obs = game.observation(game.current_player)
+        action = agent.act(obs)
+        assert action in obs.legal_actions
+        game.apply(obs.player, action)
+
+
+def test_env_rewards_add_up_to_scores():
+    env = WizardEnv(3, seed=5)
+    obs = env.reset()
+    totals = [0.0] * env.num_players
+    rng = random.Random(5)
+    while obs is not None:
+        result = env.step(rng.choice(obs.legal_actions))
+        totals = [t + r for t, r in zip(totals, result.rewards)]
+        obs = result.observation
+    assert [round(t * 100) for t in totals] == env.game.scores
+
+
+def test_trainer_runs_and_checkpoints(tmp_path):
+    config = TrainingConfig(episodes=4, player_counts=(3, 4), episodes_per_update=2,
+                            checkpoint_dir=tmp_path / "ckpt", seed=1)
+    trainer = SelfPlayTrainer(UniformPolicy(1), config)
+    trajectories, scores = trainer.play_episode()
+    assert len(trajectories) == len(scores)
+    assert all(t[-1].done for t in trajectories)
+    assert sum(sum(tr.reward for tr in t) for t in trajectories) == pytest.approx(sum(scores) / 100)
+    trainer.train()
+    assert json.loads((tmp_path / "ckpt" / "policy.json").read_text())["type"] == "uniform"
+    assert isinstance(load_policy(tmp_path / "ckpt"), UniformPolicy)
+
+
+def test_trainer_against_heuristic_opponents():
+    config = TrainingConfig(episodes=2, player_counts=(5,), opponents="heuristic", seed=2)
+    trajectories, _ = SelfPlayTrainer(UniformPolicy(2), config).play_episode()
+    assert len(trajectories) == 1
+
+
+def test_make_agent_falls_back_without_checkpoint(tmp_path):
+    assert isinstance(make_agent("model", tmp_path / "missing"), HeuristicAgent)
+    UniformPolicy().save(tmp_path / "ckpt")
+    assert isinstance(make_agent("model", tmp_path / "ckpt"), ModelAgent)
+
+
+def test_heuristic_beats_random():
+    result = evaluate([HeuristicAgent(), RandomAgent(1), RandomAgent(2), RandomAgent(3)], games=40, seed=0)
+    assert result.mean_scores[0] > max(result.mean_scores[1:])