From 39c21ab6b99279f66d2eaa2b4327a6e2c072a50c Mon Sep 17 00:00:00 2001 From: Dominik Roth Date: Thu, 14 Mar 2024 17:34:47 +0100 Subject: [PATCH] Fix problem when env is provided just as id --- metastable_baselines2/common/on_policy_algorithm.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/metastable_baselines2/common/on_policy_algorithm.py b/metastable_baselines2/common/on_policy_algorithm.py index 7194093..2af3175 100644 --- a/metastable_baselines2/common/on_policy_algorithm.py +++ b/metastable_baselines2/common/on_policy_algorithm.py @@ -92,10 +92,6 @@ class BetterOnPolicyAlgorithm(OnPolicyAlgorithm): assert not rollout_buffer_class and not rollout_buffer_kwargs - if 'dist_kwargs' not in policy_kwargs: - policy_kwargs['dist_kwargs'] = {} - policy_kwargs['dist_kwargs']['n_envs'] = len(env.envs) - super().__init__( policy=policy, env=env, @@ -122,6 +118,10 @@ class BetterOnPolicyAlgorithm(OnPolicyAlgorithm): _init_setup_model=_init_setup_model ) + if 'dist_kwargs' not in self.policy_kwargs: + self.policy_kwargs['dist_kwargs'] = {} + self.policy_kwargs['dist_kwargs']['n_envs'] = len(self.env.envs) + self.rollout_buffer_class = None self.rollout_buffer_kwargs = {}