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

Broader sampling of stds for logging with batched full covs

Browse Source
This commit is contained in:
Dominik Moritz Roth 2022-07-16 15:28:16 +02:00
parent bc0e188a0d
commit 3fa6de7e66
1 changed files with 7 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
metastable_baselines/ppo/ppo.py
View File

@ -361,17 +361,16 @@ class PPO(GaussianRolloutCollectorAuxclass, OnPolicyAlgorithm):
self.logger.record(
"train/std", th.exp(self.policy.log_std).mean().item())
elif hasattr(self.policy, "chol"):
if len(self.policy.chol.shape) == 1:
self.logger.record(
"train/std", th.mean(self.policy.chol).mean().item())
else:
if len(self.policy.chol.shape) == 2:
chol = self.policy.chol
else:
# TODO: Maybe use a broader sample?
chol = self.policy.chol[0]
if len(chol.shape) == 1:
self.logger.record(
"train/std", th.mean(chol).mean().item())
elif len(chol.shape) == 2:
self.logger.record(
"train/std", th.mean(th.sqrt(th.diagonal(chol.T @ chol, dim1=-2, dim2=-1))).mean().item())
else:
self.logger.record(
"train/std", th.mean(th.sqrt(th.diagonal(chol.mT @ chol, dim1=-2, dim2=-1))).mean().item())
self.logger.record("train/n_updates",
self._n_updates, exclude="tensorboard")