Q-Learning on a Gridworld from Absolute Zero

What You'll Master

The Gridworld Setup

Before we touch a hammer or a robot arm, we solve a simpler version of the same problem in a setting where you can hold every state in your head. That's the point of Problem 1.

The world: a 5×4 grid. The agent always starts at (0, 0) — bottom-left. There are two goal cells:

• Goal 1 at (4, 3), 7 steps away (4 right, 3 up).
• Goal 2 at (4, 0), 4 steps away (4 right).

Four actions: left, right, up, down. Moving off the edge keeps you in place but still costs you the per-step reward. Reaching either goal ends the episode.

Three reward scenarios

The reward function is parameterized by three numbers:

1[·] is the indicator function: 1 if true, 0 if false. So every step the agent gets rstep, plus a one-time bonus when it hits a goal.

Same algorithm, three different reward functions, three different optimal policies. The lesson of P1 is exactly this: the optimal policy is not determined by goal magnitude alone — the per-step reward reshapes the value landscape. We'll see why each scenario produces what it does.

Why this is the right warm-up

Tabular Q-learning has every property of "real" RL except one: function approximation. State is enumerable, actions are discrete, you can store every Q-value in a table. This means:

• No neural network. No PyTorch. Just NumPy and Python.
• You can visualize the Q-table after training and see exactly what the algorithm learned.
• The math is exact — tabular Q-learning provably converges to the optimal Q-function.

Once you internalize this version, the leap to neural Q-learning (Problem 3) and policy gradients (Problem 2) is just "replace the table with a network and add stabilization tricks."

NumPy Primer

If you know Python but NumPy is fuzzy, read this. You'll need every concept here when reading the code.

1. ndarrays are just multi-dimensional Python lists, but fast

A NumPy array (ndarray) is a contiguous block of numbers with a shape. It supports vectorized math: operations apply elementwise without writing a for loop.

import numpy as np

a = np.array([1, 2, 3])         # shape (3,)
b = np.zeros((4, 5, 4))           # shape (4, 5, 4), all zeros
c = a * 10                          # [10, 20, 30] — vectorized

2. Shapes and indexing

The Q-table in this homework has shape (height, width, 4) = (4, 5, 4). So q_table is a 4×5×4 cube of numbers. Indexing pulls out a slice:

q_table = np.zeros((4, 5, 4))           # shape (4, 5, 4)

q_table[2]                              # shape (5, 4) — one whole y-row
q_table[2, 3]                           # shape (4,)  — the 4 Q-values at one cell
q_table[2, 3, 1]                        # scalar — one specific Q-value

3. Aggregations: max, argmax, mean, sum

Two close-but-different operations:

q = np.array([0.1, 0.5, 0.3, 0.2])

np.max(q)            # 0.5  — the VALUE
np.argmax(q)         # 1    — the INDEX where the max lives

You almost always want argmax when picking actions: it returns "which action is best" (an integer), not "how good is the best action" (a float).

4. Random number generators

NumPy has two APIs for randomness. The old one (np.random.randint(...)) uses a global state — convenient but bad for reproducibility. The new one creates an explicit Generator:

rng = np.random.default_rng(seed=42)

rng.random()                # float in [0, 1)
rng.integers(4)             # int in {0, 1, 2, 3}, INCLUSIVE 0, EXCLUSIVE 4
rng.integers(low=2, high=5) # int in {2, 3, 4}

The training code creates a seeded Generator (rng = np.random.default_rng(scenario.seed)). When your code uses this rng object, runs are reproducible. If you accidentally use np.random.* directly, you bypass the seed and get different results every time.

5. The dataclass and frozen=True

The starter code uses Python dataclasses heavily:

from dataclasses import dataclass

@dataclass(frozen=True)
class Scenario:
    name: str
    goal_1_reward: float
    goal_2_reward: float
    step_reward: float
    episodes: int = 5000
    alpha: float = 0.2
    gamma: float = 0.98
    # ...

This is a structured record with fields. frozen=True makes instances immutable — you can't reassign scenario.alpha = 0.3 after construction. Useful for hyperparameters: ensures the algorithm code can't accidentally mutate them mid-training.

That's enough NumPy. Onward to the math.

RL Foundations

Before deriving Q-learning, let's nail down the language. Five concepts.

Gt = rt + γ rt+1 + γ2 rt+2 + γ3 rt+3 + ...

Why discount?

Three reasons:

1. Mathematical: without γ < 1, infinite-horizon returns can diverge. With γ < 1, the sum is bounded by rmax / (1 − γ).
2. Practical: agents should prefer rewards now over later, all else equal — matches intuition about money, food, etc.
3. Algorithmic: it makes the Bellman equation contract, which is what makes Q-learning provably converge.

The objective

RL boils down to one thing:

The expectation accounts for randomness in policy and environment. We want a policy that performs well on average, not one that got lucky once.

The Q-function is a tool to find that optimal policy. Let's build it.

The Bellman Equation

Q-learning rests on one beautiful equation. Here's how it falls out of the definitions.

From G_t to Q(s, a)

Define the Q-value as the expected return given an action:

Now the trick. Pull the first term out of the sum:

The bracketed thing in the second term is just the return from time t+1 onward — which is itself a Q-value, evaluated at the next state and the next action. So:

This is the Bellman expectation equation. The Q-value of (s, a) decomposes into the immediate reward plus the discounted Q-value at the next step.

The Bellman optimality equation

If we want the optimal Q-function (corresponding to the optimal policy), we replace "follow π" with "always pick the best action":

The maxa' says: from the next state, take whichever action gives the highest Q-value. This is what makes the equation an optimality condition — it specifies the value of always acting greedily from here forward.

Two ways to solve it

We don't know the env's transition function in general — we just observe samples by acting in it. So we use Q-learning. Next chapter.

Q-Learning, Derived

Q-learning is the Bellman optimality equation, written as an update rule.

The naive idea

If we knew the right-hand side r + γ maxa' Q*(s', a') exactly, we'd just set Q[s, a] = RHS. Done.

We don't. We have a noisy sample of the RHS, computed from one observed transition. The trick: nudge Q[s, a] toward the observed RHS, with a learning rate α controlling how much.

This is the most important equation in this homework. Memorize it.

Anatomy of the update

Read it as three pieces:

• r + γ maxa' Q(s', a') — the TD target. This is what we now think Q(s, a) should be, after seeing one more step of evidence.
• Q(s, a) — what we currently think it is.
• target − Q(s, a) — the TD error. How wrong was our current estimate? Positive: target is higher, push Q up. Negative: push Q down.

So the update reads: "new Q = old Q + α · (how wrong we were)." Or equivalently:

"New Q = blend of old Q (weight 1−α) and the new evidence (weight α)." Smoothly averaging old belief with new data. With α=0.2 (this homework's default), each update moves us 20% of the way toward the latest evidence.

The terminal-state edge case

If s' is terminal (the episode just ended), there's no future:

In code, this is an if/else. If you forget and use the second formula even on terminal steps, you'd add γ maxa' Q(goal, a') — but the Q-values at the goal cell were never observed during a "from-goal" transition because the episode ended there. They're just zeros (or worse, learned junk from spurious updates). Forgetting the if/else is a real bug, not a stylistic concern.

Why does this work?

Theorem (Watkins, 1989): if every (s, a) pair is visited infinitely often and α is decayed appropriately, tabular Q-learning provably converges to Q* — the unique fixed point of the Bellman optimality equation.

Why? At convergence, the expected TD error is zero. That happens exactly when Q satisfies the Bellman optimality equation. So the algorithm is doing stochastic gradient descent on the Bellman residual, and it converges because the Bellman operator is a contraction.

Exploration vs. Exploitation

The Q-learning update is one half of the algorithm. The other half: how do we choose which (s, a) to update? If we always pick the action with the highest current Q-value (greedy), we'll never try alternatives. We might miss a better path.

The dilemma

• Exploitation: pick the action your current Q-table says is best. Maximizes reward given what you know.
• Exploration: pick a different action, even though Q says it's worse. Risks short-term reward but might reveal the alternative is actually better.

You need both. Pure greedy and you'll get stuck in suboptimal habits. Pure random and you'll never act on what you've learned.

ε-greedy: the simplest balance

Set ε high early in training (lots of exploration) and decay it low later (mostly exploit what you've learned). This homework decays linearly from 0.4 down to 0.02 across 5000 episodes:

def epsilon_for_episode(scenario, episode_idx):
    fraction = episode_idx / max(1, scenario.episodes − 1)
    return scenario.epsilon_start + fraction * (scenario.epsilon_end − scenario.epsilon_start)

Why early exploration matters

At episode 0, the Q-table is all zeros. Argmax over zeros returns 0 (LEFT). Without exploration, the agent would pick LEFT every step, hit the left wall, never see a goal, and never learn anything. The Q-table would stay zero forever.

With ε=0.4, in 60% of states we still try-something-greedy, but in 40% we randomize. Over 20 timesteps per episode, that's plenty of stumbling-around-randomly to eventually trip into a goal. The first time it does, the goal reward propagates one cell into the Q-table. Next episode, the previous cell's Q-value is non-zero, so greedy starts working from there. It cascades.

How Learning Propagates

This is the part of Q-learning that feels almost magical when you see it for the first time. The reward only happens at terminal cells — but somehow, after enough episodes, every cell in the path knows it's on a path to reward. How?

Episode 1: first reward seen

Initially, q_table is all zeros. The agent ε-greedy-walks for 20 steps, mostly randomly. By chance, on some episode it stumbles into Goal 1 from cell (4, 2) by going UP.

The transition: (4, 2) → UP → (4, 3), reward = −1 + 10 = 9, done = True.

The TD update with done:

One cell in the table now has a non-zero Q-value. Everything else is still zero.

Episode 2: the value leaks one step backward

Suppose the agent reaches (3, 2) and takes RIGHT. Transition: (3, 2) → (4, 2), reward = −1, not done.

The TD update without done:

Now (3, 2), RIGHT has a non-zero Q. The agent now "knows" that going right from (3, 2) is on a path to reward.

Episode N: the wave reaches the start

Repeat the propagation. Each episode, the boundary of "cells with non-zero Q" expands by roughly one cell toward the start. After enough episodes, every cell on the optimal path has a positive Q-value pointing toward the goal, geometrically decayed by γ per step.

At convergence, for cells along the optimal path:

This is what makes Q-learning work. The reward at the goal is one bit of information per episode, but it propagates back through the table like a ripple, eventually labeling every relevant state with its discounted future return.

Why the speed of propagation matters

One reward per episode = one TD error per episode = one Q-value updated by a meaningful amount per episode. The propagation is "1 cell per episode" in the worst case. With 5000 episodes and a path length of 7, that's plenty.

In a 100×100 grid with sparse rewards, you'd need orders of magnitude more episodes for the wave to reach all relevant cells. This is why neural-network RL on continuous environments often needs millions of steps — the "wavefront" of learned values has a longer journey to make.

The Full Algorithm

That's it. Three nested layers (scenario × episode × step), one update equation, one exploration rule. Tabular Q-learning fits in 30 lines.

Code Tour: gridworld_q_learning.py

Five components in the starter file: action constants, Scenario dataclass, GridWorld env, helper functions, and the training loop.

Action constants

gridworld_q_learning.py:8-12LEFT = 0
RIGHT = 1
DOWN = 2
UP = 3
ACTION_NAMES = {LEFT: "left", RIGHT: "right", DOWN: "down", UP: "up"}

Just integer codes for actions. You'll see q_table[..., RIGHT] in the code, which is the same as q_table[..., 1] but reads better.

The Scenario dataclass

gridworld_q_learning.py:15-28@dataclass(frozen=True)
class Scenario:
    name: str
    goal_1_reward: float
    goal_2_reward: float
    step_reward: float
    expected_outcome: str
    episodes: int = 5000
    horizon: int = 20
    alpha: float = 0.2           # learning rate
    gamma: float = 0.98          # discount factor
    epsilon_start: float = 0.4
    epsilon_end: float = 0.02
    seed: int = 0

Bundle of hyperparameters. Three pre-built scenarios at lines 39-61 differ only in goal_1_reward, goal_2_reward, step_reward, and expected_outcome.

The GridWorld env

gridworld_q_learning.py:64-98class GridWorld:
    def __init__(self, goal_1_reward, goal_2_reward, step_reward):
        self.start = (0, 0)
        self.goal_2 = (4, 0)
        self.goal_1 = (4, 3)
        self.width = 5
        self.height = 4
        self.goal_1_reward = goal_1_reward
        self.goal_2_reward = goal_2_reward
        self.step_reward = step_reward

    def step(self, state, action):
        x, y = state
        if action == LEFT:  x = max(0, x − 1)
        if action == RIGHT: x = min(self.width − 1, x + 1)
        if action == DOWN:  y = max(0, y − 1)
        if action == UP:    y = min(self.height − 1, y + 1)
        next_state = (x, y)

        reward = self.step_reward
        done = False
        if next_state == self.goal_2:
            reward += self.goal_2_reward
            done = True
        if next_state == self.goal_1:
            reward += self.goal_1_reward
            done = True
        return next_state, reward, done

Deterministic transitions. max(0, ...) and min(width-1, ...) implement the "stay in place when moving off-edge" rule. Reaching either goal sets done=True and adds the goal bonus.

Helper functions

gridworld_q_learning.py:107-110def action_values(q_table, state):
    x, y = state
    return q_table[y, x]    # shape (4,) — the Q-values for all 4 actions at this cell

Tiny but useful. Reads "give me the row of Q-values at this state" without making the (x, y) vs (y, x) flip explicit at every call site.

The blanks you fill in

Three blanks:

1. greedy_action(q_table, state) — pick the argmax action.
2. choose_action(q_table, state, epsilon, rng) — ε-greedy version.
3. The body of the inner loop in train_q_learning — the actual Q-learning update.

Next chapter: every line you wrote, decoded.

Your Three Changes, Decoded

The centerpiece. Every line you wrote, with a linenote block per line, explaining what each NumPy/Python operation does and why.

def greedy_action(q_table, state):
    q_values = action_values(q_table, state)
    action = int(np.argmax(q_values))
    return action

def choose_action(q_table, state, epsilon, rng):
    if rng.random() < epsilon:
        action = int(rng.integers(4))
    else:
        action = greedy_action(q_table, state)
    return action

action = rng.integers(4)              # always sample, even if unused
if rng.random() < epsilon: return int(action)
...

action = choose_action(q_table, state, epsilon, rng)
next_state, reward, done = env.step(state, action)
if done:
    target = reward
else:
    target = reward + scenario.gamma * np.max(action_values(q_table, next_state))
x, y = state
q_table[y, x, action] += scenario.alpha * (target − q_table[y, x, action])
state = next_state
if done:
    break

Putting it all together

The full inner loop, distilled:

Six lines of pseudocode. ~10 lines of actual code. That's the entire engine of tabular Q-learning.

The Three Scenarios

Same algorithm, three different reward structures, three different optimal policies. The deliverable for Task 1.3 is to report which goal each scenario reaches and explain why in one sentence.

Scenario 1: r_step = −1, R₁ = 10, R₂ = 5

Compare the two paths:

Goal 1 wins by 2. The Q-table reflects this: at the start, RIGHT and UP have the highest values (both lie on a path toward Goal 1's diagonal).

Observed: trajectory through the diagonal toward Goal 1, total reward = 3. ✓

One-sentence explanation: Going to Goal 1 gives 7×(−1) + 10 = 3, while Goal 2 gives 4×(−1) + 5 = 1, so the agent learns Goal 1 is the better target.

Scenario 2: r_step = −2, R₁ = 10, R₂ = 5

Goal magnitudes unchanged. Step penalty doubled. Compare again:

Goal 2 wins by 1. The 3 extra steps to Goal 1 cost 6 (with the doubled penalty), wiping out the 5-point reward gap.

Observed: trajectory straight right to Goal 2, total reward = −3. ✓

One-sentence explanation: Now Goal 1 gives 7×(−2) + 10 = −4 but Goal 2 gives 4×(−2) + 5 = −3, so doubling the step cost flips which goal is better and the agent switches to Goal 2.

Scenario 3: r_step = +1, R₁ = R₂ = 1

Step reward is now positive. Compare again:

Stalling wins by a lot. The agent learns to walk LEFT into the wall, soaking up +1 per step for the full horizon. Reaching a goal would terminate that profitable stream after just one bonus.

Observed: trajectory of 20 LEFT actions, agent stays at (0, 0) the whole time, total reward = 20. ✓

Q-values at start are all 50.0: that's the infinite-horizon discounted return of "+1 forever," computed exactly as 1 / (1 − γ) = 1 / 0.02 = 50. Q-learning converged to the optimal value of "stay alive forever."

One-sentence explanation: Reaching the closest goal (Goal 2 in 4 steps) gives 4×(+1) + 1 = 5, but stalling against the wall for the full horizon gives 20×(+1) + 0 = 20, so the agent learns to never finish the task because the +1 per step pays more than any goal bonus.

The big-picture lesson

The Q-values tell a story

You also reported the Q-values at the start state. Look at scenario 1:

start_q_values: (1.219, 2.265, 1.219, 2.265)
                LEFT   RIGHT  DOWN   UP

LEFT and DOWN are tied at 1.219 — both move the agent into a wall (LEFT off the left edge, DOWN off the bottom edge), keeping it in place but still costing the step penalty. RIGHT and UP are tied at 2.265 — both move toward Goal 1's "diagonal." The Q-table has correctly figured out the symmetry.

For scenario 2:

start_q_values: (-4.996, -3.057, -4.996, -4.329)
                 LEFT   RIGHT   DOWN    UP

RIGHT is now distinctly better than UP — because the path to Goal 2 (going right four times) is now the optimal path. UP would head toward the now-suboptimal Goal 1. The Q-table has figured out that the policy should change.

This is why looking at the actual numbers after training is so satisfying with tabular Q-learning. You can directly inspect the values it learned and verify they encode the right preferences.

Cheat Sheet & Self-Quiz

Equations to memorize

Hyperparameter reference

Self-quiz — if you can answer these without re-reading, you're ready

1. Why does Q-learning need to track actions in the table, not just states?
2. What does the discount factor γ do, and what would happen with γ=1 in scenario 3?
3. Why is the Q-table initialized to zeros and not random values?
4. What is the difference between np.max and np.argmax? When do you want each?
5. Why does the inner loop need if done: break? Where does the break go?
6. If you forgot the if done in the TD target and always used the bootstrap formula, what would happen?
7. Why is ε high at the start of training and low at the end?
8. What does the equation Q ← Q + α (target − Q) mean intuitively?
9. Suppose r_step = 0 in scenario 3. What policy would the agent learn? Why?
10. If the agent never explored (ε = 0 throughout), what would happen?
11. What does rng.integers(4) return, and why use rng instead of np.random?
12. Why does the Q-table use shape (height, width, 4) and not (width, height, 4)?

Running it

cd /Users/ozyphus/Documents/GitHub/cs224r/homework/hw2/hw2
python gridworld_q_learning.py

Should print three blocks (one per scenario), each with expected/observed outcome, trajectory, actions, total reward, and start Q-values. If observed_outcome matches expected_outcome for all three, your implementation is correct.

You can now teach this

Three big ideas, in order of importance:

1. The Bellman equation. Q(s, a) decomposes into immediate reward plus discounted future Q. This recursion is the entire engine of TD learning — tabular, neural, on-policy, off-policy, all of them.
2. The Q-learning update. Q ← Q + α (target − Q). Same equation, scaled up to neural nets, becomes DQN, SAC, and every value-based method in modern RL.
3. Reward shaping is hard. The optimal policy is sensitive to small changes in the reward function. Scenario 3 is the canonical example of "the agent did exactly what you said, not what you wanted."

If a friend asks: "How does Q-learning work?" — you say: "It's stochastic gradient descent on the Bellman residual. You don't know the right Q-values, but every transition gives you a noisy sample of the right-hand side. Move toward that sample by α each time. Repeat enough and the Q-table converges to the optimal Q-function. The two tricks: discount future rewards by γ to make sums finite, and use ε-greedy to ensure every state-action gets sampled."

You can teach this. On to Problem 2.