Policy Gradient Methods

Jr Codex ML Notes

Level: Advanced Prerequisites: Chapter 1: RL Recap & Where It Fits in ML Time to complete: ~25 minutes

Learning a Policy Directly
Representing a Policy
The Policy Gradient Idea
REINFORCE — the Foundational Algorithm
REINFORCE Step by Step
The High-Variance Problem
Baselines — a Partial Fix
Summary & Next Steps

1. Learning a Policy Directly

Chapter 1 introduced the distinction: rather than learning Q-values and deriving a policy, policy gradient methods parameterize the policy directly and use gradient ascent (Module 3, Chapter 1's gradient descent, but maximizing instead of minimizing) to improve it.

2. Representing a Policy

import numpy as np
 
def softmax_policy(state_features, weights):
    """
    A simple LINEAR policy: computes a score for each action,
    then converts scores into a PROBABILITY DISTRIBUTION via softmax
    (directly extending Module 4, Ch.1's sigmoid to MULTIPLE classes/actions).
    """
    scores = state_features @ weights      # one score per possible action
    exp_scores = np.exp(scores - np.max(scores))      # numerical stability
    return exp_scores / np.sum(exp_scores)      # probabilities, summing to 1
 
# Example: 2 possible actions, 3 state features
weights = np.random.randn(3, 2)      # 3 features -> 2 action scores
state = np.array([1.0, 0.5, -0.2])
 
action_probabilities = softmax_policy(state, weights)
print(action_probabilities)      # e.g. [0.35, 0.65] — probability of EACH action

The Policy, Formally
─────────────────────────────────────────
  π(a | s, θ) = probability of taking action a in state s,
                  GIVEN the policy's parameters θ (weights)

  This is EXACTLY Module 4, Ch.1's logistic regression idea
  (a linear combination, squeezed into valid probabilities via
  softmax — the multi-class generalization of sigmoid), but now
  θ is trained to maximize EXPECTED REWARD instead of
  classification accuracy.
─────────────────────────────────────────

3. The Policy Gradient Idea

The core question: how should we adjust θ (the policy's parameters) to make good actions more likely and bad actions less likely?

The Policy Gradient Theorem, Intuitively
─────────────────────────────────────────
  ∇θ J(θ) ≈ E[ ∇θ log π(a|s,θ) × G ]

  "Push the policy's parameters in whatever direction
  INCREASES the log-probability of the actions taken,
  WEIGHTED by how good the outcome (G, the total reward
  received) actually turned out to be."

  If G was HIGH (a good outcome): STRENGTHEN the actions taken
  If G was LOW (a bad outcome): WEAKEN the actions taken
─────────────────────────────────────────

Directly Connects to Module 3, Chapter 1's Gradient Descent
─────────────────────────────────────────
  This is the SAME gradient-based optimization idea from linear
  regression's training — "compute a gradient, take a step in
  that direction" — just applied to POLICY PARAMETERS, with a
  loss function based on REWARD instead of squared error.
─────────────────────────────────────────

4. REINFORCE — the Foundational Algorithm

The simplest, foundational policy gradient algorithm — directly implementing the intuition from Section 3.

The REINFORCE Algorithm
─────────────────────────────────────────
  1. Run a full EPISODE using the CURRENT policy, recording
     every (state, action, reward) along the way
  2. Compute the TOTAL discounted return G from each point
     in the episode onward (directly reusing the AI Notes'
     Module 5, Ch.4 discounted return calculation)
  3. For EACH (state, action) pair in the episode, push the
     policy's parameters to make that action MORE likely if
     G was HIGH, LESS likely if G was LOW
  4. REPEAT across many episodes
─────────────────────────────────────────

5. REINFORCE Step by Step

import numpy as np
 
def compute_returns(rewards, gamma=0.99):
    """Discounted returns from each timestep onward — reusing the AI
    Notes' Module 5, Ch.4 discounting concept, applied per-timestep."""
    returns = np.zeros(len(rewards))
    running_return = 0
    for t in reversed(range(len(rewards))):
        running_return = rewards[t] + gamma * running_return
        returns[t] = running_return
    return returns
 
def reinforce_update(states, actions, rewards, weights, learning_rate=0.01, gamma=0.99):
    """One episode's worth of policy gradient update."""
    returns = compute_returns(rewards, gamma)
 
    for state, action, G in zip(states, actions, returns):
        probabilities = softmax_policy(state, weights)
 
        # Gradient of log-probability for the ACTUAL action taken
        grad_log_prob = -probabilities.copy()
        grad_log_prob[action] += 1      # standard softmax gradient form
 
        # Update weights: push toward the action taken, SCALED by G
        weights += learning_rate * G * np.outer(state, grad_log_prob)
 
    return weights
 
# Simulating one training episode (conceptual — a real environment would
# generate these through actual interaction, e.g. via the gymnasium library)
states = [np.array([1.0, 0.5, -0.2]) for _ in range(5)]
actions = [0, 1, 0, 1, 1]
rewards = [1, 1, 1, 1, 10]      # a big reward at the END of the episode
 
weights = np.random.randn(3, 2)
weights = reinforce_update(states, actions, rewards, weights)

Why This Update Makes Sense
─────────────────────────────────────────
  If a HIGH-reward episode included taking action 1 in some
  state, the update PUSHES weights to make action 1 MORE
  likely in that state next time.

  If a LOW-reward episode included the SAME action in a
  SIMILAR state, the update PUSHES weights the OPPOSITE way.

  Over MANY episodes, actions genuinely associated with HIGH
  reward become increasingly likely — directly analogous to
  how gradient descent (Module 3, Ch.1) gradually improves
  linear regression's coefficients over many iterations.
─────────────────────────────────────────

6. The High-Variance Problem

REINFORCE's most significant practical weakness: the gradient estimate can be extremely noisy, since it depends on the total return from an entire, often-long episode.

Why Variance Is So High
─────────────────────────────────────────
  A single LUCKY or UNLUCKY episode (due to randomness in the
  environment, or in the policy's own EXPLORATION, directly
  echoing the AI Notes' exploration-exploitation discussion)
  can produce a wildly different G, even for the SAME
  underlying policy.

  This means REINFORCE's gradient updates can be ERRATIC —
  directly echoing Module 1, Chapter 4's variance concept,
  but applied to gradient ESTIMATES rather than model
  predictions.
─────────────────────────────────────────

7. Baselines — a Partial Fix

A common technique to reduce REINFORCE's variance: subtract a baseline (an estimate of the expected return) from G before using it in the update — rewarding/penalizing based on how much better or worse than expected the outcome was, not its raw magnitude.

def reinforce_update_with_baseline(states, actions, rewards, weights, baseline, learning_rate=0.01, gamma=0.99):
    returns = compute_returns(rewards, gamma)
 
    for state, action, G in zip(states, actions, returns):
        advantage = G - baseline      # how much BETTER (or worse) than EXPECTED was this outcome?
 
        probabilities = softmax_policy(state, weights)
        grad_log_prob = -probabilities.copy()
        grad_log_prob[action] += 1
 
        weights += learning_rate * advantage * np.outer(state, grad_log_prob)      # use ADVANTAGE, not raw G
 
    return weights

Why Subtracting a Baseline Doesn't Bias the Result
─────────────────────────────────────────
  Mathematically, subtracting ANY baseline that doesn't
  depend on the ACTION taken leaves the EXPECTED gradient
  UNCHANGED — but it substantially reduces VARIANCE, since
  updates now reflect "better/worse than TYPICAL," not raw,
  highly variable reward magnitudes.

  A common baseline choice: the AVERAGE return observed so
  far, or (more sophisticated) a LEARNED value function —
  which is EXACTLY the "critic" that Chapter 3's actor-critic
  methods add.
─────────────────────────────────────────

This baseline idea — using a learned estimate of "expected outcome" to compute an advantage rather than raw reward — is precisely the bridge into Chapter 3's actor-critic methods, where that baseline becomes a fully-fledged, separately-trained value function.

8. Summary & Next Steps

Key Takeaways

Policy gradient methods learn a direct mapping from state to action probabilities, using gradient ascent to make good actions more likely and bad actions less likely — the same underlying optimization idea as Module 3's gradient descent.
REINFORCE, the foundational policy gradient algorithm, updates the policy after each full episode, scaling the update by the total discounted return.
REINFORCE suffers from high variance, since a single episode's return can be noisy due to environmental randomness and the policy's own exploration.
Subtracting a baseline from the return (turning it into an "advantage") reduces variance without biasing the gradient estimate — this baseline concept directly motivates Chapter 3's actor-critic methods.

Concept Check

Why does the policy gradient update push toward actions from high-return episodes and away from actions in low-return episodes?
Why does REINFORCE's gradient estimate have such high variance?
Why doesn't subtracting a baseline from the return bias the policy gradient's expected direction?

Next Chapter

→ Chapter 3: Actor-Critic Methods

Jr Codex — 1-on-1 Personalized Coaching | Back to Module Index