diff --git a/stable_baselines3/tqc/tqc.py b/stable_baselines3/tqc/tqc.py index 1f36246..312ce0c 100644 --- a/stable_baselines3/tqc/tqc.py +++ b/stable_baselines3/tqc/tqc.py @@ -277,6 +277,10 @@ class TQC(OffPolicyAlgorithm): logger.record("train/critic_loss", np.mean(critic_losses)) if len(ent_coef_losses) > 0: logger.record("train/ent_coef_loss", np.mean(ent_coef_losses)) + if hasattr(self.actor, "scale_expert"): + with th.no_grad(): + logger.record("train/scale_expert", th.sigmoid(self.actor.scale_expert).mean().item()) + logger.record("train/scale_rl", th.sigmoid(self.actor.scale_rl).mean().item()) def pretrain( self,