Bellman equation for value iteration

Background

Value iteration solves a known MDP by repeatedly applying the Bellman optimality update. The value of a state is the best expected return achievable from it: over all actions, take the one with the highest expected immediate reward plus discounted next-state value. Sweeping this update over all states until it converges yields the optimal value function $V^*$.

Problem statement

Implement bellman_update(V, transitions, gamma) performing one synchronous Bellman optimality sweep:

$$V'(s) = \max_a \sum_{(p,\,s',\,r,\,\text{done})} p\,\big(r + \gamma\,(1-\text{done})\,V(s')\big)$$

Input

• V — np.ndarray (n_states,), current value estimates.
• transitions — list indexed by state; transitions[s] is a dict action -> list of (prob, next_state, reward, done) tuples.
• gamma — float, discount factor.

Output

An np.ndarray (n_states,): the updated values new_V.

Examples

Example 1

transitions = [
  {0: [(1.0, 0, 0.0, False)], 1: [(1.0, 1, 1.0, False)]},
  {0: [(1.0, 0, 0.0, False)], 1: [(1.0, 1, 1.0, True)]},
]
V = [0.0, 0.0], gamma = 0.9
Output: [1.0, 1.0]

Explanation: from state 0 the best action (1) gives reward 1 → value 1. From state 1, action 1 gives reward 1 and ends the episode (done=True), so no future value is added → value 1.

Constraints

• For each state, take the max over actions of the expected backup.
• A transition with done=True contributes only its reward (no $\gamma V(s')$ term).
• The update is synchronous: compute new_V from the old V, don't update in place mid-sweep.

Notes

• Iterating bellman_update to a fixed point gives the optimal values $V^*$; the greedy policy w.r.t. $V^*$ is optimal.
• Replacing the max with a fixed policy's action gives policy evaluation instead of value iteration.