Deep RL: Policy Gradients and PPO 🚀

Nisha, 16, from Hyderabad had beaten Lesson 3. Her CartPole Q-agent could balance for 500 steps. She was proud — until her classmate showed her LunarLander-v2, where a rocket must fire its thrusters to land softly on a pad. The state space was continuous, the actions were continuous, and her Q-table approach was completely useless.

"Q-Learning is a value method," her teacher explained. "It learns what states are worth. Policy Gradient methods are different — they directly optimise the probability of taking good actions. That's what PPO does."

Three hours after installing stable-baselines3, Nisha's PPO agent was landing the rocket cleanly. She had moved from hand-crafted Q-tables to production-grade RL in an afternoon.

Lesson 3 covered value-based RL: learn Q(s,a), then pick the action with highest Q. This works for discrete action spaces but fails for continuous ones — you can't take argmax over infinitely many actions.

REINFORCE (Williams, 1992) is the simplest policy gradient algorithm. Collect a full episode, compute the return (sum of discounted rewards) at each step, then update the policy to make good actions more probable and bad actions less probable.

import torch, torch.nn as nn, torch.optim as optim
import gymnasium as gym, numpy as np
from torch.distributions import Categorical

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

env    = gym.make("CartPole-v1")
policy = PolicyNet(4, 2)
optim  = torch.optim.Adam(policy.parameters(), lr=1e-3)
GAMMA  = 0.99

for episode in range(2000):
    states, actions, rewards = [], [], []
    state, _ = env.reset()

    # ── Collect one episode ──────────────────────────────────
    done = False
    while not done:
        s_tensor = torch.FloatTensor(state)
        probs    = policy(s_tensor)
        dist     = Categorical(probs)
        action   = dist.sample().item()

        next_state, reward, terminated, truncated, _ = env.step(action)
        states.append(s_tensor)
        actions.append(action)
        rewards.append(reward)
        done = terminated or truncated
        state = next_state

    # ── Compute discounted returns G_t ───────────────────────
    G, returns = 0, []
    for r in reversed(rewards):
        G = r + GAMMA * G
        returns.insert(0, G)
    returns = torch.tensor(returns, dtype=torch.float32)
    returns = (returns - returns.mean()) / (returns.std() + 1e-8)  # normalise

    # ── Compute policy gradient loss ─────────────────────────
    loss = 0
    for s, a, g in zip(states, actions, returns):
        probs = policy(s)
        log_prob = Categorical(probs).log_prob(torch.tensor(a))
        loss += -log_prob * g   # negative because we maximise

    optim.zero_grad()
    loss.backward()
    optim.step()

PPO (Proximal Policy Optimisation, Schulman et al. 2017) is the most widely used deep RL algorithm. It trains on multiple mini-batches of recent experience while preventing updates that change the policy too drastically — the "proximal" (near/close) constraint.

The clipping is the key insight: if a gradient update would change the policy too much (ratio outside [0.8, 1.2]), clip it. This makes PPO conservative — it improves reliably without the instability of older policy gradient methods.

# ── PPO with Stable Baselines 3 (industry library) ──────────────
# pip install stable-baselines3[extra] gymnasium

from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
import gymnasium as gym

# ── CartPole ─────────────────────────────────────────────────────
env = make_vec_env("CartPole-v1", n_envs=4)  # 4 parallel environments
model = PPO(
    "MlpPolicy",       # multi-layer perceptron policy
    env,
    verbose=1,
    learning_rate=3e-4,
    n_steps=2048,      # steps to collect per rollout per env
    batch_size=64,
    n_epochs=10,       # PPO update epochs per rollout
    gamma=0.99,
    clip_range=0.2,    # ε — clipping parameter
    ent_coef=0.0,      # entropy bonus (encourage exploration)
)
model.learn(total_timesteps=200_000)
model.save("ppo_cartpole")

# ── LunarLander (requires continuous action support) ─────────────
from stable_baselines3 import PPO as PPO_C

lunar_env = make_vec_env("LunarLander-v2", n_envs=8)
lunar_model = PPO_C("MlpPolicy", lunar_env, verbose=1,
                     learning_rate=3e-4, n_steps=1024,
                     batch_size=64, n_epochs=4, gamma=0.999)
lunar_model.learn(total_timesteps=1_000_000)

# ── Evaluate ─────────────────────────────────────────────────────
from stable_baselines3.common.evaluation import evaluate_policy
mean_reward, std_reward = evaluate_policy(lunar_model, lunar_env, n_eval_episodes=10)
print(f"Mean reward: {mean_reward:.2f} +/- {std_reward:.2f}")
# Target: mean_reward > 200 for LunarLander-v2 = "solved"