dodox/metastable-baselines
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Allow std logging for full and diagonal cov policies

Browse Source
This commit is contained in:
Dominik Moritz Roth 2022-07-15 18:46:42 +02:00
parent 74697e8773
commit f184b88f19
1 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
metastable_baselines/ppo/ppo.py
View File

@ -361,6 +361,10 @@ class PPO(GaussianRolloutCollectorAuxclass, OnPolicyAlgorithm):
self.logger.record(
"train/std", th.exp(self.policy.log_std).mean().item())
if hasattr(self.policy, "chol"):
if len(self.policy.chol.shape) == 1:
self.logger.record(
"train/std", th.mean(self.policy.chol).mean().item())
else:
self.logger.record(
"train/std", th.mean(th.diagonal(self.policy.chol, dim1=-2, dim2=-1)).mean().item())