TD Prediction — Feedback & Solution

This is the worked solution to the Lesson 2 exercise, followed by feedback on the submitted attempt. As in Lesson 1, the mistakes are the valuable part: the TD update Owen wrote was already correct — what went wrong was how the run was measured, and that turns out to be exactly where the lesson lives.

The solution

Download lesson_2_td_prediction_solution.py

"""Lesson 2 -- TD(0) prediction on Asteroid Hop: corrected solution.

This is the cleaned-up version of the submitted r1/lesson_2_td_prediction.py.
The TD update itself was already correct; the rewrite fixes how the run is
*measured* and *reported*. See the feedback write-up (web: /lessons/2/solution)
for the full post-mortem. The substantive changes:

  * The error report now answers the assignment's actual question -- the single
    worst ABSOLUTE error and which state it is at -- instead of a per-state
    table of signed means.
  * It separates the two things a constant-alpha estimate gets wrong: the BIAS
    (mean signed error, which is small) and the NOISE FLOOR (the spread the
    estimate keeps jittering across, which never shrinks). The submission
    averaged *signed* errors, so the + and - cancelled and the headline number
    looked near-zero -- hiding the very phenomenon P1 is about.
  * It loops over EPISODES (reset to state 0 each game) rather than a flat run
    of 50,000 steps, and tracks per-state visit counts so P2 can be checked
    against the data.
  * `s - 2` on a JUMP misfire -- the Lesson 1 rule, and what the known-true
    values were computed from. An earlier "corrected" copy used `s - 3`, which
    silently changes the game out from under the ground-truth table.
"""

import random
import statistics
from collections import deque
from enum import Enum, auto

GAMMA = 0.9
ALPHA = 0.1
JUMP_SUCCESS_PROBABILITY = 0.60


class Action(Enum):
    THRUST = auto()
    JUMP = auto()


# --- The game, behind a wall ------------------------------------------------
# MODEL-FREE WALL: the learning loop below may call step() and NOTHING else.
# It never reads JUMP_SUCCESS_PROBABILITY and never consults a transition
# table. step() is the only window onto the dynamics, and each call returns a
# single *sampled* (s_next, reward, done) -- exactly what a real environment
# hands you. That wall is the definition of "model-free."
def step(s, a):  # -> (s_next, reward, done)
    if s == 10:
        raise ValueError("state 10 is terminal; no action should be applied to it")
    if a == Action.THRUST:
        if s == 9:
            return 10, 1.0, True
        return s + 1, 0.0, False
    # JUMP
    if random.random() <= JUMP_SUCCESS_PROBABILITY:
        if s >= 7:  # +3 lands on or past 10 -> win
            return 10, 1.0, True
        return s + 3, 0.0, False
    return max(0, s - 2), 0.0, False  # misfire: slide back 2, floor at 0


# The fixed policy we are EVALUATING (Lesson 1's optimal policy):
# JUMP at states 0-7, THRUST at 8-9.
policy = [Action.JUMP] * 8 + [Action.THRUST] * 2

# Ground truth from Lesson 1 value iteration -- used ONLY for scoring. The
# learner never looks at it; that would be peeking through the wall.
TRUE_V = [
    0.5219919679553231, 0.5482046378586254, 0.5706240497071444,
    0.6186571735660003, 0.6671991520804588, 0.7087165778019958,
    0.7801916910708641, 0.8551379656365434, 0.9, 1.0, 0.0,
]

EPISODES = 50_000
WINDOW = 2_000  # rolling window for the steady-state bias/noise estimate


def td_prediction():
    """Run TD(0) for `EPISODES` games under the fixed policy.

    Returns (V, visits, recent) where `recent[s]` is a rolling window of the
    most recent SIGNED errors (V[s] - TRUE_V[s]) at state s. Signed on purpose:
    the window's MEAN is the bias, its STDEV is the noise floor.
    """
    V = [0.0] * 11
    visits = [0] * 11
    recent = [deque(maxlen=WINDOW) for _ in range(11)]

    for _ in range(EPISODES):
        s = 0
        while True:
            a = policy[s]
            visits[s] += 1
            s_next, r, done = step(s, a)
            target = r if done else r + GAMMA * V[s_next]
            V[s] += ALPHA * (target - V[s])
            recent[s].append(V[s] - TRUE_V[s])
            if done:
                break
            s = s_next
    return V, visits, recent


def report(V, visits, recent):
    print(f"{'state':>5} {'learned':>9} {'true':>7} {'|err|':>7} "
          f"{'bias':>8} {'noise':>7} {'visits':>9}")
    worst_state, worst_err = -1, -1.0
    for s in range(11):
        abs_err = abs(V[s] - TRUE_V[s])
        win = recent[s]
        bias = statistics.fmean(win) if win else 0.0
        noise = statistics.pstdev(win) if len(win) > 1 else 0.0
        print(f"{s:>5} {V[s]:>9.3f} {TRUE_V[s]:>7.3f} {abs_err:>7.3f} "
              f"{bias:>+8.3f} {noise:>7.3f} {visits[s]:>9}")
        if s != 10 and abs_err > worst_err:
            worst_state, worst_err = s, abs_err

    print(f"\nWorst absolute error: {worst_err:.3f} at state {worst_state}.")
    least = min(range(10), key=lambda s: visits[s])
    print(f"Least-visited non-terminal state: {least} ({visits[least]} visits).")
    print(
        "\nP1: 'bias' shrinks toward 0 but 'noise' settles at a floor and stays "
        "there -- with a constant alpha the estimate is an exponential moving "
        "average that never stops jittering. The error does NOT reach zero."
    )
    print(
        "P2: the worst-estimated state is not the least-visited one. Under a "
        "constant alpha the noise floor is set by each state's TARGET SPREAD "
        "(how far apart the r + GAMMA*V[s'] outcomes are), not by visit count -- "
        "so states 6-7, whose JUMP can land on the +1 terminal, are noisiest."
    )


if __name__ == "__main__":
    random.seed(0)
    V, visits, recent = td_prediction()
    report(V, visits, recent)

What the run shows

50,000 episodes, α = 0.1, γ = 0.9 (one seed):

state   learned    true   |err|     bias   noise    visits
    0     0.507   0.522   0.014   +0.002   0.020    112373
    1     0.534   0.548   0.015   +0.003   0.022     29694
    2     0.576   0.571   0.006   -0.001   0.024     14327
    3     0.575   0.619   0.044   +0.003   0.027     74468
    4     0.663   0.667   0.004   +0.003   0.033     36055
    5     0.688   0.709   0.021   -0.001   0.030     17214
    6     0.726   0.780   0.054   +0.007   0.033     44774
    7     0.838   0.855   0.018   +0.002   0.040     21728
    8     0.900   0.900   0.000   -0.000   0.000     10180
    9     1.000   1.000   0.000   -0.000   0.000     36854
   10     0.000   0.000   0.000   +0.000   0.000         0

Worst absolute error: 0.054 at state 6.

Read the bias and noise columns together: the bias (mean signed error) is tiny everywhere — a few thousandths — while the noise (the spread the estimate jitters across) sits at a stubborn floor of 0.02–0.04 and refuses to shrink. The deterministic states 8 and 9 (THRUST, no randomness in the target) are the exception: zero noise, exactly right. Everything important on this page is the gap between those two columns.

Your predictions

P1 — does the error reach zero? Correct.

You predicted: "No — with the step size fixed, the expected error plateaus." That is exactly right, and it's the headline result of the whole exercise. With a constant α, V(s) is an exponentially-weighted moving average of its TD targets: the newest sample always gets weight α = 0.1, so the last ~10 visits dominate forever and older history is permanently discarded. It can never average away the sampling noise the way a true running mean (step 1/n) would.

You also asked what it plateaus at. There's a clean answer: the stationary variance of the estimate is approximately (α / (2 − α)) · Var(target). With α = 0.1 that's ≈ 0.053 · Var(target) — so the noise floor is set by α and by how spread-out the TD targets are at that state. Halve α and you roughly halve the variance (at the cost of slower convergence); let α decay toward 0 and the floor goes to 0 — which is the fix the live session covers next.

P2 — the worst state. Not what happened — and the reason is the interesting bit.

You predicted: "State 0 has the worst estimate, because the error of every other state compounds backward into it." The intuition — bootstrapping propagates error backward — is real, but it points the wrong way here. State 0 is one of the best-estimated states (|err| ≈ 0.014): it's visited most often, and, more importantly, its TD target barely moves — a JUMP from 0 lands on 3 (V ≈ 0.62) or back on 0 (V ≈ 0.52), a spread of only ~0.1.

The worst state is 6 (and 7 close behind). The lesson's own hint says "the worst estimate is the least-visited state" — and that's a genuinely deep fact, but it holds for sample-average estimators, where variance falls as 1/n. Under a constant α it doesn't hold, and this run is the counterexample: the least-visited non-terminal state is 8 (10,180 visits) yet its error is zero. Why? Because under constant α the noise floor doesn't depend on visit count at all (every state forgets its history at the same rate α) — it depends on the spread of the TD target. States 6 and 7 are the noisiest because their JUMP can land directly on the +1 terminal: their target swings between 0.9·V[9] = 0.9 (or a literal +1) and 0.9·V[4..5] ≈ 0.6 — a ~0.3 spread, three times state 0's. Big target spread, big noise floor, worst estimate. Visit count controls how fast a state converges; target spread controls where it settles.

Why your `errors` code looked wrong

Nothing crashed — the code ran and produced numbers. The problem is that the numbers contradicted the printed value table (a learned V[0] of 0.566 sitting 0.044 above the true 0.522, yet a reported "mean error" of 0.011), and that contradiction is real. Four things, roughly in order of how much they matter:

Averaging signed errors cancels the thing you're measuring. You appended V[s] - known[s] and took its mean. But under constant α the estimate oscillates around the truth, so positive and negative deviations cancel and the mean collapses to the bias (≈ 0.01) — which is small and genuinely does shrink. That made the estimates look nearly perfect while the printed V was visibly off. The quantity that does not shrink — the noise floor, the entire point of P1 — is the spread, which your mean threw away. The fix is to report both: mean (bias) and stdev (noise), or just abs(V[s] - known[s]) if you only want one number.
"sqrt error" is mislabeled — and it was your most important number. statistics.stdev is the standard deviation: the square root of the variance of the errors, not the square root of an error. It's the noise floor, the answer to P1 — but it was named like an afterthought and sat in a second column while the misleading mean got top billing.
It never answered the question that was asked. The assignment asked for the single worst absolute error and which state it's at. A per-state table of signed means and standard deviations is more elaborate but skips the one reduction requested: take abs(), then max over states. (The worked solution prints both — the table for insight, the one-line worst-error verdict for the assignment.)
You ran 50,000 steps, not 50,000 episodes. The loop was one continuous walk with a manual reset on reaching state 10, so at ~4–6 steps per game it was only ~10,000 episodes — about a fifth of the intended training. Not a bug, but it leaves more of the early transient mixed into the window. The solution loops episodes explicitly. (Minor: deque(maxlen=1000) expresses the rolling window directly, instead of the manual if len >= 1000: popleft().)

Other feedback on the submission

The model-free wall was respected — good. The learning loop calls step and only step; JUMP_SUCCESS_PROBABILITY lives on the far side of the wall and is never read by the learner. That's the one rule that defines this lesson, and it was followed.
The terminal bootstrap was handled correctly. target = r if done else r + γ·V[s'] — no value bled out of state 10, which is Lesson 1's off-by-one trap met again and avoided this time.
The misfire was s − 2, matching Lesson 1. Keep it there: the ground-truth table you score against was computed from the −2 dynamics, so a stray −3 would quietly invalidate every error number on the page.
The State/Reward/Done NewType aliases are harmless but inert — Python doesn't enforce them at runtime, and the loop uses plain ints throughout. Same note as Lesson 1: names that read like guarantees but aren't checked are worth deleting in favor of a comment.

Final notes on Lesson 2

The big idea landed: you invented TD(0) — value iteration with the model's expectation replaced by the one sample you just lived through, sweeping only the states your games actually visit. The two open threads from the lesson both resolve through this exercise's data: P1's noise floor is real and is set by α, and P2's "worst state" is subtler than "least-visited" once the step size stops decaying. Both point at the same missing piece — a step size that knows when to shrink, and a representation that doesn't need a separate cell per state.

And that second point is the wall this whole tabular approach runs into: because TD(0) keeps one number per state and bootstraps off the known value of the next state, it can never converge for games with a combinatorially huge number of second-to-last states — there are too many cells to ever fill, and you never visit the same one twice. However, as you'll see in the next lesson, you can replace the 11-entry table with a function that generalizes across states it has never seen — Continue to Lesson 3 — Neural Networks by Hand →