Policy Gradient: Complete Guide — Direct Policy Optimization

Summary

Policy Gradient is a family of reinforcement learning algorithms that directly optimizes an agent’s policy, without going through the intermediate estimation of a value function. Unlike value-based methods such as Q-learning, Policy Gradient adjusts the parameters θ of a parameterized policy π(a|s ; θ) in order to maximize the expected cumulative return. This approach offers major advantages: it works naturally with continuous action spaces, avoids discontinuities related to the max operator, and converges to optimal stochastic policies. This article presents the fundamental policy gradient theorem, the REINFORCE algorithm, variance reduction techniques such as the baseline and advantage, and provides a complete implementation in Python with PyTorch on the CartPole environment.

Mathematical Principle of Policy Gradient

Parameterized Policy

At the heart of Policy Gradient is the idea of representing the policy as a parameterized function:

$$\pi(a | s ; \theta) = P(a_t = a | s_t = s; \theta)$$

where θ denotes the set of learnable parameters (typically the weights of a neural network). The policy π assigns a probability to each action a given a state s. The goal is to find the parameters θ* that maximize an objective function J(θ), generally defined as the expected cumulative return under the policy π_θ.

The Policy Gradient Theorem

The central result of Policy Gradient is the Policy Gradient Theorem, demonstrated by Sutton and collaborators in 2000. This theorem establishes that the gradient of the objective function J(θ) can be written in the following form:

$$\nabla J(\theta) = \mathbb{E}{\tau \sim \pi\theta} \left[ \sum_{t=0}^{T} \nabla_\theta \log \pi_\theta(a_t | s_t) \cdot G_t \right]$$

where G_t represents the cumulative return (return) from time t:

$$G_t = \sum_{k=t}^{T} \gamma^{k-t} r_k$$

with γ ∈ [0, 1] the discount factor and r_k the reward obtained at step k. This theorem is remarkable because it does not require knowledge of the environment dynamics P(s’ | s, a) or the state distribution. The gradient is expressed solely in terms of the policy π_θ and observed rewards.

The intuition is elegant: if an action a_t led to a high return G_t, we reinforce the probability of choosing that action in state s_t (the term ∇θ log πθ(a_t | s_t) points in the direction that increases π_θ(a_t | s_t)). Conversely, if the return is low, the probability naturally decreases.

The REINFORCE Algorithm

REINFORCE (Williams, 1992) is the historical Policy Gradient algorithm, also known as the Monte Carlo policy gradient algorithm. Its update rule is straightforward:

$$\theta \leftarrow \theta + \alpha \cdot \nabla_\theta \log \pi_\theta(a_t | s_t) \cdot G_t$$

where α is the learning rate. The algorithm works as follows:

1. Generate a complete episode by following the current policy π_θ.
2. Compute cumulative returns G_t for each time step of the episode.
3. Compute the gradient for each step: ∇_θ log π_θ(a_t | s_t) · G_t.
4. Update parameters θ by accumulating gradients from the entire episode.

REINFORCE is a “Monte Carlo” style algorithm because it requires a complete episode before it can perform an update. This distinguishes it from temporal methods like Q-learning, which update at every step.

The Baseline: Variance Reduction

A major problem with REINFORCE is its high variance. Since the return G_t can fluctuate considerably from one episode to another, updates are noisy and learning becomes slow. The classical solution is to subtract a baseline b(s) from the return:

$$\nabla J(\theta) = \mathbb{E}{\tau} \left[ \sum \nabla_\theta \log \pi_\theta(a_t | s_t) \cdot (G_t – b(s_t)) \right]$$}^{T

This operation does not bias the gradient (the expectation remains unchanged) but can considerably reduce its variance. The most common choice for the baseline is the state value function V(s), which represents the expected average return from state s. When the baseline is well chosen, the term (G_t − V(s_t)) becomes an estimator of the advantage A(s_t, a_t) = Q(s_t, a_t) − V(s_t), measuring whether the chosen action was better or worse than average.

Other baselines are possible: a constant (average reward), a moving reference value, or even the return from another neural network. The choice depends on the context and the bias-variance tradeoff sought.

Intuition: Learning Directly “What to Do”

To properly understand the fundamental difference between Policy Gradient and value-based methods, imagine two musicians learning to improvise.

The first musician (analogous to Q-learning) starts by mentally building an enormous table: for every musical situation (the state) and every possible note (the action), they estimate the “score” that note would get. When it’s time to play, they consult this table and choose the note with the highest score. This approach works, but it’s rigid: if two notes have nearly identical scores, a small change can cause an abrupt switch from one note to the other. Moreover, in a continuous space (a glissando, for example), the number of “notes” is infinite and the table becomes unmanageable.

The second musician (analogous to Policy Gradient) has no table. They directly learn a probability distribution over possible actions. When a sequence of notes produces a satisfying result (a harmonious chord, approval from the audience), they reinforce the probabilities that led to this outcome. When it sounds wrong, they gradually decrease them. It’s learning by feel, guided by the final reward rather than an intermediate evaluation of each individual action.

Unlike Q-learning which first evaluates actions and then chooses, Policy Gradient learns directly “what to do” — like a musician who improvises without thinking about each note, guided by the feel of the result.

This fundamental difference explains why Policy Gradient excels in continuous action spaces (where Q-learning can’t simply “take the max”), why it converges to optimal stochastic policies (essential in environments where randomization is strategically advantageous), and why it’s the foundation of modern algorithms like PPO and SAC.

Python Implementation with PyTorch

Here is a complete implementation of REINFORCE on the CartPole-v1 environment from OpenAI Gym, using PyTorch as the computation framework.

Installing Dependencies

pip install torch gymnasium numpy matplotlib

Complete Code: REINFORCE on CartPole

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical
import gymnasium as gym
import matplotlib.pyplot as plt
import numpy as np


class PolicyNetwork(nn.Module):
    """
    Policy network for REINFORCE.
    Takes a state as input and returns a probability distribution
    over the discrete action space.
    """
    def __init__(self, state_dim: int, action_dim: int, hidden_size: int = 128):
        super(PolicyNetwork, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_size)
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, action_dim)

    def forward(self, state: torch.Tensor) -> torch.Tensor:
        """Returns logits for each action."""
        x = F.relu(self.fc1(state))
        x = F.relu(self.fc2(x))
        logits = self.fc3(x)
        return logits

    def select_action(self, state: torch.Tensor):
        """
        Selects an action according to the current policy.
        Returns the action, the associated log-probability, and the entropy.
        """
        logits = self.forward(state)
        dist = Categorical(logits=logits)
        action = dist.sample()
        log_prob = dist.log_prob(action)
        entropy = dist.entropy()
        return action.item(), log_prob, entropy


def compute_returns(rewards: list, gamma: float) -> torch.Tensor:
    """
    Computes discounted cumulative returns for an episode.
    G_t = r_t + gamma * r_{t+1} + gamma^2 * r_{t+2} + ...
    """
    returns = []
    R = 0.0
    for reward in reversed(rewards):
        R = reward + gamma * R
        returns.insert(0, R)
    return torch.tensor(returns, dtype=torch.float32)


def reinforce(
    env_name: str = "CartPole-v1",
    learning_rate: float = 1e-2,
    gamma: float = 0.99,
    n_episodes: int = 1000,
    hidden_size: int = 128,
    seed: int = 42
):
    """
    REINFORCE algorithm with PyTorch.

    Args:
        env_name: Name of the Gymnasium environment.
        learning_rate: Learning rate (alpha).
        gamma: Discount factor.
        n_episodes: Total number of training episodes.
        hidden_size: Size of the network's hidden layers.
        seed: Random seed for reproducibility.
    """
    # Initialization
    torch.manual_seed(seed)
    env = gym.make(env_name)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

    policy = PolicyNetwork(state_dim, action_dim, hidden_size)
    optimizer = optim.Adam(policy.parameters(), lr=learning_rate)

    rewards_history = []
    running_avg = []

    for episode in range(1, n_episodes + 1):
        state, _ = env.reset(seed=seed + episode)
        log_probs = []
        rewards = []
        entropies = []
        done = False
        total_reward = 0.0

        # === Phase 1: Collecting the complete trajectory ===
        while not done:
            state_tensor = torch.FloatTensor(state).unsqueeze(0)
            action, log_prob, entropy = policy.select_action(state_tensor)

            state, reward, terminated, truncated, _ = env.step(action)
            done = terminated or truncated

            log_probs.append(log_prob)
            rewards.append(reward)
            entropies.append(entropy)
            total_reward += reward

        # === Phase 2: Computing returns ===
        returns = compute_returns(rewards, gamma)

        # Return normalization (variance reduction)
        if len(returns) > 1:
            returns = (returns - returns.mean()) / (returns.std() + 1e-8)

        # === Phase 3: Parameter update ===
        # Negative loss since we're doing gradient descent
        policy_loss = []
        for log_prob, ret in zip(log_probs, returns):
            policy_loss.append(-log_prob * ret)

        # Adding an entropy term to encourage exploration
        entropy_bonus = -0.01 * sum(entropies)

        loss = torch.stack(policy_loss).sum() + entropy_bonus

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # Tracking
        rewards_history.append(total_reward)
        if len(rewards_history) >= 100:
            avg = sum(rewards_history[-100:]) / 100
            running_avg.append(avg)
            if episode % 50 == 0:
                print(f"Episode {episode:4d} | Reward: {total_reward:6.1f} | "
                      f"Last 100 avg: {avg:.1f} | Loss: {loss.item():.4f}")

        # Early stopping criterion
        if len(running_avg) > 0 and running_avg[-1] >= 490:
            print(f"\nConvergence reached at episode {episode}! "
                  f"Average over 100 episodes: {running_avg[-1]:.1f}")
            break

    env.close()
    return rewards_history, running_avg


if __name__ == "__main__":
    print("=" * 60)
    print("REINFORCE on CartPole-v1 — Policy Gradient (Monte Carlo)")
    print("=" * 60)

    hist, avg = reinforce(
        learning_rate=0.01,
        gamma=0.99,
        n_episodes=1000,
        hidden_size=128,
        seed=42
    )

    # Visualization
    plt.figure(figsize=(12, 4))
    plt.subplot(1, 2, 1)
    plt.plot(hist, alpha=0.3, color="gray", label="Reward per episode")
    plt.xlabel("Episode")
    plt.ylabel("Total reward")
    plt.title("Learning progress")
    plt.legend()
    plt.grid(True, alpha=0.3)

    plt.subplot(1, 2, 2)
    plt.plot(avg, color="blue", linewidth=2, label="Moving average (100 episodes)")
    plt.axhline(y=495, color="green", linestyle="--", alpha=0.5, label="Convergence threshold")
    plt.xlabel("Episode")
    plt.ylabel("Average reward")
    plt.title("Convergence")
    plt.legend()
    plt.grid(True, alpha=0.3)

    plt.tight_layout()
    plt.savefig("reinforce_cartpole_results.png", dpi=150)
    print("\nResults saved to 'reinforce_cartpole_results.png'")

Detailed Code Explanation

Network architecture: The PolicyNetwork receives a state vector of dimension 4 (cart position, velocity, pole angle, angular velocity) and returns logits for 2 actions (push left or right). The select_action function uses Categorical to sample an action according to the probability distribution, ensuring natural exploration during training.

Trajectory collection: REINFORCE is an “episode-complete” algorithm. Unlike DQN which learns at every transition, REINFORCE accumulates all log-probabilities and rewards during the entire episode before performing an update. This characteristic is both its strength (unbiased gradient estimate) and its weakness (high variance, slower learning at the start of training).

Return normalization: Before backpropagation, returns are normalized by subtracting the mean and dividing by the standard deviation. This acts as an adaptive baseline and significantly reduces gradient variance, which stabilizes learning.

Entropy term: The entropy bonus (−0.01 × H(π)) encourages the policy to remain sufficiently exploratory, preventing it from collapsing too quickly to a deterministic action. This is a very common form of regularization in modern Policy Gradient algorithms.

Hyperparameters

The choice of hyperparameters is crucial for Policy Gradient performance. Here are the main ones:

Tuning Tips

1. Start small: With CartPole, hidden_size=64 and learning_rate=0.01 work remarkably well. Don’t oversize the network unnecessarily.
2. Monitor entropy: If entropy drops too fast (the policy becomes deterministic prematurely), increase the entropy coefficient or reduce the learning rate.
3. Normalize returns: Without normalization, REINFORCE is very sensitive to extreme rewards. Normalization (z-score) is the simplest way to stabilize learning.
4. Increase gamma progressively: If the agent struggles to learn long-term behaviors, increase γ in steps (0.95 → 0.99) as training progresses.

Advantages and Limitations of Policy Gradient

Advantages

• Continuous action spaces: Unlike Q-learning which requires computing a maximum over all actions (impossible in a continuous space), Policy Gradient works naturally with continuous distributions (e.g., a multivariate Gaussian for real-valued actions).
• Optimal stochastic policies: Policy Gradient can represent and converge to stochastic policies, which is essential in partially observable environments (POMDP) where randomization is strategically necessary.
• No max discontinuity: Q-learning suffers from discontinuities when the optimal action changes abruptly. Policy Gradient evolves continuously in probability space.
• Compatibility with deep networks: Policy Gradient combines naturally with deep neural networks, enabling processing of complex observations like images or text.

Limitations

• High variance: REINFORCE gradients are inherently noisy because they depend on the complete return from an episode, which can vary considerably. This makes convergence slow and unstable without variance reduction techniques.
• Convergence to local optima: The objective function J(θ) is non-convex, and gradient algorithms can converge to suboptimal local optima, especially with poorly initialized policies.
• Sample inefficient: Each REINFORCE update benefits from only one episode. Compared to “off-policy” methods like DQN that reuse transitions via a replay buffer, “on-policy” Policy Gradient wastes a lot of data.
• Sensitivity to hyperparameters: The learning rate is particularly critical. A step that’s too large can destroy an already performing policy, while a step that’s too small makes learning endless.

These limitations have motivated the development of more sophisticated algorithms such as PPO (Proximal Policy Optimization), A3C (Asynchronous Advantage Actor-Critic), and SAC (Soft Actor-Critic), which combine the advantages of Policy Gradient with advanced stabilization mechanisms.

4 Use Cases of Policy Gradient

1. Robotics and Motor Control

Robot control in continuous spaces is one of the most classic use cases for Policy Gradient. A robotic arm with 7 joints requires a vector of 7 real values at the output (angles or torques). Value-based methods cannot efficiently handle this continuous space, while a Policy Gradient with a Gaussian distribution at the output directly models continuous commands. Work such as that of Levine et al. (2016) showed that policies learned by policy gradient enable robots to grasp objects, open doors, or perform complex manipulation tasks learning directly from sensor data.

2. Video Games and Autonomous Agents

Agents playing complex video games (StarCraft, Dota 2, Atari) frequently use variants of Policy Gradient. The most famous example is AlphaGo and its successors, which combine a policy network (selecting likely moves) with a value network (estimating winning chances). In gaming environments, the action space is often immense and stochastic policy allows richer exploration than deterministic approaches. The PPO algorithm, derived from Policy Gradient, was used to train Dota 5 agents capable of beating human professional players.

3. Algorithmic Finance and Automated Trading

In finance, trading decisions involve continuous amounts (how much to buy/sell) and precise timing (when to enter/exit a position). Policy Gradient, with its stochastic policies and ability to directly optimize complex reward functions (Sharpe ratio, maximum drawdown, cumulative return), is particularly suited to these problems. Unlike classical supervised approaches that simply predict market direction, a Policy Gradient agent learns a complete execution strategy, integrating risk management directly into its objective function.

4. Recommendation Optimization and Targeted Advertising

Modern recommendation systems (Netflix, YouTube, TikTok) optimize content selection policy in real time to maximize user engagement. Policy Gradient allows modeling this policy continuously: instead of simply classifying whether a user will like content, the algorithm learns a personalized mix of genres, durations, sources, and formats. The policy evolves over time based on implicit feedback (watch time, shares, subscriptions), which corresponds exactly to the reward-return paradigm of Policy Gradient. Advertising platforms use similar approaches to optimize real-time bidding (RTB) and maximize campaign return on investment.

See Also

• Implementation of Cramer’s Algorithm and Determinant Calculation in Python
• Optimize Your Python Code with SOP and POS: Complete Guide for Beginners