From 00595b09d866eba0123a9eba02c08f7aff675b5c Mon Sep 17 00:00:00 2001 From: mloo3 Date: Wed, 23 Sep 2020 16:40:41 -0400 Subject: [PATCH] Add actor/critic loss logging to td3 (#164) * add actor/critic loss logging to td3 * Update changelog.rst Co-authored-by: Antonin RAFFIN --- docs/misc/changelog.rst | 3 ++- stable_baselines3/td3/td3.py | 7 +++++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/docs/misc/changelog.rst b/docs/misc/changelog.rst index 220ee38..9e41842 100644 --- a/docs/misc/changelog.rst +++ b/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 diff --git a/stable_baselines3/td3/td3.py b/stable_baselines3/td3/td3.py index 543674c..e862bae 100644 --- a/stable_baselines3/td3/td3.py +++ b/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,