nashbench 0.1.0__py3-none-any.whl

nashbench/__init__.py

@@ -0,0 +1,29 @@
+"""Benchmark games and evaluation for Nash equilibrium solvers."""
+
+import importlib.metadata
+
+from nashbench.core import BOTH
+from nashbench.core import Game
+from nashbench.core import State
+from nashbench.core import TimeStep
+from nashbench.core import auto_reset
+from nashbench.exploitability import exploitability
+from nashbench.games import REGISTRY
+from nashbench.games import make
+from nashbench.policy import Policy
+from nashbench.policy import uniform_random
+
+__version__ = importlib.metadata.version("nashbench")
+
+__all__ = [
+    "BOTH",
+    "REGISTRY",
+    "Game",
+    "Policy",
+    "State",
+    "TimeStep",
+    "auto_reset",
+    "exploitability",
+    "make",
+    "uniform_random",
+]

nashbench/core.py

@@ -0,0 +1,216 @@
+"""The game API: `Game`, its `State`, and the `TimeStep` players observe."""
+
+import abc
+from collections.abc import Callable
+import dataclasses
+import functools
+from typing import dataclass_transform
+
+import jax
+import jax.numpy as jnp
+
+from nashbench import sequence_form as sequence_form_lib
+
+#: `TimeStep.current_player` when both players act at once. Equals OpenSpiel's
+#: kSimultaneousPlayerId.
+BOTH = -2
+
+
+@dataclass_transform(frozen_default=True)
+def pytree_dataclass[T](cls: type[T]) -> type[T]:
+    """Makes `cls` a frozen dataclass that JAX transformations accept."""
+    return jax.tree_util.register_dataclass(
+        dataclasses.dataclass(frozen=True)(cls)
+    )
+
+
+@pytree_dataclass
+class TimeStep:
+    """What the players see after `Game.reset` or `Game.step`.
+
+    Arrays with a leading axis of size 2 are indexed by player.
+    `observation` and `legal_action_mask` are defined only for the player to
+    act and only while `done` is false; other entries are unspecified.
+    `reward` is always defined for both players.
+
+    Attributes:
+        observation: `[2, *observation_shape]` float32 information states.
+        legal_action_mask: `[2, num_actions]` bool.
+        reward: `[2]` float32 reward of the last step.
+        done: `[]` bool, true once the episode has ended.
+        current_player: `[]` int32 player to act: 0, 1, or `BOTH`.
+    """
+
+    observation: jax.Array
+    legal_action_mask: jax.Array
+    reward: jax.Array
+    done: jax.Array
+    current_player: jax.Array
+
+
+@pytree_dataclass
+class GameState:
+    """Base class for the state of a game's rules, indexed by seat.
+
+    Attributes:
+        current_seat: `[]` int32 seat to act: 0, 1, or `BOTH`.
+        done: `[]` bool, true once the game has ended.
+    """
+
+    current_seat: jax.Array
+    done: jax.Array
+
+
+@pytree_dataclass
+class State[S: GameState]:
+    """The full state of an episode, including what players cannot see.
+
+    Attributes:
+        key: PRNG key of the next episode, used by `auto_reset`.
+        seat: `[2]` int32 permutation; `seat[p]` is the seat of player `p`.
+        game_state: The state of the game's rules.
+    """
+
+    key: jax.Array
+    seat: jax.Array
+    game_state: S
+
+
+class Game[S: GameState](abc.ABC):
+    """A two-player zero-sum game.
+
+    Players interact with a game through `reset` and `step`. At every reset,
+    a fair coin assigns the two players to the game's two seats, where seat
+    `i` is OpenSpiel's player `i`.
+
+    To add a game, subclass `Game`, set `num_actions` and `observation_shape`,
+    and implement the rules: `initial_states`, `apply_action`, `observe`,
+    `legal_action_mask`, and `returns`. The rules deal with seats only;
+    `reset` and `step` translate between seats and players.
+
+    Attributes:
+        num_actions: Number of actions.
+        observation_shape: Shape of one player's observation.
+    """
+
+    num_actions: int
+    observation_shape: tuple[int, ...]
+
+    def reset(self, key: jax.Array) -> tuple[State[S], TimeStep]:
+        """Starts an episode, assigning players to seats at random."""
+        key, seat_key, chance_key = jax.random.split(key, 3)
+        states, probs = self.initial_states()
+        i = jax.random.choice(chance_key, probs.size, p=probs)
+        state = State(
+            key=key,
+            seat=jax.random.permutation(seat_key, 2),
+            game_state=jax.tree.map(lambda x: x[i], states),
+        )
+        return state, self._timestep(state, reward=jnp.zeros(2))
+
+    def step(
+        self, state: State[S], action: jax.Array
+    ) -> tuple[State[S], TimeStep]:
+        """Advances the episode by one step.
+
+        Args:
+            state: The current state.
+            action: `[2]` int32 action of each player. Actions of players who
+                are not acting are ignored.
+
+        Returns:
+            The next state and timestep. Once the episode is done, `step`
+            leaves the state unchanged and returns zero rewards.
+        """
+        # A permutation of two elements is its own inverse, so indexing with
+        # `seat` maps player-indexed arrays to seat-indexed ones and back.
+        old = state.game_state
+        new = self.apply_action(old, action[state.seat])
+        new = jax.tree.map(lambda o, n: jnp.where(old.done, o, n), old, new)
+        reward = (self.returns(new) - self.returns(old))[state.seat]
+        state = dataclasses.replace(state, game_state=new)
+        return state, self._timestep(state, reward)
+
+    def _timestep(self, state: State[S], reward: jax.Array) -> TimeStep:
+        game_state = state.game_state
+        observe = jax.vmap(self.observe, in_axes=(None, 0))
+        legal_action_mask = jax.vmap(self.legal_action_mask, in_axes=(None, 0))
+        current_seat = game_state.current_seat
+        return TimeStep(
+            observation=observe(game_state, state.seat),
+            legal_action_mask=legal_action_mask(game_state, state.seat),
+            reward=reward.astype(jnp.float32),
+            done=game_state.done,
+            current_player=jnp.where(
+                current_seat == BOTH, BOTH, state.seat[current_seat]
+            ),
+        )
+
+    @abc.abstractmethod
+    def initial_states(self) -> tuple[S, jax.Array]:
+        """Returns every possible initial state, stacked, and its probability.
+
+        All chance events, such as card deals, happen at the start, so initial
+        states differ only in their outcomes.
+        """
+
+    @abc.abstractmethod
+    def apply_action(self, state: S, action: jax.Array) -> S:
+        """Returns the state after the seats take `action`.
+
+        Args:
+            state: The current state. If it is done, the result is discarded.
+            action: `[2]` int32 action of each seat. Turn-based games read
+                `action[state.current_seat]`.
+        """
+
+    @abc.abstractmethod
+    def observe(self, state: S, seat: jax.Array) -> jax.Array:
+        """Returns the information state of `seat` as a float32 tensor.
+
+        The tensor is a one-hot encoding of `seat` followed by OpenSpiel's
+        information-state tensor for that seat. If OpenSpiel's tensor already
+        starts with the seat, it is not repeated.
+        """
+
+    @abc.abstractmethod
+    def legal_action_mask(self, state: S, seat: jax.Array) -> jax.Array:
+        """Returns the `[num_actions]` bool mask of legal actions of `seat`."""
+
+    @abc.abstractmethod
+    def returns(self, state: S) -> jax.Array:
+        """Returns the `[2]` total reward of each seat so far."""
+
+    @functools.cached_property
+    def sequence_form(self) -> sequence_form_lib.SequenceForm:
+        """The sequence form of the game, for exact evaluation.
+
+        Built on first access and cached with the game. By default, enumerates
+        the game tree; games whose tree doesn't fit in memory override this.
+        """
+        return sequence_form_lib.enumerate_tree(self)
+
+
+def auto_reset(
+    game: Game,
+) -> Callable[[State, jax.Array], tuple[State, TimeStep]]:
+    """Returns `game.step`, modified to start a new episode when one ends.
+
+    On the step that ends an episode, the timestep keeps that step's `reward`
+    and `done`. Its other fields, and the returned state, belong to the first
+    step of the next episode.
+    """
+
+    def step(state: State, action: jax.Array) -> tuple[State, TimeStep]:
+        state, timestep = game.step(state, action)
+        next_state, next_timestep = game.reset(state.key)
+        next_timestep = dataclasses.replace(
+            next_timestep, reward=timestep.reward, done=timestep.done
+        )
+        return jax.tree.map(
+            lambda n, o: jnp.where(timestep.done, n, o),
+            (next_state, next_timestep),
+            (state, timestep),
+        )
+
+    return step

nashbench/exploitability.py

@@ -0,0 +1,90 @@
+"""Exploitability: how far a policy is from a Nash equilibrium."""
+
+import functools
+
+import jax
+import jax.numpy as jnp
+
+from nashbench import core
+from nashbench import policy as policy_lib
+from nashbench import sequence_form
+
+# Policies see information states in chunks of this size.
+_CHUNK_SIZE = 2**16
+
+
+def exploitability(game: core.Game, policy: policy_lib.Policy) -> float:
+    """Returns the exploitability of `policy` playing both seats of `game`.
+
+    Exploitability is the value a best response gains against `policy`,
+    averaged over the two seats:
+
+        (max_π u_0(π, policy) + max_π u_1(policy, π)) / 2,
+
+    where `u_i(σ_0, σ_1)` is the expected return of seat `i` when seat 0
+    plays `σ_0` and seat 1 plays `σ_1`. It is 0 exactly when `policy` is a
+    Nash equilibrium, and equals OpenSpiel's `exploitability.exploitability`
+    (NashConv / 2).
+
+    The first call for a game builds its sequence form, which takes up to a
+    minute for the largest games; later calls with the same game reuse it.
+
+    Args:
+        game: The game.
+        policy: The policy, playing both seats.
+
+    Returns:
+        The exploitability, in units of the game's rewards.
+
+    Raises:
+        ValueError: If the policy doesn't return probabilities over legal
+            actions.
+    """
+    form = game.sequence_form
+    probs = _action_probabilities(form, policy)
+    plan = _realization_plan(probs, form.parent, form.levels)
+    values = _best_response_values(
+        form.gradient(plan), form.parent, form.legal_action_mask, form.levels
+    )
+    return float(values.sum() / 2)
+
+
+def _action_probabilities(
+    form: sequence_form.SequenceForm, policy: policy_lib.Policy
+) -> jax.Array:
+    """Returns `[num_infosets, num_actions]` probabilities of `policy`."""
+    num_infosets = form.parent.size
+    chunks = []
+    for start in range(0, num_infosets, _CHUNK_SIZE):
+        ids = jnp.arange(start, min(start + _CHUNK_SIZE, num_infosets))
+        mask = form.legal_action_mask[ids]
+        chunks.append(jax.vmap(policy)(form.observe(ids), mask))
+    probs = jnp.concatenate(chunks)
+    illegal = jnp.where(form.legal_action_mask, 0.0, probs).max()
+    if illegal > 1e-6 or jnp.abs(probs.sum(1) - 1).max() > 1e-4:
+        raise ValueError(
+            "The policy must return probabilities of legal actions."
+        )
+    return probs
+
+
+@functools.partial(jax.jit, static_argnames="levels")
+def _realization_plan(probs, parent, levels):
+    """Returns the realization plan of both seats, top-down by level."""
+    plan = jnp.ones((probs.shape[0] + 1, probs.shape[1]))
+    for start, end in levels:
+        reach = plan.ravel()[parent[start:end], None] * probs[start:end]
+        plan = plan.at[1 + start : 1 + end].set(reach)
+    return plan.ravel()
+
+
+@functools.partial(jax.jit, static_argnames="levels")
+def _best_response_values(gradient, parent, legal_action_mask, levels):
+    """Returns the `[2]` value of each seat's best response, bottom-up."""
+    values = gradient.reshape(-1, legal_action_mask.shape[1])
+    for start, end in reversed(levels):
+        q = values[1 + start : 1 + end]
+        best = jnp.where(legal_action_mask[start:end], q, -jnp.inf).max(1)
+        values = values.ravel().at[parent[start:end]].add(best)
+        values = values.reshape(-1, legal_action_mask.shape[1])
+    return values[0, :2]

nashbench/games/__init__.py

@@ -0,0 +1,38 @@
+"""The benchmark games, by name."""
+
+from collections.abc import Callable
+import functools
+
+from nashbench import core
+from nashbench.games import dark_hex
+from nashbench.games import goofspiel
+from nashbench.games import kuhn_poker
+from nashbench.games import leduc_poker
+from nashbench.games import phantom_ttt
+
+#: Constructor of each game, by name.
+REGISTRY: dict[str, Callable[..., core.Game]] = {
+    "kuhn_poker": kuhn_poker.KuhnPoker,
+    "leduc_poker": leduc_poker.LeducPoker,
+    "phantom_ttt": functools.partial(phantom_ttt.PhantomTTT, abrupt=False),
+    "phantom_ttt_abrupt": functools.partial(
+        phantom_ttt.PhantomTTT, abrupt=True
+    ),
+    "dark_hex3": functools.partial(dark_hex.DarkHex3, abrupt=False),
+    "dark_hex3_abrupt": functools.partial(dark_hex.DarkHex3, abrupt=True),
+    "goofspiel": goofspiel.Goofspiel,
+}
+
+
+def make(name: str, **kwargs) -> core.Game:
+    """Returns the game named `name`, a key of `REGISTRY`.
+
+    Args:
+        name: Name of the game.
+        **kwargs: Game parameters, such as `num_cards` of `"goofspiel"`.
+    """
+    if name not in REGISTRY:
+        raise ValueError(
+            f"Unknown game {name!r}; choose from {list(REGISTRY)}."
+        )
+    return REGISTRY[name](**kwargs)

nashbench/games/_goofspiel_tree.py

@@ -0,0 +1,240 @@
+"""The sequence form of Goofspiel, from the regular structure of its tree.
+
+At turn `t`, a seat's information set is the first `t + 1` point cards, its
+first `t` bids, and who won each of the first `t` turns. Point cards are
+independent of bids, so the information sets of turn `t` are numbered
+
+    offset[t] + rank(point cards) * num_pairs[t] + pair,
+
+where `pair` numbers the (bids, winners) that some opponent bids make
+possible. A terminal history is a triple of permutations: the point cards and
+each seat's bids.
+"""
+
+import dataclasses
+import functools
+import itertools
+import math
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from nashbench import sequence_form as sequence_form_lib
+
+# Information sets per chunk when building tables.
+_CHUNK = 2**22
+
+
+def sequence_form(game) -> sequence_form_lib.SequenceForm:
+    """Returns the sequence form of a Goofspiel game."""
+    k = game.num_cards
+    tables = _Tables(k)
+    ids = jnp.arange(2 * tables.num_infosets)
+    parent, legal = (
+        jnp.concatenate(x)
+        for x in zip(
+            *(
+                _parent_and_legal(tables, ids[i : i + _CHUNK])
+                for i in range(0, ids.size, _CHUNK)
+            ),
+            strict=True,
+        )
+    )
+    levels = tuple(
+        (
+            s * tables.num_infosets + int(tables.offset[t]),
+            s * tables.num_infosets + int(tables.offset[t + 1]),
+        )
+        for t in range(k - 1)
+        for s in range(2)
+    )
+    observe = jax.jit(functools.partial(_observe, game, tables))
+    return sequence_form_lib.SequenceForm(
+        levels=levels,
+        parent=parent,
+        legal_action_mask=legal,
+        observe=observe,
+        gradient=jax.jit(functools.partial(_gradient, tables)),
+    )
+
+
+class _Tables:
+    """Small tables that number Goofspiel's information sets and leaves.
+
+    Attributes:
+        k: Number of cards.
+        perms: `[k!, k]` permutations of the cards, in lexicographic order.
+            The first `m` cards of permutation `i` have rank `i // (k - m)!`
+            among partial permutations of length `m`.
+        pair_of, pair_of_base: Maps of every turn `t`, concatenated, and where
+            each starts. A map takes `rank(bids) * 3**t + winners` to the pair,
+            or -1 if no opponent bids produce those winners. Winners are base-3
+            digits, one per turn: 0 lost, 1 tied, 2 won.
+        last_pair_of: The map of the last decision turn, as a numpy array.
+        pair_bids, pair_opponent, pair_base: Per pair of every turn,
+            concatenated, the seat's bids and opponent bids that produce its
+            winners, padded with -1, and where each turn's pairs start.
+        num_pairs, offset: Per turn, the number of pairs and the number of
+            information sets before the turn's.
+        num_infosets: Number of information sets of a seat.
+    """
+
+    def __init__(self, k):
+        self.k = k
+        perms = np.array(list(itertools.permutations(range(k))))
+        pair_of, pair_bids, pair_opponent, num_pairs = [], [], [], []
+        for t in range(k - 1):  # The last turn plays itself.
+            mine = perms[:: math.factorial(k - t), :t]
+            witness = {}
+            for rank, bids in enumerate(mine):
+                for theirs in mine:
+                    key = rank * 3**t + _code(np.sign(bids - theirs) + 1)
+                    witness.setdefault(int(key), theirs)
+            keys = sorted(witness)
+            table = np.full(len(mine) * 3**t, -1)
+            table[keys] = np.arange(len(keys))
+            pair_of.append(table)
+            for key in keys:
+                pair_bids.append(
+                    np.pad(mine[key // 3**t], (0, k - t), constant_values=-1)
+                )
+                pair_opponent.append(
+                    np.pad(witness[key], (0, k - t), constant_values=-1)
+                )
+            num_pairs.append(len(keys))
+        self.offset = np.cumsum(
+            [0] + [math.perm(k, t + 1) * n for t, n in enumerate(num_pairs)]
+        )
+        self.num_infosets = int(self.offset[-1])
+        self.num_pairs = num_pairs
+        self.perms = jnp.asarray(perms)
+        self.last_pair_of = pair_of[-1]
+        self.pair_of = jnp.asarray(np.concatenate(pair_of))
+        self.pair_of_base = jnp.asarray(
+            np.cumsum([0] + [p.size for p in pair_of])
+        )
+        self.pair_base = jnp.asarray(np.cumsum([0, *num_pairs]))
+        self.pair_bids = jnp.asarray(np.array(pair_bids))
+        self.pair_opponent = jnp.asarray(np.array(pair_opponent))
+
+    def describe(self, ids):
+        """Returns what identifies information sets, and a way to reach them.
+
+        Returns each information set's seat, turn, point-card rank, point
+        cards, bids, and opponent bids that produce its winners.
+        """
+        k, offset = self.k, jnp.asarray(self.offset)
+        seat, local = jnp.divmod(ids, self.num_infosets)
+        turn = jnp.searchsorted(offset, local, side="right") - 1
+        num_pairs = jnp.asarray(self.num_pairs)[turn]
+        rank, pair = jnp.divmod(local - offset[turn], num_pairs)
+        pair = self.pair_base[turn] + pair
+        # Completing the point cards in increasing order gives a permutation.
+        factorial = jnp.asarray([math.factorial(n) for n in range(k + 1)])
+        points = self.perms[rank * factorial[k - turn - 1]]
+        return (
+            seat,
+            turn,
+            rank,
+            points,
+            self.pair_bids[pair],
+            self.pair_opponent[pair],
+        )
+
+
+def _code(digits):
+    """Returns base-3 numbers of rows of digits, first digit most significant."""
+    return (digits * 3 ** np.arange(digits.shape[-1])[::-1]).sum(-1)
+
+
+@functools.partial(jax.jit, static_argnums=0)
+def _parent_and_legal(tables, ids):
+    """Returns the parent sequence and legal actions of information sets."""
+    k = tables.k
+    seat, turn, rank, _, bids, opponent = tables.describe(ids)
+    # The parent drops the last turn: its point card, bid, and winner.
+    previous = jnp.maximum(turn - 1, 0)
+    active = jnp.arange(k) < previous[:, None]
+    own_rank = jnp.zeros_like(ids)
+    for i in range(k):
+        smaller_unused = (jnp.arange(k) < bids[:, i : i + 1]).sum(1) - (
+            (bids[:, :i] < bids[:, i : i + 1]).sum(1)
+        )
+        own_rank = jnp.where(
+            active[:, i], own_rank * (k - i) + smaller_unused, own_rank
+        )
+    digits = jnp.where(active, jnp.sign(bids - opponent) + 1, 0)
+    power = jnp.maximum(previous[:, None] - 1 - jnp.arange(k), 0)
+    winners = (digits * 3**power).sum(1)
+    pair = tables.pair_of[
+        tables.pair_of_base[previous] + own_rank * 3**previous + winners
+    ]
+    parent_rank = rank // (k - turn)
+    offset, num_pairs = (
+        jnp.asarray(tables.offset),
+        jnp.asarray(tables.num_pairs),
+    )
+    parent = (
+        seat * tables.num_infosets
+        + offset[previous]
+        + parent_rank * num_pairs[previous]
+        + pair
+    )
+    last_bid = jnp.take_along_axis(bids, previous[:, None], 1)[:, 0]
+    parent = jnp.where(turn == 0, seat, k * (parent + 1) + last_bid)
+    legal = ~(bids[:, :, None] == jnp.arange(k)).any(1)
+    return parent.astype(jnp.int32), legal
+
+
+def _gradient(tables, plan):
+    """Returns the gradient of `plan`, summing over all terminal histories."""
+    k, perms = tables.k, tables.perms
+    t = k - 2  # Last decision turn.
+    # For every pair of bid permutations: each seat's pair and last decision,
+    # and who won each turn (+1 seat 0, -1 seat 1).
+    i0, i1 = jnp.divmod(jnp.arange(perms.shape[0] ** 2), perms.shape[0])
+    b0, b1 = perms[i0], perms[i1]
+    code = 3 ** jnp.arange(t)[::-1]
+    last_pair_of = jnp.asarray(tables.last_pair_of)
+    pair0 = last_pair_of[
+        (i0 // 2) * 3**t + ((jnp.sign(b0 - b1) + 1)[:, :t] * code).sum(1)
+    ]
+    pair1 = last_pair_of[
+        (i1 // 2) * 3**t + ((jnp.sign(b1 - b0) + 1)[:, :t] * code).sum(1)
+    ]
+    won = jnp.sign(b0 - b1)
+    offset, num_pairs = int(tables.offset[t]), tables.num_pairs[t]
+    chunk = math.gcd(8, perms.shape[0])  # Point-card permutations per step.
+
+    def body(i, gradient):
+        points = jax.lax.dynamic_slice_in_dim(perms, i * chunk, chunk) + 1
+        weight = jnp.sign(won @ points.T) / perms.shape[0]  # [pairs, chunk]
+        infoset = offset + (i * chunk + jnp.arange(chunk)) * num_pairs
+        seq0 = k * (infoset + pair0[:, None] + 1) + b0[:, t, None]
+        seq1 = (
+            k * (tables.num_infosets + infoset + pair1[:, None] + 1)
+            + b1[:, t, None]
+        )
+        gradient = gradient.at[seq0].add(weight * plan[seq1])
+        return gradient.at[seq1].add(-weight * plan[seq0])
+
+    return jax.lax.fori_loop(
+        0, perms.shape[0] // chunk, body, jnp.zeros_like(plan)
+    )
+
+
+def _observe(game, tables, ids):
+    """Returns observations of information sets, from representative states."""
+    seat, turn, _, points, bids, opponent = tables.describe(ids)
+    states = jax.tree.map(
+        lambda x: jnp.repeat(x[:1], ids.shape[0], 0), game.initial_states()[0]
+    )
+    own = jnp.arange(2)[None, :, None] == seat[:, None, None]
+    states = dataclasses.replace(
+        states,
+        point_cards=points,
+        bids=jnp.where(own, bids[:, None], opponent[:, None]),
+        turn=turn,
+    )
+    return jax.vmap(game.observe)(states, seat)