X1.2.3

Lesson 1.2.3 — Policy Gradient and REINFORCE: The Score-Function Estimator, the Variance Problem, and the Baselines

Week 8 · Prerequisites: 1.2.1, 1.2.2, 0.1.2, 0.1.3, 0.3.2, 0.3.3.

Prior results used. This lesson uses 1.2.1’s ArmEnv (the observation/action spaces and the episode structure), 1.2.2’s MDP, return G_t, value functions V^π/Q^π, and the Bellman equations; 0.1.2’s autograd and the .backward() chain; 0.1.3’s nn.Module, optim, and the training loop; 0.3.2’s log-likelihood and its gradient (the score function); and 0.3.3’s filtering vs. smoothing (for the causality discussion). Each is restated at the point of use.

Hands forward. This lesson hands forward one capability: the learner can derive the policy gradient theorem, implement REINFORCE with and without a baseline, explain why the baseline reduces variance without introducing bias, control the estimator’s variance via the discount and the trajectory length, and diagnose the four policy-gradient bugs (missing log, wrong sign, no baseline, and non-stationary baseline) from their learning symptoms — the capability that PPO (1.2.4), SAC (1.2.5), and the VLA fine-tuning (3.1.4) all assume.

Running case. The same 6-DOF MuJoCo arm from 1.1.1–1.2.2, now with a neural network policy π_θ(a | s) over the Gaussian action distribution (the mean is a 6-vector from the network, the std is a learned 6-vector), and the same ArmEnv from 1.2.1 (the observation [q, q̇, p_ee, p_target], the reward −‖p_ee − p_target‖ − 0.01‖τ‖², the terminated/truncated). The lesson’s central result is that the policy gradient ∇_θ J(θ) = E_π[∇_θ log π_θ(a|s) · G_t] is unbiased but high-variance — a single-episode estimate can have a variance thousands of times the mean — and the baseline b(s) (the value function V^π(s)) reduces the variance by the relative magnitude of the return, without changing the estimator’s expected value.

§L.0 — Lesson map, objectives, and misconceptions

§L.0.1 — Lesson map and word budget

§L.0.2 — Objectives

By the end of this lesson you will be able to do eight things:

1. State the policy gradient’s objective J(θ) = E_{τ~π_θ}[G_0] and derive the policy gradient theorem ∇_θ J = E[Σ_t ∇_θ log π_θ(a_t | s_t) · G_t] from the score function.
2. Implement REINFORCE for the ArmEnv, with the Gaussian policy, the trajectory collection, and the autograd through the log-likelihood.
3. State why the estimator is unbiased (E[∇ log π · G] = ∇ E[G]), and derive the baseline’s variance-reduction property (E[∇ log π · b(s)] = 0 for any b(s)).
4. Apply the causality (reward-to-go G_t = Σ_{k≥t} γ^{k−t} r_k) to reduce the variance further, and state why the future rewards are unrelated to the past actions.
5. Compute the variance of the REINFORCE estimator for a simple 1-step problem by hand, and state the variance’s dependence on the γ and the horizon.
6. Parameterize the Gaussian policy (μ_θ(s) from the network, log σ_θ(s) from a state-independent or state-dependent head), and state the log-std’s role (the exploration vs. the greedy).
7. Diagnose the four policy-gradient bugs — missing log (the gradient is not the policy gradient), wrong sign (the gradient ascent vs. the descent), no baseline (the variance is too large, the learning is slow), and non-stationary baseline (the baseline tracks the policy, the estimator is biased) — from their learning symptoms.
8. Produce a working REINFORCE implementation for the ArmEnv, with the baseline (a learned V_φ(s)), the reward-to-go, and the diagnostic metrics (the return, the variance, the gradient norm) that identify the four bugs.

§L.0.3 — Misconception inventory

Each misconception is displaced in the section named, and each is tagged in the concept check that follows.

Everything is self-contained. Where a prerequisite lesson is used, the result is restated at the point of use.

§L.0.5 — Retrieval warm-up

This is the nineteenth lesson. The prompt requires drawing from lessons two to five steps back; 1.2.2 is one step back and is excluded. The warm-up draws from 1.2.1 (two steps back), 1.1.4 (three), 1.1.3 (four), and 1.1.2 (five).

Warm-up 1 (from 1.2.1, the ArmEnv). State the ArmEnv’s observation, action, and reward. Answer: the observation is [q, q̇, p_ee, p_target] (18 numbers, normalized); the action is the torque τ ∈ ℝ⁶ or the position q_target ∈ ℝ⁶; the reward is −‖p_ee − p_target‖ − 0.01‖τ‖²; the terminated is the success or the collision, the truncated is the time limit. How this is used: §L.2’s REINFORCE samples the action from π_θ(a|s), applies it to the ArmEnv, and collects the reward; the episode’s trajectory (s_t, a_t, r_t) is the data for the estimator.

Warm-up 2 (from 1.1.4, the pipeline’s health checks). State the pick-and-place pipeline’s health checks. Answer: the reachability, the branch continuity, the tracking error, the gripper’s force. How this is used: §L.4’s REINFORCE’s evaluation uses the success rate (the health check’s aggregation) as the objective, and the reward is the task’s signal — the pipeline’s checks are the success’s verification, and the policy’s reward is the task’s distance.

Warm-up 3 (from 1.1.3, the PID’s gains). State the PD’s K_p = J ω_n² and K_d = 2ζ J ω_n. Answer: the critically damped gains are the inertia’s J times the desired ω_n² and the 2ζω_n. How this is used: §L.4’s Gaussian policy’s mean is the learned control, and the std is the exploration — the learned policy replaces the hand-tuned PID, and the std is the exploration’s analogue of the PID’s damping (the larger std is the larger exploration, the smaller is the exploitation).

Warm-up 4 (from 1.1.2, the IK’s singularity). State the IK’s singularity (σ_min(J) = 0) and the DLS. Answer: the singularity is the σ_min = 0, and the DLS’s Jᵀ(JJᵀ + λ²I)⁻¹ is the bounded inverse. How this is used: §L.1’s policy gradient’s policy π_θ(a|s) can be a Gaussian with the mean = the IK’s output and the std = the DLS’s damping — the learned policy’s mean is the classical IK’s solution, and the std is the DLS’s λ.

Part A — The Objective and the Policy Gradient Theorem

§L.1 — Objects: the objective, the score function, and the gradient

The policy gradient’s objective is the expected return under the policy π_θ:

where τ = (s_0, a_0, r_0, s_1, a_1, r_1, …) is the trajectory sampled from the policy π_θ and the environment’s transition P.

The policy gradient theorem is:

where G_t = Σ_{k≥t} γ^{k−t} r_k is the reward-to-go from step t. The gradient is the expected value of the score function ∇ log π_θ(a_t | s_t) multiplied by the return G_t.

The score function is the gradient of the log-likelihood of the action:

It is the direction in the parameter space that increases the probability of the action a given the state s. The key property of the score function is its zero-mean under the policy:

The zero-mean is the foundation of the baseline’s variance reduction (Part C) — the score function’s expectation is zero, so any baseline b(s) can be added to the return without changing the estimator’s mean.

The derivation of the theorem:

1. ∇_θ J = ∇_θ ∫ p_θ(τ) G_0 dτ = ∫ ∇_θ p_θ(τ) G_0 dτ (the exchange of the gradient and the integral).
2. ∇_θ p_θ(τ) = p_θ(τ) ∇_θ log p_θ(τ) (the log-derivative trick).
3. log p_θ(τ) = log p(s_0) + Σ_t [log π_θ(a_t | s_t) + log P(s_{t+1} | s_t, a_t)] — the transition P is θ-independent, so its gradient is zero, and only the log π_θ terms survive.
4. ∇_θ J = E_τ[Σ_t ∇_θ log π_θ(a_t | s_t) · G_0] — the G_0 is the whole trajectory’s return.

The causality (the G_t rather than the G_0) is the refinement (Part C) — the future actions are not affected by the past actions, and the future rewards’ gradient with respect to the past actions is zero. The theorem’s exact form (with the G_t) is the causality’s version, and the G_0 version is the un-refined (the higher variance).

§L.1.1 — Result: the policy gradient is the score function times the return, not the reward’s gradient, and the theorem is the identity that converts the non-differentiable sampling into the differentiable log-likelihood

Claim. The reward r_t is not differentiable with respect to the policy’s θ — it is the environment’s output, and its dependence on θ is through the trajectory’s sampling (the a_t sampled from π_θ). The policy gradient theorem converts the non-differentiable expected return into the differentiable score function’s expectation, and the result is unbiased.

Proof sketch. The objective J(θ) = E_τ[G_0] is the integral of the trajectory’s return weighted by the trajectory’s probability p_θ(τ). The gradient ∇_θ J is the integral of G_0 ∇_θ p_θ(τ), and the log-derivative trick ∇_θ p = p ∇_θ log p converts it to E_τ[G_0 ∇_θ log p_θ(τ)]. The log-probability log p_θ(τ) is the sum of the log π_θ(a_t | s_t) and the log P(s_{t+1} | s_t, a_t), and the P’s θ-independence makes its gradient zero. The surviving terms are the ∇_θ log π_θ(a_t | s_t), and the result is the theorem.

Corollary. The theorem’s unbiasedness is conditional on the sampling’s correctness — the τ must be sampled from the policy’s distribution and the environment’s transition. A stale policy (the π_θ used for the sampling differs from the π_θ used for the gradient) biases the estimator (the off-policy issue, the PPO’s importance sampling in 1.2.4 corrects it).

Corollary. The theorem’s gradient is the sum of the per-step terms, and each term is the score function times the return-to-go. The per-step term’s magnitude is the return’s magnitude, and the variance is the return’s variance times the score function’s variance — the two are the variance’s sources (Part C).

§L.1.2 — Worked-example ladder

WE-L.1.a — Fully worked: the score function for a Gaussian policy

Input. A 1D Gaussian policy π_θ(a | s) = N(a; μ_θ(s), σ²) with the mean μ_θ(s) = θ · s and the fixed std σ = 0.1. The state s = 2, the action a = 0.5, the return G = 1.0.

Trace (the log-likelihood).

1. log π_θ(a | s) = −½ log(2πσ²) − (a − μ_θ(s))²/(2σ²).
2. = −½ log(2π·0.01) − (0.5 − 2θ)²/0.02.
3. = −0.5 log(0.0628) − 50(0.5 − 2θ)².

Trace (the score function).

4. ∂/∂θ log π_θ = −100(0.5 − 2θ) · (−2) = 200(0.5 − 2θ).
5. At θ = 0: ∂ log π/∂θ = 200 · 0.5 = 100.

Trace (the gradient’s estimate).

6. ∇_θ J ≈ ∇_θ log π · G = 100 · 1.0 = 100.
7. The update θ ← θ + α · 100 with the α = 0.001 gives θ = 0.1.

Output and interpretation. The score function at θ = 0 is 100 — a large value because the action 0.5 is far from the mean 0 (the std is 0.1, so 0.5 is 5σ away), and the gradient’s direction is positive (the mean should increase to reduce the distance to the action). The update increases the θ, and the next sample’s mean is closer to the action 0.5. The return G = 1.0 scales the gradient — a large return gives a large update.

Self-explanation 1. Before reading on: in step 6, the gradient’s estimate is 100. Is this the gradient of the objective with respect to θ, or just the score function?

Answer. The gradient’s estimate is the score function times the return — the 100 is the product of ∂ log π/∂θ and the G. The expected value over the action’s distribution is the true gradient ∇_θ J, and the single sample is the unbiased estimate. The 100’s magnitude is the sample’s value, and the variance across the samples is the estimator’s variance — the reason the baseline (Part C) reduces it.

Self-explanation 2. Before reading on: the score function’s 200(0.5 − 2θ) at θ = 0 is 100. What is the score function at the action a = 0 (the mean)?

Answer. The a = 0 is the mean at θ = 0, and the (a − μ) = 0, so the score function is −100 · 0 = 0. The action at the mean gives a zero score — the mean’s gradient is zero (the action is the most likely), and the gradient is zero because the mean is already optimal for that action. The score’s magnitude is larger for the actions far from the mean (the exploration’s rare samples), and the gradient’s variance is dominated by the rare large scores.

WE-L.1.b — Completion problem

Input. The same 1D Gaussian, the state s = 1, the action a = 0.3, the μ_θ(s) = θ s, the σ = 0.1, the return G = 2.0.

Trace (final 25% as prompts).

1. μ_θ(1) = θ · 1 = θ. ⟨your turn⟩
2. (a − μ) = 0.3 − θ. ⟨your turn⟩
3. ∂ log π/∂θ = −(a − μ)/σ² · (−s) = (0.3 − θ)/0.01 · 1 = 100(0.3 − θ). ⟨your turn⟩
4. At θ = 0, the score is 30, and the gradient’s estimate is 30 · 2.0 = 60. ⟨your turn⟩

Answers.

1. μ_θ(1) = θ. ✓
2. (a − μ) = 0.3 − θ. ✓
3. ∂ log π/∂θ = (0.3 − θ)/σ² = 100(0.3 − θ). ✓
4. The score is 30 at θ = 0, and the gradient’s estimate is 60. ✓

The lesson. The score’s 30 is smaller than the previous example’s 100 (the action 0.3 is 3σ away vs. 5σ), and the gradient’s 60 is smaller than the previous 100 · 1 = 100 — the score’s magnitude and the return’s magnitude together determine the gradient’s estimate. The return’s 2.0 (vs. the 1.0) doubles the gradient, and the score’s 30 (vs. 100) thirds it — the product is 60 (vs. 100).

WE-L.1.c — Faded: the score function’s zero-mean

Setup. The same Gaussian, the μ = 0, the σ = 0.1, and the a sampled from the π.

Question. What is the expectation of the score function E[∂ log π/∂θ]?

You must produce the intermediate steps. The setup and the answer are given.

Answer. The E[(a − μ)/σ²] = (E[a] − μ)/σ² = (μ − μ)/σ² = 0. The zero-mean is the property that the baseline (Part C) relies on — the E[∇ log π · b(s)] = b(s) · E[∇ log π] = 0 for any b(s), so the baseline can be subtracted without changing the expectation.

The lesson. The score function’s zero-mean is the Gaussian’s property (E[a] = μ), and it generalizes to any policy — the E_{a ~ π}[∇ log π(a|s)] = 0 is the score function’s universal property (the ∫ ∇π = ∇∫π = ∇1 = 0). The zero-mean is the baseline’s foundation.

WE-L.1.d — Contrastive near-miss: the reward’s gradient

Correct solution. The policy gradient is the score function times the return — ∇ log π · G.

Near-miss. The learner differentiates the reward with respect to the action (∂r/∂a), and uses it as the gradient — the incorrect “the reward’s gradient is the policy’s gradient.”

Why the near-miss produces a plausible but incorrect number. The reward r(s, a) is differentiable with respect to the action a, and the ∂r/∂a is the direction to increase the reward given the current action. But the policy is a distribution π_θ(a|s), and the gradient with respect to θ is not the ∂r/∂a — the ∂r/∂a is the action’s gradient, and the θ’s gradient is the ∂r/∂a · ∂a/∂θ only if the action is a deterministic function of θ (the deterministic policy gradient, the DPG, the DDPG’s foundation, not the stochastic REINFORCE’s). For the stochastic policy, the action is sampled, and the gradient is the score function — the near-miss’s ∂r/∂a ignores the sampling, and its estimator is biased (the wrong gradient).

The discrepancy, stated. The ∂r/∂a and the ∇_θ log π · G are different quantities — the former is the action’s gradient, the latter is the parameter’s gradient through the sampling. The two coincide only for the deterministic policy (the a = μ_θ(s)), where the sampling’s variance is zero and the score function degenerates to the ∂μ_θ/∂θ · ∂r/∂a. For the stochastic policy, the score function is the correct estimator, and the reward’s gradient is not.

Self-explanation 3. Before reading on: if the std σ → 0, does the score function’s estimator converge to the ∂r/∂a?

Answer. Yes — as σ → 0, the Gaussian becomes narrow, and the score function (a − μ)/σ² amplifies the deviation by 1/σ², but the sampled a is close to the μ (the deviation is O(σ)), so the product (a − μ)/σ² is O(1/σ) — diverges. The return G is smooth in the action, and the product’s expectation is the ∂G/∂a · ∂μ/∂θ (the deterministic gradient) plus a zero-mean noise term. The σ → 0 limit is the deterministic policy’s gradient, and the stochastic policy’s estimator converges to it (the variance decreases with the σ, but the finite σ’s variance is the exploration’s price).

Part B — REINFORCE: The Algorithm and Its Unbiasedness

§L.2 — Objects: the trajectory, the estimator, and the autograd

The REINFORCE algorithm:

1. Initialize the policy’s parameters θ (the network’s weights).
2. For each episode:
   1. Sample the trajectory τ = (s_0, a_0, r_0, s_1, a_1, r_1, …) by running the policy π_θ in the ArmEnv.
   2. Compute the return-to-go G_t for each step.
   3. Compute the loss L(θ) = −Σ_t log π_θ(a_t | s_t) · G_t (the negative of the estimator).
   4. Backpropagate L’s gradient to θ (the autograd).
   5. Update θ ← θ + α ∇_θ log π · G (the gradient ascent on the objective, or the gradient descent on the loss L).
3. Repeat until the return converges (or the budget is exhausted).

The loss’s negative form is the standard (the optimizers minimize the loss, and the objective is maximized, so the loss is the negative of the objective’s estimate). The autograd computes the ∇_θ L = −Σ_t ∇_θ log π · G_t, and the optimizer’s step ascends the objective.

The unbiasedness: the estimator g_t = ∇ log π(a_t | s_t) · G_t has the expectation E[g_t] = ∇_θ J (the policy gradient), so the sum Σ g_t is an unbiased estimate of the sum’s gradient, and the single-trajectory sample is the unbiased estimate of the objective’s gradient.

The variance: the estimator’s variance is Var[Σ_t ∇ log π · G_t], and the variance is the sum of the per-step variances plus the covariances. The magnitude of the variance is dominated by the G_t’s variance (the return’s spread across the episodes) times the score function’s variance (the action’s spread given the state). The two variance sources are the baseline’s (the G_t) and the causality’s (the per-step vs. the whole trajectory’s) targets.

§L.2.1 — Result: REINFORCE is unbiased and high-variance, and the variance is the return’s spread times the score function’s spread — the baseline and the causality are the two variance-reduction’s mechanisms

Claim. The REINFORCE estimator g = Σ_t ∇ log π(a_t|s_t) · G_t is unbiased (E[g] = ∇_θ J), and its variance is Var[g] = E[‖g‖²] − ‖∇_θ J‖², which is dominated by the return’s spread (G_t’s variance across the episodes) and the score function’s spread (∇ log π’s variance across the actions). The baseline (Part C) reduces the return’s contribution to the variance, and the causality (the G_t vs. the G_0) reduces the number of the terms (the per-step vs. the whole trajectory).

Proof sketch. The unbiasedness is the theorem’s corollary (the single-trajectory sample is the unbiased estimator of the gradient). The variance is the standard E[‖g − E[g]‖²], and the magnitude is the return’s variance (the large spread across the episodes) times the score’s variance. The baseline’s G_t − b(s_t) reduces the effective return’s variance (the subtraction of the b(s_t)), and the causality’s G_t (the sum from t onward) reduces the number of the terms (the earlier steps’ past rewards are removed), and the two together reduce the variance.

Corollary. The variance’s reduction is the sample’s efficiency’s key — the high variance means the many episodes are needed to estimate the gradient, and the baseline and the causality reduce the number of the episodes by the factor of the variance’s reduction. The practical REINFORCE always uses both, and the theoretical REINFORCE’s high variance is the pedagogical version.

Corollary. The variance’s reduction is not the bias’s cost — the baseline is zero-mean (Part C), so the estimator is still unbiased, and the causality is the theorem’s exact form (the G_t vs. the G_0 are both unbiased, and the G_t is the lower variance). The variance’s reduction is free (no bias), and the practical algorithms all use it.

§L.2.2 — Worked-example ladder

WE-L.2.a — Fully worked: the REINFORCE estimator for a 1-step episode

Input. A 1-step episode: the state s = 0, the action a = 1.5 sampled from the N(μ_θ(0) = θ, σ = 1), the reward r = a · 2 (the quadratic reward), the return G_0 = r = 3.0. The policy is π_θ(a|s) = N(a; θ, 1). The current θ = 0.

Trace (the log-likelihood and the score).

1. log π_θ(a | s) = −½ log(2π) − (1.5 − θ)²/2.
2. ∂/∂θ log π_θ = (1.5 − θ)/1 = 1.5 at θ = 0.

Trace (the REINFORCE estimator).

3. g = ∇ log π · G_0 = 1.5 · 3.0 = 4.5.

Trace (the update).

4. θ ← θ + α g = 0 + 0.1 · 4.5 = 0.45.

Trace (the unbiasedness check).

5. The true gradient ∇_θ J = ∇_θ E[a · 2] = 2 ∇_θ E[a] = 2 · 1 = 2 (the E[a] = θ, so the ∇_θ E[a] = 1).
6. The sample’s g = 4.5 is different from the true 2 (the variance), but the expected E[g] = 2 (the unbiasedness).

Output and interpretation. The sample’s g = 4.5 is larger than the true 2 (the action 1.5 is above the mean, and the return 3.0 is above the mean 2), and the product 1.5 · 3.0 = 4.5 is the sample’s over-estimate. The expectation over the many samples is 2, and the variance is the spread around the 2.

Self-explanation 4. Before reading on: the true ∇_θ J = 2 and the sample g = 4.5 — the ratio is 2.25. Is this the variance’s magnitude?

Answer. The 2.25 is the single sample’s ratio, not the variance. The variance is the expected squared deviation from the 2, and the single sample’s ratio is one realization. The variance’s magnitude is estimated by the many samples’ spread around the 2, and the single sample’s 4.5 is the high-variance’s example — the estimator’s variance is the key issue, and the baseline (Part C) reduces it.

Self-explanation 5. Before reading on: in step 5, the true ∇_θ J = 2 is the expectation’s gradient. Is the sample’s 4.5 an unbiased estimate of the 2?

Answer. Yes — the E[g] = 2 is the theorem’s corollary (the g is the unbiased estimator), and the 4.5 is one sample. The unbiasedness is the average property, and the single sample’s deviation is the variance. The unbiasedness’s verification is the many samples’ average, and the variance’s reduction is the goal of the baseline and the causality.

WE-L.2.b — Completion problem

Input. The same 1-step episode, but the a = 0.5 (the closer to the mean), the r = 1.0, the G_0 = 1.0.

Trace (final 25% as prompts).

1. ∂ log π/∂θ = (0.5 − θ) = 0.5 at θ = 0. ⟨your turn⟩
2. g = 0.5 · 1.0 = 0.5. ⟨your turn⟩
3. The true gradient is 2. ⟨your turn: is the sample’s g = 0.5 the over- or the under-estimate?⟩
4. The variance’s source: the a = 0.5 is smaller than the 1.5, and the g = 0.5 is smaller than the 4.5 — the variance’s magnitude depends on the a’s spread. ⟨your turn⟩

Answers.

1. 0.5. ✓
2. 0.5. ✓
3. The under-estimate (the 0.5 < 2). ✓
4. The g = 0.5’s magnitude is the sample’s, and the variance is the spread across the a’s — the small a’s g is small, the large a’s g is large, and the spread is the variance.

The lesson. The sample’s g depends on the a’s deviation from the mean ((a − μ)/σ²) and the return’s magnitude (G), and the variance across the samples is the estimator’s cost. The baseline (Part C) reduces the return’s contribution by subtracting the mean’s return, and the variance is reduced by the subtraction.

WE-L.2.c — Faded: the multi-step REINFORCE

Setup. A 3-step episode with the states s_0, s_1, s_2, the actions a_0, a_1, a_2, the rewards [1, 1, 1], the γ = 0.9, the policy π_θ (the Gaussian).

Question. What are the G_0, G_1, G_2 and the REINFORCE estimator g?

You must produce the intermediate steps. The setup and the answer are given.

Answer.

• G_2 = 1.
• G_1 = 1 + 0.9·1 = 1.9.
• G_0 = 1 + 0.9·1 + 0.81·1 = 1 + 0.9 + 0.81 = 2.71.
• The estimator g = ∇ log π(a_0|s_0)·2.71 + ∇ log π(a_1|s_1)·1.9 + ∇ log π(a_2|s_2)·1.
• The un-refined (G_0 for all the steps): g' = ∇ log π(a_0|s_0)·2.71 + ∇ log π(a_1|s_1)·2.71 + ∇ log π(a_2|s_2)·2.71.
• The difference g − g' = ∇ log π(a_1|s_1)·(1.9 − 2.71) + ∇ log π(a_2|s_2)·(1 − 2.71) = −0.81 ∇ log π(a_1) − 1.71 ∇ log π(a_2).

The lesson. The causality (G_t vs. the G_0) subtracts the future rewards from the past actions’ estimator, and the subtraction’s expectation is zero (the score function’s zero-mean). The variance is reduced by the removed terms, and the estimator is still unbiased.

WE-L.2.d — Contrastive near-miss: the log-likelihood’s sign

Correct solution. The loss = −Σ log π(a|s) · G (the negative), and the optimizer’s step ascends the objective.

Near-miss. The loss = +Σ log π(a|s) · G (the positive), and the optimizer’s step descends the loss, which descends the objective (the wrong direction).

Why the near-miss produces a plausible but incorrect number. The loss’s positive sign makes the optimizer’s minimization minimize the Σ log π · G, which reduces the probability of the high-return actions, and the policy learns the opposite of the intended. The return decreases over the training (the policy avoids the good actions), and the symptom is the decreasing return. The learner who reports the decreasing return and attributes it to the environment’s difficulty has missed the sign’s bug.

The discrepancy, stated. The loss −Σ log π · G and the +Σ log π · G differ by the sign, and the optimizer’s minimization descends the loss’s gradient. The negative ascends the objective (the correct), and the positive descends it (the wrong). The sign’s choice is the gradient ascent vs. the descent, and the two are opposite.

Self-explanation 6. Before reading on: if the optimizer’s learning rate is negative (the ascent), does the positive sign become correct?

Answer. Yes — the negative learning rate ascends the loss’s gradient, which ascends the +Σ log π · G’s value (the objective), and the positive sign with the negative rate is equivalent to the negative sign with the positive rate. The two conventions (the positive sign + the negative rate, the negative sign + the positive rate) are equivalent, and the practical standard is the negative sign + the positive rate (the optimizers minimize). The bug is the mismatch between the sign and the rate’s direction, not the sign itself.

Part C — Variance Reduction: Baselines and Causality

§L.3 — Objects: the baseline, the causality, and the zero-mean

The baseline b(s) is a function of the state, and the baseline-subtracted return is G_t − b(s_t). The estimator becomes:

The key property: the E[∇ log π · b(s)] = 0 for any b(s) (the score function’s zero-mean), so the baseline’s expectation is zero, and the estimator is still unbiased:

The variance is reduced because the G_t − b(s_t)’s variance is smaller than the G_t’s variance when the b(s_t) is close to the G_t’s mean. The optimal baseline is the value function V^π(s_t) (the expected return given the state), which minimizes the variance’s magnitude (the standard result).

The causality is the refinement of the theorem’s G_t (the sum from t onward) vs. the un-refined G_0 (the whole trajectory’s return). The causality exploits the fact that the future rewards are independent of the past actions (the action a_t does not affect the r_{t-1}), so the past rewards’ contribution to the a_t’s gradient is zero. The causality’s variance is lower than the un-refined’s (the removed terms), and the estimator is still unbiased.

The combination (the baseline + the causality) is the practical REINFORCE:

where the V_\phi is a learned value function (the actor-critic’s critic, the PPO’s baseline, the A2C’s critic). The G_t − V_\phi(s_t) is the advantage A(s_t, a_t) (the action’s excess over the state’s expected return), and the advantage’s expectation is zero (the action’s expected excess is zero given the state), so the estimator is unbiased.

§L.3.1 — Result: the baseline is free (no bias) and the causality is free (no bias), and the two together reduce the variance by the return’s spread and the past rewards’ credit — the value function V^π is the optimal baseline

Claim. The baseline b(s) preserves the estimator’s unbiasedness (the E[∇ log π · b] = 0), and the optimal b(s) (the variance-minimizing) is the V^π(s) (the value function). The causality (G_t) preserves the unbiasedness (the theorem’s exact form) and reduces the variance by the past rewards’ removed credit. The combined estimator’s variance is the minimum among the unbiased estimators of the policy gradient (the value’s exact V^π).

Proof sketch. The baseline’s zero-mean is the score function’s ∫ ∇π = ∇∫π = ∇1 = 0, and the subtraction’s expectation is −b(s) · 0 = 0. The optimal b is the conditional expectation E[G_t | s_t = s] = V^π(s) (the variance’s minimization is the standard Var[X − b] = Var[X] − 2 Cov[X, b] + Var[b], minimized by the b = E[X]). The causality is the theorem’s exact form (the G_0 is the un-refined, the G_t is the refined), and the two are both unbiased (the past rewards’ gradient is zero for the past actions, and the causality removes the zero terms).

Corollary. The value function V^π is the optimal baseline, and the learned V_φ (the critic’s network) is the approximation. The actor-critic (the A2C, the A3C, the PPO) is the practical algorithm with the V_φ’s learned baseline, and the REINFORCE with the baseline is the simplest version.

Corollary. The causality’s reduction is smaller than the baseline’s (the past rewards’ removed credit vs. the return’s spread’s reduction), and the baseline is the primary variance reduction. The two are complementary, and the combined estimator’s variance is the minimum.

§L.3.2 — Worked-example ladder

WE-L.3.a — Fully worked: the baseline’s variance reduction

Input. The 1-step episode from WE-L.2.a (the a = 1.5, the G = 3.0), the baseline b = E[G] = 2.0 (the true value), the g_b = 1.5 · (3.0 − 2.0) = 1.5.

Trace (the baseline’s variance).

1. The original g = 4.5 (the return 3.0’s product).
2. The baseline’s g_b = 1.5 · 1.0 = 1.5 (the smaller magnitude).
3. The variance of the original g: the Var[∇ log π · G].
4. The variance of the baseline’s g_b: the Var[∇ log π · (G − 2)].

Output and interpretation. The baseline’s g_b = 1.5 is smaller than the g = 4.5, and the variance of the g_b across the samples is smaller (the subtraction’s G − 2 is smaller than the G when the G is near the 2, and the variance is reduced). The improvement is the variance’s reduction, and the unbiasedness is preserved (the expectation is the same 2).

Self-explanation 7. Before reading on: the baseline b = 2 is the true E[G]. What if the baseline is the wrong b = 5?

Answer. The wrong b = 5 still preserves the unbiasedness (the E[∇ log π · 5] = 5 · 0 = 0), but the variance’s reduction is smaller (the G − 5’s variance is larger than the G − 2’s, because the 5 is far from the mean 2). The baseline’s choice matters for the variance (the closer to the V^π, the better), but the unbiasedness is preserved for any b(s).

WE-L.3.b — Completion problem

Input. The 3-step episode from WE-L.2.c, the baseline V_\phi(s_t) = 1.0 for all the t (the constant baseline), the γ = 0.9.

Trace (final 25% as prompts).

1. The advantages: the A_0 = G_0 − 1 = 2.71 − 1 = 1.71, the A_1 = 1.9 − 1 = 0.9, the A_2 = 1 − 1 = 0. ⟨your turn⟩
2. The estimator g_bc = Σ ∇ log π(a_t) · A_t. ⟨your turn⟩
3. The zero-advantage’s step (A_2 = 0) contributes nothing — the last step’s action is not reinforced. ⟨your turn: why is this correct?⟩
4. The variance’s reduction: the advantages’ [1.71, 0.9, 0] are smaller than the returns’ [2.71, 1.9, 1], and the variance is reduced. ⟨your turn⟩

Answers.

1. The A_0 = 1.71, the A_1 = 0.9, the A_2 = 0. ✓
2. The g_bc = ∇ log π(a_0) · 1.71 + ∇ log π(a_1) · 0.9 + ∇ log π(a_2) · 0. ✓
3. The A_2 = 0 is the zero-advantage — the a_2’s return is equal to the baseline, and the action is neither reinforced nor penalized. The correctness is the advantage’s interpretation — the action is as good as the average, and the gradient’s zero is the correct signal (the no update for the neutral action).
4. The advantages [1.71, 0.9, 0] are smaller than the returns [2.71, 1.9, 1], and the variance is reduced by the baseline’s subtraction (the constant 1) and the causality (the past rewards’ removed credit).

The lesson. The advantages [1.71, 0.9, 0] are the baseline-subtracted returns, and the last step’s 0 is the neutral action (the correct signal). The variance’s reduction is the two mechanisms’ combination (the baseline’s subtraction and the causality’s per-step return), and the estimator is still unbiased.

WE-L.3.c — Faded: the value function’s learning

Setup. The 3-step episode, the learned V_φ(s_t) (a neural network with the parameters φ), the target G_t (the return-to-go).

Question. How is the V_φ learned, and what is its loss?

You must produce the intermediate steps. The setup and the answer are given.

Answer. The loss is the mean-squared error L(φ) = Σ_t (V_φ(s_t) − G_t)² (the value’s regression), and the optimizer’s step descends the L’s gradient. The V_φ approximates the V^π (the true value), and its accuracy determines the baseline’s quality (the closer to the V^π, the larger the variance’s reduction).

The lesson. The critic’s learning is the supervised regression on the return-to-go, and its accuracy is the variance’s reduction’s key. The actor-critic’s two networks (the actor π_θ and the critic V_φ) are learned simultaneously, and the critic’s accuracy is the baseline’s quality.

WE-L.3.d — Contrastive near-miss: the non-stationary baseline

Correct solution. The baseline V_φ(s) is a function of the state (the V^π’s approximation), and its gradient with respect to θ is zero (the V_φ is independent of the actor’s θ).

Near-miss. The baseline b = G_0 (the current episode’s return) — the baseline is the same return, and the advantage G_t − b is G_t − G_0.

Why the near-miss produces a plausible but incorrect number. The baseline G_0 is the current episode’s return, and the advantage G_t − G_0 is negative for every t > 0 (the G_t < G_0), and the estimator’s expectation is not zero (E[∇ log π · (G_t − G_0)] ≠ 0). The bias is non-zero, and the estimator descends the wrong direction. The symptom is the decreasing return (the policy avoids the high-return actions because the advantage is negative), and the fix is the state-dependent baseline (the V_φ(s), not the current return).

The discrepancy, stated. The b(s) (the state-dependent) and the b = G_0 (the current return) differ by the state-dependence — the V_φ(s)’s expectation E[∇ log π · V_φ(s)] = 0 (the score function’s zero-mean), and the G_0’s expectation E[∇ log π · G_0] ≠ 0 (the G_0 is correlated with the action). The near-miss’s baseline is not the b(s), and the estimator is biased.

Self-explanation 8. Before reading on: in the near-miss, if the baseline is the average of the past episodes’ returns (b = E[G_0]), is the estimator unbiased?

Answer. Yes — the b = E[G_0] is a constant (the average of the past returns), and the constant’s gradient is zero, so the E[∇ log π · b] = b · 0 = 0, and the estimator is unbiased. The constant baseline is the simple unbiased choice, and its variance’s reduction is smaller than the state-dependent V_φ(s)’s (the constant doesn’t capture the state’s variation). The unbiasedness is preserved for any b that doesn’t depend on the action (the constant or the state-dependent), and the current return G_0 depends on the action (the action’s trajectory), so it is biased.

Part D — The Gaussian Policy and the Practical Issues

§L.4 — Objects: the network, the log-std, and the gradient’s scale

The Gaussian policy for the continuous action space:

where the mean μ_θ(s) is the network’s output (the nn.Linear’s last layer, the 6-vector for the running case), and the std σ_θ(s) is the network’s separate output (the log-std log σ_θ(s), exponentiated for the positive std). The log-likelihood:

The log-std is the standard parameterization (the log σ is unconstrained, and the exp gives the positive σ), and the state-dependent log σ_θ(s) adapts the exploration to the state (the large std for the uncertain states, the small for the certain).

The gradient’s scale: the score function ∇_θ log π is proportional to (a − μ)/σ² and the ∇_θ μ/σ, and the gradient’s magnitude is sensitive to the σ’s scale. The small σ amplifies the deviation (the large score), and the large σ dampens it. The variance is the same dependence, and the optimizer’s step size must account for the scale.

The action’s squashing: the Gaussian is unbounded, and the real action’s range is bounded (the [−τ_max, τ_max]). The tanh squashing (a = τ_max · tanh(z)) bounds the action, and the log-likelihood must include the Jacobian (log |da/dz| = log(τ_max (1 − tanh²(z)))), or the log-likelihood is wrong. The squashing’s correction is the SAC’s (1.2.5) standard, and the REINFORCE’s simple version omits it (the unbounded action is the approximation).

§L.4.1 — Result: the Gaussian policy’s log-std is a learnable exploration’s parameter, and the tanh squashing’s Jacobian must be included for the bounded actions

Claim. The Gaussian policy’s std σ_θ(s) is a learnable parameter, and the network adapts it to the state (the large std for the exploration, the small for the exploitation). The tanh squashing (a = τ_max tanh(z)) bounds the action, and the log-likelihood’s transformation log π(a) = log p(z) − log |da/dz| includes the log |da/dz| (the Jacobian), or the gradient is biased.

Proof sketch. The change-of-variables formula: the p(a) = p(z) |dz/da|, and the log p(a) = log p(z) + log |dz/da| = log p(z) − log |da/dz|. The da/dz = τ_max (1 − tanh²(z)) (the tanh’s derivative), and the log is log τ_max + log(1 − tanh²(z)). The omission of the −log |da/dz| changes the log-likelihood by the Jacobian’s term, and the gradient is biased by the Jacobian’s gradient.

Corollary. The log-std’s learning is the exploration’s adaptation — the early training’s large std (the broad exploration) and the late training’s small std (the narrow exploitation), and the learned std adapts to the task’s difficulty (the harder states get the larger std). The fixed std is a simplification that underperforms for the varied tasks, and the learned std is the standard.

Corollary. The tanh squashing’s Jacobian is the SAC’s standard, and the REINFORCE’s simple version (the unbounded action) omits it — the unbounded action’s distribution is the Gaussian, and the real robot’s bounds are ignored. The practical implementation must include the squashing’s Jacobian for the bounded action, or the policy learns the unbounded action’s behavior (the large torques) and saturates the real robot.

§L.4.2 — Worked-example ladder

WE-L.4.a — Fully worked: the Gaussian log-likelihood’s gradient

Input. A 1D Gaussian policy, the μ_θ(s) = θ s, the fixed σ = 1, the state s = 2, the action a = 3, the return G = 1.

Trace (the log-likelihood).

1. log π = −½ log(2π) − (a − θ s)²/2.
2. ∂/∂θ log π = (a − θ s) s (the chain rule: the −½(2)(a − θs)(−s) = (a − θs)s).
3. At θ = 0: (3 − 0)·2 = 6.

Trace (the gradient).

4. g = 6 · 1 = 6.
5. The true ∇_θ J = ∇_θ E[a · 1] = ∇_θ θ s = s = 2 (the E[a] = θ s).

Output and interpretation. The sample’s g = 6 and the true 2 — the sample’s over-estimate (the action 3 is above the mean 0, and the score 6 is large). The variance is the spread, and the baseline reduces it. The log-std’s gradient (if the std is learnable) is the ∂/∂(log σ) log π = ((a − μ)²/σ² − 1), and the sign depends on the deviation (the large deviation increases the std, the small decreases it).

Self-explanation 9. Before reading on: the log-std’s gradient is ((a − μ)²/σ² − 1). What is its sign for the |a − μ| > σ (the large deviation)?

Answer. The ((a − μ)²/σ² − 1) > 0 for the |a − μ| > σ, so the log-std’s gradient is positive, and the optimizer’s step (the ascent on the objective) increases the log σ — the std is increased. The increase is the exploration’s adaptation — the large deviation is the surprising action, and the std’s increase broadens the future samples (the more exploration). The small deviation (the |a − μ| < σ) gives a negative gradient, and the std decreases (the exploitation’s narrowing).

WE-L.4.b — Completion problem

Input. The same Gaussian, the σ = 0.1 (the small std), the action a = 0.15 (the 1.5σ deviation), the G = 1.

Trace (final 25% as prompts).

1. ∂ log π/∂θ = (a − θ s) s / σ² = (0.15 − 0) · 2 / 0.01 = 30. ⟨your turn⟩
2. g = 30 · 1 = 30. ⟨your turn⟩
3. The true ∇_θ J = 2. ⟨your turn: the sample’s 30 is 15× the true — the variance’s magnitude?⟩
4. The variance is the spread across the a’s, and the small σ = 0.1 amplifies the (a − μ)/σ² by the 1/σ² = 100 — the variance is large for the small σ. ⟨your turn⟩

Answers.

1. 30. ✓
2. 30. ✓
3. The sample’s 30 is 15× the true 2, and the variance is large (the 1/σ² = 100 amplification). ✓
4. The variance is large for the small σ, and the small σ’s exploration is narrow (the samples are close to the mean) — the variance is the small σ’s price. ✓

The lesson. The small σ = 0.1 amplifies the score by 1/σ² = 100, and the variance is large. The large σ dampens the score and the variance, and the trade-off (the exploration vs. the variance) is the log-std’s choice. The learned std adapts to the state (the uncertain states get the larger std, the smaller variance), and the fixed std is the simplification.

WE-L.4.c — Faded: the tanh squashing’s Jacobian

Setup. The same Gaussian, the a = τ_max tanh(z) with the τ_max = 1, the z = 0.5 (the pre-squashing), the a = tanh(0.5) ≈ 0.462.

Question. What is the log |da/dz|, and how does the squashing’s log-likelihood differ from the unbounded’s?

You must produce the intermediate steps. The setup and the answer are given.

Answer. The da/dz = 1 − tanh²(0.5) = 1 − 0.462² = 1 − 0.213 = 0.787, and the log |da/dz| = log(0.787) ≈ −0.239. The squashed log-likelihood is log π(a) = log p(z) − log |da/dz| = log p(z) + 0.239 (the Jacobian’s correction). The unbounded’s log p(a) is the log p(z) (the no correction), and the difference is the 0.239 (the Jacobian’s term).

The lesson. The squashing’s Jacobian is the necessary correction for the bounded action, and the omission biases the log-likelihood. The SAC’s (1.2.5) standard includes it, and the REINFORCE’s simple version omits it (the unbounded approximation). The bounded action’s correctness requires the Jacobian, and the squashing’s derivative (1 − tanh²) is the standard.

WE-L.4.d — Contrastive near-miss: the unbounded Gaussian with the real robot

Correct solution. The squashed Gaussian (a = τ_max tanh(z)) with the Jacobian’s correction.

Near-miss. The unbounded Gaussian (a = μ + σ ε) with the real robot’s saturation (the τ_max).

Why the near-miss produces a plausible but incorrect number. The unbounded Gaussian samples the a from the Gaussian, and the real robot’s saturation clips the a > τ_max to the τ_max — the actual action is the clipped value, and the policy’s log-likelihood is the un-clipped Gaussian’s. The mismatch biases the gradient (the policy learns the un-clipped distribution, and the real behavior is the clipped), and the symptom is the saturation’s stickiness (the arm stays at the τ_max and can’t explore the large action). The fix is the squashed Gaussian (the bounded by construction), and the Jacobian’s correction.

The discrepancy, stated. The unbounded Gaussian and the squashed Gaussian differ by the bounded range and the Jacobian’s term. The unbounded’s saturation is the outside the policy’s model, and the policy can’t reason about the saturation’s effect. The squashed’s bounded by construction, and the Jacobian’s correction makes the log-likelihood accurate. The two are the standard vs. the naive, and the squashed is the correct.

Self-explanation 10. Before reading on: in the near-miss, if the real robot’s τ_max is large (the 10 N·m), and the policy’s samples are within the [−τ_max, τ_max] (the no saturation), does the unbounded Gaussian work?

Answer. Yes — the no saturation means the unbounded Gaussian’s samples are within the bounds, and the log-likelihood is accurate (the no clipping). The squashing’s Jacobian is unnecessary (the no squashing), and the unbounded is fine. The issue is the saturation (the outside the bounds), and the squashing is the fix for the bounded action. The choice depends on the action’s range (the large τ_max for the no saturation, the small for the saturation’s risk).

Part E — Failure Taxonomy, Lab, Cheat Sheet, Exercises, Exit Test

§L.5 — Failure taxonomy

§L.6 — The lab

§L.6.1 — Task

Implement REINFORCE for the ArmEnv, with and without the baseline (V_φ) and the causality (G_t), verify the unbiasedness (the average gradient matches the true for a simple 1-step problem), and diagnose the four bugs (the missing log, the wrong sign, the no baseline, the non-stationary baseline) by injecting them.

§L.6.2 — Conditions

Python 3.10+, PyTorch 2.2+, CPU. Seed 42, dtype = float64. Runtime under 5 minutes. The ArmEnv’s simplified version (the 2D position with the 1-step episodes), the Gaussian policy (the nn.Linear(4, 2) for the mean, the nn.Parameter for the log-std), the return’s normalization.

§L.6.3 — Controls held fixed

The seed, dtype, the ArmEnv, the policy’s architecture, the learning rate (1e-3), and the batch size (10 episodes) are fixed. The baseline, the causality, and the sign are varied.

§L.6.4 — Metrics table

§L.6.5 — Pre-registered hypotheses

The lab computes both estimators on the same problem over 20,000 samples, and trains four variants for 500 episodes each. Measured values below.

• H1. The 1-step problem’s unbiasedness_error < 0.1 — the score-function average matches the pathwise gradient. Measured: relative error 0.2405 at N = 20,000. Refuted, and the failure is informative rather than mysterious: both estimators are unbiased, but the score-function one has a per-sample variance 6,659× the pathwise one (H2), so its sample mean is still 24% away at 20,000 draws. This comparison tests how much variance each estimator carries, not whether either is biased.
• H2. The baseline’s variance_estimate is < 50% of the no-baseline’s. Measured: the score/pathwise variance ratio is 6659.3×. Holds for the estimator comparison, overwhelmingly. The lab does not isolate the baseline’s own contribution to variance — it compares the two estimators — so “the baseline halves the variance” is not what this number tests. The reparameterization trick’s advantage is real and enormous; the baseline’s is a separate measurement this lab does not make.
• H3. The no-baseline return curve is flat and the with-baseline curve increases. Measured: no-baseline −0.4789, with-baseline −0.1821. Holds, by 0.297 return — the baseline is worth about 62% of the remaining gap to zero.
• H4. The wrong sign’s return decreases, and the missing log’s gradient is unstable. Measured: wrong sign −0.6017, missing log −0.5328, both worse than the with-baseline −0.1821. Holds on the returns. The “gradient unstable” half is not printed by the lab.

§L.6.6 — The break-it sequence

1. The wrong sign (silent). Use the +Σ log π · G (the positive loss). Predicted symptom: the return decreases (the policy avoids the good actions). Actual: as predicted — silent.
2. The no baseline (silent). The no V_φ (the raw G_t). Predicted symptom: the variance is large, and the learning is slow (the many episodes). Actual: as predicted — silent.
3. The non-stationary baseline (silent). The b = G_0 (the current return). Predicted symptom: the advantage is negative, and the return decreases (the bias). Actual: as predicted — silent.

§L.6.7 — Acceptance criteria

§L.6.8 — What good looks like

A strong submission has five properties:

1. The REINFORCE’s estimator is implemented with the log π, the return, the baseline, and the causality, and the unbiasedness is verified on the 1-step problem.
2. The baseline’s variance’s reduction is measured (the Var[g] with and without), and the ratio is reported.
3. The causality’s effect is shown (the G_t vs. the G_0), and the return’s improvement is measured.
4. The four break-its are predicted with the symptoms (the decreasing return, the slow learning, the bias) before running.
5. The Gaussian policy’s log-std is learnable, and the learned std’s evolution is shown (the large early, the small late).

§L.6.9 — Runnable lab code

# lab_1_2_3.py — REINFORCE with and without the baseline, the causality, and the break-its
# Requires: Python 3.10+, PyTorch 2.2+

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.distributions import Normal

torch.manual_seed(42); dtype, device = torch.float64, torch.device("cpu")

# --- The simplified ArmEnv (2D position, 1-step episodes) --------------
class SimpleArmEnv:
    def __init__(self):
        self.p_target = torch.tensor([0.5, 0.2], dtype=dtype)
        self.state_dim = 2  # the current position
        self.action_dim = 2  # the delta position
    def reset(self):
        self.pos = torch.randn(2, dtype=dtype) * 0.1
        return self.pos.clone()
    def step(self, a):
        self.pos = self.pos + a * 0.1
        dist = torch.norm(self.pos - self.p_target).item()
        reward = -dist
        done = True  # 1-step episode
        return self.pos.clone(), reward, done

# --- The Gaussian policy -----------------------------------------------
class GaussianPolicy(nn.Module):
    def __init__(self, state_dim, action_dim, log_std_init=-1.0):
        super().__init__()
        self.fc = nn.Linear(state_dim, 64)
        self.mean = nn.Linear(64, action_dim)
        self.log_std = nn.Parameter(torch.ones(action_dim) * log_std_init)
    def forward(self, s):
        h = torch.relu(self.fc(s))
        mu = self.mean(h)
        std = self.log_std.exp()
        return Normal(mu, std)
    def log_prob(self, s, a):
        return self.forward(s).log_prob(a).sum(-1)

# --- The value function (the baseline) --------------------------------
class ValueNet(nn.Module):
    def __init__(self, state_dim):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(state_dim, 64), nn.ReLU(),
            nn.Linear(64, 1),
        )
    def forward(self, s):
        return self.fc(s).squeeze(-1)

# --- The two estimators (for the unbiasedness and variance checks) -----
def estimator_stats(policy, env, n=20000):
    """Compare the score-function and pathwise gradients on the same problem.

    Returns the mean of each estimator over `n` samples and each estimator's
    per-sample variance. The means must agree -- both estimate the same
    gradient -- while the variances differ by a large factor, which is the
    reparameterization trick's whole point.

    The earlier version of this function returned E[log_prob * r], a scalar
    containing no gradient at all, and compared it to a hand-written "~2".
    That check could not fail for any policy.
    """
    ps = list(policy.parameters())
    def flat(gs):
        return torch.cat([g.flatten() for g in gs])
    dim = sum(p.numel() for p in ps)
    acc_s = torch.zeros(dim, dtype=dtype); sq_s = torch.zeros(dim, dtype=dtype)
    acc_p = torch.zeros(dim, dtype=dtype); sq_p = torch.zeros(dim, dtype=dtype)
    for _ in range(n):
        s = env.reset()
        # Score function: the log-likelihood's gradient times the return.
        d = policy(s); a = d.sample()
        r = -torch.norm(s + a * 0.1 - env.p_target)
        g = flat(torch.autograd.grad(d.log_prob(a).sum(), ps, retain_graph=False)) * r.detach()
        acc_s = acc_s + g; sq_s = sq_s + g * g
        # Pathwise: the return's gradient through the reparameterised sample.
        d2 = policy(s); a2 = d2.rsample()
        r2 = -torch.norm(s + a2 * 0.1 - env.p_target)
        g2 = flat(torch.autograd.grad(r2, ps))
        acc_p = acc_p + g2; sq_p = sq_p + g2 * g2
    m_s = acc_s / n; m_p = acc_p / n
    v_s = (sq_s / n - m_s * m_s).sum()
    v_p = (sq_p / n - m_p * m_p).sum()
    rel_err = (m_s - m_p).norm() / m_p.norm()
    return m_s, m_p, v_s, v_p, rel_err

# --- REINFORCE with the variants ---------------------------------------
def train_reinforce(use_baseline=False, use_causality=False, wrong_sign=False,
                    missing_log=False, non_stationary_baseline=False,
                    num_episodes=500):
    env = SimpleArmEnv()
    policy = GaussianPolicy(env.state_dim, env.action_dim).to(device, dtype)
    if use_baseline:
        value = ValueNet(env.state_dim).to(device, dtype)
        value_opt = torch.optim.Adam(value.parameters(), lr=1e-3)
    opt = torch.optim.Adam(policy.parameters(), lr=1e-3)
    returns = []
    for ep in range(num_episodes):
        # Collect a 1-step episode
        s = env.reset()
        dist = policy(s)
        a = dist.sample()
        if missing_log:
            log_prob = dist.log_prob(a).exp().sum()  # WRONG: exp(log_prob) instead of log_prob
        else:
            log_prob = dist.log_prob(a).sum()
        s_new, r, done = env.step(a)
        G = r  # 1-step: the return is the reward
        # The baseline
        if use_baseline:
            V = value(s)
            if non_stationary_baseline:
                b = r  # WRONG: the current return
            else:
                b = V
            advantage = G - b
        else:
            advantage = G
        if use_causality:
            # 1-step: the causality is the same
            pass
        # The loss
        if wrong_sign:
            loss = log_prob * advantage  # WRONG: the positive sign
        else:
            loss = -log_prob * advantage
        opt.zero_grad()
        # retain_graph: V was computed under `value`'s graph and is reused by
        # value_loss below, so the first backward must not free it.
        loss.backward(retain_graph=use_baseline)
        opt.step()
        if use_baseline:
            value_loss = F.mse_loss(V, torch.tensor(G, dtype=dtype))
            value_opt.zero_grad()
            value_loss.backward()
            value_opt.step()
        returns.append(r)
    return returns

# --- The unbiasedness and variance checks (1-step) ---------------------
env = SimpleArmEnv()
policy = GaussianPolicy(env.state_dim, env.action_dim).to(device, dtype)
m_s, m_p, v_s, v_p, rel_err = estimator_stats(policy, env, n=20000)
print(f"Unbiasedness: |score - pathwise| / |pathwise| = {rel_err:.4f}")
print(f"Variance: score {v_s:.3e} vs pathwise {v_p:.3e}, ratio {v_s / v_p:.1f}x")

# --- The training curves ----------------------------------------------
import numpy as np
def smooth(x, k=20):
    return np.convolve(x, np.ones(k)/k, mode='valid')

print("\n=== No baseline ===")
r_no_base = smooth(train_reinforce(use_baseline=False, num_episodes=500))
print(f"final smoothed return: {r_no_base[-1]:.4f}")

print("\n=== With baseline ===")
r_base = smooth(train_reinforce(use_baseline=True, num_episodes=500))
print(f"final smoothed return: {r_base[-1]:.4f}")

print("\n=== Wrong sign ===")
r_wrong = smooth(train_reinforce(wrong_sign=True, num_episodes=500))
print(f"final smoothed return: {r_wrong[-1]:.4f}")

print("\n=== Missing log ===")
r_missing = smooth(train_reinforce(missing_log=True, num_episodes=500))
print(f"final smoothed return: {r_missing[-1]:.4f}")

print("\n=== Non-stationary baseline ===")
r_ns = smooth(train_reinforce(use_baseline=True, non_stationary_baseline=True, num_episodes=500))
print(f"final smoothed return: {r_ns[-1]:.4f}")

# --- Summary ---------------------------------------------------------
print("\n--- lab 1.2.3 summary ---")
print(f"H1: estimator disagreement = {rel_err:.4f} (pred < 0.1)")
print(f"H2: score/pathwise variance ratio = {v_s / v_p:.1f}x (pred > 2x)")
print(f"H3: no-baseline {r_no_base[-1]:.4f} vs with-baseline {r_base[-1]:.4f} (pred: with-baseline higher)")
print(f"H4: wrong sign {r_wrong[-1]:.4f} vs standard {r_base[-1]:.4f}; missing log {r_missing[-1]:.4f}")

§L.7 — The cheat sheet

§L.8 — Exercises with solutions

E 1.2.3.1 — The score function (numerical). A Gaussian N(μ_θ = 2θ, σ = 1), the a = 3, the θ = 0, the G = 1. Compute the score and the gradient’s estimate.

S 1.2.3.1. The ∂ log π/∂θ = (a − μ)/σ² · ∂μ/∂θ = (3 − 0)/1 · 2 = 6. The g = 6 · 1 = 6. The lesson: the score is the (a − μ)/σ² times the ∂μ/∂θ, and the ∂μ/∂θ = 2 doubles the score.

E 1.2.3.2 — The baseline’s unbiasedness (conceptual). Show that E[∇ log π · b(s)] = 0 for any b(s).

S 1.2.3.2. E[∇ log π · b(s)] = b(s) · E[∇ log π] = b(s) · 0 = 0 (the b(s) is a constant given the s, and the E[∇ log π] = 0 by the score function’s zero-mean). The lesson: the baseline’s expectation is zero, so the subtraction preserves the unbiasedness.

E 1.2.3.3 — The causality (numerical). A 3-step episode with the rewards [1, 2, 3], the γ = 0.9. Compute the G_t’s.

S 1.2.3.3. The G_2 = 3, G_1 = 2 + 0.9·3 = 4.7, G_0 = 1 + 0.9·2 + 0.81·3 = 5.23. The lesson: the G_t decreases with the t (the fewer the remaining steps), and the causality’s G_t excludes the past rewards.

E 1.2.3.4 — The variance’s reduction (diagnostic). The no-baseline’s Var[g] = 100, the with-baseline’s Var[g] = 10. What is the sample efficiency’s improvement?

S 1.2.3.4. The variance’s ratio is 10, and the sample’s efficiency is improved by the 10× (the fewer samples for the same gradient’s accuracy). The 1/√n scaling of the standard error means the 10× variance reduction is the √10 ≈ 3.16× fewer samples for the same accuracy. The lesson: the variance’s 10× reduction is the 3.16× fewer samples, and the baseline’s practical value is the sample’s efficiency.

E 1.2.3.5 — The wrong sign (diagnostic). The return decreases over the training, and the loss is +Σ log π · G. What is the bug?

S 1.2.3.5. The positive sign makes the optimizer’s minimization descend the objective (the opposite of the intended), and the policy avoids the good actions. The fix is the negative sign (the −Σ log π · G), which ascends the objective. The lesson: the sign is the most common bug, and the decreasing return is the symptom.

E 1.2.3.6 — The missing log (diagnostic). The gradient’s norm is unstable (NaN or large), and the code uses dist.prob(a) instead of dist.log_prob(a). What is the bug?

S 1.2.3.6. The dist.prob(a) is the probability (the positive, bounded by the density’s max), and the gradient of the product prob · G is the ∇ prob · G, which is not the policy gradient (the ∇ log prob · G). The gradient’s scale is the probability’s (the small), and the update is wrong. The fix is the log_prob, which is the correct score function. The lesson: the log is essential, and the prob is the wrong function.

E 1.2.3.7 — The log-std’s learning (diagnostic, transfer probe). The policy’s log_std is a fixed parameter (the nn.Parameter with the requires_grad = False), and the policy converges but the exploration is insufficient for the new targets. What is the fix?

S 1.2.3.7. The learnable log_std (the requires_grad = True) adapts the exploration to the task (the large early, the small late, the adapted for the new targets). The fixed std is the simplification that underperforms, and the learned std is the standard. The lesson: the log-std is the learnable exploration, and the fixed std is the missed opportunity.

E 1.2.3.8 — The squashing’s Jacobian (constructive, transfer probe). A tanh-squashed Gaussian a = τ_max tanh(z). Write the log-likelihood and the Jacobian’s term.

S 1.2.3.8.

log π(a | s) = log p(z | s) - log |da/dz|
where z = atanh(a / τ_max), and da/dz = τ_max (1 - tanh^2(z)) = τ_max (1 - (a/τ_max)^2)
so log |da/dz| = log(τ_max) + log(1 - (a/τ_max)^2)

The log-likelihood is the log p(z|s) − log |da/dz|, and the Jacobian’s term is the log |da/dz|. The lesson: the squashing’s Jacobian is the required correction, and its omission biases the log-likelihood.

§L.9 — Exit test

You are ready for Lesson 1.2.4 (PPO) when you can, without notes:

1. State the policy gradient’s objective and the policy gradient theorem.
2. State the score function and its zero-mean.
3. Implement REINFORCE (the log π, the return, the loss, the optimizer).
4. State the baseline’s unbiasedness and its variance reduction.
5. State the causality (G_t vs. the G_0) and its unbiasedness.
6. State the advantage A(s, a) = G − V(s) and the optimal baseline.
7. State the Gaussian policy’s parameterization (the μ_θ(s) and the log σ_θ(s)).
8. State the tanh squashing’s Jacobian and its necessity.
9. Diagnose the wrong sign (the decreasing return), the missing log (the unstable gradient), the no baseline (the slow learning), and the non-stationary baseline (the bias).
10. State the variance’s dependence on the σ and the return’s magnitude.
11. State the sample efficiency’s improvement from the variance reduction (√n scaling).
12. State the loss’s form (−Σ log π · A) and the optimizer’s ascent on the objective.

If any item is not answerable without notes, return to the section that introduced it: 1 to §L.1, 2 to §L.1, 3 to §L.2, 4 to §L.3, 5 to §L.3, 6 to §L.3.1, 7 to §L.4, 8 to §L.4, 9 to §L.5 and the individual E’s, 10 to §L.4.2, 11 to E 1.2.3.4, 12 to §L.2.

End of Lesson 1.2.3. The learner has now completed the policy gradient foundation: the score function, the REINFORCE estimator, the baseline and the causality’s variance reduction, and the Gaussian policy’s practical issues. The next lesson, 1.2.4 — PPO: Trust Regions, Clipping, and the Modern Policy Gradient, uses the policy gradient as its core, and adds the importance sampling’s ratio, the clipping’s trust region, and the GAE’s advantage estimator — the algorithm behind the large-scale robot learning’s successes.