Merge pull request pytorch#386 from r-aristov/RL-loss-normalized

Normalized loss in actor-critic and REINFORCE examples.
mpc001 · Mar 9, 2022 · 886b74e · 886b74e
2 parents 4aee9d0 + 41ae197
commit 886b74e
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 0 deletions.
diff --git a/reinforcement_learning/actor_critic.py b/reinforcement_learning/actor_critic.py
@@ -124,6 +124,9 @@ def finish_episode():
 
     # sum up all the values of policy_losses and value_losses
     loss = torch.stack(policy_losses).sum() + torch.stack(value_losses).sum()
+
+    # normalize loss by number of rewards
+    loss /= rewards.numel()
 
     # perform backprop
     loss.backward()

diff --git a/reinforcement_learning/reinforce.py b/reinforcement_learning/reinforce.py
@@ -72,6 +72,7 @@ def finish_episode():
         policy_loss.append(-log_prob * R)
     optimizer.zero_grad()
     policy_loss = torch.cat(policy_loss).sum()
+    policy_loss /= rewards.numel()
     policy_loss.backward()
     optimizer.step()
     del policy.rewards[:]