Ref/type-python

liar_dice/__init__.py

@@ -0,0 +1,38 @@
+from dataclasses import dataclass
+from typing import Literal, TypeAlias, TypedDict
+
+from torch import Tensor
+from torch.optim.optimizer import StateDict
+
+Roll: TypeAlias = tuple[int, ...]
+PlayerId: TypeAlias = Literal[0, 1]
+MatchResult: TypeAlias = Literal[-1, 1]
+
+State = Tensor  # Public state. Includes previous actions. Initialized as `zeros(self.D_PUB)`.
+Priv = Tensor
+# Private state of a player. Initialized as `zeros(self.D_PRI)`. including the perspective for the scores
+
+Variant = Literal["normal", "joker", "stairs"]
+PruneType = Literal["zero", "upper", "lower", "us", "them", "avg"]
+
+
+@dataclass
+class TrainArg:
+    d1: int
+    d2: int
+    sides: int = 6
+    variant: Variant = "normal"
+    eps: float = 1e-2
+    layers: int = 4
+    layer_size: int = 100
+    lr: float = 1e-3
+    w: float = 1e-2
+    path: str = "model.pt"
+    device: Literal["cpu", "cuda"] = "cpu"
+
+
+class CheckpointFile(TypedDict):
+    epoch: int
+    model_state_dict: StateDict
+    optimizer_state_dict: StateDict
+    args: TrainArg

code/lbr.py → liar_dice/lbr.py

@@ -2,17 +2,29 @@
 Computes the (possibly Local) Best Response strategy and exploitability of a model
 """
 
-import torch
 import argparse
-
-from snyd import Net, Game, calc_args
 from collections import Counter
+from typing import cast
+
+import torch
+
+from liar_dice import PlayerId, PruneType, Roll
+from liar_dice.snyd import Game, Net, calc_args
 
 # in normal form games, given strategy matrix A for player 1 (hand -> action_probs)
 # we will play a deterministic strategy b (hand -> action)
 
 
-def lbr(game, roll, us, prune=0, prune_type="zero"):
+def lbr(
+    game: Game,
+    roll: Roll,
+    us: PlayerId,
+    prune: float = 0,
+    prune_type: PruneType = "zero",
+) -> float:
+    """
+    Compute the Local Best Response strategy for a player in a game.
+    """
     # All possible opponent rolls - with multiplicity
     roll_cnt = Counter(game.rolls(1 - us))
     op_rolls = roll_cnt.keys()
@@ -23,7 +35,13 @@ def lbr(game, roll, us, prune=0, prune_type="zero"):
     # The opponents private state in each initial information state
     op_privs = [game.make_priv(op_roll, 1 - us) for op_roll in op_rolls]
 
-    def inner(state, reach_probs, likelihood):
+    def inner(
+        state: torch.Tensor, reach_probs: torch.Tensor, likelihood: float
+    ) -> float:
+        """
+        Compute the Local Best Response strategy for a player in a game.
+        """
+
         # basically simple max-chance algorithm.
         # possibly using our learned values for reducing the tree height
         # But I could start by making a version that just goes to the bottom
@@ -32,42 +50,50 @@ def inner(state, reach_probs, likelihood):
         cur = len(calls) % 2
 
         if likelihood < prune:
+
+            def calc_our_guess():
+                """
+                Even using our guess here, we will probably get a lower bound,
+                since we are replacing a "perfect" (best response) strategy
+                with whatever strategy the model suggests.
+                However it's not guaranteed to be (a lower bound), since the
+                value network could be overly optimistic."""
+                return game.model(game.make_priv(roll, us), state).item()
+
+            def calc_op_guess():
+                """
+                This is the same as calc_our_guess, but for the opponent.
+                Opponent guess is relative to them, not us, so it's negated.
+                """
+                return -sum(
+                    prob * game.model(op_priv, state).item()
+                    for prob, op_priv in zip(reach_probs, op_privs)
+                )
+
             # print('Pruning', likelihood)
             if prune_type == "zero":
                 return 0
             # I can return +1 or -1 here to get either an upper or lower
             # bound on exploitability
-            if prune_type == "upper":
+            elif prune_type == "upper":
                 return 1
-            if prune_type == "lower":
+            elif prune_type == "lower":
                 return -1
-            if prune_type in ["us", "avg"]:
-                # Even using our guess here, we will probably get a lower bound,
-                # since we are replacing a "perfect" (best response) strategy
-                # with whatever strategy the model suggests.
-                # However it's not guaranteed to be (a lower bound), since the
-                # value network could be overly optimistic.
-                our_guess = game.model(game.make_priv(roll, us), state).item()
-            if prune_type in ["them", "avg"]:
-                op_guess = sum(
-                    prob * game.model(op_priv, state).item()
-                    for prob, op_priv in zip(reach_probs, op_privs)
-                )
-                # Remember opponent guess is relative to them, not us
-                op_guess = -op_guess
-            if prune_type == "us":
-                return our_guess
-            if prune_type == "them":
-                return op_guess
-            if prune_type == "avg":
-                return (our_guess + op_guess) / 2
-            assert False
+            elif prune_type == "us":
+                return calc_our_guess()
+            elif prune_type == "them":
+                return calc_op_guess()
+            elif prune_type == "avg":
+                return (calc_our_guess() + calc_op_guess()) / 2
+            else:
+                raise ValueError(f"Unknown prune type: {prune_type}")
 
         if calls and calls[-1] == game.LIE_ACTION:
             prev_call = calls[-2] if len(calls) >= 2 else -1
             res = 0
             # Take the average outcome over our opponents possible rolls
             for prob, op_roll in zip(reach_probs, op_rolls):
+                prob = cast(float, prob)
                 r1, r2 = (roll, op_roll) if us == 0 else (op_roll, roll)
                 correct = game.evaluate_call(r1, r2, prev_call)
                 # If prev_call is good, and we are now, it mean we won
@@ -106,9 +132,12 @@ def inner(state, reach_probs, likelihood):
             score = 0
             for action in range(last_call + 1, game.N_ACTIONS):
                 ai = action - last_call - 1
-                pa = reach_probs @ policies[:, ai]
-                if torch.isclose(pa, torch.tensor(0.0)):
+                pa_mat = reach_probs @ policies[:, ai]
+
+                if torch.isclose(pa_mat, torch.tensor(0.0)):
                     continue
+
+                pa = float(pa_mat)
                 # Bayes: P(R|A) = P(R)P(A|R)/P(A)
                 # P(A) = P(R)P(A|R)/sum(P(A|r)P(r) for r in Rs)
                 new_reach_probs = (reach_probs * policies[:, ai]) / pa
@@ -135,7 +164,7 @@ def main():
     # options=['upper', 'lower', 'zero', 'us', 'them', 'avg'])
     args = parser.parse_args()
 
-    checkpoint = torch.load(args.path)
+    checkpoint = torch.load(args.path)  # type: ignore
     train_args = checkpoint["args"]
 
     D_PUB, D_PRI, *_ = calc_args(
@@ -147,12 +176,14 @@ def main():
         model, train_args.d1, train_args.d2, train_args.sides, train_args.variant
     )
 
-    for player in range(2):
+    for player in (0, 1):
         print(f"Testing exploitability of player {1-player}")
 
         total_val = 0
         total_cnt = 0
         for roll, cnt in Counter(game.rolls(player)).items():
+            if cnt != 1:
+                print("Warning: Duplicate roll", roll)
             print("Exploiting with roll", roll)
             total_val += cnt * lbr(game, roll, player, args.prune, args.type)
             total_cnt += cnt

code/play_model.py → liar_dice/play_model.py

@@ -2,23 +2,19 @@
 # This script allows you to play against the model from the terminal
 ################################################################################
 
-import random
-import torch
-from torch import nn
-import itertools
-import numpy as np
-import math
-from collections import Counter
 import argparse
+import random
 import re
 
-from snyd import *
+import torch
+
+from liar_dice.snyd import *
 
 parser = argparse.ArgumentParser()
 parser.add_argument("path", type=str, help="Path of model")
 args = parser.parse_args()
 
-checkpoint = torch.load(args.path, map_location=torch.device("cpu"))
+checkpoint = torch.load(args.path, map_location=torch.device(device="cpu"))  # type: ignore
 train_args = checkpoint["args"]
 
 D_PUB, D_PRI, *_ = calc_args(
@@ -30,7 +26,7 @@
 
 
 class Human:
-    def get_action(self, state):
+    def get_action(self, state: State) -> int:
         last_call = game.get_last_call(state)
         while True:
             call = input('Your call [e.g. 24 for 2 fours, or "lie" to call a bluff]: ')
@@ -53,18 +49,18 @@ def __repr__(self):
 
 
 class Robot:
-    def __init__(self, priv):
+    def __init__(self, priv: Priv):
         self.priv = priv
 
-    def get_action(self, state):
+    def get_action(self, state: State) -> int:
         last_call = game.get_last_call(state)
         return game.sample_action(self.priv, state, last_call, eps=0)
 
     def __repr__(self):
         return "robot"
 
 
-def repr_action(action):
+def repr_action(action: int) -> str:
     action = int(action)
     if action == -1:
         return "nothing"
@@ -77,9 +73,9 @@ def repr_action(action):
 
 while True:
     while (ans := input("Do you want to go first? [y/n/r] ")) not in ["y", "n", "r"]:
-        pass
+        pass  # r: bot vs bot
 
-    r1 = random.choice(list(game.rolls(0)))
+    r1 = random.choice(list(game.rolls(player=0)))
     r2 = random.choice(list(game.rolls(1)))
     privs = [game.make_priv(r1, 0), game.make_priv(r2, 1)]
     state = game.make_state()
@@ -92,10 +88,12 @@ def repr_action(action):
         players = [Robot(privs[0]), Human()]
     elif ans == "r":
         players = [Robot(privs[0]), Robot(privs[1])]
+    else:
+        raise ValueError(f"Invalid answer: {ans}")
 
     cur = 0
     while True:
-        action = players[cur].get_action(state)
+        action: int = players[cur].get_action(state)
         print()
         print(f"> The {players[cur]} called {repr_action(action)}!")