Problems
Loading...
1 / 1

Policy Gradient Loss

Reinforcement Learning
Medium

The REINFORCE algorithm uses the policy gradient theorem to directly optimize a stochastic policy. The loss is constructed so that gradient descent increases the probability of actions that led to higher-than-average returns and decreases the probability of actions that led to lower-than-average returns.

Given the log-probabilities of the actions taken, the rewards at each timestep, and a discount factor gamma, compute the policy gradient loss with a mean-return baseline.

Algorithm

  1. Compute discounted returns backward:
GT1=rT1,Gt=rt+γGt+1G_{T-1} = r_{T-1}, \quad G_t = r_t + \gamma \cdot G_{t+1}
  1. Subtract the mean return as a baseline to reduce variance:
At=GtGˉ,where Gˉ=1Tt=0T1GtA_t = G_t - \bar{G}, \quad \text{where } \bar{G} = \frac{1}{T} \sum_{t=0}^{T-1} G_t
  1. Compute the loss (negative because we want gradient ascent on expected return):
L=1Tt=0T1logπ(atst)AtL = -\frac{1}{T} \sum_{t=0}^{T-1} \log \pi(a_t | s_t) \cdot A_t
Loading visualization...

Examples

Input:

log_probs = [-1.0, -2.0, -0.5], rewards = [1, 1, 1], gamma = 1.0

Output:

0.1667

Returns: G = [3, 2, 1]. Mean = 2. Advantages = [1, 0, -1]. Loss = -(-11 + -20 + -0.5*(-1))/3 = -(- 1 + 0 + 0.5)/3 = 0.5/3.

Input:

log_probs = [-1.0, -0.5], rewards = [0, 10], gamma = 0.0

Output:

-1.25

With gamma = 0, returns equal immediate rewards: G = [0, 10]. Mean = 5. Advantages = [-5, 5]. Loss = -(-1*(-5) + -0.5*5)/2 = -(5 - 2.5)/2 = -1.25.

Hint 1

First compute returns G backward (G[T-1] = rewards[T-1], G[t] = rewards[t] + gamma * G[t+1]). Then compute mean_G = sum(G) / T. Then advantages = [g - mean_G for g in G]. Finally loss = -sum(lp * a for lp, a in zip(log_probs, advantages)) / T.

Hint 2

The negative sign is critical. Without it, gradient descent would decrease the expected return instead of increasing it.

Requirements

  • Compute discounted returns using backward recursion
  • Subtract the mean return from each return to get advantages
  • Compute loss as the negative mean of (log_prob * advantage)
  • Return a single float

Constraints

  • log_probs and rewards have the same length (at least 1)
  • log_probs contains negative floats (log of probabilities)
  • 0 <= gamma <= 1
  • Return a single float
  • Time limit: 300 ms
Try Similar Problems