DP, Monte Carlo, and TD: Three Ways to Solve the Bellman Equation

import gymnasium as gym
from gymnasium import spaces
import numpy as np
from collections import defaultdict

# ══════════════════════════════════════════════════════
# Mars Rover Gridworld: a stochastic MDP where the agent must navigate a 5×5
# grid to a science target (+10) while avoiding craters (-10).
# We build the full environment as a Gymnasium Env so that:
#   - DP can read the internal model P[s][a] directly (model-based),
#   - MC and TD can only interact via reset()/step() (model-free).
# ═══════════════════════════════════════════════════════

# World layout. Coordinates are (row, col), 0-indexed from the top-left corner.
# 5x5 board, so 25 discrete states (numbered 0..24)
GRID = 5
# science target: landing here ends the episode with +10
GOAL = (4, 4)
# landing in a crater ends the episode with -10
CRATERS = {(2, 2), (1, 3)}
# probability of slipping to EACH of the two perpendicular tiles
SLIP = 0.1
# discount: a reward k steps away is worth 0.95**k today
GAMMA = 0.95
# action -> (Δrow, Δcol): up, right, down, left
MOVES = {0: (-1, 0), 1: (0, 1), 2: (1, 0), 3: (0, -1)}
# perpendicular actions for slip: e.g. "up" can slip to "left" or "right"
PERP = {0: [3, 1], 1: [0, 2], 2: [1, 3], 3: [2, 0]}

def rc(s: int) -> tuple[int, int]:
    """Given a state index s (0 to GRID*GRID-1), returns (row, col)"""
    # divmod(a, b) -> (a // b, a % b)
    return divmod(s, GRID)

def idx(r: int, c: int) -> int:
    """Given (row, col), returns the flattened state index s"""
    return r * GRID + c

class MarsRoverEnv(gym.Env):
    def __init__(self):
        super().__init__()
        # 25 possible states
        self.observation_space = spaces.Discrete(GRID * GRID)
        # 4 possible moves
        self.action_space = spaces.Discrete(4)
        self.goal = idx(*GOAL)
        self.craters = {idx(*c) for c in CRATERS}
        # episode ends in any of these
        self.terminals = self.craters | {self.goal}
        # full transition model (only DP may read this)
        self.P = self._build_model()

    def _move(self, s: int, a: int) -> int:
        """Take one step in direction a, clamped so the rover can't walk off the grid. Takes in a state index s and an action index a, and returns the next state index."""
        # Find the current row and column of the state
        r, c = rc(s)
        # Find the delta row and column for the action
        dr, dc = MOVES[a]
        # Clamp the new row and column to the grid
        return idx(min(max(r + dr, 0), GRID - 1), min(max(c + dc, 0), GRID - 1))

    def _reward(self, s_next: int) -> float:
        """Reward depends only on which tile you land on. Takes in a next state index s_next, and returns the reward."""
        if s_next == self.goal: return 10.0
        if s_next in self.craters: return -10.0
        # every ordinary step costs 1, which nudges the rover to hurry
        return -1.0

    def _build_model(self):
        """This block builds the probability distribution over where the rover actually ends up after choosing an action, accounting for slip."""
        # This is the "god's-eye" model of the world. DP reads it directly;
        # MC and TD are not allowed to, they must learn from reset()/step() only.
        # P is a dict, with keys as state indices and values as dictionaries of action indices and lists of tuples (probability, next_state, reward, terminated).
        P = defaultdict(dict)
        # Iterate over all states
        for current_state in range(GRID * GRID):
            # For each state s, iterate over all actions a
            for action in range(4):
                # If the state is a terminal state,
                # the rover stays put and earns nothing more.
                if current_state in self.terminals:
                    P[current_state][action] = [(1.0, current_state, 0.0, True)]
                    continue

                # Initialize the outcomes dictionary - maps next_state to its probability
                outcomes: dict[int, float] = {}

                # computes the tile the rover wants to reach after taking the action
                intended = self._move(current_state, action)

                # It lands there with probability 1 - 2*SLIP
                # outcomes.get() -> gives probability of the intended outcome, or 0 if it's not in the dictionary. This accumulates into a dict instead of overwriting.
                # This matters because near walls several of these directions can resolve to the same tile (since _move clamps at the edges), and their probabilities must add up rather than clobber each other.
                outcomes[intended] = outcomes.get(intended, 0) + (1 - 2 * SLIP)

                # Add the probability of the perpendicular outcomes
                # PERP[action] are the two directions perpendicular to the intended one.
                # The rover "slips" to each of those with probability SLIP = 0.1 each.
                for slip_action in PERP[action]:
                    slipped = self._move(current_state, slip_action)
                    outcomes[slipped] = outcomes.get(slipped, 0) + SLIP

                # Build the list of transitions for this (current_state, action) pair
                transitions = []
                for next_state, probability in outcomes.items():
                    reward = self._reward(next_state)
                    terminated = next_state in self.terminals
                    transitions.append((probability, next_state, reward, terminated))
                P[current_state][action] = transitions

        return P

    def reset(self, seed=None, options=None):
        super().reset(seed=seed)
        # every episode begins in the top-left corner
        self.s = idx(0, 0)
        return self.s, {}

    def step(self, action):
        """Simulate a real interaction: sample ONE outcome from P[s][a] according to its probabilities. This is all MC and TD ever see — they never peek at the full distribution, only the single (s', r, terminated) triple that nature dealt."""

        # tr is a list of tuples (probability, next_state, reward, terminated)
        transitions = self.P[self.s][action]
        # i is the index of the chosen outcome
        i = self.np_random.choice(len(transitions), p=[t[0] for t in transitions])
        # _ is the probability of the chosen outcome
        _, s_next, reward, terminated = transitions[i]
        # Update the current state
        self.s = s_next
        # Return the next state, reward, terminated, truncated, and info
        return s_next, reward, terminated, False, {}

# Demonstrate what a policy-less agent looks like: no learning, just random moves.
# The rover has no concept of value yet — every direction is equally likely.
# This is the "before" picture; DP/MC/TD will give it purposeful behavior.
env = MarsRoverEnv()
obs, _ = env.reset(seed=42)
# reset(seed=...) only seeds the env's internal RNG used by step(); the action
# space has its OWN RNG, so we seed it too to make this trace reproducible.
env.action_space.seed(42)
done = False
steps = 0
print("Episode trace (random policy):")
while not done and steps < 20:
    # Uniform random action (no intelligence)
    a = env.action_space.sample()
    obs_next, r, terminated, truncated, _ = env.step(a)
    # Show the transition as  (row,col) [a=action] → (row,col)  and the reward received.
    # rc(obs) and rc(obs_next) are the row and column of the current and next states
    print(f"  ({rc(obs)[0]},{rc(obs)[1]}) [a={a}] → ({rc(obs_next)[0]},{rc(obs_next)[1]})  r={r:.0f}")
    obs = obs_next
    done = terminated or truncated
    steps += 1
if done:
    print(f"  Episode ended at ({rc(obs)[0]},{rc(obs)[1]})")

# P[s][a] is a list of (probability, next_state, reward, terminated).
# Print where each action could land the rover, and with what probability.
def show(state, action, label):
    print(f"{label}: from {rc(state)} taking a={action}")
    for p, s_next, r, terminated in env.P[state][action]:
        print(f"    p={p:.1f} → {rc(s_next)}  r={r:.0f}")

# Open cell: the three outcomes are three distinct tiles -> a clean 0.8 / 0.1 / 0.1.
show(idx(1, 1), 0, "Open space")     # 'up' from (1,1)
# Wall cell: 'up' is clamped back onto the same tile, so probabilities MERGE.
show(idx(0, 0), 0, "Corner (wall)")  # 'up' from (0,0)

def value_iteration(env, gamma=GAMMA, theta=1e-6):
    """Bellman optimality backup: V(s) ← max_a Σ p(s'|s,a)[r + γ·V(s')]
    Sweeps all states repeatedly until values stop changing (< theta)."""
    # start every state's value at 0 (arbitrary guess)
    V = np.zeros(env.observation_space.n)
    while True:
        # Track largest change this sweep for convergence check
        delta = 0
        for s in range(env.observation_space.n):
            v_old = V[s]
            q_values = []
            for a in range(env.action_space.n):
                # Q(s,a) = Σ p · [r + γ·V(s')·(1-done)]
                # The (1-done) factor ensures terminal states contribute 0 future value,
                # because once the episode ends there's nothing more to collect.
                q = 0
                for probability, s_next, r, done in env.P[s][a]:
                    q += probability * (r + gamma * V[s_next] * (1 - done))
                # Add the Q-value for this action to the list
                q_values.append(q)

            # Bellman optimality: the value of a state under the BEST policy is the max over actions
            V[s] = max(q_values)
            # Track the change for convergence check
            delta = max(delta, abs(v_old - V[s]))
        # Convergence: when the biggest value change in a full sweep is negligible,
        # we've found V* (guaranteed by the contraction mapping theorem).
        if delta < theta:
            break

    # Extract the greedy policy: π*(s) = argmax_a Q(s,a).
    # Given the optimal values, the best action is the one with the highest expected return.
    policy = np.zeros(env.observation_space.n, dtype=int)
    for s in range(env.observation_space.n):
        q_values = []
        for a in range(env.action_space.n):
            q = 0.0
            for p, s_next, r, done in env.P[s][a]:
                q += p * (r + gamma * V[s_next] * (1 - done))
            q_values.append(q)
        policy[s] = np.argmax(q_values)
    return V, policy

env = MarsRoverEnv()
V_dp, policy = value_iteration(env)

# Display V* as a 5×5 grid — higher values mean "closer to the goal in expectation"
print("V* (5x5 grid):")
print(V_dp.reshape(5, 5).round(2))

# Show the optimal policy as directional arrows
# This is what a perfectly rational rover would do at each cell
print("\nGreedy policy (^>v<):")
arrows = {0: "^", 1: ">", 2: "v", 3: "<"}
grid = []
for s in range(25):
    grid.append("." if s in env.terminals else arrows[policy[s]])
print(np.array(grid).reshape(5, 5))

# Experiment: what happens if we remove stochasticity entirely?
# With SLIP=0 the rover always goes exactly where it intends — no random drift
# into craters. This isolates how much of the danger comes from slip vs. proximity.
SLIP = 0.0
env_det = MarsRoverEnv()
V_det, _ = value_iteration(env_det)
SLIP = 0.1  # restore stochastic terrain for subsequent experiments
print("V* with SLIP=0 (deterministic):")
# Compare this to the stochastic grid: crater-adjacent cells should now be positive
# because the rover can simply avoid stepping into them.
print(V_det.reshape(5, 5).round(2))

def mc_prediction(env, policy, episodes=20000, gamma=GAMMA, alpha=0.02, max_steps=100):
    """Monte Carlo prediction: estimate V^π by playing full episodes and averaging
    the actual discounted returns observed from each state. No model needed —
    only the ability to run reset()/step()."""
    V = np.zeros(env.observation_space.n)
    for _ in range(episodes):
        # Phase 1: Generate a complete episode trajectory under the given policy.
        # We store (state, reward) pairs to compute returns afterward.
        episode = []
        s, _ = env.reset()
        for _ in range(max_steps):
            # follow the fixed policy we are evaluating
            a = policy[s]
            s_next, reward, terminated, truncated, _ = env.step(a)
            episode.append((s, reward))
            s = s_next
            if terminated or truncated:
                break
        # Phase 2: Compute returns backward. Starting from the final step,
        # accumulate G_t = r_{t+1} + γ·G_{t+1}. This is the TRUE discounted sum
        # of rewards from that state onward — unbiased but high-variance.
        G = 0.0
        for s, reward in reversed(episode):
            # G_t = r + γ·G_{t+1} (the recursive return)
            G = reward + gamma * G
            # Incremental mean update: nudge V(s) toward this episode's return G.
            # α controls the step size — too large = noisy, too small = slow learning.
            # MC update: V(s) ← V(s) + α·[G_t - V(s)]
            V[s] = V[s] + alpha * (G - V[s])
    return V

env = MarsRoverEnv()
_, policy = value_iteration(env)
# Show convergence progress: MC gets noisier estimates with fewer episodes
for n in [1000, 5000, 20000]:
    V_mc = mc_prediction(env, policy, episodes=n)
    print(f"MC after {n:,} episodes:  V(start)={V_mc[0]:.2f}")
print(f"\nMC final grid (20k episodes):")
print(V_mc.reshape(5, 5).round(2))

def td_prediction(env, policy, episodes=20000, gamma=GAMMA, alpha=0.05, max_steps=100):
    """TD(0) prediction: estimate V^π by bootstrapping — after each single step,
    update V(s) toward (r + γ·V(s')), using our CURRENT estimate of the next state
    as a stand-in for the true future return. No full episodes needed."""
    V = np.zeros(env.observation_space.n)
    for _ in range(episodes):
        s, _ = env.reset()
        for _ in range(max_steps):
            # follow the fixed policy we are evaluating
            a = policy[s]
            s_next, reward, terminated, truncated, _ = env.step(a)
            # Bootstrap: V(s') is our current (imperfect) estimate of future value.
            # At a terminal state, there IS no future, so bootstrap = 0.
            bootstrap = V[s_next] if not terminated else 0.0
            # TD target = r + γ·V(s'): one real reward + discounted estimate of the rest.
            # TD error δ = target - V(s): how much better/worse reality was than predicted.
            td_error = reward + gamma * bootstrap - V[s]
            # TD(0) update: V(s) ← V(s) + α·δ — nudge value by a fraction of the surprise.
            # Unlike MC, this fires EVERY step, so information propagates much faster.
            V[s] = V[s] + alpha * td_error
            s = s_next
            if terminated or truncated:
                break
    return V

env = MarsRoverEnv()
_, policy = value_iteration(env)
# Show convergence progress: TD updates at every step, so it often converges
# faster than MC even with the same number of episodes.
for n in [1000, 5000, 20000]:
    V_td = td_prediction(env, policy, episodes=n)
    print(f"TD(0) after {n:,} episodes:  V(start)={V_td[0]:.2f}")
print(f"\nTD(0) final grid (20k episodes):")
print(V_td.reshape(5, 5).round(2))