Policy Gradient Methods: Teaching Agents to Win

The Problem: Learning from Delayed Rewards

Imagine teaching a robot to play chess. It makes thousands of moves in a game lasting hours. Only at the very end does it know if it won or lost. Which moves led to victory? Which to defeat? The temporal gap between action and outcome makes learning extraordinarily difficult.

This is the fundamental challenge of reinforcement learning: learning from delayed, sparse rewards. Unlike supervised learning where you immediately know correct answers, RL agents must learn from experience, assigning credit to past actions based on future outcomes.

Policy gradient methods approach this differently than Q-learning (the other major RL approach). Instead of learning the value of state-action pairs, they directly learn the policy—the probability of each action given the state. This direct approach proves more stable and scalable.

Markov Decision Processes: The Mathematical Framework

Reinforcement learning operates within Markov Decision Processes (MDPs). An MDP consists of:

• S: State space. The robot's configuration, game board state, etc.
• A: Action space. Possible moves (move robot arm, play chess move, etc.)
• P(s'|s, a): Transition dynamics. Probability of reaching state s' from s via action a
• R(s, a): Reward function. Immediate reward for taking action a in state s
• γ: Discount factor. How much future rewards matter (0 < γ < 1). γ=0.99 means rewards 100 steps away matter 37% as much as immediate rewards

The Markov Property: The future depends only on the present state, not the history of how you got there. Chess position matters; the sequence of moves leading to it doesn't. This property enables dynamic programming and makes the problem tractable.

Return (G_t): Total discounted future reward from time t onward.

G_t = R_t + γ R_{t+1} + γ² R_{t+2} + ... = Σ γ^k R_{t+k}

The agent's goal: maximize expected return E[G_t].

Value Functions: Understanding Expected Returns

State value function V(s): Expected return starting from state s, following the policy.

V^π(s) = E[G_t | s_t = s] = E[R_t + γ V^π(s_{t+1}) | s]

This recursive relationship (Bellman equation) says: a state's value equals immediate reward plus discounted value of the next state.

Action value function (Q-function): Expected return after taking action a in state s, then following the policy.

Q^π(s, a) = E[G_t | s_t = s, a_t = a] = E[R_t + γ V^π(s_{t+1})] where s_{t+1} ~ P(·|s, a)

Advantage function A(s, a): How much better is action a than the policy's average?

A^π(s, a) = Q^π(s, a) - V^π(s)

Large positive advantage: action a is better than average. Negative: action a is worse than average. This advantage function is crucial for policy gradient methods.

Policy Gradient Theorem: The Core Insight

How do we improve the policy? We want to increase the probability of good actions and decrease the probability of bad actions. But how do we compute gradients with respect to the policy?

The policy gradient theorem provides the answer:

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

Here, J(θ) is expected return (objective), and θ are policy parameters.

Breaking this down:

log π(a|s) is the logarithm of the action probability. ∇ log π(a|s) = ∇π(a|s) / π(a|s). This ratio (gradient of probability divided by probability itself) appears naturally in policy optimization.

Multiplying by Q(s, a) (action value) weights this gradient update. If Q(s, a) is large and positive, we strongly increase log π(a|s). If negative, we decrease it. This naturally implements the credit assignment: good actions get probability boost, bad actions get reduced probability.

Why this theorem matters: It connects the intractable objective (maximize expected return) to a tractable gradient. We can sample trajectories from the environment, compute the return, and update policy parameters in the direction of the gradient.

REINFORCE: The Algorithm Derived

From the policy gradient theorem, REINFORCE is straightforward:

1. Sample trajectory τ = (s_0, a_0, r_0, s_1, a_1, r_1, ...)

2. Compute return for each time step: G_t = Σ_{k=t}^∞ γ^(k-t) r_k

3. Update policy: θ ← θ + α ∇ log π(a_t|s_t) × G_t

This is simple and unbiased (the expectation of the update is the true gradient). However, G_t has high variance. In a single trajectory, returns fluctuate wildly. This variance leads to slow, unstable learning.

Variance Reduction via Baselines: Instead of updating proportional to G_t, use the advantage:

A_t = G_t - V(s_t)

Update: θ ← θ + α ∇ log π(a_t|s_t) × A_t

This subtracts the baseline V(s_t), reducing variance without changing the expected gradient. Intuitively: instead of "was this return high?", we ask "was this return higher than expected?" This focuses on relative performance, reducing noise.

Actor-Critic Methods: Learning Value Functions

Computing accurate returns G_t requires complete trajectories. In long-horizon tasks (playing Go), you might wait thousands of steps for a return signal. Actor-critic methods learn value functions to provide immediate feedback:

• Actor: The policy network π(a|s, θ) that chooses actions
• Critic: Value network V(s, φ) that estimates expected future returns

The advantage is estimated as:

A_t ≈ R_t + γ V(s_{t+1}) - V(s_t)

One-step temporal difference (TD) error. This requires only one step of lookahead, not full trajectories.

Training procedures:

Critic update: Reduce TD error via gradient descent

φ ← φ - β ∇ (R_t + γ V(s_{t+1}) - V(s_t))²

Actor update: Increase log probability of good actions (those with positive advantage)

θ ← θ + α ∇ log π(a_t|s_t) × (R_t + γ V(s_{t+1}) - V(s_t))

The critic learns to accurately estimate value, reducing variance of advantage estimates. The actor learns to take high-advantage actions. They work synergistically.

A3C and Asynchronous Learning

Asynchronous Advantage Actor-Critic (A3C) runs multiple agents in parallel, each exploring different states. Instead of experience replay (storing and sampling old trajectories), asynchrony provides diversity.

Each agent:

1. Runs locally for n steps

2. Computes advantage estimates using its critic

3. Computes policy gradients

4. Contributes gradients to shared parameters

Asynchrony provides natural decorrelation. When agent 1 explores the environment, agents 2-16 explore different trajectories. Sharing gradients (not experience replay) keeps memory requirements low.

A3C was crucial for distributional RL: no central replay buffer bottleneck, massive parallelism. DeepMind's Atari results using A3C surpassed human performance on numerous games.

PPO: Practical and Stable

Proximal Policy Optimization (PPO) became the workhorse algorithm because it's surprisingly stable and easy to tune. The key insight: avoid policy updates that are too large.

Instead of:

L = ∇ log π(a|s) × A

PPO uses a clipped objective:

L = E[min(r_t(θ) × A, clip(r_t(θ), 1-ε, 1+ε) × A)]

where r_t(θ) = π(a|s, θ_new) / π(a|s, θ_old) is the probability ratio (new policy divided by old policy).

Interpretation: If the advantage is positive (good action), we want to increase log π(a|s), which increases r_t. The clip constrains r_t ∈ [1-ε, 1+ε], preventing the probability from changing by more than ε. If advantage is negative, we want to decrease π(a|s), constraining r_t downward.

Clipping prevents catastrophically bad updates. Early REINFORCE variants would sometimes apply massive policy updates, destabilizing training. PPO's constraint prevents this, making learning stable even with large learning rates.

Real Example: Training CartPole

CartPole is a classic RL task: balance a pole on a cart by moving left/right. State: cart position, velocity, pole angle, angular velocity (4 values). Action: move left or right (2 options).

Using policy gradient methods:

Network: Small 2-layer network. Input: 4-state dimensions. Hidden: 64 neurons. Output: 2 (logits for left/right actions).

Training loop:

For 1000 episodes:

Initialize state

For up to 200 steps:

Sample action from π(·|s)

Execute, observe reward and next state

Store transition

Compute returns for each step

Compute policy gradients

Update policy parameters

Within 100-200 episodes, the agent learns to keep the pole balanced for 200 steps (the task's maximum). This demonstrates policy gradients' effectiveness: starting from random behavior, the agent discovers a strategy through interaction.

Trust Region Methods: Understanding the Update Size

Trust region policy optimization (TRPO) and PPO both address a fundamental issue: how much should policy change per update?

TRPO explicitly constrains the KL divergence between old and new policies:

maximize E[∇ log π(a|s) × A]

subject to KL[π_old || π_new] ≤ δ

This constrains policy changes to a "trust region" where our value approximation is accurate. Updates too large violate the assumption that V(s_t) is approximately accurate for the new policy.

PPO approximates this constraint via clipping, making TRPO's theoretical guarantees slightly weaker but dramatically simpler to implement.

AlphaGo and MCTS Integration

DeepMind's AlphaGo combined policy gradients with Monte Carlo Tree Search (MCTS). MCTS is a planning algorithm that expands a game tree, evaluating positions via rollouts (random play from position to terminal state).

Integration:

Policy network: Provides prior action probabilities for MCTS, guiding which moves to explore first

Value network: Estimates winning probability from positions, replacing MCTS rollout evaluation with a learned estimate

MCTS explores with policy guidance, gathering statistics. The value network rapidly evaluates terminal-like positions. This combination enabled superhuman Go play in 2016—a breakthrough in RL history.

Later, AlphaGo Zero removed supervised learning entirely. Using only policy and value networks plus MCTS, starting from random play, it surpassed all human players and previous AlphaGo versions. This demonstrated pure RL's power: no human data required, only self-play.

Exploration-Exploitation Tradeoff

Policy gradient methods naturally balance exploration and exploitation. The policy π(a|s) assigns probability to all actions, not just the greedy best action. Even suboptimal actions are sampled.

Early in training, the policy is nearly uniform (high entropy), exploring widely. As training progresses, the policy concentrates on good actions (lower entropy), exploiting learned knowledge.

This entropy is often regularized explicitly:

Loss = -Expected Return + β Entropy(π)

The entropy term (β is hyperparameter) encourages the policy to maintain randomness. Too much entropy: the agent remains indecisive. Too little: premature convergence to suboptimal policies. Balancing this is crucial.

Multi-Task and Meta-Reinforcement Learning

Modern applications require agents solving multiple related tasks. Policy gradient methods extend naturally:

Multi-task learning: Train one policy on multiple tasks, adding the task as part of the state. The policy learns which action is appropriate given the current state and task.

Meta-RL: Train policies that quickly adapt to new tasks. The policy receives task information (context) and updates to the new task rapidly. This is learning to learn.

These extensions push RL beyond single-task agents toward more general, adaptable systems.

Convergence and Theoretical Properties

Unlike supervised learning with well-understood convergence properties, RL convergence analysis is complex. However, important results exist:

Policy Gradient Theorem: Provided the policy and value function approximations have sufficient capacity, policy gradient updates converge to local optima.

Convergence Rates: Depend on problem specifics. Linear convergence (in standard cases) means error halves every iteration. Sublinear convergence (common with function approximation) means error decreases but more slowly.

Practical convergence depends on hyperparameters, network capacity, and problem difficulty. Tuning these requires experience and intuition.

Conclusion: Learning Through Interaction

Policy gradient methods represent a fundamental approach to RL: directly optimize the policy in the direction of higher expected return. The policy gradient theorem provides the mathematical foundation. Practical algorithms (REINFORCE, A3C, PPO) add variance reduction, parallel exploration, and stability mechanisms.

From simple CartPole to complex Go, policy gradient methods have demonstrated remarkable ability. They power modern RL systems in robotics, games, and optimization. Understanding them deeply—the underlying MDP framework, the theoretical guarantees, the practical tricks—is essential for anyone working in reinforcement learning.

The fundamental insight remains: reward signals guide the agent toward better behavior. By analyzing these signals through the lens of advantage (better-than-average actions) and combining them with neural network function approximation, we create intelligent agents that learn through interaction.