REINFORCE with Baseline not Learning

Question

I have implemented REINFORCE using PyTorch and am testing it on the CartPole environment. My implementation allows for an optional baseline to be applied. At present, the baseline used is simply the mean of the returns earned during an episode.

The agent will learn a good policy when I DO NOT use a baseline, but when I apply the baseline, the agent fails to learn anything. I cannot figure out why. I have experimented with the learning rate quite a bit, but that hasn't gotten me anywhere.

I notice that the loss is always very close to zero when using the baseline, but it seems like that should be expected. When the network weights are still random, most of the actions will have a probability that is near 0.5, and thus a log probability that is close to log(0.5) ~=~ -0.7. The returns for this environment are symmetric about the mean, so the weighted sum of the centered returns should be close to zero if the weights (log probs) are nearly equal. But the loss shouldn't be exactly zero unless the probabilities are all identical, which is not the case.

Here is a link to a Colab notebook with my code: REINFORCE Implementation

And here is the code for the function that implements the training loop.

Thanks in advance. Any help you can provide would be greatly appreciated.

def train(self, episodes, lr, max_steps=None, updates=None): self.policy_net.to(device) optimizer = torch.optim.Adam(self.policy_net.parameters(), lr=lr)

for n in range(episodes):
    #--------------------------------------------
    # Generate episode and calculate returns
    #--------------------------------------------
    self.generate_episode(max_steps=max_steps)
    T = len(self.rewards)
    returns = np.zeros(T)
    Gt = 0
    for t in reversed(range(T)):
        Gt = self.rewards[t] + self.gamma * Gt
        returns[t] = Gt

    #--------------------------------------------
    # Calculate Loss
    #--------------------------------------------
    ret_tensor = torch.FloatTensor(returns).unsqueeze(1)
    if self.with_baseline:
        ret_tensor = ret_tensor - ret_tensor.mean()

    ret_tensor = ret_tensor.to(device)
    log_probs = torch.cat(self.log_probs)
    loss = - torch.sum(log_probs * ret_tensor)

    #--------------------------------------------
    # Gradient Descent
    #--------------------------------------------
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

Answer 1

Your baseline value is a mean of the returns of a single episode. This is correlated too much with the chosen actions for that episode.

Instead, use a baseline value which is a longer-running mean over many more episodes.

For instance, use something like this:

if self.with_baseline:
    baseline += 0.01 * (ret_tensor.mean() - baseline)
    ret_tensor -= baseline

and declare baseline = 0.0 at the start.