Policy Gradient Methods

Why this family exists

In an MDP, you ultimately care about behavior: which actions you take in each state. A policy is the object that produces that behavior. In policy gradient methods, the policy is parameterized and differentiable, so we can change it continuously and aim those changes toward higher return.

Instead of learning a value function first and deriving a policy from it, we optimize a policy directly:

• •Policy: $\pi_\theta(a\mid s)$, a distribution over actions given a state.
• •Objective: expected discounted return under that policy.

A standard episodic objective is

where a trajectory (rollout) is

and the trajectory distribution is induced by the environment dynamics and the policy:

The key point: θ controls the probability of your actions, and that changes which states you visit and which rewards you obtain.

What “differentiable policy” means in practice

Typically, θ parameterizes a neural network that outputs either:

• •Discrete actions: logits → softmax → categorical distribution.
• •Continuous actions: mean (and maybe log-std) of a Gaussian.

Example (discrete):

Example (continuous, diagonal Gaussian):

We then perform stochastic gradient ascent on $J(\theta)$:

The conceptual leap: “credit assignment” through log-probability

In supervised learning, you get a target label. In RL, you get rewards after decisions. The policy gradient trick ties the final outcome back to earlier action probabilities via

• •how probable the action was ($\log \pi_\theta(a\mid s)$), and
• •how good the outcome was (returns or advantages).

Intuition to hold onto:

If an action led to better-than-expected outcomes, increase its probability in that state. If it led to worse-than-expected outcomes, decrease it.

Policy gradient methods operationalize that intuition with a precise gradient estimator.

Visualization plan (interactive canvas)

To make this idea tangible, your canvas can show a tiny 2-state MDP and a 2-action policy.

Canvas panel A: “Policy sliders”

• •Let θ be a single scalar controlling a Bernoulli policy:
• •$\pi_\theta(a=1\mid s)=\sigma(\theta_s)$ for each state s.
• •Show two sliders (θ for state 0 and state 1).
• •As the learner drags θ, animate action probabilities (bar chart) shifting.

Canvas panel B: “Trajectory outcomes”

• •Sample rollouts under the current policy.
• •Show the returns $G_t$ next to each time step.

Canvas panel C: “Gradient arrows”

• •For each visited (sₜ,aₜ), display
• •$\nabla_\theta \log \pi_\theta(a_t\mid s_t)$ as an arrow on θ.
• •Multiply the arrow by an advantage estimate and show the scaled update.

This directly externalizes the algebra: gradient = (score) × (signal).

Why we need a special gradient identity

We want $\nabla_\theta J(\theta)$, but $J$ is an expectation over trajectories whose distribution depends on θ. Differentiating “through” sampling is awkward because trajectories are discrete random objects.

The score-function (a.k.a. log-derivative) trick gives a way to move the gradient inside an expectation without differentiating the environment dynamics.

The identity to remember:

This works whenever $p_\theta(x)$ is differentiable in θ and $f(x)$ is integrable.

Derivation (showing the work)

Start from:

where $R(\tau)=\sum_{t=0}^{T-1}\gamma^t r_t$.

Differentiate:

Use $\nabla P = P\nabla \log P$:

Recognize the expectation:

Now expand $\log P_\theta(\tau)$. From

take logs:

Differentiate w.r.t. θ. Only the policy terms depend on θ (environment dynamics are fixed):

So:

This already yields an unbiased estimator: sample a trajectory, compute $R(\tau)$, and push the log-prob gradients in the direction of $R(\tau)$.

Reward-to-go (better credit assignment)

Using the same return $R(\tau)$ for every time step credits early and late actions equally, even though late actions cannot affect early rewards.

A standard improvement is the reward-to-go:

Then the estimator becomes:

This is still unbiased, but typically has lower variance.

REINFORCE algorithm (Monte Carlo policy gradient)

At a high level:

1. 1)Collect trajectories using current $\pi_\theta$.
2. 2)For each time step, compute $G_t$.
3. 3)Update θ by ascending the sampled gradient.

A common minibatch form:

The geometry of the update (what the gradient does)

For a softmax policy, you can interpret $\nabla_\theta\log\pi$ as:

• •increasing parameters that make the taken action more likely,
• •decreasing parameters that make competing actions likely.

Then multiplying by $G_t$ decides direction:

• •If $G_t$ is large (good), increase probability of those actions.
• •If $G_t$ is small/negative (bad), decrease probability.

Visualization: “Score × Return” microscope

Add a per-time-step breakdown:

• •Display $\log\pi_\theta(a_t\mid s_t)$.
• •Show $\nabla_\theta\log\pi_\theta(a_t\mid s_t)$ as a vector arrow (or scalar bar if θ is 1D).
• •Multiply by $G_t$ and animate the resulting parameter step.

Learners should see that the policy gradient update is not magic—it’s a weighted push on log-probability.

Why REINFORCE struggles

REINFORCE is unbiased, but its Monte Carlo returns can have enormous variance:

• •randomness in the environment,
• •randomness in the policy,
• •long horizons with discounting.

High variance means you need many trajectories (or tiny learning rates) to make stable progress.

The central theme of modern policy gradients is:

Keep the estimator (approximately) unbiased while reducing variance.

Baselines: subtract something that doesn’t change the expectation

Key fact: for any function $b(s_t)$ that does not depend on $a_t$,

So we can subtract $b(s_t)$ inside the gradient estimator without changing its expectation:

This can drastically reduce variance when $b(s_t)$ approximates the “typical” return from $s_t$.

The most useful baseline: the value function

Choose $b(s_t)=V^\pi(s_t)$, where

Then $G_t - V^\pi(s_t)$ is an estimate of the advantage:

Advantage answers a very specific question:

Was this action better or worse than my policy’s average behavior in this state?

This is exactly the signal you want for improving a stochastic policy.

Actor-critic: two function approximators with different jobs

Actor-critic methods maintain:

• •Actor: the policy $\pi_\theta(a\mid s)$ (parameters θ)
• •Critic: a value function estimate $V_\phi(s)$ or action-value $Q_\phi(s,a)$ (parameters φ)

The critic provides a low-variance learning signal; the actor uses it to update the policy.

A common actor update uses an estimated advantage $\widehat{A}_t$:

Bootstrapping: trading bias for lower variance

Instead of Monte Carlo $G_t$, we can use TD-style targets.

One-step TD error (for value critic):

A simple advantage estimate is $\widehat{A}_t=\delta_t$.

This introduces some bias (because $V_\phi$ is approximate), but variance often drops dramatically and learning becomes faster.

Generalized Advantage Estimation (GAE)

GAE blends multi-step TD errors with an additional parameter λ that controls bias/variance.

Define TD residuals $\delta_t$ as above, then

• •λ = 0: very low variance, more bias (pure 1-step TD)
• •λ → 1: less bias, more variance (approaches Monte Carlo advantage)

Critic learning objective

If the critic is a value function $V_\phi(s)$, a typical squared-error loss is:

where $\widehat{V}_t$ might be:

• •Monte Carlo return $G_t$
• •TD target $r_t + \gamma V_\phi(s_{t+1})$
• •λ-return (related to GAE)

Comparison table (variance and bias intuition)

Visualization: “Variance comparison panel”

To address the visualization weakness explicitly, build a panel that runs the same fixed policy for many rollouts and shows the distribution of gradient estimates.

Panel design

• •Fix θ for a small MDP.
• •Sample K trajectories (e.g., K=200) and compute gradient estimates for a chosen parameter component.
• •Plot three histograms (or violin plots):

1. 1)REINFORCE with $G_t$
2. 2)Baseline with $G_t - V(s_t)$
3. 3)GAE with chosen λ

What learners should observe

• •REINFORCE histogram is wide (noisy updates).
• •Baseline narrows distribution around the same mean.
• •GAE can further narrow, depending on λ.

Add a slider for λ (0→1) and animate the histogram tightening/loosening. That makes bias/variance tradeoff visible, not just stated.

Why actor-critic is the stepping stone to modern algorithms

Many practical deep RL systems used today (including those in RLHF pipelines) rely on three ideas:

1. 1)Policy gradient objective (optimize π directly)
2. 2)Advantage-based updates (baseline/value function)
3. 3)Stabilization constraints (trust regions, clipping, KL penalties)

This node focuses on (1) and (2), which are foundational for PPO and RLHF.

A canonical on-policy actor-critic training loop

A common structure (simplified):

1. 1)Collect T steps of experience with current $\pi_\theta$.
2. 2)Fit $V_\phi$ to predict returns (or λ-returns).
3. 3)Compute advantages $\widehat{A}_t$ (often GAE).
4. 4)Update actor by maximizing:

Equivalently, minimize $-\mathcal{L}_{\text{actor}}$.

5. 5)(Often) add an entropy bonus to encourage exploration:

Where PPO fits (preview-level connection)

PPO is still a policy gradient method, but it modifies the objective so updates do not change the policy too abruptly.

A typical PPO objective uses the probability ratio

and then clips it to avoid overly large updates. Notice how everything here assumes you already understand:

• •$\pi_\theta(a\mid s)$ as a differentiable stochastic policy
• •advantage estimates $\widehat{A}_t$
• •log-probability gradients (since ratios become differences of logs)

That’s exactly why mastering policy gradients unlocks PPO and thus RLHF.

RLHF connection (conceptual)

In RLHF, you often:

• •Train a reward model from human preferences.
• •Use an RL algorithm (commonly PPO) to optimize the language model policy to maximize that learned reward, subject to constraints (e.g., KL to a reference model).

Even if the “environment” is text generation and the “reward” comes from a reward model, the policy gradient core remains:

• •actions are tokens,
• •states are partial sequences,
• •$\nabla_\theta \log \pi_\theta$ is computed by backprop through the transformer,
• •advantage estimation stabilizes learning.

Visualization: “Tiny MDP → PPO-like constraint” (bridge)

Add a simple toggle in the canvas:

• •Vanilla PG mode: $\Delta\theta \propto \nabla\log\pi \cdot \widehat{A}$
• •Constrained mode (toy PPO): if the KL between old and new policy exceeds a threshold, scale down the update.

Even if you don’t implement full clipping, just showing “unconstrained step” vs “KL-limited step” prepares learners for PPO’s motivation.

Summary of what you should be able to do after this node

• •Write down the objective $J(\theta)$ and the trajectory distribution.
• •Derive/recognize the score-function policy gradient estimator.
• •Implement REINFORCE with reward-to-go.
• •Add a baseline and explain why it doesn’t bias the gradient.
• •Build an actor-critic update with TD/GAE advantages.

Those are the conceptual and mathematical prerequisites for modern on-policy deep RL.