MDPs, states, actions, and rewards

1. Intuition

Picture a robot on a 4×4 floor. It starts in the top-left corner and wants to reach the charger in the bottom-right corner. Each move drains a little battery, so every step costs $-1$. Bumping into a wall wastes the move. The episode ends at the charger.

That is a complete RL problem. Five things define it:

The agent's behavior is a policy $\pi(a \mid s)$: for each state, a probability for each action. The goal is a policy that collects as much total reward as possible, not just the best next reward. That difference is the heart of RL: a move can look bad now (a detour) and be right overall.

The Markov property. The next state depends only on the current state and action, not on how you got there. In the grid that holds: the cell is all you need. When it fails, the fix is in the state, not the algorithm. A single frame of a video game hides which way the ball is moving, so DQN stacked the last four frames into one state. A chat model's state is the full context window for the same reason.

Discounting. A reward $k$ steps in the future is multiplied by $\gamma^k$. With $\gamma = 0.9$ a reward 10 steps away keeps about a third of its weight; 50 steps away it is nearly gone. A rough rule: $\gamma$ buys a horizon of about $1/(1-\gamma)$ steps.

Key distinction

Reward vs return vs value. A reward is the single number from one step. The return is the discounted sum of rewards from now to the end of one episode. A value is the expected return, averaged over everything the policy and world might do. Algorithms learn values; you observe rewards.

2. Visualization

Synthetic 4×4 grid, every move costs −1. Left: how fast each discount forgets the future. Middle: a random walker expects to pay a lot before stumbling onto the goal. Right: the best policy pays exactly the number of moves it needs.

Try it (predict first, then run code/26-reinforcement-learning/mdp.py):

1. In the middle panel, why are the two cells next to the goal equal (−5.8), and why is the start corner worst? Predict what happens to the start corner if you change $\gamma$ to 0.99.
2. The random policy's mean undiscounted return over 5,000 rollouts is about −60 even though the optimum is −6. Predict whether the standard deviation is small or large, then check the script output.
3. In the right panel, values go up by exactly 1 per step toward the goal. What would change if one interior cell were a wall?
4. Change the step reward to −2. Predict the new $V^*$ at the start, and whether the best policy changes.

3. The math

Symbols

Return and values

If every reward is bounded by $R_{\max}$ and $\gamma < 1$, then $\lvert G_t\rvert \le R_{\max}/(1-\gamma)$, which is where the "horizon" $1/(1-\gamma)$ comes from.

The Bellman equations

Split the return into "this step" plus "everything after", then take expectations:

The first (expectation) equation averages over the policy's actions; it defines the value of a fixed policy. The second (optimality) equation takes the best action; its solution is the value of the best policy. Policy evaluation repeats the first as an update until nothing changes. Value iteration repeats the second. Both converge for $\gamma < 1$ (and here for $\gamma = 1$ too, because every policy that matters reaches the terminal state).

Worked example 1: a discounted return

Six moves to the goal, each reward $-1$, $\gamma = 0.9$:

1. $G = -1 - 0.9 - 0.9^2 - 0.9^3 - 0.9^4 - 0.9^5$
2. Powers: $0.9^2 = 0.81$, $0.9^3 = 0.729$, $0.9^4 = 0.6561$, $0.9^5 = 0.59049$.
3. Sum: $1 + 0.9 + 0.81 + 0.729 + 0.6561 + 0.59049 = 4.68559$.
4. $G = -4.6856$. Check with the geometric-series formula: $-\frac{1 - 0.9^6}{1 - 0.9} = -\frac{1 - 0.531441}{0.1} = -4.68559$. ✓

With $\gamma = 1$ the same path returns $-6$; with $\gamma = 0.5$ it returns $-1.9688$, so a short horizon barely tells a 6-step path from a 60-step one.

Worked example 2: checking a value with the Bellman equation

Uniform random policy ($\pi = 0.25$ for each action), $\gamma = 0.9$. Take the cell $(3, 2)$, just left of the goal. The script's converged values for the cells each action reaches:

1. Sum of the last column: $-7.674 - 6.181 - 8.173 - 1.000 = -23.028$.
2. Multiply by $\pi = 0.25$: $-23.028 \times 0.25 = -5.757$.
3. That equals $V^\pi(3, 2) = -5.757$ from the table. ✓ The value is a fixed point: plug it in on the right and it comes back out on the left.

For the optimal values with $\gamma = 1$: $V^*(3,2) = \max(-1 + V^*(2,2), \dots, -1 + V^*(\text{goal})) = \max(-1 - 2, \dots, -1 + 0) = -1$, and by the same rule each step away from the goal subtracts 1, giving $V^*(0,0) = -6$.

4. Implementation

The full script is code/26-reinforcement-learning/mdp.py (NumPy only). It builds the grid, rolls out a random policy, and solves for values two ways.

Build it from scratch

Step 1: the world. The environment is one function: given a state and action, return the next state, the reward, and whether the episode ended. Writing transitions and rewards in code, not a table, is how most simulators work. Clamping the coordinates is the "bump into the wall" rule.

N = 4
STATES = [(r, c) for r in range(N) for c in range(N)]
ACTIONS = {"up": (-1, 0), "down": (1, 0), "left": (0, -1), "right": (0, 1)}
START, GOAL = (0, 0), (3, 3)

def step(s, a):
    if s == GOAL:
        return s, 0.0, True                     # absorbing terminal state
    dr, dc = ACTIONS[a]
    s2 = (min(max(s[0] + dr, 0), N - 1), min(max(s[1] + dc, 0), N - 1))
    return s2, -1.0, s2 == GOAL

Predict: what does step((0, 0), "up") return?

Step 2: the return. Working backward turns the infinite sum into one multiply-add per step, using $G_t = r_t + \gamma G_{t+1}$. Every RL codebase has this loop somewhere (it reappears as "returns-to-go" in policy gradients).

def discounted_return(rewards, gamma):
    g = 0.0
    for r in reversed(rewards):   # G_t = r_t + gamma * G_{t+1}
        g = r + gamma * g
    return g

Predict: discounted_return([0, 0, 0, 1], 0.5)?

Step 3: a policy and a rollout. A policy maps a state to a distribution over actions. A rollout samples from it until the episode ends. Averaging many rollouts estimates the value by brute force (this is the Monte Carlo method).

def uniform(_s):
    return {a: 0.25 for a in ACTIONS}

def rollout(policy, rng, max_steps=500):
    s, rewards = START, []
    for _ in range(max_steps):
        dist = policy(s)
        a = rng.choices(list(dist), weights=list(dist.values()))[0]
        s, r, done = step(s, a)
        rewards.append(r)
        if done:
            break
    return rewards

Predict: why does rollout need max_steps even though this grid always reaches the goal eventually?

Step 4: policy evaluation. Apply the Bellman expectation equation to every state, over and over, until the largest change is tiny. Updating V in place (using new values as soon as they exist) converges a bit faster than keeping two copies.

def policy_evaluation(policy, gamma, tol=1e-10):
    V = {s: 0.0 for s in STATES}
    while True:
        delta = 0.0
        for s in STATES:
            if s == GOAL:
                continue
            v = sum(p * (r + gamma * V[s2]) for a, p in policy(s).items() for s2, r, _ in [step(s, a)])
            delta, V[s] = max(delta, abs(v - V[s])), v
        if delta < tol:
            return V

Predict: with gamma=1.0 and the uniform policy, this still converges. Would it converge on a grid with no reachable goal?

Step 5: value iteration. Replace the average over the policy with a max over actions. The result is $V^*$, and acting greedily on it is optimal.

def value_iteration(gamma, tol=1e-10):
    V = {s: 0.0 for s in STATES}
    while True:
        delta = 0.0
        for s in STATES:
            if s == GOAL:
                continue
            v = max(r + gamma * V[s2] for a in ACTIONS for s2, r, _ in [step(s, a)])
            delta, V[s] = max(delta, abs(v - V[s])), v
        if delta < tol:
            return V

Predict: how many sweeps until $V^*(0,0)$ reaches −6, if each sweep goes in row-major order from (0,0)?

Use the library

Gymnasium is the standard environment interface. Its FrozenLake-v1 is a 4×4 grid MDP like ours, but slippery by default (stochastic transitions), and it exposes the transition table as env.unwrapped.P[s][a], a list of (prob, next_state, reward, done) tuples.

import gymnasium as gym

env = gym.make("FrozenLake-v1", is_slippery=True)
obs, info = env.reset(seed=0)
obs, reward, terminated, truncated, info = env.step(env.action_space.sample())
P = env.unwrapped.P          # P[s][a] -> [(prob, s2, r, done), ...]; plug into the Bellman update

The library adds what our script left implicit: separate terminated (the MDP ended) and truncated (a time limit cut it off) flags, seeded resets, and a standard step signature every RL library understands. The difference between the two flags matters: on truncation you should still bootstrap from the next state's value; on termination you should not.

Run it

uv run code/26-reinforcement-learning/mdp.py

gamma = 0.9: G = -4.6856
mean return = -59.87, std = 51.37  (optimal = -6)
Bellman check at (3, 2), gamma=0.9: ... 0.25 * sum(-1 + 0.9 V) = -5.757 = V(3, 2) = -5.757
=== V* by value iteration (gamma = 1): minus the number of moves to the goal ===
[[-6. -5. -4. -3.]
 ...
saved figures/mdps-states-actions-rewards.png

5. Engineering

Write the MDP down before writing code. For each problem, fill in five lines: state, action, transition (known or only sampled?), reward (when, and of what), and $\gamma$ or horizon.

Known model vs samples only. If you can write $P$ and $r$ as tables (small planning problems, games with known rules), dynamic programming like value iteration solves the MDP exactly. Most real problems only let you sample transitions, which is why the next lessons learn from experience.

Reward design is product design. The agent optimizes exactly the number you give it. A boat-racing agent from OpenAI (2016) learned to circle a lagoon collecting point pickups forever instead of finishing the race, because the score rewarded pickups. LLMs trained against a reward model learn to please the reward model (longer answers, confident tone) unless something checks them. Define reward from the outcome you want, and keep a held-out check the agent can't see.

Sparse rewards. If reward only comes at the very end (win or lose), a random policy may never see a nonzero signal. Options: shape the reward with intermediate signals (carefully, since shaping can be hacked), start from a supervised or imitation-learned policy (this is why LLM RL starts from an SFT model), or use a curriculum of easier tasks.

Watch out

Failure modes. A state that hides needed information (the agent can't tell two situations apart and dithers); a reward that measures a proxy; $\gamma = 1$ on a task that never ends (values become infinite); rewards on wildly different scales (+1000 for one event, ±1 for others), which make value and gradient estimates noisy.

Common mistakes

• Confusing reward with return or value, and comparing a single episode's return to a value estimate as if they should match.
• Using $\gamma = 1$ with no time limit or terminal state.
• Treating a time-limit cutoff as a real terminal state, which teaches the agent that the world ends at step 200.
• Leaving information out of the state (velocity, the last few observations, remaining budget) and blaming the algorithm.
• Changing the reward and comparing returns across runs as if the numbers meant the same thing.

6. Knowledge check

Practice exercise. An episode gives rewards $[0, 0, -1, 10]$ and $\gamma = 0.5$. Compute $G_0$, then $G_2$.

Implementation challenge. Make the grid slippery: an action moves in the intended direction with probability 0.8 and in each of the two perpendicular directions with probability 0.1. Change step to return a list of (prob, next_state, reward, done) and update value_iteration to sum over it. Does $V^*(0,0)$ get better or worse than −6, and does the best path change?

Summary

• An MDP is states, actions, transitions, rewards, and a discount. The Markov property says the state carries everything needed to predict what happens next.
• The return is the discounted sum of future rewards; $\gamma$ sets an effective horizon of about $1/(1-\gamma)$ steps.
• Values are expected returns: $V^\pi(s)$ for a state, $Q^\pi(s,a)$ for a state and first action.
• The Bellman equation splits a value into one step of reward plus the discounted value of where you land. Iterating it gives policy evaluation (average) or value iteration (max).
• Most RL failures start in the formulation: missing state, a hackable or sparse reward, or an unbounded return.

Next: Q-learning and SARSA

Related: Learning paradigms · Policy gradients and REINFORCE · Fine-tuning in practice

Go deeper, code-first (AI Engineering from Scratch): Dynamic Programming — Policy Iteration & Value Iteration

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): MDPs, States, Actions & Rewards.

Interview angle