Policy Gradient Methods

Policy gradient directly optimizes the policy by computing ∇log π(a|s) * return. Simpler than value-based methods but higher variance.

The problem, from first principles

Value-based (DQN): learn V or Q, derive policy as argmax. Policy-based: directly parameterize policy π(a|s), optimize with gradient ∇log π(a|s) * R.

From theory to code

Implement policy_gradient_loss(log_probs, advantages) which:

• Takes log probabilities of actions taken
• Takes advantage estimates (reward - baseline)
• Returns negative mean of log_prob * advantage

Constraints

• Negative loss (gradient ascent toward high-reward actions)
• Element-wise multiply log_probs and advantages
• Average across batch

Hints