REINFORCE Gradient of a Softmax Policy

mediumFREEPolicy Gradients and Advantage EstimationReinforcement LearningRL 1: Foundations

Problem Statement

Compute the REINFORCE gradient of a softmax policy with respect to its logits, for one recorded episode, with a baseline subtracted.

Background

Policy gradient methods optimise the policy directly. The policy gradient theorem, via the log-derivative trick, says

$\nabla_\theta J(\theta) = \mathbb{E}\big[ G_t\, \nabla_\theta \log \pi_\theta(A_t \mid S_t) \big]$

so all you need is the gradient of a log-probability, scaled by how good the outcome was. For a softmax policy over logits $z$ , that gradient has a famously clean closed form:

$\frac{\partial \log \pi(a \mid s)}{\partial z_i} = \mathbb{1}[i = a] - \pi(i \mid s)$

One-hot minus the probability vector. Read it as an instruction: push the taken action's logit up by 1 - pi(a), and push every other logit down by its own probability. Multiply by the return and you get "increase the logits of actions that preceded high returns". Note that these components sum to zero — a softmax gradient can only redistribute probability mass.

A baseline $b(s_t)$ subtracted from the return leaves the gradient unbiased (because $\mathbb{E}[\nabla \log \pi] = 0$ ) while cutting its variance, which is the difference between REINFORCE working and not.

$\hat{g} = \sum_{t=0}^{T-1} \big(G_t - b(s_t)\big) \big(\mathbb{1}_{A_t} - \pi_t\big)$

Your Task

Implement:

def softmax_probs(logits):
    ...

def reinforce_gradient(logits_seq, actions, rewards, gamma, baseline):
    ...

logits_seq[t] — a list of n_actions floats, the policy logits at step t.
actions[t] — the action taken at step t.
rewards[t] — the reward received at step t.
gamma — discount factor. Compute $G_t$ , the discounted return from step t onward.
baseline[t] — the baseline value subtracted from $G_t$ .

softmax_probs returns a list of n_actions floats. reinforce_gradient returns the summed gradient as a list of n_actions floats.

Input / Output Format

Nested lists of floats in; lists of floats out, rounded to 4 decimals by the grader.

Sample

logits = [[0.0, 0.0]]
print([round(p, 4) for p in softmax_probs(logits[0])])
print([round(g, 4) for g in reinforce_gradient(logits, [0], [2.0], 0.9, [0.0])])

Output:

[0.5, 0.5]
[1.0, -1.0]

The policy is uniform, so the log-prob gradient is [1, 0] - [0.5, 0.5] = [0.5, -0.5], scaled by the return 2.0.

Example

Input

softmax_probs([0.0, 0.0]) and reinforce_gradient([[0.0, 0.0]], [0], [2.0], 0.9, [0.0])

Output

[0.5, 0.5]
[1.0, -1.0]

Reasoning

Equal logits give a uniform policy [0.5, 0.5]. The log-probability gradient for the taken action 0 is one-hot minus the probabilities, [1, 0] - [0.5, 0.5] = [0.5, -0.5]. The single-step return is 2.0 with a zero baseline, so the gradient is 2.0 * [0.5, -0.5] = [1.0, -1.0].

Constraints

1 <= T <= 500, 2 <= n_actions <= 50
Subtract the max logit before exponentiating for numerical stability.
G_t is the discounted return from step t to the end of the episode.
The returned gradient components sum to (numerically) zero.
0.0 <= gamma <= 1.0
Do not round inside the functions.

Learn this in

Reinforcement LearningML Study Plan Policy Gradient MethodsReinforcement Learning Actor-Critic & PPOReinforcement Learning

More from RL 1: Foundations

GAE Over a Full Trajectoryhard PPO Clipped Objective (Single Sample)medium TD Residuals for GAEmedium Return Standardization Baselinemedium REINFORCE Loss Termmedium

Concept primer

1. Background Knowledge

Policy Gradient methods optimize the policy directly by estimating the gradient of the expected return with respect to the policy parameters. The REINFORCE algorithm is a foundational Monte Carlo policy gradient method. It relies on the log-derivative trick, which allows us to express the gradient of the expected reward as an expectation of the gradient of the log-probability of the actions taken, weighted by the return. This transforms an intractable expectation over trajectories into a sample-based estimate.

A critical component of REINFORCE is the baseline. Subtracting a state-dependent baseline $b(s_t)$ from the return $G_t$ reduces the variance of the gradient estimator without introducing bias. This is because the expected gradient of the log-probability under the policy is zero ( $\mathbb{E}[\nabla_\theta \log \pi_\theta(A_t|S_t)] = 0$ ). By centering the returns around the baseline, we ensure that only actions yielding returns significantly better or worse than the baseline contribute strongly to the gradient update, leading to more stable learning.

For a softmax policy defined by logits $z$ , the probability of action $a$ is $\pi(a|s) = \frac{e^{z_a}}{\sum_j e^{z_j}}$ . The gradient of the log-probability with respect to the logits has a clean, closed-form expression: $\frac{\partial \log \pi(a|s)}{\partial z_i} = \mathbb{1}[i=a] - \pi(i|s)$ . This means the gradient vector is the difference between a one-hot vector for the taken action and the probability distribution itself. This structure ensures that the gradient components sum to zero, effectively redistributing probability mass rather than shifting the entire distribution uniformly.

2. Algorithm Approach

The solution involves two main functions: softmax_probs and reinforce_gradient.

Softmax Calculation: Implement a numerically stable softmax function. Directly computing exponentials can lead to overflow. The standard approach is to subtract the maximum logit from all logits before exponentiating. This shifts the values without changing the resulting probabilities, keeping the exponentials within a manageable range.
Return Computation: Calculate the discounted return $G_t$ for each time step $t$ . The return is the sum of future rewards discounted by $\gamma$ . This can be computed efficiently by iterating backwards from the last time step, accumulating the discounted rewards.
Gradient Aggregation: For each time step, compute the policy gradient component $(\mathbb{1}_{A_t} - \pi_t)$ and scale it by the advantage $(G_t - b(s_t))$ . Sum these scaled gradients across all time steps to get the final gradient vector.

3. Step-by-Step Strategy

Implement softmax_probs:

Find the maximum value in the input logits list.
Subtract this maximum from each logit to ensure numerical stability.
Compute the exponential of each adjusted logit.
Sum these exponentials to get the denominator.
Divide each exponential by the sum to get the probabilities.
Return the list of probabilities.

Implement reinforce_gradient:

Initialize a gradient list of zeros with length equal to the number of actions.
Compute the discounted returns $G_t$ for all $t$ . A backward pass is efficient: start with $G_T = 0$ (or the last reward if defined), and for $t = T-1$ down to $0$ , compute $G_t = \text{rewards}[t] + \gamma \cdot G_{t+1}$ .
Iterate through each time step $t$ from $0$ to $T-1$ :
Get the logits for step $t$ and compute the probability distribution $\pi_t$ using softmax_probs.
Create a one-hot vector for the action taken at step $t$ .
Compute the log-probability gradient: one_hot - pi_t.
Compute the advantage: advantage = G_t - baseline[t].
Scale the log-probability gradient by the advantage.
Add this scaled gradient to the cumulative gradient vector.
Return the final gradient vector.

4. Common Pitfalls

Numerical Instability in Softmax: Failing to subtract the maximum logit before exponentiating can lead to inf or NaN values, especially with large logits. Always use the shift trick: $e^{z_i - \max(z)}$ .
Incorrect Return Calculation: Ensure the discount factor $\gamma$ is applied correctly. $G_t$ includes the reward at time $t$ and all future rewards. A common error is off-by-one errors in the backward accumulation loop.
Gradient Summation: Remember that the final gradient is the sum of the gradients over all time steps in the episode. Do not average them unless specified (REINFORCE typically sums).
One-Hot Vector Construction: Ensure the one-hot vector has a 1 at the index of the taken action and 0 elsewhere. Mixing up indices can lead to incorrect gradients.
Baseline Subtraction: Ensure the baseline is subtracted from the return $G_t$ , not the reward $r_t$ . The baseline is state-dependent and should be subtracted from the total return from that state onward.

5. Time & Space Complexity

Time Complexity:
softmax_probs: $O(N)$ where $N$ is the number of actions, due to finding the max, exponentiating, and summing.
reinforce_gradient: Let $T$ be the episode length and $N$ be the number of actions. Computing returns takes $O(T)$ . The main loop runs $T$ times, and each iteration involves a softmax ( $O(N)$ ) and vector operations ( $O(N)$ ). Thus, the total time complexity is $O(T \cdot N)$ .
Space Complexity:
softmax_probs: $O(N)$ to store the probabilities.
reinforce_gradient: $O(T)$ to store the returns $G_t$ (if precomputed) or $O(1)$ if computed on the fly in a backward pass. The gradient vector takes $O(N)$ space. Overall, $O(T + N)$ or $O(N)$ depending on implementation details.

solution.py

Python 3.13

Test results