HW2 — Online RL

Chapter 1: The Gridworld

Before we touch a robot arm or a neural network, we solve a simpler version of the same problem in a world you can hold entirely in your head. That is the point of Problem 1 — and it is the right warm-up because tabular Q-learning has every property of "real" RL except one: function approximation.

The world

A grid 5 wide, 4 tall. The agent always starts at (0, 0) — bottom-left. There are two goal cells:

Four actions: left, right, up, down, integer-coded {0, 1, 2, 3}. Moving off the edge keeps you in place but still costs the per-step reward — walls are not free. Reaching either goal ends the episode.

The reward function — three numbers

The reward is parameterized by three scalars. Every step the agent gets r_step; landing on a goal adds a one-time bonus:

where 𝟙[·] is the indicator — 1 if the condition holds, 0 otherwise. Three scenarios use the same algorithm and produce three different optimal policies:

Why tabular first

Because the state is enumerable and actions are discrete, you can store every value in a table — a cube of shape (height, width, 4). No neural network. No PyTorch. Just numpy. That means three superpowers:

• You can print the whole Q-table after training and see exactly what it learned.
• The math is exact: tabular Q-learning provably converges to the optimal Q-function.
• The leap to neural RL (Problems 2 and 3) is then just "replace the table with a network, and add stabilization tricks."

Chapter 2: RL Foundations

Before deriving anything, we nail the language. Five words. Each is a precise object; sloppiness here costs you the whole homework.

The return, and why we discount

The return is the sum of all future rewards, each future one shrunk by a factor of the discount γ:

Here γ = 0.98. Rewards far in the future are worth less than rewards now. Three reasons this matters:

• Mathematical. Without γ < 1, an infinite-horizon return can diverge to infinity. With γ < 1 the sum is bounded by r_max / (1 − γ).
• Practical. All else equal, agents (and people) should prefer reward now over reward later.
• Algorithmic. Discounting makes the Bellman equation a contraction — which is exactly what makes Q-learning provably converge.

The objective

All of RL boils down to one line: find the policy π that maximizes the expected return.

The expectation accounts for randomness in the policy and the environment. We want a policy that performs well on average — not one that got lucky once. The Q-function is the tool that lets us find it. That is the next chapter.

Chapter 3: The Bellman Equation

All of Q-learning rests on one recursion, and it falls out of the definitions with a single algebraic trick. Watch it happen.

From return to Q-value

Define the Q-value as the expected return if you take action a in state s, then follow π forever:

The trick. Pull the first reward out of the sum, and factor a γ out of everything that remains:

But the bracketed thing is just the return from time t+1 — which is itself a Q-value, at the next state and next action. So the Q-value refers to itself one step ahead:

This is the Bellman expectation equation. The value of (s, a) decomposes into the immediate reward plus the discounted value of where you land. That recursion is the entire engine of temporal-difference learning — tabular, neural, on-policy, off-policy, all of them.

The optimality version

If we want the optimal Q-function — the value of always acting greedily from here on — we replace "follow π" with "always take the best next action":

The maxa′ is what makes this an optimality condition: from the next state, take whichever action has the highest Q-value. This is a fixed-point equation — Q* is its unique solution. Solve it, even approximately, and you have solved the RL problem, because the optimal policy is just π*(s) = argmaxa Q*(s, a).

Chapter 4: Q-Learning, Derived

Q-learning is the Bellman optimality equation, rewritten as an update rule. Here is the whole derivation.

The naive idea, and why it fails

If we knew the right-hand side r + γ maxa′ Q*(s′, a′) exactly, we would just set Q[s, a] equal to it. Done. But we do not know it — we only have one noisy sample of it, computed from a single observed transition. So instead of overwriting, we nudge: move Q a fraction α of the way toward the observed target.

This is the most important equation in the homework. Read it as three pieces:

So the update reads: "new Q = old Q + α · (how wrong we were)." Equivalently, it is a smoothed running average:

With α = 0.2 (this homework's default), each update moves us 20% of the way toward the latest evidence, keeping 80% of the old belief. A gentle running mean over one-sample targets.

The terminal edge case — a real bug, not a nicety

If s′ is terminal (the episode just ended), there is no future. The target is just the reward:

Why it converges

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, which happens exactly when Q satisfies Bellman. The algorithm is stochastic gradient descent on the Bellman residual, and it converges because the Bellman operator is a contraction.

A worked propagation, by hand

Reward only ever happens at goals — yet after enough episodes every cell on the path "knows" it leads to reward. Here is the mechanism, step by step, with α = 0.2, γ = 0.98, r_step = −1, R₁ = 10.

Episode 1. By luck, the agent reaches (4, 2) and goes UP into Goal 1. Transition: reward = −1 + 10 = 9, done = True. Terminal target = 9:

One cell is now non-zero. Everything else is still zero.

Episode 2. The agent reaches (3, 2) and goes RIGHT to (4, 2): reward = −1, not done. Bootstrap target uses the 1.8 we just learned:

Now (3, 2) knows that going right is on a path to reward. Episode after episode, this boundary of "cells with non-zero Q" expands one cell toward the start — a bright wave spreading like dye through water. After enough episodes every cell on the optimal path points at the goal, its value geometrically decayed by γ per step.

Chapter 5: Explore vs Exploit

The Q-update is one half of the algorithm. The other half is: which (s, a) do we update? If we always take the action our current Q says is best, we will never try alternatives — and we might miss a better path forever.

The dilemma

• Exploitation: pick the action the current Q-table calls best. Maximizes reward given what you know.
• Exploration: pick a different action, even one Q calls worse. Risks short-term reward for a chance to discover something better.

You need both. Pure greedy gets stuck in bad habits; pure random never acts on what it learned.

ε-greedy: the simplest balance

Set ε high early (explore a lot) and decay it low later (mostly exploit). This homework decays linearly from 0.4 to 0.02 across 5000 episodes:

python
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)

Two subtle implementation bugs

The homework's choose_action is four lines, and two mistakes lurk:

• Argmax vs max. The greedy branch must return int(np.argmax(q_values)) — the index of the best action (an integer the env can execute), not np.max (the value, a float). Returning the value makes the env assert 1.8 in (0,1,2,3), which is False, and crash.
• Sampling order and the seeded RNG. Use the passed-in seeded rng.random() and rng.integers(4), and only sample the random action inside the explore branch. Pre-sampling on every call (even when going greedy) advances the RNG state differently and makes two logically-equivalent runs diverge from the same seed. Using global np.random.* bypasses the seed entirely and destroys reproducibility.

Chapter 6: The Reward Flip

Now the payoff of Problem 1. Same algorithm, same grid, same two goals. We change only rstep and watch three completely different optimal policies emerge — including one that refuses to do the task at all. This is the heart of the homework: the optimal policy is not determined by goal magnitude — the per-step reward reshapes the value landscape.

The arithmetic, done by hand

Both goals are reachable by an optimal-length path. Compute the total return of each path directly:

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

Goal 1 wins by 2. The far, valuable goal is worth the walk.

Scenario 2 — r_step = −2, R₁ = 10, R₂ = 5: goals unchanged, penalty doubled:

The flip. Goal 2 now wins by 1. The 3 extra steps to Goal 1 cost 6 with the doubled penalty — more than the 5-point reward gap. A single number moved, and the optimal goal switched from far to near.

Scenario 3 — r_step = +1, R₁ = R₂ = 1: the step reward is now positive:

Stalling wins by a landslide. The agent learns to walk LEFT into the wall and soak up +1 per step for all 20 steps. Reaching a goal would end that profitable stream after a single bonus. The start Q-values converge to exactly 50.0 — the geometric-sum value of "+1 forever."

Reading the actual Q-values

After training, the numbers tell the story. In Scenario 1 the start-state Q-values come out (LEFT, RIGHT, DOWN, UP) = (1.219, 2.265, 1.219, 2.265). LEFT and DOWN tie low (both walk into a wall, wasting the step); RIGHT and UP tie high (both head toward Goal 1's diagonal). In Scenario 2 they become (−4.996, −3.057, −4.996, −4.329) — now RIGHT is distinctly best and UP is worse, because the optimal path is straight right to Goal 2. The table figured out that the policy should change.

Chapter 7: PPO — Advantage & GAE

Leave the grid behind. Now the state is 39 numbers describing a robot arm, the action is a 4-D continuous vector, and the reward is sparse — 0 every step, +1 the instant a nail is driven, then done. You cannot enumerate a Q-table over a continuous action. So we switch to a policy-gradient method: PPO, the most widely used one in modern RL. It is on-policy — collect rollouts under the current policy, update, discard.

The objective and the policy gradient

We want to maximize expected discounted return, J(θ) = 𝔼τ~πθ[ Σt γt rt ]. The policy gradient theorem gives its gradient in a form we can estimate from samples:

In English: take each action you sampled and increase its log-probability in proportion to how good the future turned out after it. Good actions — push up. Bad actions — push down. The miracle of the log-derivative trick (∇ log p = (1/p) ∇ p) is that the unknown environment dynamics cancel out of the gradient. Implement just this and you have REINFORCE — correct, but painfully slow. Two problems make it slow, and PPO fixes both.

Problem 1: variance → use a baseline (advantage)

The return G_t swings wildly across rollouts. We can subtract any state-dependent baseline b(s_t) from G_t without changing the gradient's expectation (the cross-term provably has zero mean) but with a big drop in variance. The natural baseline is the state-value V(s) = 𝔼[G_t | s_t = s] — how good is this state on average? Subtracting it gives the advantage:

"How much better was action a than the policy's average action in this state?" Positive → above average, push toward it; negative → below average, push away. The clean policy gradient becomes:

Problem 2: how do we estimate A? → GAE

We do not know V, so we train a critic V_φ(s) alongside the policy. With a critic, the two extreme advantage estimators are:

Generalized Advantage Estimation (Schulman et al., 2016) blends the two with a knob λ ∈ [0, 1]. This is the exact content of your compute_gae:

Computed backwards through the trajectory. The (1 − dt) done-mask is the same terminal trick from Problem 1 — it zeros the bootstrap at episode ends. Unroll the recursion and it is a geometrically-weighted sum of future TD errors:

Each future TD error contributes less by a factor of γλ. The backward pass is just the O(T) way to compute this sum. At λ = 0 you get pure 1-step TD; at λ = 1, Monte Carlo. This homework uses λ = 0.99 — "trust the real rewards mostly, blend in the critic to dampen variance."

Chapter 8: PPO — The Clip

Vanilla policy gradient is unstable for one reason: a single big update can move the policy so far that the rollouts you collected no longer represent it, and learning collapses. PPO's whole stability mechanism is one clever objective that caps how far the policy can move per update.

The importance ratio

Define the ratio between the new and the rollout-time (old) policy for the same action:

If the policy has not changed, ρ = 1. Doubled the action's probability → ρ = 2; halved → ρ = 0.5. We compute it as exp(new_log_prob − old_log_prob), not a direct division, because densities for a 4-D continuous action can be tiny (1e-30); subtracting log-probs and exponentiating is numerically safe.

The clipped surrogate

Two terms inside the min. ρ · A is the unclipped surrogate (what vanilla PG uses). clip(ρ, 1−ε, 1+ε) · A caps ρ to [0.9, 1.1] for this homework's ε = 0.1. Then: take the elementwise min, average, negate (PyTorch minimizes; the objective is a max).

Why min — the pessimistic cap

Walk both signs of the advantage:

• A > 0 (good action). As we make it more likely, ρ rises. surr1 = ρ·A grows; surr2 is capped at (1+ε)·A. Once ρ > 1+ε, the min picks the cap — the gradient goes flat. No further reward for moving the policy further.
• A < 0 (bad action). As we make it less likely, ρ falls. Once ρ < 1−ε, surr2 = (1−ε)·A is more negative and the min picks it — flat again. No further punishment for moving further.

The min always picks the worse outcome from the optimizer's view — a pessimistic bound. The optimizer cannot game the surrogate by lunging too far in either direction. That is the entire PPO stability story.

The three extra loss terms

The full PPO loss adds three terms to the clipped surrogate:

Reverse KL to a frozen BC reference

Before RL, the policy is warm-started by behavior cloning on 20 expert demos. A frozen snapshot of that BC policy becomes the reference_actor. The reverse-KL term keeps the RL policy anchored near it:

Why: with sparse rewards, BC is the only thing that roughly knows how to do the task. Pure RL would forget it. It is called "reverse" because it is the KL from the new policy to the reference — mode-seeking, pushing the new policy to cover the reference. The off-policy problem solves the same "don't forget the demos" disease with a different antibiotic (alternating BC updates), which we meet in Chapter 9.

Chapter 9: Off-Policy & the Deadly Triad

Same robot, same hammer, same nail, same sparse reward. Completely different algorithm. PPO needed ~1,000,000 environment steps. This off-policy actor-critic reaches >90% success in ~100,000 — a 10× sample-efficiency gap. Understanding where that gap comes from is the whole lesson.

Why off-policy can reuse old data

PPO is on-policy: its gradient is an expectation over (s, a) drawn from the current policy, so stale data gives the wrong gradient — it must collect fresh rollouts and throw them away. Off-policy learns a Q-function instead, and the Bellman equation it enforces is a property of the environment, not of any policy:

As long as you have a tuple (s, a, r, s′), you can use it to enforce this constraint — regardless of which policy produced it. So you keep a replay buffer of every transition ever seen and replay it thousands of times. That is the source of the sample-efficiency win. The chef analogy: on-policy tastes a dish, writes the rating, throws the recipe away. Off-policy stores every entry in a giant cookbook and re-reads it forever.

The deadly triad

The clean story above hides a landmine. Three properties, together, make neural Q-learning diverge:

Any one is fine. Any two, usually fine. All three together — exploding Q-values and training collapse. Two concrete failure modes and their fixes:

Failure 1: the moving target. If the same critic computes both the prediction and the target, one gradient step moves both. You reach for the flag and it runs away at your own speed. Fix: a target network — a slow copy of the critic used only for targets, updated by a soft Polyak average φ̄ ← (1−τ)φ̄ + τφ with τ = 0.005. On the timescale of one gradient step, the target looks frozen. Symptom of forgetting it: the target stays at random init, Q never propagates beyond γ per step, learning crawls, the critic loss looks fine but the actor never improves.

Failure 2: maximization bias. Vanilla Q-learning's target uses maxa′ Q(s′, a′). If Q is noisy — some actions overestimated, some under, on average correct — the max systematically picks the overestimated ones. Watch:

You train Q toward an upward-biased target. Q grows. Next iteration, even more bias. Q-values explode. Fix: clipped double-Q — take the min over two critics, which counters the upward bias. That is Chapter 10.

The actor-critic fix for continuous actions

You cannot enumerate argmaxa Q(s, a) over a continuous 4-D action. So we train an actor π_θ(a | s) whose job is to output the maximizer. Its loss is simply:

"Plug the policy's chosen action into the critic; that scalar is what the actor wants to maximize" (minus to minimize). Gradient flows from Q back into θ via the reparameterization trick: sample a = tanh(μθ(s)) + 0.1·ε with ε ~ N(0,1) fixed — a is a differentiable function of θ. The critic tells you which actions are good; the actor learns to produce them; they bootstrap each other up. That ping-pong is actor-critic.

Chapter 10: Ensemble Critics & the Min-of-Two Target

The fix for maximization bias is the technical heart of Problem 3's update_critic. Three stabilization tricks stack into four lines.

Trick 1 — target network (Polyak/EMA)

Each step the target critic moves 0.5% toward the online critic. After ~200 steps it catches up, but on any single update it looks frozen — a stable regression target. Called once per critic update via utils.soft_update_params(net, target_net, tau).

Trick 2 — an ensemble of N critics

Maintain N independent critic networks, each initialized differently and fed different minibatch orderings. They disagree on out-of-distribution inputs, and that disagreement is a signal. Default N = 2; the ablation uses N = 10. With N = 10 + a high update-to-data ratio you approach REDQ ("Randomized Ensembled Double Q-learning"), a state-of-the-art sample-efficiency recipe.

Trick 3 — random pair + min for the target

When computing the TD target, do not use all N critics. Pick 2 at random, take the elementwise min:

The min counters maximization bias (two overestimates min'd together is less of an overestimate). The randomization — vs always using critics 1 and 2 — is the REDQ improvement: it forces every critic to stay reliable, otherwise critics 3..N drift because they are never used in targets.

The full target, four lines

python
with torch.no_grad():
    a_next = actor(s_next).sample(clip=stddev_clip)     # next action from current policy
    target_q_list = critic_target(s_next, a_next)        # list of N tensors
    Qi, Qj = random.sample(target_q_list, 2)              # 2 DISTINCT critics
    y = reward.unsqueeze(-1) + discount.unsqueeze(-1) * torch.min(Qi, Qj)
# loss over ALL N online critics:
critic_loss = sum(F.mse_loss(q, y) for q in critic(s, a))

Chapter 11: Showcase & Field Guide

One homework, three algorithms, one recursion underneath all of them. The Bellman equation you derived in Chapter 3 is the seed of every method in HW2. Here is the whole arc in one picture, then the cheat sheet.

The three-algorithm comparison

Cheat sheet — the equations to carry

Named failure modes — symptom → cause

Self-quiz — answer these without re-reading and you're ready

1. Why does Q-learning track actions in the table, not just states?
2. What would happen with γ = 1 in Scenario 3, and why?
3. Where does the if done: break go in the inner loop, and what breaks if it's misplaced?
4. What does λ control in GAE, and what are λ = 0 and λ = 1 called?
5. Why is the PPO ratio computed as exp of a difference, not a direct division?
6. Why does the min in the PPO surrogate stabilize training?
7. Why can off-policy reuse old data but on-policy cannot?
8. Why sample 2 critics for the target but train all N?
9. What's the difference between critic and critic_target, and when does each update?
10. Why does PPO need ~1M steps but the off-policy method ~100k?
11. What does the BC step during RL accomplish that pretraining alone can't?

Where this connects

This forge is the second stop in the CS224R arc (HW1 was imitation learning). Related on the site:

• CS224R HW1 — Imitation Learning — the first "homework as forge": BC, flow matching, DAgger.
• Q-Learning — the standalone Gleam on tabular and neural Q-learning.
• Policy Gradients — REINFORCE, baselines, and the theorem in full.
• Actor-Critic — the ping-pong of actor and critic, from scratch.
• Off-Policy Actor-Critic — replay buffers, target nets, double-Q, SAC.
• MDPs and RL Algorithms — the map of the whole territory.

Companion & practice forge for Stanford CS 224R Homework 2 (Online RL). Credits the public course materials. Implement the core math yourself; this is not a solution answer bank.