Blackjack simulator where player policy is optimized using the Monte Carlo method. As described in Chapter 5(.3) of Reinforcement Learning, an introduction by Sutton and Barto

blackjack_monte_carlo_optimizer.py

"""
Blackjack simulator where player policy is optimized using the Monte Carlo method.

As described in Chapter 5(.3) of Reinforcement Learning, an introduction by Sutton and Barto
"""
from dataclasses import dataclass, field
from typing import List, Tuple
import numpy as np

DECK = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 10, 10, 10])


@dataclass
class Person:
    """Superclass for both Dealer and Player."""
    has_usable_ace: bool = False
    total: int = 0

    def hit(self):
        """Take a card from an (infinite) deck, check if there is a usable Ace."""
        card = np.random.choice(DECK)
        # card 1 = Ace.
        if card == 1 and (self.total + 11) <= 21:
            self.has_usable_ace = True
            self.total += 11
        else:
            self.total += card
            # Spend the usable ace to decrease total if over 21
            if self.total > 21 and self.has_usable_ace:
                self.has_usable_ace = False
                self.total -= 10


@dataclass
class Player(Person):
    threshold: int = 20
    state_actions: List[Tuple[int]] = field(default_factory=list)
    first_action: bool = True

    def __post_init__(self) -> None:
        """Sample random starting state."""
        self.total = np.random.randint(12, 22)
        self.has_usable_ace = np.random.choice([True, False])

    def play(self, dealer_card_shown: int, policy: np.ndarray) -> None:
        """Apply strategy to hit or stick given total score, usable ace, dealer card shown."""
        if self.first_action:
            if np.random.random() < 0.5:
                self.state_actions.append(tuple([self.total, dealer_card_shown, int(self.has_usable_ace), 1]))
                self.hit()
            else:
                # Stick
                self.state_actions.append(tuple([self.total, dealer_card_shown, int(self.has_usable_ace), 0]))

        while self.total <= 21:
            action = policy[self.total - 12, (dealer_card_shown - 1) % 10, int(self.has_usable_ace)]

            if action == 1:
                self.state_actions.append(tuple([self.total, dealer_card_shown, int(self.has_usable_ace), 1]))
                self.hit()
            else:
                self.state_actions.append(tuple([self.total, dealer_card_shown, int(self.has_usable_ace), 0]))
                # Stick
                break


@dataclass
class Dealer(Person):
    threshold: int = 17
    showing_card: int = field(init=False)

    def __post_init__(self) -> None:
        """Take two start hits and make first card visible."""
        self.hit()
        self.showing_card = self.total
        self.hit()

    def play(self) -> None:
        """Dealer has fixed strategy, hit till past threshold."""
        while self.total < self.threshold:
            self.hit()


@dataclass
class Simulation:
    n_iter: int
    visits: np.ndarray = field(default_factory=lambda: np.zeros((10, 10, 2, 2)))
    q_table: np.ndarray = field(default_factory=lambda: np.random.uniform(-0.001, 0.001, (10, 10, 2, 2)))
    policy: np.ndarray = field(default_factory=lambda: np.zeros((10, 10, 2)))

    def __post_init__(self):
        """Set initial policy"""
        self.policy[:8, :, :] = 1

    def run(self) -> None:
        """Iterate through episodes (games) and see who won, then attribute rewards and count state visits."""
        for i in range(self.n_iter):
            p = Player()
            d = Dealer()
            p.play(d.showing_card, self.policy)

            if p.total > 21:
                reward = -1
            else:
                d.play()
                if p.total == d.total:
                    reward = 0
                elif p.total > d.total or d.total > 21:
                    reward = 1
                else:
                    reward = -1

            for state in p.state_actions:
                self.visits[state[0] - 12, (state[1] - 1) % 10, state[2], state[3]] += 1
                # Get the expected reward given the current state and action
                expected_reward = self.q_table[state[0] - 12, (state[1] - 1) % 10, state[2], state[3]]
                n_visits = self.visits[state[0] - 12, (state[1] - 1) % 10, state[2], state[3]]
                reward_effect = (reward - expected_reward) / n_visits
                self.q_table[state[0] - 12, (state[1] - 1) % 10, state[2], state[3]] += reward_effect
                # For each state, greedily select the action with the highest expected reward as the policy.
                self.policy = np.argmax(self.q_table, axis=-1)

            if (i + 1) % (self.n_iter // 20) == 0:
                print(f'Simulation {i / self.n_iter:4.0%} complete')

    @staticmethod
    def plot_heatmap(array: np.ndarray, title: str) -> None:
        import seaborn as sns
        import matplotlib.pylab as plt
        ax = sns.heatmap(np.flip(array, 0), linewidth=0.5, vmin=-1, vmax=1, center=0, cmap='coolwarm_r', annot=True, fmt=".2f")
        plt.ylabel('Player hand')
        plt.xlabel('Dealer card shown')
        plt.title(title)
        plt.xticks(np.arange(10) + 0.5, ['Ace'] + list(range(2, 11)))
        plt.yticks(np.arange(10) + 0.5, np.flip(np.arange(10) + 12))
        plt.savefig(f'{title}.png')
        plt.show()
        plt.close()

    def evaluate(self) -> None:
        """Per state, divide summed reward by visits to get average reward of state."""
        value_optimal_policy = np.max(self.q_table, axis=-1)
        self.plot_heatmap(self.policy[:, :, 0], 'Policy without usable ace')
        self.plot_heatmap(self.policy[:, :, 1], 'Policy with usable ace')
        self.plot_heatmap(self.q_table[:, :, 0, 0], 'Q table without usable ace, stick')
        self.plot_heatmap(self.q_table[:, :, 1, 0], 'Q table with usable ace, stick')
        self.plot_heatmap(self.q_table[:, :, 0, 1], 'Q table without usable ace, hit')
        self.plot_heatmap(self.q_table[:, :, 1, 1], 'Q table with usable ace, hit')
        self.plot_heatmap(value_optimal_policy[:, :, 0], 'V* without usable ace')
        self.plot_heatmap(value_optimal_policy[:, :, 1], 'V* with usable ace')


if __name__ == "__main__":
    np.random.seed(1)
    s = Simulation(15000000)
    s.run()
    s.evaluate()