MDPs, Bellman Equations & Value Iteration

Reinforcement Learning (RL) formalizes how an autonomous agent learns to make optimal decisions through trial-and-error interactions with an environment. Unlike supervised learning where target outputs $y$ are provided, the agent receives only scalar reward signals feedback.

The mathematical foundation of all reinforcement learning is the Markov Decision Process (MDP) and the recursive decomposition formalized by Richard Bellman in the Bellman Equations.

1. The Markov Decision Process (MDP) Framework#

An MDP is formally defined by the 5-tuple $(\mathcal{S}, \mathcal{A}, \mathcal{P}, \mathcal{R}, \gamma)$:

• $\mathcal{S}$: Set of all valid environment states.
• $\mathcal{A}$: Set of all valid agent actions.
• $\mathcal{P}$: State transition probability function $P(s' | s, a) = P(S_{t+1} = s' | S_t = s, A_t = a)$.
• $\mathcal{R}$: Expected reward function $R(s, a) = \mathbb{E}[R_{t+1} | S_t = s, A_t = a]$.
• $\gamma \in [0, 1)$: Discount factor balancing immediate versus delayed future rewards.

The Markov Property#

A state transition process is Markovian if the conditional probability distribution of future states depends solely upon the present state and action, completely independent of past historical trajectories:

The present state $S_t$ constitutes a sufficient statistic of the complete history.

The Cumulative Discounted Return#

The goal of the agent is to maximize the expected cumulative discounted return $G_t$:

For bounded rewards $|R_t| \le R_{\max}$, the discount factor $\gamma < 1$ guarantees that the infinite geometric series converges to a finite maximum bound:

2. Derivation of the Bellman Expectation Equations#

Let $\pi(a | s) = P(A_t = a | S_t = s)$ define a stochastic policy.

We define two fundamental value functions:

1. State-Value Function $V^\pi(s)$: Expected return starting from state $s$ under policy $\pi$:

1. Action-Value Function $Q^\pi(s, a)$: Expected return starting from state $s$, taking action $a$, and thereafter following policy $\pi$:

3. Bellman Optimality Equations & Contraction Mapping#

The optimal state-value function $V^*(s) = \max_\pi V^\pi(s)$ and optimal action-value function $Q^*(s, a) = \max_\pi Q^\pi(s, a)$ represent the maximum return achievable under any policy.

Under an optimal policy, the agent greedily selects the action that maximizes expected future return, replacing the expectation $\sum_a \pi(a|s)$ with the $\max_a$ operator:

Mathematical Proof: The Contraction Mapping Theorem#

Let $T^*$ denote the Bellman optimality operator mapping value functions $V \to V$:

We prove that $T^*$ is a contraction mapping in the infinity norm $\|u - v\|_\infty = \max_s |u(s) - v(s)|$.

For any two value functions $u$ and $v$:

Using the algebraic inequality $|\max_a f(a) - \max_a g(a)| \le \max_a |f(a) - g(a)|$:

Since $\sum_{s'} P(s' \mid s, a) = 1$:

4. From-Scratch Implementation: Zero Dependencies#

The following implementations construct an exact Value Iteration Algorithm for a deterministic GridWorld environment from scratch in pure Python and pure Rust.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
from typing import List, Dict, Tuple

class GridWorldFromScratch:
    def __init__(self, size: int = 4, gamma: float = 0.9, tol: float = 1e-4):
        self.size = size
        self.gamma = gamma
        self.tol = tol
        # Actions: Up, Down, Left, Right
        self.actions = [(-1, 0), (1, 0), (0, -1), (0, 1)]
        self.terminal_states = {(0, 0), (size - 1, size - 1)}

    def is_terminal(self, r: int, c: int) -> bool:
        return (r, c) in self.terminal_states

    def step(self, r: int, c: int, action: Tuple[int, int]) -> Tuple[int, int, float]:
        if self.is_terminal(r, c):
            return r, c, 0.0

        dr, dc = action
        next_r = max(0, min(self.size - 1, r + dr))
        next_c = max(0, min(self.size - 1, c + dc))
        reward = -1.0  # Step cost
        return next_r, next_c, reward

    def value_iteration(self) -> List[List[float]]:
        # Initialize value function table with zeros
        V = [[0.0] * self.size for _ in range(self.size)]

        while True:
            delta = 0.0
            new_V = [[0.0] * self.size for _ in range(self.size)]

            for r in range(self.size):
                for c in range(self.size):
                    if self.is_terminal(r, c):
                        continue

                    # Bellman Optimality: max over all actions [R + gamma * V(s')]
                    action_values = []
                    for act in self.actions:
                        next_r, next_c, reward = self.step(r, c, act)
                        val = reward + self.gamma * V[next_r][next_c]
                        action_values.append(val)

                    best_val = max(action_values)
                    new_V[r][c] = best_val
                    delta = max(delta, abs(V[r][c] - best_val))

            V = new_V
            if delta < self.tol:
                break

        return V

if __name__ == "__main__":
    grid = GridWorldFromScratch(size=4, gamma=0.9)
    optimal_V = grid.value_iteration()

    print("Optimal State-Value Function V*(s) for 4x4 GridWorld:")
    for row in optimal_V:
        print(" ", [round(v, 2) for v in row])

5. Active Knowledge Checks#