Fix get_environment function to work with 3rd party datasets.

og_marl/environments.py

@@ -1,28 +1,28 @@
 from og_marl.environment_wrappers.base import BaseEnvironment
 
 
-def get_environment(env_name: str, scenario: str, seed: int = 42) -> BaseEnvironment:  # noqa: C901
-    if env_name in ["smac_v1", "smac_v1_cfcql"]:
+def get_environment(source: str, env_name: str, scenario: str, seed: int = 42) -> BaseEnvironment:  # noqa: C901
+    if env_name=="smac_v1" and source in ["cfcql", "og_marl"]:
         from og_marl.environment_wrappers.smacv1 import SMACv1
 
         return SMACv1(scenario, seed=seed)
-    elif env_name == "smac_v1_omiga":
+    elif env_name == "smac_v1" and source=="omiga":
         from og_marl.environment_wrappers.smacv1_omiga import SMACv1OMIGA
 
         return SMACv1OMIGA(scenario, seed=seed)
     elif env_name == "smac_v2":
         from og_marl.environment_wrappers.smacv2 import SMACv2
 
         return SMACv2(scenario, seed=seed)
-    elif env_name == "mpe_omar":
+    elif env_name == "mpe" and source=="omar":
         from og_marl.environment_wrappers.mpe_omar import MPEOMAR
 
         return MPEOMAR(scenario, seed=seed)
-    elif env_name == "mamujoco":
+    elif env_name == "mamujoco" and source == "og_marl":
         from og_marl.environment_wrappers.mamujoco import MAMuJoCo
 
         return MAMuJoCo(scenario, seed=seed)
-    elif env_name == "mamujoco_omiga":
+    elif env_name == "mamujoco" and source=="omiga":
         from og_marl.environment_wrappers.mamujoco_omiga import MAMuJoCoOMIGA
 
         return MAMuJoCoOMIGA(scenario, seed=seed)

og_marl/tf2/systems/continuous_bc.py

@@ -183,7 +183,7 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],

og_marl/tf2/systems/discrete_bc.py

@@ -202,7 +202,7 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],

og_marl/tf2/systems/iddpg_bc.py

@@ -292,7 +292,7 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],

og_marl/tf2/systems/iddpg_cql.py

@@ -414,7 +414,7 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],

og_marl/tf2/systems/iql_bcq.py

@@ -296,7 +296,7 @@ def _tf_train_step(self, train_step: int, experience: Dict[str, Any]) -> Dict[st
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],

og_marl/tf2/systems/iql_cql.py

@@ -260,7 +260,7 @@ def _tf_train_step(self, train_step: int, experience: Dict[str, Any]) -> Dict[st
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],

og_marl/tf2/systems/maddpg_cql.py

@@ -408,7 +408,7 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],

og_marl/tf2/systems/maicq.py

@@ -318,7 +318,7 @@ def _tf_train_step(
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],

og_marl/tf2/systems/omar.py

@@ -490,7 +490,7 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],

og_marl/tf2/systems/qmix_bcq.py

@@ -333,7 +333,7 @@ def mixing(
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],

og_marl/tf2/systems/qmix_cql.py

@@ -331,7 +331,7 @@ def mixing(
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
-    env = get_environment(cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
+    env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(
         sequence_length=cfg["replay"]["sequence_length"],