Practice Formulating a Small Markov Decision Process in Python

You understand what an MDP is. Then someone hands you a task description and a blank file, and the understanding evaporates.

That freeze is not a math problem. It is a procedure problem. Definitions tell you what an MDP contains; they never tell you what order to write things down. So here is the order, and here is a small task you can finish in one sitting.

Why Formulation Feels Harder Than It Is

The weak mental model sounds like this: I need to get the tuple right. That model turns formulation into definition recitation, and recitation has no feedback loop. You cannot tell whether you are correct until someone grades you.

The stronger model: I need a specification I can run and falsify. A spec is a contract. If it is wrong, running it produces a wrong state sequence, and the wrongness is visible.

Formulation is translation, not invention. The task description already contains every element. Your job is to extract, name, and check.

Tip: The check that matters is replay. Can you walk the task description against your spec and get the same outcomes? If not, the spec is wrong — not the task.

Pick One Small Task and Write the Rules in Plain English

We need a task small enough to enumerate by hand. A 3x3 grid with a goal, a pit, and a wall exercises every MDP element without drowning you in bookkeeping.

Here are the rules, written as sentences before any code exists:

• The agent starts at the top-left cell, (0, 0).
• The agent can move up, down, left, or right.
• Moving into a wall leaves the agent where it is. The move still costs a step.
• Reaching the goal ends the episode with a reward of +1.
• Falling into the pit ends the episode with a reward of -1.
• Every other step costs -0.04.

That last number is a living cost. Without it, an agent has no reason to hurry, and a spec that "runs fine" teaches the wrong behavior.

Now decide determinism. A deterministic grid is easier, but it makes transition probabilities decorative — every row is a single 1.0. The smallest way to make probabilities real is a slip: the agent moves in the intended direction with probability 0.8, and slips into one of the two perpendicular directions with probability 0.1 each. If the slip would push the agent into a wall, it stays put.

Write this down in English first. Ambiguous rewards and unstated slip rules are the two most common sources of a spec that runs but models the wrong problem.

Extract States, Actions, and Terminal Conditions

States are positions. But there is a subtlety worth pausing on: the goal and pit coordinates are still ordinary states. What makes them terminal is a rule attached to them, not a separate kind of object.

Why does that distinction matter? Because a terminal state has no outgoing actions. If you leave the goal as a coordinate with normal transitions, your table will happily let the agent step off the goal and keep collecting rewards. The episode never ends, and the return keeps accumulating. Every downstream calculation breaks.

So keep (0, 2) and (2, 2) in your state set, and mark them terminal with a predicate — or, equivalently, give them empty transition rows. The state identity and the terminal status are two separate facts.

Actions are the fixed move set: up, down, left, right. The same four actions apply everywhere except terminal states.

Terminal conditions are what make the return finite. They are not an afterthought bolted onto the end of the spec — they are the reason the transition table stops looping.

If you have read the formal MDP map, you already know the tuple. Here you are filling in the first entries with values you can print.

Build the Transition Table Before the Code

Write transitions as (state, action) -> list of (next_state, probability) pairs. Every row must sum to 1.0. That is your first correctness check, and it catches more bugs than any other single test.

Two states by hand, before generating the rest. The wall sits at (0, 1), so a right move from (0, 0) is blocked and the agent stays put:

Read the first row carefully. The intended move hits the wall, so it self-loops at 0.8. One slip goes down to (1,0) at 0.1. The other slip goes left, which is out of bounds, so it also self-loops at 0.1. Two outcomes land on (0,0), and they aggregate to 0.9.

The second row is the same story with no escape route: up is blocked, and both perpendicular slips (left and right) are also blocked or out of bounds. Every path returns to (0,0), so the row is a single 1.0 self-loop.

Common mistake: Leaving out the self-loop on a blocked move. The row sum silently drops below 1.0, or the agent teleports somewhere it should never reach.

Terminal states get an empty transition list, not a missing key. A missing key hides bugs behind a KeyError at the worst moment. An empty list documents intent: this state has no outgoing edges, on purpose.

Encode the Specification in Python

No RL library. No framework. The point is to see the structure, not to hide it behind an API.

from collections import defaultdict

WALLS = {(0, 1), (1, 1)}
GOAL = (0, 2)
PIT = (2, 2)
TERMINALS = {GOAL, PIT}

MOVES = {
    "up": (-1, 0),
    "down": (1, 0),
    "left": (0, -1),
    "right": (0, 1),
}
PERPENDICULAR = {
    "up": ["left", "right"],
    "down": ["left", "right"],
    "left": ["up", "down"],
    "right": ["up", "down"],
}

def in_bounds(cell):
    r, c = cell
    return 0 <= r < 3 and 0 <= c < 3

def step(cell, direction):
    dr, dc = MOVES[direction]
    nxt = (cell[0] + dr, cell[1] + dc)
    if not in_bounds(nxt) or nxt in WALLS:
        return cell
    return nxt

class GridMDP:
    def __init__(self):
        self.states = [
            (r, c) for r in range(3) for c in range(3)
            if (r, c) not in WALLS
        ]
        self.actions = list(MOVES.keys())
        self.discount = 0.95
        self.start = (0, 0)

    def is_terminal(self, state):
        return state in TERMINALS

    def transitions(self, state, action):
        if self.is_terminal(state):
            return []
        outcomes = defaultdict(float)
        outcomes[step(state, action)] += 0.8
        for slip in PERPENDICULAR[action]:
            outcomes[step(state, slip)] += 0.1
        return sorted(outcomes.items())

    def reward(self, state, action, next_state):
        if next_state == GOAL:
            return 1.0
        if next_state == PIT:
            return -1.0
        return -0.04

Two design choices deserve a note.

First, the reward is keyed on (state, action, next_state), not on the action alone. This matters the moment the same action can lead to different outcomes. If you attach reward to the action, a slip into the pit pays the same as a deliberate step, and your spec quietly lies.

Second, transitions returns an empty list for terminal states. The interface stays narrow — five accessors, no more — so it is easy to test.

Verify the Spec Against the Task Description

Three checks, run in order.

Check 1: every row sums to 1.0.

mdp = GridMDP()
for s in mdp.states:
    for a in mdp.actions:
        row = mdp.transitions(s, a)
        if row:
            total = sum(p for _, p in row)
            assert abs(total - 1.0) < 1e-9, (s, a, total)

Check 2: terminal states have no outgoing transitions.

for s in mdp.states:
    if mdp.is_terminal(s):
        assert all(mdp.transitions(s, a) == [] for a in mdp.actions)

Check 3: inspect the outcome list for one state-action pair. This is not a rollout — it is a direct look at the row you just built, so you can confirm the probabilities match your hand table:

for next_state, prob in mdp.transitions((0, 0), "right"):
    print(next_state, prob)

You should see (0, 0) 0.9 and (1, 0) 0.1. If your hand table said something else, the divergence points to a specific wrong entry. That is the whole value of the exercise: a wrong spec produces a wrong row, and a wrong row is debuggable.

If you want an actual episode trace, you have to sample according to the probabilities rather than pick the first sorted outcome. A short weighted sampler makes that explicit:

import random

def sample_next(mdp, state, action):
    row = mdp.transitions(state, action)
    if not row:
        return state
    states, probs = zip(*row)
    return random.choices(states, weights=probs, k=1)[0]

state = mdp.start
for action in ["right", "right"]:
    state = sample_next(mdp, state, action)
    print(state)

Run it a few times. The path varies because the environment is stochastic — that variability is the point, not a bug.

Warning: Check the sign convention on the step cost. A step cost that is accidentally positive makes the agent love wandering, and the spec still "runs fine." Sign errors do not crash. They corrupt.

This is where most beginners stop too early. A spec that runs is not a spec that is correct.

Common Formulation Mistakes and How to Catch Them

The first two rows are the same failure in opposite directions. Too much state breaks the Markov property by making it useless; too little state breaks it by making the future depend on hidden history. Both are caught by the same question: does the current state alone predict the next one?

One Modification to Prove You Own the Spec

Change one thing. Predict the consequence before you edit the code.

Try making the move fully stochastic: 0.7 intended, 0.15 per perpendicular slip. Before you run anything, answer: which rows change? By how much? Do the row sums still hold?

Then re-run all three checks. If a row sum breaks, you learned where your abstraction leaks. If the sampled episode changes, you learned that your earlier trace was only valid under the old probabilities.

That drill is the real skill. The same procedure applies to any small sequential task you meet later: write the rules, extract the sets, build the table, encode it, verify it.

Note: This hand-built approach works for small finite MDPs. It does not scale to large or continuous state spaces — and that limitation is exactly why later methods exist.

Where This Spec Goes Next

Here is the decision rule I would hand you: if you cannot write the transition table by hand for a small task, you are not ready to hand it to an algorithm. The table is the contract. Algorithms are just ways of reading it.

Once your spec passes all three checks, you have something more valuable than a finished exercise. You have an input. A Bellman backup needs exactly this: states, actions, transition probabilities, rewards, and a discount factor. Feed your verified grid into a value iteration pass and watch the numbers settle. The formulation work you just did is what makes that next step checkable instead of mysterious.