Policy Gradient Loss

The problem, from first principles

rl-alignment-tabular-q-learning learned a Q-VALUE table directly, then acted greedily with respect to it. Policy gradient methods take a different route: train the POLICY (the action-selection strategy itself) directly via gradient ascent on expected return, using a beautifully simple update rule: "increase the log-probability of actions that led to good outcomes, decrease it for actions that led to bad ones." This question builds both pieces: the outcome signal (discounted_returns) and the loss that turns it into a gradient (policy_gradient_loss).

From theory to code

Implement discounted_returns (the REINFORCE return at every timestep, computed efficiently via one backward pass, correctly resetting at episode boundaries) and policy_gradient_loss (the basic policy-gradient loss).

Constraints

• discounted_returns(rewards, gamma, dones)[t] = rewards[t] + gamma * (1 - dones[t]) * discounted_returns[t+1], with the sequence treated as ending after the last index.
• dones[t] = 1 means episode t is the LAST step of an episode: the return computation must NOT carry any value across that boundary into the next episode.
• policy_gradient_loss(log_probs, advantages) returns -mean(log_probs * advantages).
• Increasing log_probs for an action with a POSITIVE advantage must strictly decrease the loss (the correct gradient-ascent direction).

Hints

Open one at a time. Each gives away a little more than the last.