2
$\begingroup$

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()
$\endgroup$

1 Answer 1

3
$\begingroup$

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.

$\endgroup$
3
  • 1
    $\begingroup$ Thank you. That made a HUGE difference. After implementing your solution, the agent with the baseline is learning much more rapidly than the one without the baseline. One of the books I am using at a reference suggested using the per-episode mean as a baseline, and included a plot showing that it outperformed the non-baseline version on CartPole. But, I notice that they don't include the code for how they actually implemented the baseline, so it might differ from what they stated. $\endgroup$
    – Beane
    Jul 14, 2023 at 12:56
  • $\begingroup$ Interesting answer! Can you elaborate a bit on why is bad for a baseline to be correlated with the actions? $\endgroup$ Jul 14, 2023 at 22:42
  • 1
    $\begingroup$ @LucaAnzalone - it's because the policy gradient theorem allows any offset to the gradient that is not associated with the parameters you are taking the gradient of (you could literally choose any arbitrary number or independent random source if you liked). Anything that relies on individual results from $\pi(a|s)$ is going to depend critically on the parameters of the gradient function you are refining. There is a grey area of things that are not directly single instances of $\pi(a|s)$, but still correlate heavily, which is IMO what the OP encountered. $\endgroup$ Jul 15, 2023 at 8:44

You must log in to answer this question.

Not the answer you're looking for? Browse other questions tagged .