Policy Gradient Theorem & Algorithms

Read the theorem as three questions

Let πθ(a | s) be a differentiable policy and let J(θ) be its expected discounted return. One common form of the theorem is:

∇θ J(θ) = Eπ[∇θ log πθ(a | s) · Qπ(s, a)]

The equivalent sampled-trajectory form replaces Qπ(s, a) with a return or advantage estimate. Discount and state-occupancy conventions can add a constant scaling factor, but they do not change the update direction.

The theorem does not say that every positive reward should increase an action forever. With an advantage signal, the relevant question is whether the outcome was better than the policy expected in that state.

Explore one exact policy update

Consider a policy with two actions and logits z_chosen and z_other. If the chosen action currently has probability p, its score derivatives are 1 − p for the chosen logit and −p for the other logit. A gradient-ascent step therefore changes their difference by αA, giving:

p_new = sigmoid(logit(p) + αA)

Use the explorer to test three cases: positive advantage, negative advantage, and a nearly saturated action probability. Watch both update direction and magnitude.

Derive the estimator without differentiating the environment

The derivation starts from complete trajectories τ = (s_0, a_0, r_0, …) and progressively removes terms that the policy does not control.

1. 1

   Return

   Write the objective

   Express expected return as J(θ) = Eτ~pθ(τ)[R(τ)], where trajectory probability includes the initial state, policy actions, and environment transitions.

2. 2

   Score

   Apply the log derivative

   Use ∇pθ(τ) = pθ(τ)∇log pθ(τ) to turn the gradient of an expectation into an expectation of a gradient.

3. 3

   Model-free

   Drop environment terms

   Initial-state and transition probabilities do not depend on θ, so only Σ_t log πθ(a_t | s_t) remains inside the gradient.

4. 4

   Causality

   Assign future credit

   Rewards before action a_t cannot be caused by that action. Weight its score only by return from time t onward, producing G_t.

The score-function identity is the key step:

∇θ E_{x~pθ}[f(x)] = E_{x~pθ}[f(x)∇θ log pθ(x)]

It allows samples to estimate a gradient even when f(x) is discontinuous. The policy distribution must still be differentiable with respect to θ wherever sampled actions have support.

Separate unbiasedness from variance

REINFORCE can use the sampled return G_t directly:

ĝ_t = ∇θ log πθ(a_t | s_t) · G_t

Under on-policy sampling and the usual regularity assumptions, this is an unbiased gradient estimator. It can still be impractical because each return contains variation from later actions, stochastic transitions, and reward noise.

For any action-independent baseline b(s_t):

E_{a~π}[∇θ log πθ(a | s) · b(s)] = 0

Therefore G_t − b(s_t) has the same expected policy gradient. Choosing b(s_t) ≈ Vπ(s_t) often reduces variance because the weight becomes an advantage: how much better the action performed than the expected outcome from that state.

Compare trajectory estimators

The lab keeps one rollout scenario fixed while you change the estimator, critic quality, rollout horizon, and GAE parameter λ. Inspect all three methods together; a baseline and bootstrapping solve different problems.

Choose an algorithm by its estimator and update constraint

The algorithms below share the policy-gradient direction. They differ in how they estimate action quality, when they update, and how far they permit the policy to move.

REINFORCE: full-return Monte Carlo

REINFORCE waits for sampled rewards, computes each reward-to-go G_t, and minimizes −log πθ(a_t | s_t)G_t.

• Use it for: small episodic problems, theorem verification, and estimator tests.
• Strength: simple on-policy estimator with no critic approximation.
• Cost: high variance, delayed updates, and poor sample efficiency on long trajectories.

Actor-Critic: learn the baseline

An actor chooses actions while a critic estimates Vπ(s). A one-step TD residual,

δ_t = r_t + γV(s_{t+1}) − V(s_t),

can serve as an advantage estimate. Bootstrapping lowers variance and permits short rollouts, but critic error can bias the estimate.

• Use it for: continuing tasks or online updates where complete episodes are expensive.
• Strength: more frequent, lower-variance updates than raw Monte Carlo.
• Cost: actor and critic learning can destabilize each other.

A2C with GAE: blend many TD residuals

Generalized Advantage Estimation computes:

Â_t^GAE = Σ_{k=0}^{T-t-1}(γλ)^k δ_{t+k}

Smaller λ relies more on the critic and usually lowers variance; larger λ approaches a Monte Carlo-style estimate with less bootstrap bias and more trajectory variance.

PPO: constrain reuse of on-policy data

PPO stores each action's probability under the behavior policy and compares it with the current policy using r_t(θ) = πθ(a_t | s_t) / πold(a_t | s_t). Its clipped surrogate is:

Lclip = E[min(r_tÂ_t, clip(r_t, 1 − ε, 1 + ε)Â_t)]

Clipping discourages some overly large probability-ratio changes while allowing several minibatch passes over a rollout. It does not guarantee monotonic improvement and does not replace KL, entropy, reward, or safety monitoring.

Build a training batch with explicit ownership

1. 1

   Collect

   Freeze a behavior version

   Roll out one identifiable policy version and store states, actions, rewards, termination flags, values, and old action log probabilities.

2. 2

   Estimate

   Construct targets

   Compute returns and advantages with correct terminal versus time-limit handling. Normalize advantages over an intentional batch boundary.

3. 3

   Update

   Optimize bounded losses

   Combine the actor objective, critic loss, and entropy bonus. Clip gradient norm and stop early when policy divergence exceeds the release threshold.

4. 4

   Gate

   Evaluate before recollecting

   Measure reward distributions, constraints, action entropy, critic error, and behavior slices before the new policy becomes the data collector.

The batch is a contract between rollout and optimization. Preserve these invariants:

• Store log πold(a_t | s_t) at collection time; recomputing it after policy updates corrupts PPO ratios.
• Distinguish true terminal states from time-limit truncation; bootstrap through a truncation when the task continues.
• Mask padded sequence positions before computing returns, advantages, entropy, or losses.
• Stop gradients through actor weights such as Â_t; train the critic with its own explicit target.
• Version reward code, observation transforms, action masks, policy weights, and rollout data together.

Diagnose failures from their signatures

The policy collapses to one action

• Action entropy falls quickly while reward is flat or unstable.
• Likely causes include an excessive learning rate, weak entropy regularization, reward scale spikes, or an incorrect action mask.
• Reduce update size, inspect per-action logits, and replay the exact batch that triggered collapse.

The critic looks accurate but the policy does not improve

• Value loss can be low while advantage signs are wrong on important states.
• Check target leakage, terminal masks, reward timing, and whether actor gradients were accidentally detached at the log-probability term.
• Evaluate explained variance and calibration by state slice, not value loss alone.

PPO reports clipping on most samples

• High clip fraction and KL divergence mean the optimizer is trying to move far from the behavior policy.
• Reduce learning rate or epochs, increase batch size, or stop the epoch early.
• Do not interpret clipping as proof that the final policy is safe or close everywhere.

Reward rises while user outcomes worsen

• The policy may be exploiting a proxy, simulator flaw, or evaluator blind spot.
• Maintain independent outcome metrics and adversarial behavior slices.
• Pause optimization when the reward model, environment, or policy leaves its validated operating range.

Operate policy-gradient training as a feedback system

Track enough information to separate estimator noise from policy movement and reward failure.

• Policy movement
  • Approximate KL divergence between behavior and updated policies
  • Probability ratio distribution and PPO clip fraction
  • Gradient norm, parameter-step norm, and action entropy
• Estimator health
  • Advantage mean, standard deviation, and tail percentiles
  • Critic loss, explained variance, and value calibration
  • Return variance by trajectory length and behavior slice
• Outcome and safety
  • Reward mean plus lower-tail and constraint metrics
  • Independent task-success and human-outcome measures
  • Action frequencies, invalid-action rate, and regression slices
• Reproducibility
  • Policy and critic checkpoints, optimizer state, random seeds, and environment version
  • Reward definition, rollout manifest, advantage settings, and code revision
  • A tested rollback path to the last accepted behavior policy

Use canary rollouts or shadow evaluation before broad deployment. A statistically higher mean reward is not a release decision when tail risk, constraints, or protected slices regress.

Apply the theorem where actions have delayed consequences

Language-model alignment

In RLHF-style training, token generation forms a long action sequence. A learned reward and KL penalty supply the optimization signal, while a critic or return estimator assigns credit. Continue to RLHF for the broader preference-data and reward-model pipeline.

Robotics and control

Continuous policies often parameterize distributions over motor commands. Policy gradients avoid enumerating actions, but safe exploration, simulator mismatch, and real-world sample cost dominate the system design.

Ranking and recommendation

A stochastic ranking policy can optimize delayed engagement or satisfaction signals. Off-policy logging, exposure bias, and reward-proxy failures require careful evaluation before an on-policy method is appropriate.

Across all three domains, the theorem supplies a gradient direction. It does not supply a trustworthy reward, sufficient exploration, a stable critic, or a safe release process; those remain system responsibilities.