What is Policy Gradient? The REINFORCE Algorithm Explained

Policy Gradient is a way to train an AI agent by directly making the actions that led to good results more likely, and the actions that led to bad results less likely. REINFORCE is the simplest algorithm that does this.

In this blog, we will learn about Policy Gradient and the REINFORCE algorithm. We will also see how it differs from value-based methods, how a neural network acts as a policy, why the formula has a log probability in it, how REINFORCE works step by step with a small worked example, how a baseline fixes its high variance, and how all of this connects to training LLMs.

We will cover the following:

• Value-based vs Policy-based methods
• What is a Policy?
• What is Policy Gradient?
• The Policy Gradient Formula
• Why does the log probability appear?
• The REINFORCE Algorithm
• A Worked Example
• The Problem of High Variance
• Baseline Subtraction
• Connection to LLM Training

I am Amit Shekhar, Founder @ Outcome School, I have taught and mentored many developers, and their efforts landed them high-paying tech jobs, helped many tech companies in solving their unique problems, and created many open-source libraries being used by top companies. I am passionate about sharing knowledge through open-source, blogs, and videos.

I teach AI and Machine Learning at Outcome School.

Let's get started.

In Reinforcement Learning, an agent looks at the current situation, called the state, takes an action, and gets a reward. One full game, from start to end, is called an episode. The total reward the agent collects from a step till the end of the episode is called the return. The goal of the agent is to collect as much return as possible. We have a detailed blog on Markov Decision Process (MDP) that explains how states, actions, rewards, and the return fit together.

Value-based vs Policy-based methods

There are two main families of methods to teach an agent.

Value-based methods learn a score for every action in every state. For example, Q-learning learns "how good is it to take this action in this state". Then, the agent picks the action with the highest score. Here, the agent decides indirectly, through the scores.

Policy-based methods skip the scores. They directly learn "what action should I take in this state". Here, the agent decides directly.

Policy Gradient belongs to the second family.

Let's see this as a diagram:

Value-based                  Policy-based
┌─────────────────────┐      ┌─────────────────────┐
│ State               │      │ State               │
└─────────────────────┘      └─────────────────────┘
           ↓                            ↓
┌─────────────────────┐      ┌─────────────────────┐
│ Score each action   │      │ Probability of each │
│ Left:  4.2          │      │ Left:  0.7          │
│ Right: 1.5          │      │ Right: 0.3          │
└─────────────────────┘      └─────────────────────┘
           ↓                            ↓
┌─────────────────────┐      ┌─────────────────────┐
│ Pick the highest    │      │ Sample one          │
│ Action: Left        │      │ Action: Left (70%)  │
└─────────────────────┘      └─────────────────────┘

Here, we can see that the value-based method picks the action with the highest score, while the policy-based method samples an action from the probabilities.

Policy-based methods are useful when there are too many possible actions to score one by one, or when the best behavior is to pick between actions with some probability, instead of always picking the same one.

What is a Policy?

A policy is the rule that tells the agent which action to take in a given state.

In Policy Gradient, the policy is a neural network. It takes the state as input and gives a probability for each possible action as output.

Let's say our agent is playing a simple game where it can move either Left or Right. Let's see how the policy network picks an action as a diagram:

┌─────────────────────┐
│ State (game screen) │
└─────────────────────┘
           ↓
┌─────────────────────┐
│ Policy network      │
│ (its weights)       │
└─────────────────────┘
           ↓
┌─────────────────────┐
│ Raw scores          │
│ Left:  0.85         │
│ Right: 0.00         │
└─────────────────────┘
           ↓
┌─────────────────────┐
│ Softmax             │
│ Left:  0.7          │
│ Right: 0.3          │
└─────────────────────┘
           ↓
┌─────────────────────┐
│ Sample one action   │
│ Action: Left        │
└─────────────────────┘

Here, we can see that the network looks at the game screen and first produces a raw score for each action. Then, a softmax function turns these raw scores into probabilities that add up to 1. In our example, the raw scores 0.85 and 0.00 become 0.7 and 0.3.

Then, the agent samples an action from these probabilities. Means, it picks Left 70% of the time and Right 30% of the time. This randomness lets the agent keep trying different actions.

We write the policy as pi(a|s), which we read as "the probability of taking action a in state s". The parameters of the network, which are its weights, are written as theta.

Now, the question is: how do we change theta so that the agent gets better? Here comes the Policy Gradient into the picture.

What is Policy Gradient?

Policy Gradient is a method that trains the policy network directly, by changing its weights in the direction that increases the expected return.

Here, the expected return is the average return the agent gets over many episodes.

The central idea of Policy Gradient is:

• If an action led to a high return, make that action more likely in the future.
• If an action led to a low return, make that action less likely in the future.

Let's take an example. Suppose the agent was in some state, picked Left, and finally won the game. Then, we nudge the weights so that next time, in the same state, the probability of Left goes up a little.

If the agent picked Left and lost the game, we nudge the weights so that the probability of Left goes down a little.

Let's understand the word Gradient in the name. A gradient tells us the direction in which we must change the weights to increase something. Here, we want to increase the expected return.

One more thing to notice. In normal neural network training, we go down the gradient to reduce the loss. Here, we go up the gradient to increase the return. This is called gradient ascent.

Planning your move into AI?

No matter which tech domain you work in today, with your effort, you can transition to these roles:

• AI Engineer
• Forward Deployed Engineer
• AI Agent Engineer
• Harness Engineer
• Voice AI Engineer
• AI Architect
• LLM Engineer
• LLM Inference Engineer
• Gen AI Engineer
• MLOps Engineer
• ML Engineer
• Deep Learning Engineer

Our AI and Machine Learning Program at Outcome School takes you there.

The Policy Gradient Formula

We cannot calculate the exact gradient of the expected return, because that would need every possible game. So, we estimate it from the games the agent actually plays. For one step, the estimate is as below:

gradient of J(theta) is estimated by: gradient of log pi(a|s) * G

Here,

• J(theta) is the expected return of the policy. We want to make it as large as possible.
• pi(a|s) is the probability that the policy gave to the action a it actually took in state s.
• log pi(a|s) is the log of that probability.
• gradient of log pi(a|s) is the direction in which we must change the weights to make that action more likely.
• G is the return the agent got after taking that action.

Then, we update the weights like below:

theta = theta + learning_rate * gradient of log pi(a|s) * G

Let's read this formula in plain words.

"Push the weights in the direction that makes this action more likely, and push harder if the return was bigger."

• If G is large and positive, we push strongly towards this action.
• If G is small, we push only a little.
• If G is negative, the push goes in the opposite direction, and the action becomes less likely.

This is how the formula captures the central idea.

Note: For a full episode, we add this term for every step of the episode, each step with its own return. Over many episodes, the average of these estimates points in the true gradient direction.

Why does the log probability appear?

A natural question is: why log pi(a|s) and not just pi(a|s)?

We will skip the full derivation and focus on the intuition. From calculus, the gradient of log pi(a|s) is equal to the gradient of pi(a|s) divided by pi(a|s).

gradient of log pi(a|s) = gradient of pi(a|s) / pi(a|s)

The agent samples actions from its own policy. So, actions that already have a high probability get picked very often, and they get updated again and again. Actions with a low probability get picked rarely.

If we used only the gradient of pi(a|s), the frequent actions would get far more updates just because they are picked more often, and not because they are better.

Dividing by pi(a|s) balances this. A rare action gets a bigger push when it is picked, and a common action gets a smaller push each time. For example, if an action has a probability of 0.1, its push is multiplied by 1 / 0.1 = 10. If an action has a probability of 0.9, its push is multiplied by only 1 / 0.9, which is about 1.1. This way, every action is judged by its return, and not by how often it was picked.

There is one more benefit. Probabilities of many steps get multiplied together, and the product becomes a very tiny number. The log turns this multiplication into addition, which is much easier for a computer to handle. We have a detailed blog on Log Probabilities (Logprobs) in LLMs that explains how this works step by step.

A quick note for you

No matter which tech domain you work in, get familiar with these topics:

• LLM
• RAG
• MCP
• Agent
• Fine-tuning
• Quantization

We put it all together in one video:

AI Engineering Explained: LLM, RAG, MCP, Agent, Fine-Tuning, and Quantization

No need to stop reading - bookmark it and watch later when you get time. Future you will thank you.

Now, let's get back to the topic.

The REINFORCE Algorithm

REINFORCE is the simplest Policy Gradient algorithm. It was introduced by Ronald Williams in 1992.

It repeats three steps again and again.

Step 1: Play an episode. The agent plays one full game using the current policy. We record every state, every action, and every reward.

Step 2: Compute the returns. For every step t of the episode, we calculate the return G_t, which is the total reward from that step till the end of the episode.

G_t = r_t + gamma * G_(t+1)

Here, r_t is the reward at step t, and gamma is a number between 0 and 1 called the discount factor. It makes rewards that come far in the future count a little less.

For example, suppose an episode has 3 steps with rewards 1, 0, and 10, and gamma is 0.9. We calculate the returns from the last step backwards:

G_3 = 10
G_2 = 0 + 0.9 * 10  = 9.0
G_1 = 1 + 0.9 * 9.0 = 9.1

Step 3: Update the policy. For every step, we compute gradient of log pi(a_t|s_t) * G_t, add them all up, and move the weights a small step in that direction.

Then, we go back to Step 1 with the updated policy.

Let's see the REINFORCE loop as a diagram:

┌──────────────────────────────┐
│ Step 1: Play one episode     │
│ Record every state,          │
│ action, and reward           │
└──────────────────────────────┘
                ↓
┌──────────────────────────────┐
│ Step 2: Compute returns      │
│ G_t for every step t         │
└──────────────────────────────┘
                ↓
┌──────────────────────────────┐
│ Step 3: Update the policy    │
│ theta = theta + lr * sum of  │
│ grad log pi(a_t|s_t) * G_t   │
└──────────────────────────────┘
                ↓
┌──────────────────────────────┐
│ Updated policy               │
│ Go back to Step 1            │
└──────────────────────────────┘

Here, we can see that one round of the loop uses exactly one episode. The policy plays, we score every step with its return, we update the weights, and the updated policy plays the next episode.

REINFORCE is called a Monte Carlo method. Means, it waits for the full episode to finish and uses the real returns. It does not use a second model to guess any values.

To learn Reinforcement Learning and how an agent improves its policy from rewards from the ground up, check out our AI and Machine Learning Program at Outcome School.

A Worked Example

The best way to learn this is by taking an example.

Let's take the simplest possible policy. There is only one state, and the agent has two actions: Left and Right. Just for the sake of understanding, instead of a full neural network, our policy is only two raw scores, one for each action, and softmax turns them into probabilities. These two scores are our theta.

At the start, both scores are 0. So, the probabilities are 0.5 and 0.5.

Now, the agent picks Left and gets a return G = 10. We use a learning rate of 0.1.

For a softmax policy, the gradient of log pi(Left) with respect to the scores has a simple form:

gradient of log pi(Left) = one_hot(Left) - probabilities

Here, one_hot(Left) is [1, 0], which means 1 for the action that was picked and 0 for the other. So, the gradient is [1, 0] - [0.5, 0.5] = [0.5, -0.5].

Let's see the code for this one update step:

import numpy as np

def softmax(z):
    e = np.exp(z - z.max())
    return e / e.sum()

theta = np.array([0.0, 0.0])   # scores for [Left, Right]
probs = softmax(theta)
print("Before:", probs.round(3))

action = 0                     # the agent picked Left
G = 10                         # the return it got
lr = 0.1                       # learning rate

grad_log_pi = np.eye(2)[action] - probs   # gradient of log pi(Left)
theta = theta + lr * G * grad_log_pi
print("Gradient:", grad_log_pi.round(3))
print("After:", softmax(theta).round(3))

It will print the following:

Before: [0.5 0.5]
Gradient: [ 0.5 -0.5]
After: [0.731 0.269]

Here, we can see that:

• np.eye(2)[action] gives [1, 0], which is the one-hot vector for Left.
• Inside softmax, we subtract the largest score only to avoid very large numbers. The result stays the same.
• The update is 0.1 * 10 * [0.5, -0.5] = [0.5, -0.5].
• The new scores become [0.5, -0.5].
• After softmax, the probability of Left goes up from 0.5 to 0.731.

The agent picked Left, got a high return, and now Left is more likely. This is exactly the central idea of Policy Gradient.

Now, suppose the return was G = -10 instead. The update becomes [-0.5, 0.5], and the new probabilities become 0.269 for Left and 0.731 for Right. Left becomes less likely. It works perfectly.

Stay updated: Subscribe to our newsletter to get our latest AI and Machine Learning blogs straight to your inbox.

The Problem of High Variance

REINFORCE works, but it has a big problem: high variance. Variance means how much the result jumps around from one try to another.

The return G comes from one full episode. One episode depends on many random action choices and many random events in the game. So, the same action in the same state can get a return of 50 in one episode and 5 in the next one, just because of luck later in the game.

Since the update is multiplied by G, the updates also jump around a lot. The policy gets pushed one way, then another way. Learning becomes slow and unstable, and we need a very large number of episodes.

There is one more issue. Suppose all the returns in our game are positive, for example between 8 and 12. Then, every action that was taken gets pushed up, even the bad ones. The policy still improves, because good actions get pushed a little more, but very slowly.

Baseline Subtraction

So, here comes the baseline to the rescue.

Instead of multiplying by G, we multiply by G - b, where b is a baseline. A simple choice for the baseline is the average return of recent episodes.

theta = theta + learning_rate * gradient of log pi(a|s) * (G - b)

Let's say we have 3 episodes with returns 8, 10, and 12. The average is 10, so we use 10 as the baseline.

Here, we can see that the action that did better than average gets pushed up, the action that did worse than average gets pushed down, and the average one does not move.

The best part is that subtracting a baseline does not change the average direction of the update, so the algorithm still learns the right thing. It only reduces the variance, so the learning becomes more stable.

A better baseline is the expected return from the current state, learned by a second neural network. A policy network together with this second network that judges the states is the core idea behind Actor-Critic methods, and G minus this state baseline is an estimate of the advantage, which tells how much better an action was than the usual result from that state.

If we want to go deep into Reinforcement Learning and the critic-based methods like Proximal Policy Optimization (PPO), our AI and Machine Learning Program at Outcome School covers them in depth.

Connection to LLM Training

Policy Gradient is not only for games. It is a core part of how LLMs are trained with Reinforcement Learning.

We can map an LLM to the same picture:

• Policy: the LLM itself.
• State: the prompt plus the tokens generated so far.
• Action: the next token.
• Action probabilities: the probability the LLM gives to every token in its vocabulary.
• Return: a score for the full answer, for example from a reward model or from checking whether a math answer is correct.

Each token is one action. When an answer gets a high score, we make all the tokens in that answer a little more likely. When an answer gets a low score, we make them a little less likely. This is exactly the REINFORCE idea.

Popular methods used to train LLMs, like PPO (Proximal Policy Optimization) and GRPO (Group Relative Policy Optimization), are built on top of this same Policy Gradient idea.

• PPO adds a learned baseline (a critic) and limits how much the policy can change in one update, so that training stays stable.
• GRPO generates a group of answers for the same prompt and uses the average score of the group as the baseline, so it does not need a critic network. This is the same baseline subtraction idea that we learned above.

Now we must have understood what Policy Gradient is, how the REINFORCE algorithm works, why it has high variance, how a baseline fixes it, and how it connects to training LLMs.

Frequently Asked Questions

Who introduced the REINFORCE algorithm?

Ronald Williams introduced the REINFORCE algorithm in 1992. It is the simplest Policy Gradient algorithm. It plays a full episode with the current policy, computes the return for every step, and then updates the policy so that actions with a high return become more likely.

Why is REINFORCE called a Monte Carlo method?

REINFORCE is called a Monte Carlo method because it waits for the full episode to finish and uses the real returns. It does not use a second model to guess any values. This is also why its updates have high variance, since one episode depends on many random action choices and random events.

Why does Policy Gradient use gradient ascent instead of gradient descent?

Policy Gradient uses gradient ascent because it wants to increase the expected return, not reduce a loss. In normal neural network training, we go down the gradient to reduce the loss. In Policy Gradient, we go up the gradient, so the weights move in the direction that makes high-return actions more likely.

Does subtracting a baseline change what the policy learns?

No. Subtracting a baseline does not change the average direction of the update, so the algorithm still learns the right thing. It only reduces the variance, so the learning becomes more stable. Actions that did better than the baseline get pushed up, and actions that did worse get pushed down.

How does GRPO get its baseline without a critic network?

GRPO generates a group of answers for the same prompt and uses the average score of the group as the baseline. So, it does not need a critic network. This is the same baseline subtraction idea used to reduce the high variance of REINFORCE, where the update is multiplied by the return minus the baseline.

That's it for now.

Thanks

Amit Shekhar
Founder @ Outcome School

You can connect with me on:

• X
• LinkedIn
• GitHub

Follow Outcome School on:

• X
• YouTube
• LinkedIn
• GitHub

Read all of our high-quality blogs here.