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

Enabled w2 (can now get sqrt from dist)

Browse Source
This commit is contained in:
Dominik Moritz Roth 2022-08-06 14:54:59 +02:00
parent 508ebf51f0
commit 802094a50f
2 changed files with 11 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
metastable_baselines/misc/distTools.py
View File

@ -20,13 +20,14 @@ def get_mean_and_chol(p: AnyDistribution, expand=False):
def get_mean_and_sqrt(p: UniversalGaussianDistribution):
raise Exception('Not yet implemented...')
if isinstance(p, th.distributions.Normal):
return p.mean, p.stddev
elif isinstance(p, th.distributions.MultivariateNormal):
return p.mean, p.scale_tril
elif isinstance(p, SB3_Distribution):
return get_mean_and_chol(p.distribution)
if isinstance(p, UniversalGaussianDistribution):
if not hasattr(p, 'cov_sqrt'):
raise Exception(
'Distribution was not induced from sqrt. On-demand calculation is not supported.')
else:
mean, chol = get_mean_and_chol(p)
sqrt_cov = p.cov_sqrt
return mean, sqrt_cov
else:
raise Exception('Dist-Type not implemented')

4
metastable_baselines/ppo/policies.py
View File

@ -38,6 +38,7 @@ from stable_baselines3.common.torch_layers import (
from metastable_baselines.projections.w2_projection_layer import WassersteinProjectionLayer
from ..distributions import UniversalGaussianDistribution, make_proba_distribution
from ..misc.distTools import get_mean_and_chol
class ActorCriticPolicy(BasePolicy):
@ -293,7 +294,8 @@ class ActorCriticPolicy(BasePolicy):
cov_sqrt = self.chol_net(latent_pi)
dist = self.action_dist.proba_distribution_from_sqrt(
mean_actions, cov_sqrt, latent_pi)
self.chol = dist.chol
mean, chol = get_mean_and_chol(dist, expand=False)
self.chol = chol
return dist
else:
raise Exception(