Policy Gradients, REINFORCE & Proximal Policy Optimization (PPO)

Value-based algorithms like Q-learning derive policies indirectly by estimating action values $Q(s, a)$ and selecting actions greedily. While effective for small discrete action sets, value-based approaches struggle with continuous action spaces (such as robotic torque controls) and cannot naturally represent stochastic optimal policies (such as bluffing in poker).

Policy gradient methods parameterize the policy directly as $\pi_\theta(a \mid s)$ and optimize expected cumulative return using gradient ascent. Proximal Policy Optimization (PPO) stands as the workhorse algorithm of modern deep RL and Large Language Model alignment (RLHF).

1. The Policy Gradient Theorem: Step-by-Step Derivation#

Let a trajectory $\tau$ represent a complete sequence of states and actions:

Let $R(\tau) = \sum_{t=0}^T R(s_t, a_t)$ denote the total un-discounted return of the trajectory. The expected return under policy parameters $\theta$ is:

where $P(\tau; \theta)$ is the probability density of sampling trajectory $\tau$ under policy $\pi_\theta$:

Our objective is to compute the gradient $\nabla_\theta J(\theta)$ to update parameters via gradient ascent: $\theta \leftarrow \theta + \alpha \nabla_\theta J(\theta)$.

Step 1: Differentiating the Expectation Integral#

Differentiating under the integral sign:

Notice that $\nabla_\theta P(\tau; \theta)$ is not a valid probability distribution, preventing direct Monte Carlo sampling.

Step 2: Applying the Log-Derivative Trick#

We multiply and divide by $P(\tau; \theta)$:

Substituting this back into the gradient integral:

Step 3: Expanding the Log-Probability of Trajectories#

Now consider the logarithm of the trajectory probability:

Taking the gradient with respect to $\theta$:

• $\nabla_\theta \log \rho_0(s_0) = 0$ (initial state distribution does not depend on $\theta$).
• $\nabla_\theta \log P(s_{t+1} \mid s_t, a_t) = 0$ (environmental transition dynamics do not depend on $\theta$).

Therefore:

> Crucial Insight: The environmental transition model completely vanishes from the gradient expression! We can compute unbiased analytical policy gradients without knowing transition probabilities.

Substituting back yields the fundamental Policy Gradient formula:

2. Variance Reduction: Baselines & Advantage#

While the raw REINFORCE gradient is completely unbiased, it exhibits extremely high variance because $R(\tau)$ sums rewards across the entire episode. Two mathematical techniques dramatically reduce this variance:

1. Causality (Reward-to-Go)#

Actions taken at time step $t$ cannot retroactively influence rewards received at earlier time steps $t' < t$. Therefore, we can replace the entire trajectory return $R(\tau)$ with the reward-to-go $G_t = \sum_{k=t}^T \gamma^{k-t} R_{k+1}$:

2. Baseline Invariance Proof#

We can subtract any state-dependent baseline $b(s_t)$ (such as an estimate of the state value $V(s_t)$) without introducing any bias into the gradient:

Proof:

Using $\nabla_\theta \log f(\theta) = \frac{\nabla_\theta f(\theta)}{f(\theta)}$:

Interchanging the gradient and summation:

Because probabilities must sum to 1 ($\sum_a \pi_\theta(a \mid s_t) = 1$ for all $\theta$):

By choosing $b(s_t) = V_\phi(s_t)$, the gradient term $(G_t - V(s_t))$ forms an empirical sample of the Advantage Function:

3. Proximal Policy Optimization (PPO)#

In vanilla policy gradients, a single excessively large gradient step can drastically alter $\pi_\theta$, landing the policy in a catastrophic performance collapse from which it can never collect useful trajectories again.

Trust Region Policy Optimization (TRPO) enforced a mathematical Kullback-Leibler (KL) divergence constraint: $\mathbb{E}[D_{\text{KL}}(\pi_{\theta_{\text{old}}} \parallel \pi_\theta)] \le \delta$, but required second-order conjugate gradient solvers.

In 2017, John Schulman et al. introduced Proximal Policy Optimization (PPO), which enforces trust-region behavior using a first-order clipped surrogate objective.

The Probability Ratio $r_t(\theta)$#

Define the policy probability ratio:

At the start of optimization before gradient updates, $r_t(\theta_{\text{old}}) = 1$.

The Clipped Surrogate Objective#

To discourage $r_t(\theta)$ from deviating significantly from 1, PPO clips the ratio into the interval $[1 - \epsilon, 1 + \epsilon]$ (typically $\epsilon = 0.2$):

Taking the minimum between the unclipped and clipped terms creates a pessimistic lower bound:

• When $\hat{A}_t > 0$, the action yielded higher returns than average. We want to increase its probability ($r_t > 1$), but clipping prevents updating $r_t > 1 + \epsilon$.
• When $\hat{A}_t < 0$, the action was worse than average. We want to decrease its probability ($r_t < 1$), but clipping prevents driving $r_t < 1 - \epsilon$.

4. From-Scratch Implementation: Policy Gradient (REINFORCE with Baseline)#

Let us implement a complete REINFORCE algorithm with analytical Softmax policy gradients and baseline value subtraction from first principles in Python and Rust without external libraries.

Python Implementation (Pure Standard Library)#

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
"""
REINFORCE Policy Gradient with Softmax Policy and Baseline Value Subtraction
from scratch using only the Python standard library.
"""
import math
import random
from typing import List, Tuple

class SoftmaxPolicy:
    """Linear Softmax Policy parameterized by theta: shape [num_states, num_actions]"""
    def __init__(self, num_states: int, num_actions: int, lr: float = 0.05):
        self.num_states = num_states
        self.num_actions = num_actions
        self.lr = lr
        # Initialize weights
        self.theta = [[0.0 for _ in range(num_actions)] for _ in range(num_states)]
        # Value baseline V(s)
        self.v_baseline = [0.0 for _ in range(num_states)]
        self.v_lr = 0.1

    def get_probs(self, state: int) -> List[float]:
        logits = self.theta[state]
        max_logit = max(logits)
        exp_logits = [math.exp(x - max_logit) for x in logits]
        sum_exp = sum(exp_logits)
        return [e / sum_exp for e in exp_logits]

    def select_action(self, state: int) -> int:
        probs = self.get_probs(state)
        r = random.random()
        cumulative = 0.0
        for a, p in enumerate(probs):
            cumulative += p
            if r <= cumulative:
                return a
        return self.num_actions - 1

    def update(self, trajectory: List[Tuple[int, int, float]]):
        # Compute rewards-to-go G_t
        T = len(trajectory)
        returns = [0.0] * T
        running_g = 0.0
        gamma = 0.99
        for t in reversed(range(T)):
            _, _, reward = trajectory[t]
            running_g = reward + gamma * running_g
            returns[t] = running_g

        # Update policy and baseline for each step
        for t in range(T):
            state, action, _ = trajectory[t]
            g_t = returns[t]

            # Compute advantage using baseline
            advantage = g_t - self.v_baseline[state]

            # Update baseline estimate
            self.v_baseline[state] += self.v_lr * advantage

            # Policy gradient: grad = (one_hot(a) - pi(a|s)) * advantage
            probs = self.get_probs(state)
            for a in range(self.num_actions):
                indicator = 1.0 if a == action else 0.0
                grad = (indicator - probs[a]) * advantage
                self.theta[state][a] += self.lr * grad

if __name__ == "__main__":
    policy = SoftmaxPolicy(num_states=4, num_actions=2, lr=0.05)

    # Synthetic environment rollout simulation
    for ep in range(100):
        traj = []
        state = 0
        for _ in range(10):
            action = policy.select_action(state)
            # Reward: preferring action 1 in even states, action 0 in odd states
            reward = 1.0 if (state % 2 == action) else -0.5
            traj.append((state, action, reward))
            state = (state + 1) % 4
        policy.update(traj)

    print("Learned Softmax Action Probabilities:")
    for s in range(4):
        p = policy.get_probs(s)
        print(f"State {s}: P(a=0)={p[0]:.3f}, P(a=1)={p[1]:.3f}")

Rust Implementation (Pure Standard Library)#

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
// REINFORCE Policy Gradient with Softmax Policy and Baseline Value Subtraction
// Built strictly from scratch using only Rust std.

struct SimpleRng {
    seed: u64,
}

impl SimpleRng {
    fn new(seed: u64) -> Self {
        Self { seed }
    }

    fn next_f64(&mut self) -> f64 {
        self.seed = self.seed.wrapping_mul(6364136223846793005).wrapping_add(1);
        ((self.seed >> 11) as f64) / ((1u64 << 53) as f64)
    }
}

struct SoftmaxPolicy {
    num_states: usize,
    num_actions: usize,
    lr: f64,
    v_lr: f64,
    theta: Vec<Vec<f64>>,
    v_baseline: Vec<f64>,
    rng: SimpleRng,
}

impl SoftmaxPolicy {
    fn new(num_states: usize, num_actions: usize, lr: f64, seed: u64) -> Self {
        Self {
            num_states,
            num_actions,
            lr,
            v_lr: 0.1,
            theta: vec![vec![0.0; num_actions]; num_states],
            v_baseline: vec![0.0; num_states],
            rng: SimpleRng::new(seed),
        }
    }

    fn get_probs(&self, state: usize) -> Vec<f64> {
        let logits = &self.theta[state];
        let max_l = logits.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
        let exps: Vec<f64> = logits.iter().map(|&x| (x - max_l).exp()).collect();
        let sum_exp: f64 = exps.iter().sum();
        exps.iter().map(|&e| e / sum_exp).collect()
    }

    fn select_action(&mut self, state: usize) -> usize {
        let probs = self.get_probs(state);
        let r = self.rng.next_f64();
        let mut cum = 0.0;
        for (a, &p) in probs.iter().enumerate() {
            cum += p;
            if r <= cum {
                return a;
            }
        }
        self.num_actions - 1
    }

    fn update(&mut self, trajectory: &[(usize, usize, f64)], gamma: f64) {
        let t_len = trajectory.len();
        let mut returns = vec![0.0; t_len];
        let mut running_g = 0.0;

        for t in (0..t_len).rev() {
            let (_, _, reward) = trajectory[t];
            running_g = reward + gamma * running_g;
            returns[t] = running_g;
        }

        for t in 0..t_len {
            let (state, action, _) = trajectory[t];
            let g_t = returns[t];
            let advantage = g_t - self.v_baseline[state];

            self.v_baseline[state] += self.v_lr * advantage;

            let probs = self.get_probs(state);
            for a in 0..self.num_actions {
                let indicator = if a == action { 1.0 } else { 0.0 };
                let grad = (indicator - probs[a]) * advantage;
                self.theta[state][a] += self.lr * grad;
            }
        }
    }
}

fn main() {
    let mut policy = SoftmaxPolicy::new(4, 2, 0.05, 42);

    for _ in 0..100 {
        let mut traj = Vec::new();
        let mut state = 0;
        for _ in 0..10 {
            let action = policy.select_action(state);
            let reward = if state % 2 == action { 1.0 } else { -0.5 };
            traj.push((state, action, reward));
            state = (state + 1) % 4;
        }
        policy.update(&traj, 0.99);
    }

    println!("Learned Softmax Probabilities (Rust):");
    for s in 0..4 {
        let probs = policy.get_probs(s);
        println!("State {}: P(a=0)={:.3}, P(a=1)={:.3}", s, probs[0], probs[1]);
    }
}

5. Algorithmic Comparison: Value-Based vs Policy Gradients vs PPO#

6. Knowledge Check#

question: "Why does the environmental transition probability P(s' | s, a) vanish from the Policy Gradient derivation?"

options:

• "(Recommended) Because when applying the log-derivative trick, the gradient with respect to policy parameters theta of the log transition probability is strictly zero"
• "Because the discount factor gamma cancels transition terms"
• "Because the policy is assumed to be deterministic"
• "Because the reward function R(s, a) is independent of state transitions"

explanation: "When expanding log P(tau; theta), environmental transitions log P(s_(t+1) | s_t, a_t) do not contain the policy parameter vector theta. Hence, nabla_theta log P(s_(t+1) | s_t, a_t) = 0, meaning policy gradients can be computed in entirely unknown environments."

:::

question: "In PPO, why is the clipped surrogate objective formulated as a minimum between the unclipped and clipped terms?"

options:

• "(Recommended) To create a pessimistic lower bound that prevents the policy from moving too far from the old distribution when the advantage estimate is unreliable"
• "To maximize entropy across the policy distribution"
• "To enforce a strict hard constraint on the policy's Kullback-Leibler divergence"
• "To ensure that Q-values never overestimate future returns"

explanation: "Taking the minimum min(r_t A_t, clip(r_t, 1-eps, 1+eps) A_t) constructs a pessimistic lower bound on the true expected objective, ensuring that optimistic overestimations of advantage cannot destabilize the policy."

:::