Fix: Crash when no dist_kwargs defined

2024-03-09 13:51:38 +01:00 · 2024-03-09 13:51:38 +01:00 · 66b8d21398
commit 66b8d21398
parent 8f66a34c29
1 changed files with 2 additions and 0 deletions
--- a/metastable_baselines2/common/on_policy_algorithm.py
+++ b/metastable_baselines2/common/on_policy_algorithm.py
@ -92,6 +92,8 @@ 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__(