Policy gradients and REINFORCE

What it is and when you reach for it

Reinforcement learning has two broad strategies for getting a good policy. Value-based methods, like the Q-learning family, learn how good each action is and act greedily on those estimates; the policy is a byproduct. Policy gradient methods skip the middleman: they parameterize the policy itself, π_θ(a|s), as a neural network that maps states to a distribution over actions, and they ascend the gradient of expected return with respect to θ. REINFORCE (Williams, 1992) is the original and simplest member of the family, and it is worth knowing cold because everything downstream is a variance-reduction or stability patch on top of it: actor-critic replaces the sampled return with a learned estimate, PPO clips the update so you can safely reuse data, and GRPO swaps the learned baseline for a group-mean baseline over sampled completions. You reach for policy gradients when the action space is awkward for value methods (continuous, very large, or structured, like emitting a token sequence), when you want stochastic policies for free, or when the thing you can measure is only a score at the end of a long sampled sequence, which is exactly the situation in RLHF.

The math

The RL objective

A trajectory τ = (s₀, a₀, s₁, a₁, ..., s_T) is produced by three ingredients: an initial state distribution ρ(s₀), the policy π_θ(a_t|s_t), and the environment dynamics P(s_t+1|s_t, a_t). Its probability factorizes as

p_θ(τ) = ρ(s₀) Π_t π_θ(a_t|s_t) P(s_t+1|s_t, a_t),

and the objective is the expected return J(θ) = E_τ∼pθ[R(τ)] where R(τ) = Σ_t γ^t r_t. The difficulty is that θ influences J only through which trajectories get sampled. The reward is a black box, the dynamics are a black box, and neither is differentiable. Gradient descent seems off the table until you notice that the sampling distribution itself is differentiable in θ.

The log-derivative trick, line by line

Write the expectation as an integral and push the gradient through:

∇_θJ(θ) = ∇_θ ∫ p_θ(τ) R(τ) dτ

= ∫ ∇_θ p_θ(τ) R(τ) dτ   (the gradient and integral swap; R does not depend on θ)

= ∫ p_θ(τ) [∇_θ p_θ(τ) / p_θ(τ)] R(τ) dτ   (multiply and divide by p_θ(τ))

= ∫ p_θ(τ) ∇_θ log p_θ(τ) R(τ) dτ   (the identity ∇ log p = ∇p / p, read right to left)

= E_τ∼pθ[∇_θ log p_θ(τ) R(τ)].

The multiply-and-divide step looks like algebraic sleight of hand but it is the entire method: it rewrites the gradient of an expectation as an expectation of a gradient, and expectations can be estimated by sampling. Now expand log p_θ(τ) using the factorization above:

log p_θ(τ) = log ρ(s₀) + Σ_t log π_θ(a_t|s_t) + Σ_t log P(s_t+1|s_t, a_t).

The initial-state term and the dynamics terms do not contain θ, so their gradients vanish. The unknown environment drops out of the gradient entirely, which is why policy gradients are model-free: you never need to know or learn P to compute them. What survives is the policy gradient theorem in its rawest form:

∇_θJ(θ) = E_τ[ (Σ_t ∇_θ log π_θ(a_t|s_t)) · R(τ) ].

Why rewards multiply log-probs

The estimator has a clean mechanical reading. ∇_θ log π_θ(a_t|s_t) is the direction in parameter space that increases the log probability of the action actually taken. The return R(τ) is a scalar coefficient on that direction. If the trajectory earned more reward than usual, every action along it gets pushed toward higher probability, in proportion to how good the trajectory was; if the return is negative, every action gets pushed away. It is supervised learning where the labels are your own sampled actions and the per-example weight is how well things turned out. This also exposes the method's central weakness: credit is assigned to whole trajectories, not to the individual decisions that earned it, and one scalar has to speak for hundreds of actions. The next two ideas repair that as far as it can be repaired without a learned critic.

Causality: reward-to-go

An action taken at time t cannot influence rewards collected before t. Formally, for t′ < t the cross term E[∇_θ log π_θ(a_t|s_t) · r_t′] factors over the trajectory distribution: conditioned on the past up through s_t, the reward r_t′ is already determined, and E_at∼π[∇_θ log π_θ(a_t|s_t)] = Σ_a π(a|s_t) ∇ log π(a|s_t) = ∇ Σ_a π(a|s_t) = ∇1 = 0. Past rewards contribute exactly zero to the expected gradient, but they contribute plenty of noise to any finite-sample estimate, so we delete them. Each log-prob gets multiplied only by the reward-to-go from its own timestep forward:

∇_θJ(θ) = E_τ[ Σ_t ∇_θ log π_θ(a_t|s_t) · G_t ],   G_t = Σ_t′≥t γ^t′−t r_t′.

Same expectation, strictly less variance. This is the first estimator anyone should actually implement.

Baselines: less variance, zero bias

The second repair subtracts a baseline b(s_t) from the reward-to-go, weighting each log-prob by (G_t − b(s_t)) instead of G_t. The claim that this changes nothing in expectation follows from the same identity as causality. Fix a state s and any function b that depends on the state but not the action:

E_a∼πθ[ b(s) ∇_θ log π_θ(a|s) ] = b(s) Σ_a π_θ(a|s) · ∇_θπ_θ(a|s) / π_θ(a|s) = b(s) ∇_θ Σ_a π_θ(a|s) = b(s) ∇_θ 1 = 0.

The π in the expectation cancels the π in the denominator of ∇ log π, the sum over actions collapses to the derivative of a constant, and the whole term dies. Subtracting b therefore leaves the expected gradient untouched. b(s) could not depend on the action, or the b(s, a) could not be pulled out of the sum and the cancellation would fail; that restriction is the fine print on every baseline.

Why it slashes variance is best seen with numbers. Take a two-armed bandit with a softmax policy over logits θ = [0, 0], so π = [0.5, 0.5], and suppose arm 0 pays 101 while arm 1 pays 100. For a softmax policy, ∇_θ log π(a) = onehot(a) − π. Sampling arm 0 gives the gradient estimate 101 · [0.5, −0.5] = [50.5, −50.5]; sampling arm 1 gives 100 · [−0.5, 0.5] = [−50, 50]. The true expected gradient is their average, [0.25, −0.25]: a signal of magnitude 0.25 buried under samples of magnitude 50, flipping sign from draw to draw. Now subtract the baseline b = 100.5, the mean payout. Arm 0 contributes 0.5 · [0.5, −0.5] = [0.25, −0.25]; arm 1 contributes (−0.5) · [−0.5, 0.5] = [0.25, −0.25]. Identical expectation, variance collapsed to nearly zero, and the estimator now says the interpretable thing: reinforce actions that did better than expected, suppress actions that did worse, where "expected" is the baseline. The standard choice learns b(s) ≈ V^π(s) with a small value network, making the weight G_t − V(s_t) an estimate of the advantage. Push that idea one step further, letting the critic also replace the sampled return, and you have actor-critic.

Entropy regularization

A pure policy gradient happily collapses to a deterministic policy the moment something works a little, and once π(a|s) ≈ 0 for an action, the gradient through log π can no longer resurrect it: the policy stops exploring and the estimator only sees what the policy still does. The standard countermeasure adds the policy's entropy to the objective, J(θ) + β E[H(π_θ(·|s))], with β small (10⁻² is a common default). The entropy bonus is a differentiable term you compute exactly from the action distribution, no sampling identity needed, and it applies constant gentle pressure against premature determinism. Anneal β toward zero if you eventually want a near-deterministic controller.

Implementation, twice

REINFORCE with reward-to-go, a learned value baseline, advantage normalization, and an entropy bonus, trained on CartPole-v1 through gymnasium, the standard environment API. Both versions are complete scripts: run them and they print the mean episode return per epoch. The PyTorch version leans on torch.distributions for sampling, log-probs, and entropy. The JAX version is written as pure functions over parameter pytrees: the environment-stepping loop stays outside jit because gymnasium is stateful Python, while the entire update, value fit included, is one jitted function.

import gymnasium as gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

def mlp(sizes):
    layers = []
    for i in range(len(sizes) - 1):
        layers.append(nn.Linear(sizes[i], sizes[i + 1]))
        if i < len(sizes) - 2:
            layers.append(nn.Tanh())
    return nn.Sequential(*layers)

def reward_to_go(rewards, gamma):
    """G_t = sum_{t' >= t} gamma^(t'-t) r_t', computed right to left."""
    g, out = 0.0, np.zeros(len(rewards), dtype=np.float32)
    for t in reversed(range(len(rewards))):
        g = rewards[t] + gamma * g
        out[t] = g
    return out

env = gym.make("CartPole-v1")
obs_dim = env.observation_space.shape[0]          # 4
n_act = env.action_space.n                        # 2

policy = mlp([obs_dim, 64, n_act])                # outputs logits
value = mlp([obs_dim, 64, 1])                     # baseline b(s)
opt_pi = torch.optim.Adam(policy.parameters(), lr=1e-2)
opt_v = torch.optim.Adam(value.parameters(), lr=1e-2)
gamma, ent_coef, batch_steps = 0.99, 0.01, 2000

for epoch in range(120):
    obs_buf, act_buf, ret_buf, ep_returns = [], [], [], []
    while len(obs_buf) < batch_steps:             # batch of FULL episodes
        obs, _ = env.reset()
        rewards, done = [], False
        while not done:
            with torch.no_grad():
                logits = policy(torch.as_tensor(obs, dtype=torch.float32))
            a = torch.distributions.Categorical(logits=logits).sample().item()
            obs_buf.append(obs)
            act_buf.append(a)
            obs, r, terminated, truncated, _ = env.step(a)
            rewards.append(r)
            done = terminated or truncated
        ret_buf.extend(reward_to_go(rewards, gamma))
        ep_returns.append(sum(rewards))

    obs_t = torch.as_tensor(np.array(obs_buf), dtype=torch.float32)
    act_t = torch.as_tensor(np.array(act_buf))
    ret_t = torch.as_tensor(np.array(ret_buf))

    v = value(obs_t).squeeze(-1)
    adv = ret_t - v.detach()                      # detach: baseline must not
    adv = (adv - adv.mean()) / (adv.std() + 1e-8) # receive policy gradient

    dist = torch.distributions.Categorical(logits=policy(obs_t))
    logp = dist.log_prob(act_t)
    # ascend E[logp * advantage]; entropy bonus resists early collapse
    pi_loss = -(logp * adv).mean() - ent_coef * dist.entropy().mean()
    v_loss = F.mse_loss(v, ret_t)                 # fit baseline to returns

    opt_pi.zero_grad(); pi_loss.backward(); opt_pi.step()
    opt_v.zero_grad(); v_loss.backward(); opt_v.step()
    print(f"epoch {epoch:3d}  mean episode return {np.mean(ep_returns):7.1f}")

import gymnasium as gym
import numpy as np
import jax
import jax.numpy as jnp
import optax

def init_mlp(key, sizes):
    params = []
    for m, n in zip(sizes[:-1], sizes[1:]):
        key, sub = jax.random.split(key)
        params.append({"w": jax.random.normal(sub, (m, n)) * jnp.sqrt(2.0 / m),
                       "b": jnp.zeros(n)})
    return params

def mlp_apply(params, x):
    for layer in params[:-1]:
        x = jnp.tanh(x @ layer["w"] + layer["b"])
    return x @ params[-1]["w"] + params[-1]["b"]

def reward_to_go(rewards, gamma):
    g, out = 0.0, np.zeros(len(rewards), dtype=np.float32)
    for t in reversed(range(len(rewards))):
        g = rewards[t] + gamma * g
        out[t] = g
    return out

gamma, ent_coef = 0.99, 0.01
pi_opt, v_opt = optax.adam(1e-2), optax.adam(1e-2)

@jax.jit
def action_probs(pi_params, obs):        # tiny jitted forward for sampling
    return jax.nn.softmax(mlp_apply(pi_params, obs))

def pi_loss_fn(pi_params, obs, act, adv):
    logp_all = jax.nn.log_softmax(mlp_apply(pi_params, obs))
    logp = jnp.take_along_axis(logp_all, act[:, None], axis=1)[:, 0]
    entropy = -jnp.sum(jnp.exp(logp_all) * logp_all, axis=1)
    return -jnp.mean(logp * adv) - ent_coef * jnp.mean(entropy)

def v_loss_fn(v_params, obs, ret):
    return jnp.mean((mlp_apply(v_params, obs)[:, 0] - ret) ** 2)

@jax.jit
def update_step(pi_params, v_params, pi_state, v_state, obs, act, ret):
    """One pure update step; everything stochastic stays outside jit."""
    adv = ret - mlp_apply(v_params, obs)[:, 0]    # baseline is a constant
    adv = (adv - adv.mean()) / (adv.std() + 1e-8) # w.r.t. the policy params
    pi_grads = jax.grad(pi_loss_fn)(pi_params, obs, act, adv)
    v_grads = jax.grad(v_loss_fn)(v_params, obs, ret)
    pi_up, pi_state = pi_opt.update(pi_grads, pi_state)
    v_up, v_state = v_opt.update(v_grads, v_state)
    return (optax.apply_updates(pi_params, pi_up), pi_state,
            optax.apply_updates(v_params, v_up), v_state)

env = gym.make("CartPole-v1")
key = jax.random.PRNGKey(0)
pi_params = init_mlp(jax.random.PRNGKey(1), [4, 64, 2])
v_params = init_mlp(jax.random.PRNGKey(2), [4, 64, 1])
pi_state, v_state = pi_opt.init(pi_params), v_opt.init(v_params)
rng = np.random.default_rng(0)

for epoch in range(120):
    obs_buf, act_buf, ret_buf, ep_returns = [], [], [], []
    while len(obs_buf) < 2000:                    # sampling loop: plain Python
        obs, _ = env.reset(seed=int(rng.integers(1 << 30)))
        rewards, done = [], False
        while not done:
            p = np.asarray(action_probs(pi_params, jnp.asarray(obs)))
            a = int(rng.choice(len(p), p=p))
            obs_buf.append(obs)
            act_buf.append(a)
            obs, r, terminated, truncated, _ = env.step(a)
            rewards.append(r)
            done = terminated or truncated
        ret_buf.extend(reward_to_go(rewards, gamma))
        ep_returns.append(sum(rewards))

    pi_params, pi_state, v_params, v_state = update_step(
        pi_params, v_params, pi_state, v_state,
        jnp.asarray(np.array(obs_buf), dtype=jnp.float32),
        jnp.asarray(np.array(act_buf)),
        jnp.asarray(np.array(ret_buf)))
    print(f"epoch {epoch:3d}  mean episode return {np.mean(ep_returns):7.1f}")

Two details are load-bearing. The advantage is detached from the value network (in JAX this happens for free, because the policy gradient is taken only with respect to the policy parameters): the baseline is trained by regression toward the empirical returns, never by the policy objective, or the zero-bias proof above stops applying. And the batch is made of complete episodes, because reward-to-go needs the whole tail of each episode to be observed; truncating mid-episode would silently bias G_t low for the surviving steps.

Using it on a real shape of problem

CartPole-v1 caps episodes at 500 steps with +1 reward per step, so 500 is the ceiling. With the hyperparameters above, expect the mean return to sit near 20 to 40 for the first handful of epochs, cross 100 somewhere around epoch 20 to 40, and reach the 450 to 500 range by epoch 60 to 120. The exact trajectory is seed and machine dependent, and genuinely noisy: policy gradient learning curves lurch, plateau, and occasionally regress before recovering, and a bad seed can double the time to solve. Run three seeds before concluding anything about a change. To watch the trained policy, evaluate it greedily:

env = gym.make("CartPole-v1", render_mode="human")
obs, _ = env.reset(seed=0)
done, total = False, 0.0
while not done:
    with torch.no_grad():
        logits = policy(torch.as_tensor(obs, dtype=torch.float32))
    a = int(torch.argmax(logits))     # greedy at eval time
    obs, r, terminated, truncated, _ = env.step(a)
    total += r
    done = terminated or truncated
print(total)                          # ~500 once trained

env = gym.make("CartPole-v1", render_mode="human")
obs, _ = env.reset(seed=0)
done, total = False, 0.0
while not done:
    logits = mlp_apply(pi_params, jnp.asarray(obs))
    a = int(jnp.argmax(logits))       # greedy at eval time
    obs, r, terminated, truncated, _ = env.step(a)
    total += r
    done = terminated or truncated
print(total)                          # ~500 once trained

One diagnostic worth logging alongside return: the policy entropy. Healthy runs show entropy declining smoothly from ln 2 ≈ 0.69 toward some floor as the policy commits. Entropy crashing to zero in the first few epochs means the entropy coefficient is too small or the learning rate too hot, and the run will likely stall on whatever the policy locked onto first.

Applications

Pure REINFORCE is rarely the production choice for classic control, where its sample inefficiency loses to the methods it spawned. Where it survives, and in fact dominates, is the regime it was always best suited to: a stochastic generator sampling long structured outputs, scored only at the end, with dynamics you cannot differentiate through. That is a description of RLHF. Fine-tuning a language model against a reward model treats each generated token as an action and the reward-model score as terminal reward; InstructGPT and its descendants did this with PPO, which is a policy gradient with a trust-region clamp, and the advantage it feeds on is exactly the G_t − b(s_t) quantity derived above. More recently the field has been walking back toward this page: GRPO, used to train DeepSeek's reasoning models, drops the learned value network entirely and baselines each sampled completion against the mean reward of a group of completions from the same prompt, which is REINFORCE with a Monte Carlo baseline. RLOO and similar leave-one-out estimators are the same idea. Outside of language models, the score-function estimator earns its keep wherever sampling is unavoidable: discrete latent variables in variational inference, neural combinatorial optimization, and the historical neural-architecture-search line of work all optimize non-differentiable sampled structures with exactly this gradient.

Against the real libraries

CleanRL is the reference to read next. Its philosophy is one self-contained file per algorithm variant, so the distance between this page and a research-grade implementation is inspectable line by line: the closest relatives are its single-file PPO and policy gradient scripts, which add the things that matter at scale, namely vectorized environments collecting rollouts in parallel, generalized advantage estimation instead of raw reward-to-go, learning-rate annealing, gradient-norm clipping, and tracked experiment logging. CleanRL's benchmark results are published openly, which makes it the honest yardstick for a from-scratch reimplementation.

stable-baselines3 is the production wrapper: A2C and PPO are its policy gradient members (it deliberately ships no plain REINFORCE, because with the variance reductions above you have already rebuilt most of A2C). What SB3 adds over any from-scratch script is the unglamorous reliability layer: observation normalization wrappers, correct handling of gymnasium's terminated versus truncated distinction when bootstrapping, evaluation callbacks, checkpointing, and years of issue-tracker hardening. When the goal is a controller rather than understanding, three lines of SB3 beat three hundred of yours. Gymnasium sits under all of it as the standard environment API, and its vectorized gym.vector environments are the first upgrade to make when rollout collection becomes the bottleneck. On the JAX side, optax supplies the optimizers used above.

The from-scratch version is enough whenever the environment is cheap and the point is the gradient itself: coursework, estimator research, or embedding a score-function term inside a larger differentiable system. Verification against the libraries is behavioral rather than numerical, since two correct implementations sample different trajectories: fix the environment and seeds, run CleanRL's script and yours for the same number of environment steps across several seeds each, and compare learning-curve envelopes; CartPole should be solved (mean return above 475 over 100 episodes) by both within the same order of magnitude of steps. For a sharper unit test, freeze a tiny policy, compute the analytic policy gradient on a two-state MDP by enumerating all trajectories, and check your estimator's sample mean converges to it as the sample count grows.

Traps and misconceptions

Treating the surrogate loss as a loss. The quantity -(logp * adv).mean() exists only so that its gradient equals the policy gradient at the current parameters. Its value carries no information: it can go up while the policy improves, and driving it toward any particular number means nothing. Monitor episode return and entropy, never the surrogate. Corollary: taking many gradient steps on the same batch is wrong without importance-ratio correction, because after the first step the data is no longer from the current policy; making that reuse safe is precisely the problem PPO solves.

Letting the policy gradient flow into the baseline. If the advantage is not detached, the policy objective trains the value network to minimize the very advantages the policy needs, an unstable arrangement that also voids the unbiasedness proof. The baseline is fit by regression to returns, full stop. The mirror-image mistake is using an action-dependent baseline b(s, a) with the same formula; the proof above requires b to be constant across actions in each state.

Confusing advantage normalization with the baseline. Standardizing advantages per batch is a numerical convenience that rescales the learning rate and is technically a small bias (the batch mean and std depend on the samples); the baseline is the principled variance reduction with an exact zero-bias proof. Normalization does not substitute for a baseline on problems with long horizons and diverse states, and dropping the value network because "normalization already centers things" gives back most of the variance on anything harder than CartPole.

Using the full return where reward-to-go belongs. Multiplying every log-prob by the whole-trajectory return R(τ) is unbiased but noticeably noisier, because each action is charged for rewards it could not have influenced. It is the classic silent bug: the algorithm still learns, just slower, so nothing crashes and nothing looks wrong except the sample budget.

Expecting off-policy data reuse. REINFORCE is strictly on-policy: every gradient estimate must come from trajectories sampled by the current policy, and yesterday's rollouts are garbage the moment the parameters move. This is the deepest cost of the method, and it is the axis along which the value-based family, which replays old experience freely, holds its permanent advantage in sample efficiency.