Implemented prior conditioned annealing (untested)

metastable_baselines/ppo/policies.py

@@ -25,7 +25,7 @@ from stable_baselines3.common.torch_layers import (
     MlpExtractor,
     NatureCNN,
 )
-from stable_baselines3.common.type_aliases import Schedule
+from stable_baselines3.common.type_aliases import Schedules
 
 from stable_baselines3.common.policies import BasePolicy
 from stable_baselines3.common.torch_layers import (
@@ -88,6 +88,7 @@ class ActorCriticPolicy(BasePolicy):
         activation_fn: Type[nn.Module] = nn.Tanh,
         ortho_init: bool = True,
         use_sde: bool = False,
+        use_pca: bool = False,
         std_init: float = 1.0,
         full_std: bool = True,
         sde_net_arch: Optional[List[int]] = None,
@@ -153,6 +154,7 @@ class ActorCriticPolicy(BasePolicy):
                 "sde_net_arch is deprecated and will be removed in SB3 v2.4.0.", DeprecationWarning)
 
         self.use_sde = use_sde
+        self.use_pca = use_pca
         self.dist_kwargs = dist_kwargs
 
         self.sqrt_induced_gaussian = sqrt_induced_gaussian
@@ -160,7 +162,7 @@ class ActorCriticPolicy(BasePolicy):
 
         # Action distribution
         self.action_dist = make_proba_distribution(
-            action_space, use_sde=use_sde, dist_kwargs=dist_kwargs)
+            action_space, use_sde=use_sde, use_pca=use_pca, dist_kwargs=dist_kwargs)
 
         self._build(lr_schedule)