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

Tell PCA to use mSqrt when using W2 proj

Browse Source
This commit is contained in:
Dominik Moritz Roth 2024-03-30 14:44:20 +01:00
parent 4831fbb7e1
commit ada300fa63
1 changed files with 4 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
metastable_baselines2/trpl/trpl.py
View File

@ -121,6 +121,10 @@ class TRPL(BetterOnPolicyAlgorithm):
if policy_kwargs is None: if policy_kwargs is None:
policy_kwargs = {} policy_kwargs = {}
policy_kwargs['policy_projection'] = self.projection policy_kwargs['policy_projection'] = self.projection
if 'dist_kwargs' not in policy_kwargs:
policy_kwargs['dist_kwargs'] = {}
if use_pca:
policy_kwargs['dist_kwargs']['msqrt_induces_full'] = self.projection_class == WassersteinProjectionLayer
super().__init__( super().__init__(
policy, policy,