MC vs. TD(0)

# Parameters (injected by papermill)
output_dir = "."
images_dir = "./images"
videos_dir = "./videos"
audio_dir = "./audio"
text_dir = "./text"

import os

os.makedirs(images_dir, exist_ok=True)

import numpy as np
import matplotlib.pyplot as plt

states = ["A", "B", "C", "D", "E"]
state_to_index = {s: i for i, s in enumerate(states)}
n_states = len(states)


def generate_episode(start="C"):
    state = state_to_index[start]
    episode = []
    while True:
        action = np.random.choice([-1, 1])  # left or right
        next_state = state + action
        if next_state < 0:
            return episode + [(state, 0)]  # Left terminal, reward 0
        elif next_state >= n_states:
            return episode + [(state, 1)]  # Right terminal, reward 1
        else:
            episode.append((state, 0))
            state = next_state


def mc_prediction(episodes=1000, alpha=0.1):
    V = np.full(n_states, 0.5)
    # V = np.zeros(n_states)
    for _ in range(episodes):
        episode = generate_episode()
        G = episode[-1][1]  # Only final reward matters
        for s, _ in episode:
            V[s] += alpha * (G - V[s])
    return V


def td0_prediction(episodes=10000, alpha=0.01):
    # Episode tuples follow the (S_t, R_{t+1}) convention:
    # episode[t][1] is the reward for the OUTGOING transition from episode[t][0].
    V = np.full(n_states, 0.5)
    for _ in range(episodes):
        episode = generate_episode()
        for t in range(len(episode)):
            s, r = episode[t]
            v_next = V[episode[t + 1][0]] if t + 1 < len(episode) else 0.0
            V[s] += alpha * (r + v_next - V[s])
    return V


# Run both methods (Sutton & Barto Example 6.2 uses constant-alpha MC)
alpha_mc = 0.1
alpha_td = 0.01
V_mc = mc_prediction(alpha=alpha_mc)
V_td = td0_prediction(alpha=alpha_td)

# Ground truth (analytical solution): [1/6, 2/6, 3/6, 4/6, 5/6]
true_V = np.array([1, 2, 3, 4, 5]) / 6

# Plot results
plt.plot(true_V, label="True V", linestyle="--", marker="o")
plt.plot(V_mc, label=f"MC (α={alpha_mc})", marker="x")
plt.plot(V_td, label=f"TD(0) (α={alpha_td})", marker="s")
plt.xticks(ticks=range(n_states), labels=states)
plt.ylabel("Estimated Value")
plt.title("Monte Carlo vs TD(0) on Random Walk")
plt.legend()
plt.grid(True)
plt.savefig(
    f"{images_dir}/monte_carlo_vs_td0_on_random_walk.png", dpi=150, bbox_inches="tight"
)
plt.show()

def mc_rms_history(episodes, alpha, true_V):
    V = np.full(n_states, 0.5)
    history = np.empty(episodes)
    for k in range(episodes):
        episode = generate_episode()
        G = episode[-1][1]
        for s, _ in episode:
            V[s] += alpha * (G - V[s])
        history[k] = np.sqrt(np.mean((V - true_V) ** 2))
    return history


def td0_rms_history(episodes, alpha, true_V):
    V = np.full(n_states, 0.5)
    history = np.empty(episodes)
    for k in range(episodes):
        episode = generate_episode()
        for t in range(len(episode)):
            s, r = episode[t]
            v_next = V[episode[t + 1][0]] if t + 1 < len(episode) else 0.0
            V[s] += alpha * (r + v_next - V[s])
        history[k] = np.sqrt(np.mean((V - true_V) ** 2))
    return history


def averaged_rms(method_fn, episodes, alpha, runs, true_V):
    return np.mean([method_fn(episodes, alpha, true_V) for _ in range(runs)], axis=0)


n_episodes = 100
n_runs = 100
mc_alphas = [0.01, 0.02, 0.03, 0.04]
td_alphas = [0.05, 0.10, 0.15]

plt.figure(figsize=(8, 5))
for a in mc_alphas:
    rms = averaged_rms(mc_rms_history, n_episodes, a, n_runs, true_V)
    plt.plot(rms, linestyle="--", label=f"MC \u03b1={a}")
for a in td_alphas:
    rms = averaged_rms(td0_rms_history, n_episodes, a, n_runs, true_V)
    plt.plot(rms, label=f"TD(0) \u03b1={a}")
plt.xlabel("Episodes (walks)")
plt.ylabel(f"RMS error (averaged over {n_states} states, {n_runs} runs)")
plt.title("Sutton & Barto Example 6.2 \u2014 RMS error vs. episodes")
plt.legend()
plt.grid(True)
plt.savefig(f"{images_dir}/rms_error_vs_episodes.png", dpi=150, bbox_inches="tight")
plt.show()