Should the policy parameters be updated at each time step or at the end of the episode in REINFORCE?

Question

REINFORCE is a Monte Carlo policy gradient algorithm, which updates weights (parameters) of policy network by generating episodes. Here's a pseudo-code from Sutton's book (which is same as the equation in Silver's RL note):

When I try to implement this with my own problem, I found something strange. Here's implementation from Pytorch's official GitHub:

def finish_episode():
    R = 0
    policy_loss = []
    returns = []
    for r in policy.rewards[::-1]:
        R = r + args.gamma * R
        returns.insert(0, R)
    returns = torch.tensor(returns)
    returns = (returns - returns.mean()) / (returns.std() + eps)
    for log_prob, R in zip(policy.saved_log_probs, returns):
        policy_loss.append(-log_prob * R)
    optimizer.zero_grad()
    policy_loss = torch.cat(policy_loss).sum()
    policy_loss.backward()
    optimizer.step()
    del policy.rewards[:]
    del policy.saved_log_probs[:]

I feel like there's a difference between the above two. In Sutton's pseudo-code, the algorithm updates $\theta$ for each step $t$, while the second code (PyTorch's one) accumulate loss and update $\theta$ with the summation, i.e. after each episode. I tried to search other implementation of REINFORCE, and I found that most of the implementations follow the second form, update after each generated episodes.

To check whether both give the same result, I changed the second code as

def finish_episode():
    R = 0
    policy_loss = []
    returns = []
    for r in policy.rewards[::-1]:
        R = r + args.gamma * R
        returns.insert(0, R)
    returns = torch.tensor(returns)
    returns = (returns - returns.mean()) / (returns.std() + eps)
    for log_prob, R in zip(policy.saved_log_probs, returns):
        optimizer.zero_grad()
        loss = -log_prob * R
        loss.backward()
        optimizer.step()

...

and run it, which gives different result (if my code has no problem). So they are not the same, and I think the last one is more close to the original pseudo-code of REINFORCE. What am I missing now? Is it okay because the results are approximately same? (I'm not sure about this claim)

However, in some sense, I think Pytorch's implementation is the right version of REINFORCE. In Sutton's pseudo-code, episode is generated first, so I think $\theta$ shouldn't be updated at each step and should be updated after the total loss is computed. If $\theta$ is updated at each step, then such $\theta$ might be different with the original $\theta$ that used to generate the episode.

Answer 1

The essence of your observation is that Sutton's version of REINFORCE is taking into consideration all of the trajectory to compute the returns, while in the pytorch version only the future is taken into consideration, hence going in reverse to sum the future rewards and ignore the previous rewards. The consequence is that future actions are not punished for early mistakes. The guys at OpenAI refer to this as reward-to-go, but, personally, I find that it resembles Monte Carlo On Policy Control without Exploring Starts or First Visit from Sutton's book.

You can find more on REINFORCE and Policy Gradient in Spinning Up RL: Part3: Intro to policy gradient - Don't let the past distract you.

Also, something to note is that even in Sutton's version, the whole trajectory is unrolled, i.e. the episode completes, and then the weights get updated. Otherwise it stops being a Monte Carlo Method and it becomes a TD method. In addition, you can't make a change on a single point because sampling is a non differentiable operation, instead, the gradient is estimated by collecting a lot of trajectories.