Value Iteration on a Small MDP

1. Background Knowledge

Value Iteration is a fundamental algorithm in Reinforcement Learning used to solve Markov Decision Processes (MDP). It is a specific instance of Generalized Policy Iteration (GPI), which combines policy evaluation and policy improvement into a single, unified update step. Unlike Policy Iteration, which performs a full convergence of the value function for a fixed policy before improving it, Value Iteration truncates the evaluation phase to a single sweep. This makes it computationally cheaper per iteration, though it may require more iterations to converge.

The core theoretical guarantee of Value Iteration relies on the Bellman Optimality Operator. This operator is a contraction mapping with respect to the max-norm (infinity norm) when the discount factor $\gamma$ is strictly less than 1. Specifically, applying the operator reduces the distance between the current value estimate and the true optimal value function $v_*$ by a factor of $\gamma$. This property ensures that, regardless of the initial value function, the algorithm will converge to the unique fixed point $v_*$ as the number of iterations approaches infinity.

In an MDP, the environment is defined by states, actions, transition probabilities, and rewards. The value function $v(s)$ represents the expected cumulative discounted reward starting from state $s$ and following the optimal policy thereafter. The optimal policy $\pi_*(s)$ is the deterministic action that maximizes the expected immediate reward plus the discounted value of the next state. Value Iteration computes $v_*$ first and then extracts $\pi_*$ by choosing the action that achieves the maximum in the Bellman optimality equation for each state.

2. Algorithm Approach

The approach is to iteratively update the value function using the Bellman Optimality Equation. We start with an arbitrary value function (typically all zeros) and repeatedly apply the update rule:

$v_{k+1}(s) \leftarrow \max_a \Big[ r(s,a) + \gamma \sum_{s'} p(s' \mid s,a)\, v_k(s') \Big]$

This update is synchronous, meaning we compute the new values for all states based on the old values from the previous iteration before updating any state. We continue this process until the change in values between iterations, measured by the max-norm difference $\delta = \max_s |v_{k+1}(s) - v_k(s)|$, falls below a small threshold $\theta$. Once convergence is reached, we derive the optimal policy by selecting the action that maximizes the right-hand side of the Bellman equation for each state, breaking ties by choosing the action with the lowest index.

3. Step-by-Step Strategy

1. Initialize: Create a value array v of size equal to the number of states, initialized to zeros.
2. Iterate: Loop up to max_sweeps or until convergence: a. Create a copy of the current value array v_new to store updated values. b. For each state s: i. For each action a:

• Calculate the expected return: $Q(s,a) = R[s][a] + \gamma \sum_{s'} P[s][a][s'] \cdot v[s']$.
• Note: $P[s][a]$ is a list of probabilities for transitioning to each next state $s'$. ii. Set v_new[s] to the maximum $Q(s,a)$ across all actions. c. Compute the maximum absolute difference $\delta = \max_s |v\_new[s] - v[s]|$. d. Update v = v_new. e. If $\delta < \theta$, break the loop.

3. Extract Policy: Initialize an empty policy list. For each state s: a. Recompute $Q(s,a)$ for all actions using the converged v. b. Find the action a that maximizes $Q(s,a)$. If there are ties, choose the one with the smallest index. c. Append this action to the policy list.
4. Return: Return the tuple (v, policy).

4. Common Pitfalls

• Asynchronous Updates: Do not update v[s] in-place while computing other states in the same sweep. This violates the synchronous update requirement and can lead to incorrect convergence or oscillation. Always use a separate v_new array.
• Tie-Breaking: When multiple actions yield the same maximum Q-value, the problem specifies choosing the lowest action index. Failing to handle this correctly will result in a wrong policy, even if the values are correct.
• Transition Probabilities: Ensure you correctly sum over all possible next states $s'$ weighted by their transition probabilities $p(s' \mid s,a)$. A common error is forgetting to multiply by the probability or summing incorrectly.
• Convergence Check: The stopping condition is based on the max-norm of the difference between consecutive value functions, not the sum of squared errors or any other metric. Ensure $\delta$ is computed correctly as $\max_s |v_{new}[s] - v[s]|$.
• Discount Factor: Remember that $\gamma$ must be applied to the expected future value, not the immediate reward. The formula is $r + \gamma \cdot \text{expected\_future\_value}$.

5. Time & Space Complexity

• Time Complexity: Each iteration involves computing the Q-value for every state-action pair. For $S$ states and $A$ actions, and assuming transitions to $S$ next states, one iteration takes $O(S \cdot A \cdot S) = O(S^2 A)$. The number of iterations required for convergence depends on $\gamma$ and $\theta$, but is generally bounded by $O(\log(1/\theta))$ due to the contraction property. Thus, the total time complexity is $O(S^2 A \log(1/\theta))$.
• Space Complexity: We need to store the value function v and v_new, both of size $S$. The policy is also of size $S$. The input MDP structures $P$ and $R$ take $O(S^2 A)$ space. Therefore, the auxiliary space complexity is $O(S)$, while the total space complexity including input is $O(S^2 A)$.