Iterative Policy Evaluation in a Nutshell

# Constants for actions
UP, DOWN, LEFT, RIGHT = 0, 1, 2, 3
ACTIONS = [UP, DOWN, LEFT, RIGHT]
ACTION_DELTAS = {
    UP: (-1, 0),
    DOWN: (1, 0),
    LEFT: (0, -1),
    RIGHT: (0, 1),
}

# Maze config
MAZE_ROWS = 5
MAZE_COLS = 5
WALLS = {(0, 3), (1, 1), (1, 3), (2, 1), (3, 3), (4, 0), (4, 1)}
GOAL = (0, 4)

class State:
    def __init__(self, name: int, is_terminal: bool = False):
        self.name = name
        self.is_terminal = is_terminal

class Env:
    def __init__(self):
        self.states = []
        self.state_map = {}
        for r in range(MAZE_ROWS):
            for c in range(MAZE_COLS):
                if (r, c) not in WALLS:
                    is_terminal = (r, c) == GOAL
                    state = State((r, c), is_terminal)
                    self.states.append(state)
                    self.state_map[(r, c)] = state
        self.actions = ACTIONS

    def step(self, state: State, action: int) -> List[Tuple[State, float, float]]:
        r, c = state.name
        dr, dc = ACTION_DELTAS[action]
        nr, nc = r + dr, c + dc

        # Stay in place if moving out of bounds or into a wall
        if (nr < 0 or nr >= MAZE_ROWS or
            nc < 0 or nc >= MAZE_COLS or
            (nr, nc) in WALLS):
            next_pos = (r, c)
        else:
            next_pos = (nr, nc)

        next_state = self.state_map[next_pos]
        reward = 0.0 if next_pos == GOAL else -1.0
        return [(next_state, reward, 1.0)]  # deterministic

import math
import random
from typing import List, Tuple

def iterative_policy_eval(agent: Agent, env: Env, gamma: float = 0.9, threshold: float = 1e-6):
    values = {state.name: 0.0 for state in env.states}

    while True:
        delta = 0
        for state in env.states:
            if state.is_terminal:
                continue

            old_value = values[state.name]
            new_value = 0.0

            for action in env.actions:
                action_prob = agent.policy(state, action)
                for next_state, reward, prob in env.step(state, action):
                    new_value += action_prob * prob * (reward + gamma * values[next_state.name])

            values[state.name] = new_value
            delta = max(delta, abs(old_value - new_value))

        if delta < threshold:
            break

    return values

if __name__ == '__main__':
    env = Env()
    agent = Agent()
    values = iterative_policy_eval(agent, env)

    for state in sorted(env.states, key=lambda s: s.name):
        print(f"State {state.name}: V = {values[state.name]:.2f}")

class Agent:
    def policy(self, state: State, action: int) -> float:
        return 1 if action == LEFT else 0

State (0, 0): V = -10.00
State (0, 1): V = -10.00
State (0, 2): V = -10.00
State (0, 4): V = 0.00
State (1, 0): V = -10.00
State (1, 2): V = -10.00
State (1, 4): V = -10.00
State (2, 0): V = -10.00
State (2, 2): V = -10.00
State (2, 3): V = -10.00
State (2, 4): V = -10.00
State (3, 0): V = -10.00
State (3, 1): V = -10.00
State (3, 2): V = -10.00
State (3, 4): V = -10.00
State (4, 2): V = -10.00
State (4, 3): V = -10.00
State (4, 4): V = -10.00

class Agent:
    def policy(self, state: State, action: int) -> float:
        return 1.0 / 4

State (0, 0): V = -9.96
State (0, 1): V = -9.93
State (0, 2): V = -9.87
State (0, 4): V = 0.00
State (1, 0): V = -9.96
State (1, 2): V = -9.76
State (1, 4): V = -4.53
State (2, 0): V = -9.96
State (2, 2): V = -9.54
State (2, 3): V = -8.89
State (2, 4): V = -7.75
State (3, 0): V = -9.93
State (3, 1): V = -9.87
State (3, 2): V = -9.76
State (3, 4): V = -8.82
State (4, 2): V = -9.75
State (4, 3): V = -9.64
State (4, 4): V = -9.37

class Agent:
    def policy(self, state: State, action: int) -> float:
        if action == UP and state.name == (1,4): return 1
        if state.name == (1,4): return 0
        return 1.0 / 4

State (0, 0): V = -9.92
State (0, 1): V = -9.87
State (0, 2): V = -9.77
State (0, 4): V = 0.00
State (1, 0): V = -9.94
State (1, 2): V = -9.56
State (1, 4): V = 0.00
State (2, 0): V = -9.92
State (2, 2): V = -9.15
State (2, 3): V = -7.97
State (2, 4): V = -5.88
State (3, 0): V = -9.87
State (3, 1): V = -9.77
State (3, 2): V = -9.56
State (3, 4): V = -7.85
State (4, 2): V = -9.55
State (4, 3): V = -9.35
State (4, 4): V = -8.85

class Agent:
    def __init__(self):
        self.policy_map = {
            (0,0): RIGHT,
            (0,1): RIGHT,
            (0,2): DOWN,
            (1,2): DOWN,
            (2,2): RIGHT,
            (2,3): RIGHT,
            (2,4): UP,
            (1,4): UP
        }

    def policy(self, state, action):
        # If state is a State object, extract the tuple
        if hasattr(state, "name"):
            state = state.name

        optimal_action = self.policy_map.get(state)
        if optimal_action is None:
            return 0
        return 1 if action == optimal_action else 0
        ```

* This policy maximizes expected rewards by minimizing steps to the goal.


## How its used in RL

* The **value function** measures how good it is to be in a state, under a given policy.
* **Iterative policy evaluation** refines these values over time.
* By updating the policy to choose actions that **maximize future rewards**, we can move toward optimal behavior.
* **Exercise**: Use the Bellman Optimality Equation to solve this using a system of linear equations. Do you get the same values for each policy?