A frozen dataset of someone else's wandering — two half-paths that never reach the goal. Plain imitation is stuck copying them. But a value function can stitch the good halves into a full path. Two algorithms do it: AWAC leans the policy toward the data's good actions; IQL learns a value that is a soft-max over in-distribution actions — without ever asking a question the data can't answer.
A companion & practice forge for Stanford's CS 224R Homework 3 (Offline RL). It credits the public course materials and teaches you to implement the core math yourself — it is not a copy-paste answer bank. The real starter-code contracts you meet here (the double-Q TD target, the AWAC exp-weighted actor, the IQL expectile_loss, update_v, update_q, the advantage estimate) are the ones from the homework, shrunk to a scale you can run in your browser with numpy alone — no torch, no D4RL, no cluster.
You are handed a box of driving logs and told to build a good driver. You cannot drive yourself — no test track, no simulator, no do-overs. Just the logs. That is offline reinforcement learning: learn a policy from a fixed dataset of someone else's experience, with no way to try things out.
Here is the twist that makes it interesting. Imagine a tiny puck — a PointMass — that must navigate to a goal. Your dataset contains two kinds of trip, and neither one reaches the goal:
Copy trip A and you get halfway, then stall. Copy trip B and you never leave the middle. But look closer: A's good first half ends near where B's good second half begins. Splice them — A's first half, then B's second half — and you have a full path to the goal that appears in no single trip in the dataset. That splice is called stitching, and it is the whole point of offline RL.
The start is bottom-left, the goal is top-right. Trip A (teal) does the first leg well then stalls; Trip B (purple) starts mid-way and finishes well. Toggle Stitch and watch the good halves join into a full route the data never contained.
The catch: learning a value function from a fixed dataset is dangerous. The moment you ask "what's the value of an action the data never tried?", a neural network will happily make up a number — usually too big — and the whole thing blows up. The two algorithms of HW3, AWAC and IQL, are two different disciplines for learning a value function without ever getting fooled by those made-up numbers. That is the story of the next ten chapters.
Before AWAC or IQL, we pin down the world and the data. HW3 uses a PointMass navigation task (and, in the full assignment, an 8-DOF AntMaze). PointMass is the right thing to hold in your head because it is 2-D: you can literally draw the trajectories and see the stitch.
A puck lives in a 2-D arena. Its state is its position (and velocity), a handful of numbers. Its action is a continuous 2-D push — a little force vector, each component in [−1, 1]. The reward is sparse: you get 0 at the goal and −1 every other step. So the return of an episode is just "minus the number of steps to reach the goal" — a fast route scores near zero, a wandering route scores very negative, and never reaching the goal scores the worst of all.
| Piece | PointMass | Why it matters |
|---|---|---|
| State s | puck position (+ velocity), continuous | Continuous → can't enumerate; need a function approximator, not a table. |
| Action a | 2-D force in [−1,1]², continuous | Infinitely many actions → you can't argmax over them; the OOD problem lives here. |
| Reward r | 0 at goal, −1 otherwise (sparse) | Sparse → can't be gamed by loitering; the only way to score is to finish. |
| Discount γ | 0.99 | Future reward matters almost as much as now — long-horizon credit. |
The defining constraint: during training you never touch the environment. No rollouts, no exploration. You read from a fixed set of transitions:
collected by some other policy — a mixture of decent and mediocre drivers. It never grows. If a transition is not in D, you will never see it.
HW3 ships a curated pointmass_stitching_dataset.npz designed to test exactly one thing. Its trajectories have:
| Statistic | Value | Meaning |
|---|---|---|
| Best single-trajectory return | −46 | The best any one trip in the box achieves. Behavior cloning is capped here. |
| Average trajectory return | −104 | Most trips are mediocre — they wander. |
The test: train an offline RL algorithm and beat −46. If your learned policy returns better than the best single trajectory in the dataset, it must have composed pieces of several trips — it stitched. Behavior cloning (even filtered to the top 10% of trips) can, by definition, only reach −46. Beating it is the proof.
A sample of the dataset's trajectories over the arena. Most (dim) wander and stall. Two are highlighted: the good first half and the good second half whose ends nearly touch. Drag the slider to see how few trips ever get close to the goal at all.
Here is the natural idea, and why it fails so badly it broke offline RL for years. You have a fixed dataset. Just run your favorite off-policy algorithm on it — train a Q-network with Bellman TD, extract the greedy policy. What could go wrong?
A Q-network is a neural net. Feed it a state and an action, it returns a number. For state-action pairs in the dataset, that number is anchored by training data — it's trustworthy. But the network is a smooth function over all actions, including the infinitely many it has never seen. There, it is pure extrapolation. Nothing constrains what Q outputs at an out-of-distribution (OOD) action — not the data, not the loss. And extrapolating neural nets tend to guess too high.
Now watch the loop. The policy's job is to find the action with the highest Q. If Q is accidentally overestimated at some OOD action a*, the policy happily moves toward a*. The critic then trains its Bellman targets using a* — but there is no real signal that a* is bad, so the inflated value persists. The policy chases it harder. Q climbs further. This is the actor-critic feedback loop, and left alone it diverges: Q-values run off to infinity, returns collapse.
The horizontal axis is a 1-D action. The data cloud (teal band) covers a middle region; the true value (dashed) is flat. Drag training forward: inside the data, the fitted Q stays honest, but outside it extrapolates upward, and the greedy policy (arrow) chases the highest point — further and further OOD as the bump grows.
Modern offline RL has two families of fixes, and AWAC and IQL are one from each:
| Family | The discipline | HW3's example |
|---|---|---|
| Constrain the policy | Keep the policy close to the data distribution, so it never proposes an OOD action for the critic to be fooled by. | AWAC |
| Constrain the value | Change value learning so it never queries Q at an OOD action in the first place. | IQL |
AWAC extracts the policy by weighted behavior cloning — it only ever imitates dataset actions, so it stays in-distribution. IQL introduces a state-only value function trained by expectile regression, so its Bellman target uses V(s′) and never needs an action at all. Same disease, two cures. We build both.
Both algorithms rest on the same two objects: the Q-function and the value function. We recap them precisely, because the entire art of offline RL is computing them without asking a question the data can't answer.
The Q-function Q(s, a) is the expected return if you take action a in state s, then act well forever. The value function V(s) is the expected return from state s under the policy — it does not fix an action:
Read V(s) as "how good is being here, if I act as my policy would." The advantage ties them together:
Positive advantage: this specific action was better than the state's average. Negative: worse. Advantage is the clean signal both AWAC and IQL use to decide which dataset actions to imitate more strongly — it strips out the "how good is this state" baseline and leaves only "how good was this action, here."
V(s) — a state-dependent baseline — centers the signal on zero, isolating the per-action difference. Same role the value baseline plays in a policy gradient.The Q-function refers to itself one step ahead. The Bellman target for a transition (s, a, r, s′, d) is:
The (1 − d) mask is the terminal case: if the episode ended (d = 1), there is no future, so the target is just the reward. Then Q regresses toward y by squared error. This much is HW2. The entire question of offline RL is: how do you compute "value of the next state" without querying Q at an OOD action?
| Approach | Value of s′ | OOD risk |
|---|---|---|
| Naive Q-learning | maxa′ Q(s′, a′) | Catastrophic — the max searches OOD actions and finds the overestimates. |
| AWAC | min(Q̄1, Q̄2)(s′, a′), a′ ~ π | Low — the policy is kept near the data, so a′ is roughly in-distribution. |
| IQL | V(s′) — a state-only function | None — no action is chosen, so no OOD query is possible. |
The arena, coarsely gridded. The goal cell has value 0; every other cell is −1 per step. Press step to run Bellman backups: value ripples outward from the goal, so every cell learns "how far am I from home." Because value is state-conditioned, a cell shared by two half-trips inherits the better continuation — that is stitching.
s appears in the wandering tail of trip A and in the good middle of trip B. The Q-function does not care which trip a transition came from — it just learns "from s, the action that leads toward the goal has the higher value." So it silently adopts B's good continuation at s, even though no full trajectory did A's first half followed by B's second. The value function is trajectory-blind and state-aware — and that is exactly what lets it splice.AWAC — Advantage-Weighted Actor-Critic (Nair et al., 2020) — is the "constrain the policy" answer. It has two halves: a critic that learns Q by Bellman TD, and an actor extracted by advantage-weighted regression. This chapter builds the critic; the next builds the actor.
AWAC's critic is exactly the off-policy critic from HW2 — two independent Q-networks with slow-moving target copies. For a batch of transitions (s, a, r, s′, d), the target is:
and the loss trains both online critics toward it:
Three pieces, each with a job:
| Piece | Name | Why it's there |
|---|---|---|
two Q̄1, Q̄2 | target networks | Slow, Polyak-averaged copies of the online critics. A regression target must be stable; if it moved every step you'd chase your own tail and diverge. |
min(·, ·) | clipped double-Q | The max over noisy estimates is biased upward (maximization bias). Taking the min of two independent critics is a pessimistic counter-bias — crucial when you can't reality-check. |
(1 − d) | terminal mask | If s′ is terminal, no future exists — the target is just r. Forget this and you bootstrap off a non-existent future and corrupt training. |
a′ comes from — the one place the critic touches the actorTo evaluate min(Q̄1, Q̄2)(s′, a′) you need an action a′ at the next state. AWAC samples it from the current actor: a′ ~ π(·|s′) (gradients off — it's used only in the target). This is what makes AWAC an actor-critic rather than pure offline Q-learning.
a′ is only approximately in-distribution. AWAC keeps its actor near the data (next chapter), so a′ is usually fine — but the actor is a Gaussian, and its tails sample actions that wander slightly outside the data's support. The target Q must then evaluate those slightly-OOD actions, and small errors leak in. AWAC works because the behavior-cloning anchor is strong enough to keep the drift small. With a weaker anchor or a smaller dataset, this coupling can still break. IQL removes it entirely — hold that thought for Chapter 6.no_grad disciplineTwo rules that are bugs if broken. First, the target is computed inside a stop-gradient block: gradients must flow only through the prediction Q(s,a), never through the target. Forget it and gradients flow into the target networks, the target moves while you chase it, and the loss explodes to millions. Second, the (1 − d) mask — miss it and terminal states bootstrap off garbage. You will write this exact target as Kernel 1 in the Studio.
Two critics each estimate the same next-state value (true value = 1.0, dashed) with independent noise. The max systematically overshoots — it selects whichever critic got lucky. The min of the two sits at or below the truth: a deliberate pessimism that keeps the offline target from running away. Reshuffle to see it hold across draws.
min(Q̄1, Q̄2) in the Bellman target instead of a single Q or the max?The critic gives us Q. Now the actor. This is where AWAC earns its name and does its stitching — and where the "constrain the policy" discipline lives.
The AWAC actor is trained by advantage-weighted regression: imitate the dataset, but weight each state-action pair by the exponential of its advantage.
Read it piece by piece:
log πψ(a|s) — the log-probability the policy assigns to the dataset's action. This is ordinary behavior cloning.(s, a) ~ D — the action comes from the data, never from the policy. This is the whole trick: the actor is only ever asked "given someone took a at s, raise log π(a|s) by how much?" It never proposes its own action, so it cannot drift OOD.A(s, a) = Q(s, a) − V(s) — the advantage. In AWAC, V(s) is a single-sample estimate: sample aπ ~ π(·|s) and use Q(s, aπ).exp(A / λ) — the weight. Always positive; larger when the action beat the state's average. λ is a temperature.exp(A/λ) factor is the only difference: it up-weights the dataset's better-than-average actions and down-weights the worse ones. So AWAC imitates the data (staying in-distribution) but leans toward its good parts (improving over it). Best of both. And because it up-weights the good action at each state regardless of which trajectory it came from, this is where stitching happens.The exponential is not arbitrary — it is the closed-form solution to "maximize expected advantage while staying within a KL-ball of the data policy." The temperature λ is the Lagrange multiplier of that KL constraint:
| λ | exp(A/λ) behavior | Resulting policy |
|---|---|---|
| → 0 (small) | sharp spike on the single best action | near-argmax over dataset actions — aggressive, risks overfit |
| moderate (~0.3–1) | gentle reweighting | tilted BC — the sweet spot AWAC uses |
| → ∞ (large) | all weights ≈ 1 | plain BC — safe but no improvement over the data |
With small λ and a large advantage, exp(A/λ) can overflow — exp(20) is already 4.8×10⁸. So the weights are clamped: exp_weights.clamp(max=50). One outlier capped at 50 in a batch of 256 is harmless; an uncapped exp(100) would swamp everything. You will write both the advantage estimate and the exp-weight as Kernels 2 and 3.
Each bar is a dataset action at one state; its height is the exponential weight exp(A/λ) on its behavior-cloning term. Green = better than average (up-weighted), red = worse (down-weighted). Drag λ: small makes it near-argmax on the best action; large flattens toward plain BC (all weights → 1).
(s, a) pairs, so its probability mass stays anchored on data actions. That means when the critic samples a′ ~ π(·|s′) for its target (Chapter 4), those samples are approximately in-distribution too. The actor keeps the critic honest, the critic keeps the actor pointed at good actions — a cycle that holds as long as the BC anchor is strong.(s, a) pairs from the dataset, never on (s, π(s))?AWAC works, but it has that one soft spot: its critic target still samples a′ ~ π, and the actor's Gaussian tails wander slightly OOD. IQL — Implicit Q-Learning (Kostrikov et al., 2021) — removes the soft spot entirely with one idea: never query Q at an action the policy chose. Its Bellman target uses a state-only value function V(s′), so no action is ever chosen in the value loop. The question becomes: how do you train that V? The answer is IQL's signature: expectile regression.
If you regress V(s) toward the Q-values of the dataset's actions with ordinary squared error, the minimizer is the mean:
But the mean is too pessimistic. The dataset at a state mixes good actions and mediocre ones; the mean dilutes the good with the bad. We don't want the average dataset action's value — we want the value of the good dataset actions. Not the mean, and not the max (the max would extrapolate OOD, the very thing we're avoiding). Something between the mean and the max.
The expectile is that "something between." It is a quantile-like statistic computed with squared — not absolute — error, which keeps it smooth for gradient descent. The expectile loss with parameter τ ∈ (0,1) is:
The 𝟙[u<0] is the indicator — 1 when u is negative, 0 otherwise. So the weight in front of u² is:
u > 0 (Q above V — V is too low): weight is τ;u < 0 (Q below V — V is too high): weight is 1 − τ.With τ = 0.9, being too low is penalized 9× more than being too high. So to minimize loss, V gets pushed up — until only about 10% of the Q-distribution sits above it. That is the upper expectile: a soft-max over the in-distribution actions.
τ = 0.9 learns a "soft max" of Q over the dataset's actions — optimistic enough to represent the good actions' value, conservative enough to never extrapolate beyond what was actually tried. At τ = 0.5 the two weights are equal and you're back to MSE (the mean). Sweep τ toward 1 and V rises toward the in-distribution max.The dots are Q-values of the dataset's actions at one state (spread along the axis). The fitted V (vertical line) is the τ-expectile. At τ = 0.5 it lands on the mean (dashed). Drag τ up and V rises toward the max — the value of the good in-distribution actions, never past them. The shading shows the asymmetric penalty: overshooting the fit (Q above V) costs τ, undershooting costs 1−τ.
Once V is trained by expectile regression on dataset (s, a) pairs, IQL's Q-target becomes:
Notice what's gone: no a′, no policy sample, no min(Q, Q)(s′, a′). V(s′) is a function of the next state alone. Every Q-value in the entire IQL pipeline is evaluated only at in-distribution actions. OOD overestimation is not "reduced" — it is structurally impossible.
V with expectile regression (τ near 0.9) instead of plain MSE?With the expectile V in hand, the rest of IQL falls out cleanly. IQL runs three updates per step, in a deliberate order, over three network roles.
| Network | Role | Trained by |
|---|---|---|
| Q1, Q2 (+ targets Q̄1, Q̄2) | estimate Q(s,a) at dataset actions | TD toward y = r + γ(1−d)V(s′) |
| V (state-only) | upper expectile of Q over dataset actions | expectile regression toward Q̄(s,a) |
AWAC had two roles (actor + critic). IQL adds the V-network as its own component — that's the extra network, and it's what buys the total OOD-freedom.
V regresses (by the expectile loss) toward the target Q at the dataset action:
Why target Q, not online Q? Same reason as any target network: if V regressed onto the online Q it would chase a target that moves every step. The slow target Q is stable on the timescale of one V update. You'll write this as Kernel 5 (using your expectile loss from Kernel 4).
Now the payoff. The Q-target uses V(s′) directly:
This is the same terminal-masked, stop-gradient Bellman target from Chapter 4 — but with the whole min(Q, Q)(s′, a′) machinery replaced by a single call to V(s′). No policy sample. No OOD query. You'll write this as Kernel 6.
Here is the elegant part: IQL extracts its policy exactly like AWAC does. Advantage-weighted regression, same equation:
The only difference from AWAC: how the advantage is computed. AWAC used a single-sample V(s) ≈ Q(s, aπ) — noisy. IQL has a clean, learned V-network, so A = Q(s,a) − V(s) is a smooth, low-variance signal. That's a big practical win: the actor gets a cleaner tilt and converges more reliably.
V using the current target Q; (2) extract the actor using the just-updated V; (3) update Q bootstrapping on V(s′); (4) Polyak-update the target Qs. Each reads the freshest available value and moves one step. And crucially: nowhere does IQL sample a′ ~ π(·|s′). The policy never enters the value loops at all — it only consumes the trained Q and V to decide which dataset actions to imitate.The three updates, animated. V regresses toward the expectile of target Q; Q regresses toward r + γV(s′); the policy reads Q and V for its advantage. Watch the arrows: the policy consumes values but never feeds an action back into them — the loop is policy-free, so no OOD query can occur.
a′ ~ π(·|s′) sampled inside the value-learning loop?Both algorithms extract their policy the same way — advantage-weighted regression. The entire difference is in how they compute Q and V, and that difference decides which one wins where.
| Component | AWAC | IQL |
|---|---|---|
| Actor loss | Identical: −mean( log π(a|s) · exp(A/λ) ) | |
| Q-target | r + γ min(Q̄1,Q̄2)(s′,a′), a′ ~ π | r + γ V(s′) — no actor |
| V estimate | Q(s, aπ) — single MC sample | separate V-net, expectile regression |
| Networks | actor + 2Q + 2 target Q = 5 | actor + 2Q + 2 target Q + V = 6 |
| OOD risk | possible if the actor drifts | none — V never queries OOD |
| Advantage variance | high (one sample) | low (smooth V-net) |
| Knobs | λ only | τ (expectile) and λ |
| Wins on | easier tasks, smaller datasets | harder, longer-horizon, sparse-reward tasks |
On a short task, AWAC's soft spot barely matters — a little OOD leakage, a little advantage noise, no big deal. But on a long-horizon task like a medium maze, TD errors compound over many bootstrap steps. AWAC's actor can drift over those many steps, its critic samples slightly-OOD a′, and the error accumulates. IQL never queries an OOD action anywhere in the value pipeline, so its value estimates stay reliable no matter how long the horizon. That's the whole asymptotic argument, in three sentences:
a′ ~ π(·|s′), which can drift outside the data as training progresses, letting OOD evaluation errors enter and propagate. (2) IQL replaces that sample with a learned V(s′) trained by expectile regression on dataset pairs, so every Q-query in the whole pipeline is at an in-distribution action — eliminating OOD overestimation by construction. (3) On longer-horizon tasks where TD errors compound over many bootstraps, IQL's stricter OOD avoidance gives more reliable values and substantially better final performance.Tuning IQL's expectile is a real deliverable. The pattern:
| τ | V learns | Effect |
|---|---|---|
| 0.2 | lower expectile (pessimistic) | advantages diffuse; policy collapses toward plain BC |
| 0.5 | the mean (MSE) | AWAC-like, but with a smoother mean-V |
| 0.7–0.9 | upper expectile | sharp advantage around 0; good actions heavily up-weighted; policy improves over data |
| 0.99 | near-max | can go unstable — V tracks noise if the Q-distribution is heavy-tailed |
On the umaze task, τ = 0.9 typically beats τ = 0.2: the upper expectile better represents what's achievable from each state, which sharpens the advantage signal and speeds learning.
Eval return vs training, on two task difficulties. On the easy task both reach the ceiling. Toggle to the hard task: AWAC plateaus below IQL as its OOD leakage compounds over the long horizon. The dashed line is the dataset's best trajectory — anything above it is stitching.
A = Q(s,a) − V(s) have lower variance than AWAC's?Everything comes together here. This is the payoff of the whole homework: a learned policy that stitches two half-paths into a full route to the goal — a route that appears in no single trajectory of the dataset.
The arena holds the same two highlighted trips from Chapter 0: trip A does the first leg well then stalls; trip B starts mid-way and finishes well. Behavior cloning is stuck — it can only replay whole trips. But IQL learns a value function whose contours point "downhill toward the goal" at every state. Where trip A stalls, its state is also visited by trip B, and there the value function has learned B's good continuation. So the extracted policy follows A's first half, then — at the shared state — peels off onto B's good second half. That is the stitch, executed.
Press Play. First watch Behavior Cloning replay trip A and stall (it can't do better than a single trip). Then watch IQL: the value contours (faint) push it along A's good first half, then hand off to B's good second half at the shared state — reaching the goal with a return better than the dataset's best. Toggle the value field on to see why the policy turns where it does.
And the mechanism is exactly the value function from Chapter 3, learned safely by the expectile trick from Chapter 6, extracted by the advantage-weighting from Chapters 5 and 7. Every piece you built is on screen at once.
The debugging reality of offline RL. When training misbehaves, the symptom usually points straight at one of the ideas in this lesson. Read the symptom, find the cause.
| Symptom | Likely cause | Fix |
|---|---|---|
| Q-values climb without bound; eval return crashes | OOD overestimation — the critic is being queried at unseen actions (or you dropped the min / used max) | Use min(Q̄1,Q̄2) (AWAC) or switch to IQL's V-bootstrap; verify no OOD query |
| Critic loss explodes to 1e6+ | Missing no_grad on the target — gradients leak into the target net, target moves while chasing | Wrap the target in torch.no_grad() |
| Q near goal creeps up; erratic near terminals | Missing (1−d) mask — bootstrapping off a non-existent future | Add the terminal mask to the Bellman target |
| Policy just reproduces the data; no improvement | Expectile τ too low (near 0.5) → V is the mean → diffuse advantages → plain BC. Or λ too large. | Raise τ toward 0.7–0.9; lower λ |
| V unstable, tracks noise | Expectile τ too high (near 0.99) → V chases the max of a noisy Q-distribution | Back τ off to ~0.9 |
| Actor loss is positive | BC term sign flipped — you're pushing away from the demos | The loss is −mean(logπ · w); check the minus |
inf/nan in the actor loss | exp(A/λ) overflowed — small λ, large advantage | Clamp the weights: exp_weights.clamp(max=50) |
| Shape-mismatch silent garbage (loss "fine" but never learns) | [B] reward broadcast against a [B,1] Q → [B,B] | .squeeze(-1) the V/Q outputs to [B]; print shapes |
exp(A/λ)).V that is a soft-max over in-distribution actions, so the Q-target uses V(s′) and never queries an OOD action — OOD-free by construction, and the reason it wins the hard tasks.This lesson completes the CS224R homework set. It builds on HW1 (Imitation Learning: BC, Flow Matching, DAgger) and HW2 (Online RL: Q-learning, PPO, off-policy actor-critic) — the Bellman recursion you saw scaled to neural nets there is exactly the engine that, here, learns a value function safely from a fixed dataset. Related Gleams: Offline RL, Off-Policy Actor-Critic, RL Algorithms.