M6-T2 — Policy Gradients and REINFORCE

Learning objectives

After this lesson, you should be able to:

• Define the objective J(θ) and derive the policy gradient estimator step by step with the log-derivative trick.
• State the policy gradient theorem accurately and write REINFORCE for a Gaussian policy.
• Check a gradient estimate against finite differences on a one-parameter policy.
• Explain, with measured numbers, why REINFORCE has high variance and why it often fails to find the goal on Scout’s task within a modest budget.

The intuition

Suppose Scout’s velocity command is a Gaussian centred on θ, and the reward is highest when the command equals 1. Scout tries four commands. One of them, 0.9, got a reward very close to the best; another, −0.1, got a poor one. Where should the mean move? Towards the actions that did well, and away from the ones that did badly, weighted by how well and how badly. That is the whole idea.

The score of a sampled action, the derivative of its log density with respect to the mean, says which way the mean has to move to make that action more likely. For a Gaussian it is (a − θ)/σ²: an action above the mean has a positive score, and moving the mean up makes it more likely. Multiply each action’s score by the reward it earned and average, and the sum points uphill on average.

The mathematics

The objective

A trajectory is τ = (s₀, a₀, r₀, s₁, a₁, r₁, …). A policy π_θ with parameters θ makes some trajectories more likely than others. The quantity to maximise is the expected return:

The initial-state distribution ρ and the dynamics P do not depend on θ; only the policy factors do.

The log-derivative trick, step by step

1. Write the gradient of the expectation: ∇J = ∇ Σ_τ p_θ(τ) R(τ) = Σ_τ ∇p_θ(τ) R(τ), because R(τ) does not depend on θ. (For continuous variables replace the sum by an integral.)
2. Use ∇p = p ∇log p, which is just the chain rule for the logarithm: ∇p_θ(τ) = p_θ(τ) · ∇log p_θ(τ). Then ∇J = Σ_τ p_θ(τ) ∇log p_θ(τ) R(τ) = E[ ∇log p_θ(τ) · R(τ) ]. The gradient has become an expectation we can estimate by sampling.
3. Take the log of the product: log p_θ(τ) = log ρ(s₀) + Σ_t log π_θ(a_t | s_t) + Σ_t log P(s_t+1 | s_t, a_t). The first and third parts do not contain θ, so ∇log p_θ(τ) = Σ_t ∇log π_θ(a_t | s_t). The robot’s dynamics, which we may not know, drop out.
4. So ∇J(θ) = E[ ( Σ_t ∇log π_θ(a_t | s_t) ) · R(τ) ].
5. Causality: the reward r_t′ is not influenced by an action taken later, at t > t′. For such pairs E[ ∇log π(a_t | s_t) · r_t′ ] = 0 (the same argument as the baseline proof in lesson M6-T3), so those terms can be dropped. Each action is credited only with the rewards that follow it.

The policy gradient theorem

The statement in Sutton and Barto (section 13.2, the episodic case) says that ∇J(θ) is proportional to Σ_s μ(s) Σ_a q_π(s, a) ∇π(a | s; θ), where μ(s) is the on-policy distribution of states. What it buys is that the gradient has no term for how the state distribution changes with θ, even though it does change. The estimator above is its sampled form with G_t in place of q_π(s_t, a_t). One honest caveat: with discounting the exact gradient also carries a factor γ^t on the t-th term. Almost every implementation, including the trainer in this module, drops it and uses the undiscounted state distribution. The resulting direction is not the exact gradient of the discounted objective but is the common practice.

Gradients for two policy families

In the trainer, one backward pass through the actor network per step with output signal −advantage × z/σ gives the gradient for the whole mean network; the code is the chain rule applied to the formulas above.

Why continuous actions are natural here

Value-based control needs max_a Q(s, a); with a continuous velocity this is itself an optimisation at every step. A policy-gradient method needs only to sample an action and differentiate the log density of what was sampled. Both are closed-form for a Gaussian. The price: a Gaussian has one peak, so it cannot hold two separate good actions (left of the pillar and right of the pillar), and noise samples that fall outside [−1, 1] are clipped by the world while their density is scored as if they were not.

A worked example

Every number below was computed (by closed form, by numerical integration, by simulation, or with RLCore.trainPolicyGradient).

1. A hand-checkable gradient. The policy is a ~ Normal(θ, 0.5²) with θ = 0.2 and the one-step reward is r(a) = −(a − 1)². Expected reward: J(θ) = E[r] = −((θ − 1)² + σ²) = −(0.64 + 0.25) = −0.89. Its derivative is dJ/dθ = −2(θ − 1) = 1.6. Finite-difference check: with h = 10⁻⁴, (J(θ + h) − J(θ − h)) / 2h = 1.6000. A numerical integral of r(a) × density over a, with the same central difference, gives 1.6000 as well.

The REINFORCE estimator uses the score (a − θ)/σ² = (a − 0.2)/0.25 and four sampled actions:

The estimate is 0.243: the right sign, about one seventh of the true size. Averaging over 20,000 independent batches of four samples, the estimator has mean 1.619 (unbiased, up to simulation noise) and a standard deviation of 1.93 per batch. A single batch of four is little better than a coin flip about the size of the step.

2. A baseline. Subtract a constant b from every reward before multiplying by the score. With the batch mean of the four rewards, b = −0.4675, the same four samples give 0.664. With b equal to the true J = −0.89, the 20,000-batch standard deviation drops from 1.93 to 1.39 and the mean stays 1.597. Subtracting the batch’s own mean, though, gives an average of 1.197 instead of 1.6, which is 3/4 of the truth: using the same samples to compute b and the gradient correlates them, and the bias is a factor (N − 1)/N for N = 4. Lesson M6-T3 proves that a baseline fixed in advance is unbiased and why the batch mean is not.

3. REINFORCE on Scout. trainPolicyGradient({algo: 'reinforce'}): returns-to-go (γ = 0.99), no critic, one gradient step per batch, learning rate 3e-3, Adam. Like every algorithm in this trainer it normalises the weights of a batch to mean 0 and standard deviation 1 before the update (lesson M6-T3), so it carries a batch-mean baseline of the kind criticised in example 2; with thousands of steps per batch that bias should be small, but the variance problem remains. Mean-policy success is measured on 50 fresh starts. Ten seeds (five for the longest run):

On the default budget REINFORCE learns the task in 3 of 10 seeds, and seed 1 scores 0: its mean return stays at −4.0, which is 80 steps × −0.05 and means Scout never reaches the goal and never crashes. Five times more episodes lifts the mean success to 0.38, not to 1.0: the seeds that fail at 720 episodes mostly still fail at 3,600 and at 7,200 (seeds 1, 2, 3, 5 never leave −4.0 in the 600-iteration run). So the budget is not the whole story.

What the logs show. An untrained stochastic policy reached the goal in 0 of 3,000 simulated episodes (95 % Wilson upper limit 0.13 %), so most batches never see the +10 bonus. All the learning signal is then the −0.05 per step and the −5 crash. In the failing seeds the trained mean policy drives Scout into a wall and keeps it there: at the end of 8 evaluation episodes the position was (0, 2.1) in every episode of seed 1 and (0, 6) in every episode of seed 3, all 80 steps. This is the sparse-reward exploration problem of lesson M5-T3 in a continuous setting. I did not test whether an entropy bonus, an easier start region or a shaped reward cures it in this trainer; those are the usual remedies and the labs of this module let you try some of them.

4. The variance of the gradient, measured. For the actor’s 226 parameters, 300 batches of 12 episodes were drawn from one fixed policy, and the REINFORCE gradient was computed from each with three choices of what multiplies the score (a separate measurement script, not the trainer, so that the unnormalised estimator could be studied). The table gives the square root of the total variance across batches (summed over the 226 parameters), and the mean cosine between a batch’s gradient and the best available direction, the average of the 300 learned-baseline gradients at the 25-iteration checkpoint.

A fitted baseline cuts the noise by a factor 2.7 for the untrained actor and 3.8 for the checkpoint. A random direction in 226 dimensions has a cosine of about ±1/√226 = ±0.07 with any fixed direction, so a raw REINFORCE batch (0.02) is indistinguishable from random in direction. The critic baseline raises that to 0.29: clearly informative but still mostly noise. Caveats: one policy per column, 300 batches, so the reference direction is itself noisy; for the untrained actor its error is as large as its length, so only the noise column is reported there. The critic used as baseline was fitted separately, on 200 independent episodes (R² of 0.45 on its fit for the untrained actor and 0.74 at the checkpoint), so it is independent of the batch and does not bias the gradient.

Engineering implications and limitations

Why REINFORCE is rarely used as it stands

It is the cleanest idea in the module, and it has three practical problems visible above: the estimator is very noisy, which means small learning rates and many episodes; it uses returns of whole episodes, so the credit for a good action is spread across every step that follows; and it needs the goal to be reached by chance before anything improves. The rest of the module is the repair kit: a baseline and advantages (lesson M6-T3), a learned critic (lesson M6-T4), and a limit on how far one update may move the policy (lesson M6-T5).

Implementation notes for the Gaussian actor

• Store the sampled action with the batch, not just the observation: the gradient needs the exact a that was taken, before clipping.
• Use returns-to-go, not the whole-episode return, for each step (the causality step above); this removes some noise for free.
• Compute the gradient of the sum or the mean of log π × weight, and decide once whether episodes are weighted equally or steps are; with episodes of different lengths the choice matters.
• Learn log σ, not σ, so that it stays positive; bound it (this trainer keeps σ between 0.05 and 2).
• Check the gradient code on a tiny case first, as the lesson did: finite differences against the analytic score on a one-parameter policy catch sign and factor errors.

Report several seeds

A single REINFORCE run on this task would have told a very different story for seed 4 (success 0.96) and seed 1 (success 0.00). Report the number of seeds that learned the task and a mean with an interval (lesson M1-T5; RLCore.meanCI gives a Student-t interval), never one curve.

Lab M6-L1 (compute-returns-and-advantages) practises the returns-to-go calculation with correct handling of episode boundaries. Lab M6-L3 (tune-a-ppo-run) lets you set the algorithm to REINFORCE and compare it with A2C and PPO under the same budget.

Common misconceptions

Engineering takeaways

1. J(θ) is the expected return, and ∇J = E[ Σ_t ∇log π(a_t | s_t) · G_t ]: the log-derivative trick turns a gradient of an expectation into an expectation of a gradient, with the dynamics gone.
2. REINFORCE averages score × return-to-go over sampled episodes. In the one-parameter check the true gradient is 1.6 (confirmed by finite differences) and a four-sample estimate was 0.243.
3. The estimator is unbiased but noisy: at 12 episodes per batch the 226-parameter gradient of an untrained Scout actor has a mean cosine of 0.02 with the reference direction, no better than chance.
4. On Scout’s task REINFORCE learned the goal in only 3 of 10 seeds in 720 episodes, and 5 to 10 times more episodes helped a little: when the reward is not seen, there is no gradient to follow.
5. A Gaussian policy makes continuous actions easy to sample and differentiate, at the price of a single peak and clipped noise.

References and further reading

Williams introduced REINFORCE; Sutton and Barto give the policy gradient theorem and the textbook treatment. The theorem’s original statement for function approximation is the Sutton, McAllester, Singh and Mansour paper.

• R. S. Sutton and A. G. Barto, Reinforcement Learning: An Introduction, 2nd edition, MIT Press, 2018 (chapter 13: policy gradient methods, including the policy gradient theorem and REINFORCE).
• R. J. Williams, “Simple statistical gradient-following algorithms for connectionist reinforcement learning”, Machine Learning 8(3–4), 229–256, 1992.
• R. S. Sutton, D. McAllester, S. Singh and Y. Mansour, “Policy gradient methods for reinforcement learning with function approximation”, Advances in Neural Information Processing Systems 12, 2000 (NIPS 1999).

Knowledge check

Three conceptual questions. Write an answer, then reveal the explanation. Your answers stay in your browser.

What’s next

Next: the 226-parameter gradient is mostly noise. Lesson M6-T3 shows that subtracting a baseline leaves the gradient unbiased, defines the advantage, and builds generalised advantage estimation.

Related content

Nothing below is required to finish this lesson.

• Before
  M6-T1: Function Approximation and Deep RL

  Explain why tables stop working, what changes when a network represents a value or a policy, and what can go wrong.

• Lab
  M6-L1: Compute Returns and Advantages

  The advantage estimates are wrong at episode boundaries: write discounted returns and GAE that handle terminal states.