Value Iteration Step - Reinforcement Learning

Value Iteration Step

What you'll build

One synchronous sweep of value iteration: back up every state with the Bellman optimality operator, using a known model of the MDP.

Background

When transition probabilities and rewards are known, dynamic programming can compute optimal values directly. Value iteration applies the Bellman optimality backup to all states and repeats until V stops changing; the greedy policy with respect to the result is optimal. Q-learning is the sample-based, model-free version of this backup.

Formula

For every state s:

Q(s, a)   = Σ_{(p, r, s')} p · (r + γ · V_old[s'])
V_new[s]  = max_a Q(s, a)

Algorithm (one sweep)

1. Start from V_new = zeros(n_states).
2. For each state s that has an entry in transitions, compute Q(s, a) for each action a from V_old (the input), and set V_new[s] = max_a Q(s, a).
3. States with no entry in transitions are terminal: V_new[s] = 0.
4. Return V_new.

Inputs

• V: 1-D array of current values, length n_states
• transitions: dict, transitions[s][a] = list of (prob, reward, next_state) tuples
• gamma: discount factor in [0, 1]

Expected return format

The new value array (a list also passes), rounded to 1 decimal.

Edge cases and conventions

• Synchronous update: every Q(s, a) reads the old V. Don't let V_new[0] leak into the backup of state 1 within the same sweep.
• A state missing from transitions gets 0.

Constraints

NumPy only. Loops over states and actions are fine.

Follow-up

• In-place (Gauss–Seidel) value iteration uses fresh values as soon as they exist. Does it still converge? Is it faster?
• Compare value iteration with policy iteration. Cost per sweep: O(|S|²·|A|).

Input: V = [0, 0], transitions = {0: {0: [(1.0, 1.0, 0)], 1: [(1.0, 5.0, 1)]}, 1: {0: [(1.0, 2.0, 0)], 1: [(1.0, 0.0, 1)]}}, gamma = 0.9
Output: [5.0, 2.0]

Input: Same MDP, V = [5, 2]
Output: [6.8, 6.5]

import numpy as np

def value_iteration_step(V: np.ndarray, transitions: dict,
                         gamma: float) -> np.ndarray:
    """
    Perform one synchronous sweep of value iteration.

    Args:
        V: (n_states,) current values
        transitions: transitions[s][a] = [(prob, reward, next_state), ...]
        gamma: discount factor

    Returns:
        (n_states,) new values, V_new[s] = max_a Q(s, a), rounded to 1 decimal
        (states without transitions are terminal and get 0)
    """
    V = np.array(V, dtype=float)
    n_states = len(V)
    V_new = np.zeros(n_states)

    # Your code here
    pass