pass entire experiment to model

facebookresearch · Oct 20, 2023 · 0f299ca · 0f299ca
1 parent 0b8436c
commit 0f299ca
Show file tree

Hide file tree

Showing 11 changed files with 32 additions and 23 deletions.
diff --git a/benchmarl/algorithms/iddpg.py b/benchmarl/algorithms/iddpg.py
@@ -101,6 +101,7 @@ def _get_policy_for_loss(
                 centralised=False,
                 share_params=self.experiment_config.share_policy_params,
                 device=self.device,
+                experiment=self.experiment,
             )
 
             policy = ProbabilisticActor(
@@ -217,6 +218,7 @@ def get_value_module(self, group: str) -> TensorDictModule:
                 agent_group=group,
                 share_params=self.share_param_critic,
                 device=self.device,
+                experiment=self.experiment,
             )
         )
 

diff --git a/benchmarl/algorithms/ippo.py b/benchmarl/algorithms/ippo.py
@@ -48,7 +48,6 @@ def __init__(
     def _get_loss(
         self, group: str, policy_for_loss: TensorDictModule, continuous: bool
     ) -> Tuple[LossModule, bool]:
-
         # Loss
         loss_module = ClipPPOLoss(
             actor=policy_for_loss,
@@ -83,7 +82,6 @@ def _get_parameters(self, group: str, loss: ClipPPOLoss) -> Dict[str, Iterable]:
     def _get_policy_for_loss(
         self, group: str, model_config: ModelConfig, continuous: bool
     ) -> TensorDictModule:
-
         n_agents = len(self.group_map[group])
         if continuous:
             logits_shape = list(self.action_spec[group, "action"].shape)
@@ -124,6 +122,7 @@ def _get_policy_for_loss(
             centralised=False,
             share_params=self.experiment_config.share_policy_params,
             device=self.device,
+            experiment=self.experiment,
         )
 
         if continuous:
@@ -261,14 +260,14 @@ def get_critic(self, group: str) -> TensorDictModule:
             agent_group=group,
             share_params=self.share_param_critic,
             device=self.device,
+            experiment=self.experiment,
         )
 
         return value_module
 
 
 @dataclass
 class IppoConfig(AlgorithmConfig):
-
     share_param_critic: bool = MISSING
     clip_epsilon: float = MISSING
     entropy_coef: float = MISSING

diff --git a/benchmarl/algorithms/iql.py b/benchmarl/algorithms/iql.py
@@ -62,7 +62,6 @@ def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
     def _get_policy_for_loss(
         self, group: str, model_config: ModelConfig, continuous: bool
     ) -> TensorDictModule:
-
         n_agents = len(self.group_map[group])
         logits_shape = [
             *self.action_spec[group, "action"].shape,
@@ -99,6 +98,7 @@ def _get_policy_for_loss(
             centralised=False,
             share_params=self.experiment_config.share_policy_params,
             device=self.device,
+            experiment=self.experiment,
         )
         if self.action_mask_spec is not None:
             action_mask_key = (group, "action_mask")
@@ -175,7 +175,6 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
 
 @dataclass
 class IqlConfig(AlgorithmConfig):
-
     delay_value: bool = MISSING
     loss_function: str = MISSING
 

diff --git a/benchmarl/algorithms/isac.py b/benchmarl/algorithms/isac.py
@@ -126,7 +126,6 @@ def _get_parameters(self, group: str, loss: ClipPPOLoss) -> Dict[str, Iterable]:
     def _get_policy_for_loss(
         self, group: str, model_config: ModelConfig, continuous: bool
     ) -> TensorDictModule:
-
         n_agents = len(self.group_map[group])
         if continuous:
             logits_shape = list(self.action_spec[group, "action"].shape)
@@ -167,6 +166,7 @@ def _get_policy_for_loss(
             centralised=False,
             share_params=self.experiment_config.share_policy_params,
             device=self.device,
+            experiment=self.experiment,
         )
 
         if continuous:
@@ -291,6 +291,7 @@ def get_discrete_value_module(self, group: str) -> TensorDictModule:
             agent_group=group,
             share_params=self.share_param_critic,
             device=self.device,
+            experiment=self.experiment,
         )
 
         return value_module
@@ -346,6 +347,7 @@ def get_continuous_value_module(self, group: str) -> TensorDictModule:
                 agent_group=group,
                 share_params=self.share_param_critic,
                 device=self.device,
+                experiment=self.experiment,
             )
         )
 
@@ -354,7 +356,6 @@ def get_continuous_value_module(self, group: str) -> TensorDictModule:
 
 @dataclass
 class IsacConfig(AlgorithmConfig):
-
     share_param_critic: bool = MISSING
 
     num_qvalue_nets: int = MISSING

diff --git a/benchmarl/algorithms/maddpg.py b/benchmarl/algorithms/maddpg.py
@@ -62,7 +62,6 @@ def _get_loss(
             )
 
     def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
-
         return {
             "loss_actor": list(loss.actor_network_params.flatten_keys().values()),
             "loss_value": list(loss.value_network_params.flatten_keys().values()),
@@ -103,6 +102,7 @@ def _get_policy_for_loss(
                 centralised=False,
                 share_params=self.experiment_config.share_policy_params,
                 device=self.device,
+                experiment=self.experiment,
             )
 
             policy = ProbabilisticActor(
@@ -222,11 +222,11 @@ def get_value_module(self, group: str) -> TensorDictModule:
                     agent_group=group,
                     share_params=self.share_param_critic,
                     device=self.device,
+                    experiment=self.experiment,
                 )
             )
 
         else:
-
             modules.append(
                 TensorDictModule(
                     lambda obs, action: torch.cat([obs, action], dim=-1),
@@ -263,6 +263,7 @@ def get_value_module(self, group: str) -> TensorDictModule:
                     agent_group=group,
                     share_params=self.share_param_critic,
                     device=self.device,
+                    experiment=self.experiment,
                 )
             )
 
@@ -282,7 +283,6 @@ def get_value_module(self, group: str) -> TensorDictModule:
 
 @dataclass
 class MaddpgConfig(AlgorithmConfig):
-
     share_param_critic: bool = MISSING
 
     loss_function: str = MISSING

diff --git a/benchmarl/algorithms/mappo.py b/benchmarl/algorithms/mappo.py
@@ -123,6 +123,7 @@ def _get_policy_for_loss(
             centralised=False,
             share_params=self.experiment_config.share_policy_params,
             device=self.device,
+            experiment=self.experiment,
         )
 
         if continuous:
@@ -258,6 +259,7 @@ def get_critic(self, group: str) -> TensorDictModule:
                 agent_group=group,
                 share_params=self.share_param_critic,
                 device=self.device,
+                experiment=self.experiment,
             )
 
         else:
@@ -282,6 +284,7 @@ def get_critic(self, group: str) -> TensorDictModule:
                 agent_group=group,
                 share_params=self.share_param_critic,
                 device=self.device,
+                experiment=self.experiment,
             )
         if self.share_param_critic:
             expand_module = TensorDictModule(

diff --git a/benchmarl/algorithms/masac.py b/benchmarl/algorithms/masac.py
@@ -121,7 +121,6 @@ def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
     def _get_policy_for_loss(
         self, group: str, model_config: ModelConfig, continuous: bool
     ) -> TensorDictModule:
-
         n_agents = len(self.group_map[group])
         if continuous:
             logits_shape = list(self.action_spec[group, "action"].shape)
@@ -162,6 +161,7 @@ def _get_policy_for_loss(
             centralised=False,
             share_params=self.experiment_config.share_policy_params,
             device=self.device,
+            experiment=self.experiment,
         )
 
         if continuous:
@@ -279,6 +279,7 @@ def get_discrete_value_module(self, group: str) -> TensorDictModule:
                 agent_group=group,
                 share_params=self.share_param_critic,
                 device=self.device,
+                experiment=self.experiment,
             )
 
         else:
@@ -303,6 +304,7 @@ def get_discrete_value_module(self, group: str) -> TensorDictModule:
                 agent_group=group,
                 share_params=self.share_param_critic,
                 device=self.device,
+                experiment=self.experiment,
             )
         if self.share_param_critic:
             expand_module = TensorDictModule(
@@ -369,11 +371,11 @@ def get_continuous_value_module(self, group: str) -> TensorDictModule:
                     agent_group=group,
                     share_params=self.share_param_critic,
                     device=self.device,
+                    experiment=self.experiment,
                 )
             )
 
         else:
-
             modules.append(
                 TensorDictModule(
                     lambda obs, action: torch.cat([obs, action], dim=-1),
@@ -410,6 +412,7 @@ def get_continuous_value_module(self, group: str) -> TensorDictModule:
                     agent_group=group,
                     share_params=self.share_param_critic,
                     device=self.device,
+                    experiment=self.experiment,
                 )
             )
 
@@ -429,7 +432,6 @@ def get_continuous_value_module(self, group: str) -> TensorDictModule:
 
 @dataclass
 class MasacConfig(AlgorithmConfig):
-
     share_param_critic: bool = MISSING
 
     num_qvalue_nets: int = MISSING

diff --git a/benchmarl/algorithms/qmix.py b/benchmarl/algorithms/qmix.py
@@ -67,7 +67,6 @@ def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
     def _get_policy_for_loss(
         self, group: str, model_config: ModelConfig, continuous: bool
     ) -> TensorDictModule:
-
         n_agents = len(self.group_map[group])
         logits_shape = [
             *self.action_spec[group, "action"].shape,
@@ -104,6 +103,7 @@ def _get_policy_for_loss(
             centralised=False,
             share_params=self.experiment_config.share_policy_params,
             device=self.device,
+            experiment=self.experiment,
         )
         if self.action_mask_spec is not None:
             action_mask_key = (group, "action_mask")
@@ -175,7 +175,6 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
     #####################
 
     def get_mixer(self, group: str) -> TensorDictModule:
-
         n_agents = len(self.group_map[group])
 
         if self.state_spec is not None:
@@ -201,7 +200,6 @@ def get_mixer(self, group: str) -> TensorDictModule:
 
 @dataclass
 class QmixConfig(AlgorithmConfig):
-
     mixing_embed_dim: int = MISSING
     delay_value: bool = MISSING
     loss_function: str = MISSING

diff --git a/benchmarl/algorithms/vdn.py b/benchmarl/algorithms/vdn.py
@@ -59,15 +59,13 @@ def _get_loss(
             return loss_module, True
 
     def _get_parameters(self, group: str, loss: LossModule) -> Dict[str, Iterable]:
-
         return {
             "loss": loss.parameters(),
         }
 
     def _get_policy_for_loss(
         self, group: str, model_config: ModelConfig, continuous: bool
     ) -> TensorDictModule:
-
         n_agents = len(self.group_map[group])
         logits_shape = [
             *self.action_spec[group, "action"].shape,
@@ -104,6 +102,7 @@ def _get_policy_for_loss(
             centralised=False,
             share_params=self.experiment_config.share_policy_params,
             device=self.device,
+            experiment=self.experiment,
         )
         if self.action_mask_spec is not None:
             action_mask_key = (group, "action_mask")
@@ -175,7 +174,6 @@ def process_batch(self, group: str, batch: TensorDictBase) -> TensorDictBase:
     #####################
 
     def get_mixer(self, group: str) -> TensorDictModule:
-
         n_agents = len(self.group_map[group])
         mixer = TensorDictModule(
             module=VDNMixer(
@@ -191,7 +189,6 @@ def get_mixer(self, group: str) -> TensorDictModule:
 
 @dataclass
 class VdnConfig(AlgorithmConfig):
-
     delay_value: bool = MISSING
     loss_function: str = MISSING