# 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

```mermaid
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/`)

```python
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

```python
@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`:

   ```python
   @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, the
container is capped at `AI_CPUS` (1 CPU by default) and the server waits at
most `AI_DECISION_TIMEOUT` (5 s). If the model is too slow the service
automatically answers with the heuristic instead (see below), so measure the
decision time before deploying.

## Training (`training/`)

### `WizardEnv`

A turn-based multi-agent environment around `WizardGame`:

```python
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, fallback=…, workers=…, max_queue=…, degrade_queue=…).run_forever()`:

- connects to `SERVER_URL` with `Authorization: Bearer <AI_SHARED_SECRET>`,
  sends `hello`, then answers `decide` requests (protocol in
  [chapter 5](05-protocol.md#ai-protocol));
- each request runs `agent.act` in a worker thread (`asyncio.to_thread`), so a
  slow model does not block the websocket;
- **bounded work** (`handle_request`):
  - at most `AI_WORKERS` (default 1) decisions are computed at the same time —
    agents must be thread safe if you raise it;
  - at most `AI_MAX_QUEUE` (64) requests may be in progress; more are answered
    with `error` "overloaded" immediately;
  - a request whose server `timeout` passed while it waited for a worker is
    answered with `error` "expired" without computing it;
- **degradation**: when `AI_DEGRADE_QUEUE` (8) or more other requests are
  waiting, or the primary agent's average decision time is more than half of
  the remaining time, the cheap fallback agent (`HeuristicAgent`) decides
  instead. `serve` sets the heuristic as fallback automatically unless the
  primary agent already is the heuristic. The reply's `agent` field says which
  one decided;
- `OMP_NUM_THREADS` / `MKL_NUM_THREADS` are set to `AI_WORKERS` in
  `compose.yaml`, so numeric libraries don't start a thread per host core;
- 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

```bash
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`, `AI_WORKERS` (1), `AI_MAX_QUEUE` (64),
`AI_DEGRADE_QUEUE` (8).

## 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, and the
service limits (queue full → overloaded, expired requests, degradation to the
heuristic when busy or when the model is too slow).