<jupyter_start><jupyter_text>`pymdptoolbox` demo<jupyter_code>import warnings

from mdptoolbox import mdp
import numpy as np
import scipy as sp
import matplotlib.pyplot as plt<jupyter_output><empty_output><jupyter_text>The problem* You have a 20-sided die, and you get to roll repeatedly until the sum of your rolls either gets as close as possible to 21 or you bust.* Your score is the numerical value of the sum of your rolls; if you bust, you get zero.* What is the optimal strategy? ![state diagram](d20_mdp.png)  The solution Let's look at what we have to deal with:* State space is 23-dimensional (sum of rolls can be 0-21 inclusive, plus the terminal state)* Action space is 2-dimensional (roll/stay)* State transitions are stochastic; requires transition matrix $T(s^\prime;s,a)$* $T$ is mildly sparse (some transitions like 9->5 or 0->21 are impossible)* Rewards depend on both state and action taken from that state, but are not stochastic (only ever get positive reward when choosing "stay") We're going to use the [*value iteration*](https://pymdptoolbox.readthedocs.io/en/latest/api/mdp.htmlmdptoolbox.mdp.ValueIteration) algorithm. Looking at the documentation, we can see that it requires as input a transition matrx, a reward matrix, and a discount factor (we will use $\gamma = 1$). Let's first specify the transition "matrix". It's going to be a 3-dimensional tensor of shape $(|\mathcal{A}|,|\mathcal{S}|,|\mathcal{S}|) = (2, 23, 23)$. Most entries are probably zero, so let's start with a zero matrix and fill in the blanks. I'm going reserve the very last state (the 23rd entry) for the terminal state.<jupyter_code>def make_transition_matrix(n_sides=20, max_score=21):
    """Constructs the transition matrix for the MDP
    
    Arguments:
        n_sides: number of sides on the die being rolled
        max_score: the maximum score of the game before going bust
    Returns:
        np.ndarray: array of shape (A,S,S), where A=2, and S=max_score+2
            representing the transition matrix for the MDP
    """
    A = 2
    S = max_score + 2
    T = np.zeros(shape=(A, S, S))
    p = 1/n_sides

    # All the "roll" action transitions
    # First, the transition from state s to any non terminal state s' has probability
    # 1/n_sides unless s' <= s or s' > s + n_sides
    for s in range(0, S-1):
        for sprime in range(s+1, S-1):
            if sprime <= s + n_sides:
                T[0,s,sprime] = p
    # The rows of T[0] must all sum to one, so all the remaining probability goes to
    # the terminal state
    for s in range(0, S-1):
        T[0,s,S-1] = 1 - T[0,s].sum()
    # It is impossible to transition out of the terminal state; it is "absorbing"
    T[0,S-1,S-1] = 1

    # All the "stay" action transitions
    # This one is simple - all "stay" transitions dump you in the terminal state,
    # regardless of starting state
    T[1,:,S-1] = 1
    
    T[T<0] = 0 # There may be some very small negative probabilities due to rounding
               # errors - this fixes errythang
    
    return T
# Take a peek at a smaller version 
T = make_transition_matrix(n_sides=4, max_score=5)
print("roll transitions:")
print(T[0])
print("\nstay transitions:")
print(T[1])<jupyter_output><empty_output><jupyter_text>Now let's build the reward matrix. This is going to be a tensor of shape $(|\mathcal{S}|,|\mathcal{A}|) = (23,2)$. This one is even simpler than the transition matrix because only "stay" actions generate nonzero rewards, which are equal to the index of the state itself.<jupyter_code>def make_reward_matrix(max_score=21):
    """Create the reward matrix for the MDP.
    
    Arguments:
        max_score: the maximum score of the game before going bust
    Returns:
        np.ndarray: array of shape (S,A), where A=2, and S=max_score+2
            representing the reward matrix for the MDP
    """
    A = 2
    S = max_score + 2
    R = np.zeros(shape=(S, A))
    
    # Only need to create rewards for the "stay" action
    # Rewards are equal to the state index, except for the terminal state, which
    # always returns zero
    for s in range(0, S-1):
        R[s,1] = s
    
    return R
# Take a peek at a smaller version
R = make_reward_matrix(max_score=5)
print("roll rewards:")
print(R[:,0])
print("\nstay rewards:")
print(R[:,1])<jupyter_output><empty_output><jupyter_text>The algorithm Alright, now that we have the transition and reward matrices, our MDP is completely defined, and we can use the `pymdptoolbox` to help us figure out the optimal policy/strategy.<jupyter_code>n_sides = 20
max_score = 21

T = make_transition_matrix(n_sides, max_score)
R = make_reward_matrix(max_score)

model = mdp.ValueIteration(
    transitions=T,
    reward=R,
    discount=1,
    epsilon=0.001,
    max_iter=1000,
)
model.setVerbose()
model.run()
print(f"Algorithm finished running in {model.time:.2e} seconds")<jupyter_output><empty_output><jupyter_text>That ran pretty fast, didn't it? Unfortunately most realistic MDP problems have millions or billions of possible states (or more!), so this doesn't really scale very well. But it works for our small problem very well.  The results Now let's analyze the results. The `ValueIteration` object gives us easy access to the optimal value function and policy.<jupyter_code>plt.plot(model.V, marker='o')
x = np.linspace(0, max_score, 10)
plt.plot(x, x, linestyle="--", color='black')
ticks = list(range(0, max_score+1, 5)) + [max_score+1]
labels = [str(x) for x in ticks[:-1]] + ["\u2205"]
plt.xticks(ticks, labels)
plt.xlim(-1, max_score+2)
plt.xlabel("State sum of rolls $s$")
plt.ylabel("State value $V$")
plt.title("MDP optimal value function $V^*(s)$")
plt.show()
plt.plot(model.policy, marker='o')
ticks = list(range(0, max_score+1, 5)) + [max_score+1]
labels = [str(x) for x in ticks[:-1]] + ["\u2205"]
plt.xticks(ticks, labels)
plt.xlim(-1, max_score+2)
ticks = [0, 1]
labels = ["roll", "stay"]
plt.yticks(ticks, labels)
plt.ylim(-0.25, 1.25)
plt.xlabel("State sum of rolls $s$")
plt.ylabel("Policy $\pi$")
plt.title("MDP optimal policy $\pi^*(s)$")
plt.show()<jupyter_output><empty_output><jupyter_text>Looks like the optimal policy is to keep rolling until the sum gets to 10. This is why $V(s) = s$ for $s>=10$ (black dashed line); because that's the score you end up with when following this policy. For $s<10$, it's actually a bit higher than $s$ because you get an opportunity to roll again to get a higher score, and the sum is low enough that your chances of busting are relatively low. We can see the slope is positive for $s \le 21 - 20 = 1$ because it's impossible to bust below that point, but the slope becomes negative between $1 \le s \le 10$ because you're more likely to bust the higher you get. We can also calculate the state distribution $\rho_\pi(s_0 \rightarrow s,t)$, which tells us the probability to be in any one of the states $s$ after a time $t$ when starting from state $s_0$:$$\rho_\pi(s_0 \rightarrow s,t) = \sum_{s^\prime} T(s;s^\prime,\pi(s^\prime)) \rho_\pi(s_0 \rightarrow s^\prime, t-1) \\\text{where }\rho_\pi(s_0 \rightarrow s, 0) = \delta_{s, s_0}$$<jupyter_code>def calculate_state_distribution(policy, T, t_max=10):
    S = len(policy)
    # Reduce transition matrix to T(s';s) since policy is fixed
    T_ = np.zeros(shape=(S, S))
    for s in range(S):
        for sprime in range(S):
            T_[s,sprime] = T[policy[s],s,sprime]
    T = T_
    
    # Initialize rho
    rho = np.zeros(shape=(S, S, t_max+1))
    for s in range(0, S):
        rho[s,s,0] = 1
        
    # Use the iterative update equation
    for t in range(1, t_max+1):
        rho[:,:,t] = np.einsum("ji,kj->ki", T, rho[:,:,t-1])
    
    return rho
rho = calculate_state_distribution(model.policy, T, 5)
with warnings.catch_warnings():
    warnings.simplefilter('ignore') # Ignore the divide by zero error from taking log(0)
    plt.imshow(np.log10(rho[0].T), cmap='viridis')
cbar = plt.colorbar(shrink=0.35, aspect=9)
cbar.ax.set_title(r"$\log_{10}(\rho)$")
ticks = list(range(0, max_score+1, 5)) + [max_score+1]
labels = [str(x) for x in ticks[:-1]] + ["\u2205"]
plt.xticks(ticks, labels)
plt.xlabel("State sum of rolls $s$")
plt.ylabel("Number of rolls/turns $t$")
plt.title(r"Optimal state distribution $\rho_{\pi^*}(s_0\rightarrow s;t)$")
plt.subplots_adjust(right=2, top=2)
plt.show()<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12094748175619094, 'fraction_numerical': 0.021174781672448857, 'mean_word_length': 3.8746355685131197, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 30, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '13930', 'license': 'MIT', 'n_tokens_mistral': 2918, 'n_tokens_neox': 2756, 'n_words': 1029, 'path': 'mdp_demo.ipynb', 'repo_name': 'pydatacharlotte/04-Markov-decision-processes', 'size': '8359'}