Wizard/wizard_ai/wizard_ai/model/encoding.py

"""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.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 + 1 + MAX_PLAYERS + MAX_PLAYERS
    + DECK_SIZE * 3 + (len(SUITS) + 1) * 2 + DECK_SIZE
    + 5
    + MAX_OPPONENTS * 7
    + 1
)