Bellman Expectation (State Value) - Reinforcement Learning

Bellman Expectation (State Value)

What you'll build

A one-state backup: compute V(s) from the possible transitions out of s and the current estimates of the next-state values.

Background

The Bellman expectation equation says a state's value is the expected immediate reward plus the discounted value of where you land. It is the consistency condition behind policy evaluation, TD learning and every critic in actor-critic methods.

Formula

For a policy π, the general form is

V(s) = Σ_a π(a|s) · Σ_{s'} P(s'|s,a) · [R(s,a,s') + γ·V(s')]

Here the policy is already folded in, so you get a flat list of outcomes, each with its total probability p = π(a|s)·P(s'|s,a):

V(s) = Σ_i p_i · (r_i + γ · V[s'_i])

Inputs

• transitions: list of (probability, reward, next_state_index) tuples; probabilities sum to 1. The same next state may appear more than once (different rewards).
• gamma: discount factor in [0, 1]
• next_values: 1-D array, next_values[j] is the current estimate V(s_j)

Expected return format

A Python float: V(s) rounded to 2 decimals.

Edge cases and conventions

• gamma = 0 gives the expected immediate reward.
• Only the states that appear in transitions matter; other entries of next_values are ignored.

Constraints

NumPy only. A loop over transitions is fine; a vectorized version gathers next_values[next_states].

Follow-up

• Applying this backup to every state repeatedly is iterative policy evaluation. Why does it converge for γ < 1? (It is a γ-contraction.)
• Write the same equation in matrix form, V = R + γPV, and solve it in closed form.

Input: transitions = [(0.5, 1, 0), (0.5, 2, 1)], gamma = 0.9, next_values = [5, 10]
Output: 8.25

Input: transitions = [(1.0, 5, 0)], gamma = 0.9, next_values = [10]
Output: 14.0

import numpy as np

def bellman_value(transitions: list, gamma: float,
                  next_values: np.ndarray) -> float:
    """
    Compute V(s) with the Bellman expectation equation.

    Args:
        transitions: list of (probability, reward, next_state_index) tuples
        gamma: discount factor
        next_values: (n_states,) current estimates V(s')

    Returns:
        V(s) = sum_i p_i * (r_i + gamma * V[s'_i]), rounded to 2 decimals
    """
    next_values = np.array(next_values)
    # Your code here
    pass