Reinforcement Learning Optimization Alias: Policy Optimization

Policy Gradient Methods

A class of reinforcement learning algorithms that optimize the parameters of a policy directly by performing gradient ascent on the expected return.

In One Sentence

"Instead of building a value-based map of the environment to derive actions, policy gradient methods directly parameterize the agent's behavior and iteratively adjust it to favor actions that yielded higher-than-expected rewards."

Intuition

Imagine learning to play golf. A value-based approach (like Q-learning) would try to calculate the exact numerical value of every square inch of the golf course, figuring out the best next position from anywhere you stand.

A policy gradient approach bypasses this entirely. It focuses directly on your swing (the policy). You stand at the tee, try a swing, and observe the outcome. If the ball lands near the green (high reward), your brain reinforces the exact muscle movements you just made, making that swing more likely in the future. If the ball flies into a lake (low reward), you suppress those movements.

Mathematically, we scale the direction of our policy's parameter update by the quality of the outcome. If an action was better than our baseline expectation, we push the policy's probability distribution toward that action; if it was worse, we push it away.

Formal Definition

Let $\pi_\theta(a | s)$ be a stochastic policy parameterized by $\theta$, representing the probability of taking action $a$ in state $s$. Our goal is to maximize the expected cumulative reward over trajectories $\tau = (s_0, a_0, s_1, a_1, \dots)$:

$$J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta}[R(\tau)]$$

Using the Policy Gradient Theorem, we can compute the gradient of this objective with respect to the parameters $\theta$ without knowing the transition dynamics of the environment:

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

Where $\Psi_t$ is an estimator of the trajectory's quality. Depending on the algorithm, $\Psi_t$ can be:

Total return: $G_t = \sum_{k=t}^T \gamma^{k-t} r_k$ (REINFORCE)
State-action value: $Q^{\pi}(s_t, a_t)$ (Actor-Critic)
Advantage: $A^{\pi}(s_t, a_t) = Q^{\pi}(s_t, a_t) - V^{\pi}(s_t)$ (Advantage Actor-Critic)

Interactive Policy Update Simulator

Visualize how a single-dimensional Gaussian policy $\pi_\theta(a) = \mathcal{N}(\mu, \sigma^2)$ updates its parameters based on a sampled action and the resulting advantage.

Simulator Controls

Sampled Action ($a$) 1.0

Observed Reward ($R$) 1.5

Baseline Expectation ($V(s)$) 0.8

Learning Rate ($\eta$) 0.2

Advantage ($A = R - V$) +0.70

Policy Mean ($\mu$) 0.00

Policy Std Dev ($\sigma$) 0.80

When to use

✓ Continuous Action Spaces: Highly natural for robotics or control tasks where actions are real-valued vectors.
✓ Stochastic Policies: Essential when the optimal strategy is probabilistic (e.g., poker, rock-paper-scissors, or partially observable states).
✓ Differentiable Objectives: When you want to directly optimize the metric of interest (expected return) rather than Bellman errors.

When NOT to use

✗ Sample Inefficiency: They typically require millions of environment interactions compared to off-policy value-based methods.
✗ High Variance: Gradients can exhibit extreme variance, occasionally leading to catastrophic policy collapse during training.
✗ Local Optima: Susceptible to converging to sub-optimal deterministic policies early in training if exploration is not managed.

Evolutionary Timeline

Feb 2015

Trust Region Policy Optimization (TRPO)

Schulman et al. (arXiv:1502.05477)

Introduced mathematical guarantees for monotonic policy improvement by constraining updates using Kullback-Leibler (KL) divergence, preventing destabilizing updates.

Sep 2015

Deep Deterministic Policy Gradients (DDPG)

Lillicrap et al. (arXiv:1509.02971)

Adapted deep Q-learning principles to continuous action spaces using a deterministic policy gradient framework, allowing off-policy actor-critic learning.

Jul 2017

Proximal Policy Optimization (PPO)

Schulman et al. (arXiv:1707.06347)

Simplified the trust-region concept with a clipped surrogate objective function, enabling multiple epochs of minibatch updates while maintaining training stability.

Jan 2018

Soft Actor-Critic (SAC)

Haarnoja et al. (arXiv:1801.01290)

Proposed an off-policy maximum entropy actor-critic framework, optimizing both expected reward and policy entropy to maximize exploration and robustness.

Implementation (REINFORCE with Baseline)

reinforce_baseline.py PyTorch

import torch
import torch.nn as nn
import torch.optim as optim

class PolicyNetwork(nn.Module):
    def __init__(self, state_dim, action_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim),
            nn.Softmax(dim=-1)
        )
        
    def forward(self, state):
        return self.net(state)

def update_policy(policy, optimizer, states, actions, rewards, baseline_values, lr=1e-3):
    # Calculate discounted returns
    discounted_returns = []
    G = 0
    for r in reversed(rewards):
        G = r + 0.99 * G
        discounted_returns.insert(0, G)
    
    returns = torch.tensor(discounted_returns, dtype=torch.float32)
    baselines = torch.tensor(baseline_values, dtype=torch.float32)
    
    # Calculate advantages
    advantages = returns - baselines
    
    # Compute loss
    loss = 0
    for state, action, advantage in zip(states, actions, advantages):
        probs = policy(state)
        dist = torch.distributions.Categorical(probs)
        log_prob = dist.log_prob(action)
        
        # Negative sign for gradient ascent
        loss += -log_prob * advantage
        
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Common Variants

REINFORCE (Vanilla PG)

The simplest form of policy gradient. Uses Monte Carlo rollouts to estimate trajectory returns directly, resulting in high variance but unbiased updates.

Actor-Critic (A2C/A3C)

Introduces a learned value function (Critic) to evaluate actions, replacing high-variance returns with lower-variance advantage estimates.

PPO (Proximal Policy Optimization)

The modern standard. Employs a clipped surrogate objective to constrain policy updates, allowing multiple epochs of minibatch optimization safely.

SAC (Soft Actor-Critic)

An off-policy maximum entropy variant that optimizes for both expected reward and policy randomness, maximizing exploration and robustness.

Representative Flash Papers

While our current corpus focuses heavily on representation learning and transformers, optimization techniques like Adam are universally used to stabilize policy gradient methods.

Optimization Core

Adam: A Method for Stochastic Optimization

Policy gradient algorithms suffer from highly non-stationary data distributions. Adam's adaptive learning rates are critical for stabilizing policy updates in PPO, TRPO, and SAC.

Grounding References

[1] Schulman et al., "Proximal Policy Optimization Algorithms", arXiv:1707.06347
[2] Schulman et al., "Trust Region Policy Optimization", arXiv:1502.05477
[3] Lillicrap et al., "Continuous control with deep reinforcement learning", arXiv:1509.02971
[4] Haarnoja et al., "Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor", arXiv:1801.01290