first prototype

wizard_ai/.dockerignore

@@ -0,0 +1,5 @@
+**/__pycache__
+**/*.pyc
+.venv
+.pytest_cache
+tests

wizard_ai/.gitignore

@@ -0,0 +1,6 @@
+__pycache__/
+*.pyc
+.venv/
+.pytest_cache/
+models/*
+!models/.gitkeep

wizard_ai/Dockerfile

@@ -0,0 +1,27 @@
+# Build with the game logic as additional build context, e.g.
+#   docker build --build-context game_logic=../wizard_server/game_logic .
+# (docker compose does this automatically, see compose.yaml)
+FROM python:3.13-slim
+
+ENV PYTHONDONTWRITEBYTECODE=1 \
+    PYTHONUNBUFFERED=1 \
+    PIP_NO_CACHE_DIR=1 \
+    PIP_DISABLE_PIP_VERSION_CHECK=1
+
+WORKDIR /srv
+
+COPY requirements.txt .
+RUN pip install -r requirements.txt
+
+# Shared rules engine from the game server.
+COPY --from=game_logic . /tmp/game_logic
+RUN pip install /tmp/game_logic && rm -rf /tmp/game_logic
+
+COPY wizard_ai ./wizard_ai
+COPY models ./models
+
+RUN useradd --system --no-create-home wizard
+USER wizard
+
+ENV MODEL_PATH=/srv/models/latest
+CMD ["python", "-m", "wizard_ai", "serve"]

wizard_ai/README.md

@@ -0,0 +1,27 @@
+# wizard_ai
+
+AI player for the Wizard server. The actual RL model is still to be implemented; this package
+provides everything around it.
+
+- `agents/` — `Agent` interface (`act(observation) -> action`) plus `RandomAgent`,
+  `HeuristicAgent` (baseline / fallback) and `ModelAgent` (wraps a `Policy`)
+- `model/encoding.py` — observation → fixed-size feature vector (`FEATURE_SIZE`), and the fixed
+  action space (21 bids + 60 cards + 4 trump suits) with legal-action masks; works for 3–6 players
+- `model/policy.py` — **the place to plug in the model**: subclass `Policy`, implement
+  `select`, `update`, `save_weights`/`load`, register with `@register_policy`.
+  `UniformPolicy` is a placeholder so the pipeline runs end to end.
+- `training/env.py` — turn-based multi-agent environment around `wizard_game`
+  (reward = score change per round / 100)
+- `training/self_play.py` — self-play loop collecting per-player trajectories and calling
+  `Policy.update`, with checkpointing
+- `training/evaluate.py` — play agents against each other (mean score, win rate)
+- `service/client.py` — production: connects to `ws://wizard_server:8083/internal/ai`,
+  answers `decide` requests, reconnects automatically
+
+```bash
+pip install -r requirements-dev.txt
+python -m wizard_ai train --episodes 5000 --out models/latest
+python -m wizard_ai evaluate --agents model,heuristic,random,random --games 200
+python -m wizard_ai serve   # env: SERVER_URL, AI_SHARED_SECRET, AGENT, MODEL_PATH
+pytest
+```

wizard_ai/pytest.ini

@@ -0,0 +1,2 @@
+[pytest]
+testpaths = tests

wizard_ai/requirements-dev.txt

@@ -0,0 +1,3 @@
+-r requirements.txt
+-e ../wizard_server/game_logic
+pytest>=8

wizard_ai/requirements.txt

@@ -0,0 +1 @@
+websockets>=14,<18

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:])

wizard_ai/wizard_ai/__init__.py

@@ -0,0 +1,7 @@
+"""AI player for the Wizard game server.
+
+- ``agents``:   the ``Agent`` interface and implementations
+- ``model``:    observation/action encoding and the (RL) policy interface
+- ``training``: self-play environment, training loop and evaluation
+- ``service``:  websocket client that serves decisions to the game server
+"""

wizard_ai/wizard_ai/__main__.py

@@ -0,0 +1,83 @@
+"""Command line entry point.
+
+    python -m wizard_ai serve                      # production: connect to the game server
+    python -m wizard_ai train --episodes 5000      # self-play training
+    python -m wizard_ai evaluate --agents model,heuristic,random,random
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+from pathlib import Path
+
+from wizard_game import GameOptions
+
+from .agents import AGENT_NAMES, make_agent
+from .config import Config
+
+
+def main() -> None:
+    config = Config()
+    parser = argparse.ArgumentParser(prog="wizard_ai")
+    parser.add_argument("--log-level", default=config.log_level)
+    commands = parser.add_subparsers(dest="command", required=True)
+
+    serve = commands.add_parser("serve", help="serve AI decisions to the game server")
+    serve.add_argument("--server-url", default=config.server_url)
+    serve.add_argument("--agent", default=config.agent, choices=AGENT_NAMES)
+    serve.add_argument("--model-path", default=config.model_path)
+
+    train = commands.add_parser("train", help="train the policy with self-play")
+    train.add_argument("--episodes", type=int, default=1000)
+    train.add_argument("--players", default="3,4,5,6", help="comma separated player counts")
+    train.add_argument("--opponents", default="self", choices=("self", "heuristic", "random"))
+    train.add_argument("--not-equal", action="store_true")
+    train.add_argument("--hidden-bids", action="store_true")
+    train.add_argument("--resume", action="store_true", help="continue from --out if it exists")
+    train.add_argument("--out", default=config.model_path, help="checkpoint directory")
+    train.add_argument("--seed", type=int)
+
+    evaluation = commands.add_parser("evaluate", help="let agents play against each other")
+    evaluation.add_argument("--agents", default="model,heuristic,random,random",
+                            help=f"comma separated, 3-6 of {AGENT_NAMES}")
+    evaluation.add_argument("--games", type=int, default=200)
+    evaluation.add_argument("--model-path", default=config.model_path)
+    evaluation.add_argument("--seed", type=int)
+
+    args = parser.parse_args()
+    logging.basicConfig(level=args.log_level.upper(), format="%(asctime)s %(levelname)s %(name)s: %(message)s")
+
+    if args.command == "serve":
+        from .service import AIService
+
+        agent = make_agent(args.agent, args.model_path)
+        asyncio.run(AIService(args.server_url, agent, config.secret).run_forever())
+
+    elif args.command == "train":
+        from .model import UniformPolicy, checkpoint_exists, load_policy
+        from .training import SelfPlayTrainer, TrainingConfig
+
+        out = Path(args.out)
+        # TODO: replace UniformPolicy with the real model once it is implemented.
+        policy = load_policy(out) if args.resume and checkpoint_exists(out) else UniformPolicy(args.seed)
+        training_config = TrainingConfig(
+            episodes=args.episodes,
+            player_counts=tuple(int(n) for n in args.players.split(",")),
+            options=GameOptions(not_equal=args.not_equal, hidden_bids=args.hidden_bids),
+            opponents=args.opponents,
+            checkpoint_dir=out,
+            seed=args.seed,
+        )
+        SelfPlayTrainer(policy, training_config).train()
+
+    elif args.command == "evaluate":
+        from .training import evaluate
+
+        agents = [make_agent(name.strip(), args.model_path) for name in args.agents.split(",")]
+        print(evaluate(agents, args.games, seed=args.seed))
+
+
+if __name__ == "__main__":
+    main()

wizard_ai/wizard_ai/agents/__init__.py

@@ -0,0 +1,36 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+from ..model.policy import checkpoint_exists, load_policy
+from .base import Agent
+from .heuristic_agent import HeuristicAgent
+from .model_agent import ModelAgent
+from .random_agent import RandomAgent
+
+logger = logging.getLogger(__name__)
+
+AGENT_NAMES = ("model", "heuristic", "random")
+
+
+def make_agent(name: str, model_path: str | Path | None = None) -> Agent:
+    """Create an agent by name.
+
+    ``model`` loads the policy checkpoint at ``model_path``. As long as no
+    trained model exists, it falls back to the heuristic agent.
+    """
+    if name == "random":
+        return RandomAgent()
+    if name == "heuristic":
+        return HeuristicAgent()
+    if name == "model":
+        if checkpoint_exists(model_path):
+            logger.info("Loading policy from %s", model_path)
+            return ModelAgent(load_policy(model_path))
+        logger.warning("No model checkpoint at %s, using the heuristic agent", model_path)
+        return HeuristicAgent()
+    raise ValueError(f"Unknown agent {name!r}, expected one of {AGENT_NAMES}")
+
+
+__all__ = ["AGENT_NAMES", "Agent", "HeuristicAgent", "ModelAgent", "RandomAgent", "make_agent"]

wizard_ai/wizard_ai/agents/base.py

@@ -0,0 +1,26 @@
+"""The interface every AI player implements."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+
+from wizard_game import Action, Observation
+
+
+class Agent(ABC):
+    """Chooses actions from observations.
+
+    The same interface is used during training (observations come straight
+    from ``wizard_game.WizardGame``) and in production (observations arrive
+    from the game server via websocket), so a trained agent can be deployed
+    without changes.
+    """
+
+    name: str = "agent"
+
+    @abstractmethod
+    def act(self, observation: Observation) -> Action:
+        """Return one of ``observation.legal_actions``."""
+
+    def reset(self) -> None:  # noqa: B027 - optional hook
+        """Called before a new game starts (training only)."""

wizard_ai/wizard_ai/agents/heuristic_agent.py

@@ -0,0 +1,96 @@
+"""A simple rule based player.
+
+It serves as a sparring partner / benchmark for reinforcement learning and
+as the fallback in production while no trained model is available.
+"""
+
+from __future__ import annotations
+
+from collections import Counter
+
+from wizard_game import SUITS, Action, Bid, Card, ChooseTrump, Observation, Phase, PlayCard, Suit
+from wizard_game.rules import trick_winner
+
+from .base import Agent
+
+
+class HeuristicAgent(Agent):
+    name = "heuristic"
+
+    def act(self, observation: Observation) -> Action:
+        if observation.phase is Phase.TRUMP_SELECTION:
+            return ChooseTrump(self._best_trump(observation.hand))
+        if observation.phase is Phase.BIDDING:
+            return self._bid(observation)
+        return self._play(observation)
+
+    # --- Trump & bidding ----------------------------------------------------
+
+    @staticmethod
+    def _best_trump(hand: tuple[Card, ...]) -> Suit:
+        strength = Counter({suit: 0.0 for suit in SUITS})
+        for card in hand:
+            if card.is_number and card.suit is not None:
+                strength[card.suit] += 1 + card.rank / 13
+        return strength.most_common(1)[0][0]
+
+    @staticmethod
+    def _expected_tricks(observation: Observation) -> float:
+        trump = observation.trump_suit
+        opponents = observation.num_players - 1
+        expected = 0.0
+        for card in observation.hand:
+            if card.is_wizard:
+                expected += 0.95
+            elif card.is_number and card.suit is trump:
+                expected += 0.35 + 0.05 * card.rank
+            elif card.is_number and card.rank >= 12:
+                # High side suit cards are good but get trumped more often
+                # with more opponents.
+                expected += max(0.1, 0.75 - 0.1 * opponents) * (card.rank - 10) / 3
+        return expected
+
+    def _bid(self, observation: Observation) -> Action:
+        target = round(self._expected_tricks(observation))
+        bids = [action for action in observation.legal_actions if isinstance(action, Bid)]
+        return min(bids, key=lambda bid: (abs(bid.value - target), bid.value))
+
+    # --- Playing ------------------------------------------------------------
+
+    def _play(self, observation: Observation) -> Action:
+        candidates = [a for a in observation.legal_actions if isinstance(a, PlayCard)]
+        me = observation.player
+        bid = observation.bids[me] or 0
+        wants_tricks = observation.tricks_won[me] < bid
+        trick = [played.card for played in observation.current_trick]
+        is_last = len(trick) == observation.num_players - 1
+
+        winning = [a for a in candidates if self._wins(trick, a.card, observation.trump_suit)]
+        losing = [a for a in candidates if a not in winning]
+
+        if wants_tricks:
+            if winning:
+                # Last to play: win as cheaply as possible, otherwise go strong.
+                chooser = min if is_last else max
+                return chooser(winning, key=lambda a: self._strength(a.card, observation.trump_suit))
+            return min(candidates, key=lambda a: self._strength(a.card, observation.trump_suit))
+
+        if losing:
+            # Get rid of the strongest card that still loses.
+            return max(losing, key=lambda a: self._strength(a.card, observation.trump_suit))
+        return min(candidates, key=lambda a: self._strength(a.card, observation.trump_suit))
+
+    @staticmethod
+    def _wins(trick: list[Card], card: Card, trump: Suit | None) -> bool:
+        """Whether ``card`` would currently win the trick (ignores later players)."""
+        return trick_winner([*trick, card], trump) == len(trick)
+
+    @staticmethod
+    def _strength(card: Card, trump: Suit | None) -> float:
+        if card.is_wizard:
+            return 100
+        if card.is_jester:
+            return -1
+        if card.suit is trump:
+            return 50 + card.rank
+        return card.rank

wizard_ai/wizard_ai/agents/model_agent.py

@@ -0,0 +1,33 @@
+from __future__ import annotations
+
+from collections.abc import Callable
+
+from wizard_game import Action, Observation
+
+from ..model.encoding import ActionSpace, encode_observation
+from ..model.policy import Policy
+from .base import Agent
+
+# Called with (observation, features, mask, action_index) for every decision.
+DecisionRecorder = Callable[[Observation, list[float], list[bool], int], None]
+
+
+class ModelAgent(Agent):
+    """Agent backed by a learned ``Policy``."""
+
+    name = "model"
+
+    def __init__(self, policy: Policy, *, explore: bool = False, recorder: DecisionRecorder | None = None) -> None:
+        self.policy = policy
+        self.explore = explore
+        self.recorder = recorder
+
+    def act(self, observation: Observation) -> Action:
+        features = encode_observation(observation)
+        mask = ActionSpace.mask(observation)
+        index = self.policy.select(features, mask, explore=self.explore)
+        if not mask[index]:
+            raise ValueError(f"Policy selected illegal action index {index}")
+        if self.recorder is not None:
+            self.recorder(observation, features, mask, index)
+        return ActionSpace.action(index)

wizard_ai/wizard_ai/agents/random_agent.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import random
+
+from wizard_game import Action, Observation
+
+from .base import Agent
+
+
+class RandomAgent(Agent):
+    """Plays a uniformly random legal action. Useful as a training baseline."""
+
+    name = "random"
+
+    def __init__(self, seed: int | None = None) -> None:
+        self._rng = random.Random(seed)
+
+    def act(self, observation: Observation) -> Action:
+        return self._rng.choice(observation.legal_actions)

wizard_ai/wizard_ai/config.py

@@ -0,0 +1,19 @@
+"""Configuration of the AI service, read from environment variables."""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True)
+class Config:
+    # Websocket endpoint of the game server for AI players.
+    server_url: str = os.environ.get("SERVER_URL", "ws://localhost:8000/internal/ai")
+    # Must match AI_SHARED_SECRET of the game server.
+    secret: str = os.environ.get("AI_SHARED_SECRET", "")
+    # "model", "heuristic" or "random".
+    agent: str = os.environ.get("AGENT", "model")
+    # Checkpoint directory of the trained policy.
+    model_path: str = os.environ.get("MODEL_PATH", "models/latest")
+    log_level: str = os.environ.get("LOG_LEVEL", "INFO")

wizard_ai/wizard_ai/model/__init__.py

@@ -0,0 +1,14 @@
+from .encoding import FEATURE_SIZE, ActionSpace, encode_observation
+from .policy import Policy, Transition, UniformPolicy, checkpoint_exists, load_policy, register_policy
+
+__all__ = [
+    "FEATURE_SIZE",
+    "ActionSpace",
+    "Policy",
+    "Transition",
+    "UniformPolicy",
+    "checkpoint_exists",
+    "encode_observation",
+    "load_policy",
+    "register_policy",
+]

wizard_ai/wizard_ai/model/encoding.py

@@ -0,0 +1,158 @@
+"""Conversion between game objects and model inputs / outputs.
+
+``encode_observation`` turns an ``Observation`` into a fixed size feature
+vector and ``ActionSpace`` maps actions to indices of a fixed size output
+layer (plus a mask of the currently legal actions).
+
+Everything is plain Python lists so the module has no dependencies; convert
+to numpy / torch tensors in the model implementation.
+"""
+
+from __future__ import annotations
+
+from wizard_game import (
+    DECK_SIZE,
+    MAX_PLAYERS,
+    MIN_PLAYERS,
+    SUITS,
+    Action,
+    Bid,
+    ChooseTrump,
+    Observation,
+    Phase,
+    PlayCard,
+    Suit,
+    card_by_id,
+    total_rounds,
+)
+
+MAX_BID = total_rounds(MIN_PLAYERS)  # 20 cards in the last round with 3 players
+MAX_OPPONENTS = MAX_PLAYERS - 1
+SCORE_SCALE = 100.0
+
+
+class ActionSpace:
+    """Fixed action indexing: bids 0-20, then the 60 cards, then 4 trump suits."""
+
+    BID_OFFSET = 0
+    CARD_OFFSET = BID_OFFSET + MAX_BID + 1
+    TRUMP_OFFSET = CARD_OFFSET + DECK_SIZE
+    SIZE = TRUMP_OFFSET + len(SUITS)
+
+    @classmethod
+    def index(cls, action: Action) -> int:
+        if isinstance(action, Bid):
+            return cls.BID_OFFSET + action.value
+        if isinstance(action, PlayCard):
+            return cls.CARD_OFFSET + action.card.id
+        if isinstance(action, ChooseTrump):
+            return cls.TRUMP_OFFSET + action.suit.index
+        raise TypeError(f"Unknown action {action!r}")
+
+    @classmethod
+    def action(cls, index: int) -> Action:
+        if not 0 <= index < cls.SIZE:
+            raise ValueError(f"Action index out of range: {index}")
+        if index >= cls.TRUMP_OFFSET:
+            return ChooseTrump(SUITS[index - cls.TRUMP_OFFSET])
+        if index >= cls.CARD_OFFSET:
+            return PlayCard(card_by_id(index - cls.CARD_OFFSET))
+        return Bid(index - cls.BID_OFFSET)
+
+    @classmethod
+    def mask(cls, observation: Observation) -> list[bool]:
+        mask = [False] * cls.SIZE
+        for action in observation.legal_actions:
+            mask[cls.index(action)] = True
+        return mask
+
+
+def _one_hot(size: int, index: int | None) -> list[float]:
+    vector = [0.0] * size
+    if index is not None:
+        vector[index] = 1.0
+    return vector
+
+
+def _suit_one_hot(suit: Suit | None) -> list[float]:
+    # Last slot means "no suit".
+    return _one_hot(len(SUITS) + 1, suit.index if suit is not None else len(SUITS))
+
+
+def _cards(cards) -> list[float]:
+    vector = [0.0] * DECK_SIZE
+    for card in cards:
+        vector[card.id] = 1.0
+    return vector
+
+
+PHASES = (Phase.TRUMP_SELECTION, Phase.BIDDING, Phase.PLAYING)
+
+
+def encode_observation(obs: Observation) -> list[float]:
+    """Encode an observation as a vector of length ``FEATURE_SIZE``.
+
+    Opponent features are ordered relative to the observing player (next
+    player first) and zero padded to ``MAX_OPPONENTS``, so the same model
+    can play with 3 to 6 players.
+    """
+    me = obs.player
+    n = obs.num_players
+    features: list[float] = []
+
+    # Game context
+    features += _one_hot(MAX_PLAYERS - MIN_PLAYERS + 1, n - MIN_PLAYERS)
+    features += _one_hot(len(PHASES), PHASES.index(obs.phase) if obs.phase in PHASES else None)
+    features += [obs.round_number / MAX_BID, obs.round_number / obs.total_rounds]
+    features += [float(obs.options.not_equal), float(obs.options.hidden_bids)]
+    features += _one_hot(MAX_PLAYERS, (me - obs.dealer) % n)  # seat relative to dealer
+    features += _one_hot(MAX_PLAYERS, len(obs.current_trick))  # position in the trick
+
+    # Cards
+    features += _cards(obs.hand)
+    features += _cards(obs.played_cards)
+    features += _cards(played.card for played in obs.current_trick)
+    features += _suit_one_hot(obs.trump_suit)
+    features += _suit_one_hot(obs.lead_suit)
+    features += _one_hot(DECK_SIZE, obs.trump_card.id if obs.trump_card else None)
+
+    # Own bid and tricks
+    own_bid = obs.bids[me]
+    features += [
+        float(own_bid is not None),
+        (own_bid or 0) / MAX_BID,
+        obs.tricks_won[me] / MAX_BID,
+        ((own_bid or 0) - obs.tricks_won[me]) / MAX_BID,
+        obs.scores[me] / SCORE_SCALE,
+    ]
+
+    # Opponents, relative to me
+    for offset in range(1, MAX_OPPONENTS + 1):
+        if offset < n:
+            p = (me + offset) % n
+            bid = obs.bids[p]
+            features += [
+                1.0,
+                float(obs.has_bid[p]),
+                float(bid is not None),
+                (bid or 0) / MAX_BID,
+                obs.tricks_won[p] / MAX_BID,
+                (obs.scores[p] - obs.scores[me]) / SCORE_SCALE,
+                float(any(played.player == p for played in obs.current_trick)),
+            ]
+        else:
+            features += [0.0] * 7
+
+    # Sum of known bids vs. tricks available
+    known_bids = sum(bid for bid in obs.bids if bid is not None)
+    features += [(known_bids - obs.round_number) / MAX_BID]
+    return features
+
+
+FEATURE_SIZE = (
+    (MAX_PLAYERS - MIN_PLAYERS + 1) + len(PHASES) + 2 + 2 + MAX_PLAYERS + MAX_PLAYERS
+    + DECK_SIZE * 3 + (len(SUITS) + 1) * 2 + DECK_SIZE
+    + 5
+    + MAX_OPPONENTS * 7
+    + 1
+)

wizard_ai/wizard_ai/model/policy.py

@@ -0,0 +1,126 @@
+"""Policy (model) interface and checkpoint handling.
+
+This is the place where the actual reinforcement learning model will be
+plugged in. A policy maps an encoded observation plus a legal action mask
+(see ``encoding``) to an action index, and can be updated from experience
+collected in self-play (see ``training.self_play``).
+
+To add a model, subclass ``Policy``, implement ``select``, ``update``,
+``save`` and ``load`` and decorate the class with ``@register_policy``. A
+checkpoint is a directory containing ``policy.json`` (with the registered
+``type``) plus whatever files the policy writes (e.g. ``weights.pt``).
+"""
+
+from __future__ import annotations
+
+import json
+import random
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, ClassVar
+
+from .encoding import ActionSpace, FEATURE_SIZE
+
+METADATA_FILE = "policy.json"
+
+
+@dataclass(slots=True)
+class Transition:
+    """One decision of one player, collected during self-play."""
+
+    features: list[float]
+    mask: list[bool]
+    action: int
+    # Reward observed after this decision (score change of a finished round,
+    # scaled). Rewards are attached to the player's most recent decision.
+    reward: float = 0.0
+    # True for the player's last decision of the game.
+    done: bool = False
+
+
+class Policy(ABC):
+    type_name: ClassVar[str]
+
+    feature_size: int = FEATURE_SIZE
+    action_size: int = ActionSpace.SIZE
+
+    @abstractmethod
+    def select(self, features: list[float], mask: list[bool], *, explore: bool = False) -> int:
+        """Return the index of a legal action (``mask[index]`` is True)."""
+
+    @abstractmethod
+    def update(self, episodes: list[list[Transition]]) -> dict[str, float]:
+        """Improve the policy from trajectories (one list per player and game).
+
+        Returns metrics for logging.
+        """
+
+    def save(self, directory: str | Path) -> None:
+        directory = Path(directory)
+        directory.mkdir(parents=True, exist_ok=True)
+        metadata = {"type": self.type_name, "config": self.config()}
+        (directory / METADATA_FILE).write_text(json.dumps(metadata, indent=2))
+        self.save_weights(directory)
+
+    def config(self) -> dict[str, Any]:
+        """Constructor arguments stored in the checkpoint metadata."""
+        return {}
+
+    def save_weights(self, directory: Path) -> None:  # noqa: B027 - optional hook
+        """Write model parameters into ``directory``."""
+
+    @classmethod
+    def load(cls, directory: Path, config: dict[str, Any]) -> Policy:
+        return cls(**config)
+
+
+_REGISTRY: dict[str, type[Policy]] = {}
+
+
+def register_policy(cls: type[Policy]) -> type[Policy]:
+    _REGISTRY[cls.type_name] = cls
+    return cls
+
+
+def load_policy(directory: str | Path) -> Policy:
+    directory = Path(directory)
+    metadata = json.loads((directory / METADATA_FILE).read_text())
+    policy_type = metadata.get("type")
+    if policy_type not in _REGISTRY:
+        raise ValueError(f"Unknown policy type {policy_type!r} in {directory}")
+    return _REGISTRY[policy_type].load(directory, metadata.get("config", {}))
+
+
+def checkpoint_exists(directory: str | Path | None) -> bool:
+    return directory is not None and (Path(directory) / METADATA_FILE).is_file()
+
+
+@register_policy
+class UniformPolicy(Policy):
+    """Placeholder model: picks a uniformly random legal action.
+
+    Replace with the real (neural network) policy; it only exists so that
+    the training and serving pipeline can be run end to end.
+    """
+
+    type_name = "uniform"
+
+    def __init__(self, seed: int | None = None) -> None:
+        self._seed = seed
+        self._rng = random.Random(seed)
+
+    def select(self, features: list[float], mask: list[bool], *, explore: bool = False) -> int:
+        legal = [index for index, allowed in enumerate(mask) if allowed]
+        return self._rng.choice(legal)
+
+    def update(self, episodes: list[list[Transition]]) -> dict[str, float]:
+        transitions = sum(len(episode) for episode in episodes)
+        total_reward = sum(t.reward for episode in episodes for t in episode)
+        return {
+            "transitions": float(transitions),
+            "mean_episode_reward": total_reward / max(1, len(episodes)),
+        }
+
+    def config(self) -> dict[str, Any]:
+        return {"seed": self._seed}

wizard_ai/wizard_ai/service/__init__.py

@@ -0,0 +1,3 @@
+from .client import AIService
+
+__all__ = ["AIService"]

wizard_ai/wizard_ai/service/client.py

@@ -0,0 +1,75 @@
+"""Production mode: serve decisions to the game server over a websocket.
+
+The service connects to the server's ``/internal/ai`` endpoint and answers
+``decide`` requests (see ``wizard_server/app/ai_bridge.py``). It reconnects
+automatically if the connection drops.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from typing import Any
+
+from websockets.asyncio.client import ClientConnection, connect
+from websockets.exceptions import ConnectionClosed, InvalidHandshake, InvalidStatus
+
+from wizard_game import Observation
+
+from ..agents import Agent
+
+logger = logging.getLogger(__name__)
+
+MAX_BACKOFF = 30.0
+
+
+class AIService:
+    def __init__(self, server_url: str, agent: Agent, secret: str = "") -> None:
+        self.server_url = server_url
+        self.agent = agent
+        self.secret = secret
+        # Inference runs in worker threads; the lock keeps agents that are not
+        # thread safe (e.g. stateful models) consistent.
+        self._agent_lock = asyncio.Lock()
+
+    async def run_forever(self) -> None:
+        backoff = 1.0
+        while True:
+            try:
+                await self._run_once()
+                backoff = 1.0
+            except (OSError, ConnectionClosed, InvalidHandshake, TimeoutError) as exc:
+                if isinstance(exc, InvalidStatus) and exc.response.status_code in (401, 403):
+                    logger.error("Server rejected the connection, check AI_SHARED_SECRET")
+                logger.warning("Connection to %s lost (%s), retrying in %.0fs", self.server_url, exc, backoff)
+            await asyncio.sleep(backoff)
+            backoff = min(backoff * 2, MAX_BACKOFF)
+
+    async def _run_once(self) -> None:
+        headers = {"Authorization": f"Bearer {self.secret}"} if self.secret else {}
+        async with connect(self.server_url, additional_headers=headers, open_timeout=10) as websocket:
+            logger.info("Connected to %s as '%s' agent", self.server_url, self.agent.name)
+            await websocket.send(json.dumps({"type": "hello", "agent": self.agent.name}))
+            tasks: set[asyncio.Task] = set()
+            async for raw in websocket:
+                message = json.loads(raw)
+                if message.get("type") == "decide":
+                    task = asyncio.create_task(self._decide(websocket, message))
+                    tasks.add(task)
+                    task.add_done_callback(tasks.discard)
+
+    async def _decide(self, websocket: ClientConnection, message: dict[str, Any]) -> None:
+        request_id = message.get("request_id")
+        try:
+            observation = Observation.from_dict(message["observation"])
+            async with self._agent_lock:
+                action = await asyncio.to_thread(self.agent.act, observation)
+            reply = {"type": "action", "request_id": request_id, "action": action.to_dict()}
+        except Exception as exc:  # report instead of crashing; the server falls back
+            logger.exception("Failed to decide request %s", request_id)
+            reply = {"type": "error", "request_id": request_id, "message": str(exc)}
+        try:
+            await websocket.send(json.dumps(reply))
+        except ConnectionClosed:
+            pass

wizard_ai/wizard_ai/training/__init__.py

@@ -0,0 +1,5 @@
+from .env import StepResult, WizardEnv
+from .evaluate import EvaluationResult, evaluate
+from .self_play import SelfPlayTrainer, TrainingConfig
+
+__all__ = ["EvaluationResult", "SelfPlayTrainer", "StepResult", "TrainingConfig", "WizardEnv", "evaluate"]

wizard_ai/wizard_ai/training/env.py

@@ -0,0 +1,72 @@
+"""Multi-agent training environment around ``wizard_game.WizardGame``."""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+
+from wizard_game import Action, Event, GameOptions, Observation, RoundCompleted, WizardGame
+
+from ..model.encoding import SCORE_SCALE
+
+
+@dataclass(slots=True)
+class StepResult:
+    # Observation of the player who has to act next (None once the game is over).
+    observation: Observation | None
+    # Reward per player for this step: the scaled score change at the end of
+    # a round, 0 otherwise.
+    rewards: list[float]
+    done: bool
+    events: list[Event]
+
+
+class WizardEnv:
+    """Turn based environment: each step is one decision of the current player.
+
+    ``num_players`` may be a fixed number or a tuple of player counts to
+    sample from on every reset, so a single model learns all table sizes.
+    """
+
+    def __init__(
+        self,
+        num_players: int | tuple[int, ...] = 4,
+        options: GameOptions | None = None,
+        seed: int | None = None,
+    ) -> None:
+        self._player_counts = (num_players,) if isinstance(num_players, int) else tuple(num_players)
+        self._options = options or GameOptions()
+        self._rng = random.Random(seed)
+        self.game: WizardGame | None = None
+
+    @property
+    def current_player(self) -> int | None:
+        return self.game.current_player if self.game else None
+
+    @property
+    def num_players(self) -> int:
+        assert self.game is not None, "call reset() first"
+        return self.game.num_players
+
+    def reset(self) -> Observation:
+        num_players = self._rng.choice(self._player_counts)
+        self.game = WizardGame(num_players, self._options, rng=random.Random(self._rng.getrandbits(64)))
+        return self.observe(self.game.current_player)
+
+    def observe(self, player: int) -> Observation:
+        assert self.game is not None, "call reset() first"
+        return self.game.observation(player)
+
+    def step(self, action: Action) -> StepResult:
+        game = self.game
+        assert game is not None and game.current_player is not None, "game is over, call reset()"
+        events = game.apply(game.current_player, action)
+
+        rewards = [0.0] * game.num_players
+        for event in events:
+            if isinstance(event, RoundCompleted):
+                rewards = [change / SCORE_SCALE for change in event.result.score_changes]
+
+        done = game.is_finished
+        observation = None if done else self.observe(game.current_player)
+        return StepResult(observation, rewards, done, events)

wizard_ai/wizard_ai/training/evaluate.py

@@ -0,0 +1,61 @@
+"""Let agents play against each other to measure their strength."""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+
+from wizard_game import GameOptions, WizardGame
+
+from ..agents import Agent
+
+
+@dataclass
+class EvaluationResult:
+    names: list[str]
+    games: int
+    mean_scores: list[float]
+    win_rates: list[float]
+
+    def __str__(self) -> str:
+        rows = [f"{'agent':<14}{'mean score':>12}{'win rate':>10}"]
+        for name, score, win_rate in zip(self.names, self.mean_scores, self.win_rates):
+            rows.append(f"{name:<14}{score:>12.1f}{win_rate:>10.1%}")
+        return f"{self.games} games\n" + "\n".join(rows)
+
+
+def evaluate(
+    agents: list[Agent],
+    games: int,
+    options: GameOptions | None = None,
+    seed: int | None = None,
+) -> EvaluationResult:
+    """Play ``games`` games; seats are rotated so every agent starts everywhere.
+
+    Ties count as a win for every tied agent.
+    """
+    rng = random.Random(seed)
+    n = len(agents)
+    totals = [0.0] * n
+    wins = [0.0] * n
+    for game_index in range(games):
+        shift = game_index % n
+        seating = [agents[(seat + shift) % n] for seat in range(n)]
+        for agent in seating:
+            agent.reset()
+        game = WizardGame(n, options, rng=random.Random(rng.getrandbits(64)))
+        while not game.is_finished:
+            player = game.current_player
+            assert player is not None
+            game.apply(player, seating[player].act(game.observation(player)))
+        winners = set(game.winners())
+        for seat in range(n):
+            agent_index = (seat + shift) % n
+            totals[agent_index] += game.scores[seat]
+            wins[agent_index] += seat in winners
+    return EvaluationResult(
+        names=[agent.name for agent in agents],
+        games=games,
+        mean_scores=[total / games for total in totals],
+        win_rates=[win / games for win in wins],
+    )

wizard_ai/wizard_ai/training/self_play.py

@@ -0,0 +1,102 @@
+"""Self-play training loop.
+
+The trainer lets the learning policy play against copies of itself (or
+against baseline agents), collects one trajectory of ``Transition`` objects
+per learning seat and game, and hands batches of trajectories to
+``Policy.update``. The learning algorithm itself (PPO, DQN, ...) lives in the
+policy implementation.
+"""
+
+from __future__ import annotations
+
+import logging
+import random
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from wizard_game import GameOptions, Observation
+
+from ..agents import Agent, HeuristicAgent, ModelAgent, RandomAgent
+from ..model.policy import Policy, Transition
+from .env import WizardEnv
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TrainingConfig:
+    episodes: int = 1000
+    player_counts: tuple[int, ...] = (3, 4, 5, 6)
+    options: GameOptions = field(default_factory=GameOptions)
+    # Opponents of the learning seat: "self" (all seats learn), "heuristic" or "random".
+    opponents: str = "self"
+    episodes_per_update: int = 16
+    checkpoint_dir: Path | None = None
+    checkpoint_every: int = 500
+    log_every: int = 100
+    seed: int | None = None
+
+
+class SelfPlayTrainer:
+    def __init__(self, policy: Policy, config: TrainingConfig) -> None:
+        self.policy = policy
+        self.config = config
+        self.env = WizardEnv(config.player_counts, config.options, seed=config.seed)
+        self._rng = random.Random(config.seed)
+
+    def play_episode(self) -> tuple[list[list[Transition]], list[int]]:
+        """Play one game. Returns the learners' trajectories and the final scores."""
+        observation: Observation | None = self.env.reset()
+        num_players = self.env.num_players
+        learners = self._learning_seats(num_players)
+        trajectories: dict[int, list[Transition]] = {seat: [] for seat in learners}
+
+        def record(obs: Observation, features: list[float], mask: list[bool], index: int) -> None:
+            trajectories[obs.player].append(Transition(features, mask, index))
+
+        learner = ModelAgent(self.policy, explore=True, recorder=record)
+        agents: list[Agent] = [learner if seat in learners else self._opponent() for seat in range(num_players)]
+
+        while observation is not None:
+            action = agents[observation.player].act(observation)
+            result = self.env.step(action)
+            for seat, reward in enumerate(result.rewards):
+                if reward and seat in trajectories and trajectories[seat]:
+                    trajectories[seat][-1].reward += reward
+            observation = result.observation
+
+        for trajectory in trajectories.values():
+            if trajectory:
+                trajectory[-1].done = True
+        assert self.env.game is not None
+        return list(trajectories.values()), list(self.env.game.scores)
+
+    def train(self) -> None:
+        config = self.config
+        batch: list[list[Transition]] = []
+        for episode in range(1, config.episodes + 1):
+            trajectories, _ = self.play_episode()
+            batch.extend(trajectories)
+
+            if episode % config.episodes_per_update == 0 or episode == config.episodes:
+                metrics = self.policy.update(batch)
+                batch = []
+                if episode % config.log_every < config.episodes_per_update or episode == config.episodes:
+                    summary = ", ".join(f"{key}={value:.3f}" for key, value in metrics.items())
+                    logger.info("episode %d/%d: %s", episode, config.episodes, summary)
+
+            if config.checkpoint_dir and (episode % config.checkpoint_every == 0 or episode == config.episodes):
+                self.policy.save(config.checkpoint_dir)
+                logger.info("Saved checkpoint to %s", config.checkpoint_dir)
+
+    def _learning_seats(self, num_players: int) -> set[int]:
+        if self.config.opponents == "self":
+            return set(range(num_players))
+        return {self._rng.randrange(num_players)}
+
+    def _opponent(self) -> Agent:
+        if self.config.opponents == "heuristic":
+            return HeuristicAgent()
+        if self.config.opponents == "random":
+            return RandomAgent(self._rng.getrandbits(32))
+        raise ValueError(f"Unknown opponents {self.config.opponents!r}")