Add actor/critic loss logging to td3 (#164)

* add actor/critic loss logging to td3

* Update changelog.rst

Co-authored-by: Antonin RAFFIN <antonin.raffin@ensta.org>

docs/misc/changelog.rst

@@ -16,6 +16,7 @@ New Features:
 - Added ``StopTrainingOnMaxEpisodes`` to callback collection (@xicocaio)
 - Added ``device`` keyword argument to ``BaseAlgorithm.load()`` (@liorcohen5)
 - Callbacks have access to rollout collection locals as in SB2. (@PartiallyTyped)
+- Added actor/critic loss logging for TD3. (@mloo3)
 
 Bug Fixes:
 ^^^^^^^^^^
@@ -402,4 +403,4 @@ And all the contributors:
 @MarvineGothic @jdossgollin @SyllogismRXS @rusu24edward @jbulow @Antymon @seheevic @justinkterry @edbeeching
 @flodorner @KuKuXia @NeoExtended @PartiallyTyped @mmcenta @richardwu @kinalmehta @rolandgvc @tkelestemur @mloo3
 @tirafesi @blurLake @koulakis @joeljosephjin @shwang @rk37 @andyshih12 @RaphaelWag @xicocaio
-@diditforlulz273 @liorcohen5 @ManifoldFR
+@diditforlulz273 @liorcohen5 @ManifoldFR @mloo3

stable_baselines3/td3/td3.py

@@ -1,5 +1,6 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union
 
+import numpy as np
 import torch as th
 from torch.nn import functional as F
 
@@ -130,6 +131,8 @@ class TD3(OffPolicyAlgorithm):
         # Update learning rate according to lr schedule
         self._update_learning_rate([self.actor.optimizer, self.critic.optimizer])
 
+        actor_losses, critic_losses = [], []
+
         for gradient_step in range(gradient_steps):
 
             # Sample replay buffer
@@ -151,6 +154,7 @@ class TD3(OffPolicyAlgorithm):
 
             # Compute critic loss
             critic_loss = sum([F.mse_loss(current_q, target_q) for current_q in current_q_estimates])
+            critic_losses.append(critic_loss.item())
 
             # Optimize the critics
             self.critic.optimizer.zero_grad()
@@ -161,6 +165,7 @@ class TD3(OffPolicyAlgorithm):
             if gradient_step % self.policy_delay == 0:
                 # Compute actor loss
                 actor_loss = -self.critic.q1_forward(replay_data.observations, self.actor(replay_data.observations)).mean()
+                actor_losses.append(actor_loss.item())
 
                 # Optimize the actor
                 self.actor.optimizer.zero_grad()
@@ -172,6 +177,8 @@ class TD3(OffPolicyAlgorithm):
 
         self._n_updates += gradient_steps
         logger.record("train/n_updates", self._n_updates, exclude="tensorboard")
+        logger.record("train/actor_loss", np.mean(actor_losses))
+        logger.record("train/critic_loss", np.mean(critic_losses))
 
     def learn(
         self,