Policy Gradient and Actor-Critic Explained

Disclosure: As an Amazon Associate, I earn from qualifying purchases. Some links in this post are affiliate links — they cost you nothing extra.


.rl-series-nav{margin:2rem 0;padding:1.5rem;border:1px solid #e2e8f0;border-radius:10px;background:#f8fafc}.rl-series-nav h4{margin:0 0 .8rem;color:#6c5ce7;font-size:1.1rem}.rl-series-nav ul{list-style:none;padding:0;margin:0}.rl-series-nav li{padding:.3rem 0}.rl-series-nav a{color:#4a5568;text-decoration:none;font-size:.9rem}.rl-series-nav a:hover{color:#6c5ce7}.rl-series-nav .current{font-weight:700;color:#6c5ce7}

This is Part 3 of our 5-part Reinforcement Learning series. We’re leaving value-based methods behind and learning to optimize policies directly.

Series Overview:
– Part 1: RL Basics — MDP, Bellman Equation, Value Functions
– Part 2: From Q-Learning to DQN
Part 3: Policy Gradient Methods (You are here)
– Part 4: PPO — The Industry Standard
– Part 5: SAC — Mastering Continuous Control


Why Policy Gradients?

DQN learns a value function and derives a policy from it. This works for discrete actions, but what about continuous control — steering angles, joint torques, or portfolio allocations?

Policy gradient methods take a different approach: directly parameterize and optimize the policy.

Approach Learns Action Space Example
Value-based (DQN) Q(s,a) → derive π Discrete only Atari games
Policy gradient π(a|s) directly Discrete or continuous Robotics, control

The Policy Gradient Theorem

We want to find policy parameters θ that maximize expected return:

J(θ) = E_π [Σ γ^t · r_t]

The policy gradient theorem gives us the gradient:

∇J(θ) = E_π [ ∇log π_θ(a|s) · Q_π(s, a) ]

Intuitively: increase the probability of actions that lead to high returns, decrease the probability of actions that lead to low returns. The log π term handles the “how much to adjust” and Q handles “in which direction.”

Wrapping your head around policy gradient math at 3am? Dark Chocolate Espresso Beans make the derivative of log π taste a lot sweeter.

Enjoying this article? Get more like it delivered to your inbox. Subscribe to the newsletter

REINFORCE: The Simplest Policy Gradient

REINFORCE (Williams, 1992) replaces Q(s,a) with the actual episode return G_t:

∇J(θ) ≈ Σ_t ∇log π_θ(a_t|s_t) · G_t

where G_t = r_t + γr_{t+1} + γ²r_{t+2} + …

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import gymnasium as gym


class PolicyNetwork(nn.Module):
    """Stochastic policy network for discrete actions."""
    def __init__(self, state_dim, action_dim, hidden=128):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, action_dim),
            nn.Softmax(dim=-1),
        )

    def forward(self, x):
        return self.net(x)


def reinforce(env_name="CartPole-v1", episodes=1000, gamma=0.99, lr=1e-3):
    env = gym.make(env_name)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

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

    rewards_history = []

    for ep in range(episodes):
        state, _ = env.reset()
        log_probs = []
        rewards = []

        # Collect one full episode
        while True:
            state_t = torch.FloatTensor(state)
            probs = policy(state_t)
            dist = torch.distributions.Categorical(probs)
            action = dist.sample()
            log_probs.append(dist.log_prob(action))

            next_state, reward, terminated, truncated, _ = env.step(action.item())
            rewards.append(reward)
            state = next_state

            if terminated or truncated:
                break

        # Compute discounted returns
        returns = []
        G = 0
        for r in reversed(rewards):
            G = r + gamma * G
            returns.insert(0, G)

        returns = torch.FloatTensor(returns)
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)  # Normalize

        # Policy gradient update
        loss = 0
        for log_prob, G in zip(log_probs, returns):
            loss -= log_prob * G  # Negative because we maximize

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

        total = sum(rewards)
        rewards_history.append(total)

        if (ep + 1) % 100 == 0:
            avg = np.mean(rewards_history[-100:])
            print(f"Episode {ep+1} | Avg Reward: {avg:.1f}")

    env.close()
    return policy, rewards_history

policy, rewards = reinforce()

The Variance Problem

REINFORCE works but it’s high variance. Why? Because G_t includes all future randomness — the same action in the same state might get very different returns across episodes. This makes learning noisy and slow.

Solution: Subtract a baseline. If we subtract a value that doesn’t depend on the action, the gradient is still unbiased but has lower variance:

∇J(θ) ≈ Σ_t ∇log π_θ(a_t|s_t) · (G_t - b)

The optimal baseline is V(s) — the expected return from that state. The difference A(s,a) = Q(s,a) – V(s) is called the advantage — how much better this action is compared to average.

Actor-Critic: Best of Both Worlds

Instead of using full episode returns, Actor-Critic methods learn two networks simultaneously:

  • Actor: The policy network π_θ(a|s) — decides what to do
  • Critic: The value network V_φ(s) — evaluates how good the current state is

The advantage is estimated as:

A(s, a) ≈ r + γ · V(s') - V(s)    (TD advantage)

This is a one-step estimate — much lower variance than full returns, at the cost of some bias.

import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
import gymnasium as gym


class ActorCritic(nn.Module):
    """Combined actor-critic with shared feature layers."""
    def __init__(self, state_dim, action_dim, hidden=128):
        super().__init__()
        # Shared feature extractor
        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.ReLU(),
        )
        # Actor head (policy)
        self.actor = nn.Sequential(
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, action_dim),
            nn.Softmax(dim=-1),
        )
        # Critic head (value)
        self.critic = nn.Sequential(
            nn.Linear(hidden, hidden),
            nn.ReLU(),
            nn.Linear(hidden, 1),
        )

    def forward(self, x):
        features = self.shared(x)
        return self.actor(features), self.critic(features)


def actor_critic(env_name="CartPole-v1", episodes=1000, gamma=0.99, lr=3e-4):
    env = gym.make(env_name)
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

    model = ActorCritic(state_dim, action_dim)
    optimizer = optim.Adam(model.parameters(), lr=lr)

    rewards_history = []

    for ep in range(episodes):
        state, _ = env.reset()
        total_reward = 0
        log_probs, values, rewards, dones = [], [], [], []

        while True:
            state_t = torch.FloatTensor(state)
            probs, value = model(state_t)
            dist = torch.distributions.Categorical(probs)
            action = dist.sample()

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

            log_probs.append(dist.log_prob(action))
            values.append(value.squeeze())
            rewards.append(reward)
            dones.append(done)

            state = next_state
            total_reward += reward

            if done:
                break

        # Compute advantages using TD residuals
        advantages = []
        returns = []
        next_value = 0  # Terminal state value

        for t in reversed(range(len(rewards))):
            if dones[t]:
                next_value = 0
            R = rewards[t] + gamma * next_value
            advantage = R - values[t].item()
            returns.append(R)
            advantages.append(advantage)
            next_value = values[t].item()

        advantages = torch.FloatTensor(advantages[::-1])
        returns = torch.FloatTensor(returns[::-1])
        advantages = (advantages - advantages.mean()) / (advantages.std() + 1e-8)

        log_probs = torch.stack(log_probs)
        values = torch.stack(values)

        # Actor loss: policy gradient with advantage
        actor_loss = -(log_probs * advantages.detach()).mean()

        # Critic loss: MSE between predicted and actual returns
        critic_loss = nn.MSELoss()(values, returns.detach())

        # Combined loss
        loss = actor_loss + 0.5 * critic_loss

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

        rewards_history.append(total_reward)

        if (ep + 1) % 100 == 0:
            avg = np.mean(rewards_history[-100:])
            print(f"Episode {ep+1} | Avg: {avg:.1f}")

    env.close()
    return model, rewards_history

model, rewards = actor_critic()

Generalized Advantage Estimation (GAE)

The one-step TD advantage r + γV(s') - V(s) is low variance but biased. Full returns are unbiased but high variance. GAE (Schulman, 2016) provides a smooth interpolation using parameter λ:

δ_t = r_t + γV(s_{t+1}) - V(s_t)           (TD error)

A_t^GAE = δ_t + (γλ)δ_{t+1} + (γλ)²δ_{t+2} + ...
λ value Behavior
λ = 0 Pure TD (low variance, high bias)
λ = 1 Monte Carlo returns (high variance, no bias)
λ = 0.95 Common sweet spot
def compute_gae(rewards, values, dones, gamma=0.99, lam=0.95):
    """Compute Generalized Advantage Estimation."""
    advantages = []
    gae = 0

    for t in reversed(range(len(rewards))):
        if t == len(rewards) - 1:
            next_value = 0
        else:
            next_value = values[t + 1]

        if dones[t]:
            next_value = 0
            gae = 0

        delta = rewards[t] + gamma * next_value - values[t]
        gae = delta + gamma * lam * gae
        advantages.insert(0, gae)

    return advantages

GAE is used in essentially every modern policy gradient algorithm, including PPO (Part 4).

Continuous Action Spaces

For continuous actions, we output a Gaussian distribution instead of discrete probabilities:

class ContinuousPolicy(nn.Module):
    """Policy for continuous action spaces."""
    def __init__(self, state_dim, action_dim, hidden=128):
        super().__init__()
        self.shared = nn.Sequential(
            nn.Linear(state_dim, hidden),
            nn.ReLU(),
            nn.Linear(hidden, hidden),
            nn.ReLU(),
        )
        self.mean = nn.Linear(hidden, action_dim)
        self.log_std = nn.Parameter(torch.zeros(action_dim))

    def forward(self, state):
        features = self.shared(state)
        mean = self.mean(features)
        std = self.log_std.exp()
        return torch.distributions.Normal(mean, std)

    def act(self, state):
        dist = self.forward(state)
        action = dist.sample()
        log_prob = dist.log_prob(action).sum(dim=-1)
        return action, log_prob

The network outputs mean μ and we learn a separate log standard deviation. Actions are sampled from N(μ, σ²), and log_prob gives us ∇log π for the policy gradient.

Algorithm Comparison

Method Variance Bias Sample Efficiency Implementation
REINFORCE High None Low (full episodes) Simple
Actor-Critic (1-step) Low Some Medium Moderate
Actor-Critic + GAE Tunable Tunable Medium Moderate
A2C (synchronous) Low-Med Some Medium Moderate
A3C (async parallel) Low-Med Some High (parallel) Complex

Key Takeaways

  1. Policy gradients optimize the policy directly, enabling continuous actions
  2. REINFORCE is simple but high variance — normalize returns as a minimum
  3. Actor-Critic combines policy and value learning for lower variance
  4. GAE provides a flexible bias-variance tradeoff via λ
  5. Continuous policies use Gaussian distributions parameterized by neural networks

What’s Next

In Part 4, we tackle PPO (Proximal Policy Optimization) — the algorithm that made policy gradients practical for production. We’ll see how a simple clipping trick prevents catastrophic policy updates and why PPO became the default choice for everything from game AI to RLHF in ChatGPT.

FAQ

What’s the difference between A2C and A3C?

A3C (Asynchronous Advantage Actor-Critic) runs multiple workers in parallel, each with its own copy of the environment, updating a shared model asynchronously. A2C is the synchronous version — all workers collect experience, then update together. A2C is simpler to implement, easier to debug, and in practice performs just as well. Most modern codebases use A2C or PPO instead of A3C.

Why do we normalize advantages?

Normalizing advantages (subtracting mean, dividing by std) prevents the gradient from being dominated by outlier episodes. Without normalization, a single very high or very low return can cause a large, destabilizing update. It also helps maintain a consistent learning speed as the agent improves and returns change scale.

How is policy gradient different from evolutionary strategies?

Both optimize a policy, but policy gradients use backpropagation through the policy’s log-probabilities — exploiting the structure of the problem. Evolutionary strategies treat the policy as a black box and use population-based search. Policy gradients are far more sample-efficient for problems where gradients are informative, but evolutionary methods can work in non-differentiable settings and are embarrassingly parallel.

Did you find this helpful?

Your support keeps this blog running and ad-free content coming.

☕ Buy me a coffee
TODAY 372 | TOTAL 113,648