REINFORCE Policy Gradient - Reinforcement Learning

REINFORCE Policy Gradient

What you'll build

The REINFORCE (Monte Carlo policy gradient) estimate ∇_θ J for a tabular softmax policy, from one recorded episode.

Background

Value-based methods learn Q and act greedily; policy-gradient methods adjust the policy's parameters directly in the direction that makes high-return actions more likely. REINFORCE (Williams, 1992) is the simplest such method and the ancestor of actor-critic, PPO and the RL used to fine-tune LLMs.

Formula

Policy gradient theorem, estimated from one episode:

∇_θ J ≈ Σ_t G_t · ∇_θ log π(a_t | s_t)
G_t   = r_t + γ·r_{t+1} + γ²·r_{t+2} + ...      (return from step t to the end)

Tabular softmax policy, with θ of shape (n_states, n_actions):

π(·|s) = softmax(θ[s])
∇_{θ[s]} log π(a|s) = one_hot(a) - π(·|s)       (zero for every other row)

Algorithm

1. Compute G_t for every step with one backward pass: G = r_t + γ·G.
2. For each step t with state s_t and action a_t: add G_t · (one_hot(a_t) - softmax(θ[s_t])) to row s_t of the gradient.
3. Return the accumulated gradient.

Inputs

• theta: (n_states, n_actions) float array of logits
• episode: list of (state, action, reward) tuples, where reward is the reward received after taking that action
• gamma: discount factor

Expected return format

A NumPy array with the same shape as theta, rounded to 2 decimals. It is the ascent direction (gradient of J, not of a loss): θ ← θ + lr · gradient improves the policy.

Edge cases and conventions

• An empty episode gives an all-zero gradient.
• A state visited several times accumulates one contribution per visit; unvisited rows stay 0.
• Use the softmax with the max subtracted (the provided helper) for stability.

Constraints

NumPy only. Compute probabilities from theta as given; don't update theta during the episode.

Follow-up

• Subtracting a baseline b(s) from G_t keeps the estimate unbiased and lowers variance. Why is the bias still zero?
• Why does this estimator have high variance, and how do GAE and PPO address it?

Input: theta = zeros((2, 2)), episode = [(0, 1, 1.0), (1, 0, 1.0)], gamma = 0.9
Output: [[-0.95, 0.95], [0.5, -0.5]]

Input: theta = zeros((2, 2)), episode = [(0, 0, 5.0)], gamma = 0.9
Output: [[2.5, -2.5], [0.0, 0.0]]

import numpy as np

def softmax(x):
    """Compute softmax probabilities."""
    exp_x = np.exp(x - np.max(x))
    return exp_x / np.sum(exp_x)

def reinforce_gradient(theta: np.ndarray, episode: list,
                       gamma: float) -> np.ndarray:
    """
    Compute the REINFORCE policy gradient for a tabular softmax policy.

    Args:
        theta: (n_states, n_actions) policy logits
        episode: list of (state, action, reward) tuples
        gamma: discount factor

    Returns:
        (n_states, n_actions) gradient of sum_t G_t * log pi(a_t|s_t),
        rounded to 2 decimals
    """
    theta = np.array(theta, dtype=float)
    n_states, n_actions = theta.shape
    gradient = np.zeros_like(theta)

    # Your code here
    pass