Module 2. State, Value, and Bellman Equations

from __future__ import annotations

TEXT = {'profile': 'Profile: {mode}; independent seeds: {seeds}; CliffWalking episodes: {episodes}', 'ordinary_transition': 'ordinary transition', 'true_terminal': 'true terminal', 'external_truncation': 'external truncation', 'gridworld_values': 'Exact random-policy values in Gridworld:', 'lake_sweeps': 'Value-iteration sweeps', 'lake_values': 'Optimal FrozenLake values:', 'lake_policy': 'Greedy actions (0=left, 1=down, 2=right, 3=up):', 'lake_panel_titles': ['V₀', 'after 1 sweep', 'after 5 sweeps', 'converged'], 'lake_plot_title': 'Value propagation through FrozenLake', 'corridor_truth': 'Exact corridor values:', 'mc': 'first-visit MC', 'td0': 'TD(0), α=0.1', 'td3': '3-step TD, α=0.1', 'episodes': 'Episodes (log scale)', 'rmse': 'RMSE over non-terminal states', 'prediction_title': 'Corridor prediction: median error and interquartile range', 'prediction_last_budget': 'Error at a budget of {budget} episodes:', 'training_episode': 'Training episode', 'rolling_return': 'Rolling mean return (window 25)', 'cliff_learning_title': 'CliffWalking learning across random seeds', 'cliff_summary': 'Run summary:', 'train_tail': 'median training return over the last 100 episodes={value:7.2f}', 'greedy_summary': 'median completed greedy-route return={median:5.1f}; range=[{low:.0f}, {high:.0f}]; failed to reach the goal before the safety cap: {failed}', 'greedy_none': 'no greedy rollout reached the goal before the safety cap; runs: {failed}', 'greedy_return': 'greedy-route return', 'route_cap': 'goal not reached within {steps} steps', 'cliff_routes_title': 'Greedy routes after one reproducible run (seed=0)', 'prefix_values': 'Values in the prefix tree:', 'empty_prefix': '<empty>'}



from dataclasses import dataclass
import math
from typing import Literal

import matplotlib.pyplot as plt
from matplotlib.patches import Circle, Polygon, Rectangle
import numpy as np

FAST_MODE = True
N_SEEDS = 20 if FAST_MODE else 100
CLIFF_EPISODES = 500 if FAST_MODE else 2_000
PREDICTION_BUDGETS = np.asarray([20, 50, 100, 250, 500, 1_000, 2_500])
MASTER_SEED = 20260804

PALETTE = {
    "policy": "#0072B2",
    "reward": "#E69F00",
    "environment": "#009E73",
    "danger": "#D55E00",
    "value": "#CC79A7",
    "data": "#56B4E9",
    "reference": "#7F7F7F",
}

print(TEXT["profile"].format(
    mode="FAST/CPU" if FAST_MODE else "FULL/CPU",
    seeds=N_SEEDS,
    episodes=CLIFF_EPISODES,
))

def td_target(
    reward: float,
    next_value: float,
    gamma: float,
    *,
    terminated: bool,
    truncated: bool = False,
) -> float:
    if not all(math.isfinite(x) for x in (reward, next_value, gamma)):
        raise ValueError("reward, next_value, and gamma must be finite")
    if not 0.0 <= gamma <= 1.0:
        raise ValueError("gamma must be in [0, 1]")
    if not isinstance(terminated, (bool, np.bool_)):
        raise TypeError("terminated must be boolean")
    if not isinstance(truncated, (bool, np.bool_)):
        raise TypeError("truncated must be boolean")
    return float(reward) if bool(terminated) else float(reward + gamma * next_value)

examples = {
    TEXT["ordinary_transition"]: td_target(1.0, 2.0, 0.9, terminated=False),
    TEXT["true_terminal"]: td_target(1.0, 2.0, 0.9, terminated=True),
    TEXT["external_truncation"]: td_target(1.0, 2.0, 0.9, terminated=False, truncated=True),
}
for name, value in examples.items():
    print(f"{name:>28}: {value:.3f}")

assert examples[TEXT["ordinary_transition"]] == 2.8
assert examples[TEXT["true_terminal"]] == 1.0
assert examples[TEXT["external_truncation"]] == 2.8

def validate_model(transitions, rewards, terminals):
    p = np.asarray(transitions, dtype=np.float64)
    r = np.asarray(rewards, dtype=np.float64)
    terminal = np.asarray(terminals, dtype=bool)
    if p.ndim != 3 or p.shape != r.shape or p.shape[0] != p.shape[2]:
        raise ValueError("P and R must have shape (n_states, n_actions, n_states)")
    if terminal.shape != (p.shape[0],):
        raise ValueError("terminals has the wrong shape")
    if np.any(p < 0.0) or not np.allclose(p.sum(axis=2), 1.0, atol=1e-12, rtol=0):
        raise ValueError("each P[s,a,:] must be a probability distribution")
    return p, r, terminal


def validate_policy(policy, n_states, n_actions):
    pi = np.asarray(policy, dtype=np.float64)
    if pi.shape != (n_states, n_actions):
        raise ValueError("policy has the wrong shape")
    if np.any(pi < 0.0) or not np.allclose(pi.sum(axis=1), 1.0, atol=1e-12, rtol=0):
        raise ValueError("each policy row must be a probability distribution")
    return pi


def q_from_v(transitions, rewards, values, gamma, terminals):
    p, r, terminal = validate_model(transitions, rewards, terminals)
    v = np.asarray(values, dtype=np.float64)
    if v.shape != (p.shape[0],):
        raise ValueError("values has the wrong shape")
    continuation = v.copy()
    continuation[terminal] = 0.0
    q = np.sum(p * (r + gamma * continuation[None, None, :]), axis=2)
    q[terminal] = 0.0
    return q


def exact_policy_value(transitions, rewards, policy, gamma, terminals):
    p, r, terminal = validate_model(transitions, rewards, terminals)
    pi = validate_policy(policy, p.shape[0], p.shape[1])
    p_pi = np.einsum("sa,san->sn", pi, p)
    r_pi = np.einsum("sa,san,san->s", pi, p, r)
    nonterminal = np.flatnonzero(~terminal)
    values = np.zeros(p.shape[0], dtype=np.float64)
    system = np.eye(nonterminal.size) - gamma * p_pi[np.ix_(nonterminal, nonterminal)]
    values[nonterminal] = np.linalg.solve(system, r_pi[nonterminal])
    return values


def value_iteration_trace(transitions, rewards, gamma, terminals, *, tolerance=1e-10, max_sweeps=100_000):
    p, r, terminal = validate_model(transitions, rewards, terminals)
    values = np.zeros(p.shape[0], dtype=np.float64)
    trace = [values.copy()]
    for sweep in range(1, max_sweeps + 1):
        updated = np.max(q_from_v(p, r, values, gamma, terminal), axis=1)
        updated[terminal] = 0.0
        trace.append(updated.copy())
        if np.max(np.abs(updated - values)) < tolerance:
            return np.stack(trace), sweep
        values = updated
    raise RuntimeError("value iteration did not converge")


def greedy_policy(transitions, rewards, values, gamma, terminals):
    q = q_from_v(transitions, rewards, values, gamma, terminals)
    actions = np.argmax(q, axis=1)
    actions[np.asarray(terminals, dtype=bool)] = 0
    return actions


def build_gridworld_4x4():
    n = 4
    terminals = np.zeros(n * n, dtype=bool)
    terminals[[0, n * n - 1]] = True
    p = np.zeros((n * n, 4, n * n), dtype=np.float64)
    r = np.zeros_like(p)
    for state in range(n * n):
        for action in range(4):
            if terminals[state]:
                p[state, action, state] = 1.0
                continue
            row, col = divmod(state, n)
            if action == 0:
                row = max(row - 1, 0)
            elif action == 1:
                col = min(col + 1, n - 1)
            elif action == 2:
                row = min(row + 1, n - 1)
            else:
                col = max(col - 1, 0)
            nxt = row * n + col
            p[state, action, nxt] = 1.0
            r[state, action, nxt] = -1.0
    return p, r, terminals


def build_frozen_lake_4x4():
    layout = ("SFFF", "FHFH", "FFFH", "HFFG")
    n = 4
    holes = {row * n + col for row in range(n) for col in range(n) if layout[row][col] == "H"}
    goal = n * n - 1
    terminals = np.zeros(n * n, dtype=bool)
    terminals[list(holes | {goal})] = True
    p = np.zeros((n * n, 4, n * n), dtype=np.float64)
    r = np.zeros_like(p)

    def move(state, action):
        row, col = divmod(state, n)
        if action == 0:
            col = max(col - 1, 0)
        elif action == 1:
            row = min(row + 1, n - 1)
        elif action == 2:
            col = min(col + 1, n - 1)
        else:
            row = max(row - 1, 0)
        return row * n + col

    for state in range(n * n):
        for action in range(4):
            if terminals[state]:
                p[state, action, state] = 1.0
                continue
            for actual in ((action - 1) % 4, action, (action + 1) % 4):
                nxt = move(state, actual)
                p[state, action, nxt] += 1.0 / 3.0
                if nxt == goal:
                    r[state, action, nxt] = 1.0
    return p, r, terminals

grid_p, grid_r, grid_terminal = build_gridworld_4x4()
uniform_policy = np.full((16, 4), 0.25)
grid_values = exact_policy_value(grid_p, grid_r, uniform_policy, gamma=1.0, terminals=grid_terminal)
print(TEXT["gridworld_values"])
print(np.array2string(grid_values.reshape(4, 4), precision=0, suppress_small=True))

expected_grid = np.asarray([
    0, -14, -20, -22,
    -14, -18, -20, -20,
    -20, -20, -18, -14,
    -22, -20, -14, 0,
], dtype=np.float64)
np.testing.assert_allclose(grid_values, expected_grid, atol=2e-12, rtol=0)

lake_p, lake_r, lake_terminal = build_frozen_lake_4x4()
lake_trace, lake_sweeps = value_iteration_trace(lake_p, lake_r, gamma=0.9, terminals=lake_terminal)
lake_values = lake_trace[-1]
lake_policy = greedy_policy(lake_p, lake_r, lake_values, gamma=0.9, terminals=lake_terminal)
print(f"\n{TEXT['lake_sweeps']}: {lake_sweeps}")
print(TEXT["lake_values"])
print(np.array2string(lake_values.reshape(4, 4), precision=4, suppress_small=True))
print(TEXT["lake_policy"])
print(lake_policy.reshape(4, 4))

# Iteration counts and cell values depend on gamma and stopping tolerance.
assert lake_sweeps >= 1 and np.all(np.isfinite(lake_values))
assert np.all(lake_values[lake_terminal] == 0.0)

layout = ("SFFF", "FHFH", "FFFH", "HFFG")
indices = (0, min(1, lake_sweeps), min(5, lake_sweeps), lake_sweeps)
titles = TEXT["lake_panel_titles"]

fig, axes = plt.subplots(1, 4, figsize=(14, 4.2))
fig.suptitle(TEXT["lake_plot_title"], fontsize=15)
for ax, index, title in zip(axes, indices, titles, strict=True):
    values = lake_trace[index]
    ax.set_xlim(0, 4)
    ax.set_ylim(4, 0)
    ax.set_aspect("equal")
    ax.set_xticks([])
    ax.set_yticks([])
    ax.set_title(title)
    vmax = max(float(values.max()), 1e-12)
    for row in range(4):
        for col in range(4):
            state = row * 4 + col
            token = layout[row][col]
            if token == "H":
                face = "#F4C7B5"
            elif token == "G":
                face = "#A8E3D1"
            else:
                intensity = values[state] / vmax if vmax > 0 else 0.0
                face = (0.98 - 0.18 * intensity, 0.94 - 0.25 * intensity, 0.98 - 0.08 * intensity)
            ax.add_patch(Rectangle((col, row), 1, 1, facecolor=face, edgecolor="white", linewidth=2))
            if token == "H":
                ax.text(col + 0.5, row + 0.35, "×", ha="center", va="center", fontsize=18, color=PALETTE["danger"])
            elif token == "G":
                ax.text(col + 0.5, row + 0.33, "⚑", ha="center", va="center", fontsize=16, color=PALETTE["environment"])
            elif token == "S":
                ax.add_patch(Circle((col + 0.22, row + 0.22), 0.10, facecolor=PALETTE["policy"], edgecolor="white"))
            ax.text(col + 0.5, row + 0.72, f"{values[state]:.3f}", ha="center", va="center", fontsize=8.5)
plt.tight_layout()
plt.show()

@dataclass
class CorridorEnv:
    n_cells: int = 5

    def __post_init__(self):
        if self.n_cells < 3 or self.n_cells % 2 == 0:
            raise ValueError("n_cells must be an odd integer >= 3")
        self.state = self.n_cells // 2

    def reset(self):
        self.state = self.n_cells // 2
        return self.state

    def step(self, action):
        if action not in (0, 1):
            raise ValueError("action must be 0 or 1")
        self.state += -1 if action == 0 else 1
        terminated = self.state in (0, self.n_cells - 1)
        reward = 1.0 if self.state == self.n_cells - 1 else 0.0
        return self.state, reward, terminated


def exact_corridor_values(n_cells=5):
    values = np.linspace(0.0, 1.0, n_cells)
    values[-1] = 0.0
    return values


def corridor_episode(rng, env):
    states, rewards = [], []
    state = env.reset()
    while True:
        states.append(state)
        action = int(rng.integers(0, 2))
        _, reward, terminated = env.step(action)
        rewards.append(reward)
        if terminated:
            return states, rewards
        state = env.state


def mc_prediction(n_episodes, *, seed, gamma=1.0):
    rng = np.random.default_rng(seed)
    env = CorridorEnv()
    values = np.zeros(env.n_cells)
    counts = np.zeros(env.n_cells, dtype=np.int64)
    for _ in range(n_episodes):
        states, rewards = corridor_episode(rng, env)
        g = 0.0
        returns = np.empty(len(rewards))
        for t in range(len(rewards) - 1, -1, -1):
            g = rewards[t] + gamma * g
            returns[t] = g
        seen = set()
        for state, observed_return in zip(states, returns, strict=True):
            if state in seen:
                continue
            seen.add(state)
            counts[state] += 1
            values[state] += (observed_return - values[state]) / counts[state]
    return values


def td_prediction(n_episodes, *, seed, alpha=0.1, gamma=1.0):
    rng = np.random.default_rng(seed)
    env = CorridorEnv()
    values = np.zeros(env.n_cells)
    for _ in range(n_episodes):
        state = env.reset()
        while True:
            action = int(rng.integers(0, 2))
            nxt, reward, terminated = env.step(action)
            target = td_target(reward, values[nxt], gamma, terminated=terminated)
            values[state] += alpha * (target - values[state])
            state = nxt
            if terminated:
                break
    return values


def n_step_prediction(n_episodes, *, n_steps, seed, alpha=0.1, gamma=1.0):
    rng = np.random.default_rng(seed)
    env = CorridorEnv()
    values = np.zeros(env.n_cells)
    for _ in range(n_episodes):
        states, rewards = corridor_episode(rng, env)
        horizon = len(rewards)
        for t, state in enumerate(states):
            end = min(t + n_steps, horizon)
            target = 0.0
            discount = 1.0
            for reward in rewards[t:end]:
                target += discount * reward
                discount *= gamma
            if end < horizon:
                target += discount * values[states[end]]
            values[state] += alpha * (target - values[state])
    return values

truth = exact_corridor_values()
print(TEXT["corridor_truth"], truth)

def rmse(values, truth):
    return float(np.sqrt(np.mean((values[1:-1] - truth[1:-1]) ** 2)))

methods = {
    TEXT["mc"]: lambda budget, seed: mc_prediction(budget, seed=seed),
    TEXT["td0"]: lambda budget, seed: td_prediction(budget, seed=seed, alpha=0.1),
    TEXT["td3"]: lambda budget, seed: n_step_prediction(budget, n_steps=3, seed=seed, alpha=0.1),
}
seed_sequence = np.random.SeedSequence(MASTER_SEED)
seeds = [int(s.generate_state(1)[0]) for s in seed_sequence.spawn(N_SEEDS)]
curves = {}
for name, estimator in methods.items():
    samples = np.empty((N_SEEDS, len(PREDICTION_BUDGETS)))
    for i, seed in enumerate(seeds):
        for j, budget in enumerate(PREDICTION_BUDGETS):
            samples[i, j] = rmse(estimator(int(budget), seed), truth)
    curves[name] = samples

plt.figure(figsize=(9.2, 5.4))
for name, samples in curves.items():
    median = np.median(samples, axis=0)
    low, high = np.quantile(samples, [0.25, 0.75], axis=0)
    plt.plot(PREDICTION_BUDGETS, median, marker="o", label=name)
    plt.fill_between(PREDICTION_BUDGETS, low, high, alpha=0.16)
plt.xscale("log")
plt.xlabel(TEXT["episodes"])
plt.ylabel(TEXT["rmse"])
plt.title(TEXT["prediction_title"])
plt.legend()
plt.grid(alpha=0.25)
plt.show()

print(TEXT["prediction_last_budget"].format(budget=int(PREDICTION_BUDGETS[-1])))
for name, samples in curves.items():
    print(f"  {name:>24}: median={np.median(samples[:, -1]):.4f}, IQR=[{np.quantile(samples[:, -1], .25):.4f}, {np.quantile(samples[:, -1], .75):.4f}]")

@dataclass
class CliffEnv:
    rows: int = 4
    cols: int = 12

    def __post_init__(self):
        self.start = (self.rows - 1, 0)
        self.goal = (self.rows - 1, self.cols - 1)
        self.cliff = {(self.rows - 1, col) for col in range(1, self.cols - 1)}
        self.position = self.start

    @property
    def n_states(self):
        return self.rows * self.cols

    @property
    def n_actions(self):
        return 4

    def _state(self, position):
        return position[0] * self.cols + position[1]

    def reset(self):
        self.position = self.start
        return self._state(self.position)

    def step(self, action):
        row, col = self.position
        if action == 0:
            row = max(row - 1, 0)
        elif action == 1:
            col = min(col + 1, self.cols - 1)
        elif action == 2:
            row = min(row + 1, self.rows - 1)
        elif action == 3:
            col = max(col - 1, 0)
        else:
            raise ValueError("action must be in {0,1,2,3}")
        candidate = (row, col)
        if candidate in self.cliff:
            self.position = self.start
            return self._state(self.position), -100.0, False
        self.position = candidate
        terminated = candidate == self.goal
        return self._state(candidate), -1.0, terminated


def epsilon_greedy(q, state, epsilon, rng):
    if rng.random() < epsilon:
        return int(rng.integers(q.shape[1]))
    return int(np.argmax(q[state]))


def train_cliff(method: Literal["q_learning", "sarsa", "double_q"], *, seed, n_episodes, alpha=0.5, gamma=1.0, epsilon=0.1, max_steps=1_000):
    rng = np.random.default_rng(seed)
    env = CliffEnv()
    q_a = np.zeros((env.n_states, env.n_actions))
    q_b = np.zeros_like(q_a)
    returns = np.zeros(n_episodes)

    for episode in range(n_episodes):
        state = env.reset()
        total = 0.0
        action = epsilon_greedy(q_a, state, epsilon, rng) if method == "sarsa" else None
        for _ in range(max_steps):
            if method == "q_learning":
                action = epsilon_greedy(q_a, state, epsilon, rng)
            elif method == "double_q":
                action = epsilon_greedy(q_a + q_b, state, epsilon, rng)
            assert action is not None
            nxt, reward, terminated = env.step(action)
            total += reward

            if method == "q_learning":
                target = td_target(reward, np.max(q_a[nxt]), gamma, terminated=terminated)
                q_a[state, action] += alpha * (target - q_a[state, action])
                state = nxt
            elif method == "sarsa":
                if terminated:
                    target = reward
                    q_a[state, action] += alpha * (target - q_a[state, action])
                else:
                    next_action = epsilon_greedy(q_a, nxt, epsilon, rng)
                    target = reward + gamma * q_a[nxt, next_action]
                    q_a[state, action] += alpha * (target - q_a[state, action])
                    state, action = nxt, next_action
            else:
                update_a = bool(rng.integers(0, 2))
                selector, evaluator = (q_a, q_b) if update_a else (q_b, q_a)
                selected = int(np.argmax(selector[nxt]))
                target = td_target(reward, evaluator[nxt, selected], gamma, terminated=terminated)
                selector[state, action] += alpha * (target - selector[state, action])
                state = nxt

            if terminated:
                break
        returns[episode] = total
    return q_a + q_b if method == "double_q" else q_a, returns


def greedy_rollout(q, max_steps=1_000):
    """Run a greedy policy and report whether the goal was reached before the safety cap."""
    env = CliffEnv()
    state = env.reset()
    path = [state]
    total = 0.0
    for _ in range(max_steps):
        state, reward, terminated = env.step(int(np.argmax(q[state])))
        total += reward
        path.append(state)
        if terminated:
            return total, path, True
    return total, path, False


def rolling_mean(values, window=25):
    kernel = np.ones(window) / window
    padded = np.pad(values, (window - 1, 0), mode="edge")
    return np.convolve(padded, kernel, mode="valid")

control_methods = ["q_learning", "sarsa", "double_q"]
control_labels = {
    "q_learning": "Q-learning",
    "sarsa": "SARSA",
    "double_q": "Double Q-learning",
}
control_runs = {}
control_greedy = {}
control_greedy_terminated = {}
for method in control_methods:
    returns = np.empty((N_SEEDS, CLIFF_EPISODES))
    greedy_returns = np.full(N_SEEDS, np.nan)
    terminated_flags = np.zeros(N_SEEDS, dtype=bool)
    for i, seed in enumerate(seeds):
        q, run_returns = train_cliff(method, seed=seed, n_episodes=CLIFF_EPISODES)
        returns[i] = run_returns
        greedy_return, _, reached_goal = greedy_rollout(q)
        terminated_flags[i] = reached_goal
        if reached_goal:
            greedy_returns[i] = greedy_return
    control_runs[method] = returns
    control_greedy[method] = greedy_returns
    control_greedy_terminated[method] = terminated_flags

plt.figure(figsize=(9.5, 5.4))
for method in control_methods:
    smoothed = np.vstack([rolling_mean(row) for row in control_runs[method]])
    median = np.median(smoothed, axis=0)
    low, high = np.quantile(smoothed, [0.25, 0.75], axis=0)
    x = np.arange(1, CLIFF_EPISODES + 1)
    plt.plot(x, median, label=control_labels[method])
    plt.fill_between(x, low, high, alpha=0.14)
plt.xlabel(TEXT["training_episode"])
plt.ylabel(TEXT["rolling_return"])
plt.title(TEXT["cliff_learning_title"])
plt.ylim(-160, -10)
plt.legend()
plt.grid(alpha=0.25)
plt.show()

print(TEXT["cliff_summary"])
for method in control_methods:
    train_tail = control_runs[method][:, -100:].mean(axis=1)
    greedy = control_greedy[method]
    finite = greedy[np.isfinite(greedy)]
    failed = int((~control_greedy_terminated[method]).sum())
    if finite.size:
        greedy_summary = TEXT["greedy_summary"].format(
            median=np.median(finite),
            low=finite.min(),
            high=finite.max(),
            failed=failed,
        )
    else:
        greedy_summary = TEXT["greedy_none"].format(failed=failed)
    print(
        f"  {control_labels[method]:>18}: "
        f"{TEXT['train_tail'].format(value=np.median(train_tail))}; "
        f"{greedy_summary}"
    )

representative_seed = 0
fig, axes = plt.subplots(3, 1, figsize=(11.5, 7.2))
for ax, method in zip(axes, control_methods, strict=True):
    q, _ = train_cliff(method, seed=representative_seed, n_episodes=CLIFF_EPISODES)
    greedy_return, path, reached_goal = greedy_rollout(q)
    ax.set_xlim(0, 12)
    ax.set_ylim(4, 0)
    ax.set_aspect("equal")
    ax.set_xticks([])
    ax.set_yticks([])
    for row in range(4):
        for col in range(12):
            state = row * 12 + col
            if 37 <= state <= 46:
                face = "#F4C7B5"
            elif state == 36:
                face = "#CFEAF8"
            elif state == 47:
                face = "#C7EADD"
            else:
                face = "white"
            ax.add_patch(Rectangle((col, row), 1, 1, facecolor=face, edgecolor="#DDDDDD"))
    xy = np.asarray([(state % 12 + 0.5, state // 12 + 0.5) for state in path])
    ax.plot(xy[:, 0], xy[:, 1], linewidth=2.5, marker="o", markersize=3.5)
    route_note = (
        TEXT["greedy_return"] + f" {greedy_return:.0f}"
        if reached_goal
        else TEXT["route_cap"].format(steps=len(path) - 1)
    )
    ax.set_title(f"{control_labels[method]} — {route_note}")
fig.suptitle(TEXT["cliff_routes_title"], fontsize=15)
plt.tight_layout()
plt.show()

# A tiny deterministic prefix tree. Rewards live on transitions into completed responses.
tree = {
    "": {"because": "because", "guess": "guess"},
    "because": {"check": "because check", "skip": "because skip"},
    "guess": {"check": "guess check", "skip": "guess skip"},
}
terminal_reward = {
    "because check": 1.0,
    "because skip": 0.4,
    "guess check": 0.2,
    "guess skip": -0.2,
}
policy = {
    "": {"because": 0.55, "guess": 0.45},
    "because": {"check": 0.70, "skip": 0.30},
    "guess": {"check": 0.25, "skip": 0.75},
}

def prefix_values(tree, terminal_reward, policy, gamma=1.0):
    v = {state: 0.0 for state in terminal_reward}
    q = {}
    for state in sorted(tree, key=len, reverse=True):
        q[state] = {}
        for action, nxt in tree[state].items():
            immediate = terminal_reward.get(nxt, 0.0)
            continuation = 0.0 if nxt in terminal_reward else v[nxt]
            q[state][action] = immediate + gamma * continuation
        v[state] = sum(policy[state][a] * q[state][a] for a in q[state])
    return v, q

prefix_v, prefix_q = prefix_values(tree, terminal_reward, policy)
print(TEXT["prefix_values"])
for state in ("", "because", "guess"):
    label = TEXT["empty_prefix"] if state == "" else state
    print(f"  {label:>12}: V={prefix_v[state]:.3f}; Q={prefix_q[state]}")

assert all(prefix_v[state] == 0.0 for state in terminal_reward)
for state in tree:
    assert np.isclose(prefix_v[state], sum(policy[state][a] * prefix_q[state][a] for a in tree[state]))