Markov Decision Processes

The Setup

Imagine you are navigating a gridworld — a 4×4 grid of cells. Some cells have rewards (a treasure chest, a pit of fire). At each step you can move up, down, left, or right. The floor is slippery: when you try to move right, 80% of the time you succeed, but 10% of the time you slide up and 10% you slide down. At the end of some path through this grid, you reach a terminal state and collect your cumulative reward.

This is not a simple optimization — you cannot just pick the action with the highest immediate reward, because an action that looks bad now might lead to a great state later. And the environment is stochastic, so you can't plan a deterministic path; you need to decide what to do in each possible state you might land in.

The key insight: a good decision rule must be a function of the current state, not just a fixed sequence of actions. Such a function is called a policy.

Why Single-Step Isn't Enough

In a single-step ("bandit") problem, you pull an arm, get a reward, done. But in sequential problems, today's action affects tomorrow's state. This temporal dependency is the distinguishing feature. To handle it, we need a model that tracks:

• What states exist
• What actions are available
• How actions cause state transitions
• What rewards are received
• How to aggregate rewards over time

MDPs provide exactly this model.

# Gridworld as intuition pump
# States: (row, col) in a 4x4 grid
# Actions: UP, DOWN, LEFT, RIGHT
# Transition: intended direction with prob 0.8, perpendicular with 0.1 each
# Reward: +1 at goal (3,3), -1 at pit (3,1), 0 elsewhere
# Terminal: goal and pit are absorbing states

GRID_SIZE = 4
GOAL = (3, 3)
PIT  = (3, 1)
REWARDS = {GOAL: +1.0, PIT: -1.0}

def is_terminal(s):
    return s in (GOAL, PIT)

Five Components

A Markov Decision Process is a 5-tuple $(\mathcal{S}, \mathcal{A}, T, R, \gamma)$:

Some formulations write $R(s, a)$ or $R(s)$ — these are simplifications of the general $R(s, a, s')$.

The Markov Property

The Markov property is the key modeling assumption: the future depends on the past only through the current state. Formally: $$P(s_{t+1} \mid s_t, a_t, s_{t-1}, a_{t-1}, \ldots, s_0, a_0) = P(s_{t+1} \mid s_t, a_t) = T(s_t, a_t, s_{t+1})$$

This says: if you know where you are now, knowing your entire history gives you zero additional predictive power. The state must be a sufficient statistic for the future. This is why defining state correctly is the art of MDP modeling — a poorly defined state that omits relevant history violates the Markov property and will give wrong answers.

The Discount Factor

The discount factor $\gamma$ captures the time value of reward. A reward received now is worth more than the same reward received later (for $\gamma < 1$). Mathematically, it ensures the infinite sum of rewards converges: $$G_t = \sum_{k=0}^{\infty} \gamma^k R_{t+k+1}$$

If $\gamma = 0$: agent is myopic, only cares about immediate reward.
If $\gamma \to 1$: agent cares equally about all future rewards (requires finite episodes for convergence).
Typically $\gamma = 0.9$ or $0.99$.

Analogy: Think of $\gamma$ as an inflation rate. A dollar today is worth more than a dollar in a year. Setting $\gamma = 0.9$ means a reward 10 steps away is worth $0.9^{10} \approx 0.35$ of its face value today.

from dataclasses import dataclass, field
from typing import Dict, Tuple, List

State  = int   # or any hashable
Action = int

@dataclass
class MDP:
    states:   List[State]
    actions:  List[Action]
    # T[s][a] -> list of (prob, next_state, reward)
    T: Dict[State, Dict[Action, List[Tuple[float, State, float]]]]
    gamma: float = 0.9

    def transitions(self, s: State, a: Action):
        """Return list of (probability, next_state, reward)."""
        return self.T[s][a]

What Is a Policy?

A policy $\pi$ is a rule for selecting actions. The most general form is a stochastic policy: $\pi(a \mid s) = P(A_t = a \mid S_t = s)$. A deterministic policy is a special case: $\pi(s) = a$.

For fully observable MDPs (which is what we study here), it is sufficient to consider stationary deterministic policies: the action depends only on the current state, not on time or history. This is a deep theorem — the optimal policy for a discounted infinite-horizon MDP is always stationary and deterministic. You don't need to remember the past or act differently at different time steps.

The Return and the Objective

Given a policy $\pi$, starting from state $s_0$, the agent follows a trajectory $s_0, a_0, s_1, a_1, \ldots$ The return is the discounted cumulative reward: $$G_0 = \sum_{t=0}^{\infty} \gamma^t R_{t+1}$$ The return is a random variable (because transitions are stochastic). The objective is to find a policy that maximizes the expected return from every state: $$\pi^* = \arg\max_\pi ; \mathbb{E}\pi\left[\sum{t=0}^{\infty} \gamma^t R_{t+1} \mid S_0 = s\right] \quad \forall s \in \mathcal{S}$$

Key Insight: A Single Policy Can Simultaneously Maximize From Every State

A remarkable (and non-obvious) fact: there exists a single policy $\pi^*$ that is simultaneously optimal starting from every state. You don't need different policies for different starting states. This is unique to MDPs and is what makes them tractable.

import random

def run_episode(mdp: MDP, policy: dict, start: State, max_steps=1000) -> float:
    """
    Simulate one episode under a deterministic policy.
    Returns discounted total reward.
    """
    s = start
    total_reward = 0.0
    discount = 1.0
    for _ in range(max_steps):
        a = policy[s]
        outcomes = mdp.transitions(s, a)  # [(prob, s', r), ...]
        probs    = [o[0] for o in outcomes]
        idx      = random.choices(range(len(outcomes)), weights=probs)[0]
        _, s_next, r = outcomes[idx]
        total_reward += discount * r
        discount     *= mdp.gamma
        s = s_next
    return total_reward

The State-Value Function

The state-value function (or just value function) $V^\pi(s)$ under policy $\pi$ is the expected return starting from state $s$ and following $\pi$ thereafter: $$V^\pi(s) = \mathbb{E}\pi\left[\sum{t=0}^{\infty} \gamma^t R_{t+1} \mid S_0 = s\right]$$

The optimal value function $V^(s)$ is the maximum over all policies: $$V^(s) = \max_\pi V^\pi(s)$$

If you know $V^*$, extracting the optimal policy is easy — at each state, pick the action that leads to the highest expected value.

The Action-Value Function (Q-function)

The action-value function (Q-function) $Q^\pi(s, a)$ is the expected return starting from state $s$, taking action $a$ (which may or may not be what $\pi$ would do), then following $\pi$: $$Q^\pi(s, a) = \mathbb{E}\pi\left[\sum{t=0}^{\infty} \gamma^t R_{t+1} \mid S_0 = s, A_0 = a\right]$$

The optimal Q-function: $$Q^(s, a) = \mathbb{E}\left[R_{t+1} + \gamma \max_{a'} Q^(S_{t+1}, a') \mid S_t = s, A_t = a\right]$$

Why Q* Is Critical for Model-Free RL

$V^$ is beautiful for analysis, but to extract the optimal policy from $V^$ you need to know the transition model $T$ (to compute the expected next-state value). $Q^$ encodes the action quality directly — to find the optimal action at state $s$, just compute $\arg\max_a Q^(s, a)$, no model needed. This is why Q-learning (a model-free RL algorithm) targets $Q^*$.

Relationship

$$V^(s) = \max_a Q^(s, a)$$ $$Q^(s, a) = \sum_{s'} T(s, a, s') \left[R(s, a, s') + \gamma V^(s')\right]$$

# V* to optimal policy extraction (requires model T)
def extract_policy(mdp: MDP, V: dict) -> dict:
    policy = {}
    for s in mdp.states:
        best_a, best_q = None, float('-inf')
        for a in mdp.actions:
            q = sum(p * (r + mdp.gamma * V[s_next])
                    for p, s_next, r in mdp.transitions(s, a))
            if q > best_q:
                best_q, best_a = q, a
        policy[s] = best_a
    return policy

# Q* to optimal policy (model-free — no T needed at decision time)
def extract_policy_from_Q(Q: dict) -> dict:
    # Q[(s,a)] -> value
    states  = set(s for s, a in Q)
    actions = set(a for s, a in Q)
    return {s: max(actions, key=lambda a: Q[(s, a)]) for s in states}

Recursive Structure of Value

The key mathematical insight behind MDPs is that the value function has a recursive structure. The value of being in state $s$ under policy $\pi$ equals the immediate reward you expect plus the discounted value of the next state:

$$V^\pi(s) = \sum_a \pi(a|s) \sum_{s'} T(s,a,s') \left[R(s,a,s') + \gamma V^\pi(s')\right]$$

This is the Bellman expectation equation. It says: the value of a state is the expected reward from taking your policy's action, plus the discounted expected value of where you land.

Bellman Optimality Equation

For the optimal value function, the Bellman optimality equation replaces the policy average with a max:

$$\boxed{V^(s) = \max_a \sum_{s'} T(s,a,s') \left[R(s,a,s') + \gamma V^(s')\right]}$$

This is one of the most important equations in all of AI. Let's unpack it:

• For each action $a$ you could take from $s$...
• Sum over all possible next states $s'$ (weighted by transition probability $T(s,a,s')$)...
• The value of landing in $s'$ is $R(s,a,s')$ (immediate) plus $\gamma V^*(s')$ (discounted future optimal value).
• Take the max over actions — because at $s$, the optimal agent acts optimally.

Why Is This Useful?

The Bellman optimality equation is a system of equations — one per state. There are $|\mathcal{S}|$ unknowns ($V^*(s)$ for each state) and $|\mathcal{S}|$ equations. If the equations were linear, you could solve them directly. They're not linear (because of the max), but they define a contraction mapping — a key fact that guarantees iterative methods converge.

The Contraction Mapping Theorem

Define the Bellman optimality operator $\mathcal{T}$: $$(\mathcal{T}V)(s) = \max_a \sum_{s'} T(s,a,s')\left[R(s,a,s') + \gamma V(s')\right]$$

$V^$ is the unique fixed point of $\mathcal{T}$: $\mathcal{T}V^ = V^*$. Moreover, $\mathcal{T}$ is a $\gamma$-contraction in the $\ell^\infty$ norm: $$|\mathcal{T}V - \mathcal{T}U|\infty \leq \gamma |V - U|\infty$$

By the Banach fixed-point theorem, repeated application of $\mathcal{T}$ to any initial $V_0$ converges to $V^*$ at a geometric rate $\gamma^k$. This is the theoretical foundation of value iteration.

# Bellman operator applied to a value function V
# Returns the "one-step improved" value function
def bellman_operator(mdp: MDP, V: dict) -> dict:
    V_new = {}
    for s in mdp.states:
        # max over actions of Q(s, a)
        V_new[s] = max(
            sum(p * (r + mdp.gamma * V[s_next])
                for p, s_next, r in mdp.transitions(s, a))
            for a in mdp.actions
        )
    return V_new

# The contraction: ||TV - TU||_inf <= gamma * ||V - U||_inf
# guarantees geometric convergence to V*

Algorithm

Value Iteration (Bellman, 1957) directly exploits the contraction property. It initializes $V$ arbitrarily and repeatedly applies the Bellman operator until convergence:

Initialize V(s) = 0 for all s
Repeat:
    For each state s:
        V(s) ← max_a Σ_{s'} T(s,a,s') [R(s,a,s') + γ V(s')]
Until max_s |V_new(s) - V_old(s)| < ε
Extract policy: π(s) = argmax_a Σ_{s'} T(s,a,s') [R(s,a,s') + γ V(s')]

Convergence Guarantee

After $k$ iterations: $$|V_k - V^*|\infty \leq \frac{\gamma^k}{1-\gamma} |V_1 - V_0|\infty$$

To achieve $|V_k - V^*|_\infty < \epsilon$, you need $k \geq \frac{\log(1/\epsilon(1-\gamma))}{\log(1/\gamma)}$ iterations.

Complexity

Each iteration scans all states and actions:

• Time per iteration: $O(|\mathcal{S}|^2 |\mathcal{A}|)$ in the worst case (each transition can reach any state)
• Memory: $O(|\mathcal{S}||\mathcal{A}|)$ for the transition table
• Number of iterations: Polynomial in $\frac{1}{1-\gamma}$ and $\log(1/\epsilon)$

This is polynomial time in the size of the MDP description. Value iteration is one of the few algorithms in AI with such clean guarantees.

Asynchronous Value Iteration

You don't need to update all states simultaneously in each sweep. Asynchronous value iteration updates states one at a time (in any order). As long as every state is updated infinitely often, the algorithm still converges to $V^*$. This is important for large state spaces where a full sweep is expensive — you can prioritize states that are likely to change the most (prioritized sweeping).

def value_iteration(mdp: MDP, epsilon: float = 1e-6) -> dict:
    """
    Solve MDP exactly via value iteration.
    Returns V* as a dict {state: value}.
    """
    V = {s: 0.0 for s in mdp.states}

    while True:
        delta = 0.0
        V_new = {}
        for s in mdp.states:
            # Bellman optimality backup
            v = max(
                sum(p * (r + mdp.gamma * V[s_next])
                    for p, s_next, r in mdp.transitions(s, a))
                for a in mdp.actions
            )
            V_new[s] = v
            delta = max(delta, abs(v - V[s]))
        V = V_new
        if delta < epsilon:
            break

    # Extract optimal policy
    policy = extract_policy(mdp, V)
    return V, policy

# Convergence criterion: max change across all states < epsilon
# Error bound: ||V_k - V*||_inf <= gamma/(1-gamma) * delta_k

A Different Approach

Policy Iteration (Howard, 1960) interleaves two phases:

1. Policy Evaluation: Given current policy $\pi$, compute $V^\pi$ exactly (or approximately).
2. Policy Improvement: Given $V^\pi$, compute a better (greedy) policy $\pi'$.

Repeat until the policy doesn't change.

Policy Evaluation

For a fixed policy $\pi$, the Bellman expectation equations become a linear system: $$V^\pi(s) = \sum_{s'} T(s, \pi(s), s') \left[R(s, \pi(s), s') + \gamma V^\pi(s')\right]$$

This is $|\mathcal{S}|$ linear equations in $|\mathcal{S}|$ unknowns. You can solve it exactly with linear algebra ($O(|\mathcal{S}|^3)$ with Gaussian elimination), or iteratively by applying the Bellman expectation operator repeatedly (converges since $\gamma < 1$).

Policy Improvement Theorem

Given $V^\pi$, define the greedy policy: $$\pi'(s) = \arg\max_a \sum_{s'} T(s,a,s')\left[R(s,a,s') + \gamma V^\pi(s')\right]$$

Theorem: $V^{\pi'}(s) \geq V^\pi(s)$ for all $s$, with strict inequality somewhere unless $\pi = \pi^*$.

This is the Policy Improvement Theorem — switching to the greedy policy with respect to $V^\pi$ can only make things better or keep them the same. When the policy stops changing, it must be optimal.

Convergence: Finite Steps to Optimality

Since the number of deterministic policies is $|\mathcal{A}|^{|\mathcal{S}|}$ (finite), and each policy iteration step strictly improves the policy (or stops), policy iteration terminates in finite steps — typically far fewer than value iteration. In practice, policy iteration often converges in 5–20 iterations for moderate MDPs, while value iteration might need hundreds.

Modified Policy Iteration

A hybrid: instead of solving the linear system exactly in the evaluation step, do only $k$ steps of iterative evaluation. With $k=1$, you recover value iteration. With $k \to \infty$, you get full policy iteration. $k = 5$–$20$ often hits the sweet spot in practice.

import numpy as np

def policy_evaluation(mdp: MDP, policy: dict, epsilon=1e-6) -> dict:
    """Iterative policy evaluation for deterministic policy."""
    V = {s: 0.0 for s in mdp.states}
    while True:
        delta = 0.0
        for s in mdp.states:
            v = sum(p * (r + mdp.gamma * V[s_next])
                    for p, s_next, r in mdp.transitions(s, policy[s]))
            delta = max(delta, abs(v - V[s]))
            V[s] = v
        if delta < epsilon:
            return V

def policy_iteration(mdp: MDP) -> dict:
    """Solve MDP via policy iteration. Returns optimal policy."""
    # Start with any policy (e.g., always take first action)
    policy = {s: mdp.actions[0] for s in mdp.states}

    while True:
        V = policy_evaluation(mdp, policy)
        policy_new = extract_policy(mdp, V)   # greedy improvement

        if policy_new == policy:
            return policy, V  # converged to optimal policy
        policy = policy_new

Side-by-Side Comparison

When to Use Which

• Value iteration is simpler to implement, naturally asynchronous, and easy to parallelize. It's preferred when states are numerous and you can afford approximate convergence.
• Policy iteration converges in fewer outer iterations but each iteration is more expensive. It's preferred for moderate-sized MDPs where you want the exact optimal policy without many sweeps.
• Linear programming is a third exact method: solve $V^* = \arg\min_V \sum_s V(s)$ subject to $V(s) \geq \sum_{s'} T(s,a,s')[R(s,a,s') + \gamma V(s')]$ for all $s, a$. It's $O(|\mathcal{S}|^{3.5})$ with interior-point methods — competitive for very small state spaces.

Practical Reality

For real-world problems, the state space is often enormous (the classic curse of dimensionality: a robot arm with 10 joints, each with 1000 positions, has $10^{30}$ states — table-based methods are hopeless). This is where approximate dynamic programming, function approximation, and reinforcement learning enter.

Why Exact Methods Break Down

Both value iteration and policy iteration require:

1. An explicit representation of $T(s, a, s')$ for all $(s, a, s')$ triples — the model.
2. A table for $V(s)$ or $Q(s, a)$ for all states — tabular representation.

Both assumptions fail in practice. For a robot, the state is a continuous vector of joint angles and velocities. For a game of Go, the state space is $3^{361} \approx 10^{172}$. For language modeling, the state is a sequence of tokens — essentially infinite.

The solution path splits into two directions:

• Model-based planning with function approximation: Approximate $T$ and $R$ with learned models, and approximate $V^$ or $Q^$ with neural networks.
• Model-free reinforcement learning: Don't learn $T$ at all. Instead, estimate $V^$ or $Q^$ directly from experience (sampled transitions).

What RL Assumes

In reinforcement learning, the agent does NOT know $T$ or $R$. It only sees:

• The current state $s_t$
• The reward $r_{t+1}$ received after taking action $a_t$
• The next state $s_{t+1}$

Despite this, the goal is still to solve the same MDP. The Bellman equations still define the objective. The algorithms are stochastic approximation versions of value/policy iteration:

The Q-learning Update (Bridge Equation)

The canonical RL update that mirrors the Bellman optimality equation: $$Q(s, a) \leftarrow Q(s, a) + \alpha \left[ r + \gamma \max_{a'} Q(s', a') - Q(s, a) \right]$$

The term in brackets is the Temporal Difference (TD) error — the difference between the current estimate $Q(s,a)$ and the one-step Bellman target $r + \gamma \max_{a'} Q(s', a')$. This is a stochastic approximation to applying $\mathcal{T}$ — the Bellman operator — using a single sampled transition instead of the full expectation over all next states.

# Q-learning: model-free RL that estimates Q* via TD updates
import random

def q_learning(env, states, actions, gamma=0.9, alpha=0.1,
               epsilon=0.1, episodes=10000):
    """
    env.step(s, a) -> (s_next, reward, done)
    Returns learned Q table.
    """
    Q = {(s, a): 0.0 for s in states for a in actions}

    for _ in range(episodes):
        s = env.reset()
        done = False
        while not done:
            # epsilon-greedy exploration
            if random.random() < epsilon:
                a = random.choice(actions)
            else:
                a = max(actions, key=lambda a: Q[(s, a)])

            s_next, r, done = env.step(s, a)

            # Bellman target (uses max over next actions)
            target = r + gamma * max(Q[(s_next, a_)] for a_ in actions)

            # TD update — stochastic approximation of Bellman operator
            Q[(s, a)] += alpha * (target - Q[(s, a)])

            s = s_next
    return Q

The Problem

Consider a 3-state MDP: states ${A, B, C}$ where $C$ is terminal (absorbing, $V^*(C) = 0$). Actions: ${left, right}$. Transitions are deterministic. Rewards are received on leaving a state. $\gamma = 0.9$.

We want $V^(A)$, $V^(B)$, $V^*(C)$.

Value Iteration

Initialize: $V_0 = {A: 0, B: 0, C: 0}$

Iteration 1:

• $V_1(C) = 0$ (terminal)
• $V_1(B) = \max[right: 10 + 0.9 \cdot 0, ; left: 0 + 0.9 \cdot 0] = \max[10, 0] = 10$
• $V_1(A) = \max[right: 0 + 0.9 \cdot 0, ; left: -1 + 0.9 \cdot 0] = \max[0, -1] = 0$

Iteration 2:

• $V_2(C) = 0$
• $V_2(B) = \max[right: 10 + 0.9 \cdot 0, ; left: 0 + 0.9 \cdot 0] = 10$ (no change)
• $V_2(A) = \max[right: 0 + 0.9 \cdot 10, ; left: -1 + 0.9 \cdot 0] = \max[9, -1] = 9$

Iteration 3:

• $V_3(A) = \max[right: 0 + 0.9 \cdot 10, ; left: -1 + 0.9 \cdot 9] = \max[9, 7.1] = 9$ (no change)

Converged! $V^* = {A: 9, B: 10, C: 0}$.

Optimal Policy

$\pi^(A) = right$, $\pi^(B) = right$. The agent goes $A \to B \to C$, collecting reward $0 + 10 = 10$, discounted: $0 + 0.9 \cdot 10 = 9$ starting from A. Confirmed.

Intuition Check

Why does $V^*(A) = 9$ and not $10$? Because to get the $+10$ reward at $B$, the agent must first travel from $A$ to $B$ (one step, zero reward), and the future is discounted: $\gamma^1 \times 10 = 0.9 \times 10 = 9$.

# Verify by hand calculation
mdp_states  = ['A', 'B', 'C']
mdp_actions = ['left', 'right']

# T[s][a] = [(prob, s', r)]
T = {
    'A': {'right': [(1.0, 'B',  0.0)],
          'left':  [(1.0, 'A', -1.0)]},
    'B': {'right': [(1.0, 'C', 10.0)],
          'left':  [(1.0, 'A',  0.0)]},
    'C': {'right': [(1.0, 'C',  0.0)],   # absorbing
          'left':  [(1.0, 'C',  0.0)]},
}

V = {'A': 0.0, 'B': 0.0, 'C': 0.0}
gamma = 0.9

for iteration in range(10):
    V_new = {}
    for s in mdp_states:
        V_new[s] = max(
            sum(p * (r + gamma * V[s2])
                for p, s2, r in T[s][a])
            for a in mdp_actions
        )
    V = V_new
    print(f"Iter {iteration+1}: {V}")

# Output:
# Iter 1: {'A': 0.0,  'B': 10.0, 'C': 0.0}
# Iter 2: {'A': 9.0,  'B': 10.0, 'C': 0.0}
# Iter 3: {'A': 9.0,  'B': 10.0, 'C': 0.0}  <- converged

The Fully Observable Assumption

Everything above assumed the agent observes the true state $s_t$ directly. This is the fully observable MDP. But real environments are often only partially observable: a robot doesn't know its exact position (sensor noise), a card game player doesn't see opponents' hands, a medical AI doesn't know the patient's true underlying health state.

The POMDP Extension

A Partially Observable MDP (POMDP) adds:

• An observation space $\mathcal{O}$
• An observation function $O(s', a, o) = P(O_{t+1} = o \mid S_{t+1} = s', A_t = a)$

The agent receives observation $o_t$ instead of the true state $s_t$. The optimal strategy must maintain a belief state $b_t$ — a probability distribution over states: $b_t(s) = P(S_t = s \mid o_1, a_1, o_2, a_2, \ldots, o_t)$.

Critically, the belief state IS a sufficient statistic for the future: $b_t$ is Markov over the history. So POMDPs can be re-framed as MDPs over the belief space — but the belief space is a continuous simplex of dimension $|\mathcal{S}| - 1$, making exact solution exponentially hard in general (PSPACE-complete).

Why This Matters

POMDPs are the principled model for real-world AI agents. Modern deep RL handles partial observability through recurrent neural networks (LSTM/GRU state as implicit belief) or explicit memory mechanisms. Understanding that the POMDP belief update is Bayesian inference — and that the Bellman equation still applies to beliefs — connects MDP theory all the way to modern transformers-as-agents research.

The takeaway: full MDPs are the essential building block, but real problems demand extending the framework — and the structure (Bellman equations, value functions, policy optimization) carries over intact.

# Belief state update (Bayesian filter) for a POMDP
def belief_update(belief: dict, a, o,
                 T: dict,   # T[s][a] = [(p, s', r)]
                 O: dict    # O[s'][a][o] = P(o | s', a)
                ) -> dict:
    """
    Update belief b over states given action a and observation o.
    belief: {state: probability}
    Returns new normalized belief.
    """
    new_belief = {}
    for s_next in belief:
        # P(s' | b, a) = Σ_s T(s,a,s') * b(s)
        pred = sum(p * belief[s]
                   for s in belief
                   for p, sp, _ in T[s][a]
                   if sp == s_next)
        # P(o | s', a) * P(s' | b, a)
        new_belief[s_next] = O[s_next][a][o] * pred

    # Normalize
    total = sum(new_belief.values())
    return {s: v / total for s, v in new_belief.items()}