The transition matrix is a 3-D array

An MDP's dynamics live in one object: T with shape (S, A, S′), where T[s, a, s′] = P(s′ | s, a). It's not really a matrix — it's a brick of probabilities, and picking a state and an action pulls out one rod of possibilities T[s, a, :] that sums to 1. Once you can see the rod, every RL algorithm is a few lines of numpy. World here: a 5-square hallway, actions Left / Stay / Right; you do what you chose with p = 0.6, each other move with p = 0.2, and pushing off an end collapses into staying. All indices are 0-based, exactly as in numpy.

T.shape = (5, 3, 5) · T[s, a, s′] = P(s′ | s, a) · Σ_s′ T[s, a, s′] = 1 for every (s, a)

The world — hallway squares 0…4

The agent (●) sits on the selected s; the arrow is the selected action; the numbers are where you actually end up — the rod T[s, a, :], drawn on the world itself.

entry T[s,a,s′] (cube volume ∝ probability) selected rod T[s,a,:] exact zero

Inspector — pick (s, a), get the rod

state s (which square you're on)

action a

Access patterns — what each slice is

T[s, a, s2]  # scalar — one probability
T[s, a]      # shape (5,)  — the rod: where can I land?
T[s]         # shape (3,5) — this square's options
T[:, a]      # shape (5,5) — one action's Markov matrix
T.sum(axis=2) # all ones — every rod is a distribution

Building T — the walls are one clip

T = np.zeros((5, 3, 5))
MOVES = [-1, 0, +1]          # what L, S, R try to do
for s in range(5):
    for a in range(3):
        for m, move in enumerate(MOVES):
            p = 0.6 if m == a else 0.2
            sp = np.clip(s + move, 0, 4)
            T[s, a, sp] += p # += : off-grid mass
                             #      collapses into "stay"

The += is the entire wall rule. At s = 0, "left" clips back to 0, so its 0.6 lands on the same square that "stay" feeds — that's where the 0.8 comes from. No special cases, and rows sum to 1 by construction.

From T to V* and Q* — the whole algorithm

This is why the tensor layout matters: value iteration is three lines, and each line is something you can point at in the picture above.

R = np.array([0., 0., 0., 0., 1.])   # paid for acting from a square
V = np.zeros(5)
while True:
    Q = R[:, None] + gamma * (T @ V) # (5,3,5)@(5,) → (5,3):  every
                                     # Q[s,a] = R[s] + γ·(rod · V)
    V_new = Q.max(axis=1)            # best action per state → V
    if np.abs(V_new - V).max() < 1e-6: break
    V = V_new
policy = Q.argmax(axis=1)            # π: which action won each row

T @ V contracts the s′ axis. numpy matmul dots the last axis of T against V — it multiplies every rod by the value of where it lands and sums. One @ does all 15 expectations at once. The max and argmax then run along the action axis (axis=1) — that's the Bellman update and the greedy policy, verbatim.

Run it — live, on this T

discount γ 0.9

s	V[s]	Q[s,0] L	Q[s,1] S	Q[s,2] R	π[s]

The outlined Q cell is the selected (s, a) from the inspector. Its dot product uses exactly the highlighted rod — same object in the lattice, the numpy repr, and here.