Policy Gradient Computation via REINFORCE

In reinforcement learning, policy gradient methods represent a powerful family of algorithms that directly optimize the policy parameters to maximize expected cumulative rewards. Unlike value-based methods that learn action values and derive policies indirectly, policy gradient approaches adjust parameters in the direction that increases the probability of high-reward trajectories.

The REINFORCE algorithm (REward Increment = Nonnegative Factor × Offset Reinforcement × Characteristic Eligibility) is the foundational Monte Carlo policy gradient method. It uses complete episode trajectories to estimate gradients, making it conceptually simple yet surprisingly effective for many problems.

Policy Representation

The policy is parameterized by a 2D array θ (theta) of shape (num_states × num_actions). For each state s, the probability of selecting action a is computed using the softmax function over the corresponding row:

$$\pi_\theta(a|s) = \frac{e^{\theta_{s,a}}}{\sum_{a'} e^{\theta_{s,a'}}}$$

This ensures that action probabilities for each state sum to 1 and are always positive.

The Policy Gradient Theorem

The REINFORCE algorithm leverages the policy gradient theorem, which states that the gradient of the expected return with respect to policy parameters can be written as:

$$ abla_\theta J(\theta) = \mathbb{E}\pi \left[ \sum{t=0}^{T-1} G_t \cdot abla_\theta \log \pi_\theta(a_t|s_t) \right]$$

Where:

• Gₜ is the return (cumulative future reward) from time step t onward: $G_t = \sum_{k=t}^{T-1} r_k$
• ∇θ log πθ(aₜ|sₜ) is the score function (gradient of log-probability of the action taken)

Score Function for Softmax Policy

For a softmax policy, the gradient of the log-probability with respect to θ has a particularly elegant form. For the action a taken in state s:

$$ abla_{\theta_{s,a'}} \log \pi_\theta(a|s) = \mathbb{1}[a' = a] - \pi_\theta(a'|s)$$

This means:

• For the chosen action: gradient = (1 - π(a|s))
• For other actions: gradient = -π(a'|s)

The gradient is localized to state s; all other rows of θ have zero gradient for this transition.

Your Task

Implement the REINFORCE policy gradient estimator. Given the policy parameters θ and a list of episodes (each containing state-action-reward transitions), compute the average gradient across all episodes. Each transition contributes its return (sum of rewards from that step to the end of the episode) multiplied by the score function gradient.

Algorithm Steps:

1. For each episode, compute the return Gₜ for each time step t
2. Compute the softmax policy probabilities for each visited state
3. Calculate the score function gradient for each transition
4. Weight each gradient by its corresponding return Gₜ
5. Sum gradients across all transitions and average across episodes

Return Format: A 2D list of the same shape as θ, with values rounded to 2 decimal places.