REINFORCE Gradient of a Softmax Policy

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.