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.
• Under the usual visitation, step-size, and finite-table assumptions, tabular Q-learning converges to the optimal Q-function.
• Neural RL replaces explicit tables with learned functions and introduces separate policy objectives, replay choices, and stabilization methods.

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 corresponding discounted Bellman operator a contraction; convergence of a particular learner still needs its own assumptions.

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 π thereafter. The expectation is over the environment and policy randomness:

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

Convergence contract. With a finite Markov decision process, sufficient visits to every relevant (s, a) pair, bounded rewards, and an appropriate decaying step size, the standard tabular Q-learning theorem gives convergence to Q*. The update is a semi-gradient stochastic approximation to a moving Bellman target; it is useful to understand the target as a noisy sample of the recursion, rather than as ordinary gradient descent on a fixed Bellman-residual objective. The contraction and step-size assumptions are doing real work.

A worked propagation, by hand

Positive goal bonuses arrive only at goals, while each move still carries the step reward — yet after enough episodes every cell on the path can "know" it leads to reward. Here is the mechanism, step by step, with α = 0.2, γ = 0.98, r_step = −1, R₁ = 10.

Controlled backup 1. Start with a zero table and isolate a transition from (4, 2) going UP into Goal 1. Transition: reward = −1 + 10 = 9, done = True. Terminal target = 9:

In this isolated fixture, only that entry has changed. During an actual episode, earlier negative step rewards would also update visited entries.

Controlled backup 2. Next isolate a transition from (3, 2) going RIGHT to (4, 2): reward = −1, not done. Bootstrap target uses the 1.8 we just learned:

Now (3, 2) has evidence that going right is on a path to reward. Episode after episode, updates can spread that evidence toward the start — a bright wave spreading like dye through water. In this controlled corridor fixture, repeated backups approach the discounted path values; stochastic visits, step costs, terminal handling, and finite episodes determine the actual table.

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 can win in the continuing version of the task: the agent can collect +1 per step without ending the stream. A finite 20-step rollout collects only the discounted 20-step sum, while the continuing derivation is 50. Reaching a goal may still be optimal if the environment terminates or assigns a larger terminal return; the result depends on the full reward and termination contract.

Reading illustrative Q-values

For a representative illustrative snapshot, the numbers tell the story. In Scenario 1 the start-state Q-values are shown as (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 head toward the far goal by different staircase paths. In Scenario 2 an illustrative snapshot is (−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. These displayed values explain the policy flip; they are fixtures for intuition, not a benchmark claim.

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. It is on-policy: collect a rollout under a frozen old-policy snapshot, reuse that batch for the configured few epochs (three in the starter), then collect a fresh rollout after the policy update.

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) and a well-chosen baseline can reduce variance; an arbitrary baseline need not. 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 recurrence implemented by your compute_gae, together with its rollout masks:

Computed backwards through the trajectory. The (1 − dt) done-mask zeros the bootstrap at true episode ends; a timeout or environment reset needs the separate mask required by the rollout contract. 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 minus the baseline on a complete terminated trajectory; truncated rollouts can retain a final-value bootstrap. 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 can be unstable when a single big update moves the policy far from the rollout distribution. PPO uses an objective that removes incentive for some large ratio changes per batch; it is a trust-region-like heuristic, not a hard bound on every policy update or ratio.

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. This sample supplies no additional surrogate reward for increasing the probability 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. The objective removes additional surrogate reward for reducing the probability further.

The min always picks the worse outcome from the optimizer's view — a pessimistic batch surrogate. It can reduce the incentive to lunge past the clip, while diagnostics and the KL term still matter because the objective does not guarantee every update stays inside the interval.

The three extra loss terms

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

Distribution caveat. The starter clamps Gaussian action samples but inherits the ordinary Normal log-density. That density is not the exact distribution of the clamped output, which has boundary mass. Its PPO ratio and KL calculation therefore follow the implementation convention; they are not exact importance ratios or KL values for the bounded executed-action distribution.

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 a useful starting policy that pure RL can drift away from. It is called "reverse" because it is the KL from the new policy to the reference; that direction has different support and mode effects from the forward KL, so describe the direction rather than promising a universal mode-seeking behavior. The off-policy problem addresses demo retention with 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. The assignment uses on-policy and off-policy setups to study a sample-reuse tradeoff; any step counts or success rates depend on seeds, implementation, and evaluation. The curves later in this page are illustrative schematics, not measured results.

Why off-policy can reuse old data

PPO is on-policy: its gradient is an expectation over (s, a) drawn from the rollout-time policy, so stale data eventually gives the wrong gradient even though one batch can be reused for a few PPO epochs. Off-policy learns a Q-function whose Bellman target is defined by the transition dynamics plus the policy used for the next action:

A transition can remain useful under a different behavior policy when the critic and actor target are defined consistently, subject to coverage and distribution-shift limits. The finite replay buffer stores recent transitions and re-samples them several times; it does not keep every transition forever. That reuse can improve sample efficiency, while stale or narrow data can bias learning.

The deadly triad

The clean story above hides a landmine. Three properties together create a well-known risk of instability in neural Q-learning:

The combination can amplify errors; it does not guarantee divergence on every task or seed. 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. Forgetting it can produce stale or rapidly moving targets; inspect target and online predictions together rather than assuming one fixed symptom.

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 tends to pick an overestimated one. Watch:

You train Q toward an upward-biased target, so overestimation can compound across backups. Fix: clipped double-Q — take the min over two critics, which reduces upward bias on average at the cost of possible underestimation. 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 a reparameterized differentiable sample: the starter uses a clamped/TruncatedNormal-style action distribution with fixed Gaussian noise for the sample path. The executed action is therefore noisy during training, while the actor remains differentiable. 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 updates about 36.7% of a fixed initial gap remains; the response half-life is about 138 updates. It deliberately lags the online critic. Called once per critic update via utils.soft_update_params(net, target_net, tau).

Trick 2 — an ensemble of N critics

Maintain N independently initialized critic networks. In the starter, the ensemble evaluates the same sampled batch and its members can disagree on out-of-distribution inputs; that disagreement is a diagnostic. Default N = 2; the ablation uses N = 10. With N = 10 plus a high update-to-data ratio you get a REDQ-like configuration; its stability and sample efficiency are empirical properties to measure, not guarantees.

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 often counters maximization bias when critics are noisy, though it can under-estimate. The randomization — vs always using critics 1 and 2 — is a REDQ-style ingredient: the loss still trains every critic, while random pairs vary the target signal and make ensemble coverage measurable.

The target and replay shape contract

The equation above is the one-step case. The starter uses three-step replay: reward already contains the discounted reward sum and discount already includes the accumulated gamma factors and terminal mask. Use them directly; multiplying by gamma again is a bug.

python
def batch_col(x):
    x = torch.as_tensor(x)
    if x.ndim == 1: x = x[:, None]
    assert x.ndim == 2 and x.shape[1] == 1
    return x
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 = batch_col(reward) + batch_col(discount) * torch.min(Qi, Qj)
    assert y.shape == Qi.shape
# 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 algorithm families, and related value and policy ideas underneath them. The Bellman equation you derived in Chapter 3 is central to the tabular and off-policy critics; PPO adds a policy-gradient objective with GAE and a value baseline. Here is the 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. How does off-policy replay differ from the limited reuse of a fresh PPO rollout batch?
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. What tradeoffs could make a replay-based method use fewer environment samples than a fresh-rollout method? Treat any exact step counts as experiment-specific.
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.