bugs

facebookresearch · Nov 6, 2024 · ecf6c64 · ecf6c64
1 parent bceb7e6
commit ecf6c64
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/benchmarl/experiment/experiment.py b/benchmarl/experiment/experiment.py
@@ -17,6 +17,7 @@
 from typing import Any, Dict, List, Optional
 
 import torch
+from algorithms import IppoConfig, MappoConfig
 from tensordict import TensorDictBase
 from tensordict.nn import TensorDictSequential
 from torchrl.collectors import SyncDataCollector
@@ -374,7 +375,7 @@ def _perfrom_checks(self):
                             " layer of sequence models"
                         )
 
-        if self.algorithm_name in ("mappo", "ippo"):
+        if self.algorithm_config in (MappoConfig, IppoConfig):
             critic_model_config = self.critic_model_config
             if isinstance(critic_model_config, SequenceModelConfig):
                 critic_model_config = self.critic_model_config.model_configs[0]