CS5446 8. Value function approximation

import numpy as np

H, W = 5, 5
gamma = 0.9

ACTIONS = {
    0: (-1, 0),  # up
    1: (1, 0),  # down
    2: (0, -1),  # left
    3: (0, 1),  # right
    4: (0, 0),  # stay
}
pi = 1.0 / len(ACTIONS)  # 0.2

# 1-based coordinates
forbidden = {(2, 2), (2, 3), (3, 3), (4, 2), (4, 4), (5, 2)}
target = (4, 3)

r_forbidden = -1.0
r_boundary = -1.0
r_target = 1.0


def idx(r, c):  # r,c are 1-based
    return (r - 1) * W + (c - 1)


def step(r, c, a):
    # direction row, column
    dr, dc = ACTIONS[a]
    # next row, column
    nr, nc = r + dr, c + dc

    # boundary -> stay (and get boundary reward)
    if nr < 1 or nr > H or nc < 1 or nc > W:
        return (r, c), r_boundary

    ns = (nr, nc)

    # reward by landing cell
    if ns == target:
        return ns, r_target
    if ns in forbidden:
        return ns, r_forbidden

    return ns, 0.0  # normal cell reward


S = H * W
# state transition probability matrix
P_pi = np.zeros((S, S), dtype=float)
# state reward vector
r_pi = np.zeros(S, dtype=float)

for r in range(1, H + 1):
    for c in range(1, W + 1):
        s = idx(r, c)
        for a in ACTIONS:
            (nr, nc), reward = step(r, c, a)
            sp = idx(nr, nc)
            P_pi[s, sp] += pi
            r_pi[s] += pi * reward

I = np.eye(S)
A = I - gamma * P_pi
# Solve (I - gamma P) v = r
v = np.linalg.solve(A, r_pi)
# print as 5x5
V = v.reshape(H, W)
np.set_printoptions(precision=3, suppress=True)
print("\n=== V (5x5) ===")
print(V)

import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D  # noqa: F401

# =========================
# 1) Gridworld definition (keep yours)
# =========================
H, W = 5, 5
gamma = 0.9

ACTIONS = {
    0: (-1, 0),  # up
    1: (1, 0),  # down
    2: (0, -1),  # left
    3: (0, 1),  # right
    4: (0, 0),  # stay
}
A = len(ACTIONS)

r_forbidden = -1.0
r_boundary = -1.0
r_target = 1.0

FORBIDDEN = {(2, 2), (2, 3), (3, 3), (4, 2), (4, 4), (5, 2)}
TARGET = (4, 3)


def step(state_rc, action):
    r, c = state_rc
    dr, dc = ACTIONS[action]
    nr, nc = r + dr, c + dc

    if nr < 1 or nr > H or nc < 1 or nc > W:
        return (r, c), r_boundary

    ns = (nr, nc)
    if ns == TARGET:
        return ns, r_target
    if ns in FORBIDDEN:
        return ns, r_forbidden
    return ns, 0.0


# =========================================
# 2) Polynomial features phi(s) up to degree K
# =========================================
def make_poly_phi(degree: int):
    """
    Build polynomial features in (rn, cn) up to total degree=degree.
    Example:
      degree=1 -> [1, rn, cn]
      degree=2 -> [1, rn, cn, rn^2, rn*cn, cn^2]
      degree=3 -> add rn^3, rn^2*cn, rn*cn^2, cn^3
    """
    assert degree >= 0

    def phi(state_rc):
        r, c = state_rc
        rn = (r - 1) / (H - 1)  # [0,1]
        cn = (c - 1) / (W - 1)  # [0,1]

        feats = [1.0]
        # total degree d: rn^i * cn^j with i+j=d
        for d in range(1, degree + 1):
            for i in range(d, -1, -1):
                j = d - i
                feats.append((rn**i) * (cn**j))
        return np.asarray(feats, dtype=np.float64)

    # feature dimension
    dim = 1 + sum((d + 1) for d in range(1, degree + 1))  # 1 + (2+3+...+degree+1)
    return phi, dim


def v_hat(state_rc, w, phi_fn):
    return float(np.dot(w, phi_fn(state_rc)))


# =========================================
# 3) TD(0) trainer
# =========================================
rng = np.random.default_rng(0)

num_episodes = 500
steps_per_episode = 500
alpha0 = 0.1


def alpha(t):
    return alpha0 / np.sqrt(1.0 + t / 2000.0)


def train_td0(phi_fn, d, seed=0):
    rng_local = np.random.default_rng(seed)
    w = np.zeros(d, dtype=np.float64)

    td_err_log = []
    t_global = 0

    for ep in range(num_episodes):
        s = (int(rng_local.integers(1, H + 1)), int(rng_local.integers(1, W + 1)))

        for _ in range(steps_per_episode):
            a = int(rng_local.integers(0, A))
            s_next, r = step(s, a)

            v_s = v_hat(s, w, phi_fn)
            v_sn = v_hat(s_next, w, phi_fn)
            delta = r + gamma * v_sn - v_s

            w += alpha(t_global) * delta * phi_fn(s)

            if t_global % 1000 == 0:
                td_err_log.append(abs(delta))

            s = s_next
            t_global += 1

    return w, np.asarray(td_err_log, dtype=np.float64)


def compute_V(w, phi_fn):
    V = np.zeros((H, W), dtype=np.float64)
    for r in range(1, H + 1):
        for c in range(1, W + 1):
            V[r - 1, c - 1] = v_hat((r, c), w, phi_fn)
    return V


# =========================================
# 4) Train for degree 1/2/3, compute V
# =========================================
degrees = [1, 2, 3]
results = {}

for deg in degrees:
    phi_fn, d = make_poly_phi(deg)
    w, td_err = train_td0(phi_fn, d, seed=0)
    V = compute_V(w, phi_fn)
    results[deg] = {"phi": phi_fn, "d": d, "w": w, "V": V, "td_err": td_err}

    print(f"\n=== degree={deg} ===")
    print("d =", d)
    print("w =", w)
    print("Estimated V =\n", np.round(V, 3))


# =========================================
# 5) 3D visualization with gradient color + colorbar
# =========================================
rows = np.arange(1, H + 1)
cols = np.arange(1, W + 1)
C, R = np.meshgrid(cols, rows)  # X=col, Y=row

# optional: use same color scale across degrees for easy comparison
all_vals = np.concatenate([results[d]["V"].ravel() for d in degrees])
vmin, vmax = float(all_vals.min()), float(all_vals.max())

for deg in degrees:
    Z = results[deg]["V"]

    fig = plt.figure(figsize=(9, 6))
    ax = fig.add_subplot(111, projection="3d")

    surf = ax.plot_surface(
        C,
        R,
        Z,
        rstride=1,
        cstride=1,
        linewidth=0,
        antialiased=True,
        cmap="viridis",  # 渐变色
        vmin=vmin,
        vmax=vmax,  # 统一色阶（想各自自适应就删掉这行）
    )

    # colorbar（显示渐变对应的数值）
    fig.colorbar(surf, ax=ax, shrink=0.7, pad=0.1, label="v_hat")

    ax.set_xlabel("col")
    ax.set_ylabel("row")
    ax.set_zlabel("v_hat(row,col)")
    ax.set_title(f"TD(0) + Poly Features (degree={deg})")
    ax.set_box_aspect((1.0, 1.0, 0.5))
    ax.invert_yaxis()
    ax.view_init(elev=30, azim=225)

    plt.savefig(f"value_surface_deg{deg}.png", dpi=300, bbox_inches="tight")
    plt.show()