diff --git a/examples/pytorch/pytorch-video-poker-bot/.gitignore b/examples/pytorch/pytorch-video-poker-bot/.gitignore new file mode 100644 index 00000000..8073a10a --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/.gitignore @@ -0,0 +1,10 @@ +.venv/ +__pycache__/ +*.py[cod] +.DS_Store +checkpoints/ +wandb/ +.env +*.pt +data/*.npz +artifacts/ diff --git a/examples/pytorch/pytorch-video-poker-bot/README.md b/examples/pytorch/pytorch-video-poker-bot/README.md new file mode 100644 index 00000000..7c65a0a5 --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/README.md @@ -0,0 +1,136 @@ +# Train a video poker bot with PyTorch and W&B + +This example trains a small PyTorch network to play video poker and tracks +it with [Weights & Biases](https://wandb.ai). The poker is just a stand-in +task. The point is to see a complete W&B run: + +1. **`wandb.init()`** starts a run, puts it in a group (`train` or `eval`), + and records its config (hyperparameters and dataset size). +2. **`run.log()`** sends metrics every epoch, so you can watch training live. +3. **`run.log_artifact()`** uploads the trained checkpoint as a versioned + model artifact. +4. **Finish**: `train.py` opens the run with `with wandb.init(...) as run:`, + so the run finishes automatically when the block exits. Without the + `with` block, you would call `run.finish()` yourself. + +`evaluate.py` then starts a second run (`job_type="evaluation"`). It fetches +the model artifact from W&B with **`run.use_artifact()`**, which links the +two runs, and logs its score every 100 hands as it plays. + +## How video poker works + +You are dealt five cards. You choose which of them to hold: none, all five, +or any mix in between (32 possible combinations). The cards you don't hold +are replaced with new ones from the deck, and that is your final hand. Its +reward is looked up in the paytable. + +This example plays **9/6 Jacks or Better**, betting 5 credits per hand: + +| Final hand | Reward (credits) | +| --- | --- | +| Royal flush | 4000 | +| Straight flush | 250 | +| Four of a kind | 125 | +| Full house | 45 | +| Flush | 30 | +| Straight | 20 | +| Three of a kind | 15 | +| Two pair | 10 | +| Pair of jacks or better | 5 | +| Anything else | 0 | + +The name "9/6" comes from the full house paying 9 and the flush paying 6 for +each credit bet. With optimal holds, the rewards average about 99.5% of the +credits wagered. + +The dataset gives each dealt hand the exact expected reward of all 32 hold +choices. The network learns to predict those values, and at play time it +makes the hold with the highest prediction. + +## Setup + +```bash +cd examples/pytorch/pytorch-video-poker-bot +python3 -m venv .venv +source .venv/bin/activate +pip install -r requirements.txt +wandb login +``` + +## Run it + +### 1. Generate the dataset + +```bash +python generate_dataset.py --output data/hands.npz +``` + +This deals 1,000,000 random hands and labels each one. It takes about two +minutes and writes a ~44 MB file. For a quick test, add `--hands 20000`. + +### 2. Train + +```bash +python train.py \ + --dataset data/hands.npz \ + --checkpoint checkpoints/jacks_or_better_network.pt \ + --artifact-name jacks-or-better-network \ + --project video-poker \ + --run-name train-baseline +``` + +This takes under a minute on a laptop CPU. After each epoch it saves the +checkpoint if validation improved, and at the end it uploads the checkpoint +to W&B as the model artifact `jacks-or-better-network`. Optional flags: +`--entity` (a W&B team or user; defaults to your default entity), +`--epochs` (default 20), `--lr`, `--batch-size`, `--hidden-size`, `--seed`. + +Without `--seed`, each run picks a random seed. The seed it used is saved +in the run's config in W&B, so you can repeat any run exactly by passing +that value. + +### 3. Evaluate + +```bash +python evaluate.py \ + --artifact jacks-or-better-network:latest \ + --project video-poker \ + --run-name eval-baseline +``` + +This downloads the latest version of the model artifact and plays 100,000 +hands with it (change with `--hands`), which takes about 10 seconds. It logs +to a separate run in the same project. Like training, it deals with a random +seed unless you pass `--seed`, so each evaluation plays different hands. +If you trained under a team with `--entity`, pass the same `--entity` here. + +## What you will see in W&B + +- **Training run**: config on the Overview tab, and these charts per epoch: + - `train_loss` and `val_loss` + - `val_optimal_action_pct`: how often the network makes the optimal hold + - `val_mean_regret`: the expected reward per credit bet that the network's + holds give up compared with the optimal holds, on average. Lower is + better, and the saved checkpoint is the epoch with the lowest value. + - `val_expected_return_pct`: expected rewards as a percentage of credits + wagered +- **Evaluation run**: running totals logged every 100 hands, so you can + watch the score settle as more hands are played: `wagered` (credits bet), + `payout` (credits rewarded), `profit` (the difference) and `return_pct` + (rewards as a percentage of credits wagered). +- **Model artifact**: every training run adds a new version of + `jacks-or-better-network`. Its Lineage tab shows which training run + produced each version and which evaluation runs used it. + +## Files + +| File | What it does | +| --- | --- | +| `generate_dataset.py` | Builds the labeled training dataset. | +| `train.py` | Trains the network and logs the run to W&B. | +| `evaluate.py` | Fetches the trained model from W&B, plays hands with it, and logs the score. | +| `dataset.py` | Loads a generated dataset for training. | +| `model.py` | The network, hand encoding, training and validation steps, and checkpoint save/load. | +| `game.py` | Cards, deck, dealing and drawing. | +| `jacks_or_better.py` | Hand rankings and the paytable. | +| `ev.py` | Calculates the exact expected reward of each hold choice. | diff --git a/examples/pytorch/pytorch-video-poker-bot/dataset.py b/examples/pytorch/pytorch-video-poker-bot/dataset.py new file mode 100644 index 00000000..7e09d257 --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/dataset.py @@ -0,0 +1,48 @@ +"""Load a dataset written by generate_dataset.py.""" + +from __future__ import annotations + +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +from game import CARDS_PER_HAND, NUM_HOLDS +from model import encode_cards + + +@dataclass(frozen=True) +class HoldDataset: + cards: np.ndarray # (N, 5) card codes, sorted + targets: np.ndarray # (N, 32) expected profit per credit bet for each hold + + @property + def states(self) -> np.ndarray: + return encode_cards(self.cards) + + def split_samples( + self, + *, + seed: int, + validation_fraction: float = 0.1, + ) -> tuple[np.ndarray, np.ndarray]: + """Shuffle into train / validation (90/10).""" + rng = np.random.default_rng(seed) + order = rng.permutation(len(self.cards)) + val_start = len(order) - int(len(order) * validation_fraction) + return order[:val_start], order[val_start:] + + +def load_dataset(path: Path) -> HoldDataset: + if not path.exists(): + raise FileNotFoundError(f"{path} not found; create it with generate_dataset.py --output {path}") + with np.load(path, allow_pickle=False) as payload: + cards = payload["cards"] + targets = payload["targets"] + + if cards.ndim != 2 or cards.shape[1] != CARDS_PER_HAND: + raise ValueError(f"invalid cards shape: {cards.shape}") + if targets.shape != (len(cards), NUM_HOLDS): + raise ValueError(f"invalid targets shape: {targets.shape}") + + return HoldDataset(cards=cards, targets=targets) diff --git a/examples/pytorch/pytorch-video-poker-bot/ev.py b/examples/pytorch/pytorch-video-poker-bot/ev.py new file mode 100644 index 00000000..eeabbe56 --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/ev.py @@ -0,0 +1,134 @@ +"""Exact expected profit of every hold choice, found by counting, not simulating. + +Once per process, `build_tables` ranks all 2,598,960 five-card hands with the +Jacks or Better classifier. For every group of 0 to 4 cards, it counts how many +five-card hands containing that group land in each hand rank. + +For a dealt hand, `hold_expected_values` looks up those counts for the hand's +32 subsets of cards. The counts include hands that reuse discarded cards, which +can't be drawn again, so inclusion-exclusion over the discards removes them. +Averaging the rewards over every possible draw and subtracting the bet gives +the expected profit of each hold. +""" + +from __future__ import annotations + +from itertools import chain, combinations +from math import comb +from typing import Sequence + +import numpy as np + +from game import BET, CARDS_PER_HAND, DECK_SIZE, NUM_HOLDS, Card, Paytable, hand_to_codes +from jacks_or_better import NUM_RANK_CLASSES, RANK_CLASSES, classify + +TOTAL_HANDS = comb(DECK_SIZE, CARDS_PER_HAND) +ALL_POSITIONS = NUM_HOLDS - 1 # every card position set +CARDS_LEFT_IN_DECK = DECK_SIZE - CARDS_PER_HAND + +_BINOMIAL = np.array( + [[comb(n, k) for k in range(CARDS_PER_HAND + 1)] for n in range(DECK_SIZE + 1)], + dtype=np.int64, +) + + +def colex_index(sorted_codes: np.ndarray, size: int) -> np.ndarray: + index = np.zeros(sorted_codes.shape[0], dtype=np.int64) + for position in range(size): + index += _BINOMIAL[sorted_codes[:, position], position + 1] + return index + + +def colex_index_scalar(sorted_codes: Sequence[int]) -> int: + return sum(comb(code, position + 1) for position, code in enumerate(sorted_codes)) + + +class EVTables: + def __init__(self, subset_counts: Sequence[np.ndarray], hand_ranks: np.ndarray) -> None: + self._subset_counts = tuple(subset_counts) + self._hand_ranks = hand_ranks + + def counts_for(self, sorted_codes: Sequence[int]) -> np.ndarray: + size = len(sorted_codes) + index = colex_index_scalar(sorted_codes) + if size == CARDS_PER_HAND: + counts = np.zeros(NUM_RANK_CLASSES, dtype=np.int64) + counts[int(self._hand_ranks[index])] = 1 + return counts + return self._subset_counts[size][index].astype(np.int64) + + +def build_tables() -> EVTables: + flat = np.fromiter( + chain.from_iterable(combinations(range(DECK_SIZE), CARDS_PER_HAND)), + dtype=np.int8, + count=TOTAL_HANDS * CARDS_PER_HAND, + ) + hands = flat.reshape(TOTAL_HANDS, CARDS_PER_HAND) + ranks = classify(hands).astype(np.int64) + + hand_ranks = np.zeros(TOTAL_HANDS, dtype=np.uint8) + hand_ranks[colex_index(hands, CARDS_PER_HAND)] = ranks.astype(np.uint8) + + subset_counts: list[np.ndarray] = [] + for size in range(CARDS_PER_HAND): + rows = comb(DECK_SIZE, size) + table = np.zeros(rows * NUM_RANK_CLASSES, dtype=np.int64) + for positions in combinations(range(CARDS_PER_HAND), size): + if size == 0: + index = np.zeros(TOTAL_HANDS, dtype=np.int64) + else: + index = colex_index(hands[:, positions], size) + table += np.bincount( + index * NUM_RANK_CLASSES + ranks, + minlength=rows * NUM_RANK_CLASSES, + ) + subset_counts.append(table.reshape(rows, NUM_RANK_CLASSES).astype(np.uint32)) + return EVTables(subset_counts, hand_ranks) + + +_TABLES: EVTables | None = None + + +def get_tables() -> EVTables: + """Build EV tables once per process; reuse in memory afterward.""" + global _TABLES + if _TABLES is None: + _TABLES = build_tables() + return _TABLES + + +def rank_payouts(paytable: Paytable) -> np.ndarray: + """Reward for each hand rank, in the order `classify` numbers them.""" + return np.array([paytable.payout_for(rank) for rank in RANK_CLASSES], dtype=np.int64) + + +def hold_expected_values(hand: Sequence[Card], payouts: np.ndarray) -> list[float]: + """Exact expected profit (in credits) for every hold mask. + + `payouts` comes from `rank_payouts`; build it once and reuse it. + """ + tables = get_tables() + codes = hand_to_codes(hand) + + subset_payout = [0] * NUM_HOLDS + for subset in range(NUM_HOLDS): + sorted_codes = sorted( + codes[position] for position in range(CARDS_PER_HAND) if subset & (1 << position) + ) + subset_payout[subset] = int(tables.counts_for(sorted_codes) @ payouts) + + expected_values: list[float] = [] + for hold_mask in range(NUM_HOLDS): + discards = ALL_POSITIONS ^ hold_mask + total = 0 + subset = discards + while True: + sign = -1 if bin(subset).count("1") % 2 else 1 + total += sign * subset_payout[hold_mask | subset] + if subset == 0: + break + subset = (subset - 1) & discards + draws = comb(CARDS_LEFT_IN_DECK, bin(discards).count("1")) + expected_values.append(total / draws - BET) + return expected_values diff --git a/examples/pytorch/pytorch-video-poker-bot/evaluate.py b/examples/pytorch/pytorch-video-poker-bot/evaluate.py new file mode 100644 index 00000000..d87bc2f1 --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/evaluate.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +"""Fetch a trained network from W&B and play hands with it, logging as it goes.""" + +from __future__ import annotations + +import argparse +import random +from pathlib import Path + +import wandb + +from game import BET, CARDS_PER_HAND, VideoPokerGame, deal, make_deck, shuffle_deck +from jacks_or_better import make_game +from model import Network, choose_hold, load_checkpoint + + +def play_hands( + run: wandb.Run, + model: Network, + game: VideoPokerGame, + *, + hands: int, + seed: int, + log_every: int = 100, +) -> None: + """Play greedy hands, logging and printing running totals every `log_every` hands.""" + rng = random.Random(seed) + wagered = 0 + payout = 0 + for hand_num in range(1, hands + 1): + deck = make_deck() + shuffle_deck(deck, rng=rng) + dealt = deal(deck, CARDS_PER_HAND) + payout += game.play_hand(deck, dealt, choose_hold(model, dealt)) + wagered += BET + if hand_num % log_every == 0 or hand_num == hands: + return_pct = (payout / wagered) * 100 + run.log( + { + "hands_played": hand_num, + "wagered": wagered, + "payout": payout, + "profit": payout - wagered, + "return_pct": return_pct, + }, + step=hand_num, + ) + print(f"hand {hand_num:,}/{hands:,} return={return_pct:.2f}% profit={payout - wagered:,}", flush=True) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--artifact", required=True, help="model artifact from train.py, e.g. my-model:latest") + parser.add_argument("--entity", help="W&B team or user to log to (default: your default entity)") + parser.add_argument("--project", required=True, help="W&B project to log to") + parser.add_argument("--run-name", required=True, help="name for this W&B run") + parser.add_argument("--hands", type=int, default=100_000) + parser.add_argument("--seed", type=int, help="default: random; the seed used is logged to W&B") + args = parser.parse_args() + if args.seed is None: + args.seed = random.randrange(2**32) + + with wandb.init( + entity=args.entity, + project=args.project, + name=args.run_name, + job_type="evaluation", + group="eval", + config={"artifact": args.artifact, "hands": args.hands, "seed": args.seed}, + ) as run: + # use_artifact records that this run consumed the model, linking it to the training run. + checkpoint = Path(run.use_artifact(args.artifact).file()) + model = load_checkpoint(checkpoint) + play_hands(run, model, make_game(), hands=args.hands, seed=args.seed) + + +if __name__ == "__main__": + main() diff --git a/examples/pytorch/pytorch-video-poker-bot/game.py b/examples/pytorch/pytorch-video-poker-bot/game.py new file mode 100644 index 00000000..6b0580b9 --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/game.py @@ -0,0 +1,113 @@ +"""Video poker basics: cards, the deck, dealing and drawing, hand ranks, and +the paytable type. The Jacks or Better rules and rewards are in +jacks_or_better.py. + +The bet is always five credits, baked into BET below. +""" + +from __future__ import annotations + +import random +from dataclasses import dataclass, field +from enum import Enum +from typing import Callable, Mapping, Sequence + +BET = 5 # credits bet per hand +CARDS_PER_HAND = 5 +DECK_SIZE = 52 +NUM_HOLDS = 1 << CARDS_PER_HAND # 32 ways to choose which cards to hold + +RANKS = tuple(range(2, 15)) # 2..14, Ace high +SUITS = (0, 1, 2, 3) +RANK_CHARS = {11: "J", 12: "Q", 13: "K", 14: "A"} +SUIT_CHARS = ("s", "h", "d", "c") + + +@dataclass(frozen=True) +class Card: + rank: int + suit: int + + def __post_init__(self) -> None: + if self.rank not in RANKS: + raise ValueError(f"invalid rank: {self.rank}") + if self.suit not in SUITS: + raise ValueError(f"invalid suit: {self.suit}") + + def __str__(self) -> str: + rank = str(self.rank) if self.rank <= 10 else RANK_CHARS[self.rank] + return f"{rank}{SUIT_CHARS[self.suit]}" + + +def make_deck() -> list[Card]: + return [Card(rank=rank, suit=suit) for suit in SUITS for rank in RANKS] + + +def card_to_code(card: Card) -> int: + """Pack a card into 0..51 (rank-major).""" + return (card.rank - 2) * 4 + card.suit + + +def hand_to_codes(hand: Sequence[Card]) -> list[int]: + return [card_to_code(card) for card in hand] + + +def shuffle_deck(deck: list[Card], rng: random.Random | None = None) -> None: + (rng or random).shuffle(deck) + + +def deal(deck: list[Card], count: int) -> list[Card]: + if count > len(deck): + raise ValueError("not enough cards left in deck") + dealt = deck[:count] + del deck[:count] + return dealt + + +def apply_hold(hand: Sequence[Card], hold_mask: int, deck: list[Card]) -> list[Card]: + if not 0 <= hold_mask < NUM_HOLDS: + raise ValueError(f"hold_mask must be 0-{NUM_HOLDS - 1}") + held = [card for index, card in enumerate(hand) if hold_mask & (1 << index)] + draw_count = CARDS_PER_HAND - len(held) + return held + (deal(deck, draw_count) if draw_count else []) + + +class HandRank(Enum): + ROYAL_FLUSH = "royal_flush" + STRAIGHT_FLUSH = "straight_flush" + FOUR_OF_A_KIND = "four_of_a_kind" + FULL_HOUSE = "full_house" + FLUSH = "flush" + STRAIGHT = "straight" + THREE_OF_A_KIND = "three_of_a_kind" + TWO_PAIR = "two_pair" + JACKS_OR_BETTER = "jacks_or_better" + NOTHING = "nothing" + + +@dataclass(frozen=True) +class EvaluatedHand: + rank: HandRank + + +@dataclass(frozen=True) +class Paytable: + """Maps a hand rank to total credits paid (for a fixed BET).""" + + payouts: Mapping[HandRank, int] = field(default_factory=dict) + + def payout_for(self, rank: HandRank) -> int: + return self.payouts.get(rank, 0) + + +@dataclass +class VideoPokerGame: + paytable: Paytable + evaluate: Callable[[Sequence[Card]], EvaluatedHand] + + def play_hand(self, deck: list[Card], hand: Sequence[Card], hold_mask: int) -> int: + """Hold, draw from `deck`, and return the reward in credits.""" + if len(hand) != CARDS_PER_HAND: + raise ValueError("expected a 5-card hand") + final = apply_hold(hand, hold_mask, deck) + return self.paytable.payout_for(self.evaluate(final).rank) diff --git a/examples/pytorch/pytorch-video-poker-bot/generate_dataset.py b/examples/pytorch/pytorch-video-poker-bot/generate_dataset.py new file mode 100644 index 00000000..36e5c972 --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/generate_dataset.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 +"""Generate the training dataset: random deals labeled with the exact expected +profit of every hold, per credit bet.""" + +from __future__ import annotations + +import argparse +import random +import time +from pathlib import Path + +import numpy as np + +from ev import hold_expected_values, rank_payouts +from game import BET, CARDS_PER_HAND, NUM_HOLDS, card_to_code, make_deck +from jacks_or_better import PAYTABLE + + +def generate_dataset(path: Path, *, hands: int, seed: int) -> None: + deck = make_deck() + rng = random.Random(seed) + payouts = rank_payouts(PAYTABLE) + cards = np.empty((hands, CARDS_PER_HAND), dtype=np.uint8) + targets = np.empty((hands, NUM_HOLDS), dtype=np.float32) + + started = time.perf_counter() + print(f"labeling {hands:,} random deals", flush=True) + for index in range(hands): + # Sorted by card code: the same order model.choose_hold feeds the network. + hand = sorted(rng.sample(deck, CARDS_PER_HAND), key=card_to_code) + cards[index] = [card_to_code(card) for card in hand] + targets[index] = np.asarray(hold_expected_values(hand, payouts)) / BET + if (index + 1) % 100_000 == 0 or index + 1 == hands: + print(f"labeled {index + 1:,}/{hands:,}", flush=True) + + path.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed(path, cards=cards, targets=targets) + print(f"wrote {path} ({path.stat().st_size / 1e6:.1f} MB) in {time.perf_counter() - started:.1f}s") + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output", type=Path, required=True, help="where to write the dataset (.npz)") + parser.add_argument("--hands", type=int, default=1_000_000) + parser.add_argument("--seed", type=int, default=42) + args = parser.parse_args() + generate_dataset(args.output, hands=args.hands, seed=args.seed) + + +if __name__ == "__main__": + main() diff --git a/examples/pytorch/pytorch-video-poker-bot/jacks_or_better.py b/examples/pytorch/pytorch-video-poker-bot/jacks_or_better.py new file mode 100644 index 00000000..575872f9 --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/jacks_or_better.py @@ -0,0 +1,114 @@ +"""9/6 Jacks or Better: ranking, paytable, and game factory. + +`classify` is the single source of truth for hand ranks (vectorized). +`evaluate_hand` uses it to rank a single hand. +""" + +from __future__ import annotations + +from typing import Sequence + +import numpy as np + +from game import ( + CARDS_PER_HAND, + Card, + EvaluatedHand, + HandRank, + Paytable, + VideoPokerGame, + hand_to_codes, +) + +# Absolute credits paid at BET=5 (royal includes the max-coin bonus). +PAYTABLE = Paytable( + payouts={ + HandRank.ROYAL_FLUSH: 4000, + HandRank.STRAIGHT_FLUSH: 250, + HandRank.FOUR_OF_A_KIND: 125, + HandRank.FULL_HOUSE: 45, + HandRank.FLUSH: 30, + HandRank.STRAIGHT: 20, + HandRank.THREE_OF_A_KIND: 15, + HandRank.TWO_PAIR: 10, + HandRank.JACKS_OR_BETTER: 5, + HandRank.NOTHING: 0, + }, +) + +# Index into this tuple is what `classify` returns. +RANK_CLASSES: tuple[HandRank, ...] = ( + HandRank.NOTHING, + HandRank.JACKS_OR_BETTER, + HandRank.TWO_PAIR, + HandRank.THREE_OF_A_KIND, + HandRank.STRAIGHT, + HandRank.FLUSH, + HandRank.FULL_HOUSE, + HandRank.FOUR_OF_A_KIND, + HandRank.STRAIGHT_FLUSH, + HandRank.ROYAL_FLUSH, +) +NUM_RANK_CLASSES = len(RANK_CLASSES) + +_TEN_RANK = 8 +_JACK_RANK = 9 +_ACE_RANK = 12 + + +def classify(hands: np.ndarray) -> np.ndarray: + """Rank many hands at once. Each row is five card codes; result is RANK_CLASSES index.""" + codes = np.asarray(hands) + if codes.ndim != 2 or codes.shape[1] != CARDS_PER_HAND: + raise ValueError("expected an (N, 5) array of card codes") + codes = np.sort(codes, axis=1) + + ranks = codes // 4 + suits = codes % 4 + r0, r1, r2, r3, r4 = (ranks[:, i] for i in range(5)) + + is_flush = (suits == suits[:, :1]).all(axis=1) + distinct = (r0 < r1) & (r1 < r2) & (r2 < r3) & (r3 < r4) + is_run = distinct & ((r4 - r0) == 4) + is_wheel = distinct & (r3 == 3) & (r4 == _ACE_RANK) + is_straight = is_run | is_wheel + is_royal = is_run & (r0 == _TEN_RANK) + + pair_01, pair_12, pair_23, pair_34 = r0 == r1, r1 == r2, r2 == r3, r3 == r4 + adjacent = ( + pair_01.astype(np.int8) + + pair_12.astype(np.int8) + + pair_23.astype(np.int8) + + pair_34.astype(np.int8) + ) + has_quads = (r0 == r3) | (r1 == r4) + has_trips = (r0 == r2) | (r1 == r3) | (r2 == r4) + is_full_house = (adjacent == 3) & ~has_quads + is_trips = (adjacent == 2) & has_trips + is_two_pair = (adjacent == 2) & ~has_trips + pair_rank = np.where(pair_01, r0, np.where(pair_12, r1, np.where(pair_23, r2, r3))) + is_high_pair = (adjacent == 1) & (pair_rank >= _JACK_RANK) + + out = np.zeros(codes.shape[0], dtype=np.uint8) + out[is_high_pair] = RANK_CLASSES.index(HandRank.JACKS_OR_BETTER) + out[is_two_pair] = RANK_CLASSES.index(HandRank.TWO_PAIR) + out[is_trips] = RANK_CLASSES.index(HandRank.THREE_OF_A_KIND) + out[is_straight] = RANK_CLASSES.index(HandRank.STRAIGHT) + out[is_flush] = RANK_CLASSES.index(HandRank.FLUSH) + out[is_full_house] = RANK_CLASSES.index(HandRank.FULL_HOUSE) + out[has_quads] = RANK_CLASSES.index(HandRank.FOUR_OF_A_KIND) + out[is_flush & is_straight] = RANK_CLASSES.index(HandRank.STRAIGHT_FLUSH) + out[is_flush & is_royal] = RANK_CLASSES.index(HandRank.ROYAL_FLUSH) + return out + + +def evaluate_hand(cards: Sequence[Card]) -> EvaluatedHand: + """Classify one hand; translate the fast ordinal into EvaluatedHand.""" + if len(cards) != CARDS_PER_HAND: + raise ValueError("expected exactly 5 cards") + ordinal = int(classify(np.asarray([hand_to_codes(cards)], dtype=np.uint8))[0]) + return EvaluatedHand(RANK_CLASSES[ordinal]) + + +def make_game() -> VideoPokerGame: + return VideoPokerGame(paytable=PAYTABLE, evaluate=evaluate_hand) diff --git a/examples/pytorch/pytorch-video-poker-bot/model.py b/examples/pytorch/pytorch-video-poker-bot/model.py new file mode 100644 index 00000000..de9f2e75 --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/model.py @@ -0,0 +1,146 @@ +"""The network: encode hands, score the 32 holds, train, validate, save and load. + +Everything runs on CPU; the model is tiny. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Sequence + +import numpy as np +import torch +from torch import nn +from torch.nn import functional as F + +from game import CARDS_PER_HAND, NUM_HOLDS, Card, card_to_code + +RANK_FEATURES = 13 +SUIT_FEATURES = 4 +CARD_FEATURES = RANK_FEATURES + SUIT_FEATURES +INPUT_SIZE = CARDS_PER_HAND * CARD_FEATURES + + +def encode_cards(cards: np.ndarray) -> np.ndarray: + """One-hot encode (N, 5) card codes into (N, INPUT_SIZE) network inputs.""" + card_codes = np.asarray(cards) + if card_codes.ndim != 2 or card_codes.shape[1] != CARDS_PER_HAND: + raise ValueError("cards must have shape (N, 5)") + ranks = card_codes // 4 + suits = card_codes % 4 + rows = np.arange(len(card_codes))[:, None] + positions = np.arange(CARDS_PER_HAND)[None, :] + encoded = np.zeros((len(card_codes), CARDS_PER_HAND, CARD_FEATURES), dtype=np.float32) + encoded[rows, positions, ranks] = 1.0 + encoded[rows, positions, RANK_FEATURES + suits] = 1.0 + return encoded.reshape(len(card_codes), INPUT_SIZE) + + +class Network(nn.Module): + """MLP that scores each of the 32 possible hold patterns.""" + + def __init__(self, hidden_size: int) -> None: + super().__init__() + self.hidden_size = hidden_size + self.layers = nn.Sequential( + nn.Linear(INPUT_SIZE, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, hidden_size), + nn.ReLU(), + nn.Linear(hidden_size, NUM_HOLDS), + ) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + return self.layers(inputs) + + +def save_checkpoint(path: Path, model: Network) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + torch.save({"hidden_size": model.hidden_size, "weights": model.state_dict()}, path) + + +def load_checkpoint(path: Path) -> Network: + payload = torch.load(path, map_location="cpu", weights_only=True) + model = Network(payload["hidden_size"]) + model.load_state_dict(payload["weights"]) + model.eval() + return model + + +def choose_hold(model: Network, hand: Sequence[Card]) -> int: + """Pick the highest-scoring hold; returns a hold mask over `hand` as dealt. + + The network sees cards sorted by card code, the same order the dataset uses. + """ + order = sorted(range(CARDS_PER_HAND), key=lambda i: card_to_code(hand[i])) + codes = np.array([[card_to_code(hand[i]) for i in order]]) + with torch.no_grad(): + action = int(model(torch.from_numpy(encode_cards(codes))).argmax(dim=1).item()) + return sum(1 << order[bit] for bit in range(CARDS_PER_HAND) if action & (1 << bit)) + + +def train_epoch( + model: Network, + optimizer: torch.optim.Optimizer, + *, + states: np.ndarray, + targets: np.ndarray, + train_ids: np.ndarray, + batch_size: int, + rng: np.random.Generator, +) -> float: + model.train() + order = rng.permutation(len(train_ids)) + total_loss = 0.0 + n = 0 + for start in range(0, len(order), batch_size): + ids = train_ids[order[start : start + batch_size]] + loss = F.smooth_l1_loss( + model(torch.from_numpy(states[ids])), + torch.from_numpy(targets[ids]), + ) + optimizer.zero_grad() + loss.backward() + nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) + optimizer.step() + total_loss += float(loss.item()) * len(ids) + n += len(ids) + return total_loss / n + + +@torch.no_grad() +def validation_metrics( + model: Network, + *, + states: np.ndarray, + targets: np.ndarray, + sample_ids: np.ndarray, + batch_size: int, +) -> dict[str, float]: + """Loss, % optimal holds, mean regret, and expected return vs the EV labels.""" + model.eval() + total_loss = total_regret = total_policy = 0.0 + optimal = n = 0 + for start in range(0, len(sample_ids), batch_size): + ids = sample_ids[start : start + batch_size] + batch = targets[ids] + preds = model(torch.from_numpy(states[ids])) + total_loss += float( + F.smooth_l1_loss( + preds, torch.from_numpy(batch), reduction="sum" + ).item() + ) + chosen = preds.argmax(dim=1).numpy() + best = batch.max(axis=1) + chosen_values = batch[np.arange(len(ids)), chosen] + regrets = best - chosen_values + total_regret += float(regrets.sum()) + total_policy += float(chosen_values.sum()) + optimal += int(np.count_nonzero(regrets <= 1e-7)) + n += len(ids) + return { + "loss": total_loss / (n * targets.shape[1]), + "optimal_action_pct": (optimal / n) * 100, + "mean_regret": total_regret / n, + "expected_return_pct": (1.0 + total_policy / n) * 100, + } diff --git a/examples/pytorch/pytorch-video-poker-bot/requirements.txt b/examples/pytorch/pytorch-video-poker-bot/requirements.txt new file mode 100644 index 00000000..be506701 --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/requirements.txt @@ -0,0 +1,3 @@ +numpy +torch +wandb diff --git a/examples/pytorch/pytorch-video-poker-bot/train.py b/examples/pytorch/pytorch-video-poker-bot/train.py new file mode 100644 index 00000000..090a0fa2 --- /dev/null +++ b/examples/pytorch/pytorch-video-poker-bot/train.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python3 +"""Train the network to predict each hold's exact expected profit and log the run to W&B.""" + +from __future__ import annotations + +import argparse +import random +from pathlib import Path + +import numpy as np +import torch +import wandb + +from dataset import load_dataset +from model import Network, save_checkpoint, train_epoch, validation_metrics + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--dataset", type=Path, required=True, help=".npz from generate_dataset.py") + parser.add_argument("--checkpoint", type=Path, required=True, help="where to save the trained model (.pt)") + parser.add_argument("--artifact-name", required=True, help="W&B model artifact to upload the checkpoint as") + parser.add_argument("--entity", help="W&B team or user to log to (default: your default entity)") + parser.add_argument("--project", required=True, help="W&B project to log to") + parser.add_argument("--run-name", required=True, help="name for this W&B run") + parser.add_argument("--epochs", type=int, default=20) + parser.add_argument("--batch-size", type=int, default=1024) + parser.add_argument("--lr", type=float, default=1e-3) + parser.add_argument("--hidden-size", type=int, default=256) + parser.add_argument("--seed", type=int, help="default: random; the seed used is logged to W&B") + args = parser.parse_args() + if args.seed is None: + args.seed = random.randrange(2**32) + + torch.manual_seed(args.seed) + + data = load_dataset(args.dataset) + states = data.states + train_ids, val_ids = data.split_samples(seed=args.seed) + model = Network(args.hidden_size) + optimizer = torch.optim.Adam(model.parameters(), lr=args.lr) + rng = np.random.default_rng(args.seed) + + config = { + "epochs": args.epochs, + "batch_size": args.batch_size, + "lr": args.lr, + "hidden_size": args.hidden_size, + "seed": args.seed, + "hands": len(data.cards), + } + + # One place for init → log → finish (finish runs automatically on exit). + with wandb.init( + entity=args.entity, + project=args.project, + name=args.run_name, + group="train", + config=config, + ) as run: + best_regret = float("inf") + for epoch in range(1, args.epochs + 1): + train_loss = train_epoch( + model, + optimizer, + states=states, + targets=data.targets, + train_ids=train_ids, + batch_size=args.batch_size, + rng=rng, + ) + val = validation_metrics( + model, + states=states, + targets=data.targets, + sample_ids=val_ids, + batch_size=args.batch_size, + ) + run.log( + { + "epoch": epoch, + "train_loss": train_loss, + "val_loss": val["loss"], + "val_optimal_action_pct": val["optimal_action_pct"], + "val_mean_regret": val["mean_regret"], + "val_expected_return_pct": val["expected_return_pct"], + }, + step=epoch, + ) + print( + f"epoch {epoch}/{args.epochs} " + f"loss={train_loss:.4f} " + f"val_optimal={val['optimal_action_pct']:.1f}% " + f"val_regret={val['mean_regret']:.4f} " + f"val_return={val['expected_return_pct']:.2f}%" + ) + if val["mean_regret"] < best_regret: + best_regret = val["mean_regret"] + save_checkpoint(args.checkpoint, model) + + artifact = wandb.Artifact(args.artifact_name, type="model") + artifact.add_file(str(args.checkpoint)) + run.log_artifact(artifact) + print(f"saved {args.checkpoint} and logged artifact {args.artifact_name}") + + +if __name__ == "__main__": + main()