Module 01: Math Foundations for RL

Before we start training agents, we need the mathematical language that makes RL precise. This module covers MDPs, value functions, and the Bellman equations - the backbone of everything that follows.

Learning Objectives

1. Understand Markov Decision Processes (MDPs) formally.
2. Learn the Bellman Expectation & Optimality equations.
3. Implement Value Iteration and Policy Iteration.
4. Solve GridWorld and FrozenLake with dynamic programming.

---

1. What is an MDP?

An MDP (Markov Decision Process) is defined as a 5-tuple:

$$\mathcal{M} = \langle S, A, P, R, \gamma \rangle$$

Symbol	Meaning
$S$	Set of states
$A$	Set of actions
$P(s' \mid s, a)$	Transition probability (next state distribution)
$R(s, a, s')$	Reward function
$\gamma$	Discount factor ($0 \leq \gamma < 1$)

Goal of RL: Find a policy $\pi(a \mid s)$ that maximizes expected discounted return:

$$G_t = \sum_{k=0}^\infty \gamma^k R_{t+k+1}$$

Real-world analogy

Think of a self-driving car. The state is what the sensors see (road, cars, pedestrians). The action is steering/braking/accelerating. The transition is physics + other drivers' behavior. The reward is: +1 for safe driving, -1000 for hitting something, +10 for reaching the destination efficiently. The discount $\gamma$ means the car cares more about not hitting the pedestrian right now than saving 30 seconds at a traffic light in 5 minutes.

---

2. Value Functions & Bellman Equations

State-value function

Under policy $\pi$:

$$V^\pi(s) = \mathbb{E}_\pi [ G_t \mid S_t = s ]$$

"How good is it to be in state $s$ if I follow policy $\pi$?"

Action-value function

$$Q^\pi(s,a) = \mathbb{E}_\pi [ G_t \mid S_t = s, A_t = a ]$$

"How good is it to take action $a$ in state $s$ and then follow $\pi$?"

Bellman Expectation Equation

$$V^\pi(s) = \sum_a \pi(a \mid s) \sum_{s'} P(s' \mid s,a) \big[ R(s,a,s') + \gamma V^\pi(s') \big]$$

This is recursive - the value of a state depends on the values of its successor states. It's turtles all the way down, except the turtles converge.

Bellman Optimality Equation

$$V^*(s) = \max_a \sum_{s'} P(s' \mid s,a) \big[ R(s,a,s') + \gamma V^*(s') \big]$$

The optimal value function picks the best action at every state, not just whatever $\pi$ says.

---

3. Value Iteration

Pseudocode:

Initialize V(s) arbitrarily
Repeat until convergence:
    For each state s:
        V(s) <- max_a Sum_{s'} P(s'|s,a) [ R(s,a,s') + gamma * V(s') ]

Then derive the optimal policy:

$$\pi^*(s) = \arg\max_a \sum_{s'} P(s' \mid s,a) [ R(s,a,s') + \gamma V(s') ]$$

---

4. Example: NumPy GridWorld

A small 4x4 grid (like Sutton & Barto Chapter 4). The agent starts anywhere, tries to reach a terminal state. Reward = -1 per step, 0 at terminal.

import numpy as np


class GridWorld:
    def __init__(
        self,
        size: int = 4,
        terminal_states: list[int] | None = None,
        gamma: float = 1.0,
    ) -> None:
        self.size = size
        self.n_states = size * size
        self.n_actions = 4  # up, right, down, left
        self.terminal_states = terminal_states or [0, 15]
        self.gamma = gamma

    def step(self, state: int, action: int) -> tuple[int, float]:
        if state in self.terminal_states:
            return state, 0.0

        row, col = divmod(state, self.size)
        if action == 0:    # up
            row = max(row - 1, 0)
        elif action == 1:  # right
            col = min(col + 1, self.size - 1)
        elif action == 2:  # down
            row = min(row + 1, self.size - 1)
        elif action == 3:  # left
            col = max(col - 1, 0)

        next_state = row * self.size + col
        reward = -1.0
        return next_state, reward


def value_iteration(env: GridWorld, theta: float = 1e-4) -> tuple[np.ndarray, np.ndarray]:
    V = np.zeros(env.n_states)

    while True:
        delta = 0.0
        for s in range(env.n_states):
            if s in env.terminal_states:
                continue
            v = V[s]
            q_values = []
            for a in range(env.n_actions):
                s_next, r = env.step(s, a)
                q_values.append(r + env.gamma * V[s_next])
            V[s] = max(q_values)
            delta = max(delta, abs(v - V[s]))
        if delta < theta:
            break

    # Extract policy
    policy = np.zeros([env.n_states, env.n_actions])
    for s in range(env.n_states):
        if s in env.terminal_states:
            continue
        q_values = []
        for a in range(env.n_actions):
            s_next, r = env.step(s, a)
            q_values.append(r + env.gamma * V[s_next])
        best_a = np.argmax(q_values)
        policy[s, best_a] = 1.0

    return V.reshape(env.size, env.size), policy

Run it:

env = GridWorld()
V, policy = value_iteration(env)
print("Optimal Value Function:")
print(V)

arrows = {0: "^", 1: ">", 2: "v", 3: "<"}
policy_arrows = []
for s in range(env.n_states):
    if s in env.terminal_states:
        policy_arrows.append("T")
    else:
        policy_arrows.append(arrows[np.argmax(policy[s])])
print("\nOptimal Policy:")
print(np.array(policy_arrows).reshape(env.size, env.size))

---

5. Example: Gymnasium FrozenLake

Value iteration on FrozenLake-v1 - a discrete MDP with slippery ice (stochastic transitions):

import gymnasium as gym
import numpy as np

env = gym.make("FrozenLake-v1", is_slippery=True)
n_states = env.observation_space.n
n_actions = env.action_space.n


def value_iteration_frozenlake(
    env: gym.Env,
    gamma: float = 0.99,
    theta: float = 1e-8,
) -> tuple[np.ndarray, np.ndarray]:
    V = np.zeros(n_states)

    while True:
        delta = 0.0
        for s in range(n_states):
            v = V[s]
            q_values = []
            for a in range(n_actions):
                q = 0.0
                for prob, s_next, reward, terminated in env.P[s][a]:
                    q += prob * (reward + gamma * V[s_next])
                q_values.append(q)
            V[s] = max(q_values)
            delta = max(delta, abs(v - V[s]))
        if delta < theta:
            break

    # Derive policy
    policy = np.zeros([n_states, n_actions])
    for s in range(n_states):
        q_values = []
        for a in range(n_actions):
            q = 0.0
            for prob, s_next, reward, terminated in env.P[s][a]:
                q += prob * (reward + gamma * V[s_next])
            q_values.append(q)
        best_a = np.argmax(q_values)
        policy[s, best_a] = 1.0

    return V, policy


V, policy = value_iteration_frozenlake(env)
print("Optimal State Values:", V.reshape(4, 4))

note

FrozenLake with is_slippery=True is stochastic - the agent only goes where it intended 1/3 of the time. This is why the optimal policy looks weird. It's not confused; it's accounting for the ice.

---

6. Policy Iteration

An alternative to value iteration that alternates between two steps:

1. Policy Evaluation: Given a fixed policy $\pi$, compute $V^\pi$ by iterating:

$$V^\pi(s) = \sum_a \pi(a \mid s)\sum_{s'} P(s' \mid s,a)[R(s,a,s') + \gamma V^\pi(s')]$$

2. Policy Improvement: For each state, pick the greedy action:

$$\pi'(s) = \arg\max_a \sum_{s'} P(s' \mid s,a) [R(s,a,s') + \gamma V^\pi(s')]$$

3. If the policy didn't change, stop. You've found $\pi^*$.

def policy_evaluation(
    policy: np.ndarray,
    env: GridWorld,
    V: np.ndarray,
    theta: float = 1e-4,
) -> np.ndarray:
    while True:
        delta = 0.0
        for s in range(env.n_states):
            if s in env.terminal_states:
                continue
            v = V[s]
            new_v = 0.0
            for a, action_prob in enumerate(policy[s]):
                s_next, r = env.step(s, a)
                new_v += action_prob * (r + env.gamma * V[s_next])
            V[s] = new_v
            delta = max(delta, abs(v - new_v))
        if delta < theta:
            break
    return V


def policy_improvement(V: np.ndarray, env: GridWorld) -> np.ndarray:
    policy = np.zeros([env.n_states, env.n_actions])
    for s in range(env.n_states):
        if s in env.terminal_states:
            continue
        q_values = []
        for a in range(env.n_actions):
            s_next, r = env.step(s, a)
            q_values.append(r + env.gamma * V[s_next])
        best_a = np.argmax(q_values)
        policy[s, best_a] = 1.0
    return policy


def policy_iteration(env: GridWorld) -> tuple[np.ndarray, np.ndarray]:
    # Start with uniform random policy
    policy = np.ones([env.n_states, env.n_actions]) / env.n_actions
    V = np.zeros(env.n_states)

    while True:
        V = policy_evaluation(policy, env, V)
        new_policy = policy_improvement(V, env)
        if np.array_equal(new_policy, policy):
            break
        policy = new_policy

    return V.reshape(env.size, env.size), policy

---

7. Key Takeaways

Algorithm	How it works	When to use
Value Iteration	Directly applies Bellman optimality updates	Faster in practice for small MDPs
Policy Iteration	Alternates evaluation + improvement	Gives insight into policy/value interaction

Both converge to the optimal policy $\pi^*$ in finite MDPs. They require knowing the transition model $P$ - which we usually don't have. That's why the next modules teach model-free methods.

---

Milestone Checklist

• Can define an MDP formally (all 5 components)
• Can explain and derive the Bellman expectation & optimality equations
• Implemented Value Iteration from scratch
• Implemented Policy Iteration from scratch
• Solved both GridWorld and FrozenLake

---

Exercises

1. Modify GridWorld: Add a "lava" state at position 5 that gives -10 reward and terminates the episode. How does the optimal policy change?
2. Stochastic GridWorld: Make actions succeed only 80% of the time (20% random). Re-run value iteration. What happens to the value function?
3. Discount factor sweep: Run value iteration on FrozenLake with $\gamma \in \{0.5, 0.9, 0.99, 0.999\}$. Plot the value functions. Why does $\gamma$ close to 1 give higher values?
4. Prove it: Show on paper that Value Iteration is equivalent to Policy Iteration where policy evaluation uses only a single sweep. (Hint: think about the number of evaluation sweeps.)