REINFORCE Algorithm Primer

REINFORCE is the basic policy gradient algorithm. Policy gradient methods try to increase the probabilities of actions that lead to higher return and decrease the probabilities of actions that lead to lower return.

REINFORCE is (natively):

• an online reinforcement learning algorithm as it uses new data generated by the learned policy (the policy interacts with the environment while learning)
• on-policy as the trajectories used for the update are sampled from the current policy
• model-free as it learns directly from interacting with the environment without building an internal transition or reward model of the environment

1. Problem Setting

We aim to learn a policy parameterized by . The policy can be conditioned on an observation in a partially-observed setting. At each timestep, .

A trajectory is a sequence of states/observations and actions:

We have a reward function which dictates how good taking action at state is in the environment. This assumes the reward function is not explicitly changing across time.

The objective is to maximise expected total return:

The probability of a trajectory under the current policy is

Here, is the discount factor.

2. Monte Carlo Approximation of the Objective

The expectation over all possible trajectories is generally impossible to compute exactly as an integral. Instead, we sample trajectories using the current policy,

and approximate the expectation with the sample mean:

We will switch back to integrals for now and evaluate over all possible trajectories to derive the policy gradient, then revert to this Monte Carlo approximation at the end.

3. Deriving the Policy Gradient

How do we change so that this expected return increases? Let’s derive the policy gradient. We want to compute .

For simplicity, define the trajectory return as

Then our objective can be written as

Writing the expectation as an integral gives

To take the gradient with respect to our parameters , we differentiate with respect to :

As it is, is not a valid probability distribution, so this cannot be rewritten directly as an expectation. We therefore use the log-derivative trick:

Substituting this into the gradient gives

Because the integral now contains , we can rewrite it as an expectation:

We now need to evaluate .

Recall that

Taking the logarithm turns the products into sums:

Taking the gradient with respect to eliminates the initial state distribution and the environment dynamics, as neither depends on :

Therefore,

This is the basic REINFORCE gradient.

4. Monte Carlo Approximation of the Gradient

We cannot compute this expectation over all possible trajectories, so we approximate it using sampled trajectories:

Intuition

REINFORCE is behavior cloning with a weight on the log-probability gradient in the form of the return.

The gradient in the behavior cloning (supervised learning) setup is simply

REINFORCE adds the weight :

which increases the likelihood of actions you took in trajectories that returned high reward and decreases the likelihood of actions you took in trajectories that returned low reward.

5. Improving REINFORCE

Without introducing a new algorithm, REINFORCE can be improved with two simple modifications.

5.1 Reward-to-Go

If you look to our REINFORCE gradient update, we are weighing the likelihood of actions taken by the return of the entire trajecotry (reward of the trajectory from the first state to the final state in the trajectory). Intuitively, this dilutes and weakens our learning signal (particularly for later actions in the episode). Why should the likelihood of an action be weighted by rewards that preceded that action being taken in our gradient?

Therefore, one simple modification of our REINFORCE gradient is to weigh the likelihood of an action by the rewards that follow that action being taken. This changes our gradient update to:

5.2 Baseline

Another simple trick to reduce the variance of the policy-gradients is to introduce a baseline for our return.

Given that we have re-defined our weight to be based on the rewards that follow an action being taken: , we can define a batch-level baseline that centers these reward-to-go values across the batch:

Then our policy gradient can become:

I need to confirm this. Most REINFORCE discussions define the baseline by averaging total returns across a batch of trajectories. Here, we are trying to factor in our reward-to-formulation and define the baseline across a batch as well as across the timesteps of those trajectories. Either way, this is not very important as this baseline will simply become the value function when we move to actor-critic architectures.

6. Algorithm and Implementation

6.1 Implementation Detail

One limitation of the algorithm we presented above is the need to conduct N*T backward passes as form our policy-gradient estiamte as a summation over all episodes in our buffer as well as all timesteps of those episodes. This is computationally inefficient.

Ideally, we want to be able to conduct one backward pass by using automatic differentiation on the full objective. By using the fact that differentiation is linear, we can define a scalar surrogate loss whose gradient is equal to the REINFORCE policy gradient we presented above.

Starting from

define the surrogate objective

where the sampled returns and baseline are treated as constants with respect to the policy parameters:

Then, by linearity of differentiation,

Thus, all timestep-level policy-gradient contributions can be accumulated into one scalar loss and computed with a single backward pass.

Additionally, the form of depends on the policy distribution. For a categorical policy, its negative is equivalent to cross-entropy with the sampled action as the target. For a fixed-variance Gaussian policy, its negative is proportional to the squared error between the sampled action and the policy mean.

7. Can we make this off-policy?

One issue that arises with policy gradients is the sampling process. Step 1 of our algorithm collects trajectories from but is updated by the end of Step 5, which forces us to resample new trajectories (this is often the most expensive part of learning in these online settings). This is especially because we often require thousands (if not millions!) of steps/paramter updates to train neural networks.

This arises because policy gradients is an on-policy algorithm, which learns from data generated by the same policy that is currently being optimized. Wheras off-policy algorithms can learn from data generated by a different or older policy. We can exploit this “older policy” feature of off-policy algorithms to improve policy gradients in order to “use more” of our batch of data (learn more from single batch of data without having to immediately resample new trajectories after each paramter update).

7.1 Off-Policy Version of Policy Gradient

The core motivation of this off-policy variant is being able to use samples from a previous policy, .

We now distinguish between our past (often called, behaviour) policy which generated the data we have in our buffer, and our proposed (current) policy , which we want to evaluate or improve.

So far, our on-policy’s objective was based on what would happen under the proposed policy (which made sense when our trajectories were sampled from that same policy). This no longer applies here, which means we need to change our objective.

Importance Sampling

The generic importance-sampling identity is

In our policy-gradient setting, the equivalent distributions are and , and our function is the expected return .

Therefore,

This means we can use trajectories generated by the old policy , but we must reweight each trajectory according to how likely that trajectory would have been under the proposed policy relative to how likely it was under . **It is important to ensure the support of our past policy includes high probability regions of our proposed policy (so that fraction does not explode).

We can simplify the importance ratio using the trajectory probabilities:

The initial-state distribution and environment dynamics are identical under both policies, so they cancel:

Therefore,

Off-Policy Policy Gradient

Recall our REINFORCE gradient with reward-to-go and a baseline:

If our trajectories instead come from the past policy , importance sampling allows us to rewrite the expectation under .

A simple trajectory-wise importance-sampling version is

Cumulative Importance Ratio:

However, for the gradient contribution at timestep , we do not need to correct for actions taken after timestep . The state and action were produced only by the sequence of decisions up to timestep .

We can therefore define the cumulative importance ratio

This gives the per-decision importance-sampling policy gradient

The cumulative ratio is important because the difference between the old and new policies affects more than just the probability of taking the current action . It also affects the probability of having reached the current state in the first place. In this sense, cumulative importance sampling corrects both the state-visitation mismatch and the action-probability mismatch between the behaviour policy and the current policy.

However, there is an important practical limitation. Products of importance ratios can have extremely high variance, particularly at longer time horizons (). Conversely, ratios smaller than one can rapidly drive the weight towards zero. So cumulative importance sampling gives us a relatively principled correction for off-policy data, but potentially at the cost of very high variance.

Approximating the State-Visitation Correction

Instead of correcting the probability of the whole trajectory prefix, we can look directly at the distribution over the state-action pair that appears in the policy-gradient update.

Under the current policy,

while under the behaviour policy,

Therefore,

The first term, which we call the *state-visition ratio, corrects for how likely each policy is to reach (distribution of states at timestep ), while the second corrects for how likely each policy is to take at .

To connect this to what we did with cumulative importance sampling, the state-visition ratio at is the average cumulative importance weight of all behaviour-policy trajectory prefixes that reach .

The difficulty is that the state-visitation ratio

is hard to compute directly. A common approximation is to assume the old and new policies are close enough that , so that .

This leaves only the single-step action ratio and gives

This is less exact than cumulative importance sampling, but has much lower variance because it avoids multiplying importance ratios across the trajectory history.

References

• Why is it better to subtract a baseline in REINFORCE?
• Stanford CS224R Deep Reinforcement Learning | Spring 2025 | Lecture 3: Policy Gradients