Wizard/docs/07-ai.md

7. AI player (wizard_ai)

The AI package contains everything around a reinforcement-learning player: the agent interface, a feature encoding, a pluggable policy (model) interface, a training environment with a self-play loop, an evaluation tool, and the production websocket service. The actual learning model is not implemented yet: UniformPolicy is a random placeholder, and until a trained checkpoint exists production uses the rule-based HeuristicAgent.

Package layout

wizard_ai/wizard_ai/
├── __main__.py            # CLI: serve | train | evaluate
├── config.py              # env vars: SERVER_URL, AI_SHARED_SECRET, AGENT, MODEL_PATH, LOG_LEVEL
├── agents/
│   ├── base.py            # Agent interface
│   ├── random_agent.py    # RandomAgent
│   ├── heuristic_agent.py # HeuristicAgent (baseline + production fallback)
│   ├── model_agent.py     # ModelAgent: Agent backed by a Policy
│   └── __init__.py        # make_agent(name, model_path)
├── model/
│   ├── encoding.py        # encode_observation(), FEATURE_SIZE, ActionSpace
│   └── policy.py          # Policy interface, Transition, checkpoints, UniformPolicy
├── training/
│   ├── env.py             # WizardEnv (multi-agent environment)
│   ├── self_play.py       # TrainingConfig, SelfPlayTrainer
│   └── evaluate.py        # evaluate(agents, games)
└── service/client.py      # AIService: production websocket client

How the pieces fit together

flowchart TB
    subgraph production
        SRV[game server] -- "decide (observation JSON)" --> SVC[AIService]
        SVC -- "Observation.from_dict" --> AG["Agent.act"]
        AG -- Action --> SVC -- "action JSON" --> SRV
    end
    subgraph training
        ENV["WizardEnv<br/>wraps WizardGame"] -- Observation --> MA[ModelAgent]
        MA -- "encode_observation + ActionSpace.mask" --> POL["Policy.select"]
        POL -- "action index" --> MA
        MA -- Action --> ENV
        MA -- "records Transition" --> TR[SelfPlayTrainer]
        ENV -- rewards --> TR
        TR -- trajectories --> UPD["Policy.update"]
        TR -- save --> CK[("models/latest/")]
    end
    CK -- load_policy --> AG

The key idea: an Agent maps an Observation to an Action, and the observation is identical in training (straight from WizardGame) and in production (JSON from the server, decoded with Observation.from_dict). A trained agent can be deployed without any change.

Agents (agents/)

class Agent(ABC):
    name: str
    def act(self, observation: Observation) -> Action: ...   # must return one of observation.legal_actions
    def reset(self) -> None: ...                              # optional, called before each game in evaluation

Agent	Behaviour
RandomAgent	uniformly random legal action
HeuristicAgent	Bidding: estimates expected tricks (wizard ≈ 0.95, trump ≈ 0.35 + 0.05·rank, high side cards 12–13 count less with more opponents), bids the closest legal value. Playing: if it still needs tricks, it wins with the cheapest winning card when last to play, otherwise with the strongest (and plays its weakest card if it cannot win); if it needs no more tricks, it throws the strongest card that still loses, else its weakest card.
ModelAgent(policy, explore=False, recorder=None)	encodes the observation, asks the policy for an action index, converts it back to an Action; optionally records every decision (used by the trainer)

make_agent(name, model_path): "random", "heuristic" or "model". For "model" it loads the checkpoint at model_path if it exists, otherwise it logs a warning and returns a HeuristicAgent.

Baseline strength (100 games, evaluate): heuristic vs. three random agents → heuristic ≈ +220 average score, 100 % win rate; random agents ≈ −350.

Encoding (model/encoding.py)

Features — encode_observation(obs) -> list[float] (length FEATURE_SIZE = 312)

Plain Python floats (convert to numpy/torch in the model). Opponents are ordered relative to the observer (next player first) and zero-padded to 5, so one model works for 3–6 players.

Index	Size	Feature
0–3	4	number of players one-hot (3, 4, 5, 6)
4–5	2	phase one-hot (bidding, playing)
6	1	round number / 20
7	1	round number / total rounds (game progress)
8	1	hidden bids option
9–14	6	own seat relative to the dealer, one-hot
15–20	6	own position in the current trick (0 = leading), one-hot
21–80	60	own hand (card id one-hot)
81–140	60	all cards played this round
141–200	60	cards in the current trick
201–205	5	trump suit one-hot (4 suits + none)
206–210	5	lead suit one-hot (4 suits + none)
211–270	60	trump card one-hot (none = all zero)
271–275	5	own: has bid, bid/20, tricks/20, (bid − tricks)/20, score/100
276–310	35	5 opponent slots × 7: present, has bid, bid visible, bid/20, tricks/20, (their score − my score)/100, already played in this trick
311	1	(sum of known bids − round number)/20

If you change the encoding, FEATURE_SIZE must be updated (a test checks the length for all player counts) and existing checkpoints become incompatible.

Actions — ActionSpace (size 81)

Index	Action
0–20	Bid(0) … Bid(20)
21–80	PlayCard(card id 0…59)

ActionSpace.index(action), ActionSpace.action(index) and ActionSpace.mask(observation) (list of 81 booleans, True = legal). A model should give illegal actions zero probability using the mask.

Policy interface (model/policy.py) — where the model goes

@dataclass
class Transition:
    features: list[float]   # encode_observation(obs)
    mask: list[bool]        # ActionSpace.mask(obs)
    action: int             # chosen action index
    reward: float = 0.0     # scaled score change, attached to the player's last decision of a round
    done: bool = False      # True for the player's last decision of the game

class Policy(ABC):
    type_name: ClassVar[str]                      # key in the checkpoint metadata
    def select(self, features, mask, *, explore=False) -> int: ...
    def update(self, episodes: list[list[Transition]]) -> dict[str, float]: ...
    def config(self) -> dict: ...                 # constructor kwargs stored in the checkpoint
    def save_weights(self, directory: Path): ...  # write parameters
    @classmethod
    def load(cls, directory: Path, config: dict) -> Policy: ...

Checkpoints are directories: policy.save(dir) writes dir/policy.json = {"type": type_name, "config": {...}} and calls save_weights(dir). load_policy(dir) reads the type, finds the class registered with @register_policy, and calls its load.

Implementing the real model — step by step

1. Add torch (or another library) to wizard_ai/requirements.txt.

2. Create e.g. wizard_ai/wizard_ai/model/ppo_policy.py:

   @register_policy
   class PPOPolicy(Policy):
       type_name = "ppo"
   
       def __init__(self, hidden: int = 256, lr: float = 3e-4):
           self.hidden, self.lr = hidden, lr
           self.net = ...  # FEATURE_SIZE -> hidden -> (ActionSpace.SIZE logits, value)
   
       def select(self, features, mask, *, explore=False):
           logits = self.net(torch.tensor(features))[0]
           logits[~torch.tensor(mask)] = -inf           # only legal actions
           dist = Categorical(logits=logits)
           return int(dist.sample() if explore else logits.argmax())
   
       def update(self, episodes):
           # compute returns/advantages from Transition.reward/done, run PPO epochs
           return {"loss": ..., "mean_episode_reward": ...}
   
       def config(self):
           return {"hidden": self.hidden, "lr": self.lr}
   
       def save_weights(self, directory):
           torch.save(self.net.state_dict(), directory / "weights.pt")
   
       @classmethod
       def load(cls, directory, config):
           policy = cls(**config)
           policy.net.load_state_dict(torch.load(directory / "weights.pt"))
           return policy
   

   For PPO you will likely want to store log-probabilities and value estimates per decision; extend Transition (or keep them inside the policy keyed by decision order) as needed.

3. Import the module in model/__init__.py so @register_policy runs.

4. In __main__.py (train command) replace UniformPolicy(args.seed) with your policy (the line is marked TODO).

5. Train: python -m wizard_ai train --episodes 50000 --out models/latest.

6. Compare: python -m wizard_ai evaluate --agents model,heuristic,heuristic,heuristic.

7. Deploy: the checkpoint directory wizard_ai/models/latest/ is mounted into the container (compose.yaml); restart wizard_ai. With AGENT=model (default) it is loaded automatically.

Model size matters: the service answers every AI move in a worker thread and the server waits at most AI_DECISION_TIMEOUT (5 s).

Training (training/)

WizardEnv

A turn-based multi-agent environment around WizardGame:

env = WizardEnv(num_players=(3, 4, 5, 6), options=GameOptions(), seed=1)
obs = env.reset()                  # observation of the player to act; samples a player count
while obs is not None:
    result = env.step(action)      # action for obs.player
    result.rewards                 # per seat: score change / 100 when a round ended, else 0
    result.done, result.events
    obs = result.observation       # next player to act (None when the game is over)

The sum of all rewards of a seat over a game equals its final score / 100 (tested).

SelfPlayTrainer

TrainingConfig fields: episodes, player_counts, options, opponents ("self" = every seat is the learning policy, "heuristic" / "random" = one learning seat against those), episodes_per_update, checkpoint_dir, checkpoint_every, log_every, seed.

Per episode (= one full game) it:

1. resets the env (random player count),
2. lets a ModelAgent(policy, explore=True, recorder=…) play the learning seats, recording one Transition per decision,
3. adds each round's reward to that seat's most recent transition, and marks each seat's last transition done,
4. collects the trajectories (one list per learning seat and game).

Every episodes_per_update episodes it calls policy.update(batch), logs the returned metrics, and saves a checkpoint every checkpoint_every episodes and at the end.

evaluate(agents, games, options=None, seed=None)

Plays games games with len(agents) players, rotating seats so every agent starts in every position, and returns mean score and win rate per agent (ties count as a win for everyone tied). Seeded for fair comparisons.

Production service (service/client.py)

AIService(server_url, agent, secret).run_forever():

• connects to SERVER_URL with Authorization: Bearer <AI_SHARED_SECRET>, sends hello, then answers decide requests (protocol in chapter 5);
• each request runs agent.act in a worker thread (asyncio.to_thread) behind a lock, so a slow model does not block the websocket and non-thread-safe agents stay consistent;
• errors while deciding are reported back as error (the server then plays a random move) instead of crashing;
• reconnects forever with exponential backoff (1 s … 30 s); logs a hint if the server rejects the secret.

Command line

cd wizard_ai
python -m wizard_ai serve [--server-url …] [--agent model|heuristic|random] [--model-path …]
python -m wizard_ai train --episodes 5000 [--players 3,4,5,6] [--opponents self|heuristic|random]
                          [--hidden-bids] [--out models/latest] [--resume] [--seed N]
python -m wizard_ai evaluate --agents model,heuristic,random,random --games 200 [--seed N]

Environment variables (config.py): SERVER_URL (default ws://localhost:8000/internal/ai; in Docker ws://wizard_server:8083/internal/ai), AI_SHARED_SECRET, AGENT (default model), MODEL_PATH (default models/latest, in Docker /srv/models/latest), LOG_LEVEL.

Tests (wizard_ai/tests/test_ai.py)

Feature length and masks for 3–6 players, action-space round trip, all agents play legal full games, env rewards add up to scores, trainer runs and writes a loadable checkpoint, make_agent fallback, heuristic beats random.