From e14aae90ac8fec91329d2b566fb44686cabdd422 Mon Sep 17 00:00:00 2001 From: Dominik Roth Date: Sat, 9 Mar 2024 13:55:04 +0100 Subject: [PATCH] Fix: drop n_envs for DIstributions that dont need to know --- metastable_baselines2/common/policies.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/metastable_baselines2/common/policies.py b/metastable_baselines2/common/policies.py index 3923eb6..43e6a54 100644 --- a/metastable_baselines2/common/policies.py +++ b/metastable_baselines2/common/policies.py @@ -530,7 +530,7 @@ class ActorCriticPolicy(BasePolicy): self.n_envs = dist_kwargs.pop('n_envs', 1) # Action distribution - self.action_dist = make_proba_distribution(action_space, self.n_envs, use_sde=use_sde, use_pca=use_pca, dist_kwargs=dist_kwargs) + self.action_dist = make_proba_distribution(action_space, self.n_envs, use_sde=use_sde, use_pca=use_pca, n_envs=self.n_envs, dist_kwargs=dist_kwargs) self._build(lr_schedule)