REINFORCE Algorithm

The Policy Gradient Theorem provides a way to calculate the gradient of the expected return with respect to the policy parameters $\theta$. To apply this theory, a practical algorithm is required to estimate and use this gradient. The REINFORCE algorithm, also known as Monte Carlo Policy Gradient, offers precisely such an approach. It learns directly from complete episodes of experience, sampled using the current policy $\pi_\theta$.

The core idea is elegant: we want to adjust the policy parameters $\theta$ such that actions leading to higher overall returns become more probable, and actions leading to lower returns become less probable. REINFORCE achieves this by sampling trajectories and using the actual returns experienced within those trajectories to guide the parameter updates.

Recall the policy gradient theorem:

$$\nabla_\theta J(\theta) = \mathbb{E}_{\tau \sim \pi_\theta} \left[ \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t | s_t) G_t \right]$$

Here, $J(\theta)$ is the expected total return, $\tau$ represents a trajectory (episode) sampled according to the policy $\pi_\theta$, $G_t$ is the total discounted return starting from time step $t$, and $\nabla_\theta \log \pi_\theta(a_t | s_t)$ is the gradient of the log-probability of taking action $a_t$ in state $s_t$. This term, often called the "score function", indicates the direction in parameter space that increases the probability of the action $a_t$ taken at state $s_t$.

Since we usually cannot compute the expectation $\mathbb{E}_{\tau \sim \pi_\theta}[\cdot]$ directly (as it requires integrating over all possible trajectories), REINFORCE uses Monte Carlo estimation. We run the current policy $\pi_\theta$ to generate one or more complete episodes. For each episode, we calculate the return $G_t$ for every time step $t$. Then, we approximate the gradient using the sampled data.

The REINFORCE Algorithm

The algorithm proceeds as follows:

1. Initialize: Start with randomly initialized policy parameters $\theta$.
2. Loop (for each episode): a. Generate Trajectory: Run the policy $\pi_\theta$ in the environment to collect a full episode: $\tau = (s_0, a_0, r_1, s_1, a_1, r_2, ..., s_{T-1}, a_{T-1}, r_T)$. Here, $T$ is the terminal time step of the episode. b. Loop (for each time step $t$ from $0$ to $T-1$): i. Calculate Return: Compute the return from time step $t$ onwards: $G_t = \sum_{k=t+1}^{T} \gamma^{k-t-1} r_k$. Note that $G_t$ is the actual, observed discounted return following state $s_t$ and action $a_t$ in this specific trajectory. ii. Calculate Update: Determine the update direction for this step: $\nabla_\theta \log \pi_\theta(a_t | s_t) G_t$. c. Update Parameters: Update the policy parameters using the accumulated gradients from the episode, typically using stochastic gradient ascent: $\theta \leftarrow \theta + \alpha \sum_{t=0}^{T-1} \nabla_\theta \log \pi_\theta(a_t | s_t) G_t$ where $\alpha$ is the learning rate. The update can also be performed after each step or after a batch of episodes.

Understanding the Update

Let's break down the update step $\theta \leftarrow \theta + \alpha \nabla_\theta \log \pi_\theta(a_t | s_t) G_t$.

• $\nabla_\theta \log \pi_\theta(a_t | s_t)$ (Score Function): This vector points in the direction in parameter space ($\theta$) that most increases the log-probability of the action $a_t$ that was actually taken in state $s_t$ during the sampled episode.
• $G_t$ (Return): This scalar value represents the quality of the trajectory following the state-action pair $(s_t, a_t)$. It acts as a weighting factor for the score function.
  • If $G_t$ is large and positive, the update pushes $\theta$ significantly in the direction that makes $a_t$ more likely in state $s_t$. The algorithm "reinforces" this action because it led to a good outcome.
  • If $G_t$ is large and negative (or small and positive), the update still pushes $\theta$ in the direction indicated by the score function, but the magnitude is smaller, or it might even push away from that direction if $G_t$ is negative. The algorithm weakly reinforces or even discourages this action because it led to a poor outcome.

Essentially, REINFORCE correlates the direction of steepest ascent for the action probability ($\nabla_\theta \log \pi_\theta$) with the return ($G_t$) that followed that action. Actions associated with higher returns get "boosted" more strongly.

The following diagram outlines the information flow within a single REINFORCE update step based on a sampled episode.

Flow diagram illustrating how information from a sampled trajectory is used to calculate the return and score function, which are then combined to update the policy parameters in REINFORCE.

Strengths and Weaknesses

REINFORCE offers several advantages:

• Simplicity: The core concept and update rule are relatively straightforward.
• Direct Optimization: It directly optimizes the policy for the objective function (expected return).
• Compatibility: Works naturally with both discrete and continuous action spaces (by parameterizing appropriate probability distributions). It can also learn optimal stochastic policies.

However, it also has significant drawbacks:

• High Variance: The Monte Carlo estimation of the return $G_t$ based on single trajectories can be extremely noisy. Different episodes can lead to significantly different returns, even if the policy changes slightly, resulting in high variance in the gradient estimate $\nabla_\theta J(\theta)$. This high variance makes learning unstable and slow.
• Sample Inefficiency: It typically requires many episodes to converge because of the noisy gradient estimates.
• Requires Complete Episodes: Updates are performed only after an episode finishes, which can be problematic in very long or non-terminating tasks.

The high variance is the most pressing issue with the basic REINFORCE algorithm. In the next section, we will discuss how introducing baselines can help mitigate this problem significantly, leading to faster and more stable learning.

Implementation Approaches

When implementing REINFORCE, especially with neural networks:

• Policy Representation: The policy $\pi_\theta(a|s)$ is typically represented by a neural network that takes the state $s$ as input.
  • For discrete actions, the network usually outputs probabilities for each action (e.g., using a softmax output layer). $\log \pi_\theta(a_t | s_t)$ is then the log of the output probability corresponding to the action $a_t$ that was taken.
  • For continuous actions, the network might output the parameters of a probability distribution (e.g., the mean $\mu$ and standard deviation $\sigma$ of a Gaussian distribution). The action $a_t$ is sampled from $\mathcal{N}(\mu(s_t; \theta), \sigma(s_t; \theta))$, and $\log \pi_\theta(a_t | s_t)$ is the log-probability density function (PDF) of the chosen action under that distribution.
• Gradient Calculation: Deep learning libraries (like TensorFlow or PyTorch) handle the calculation of $\nabla_\theta \log \pi_\theta(a_t | s_t)$ via automatic differentiation. You typically define a loss function whose gradient matches the policy gradient update. A common approach is to use a surrogate loss like $- \log \pi_\theta(a_t | s_t) G_t$. Minimizing this loss using gradient descent is equivalent to performing gradient ascent on the objective $J(\theta)$ using the term $\nabla_\theta \log \pi_\theta(a_t | s_t) G_t$. Remember that $G_t$ is treated as a constant factor when computing the gradient for the update at step $t$.