Make jax run on CPU so it does not conflict with TF.

og_marl/tf2_systems/offline/continuous_bc.py

@@ -17,6 +17,7 @@
 
 import hydra
 import numpy as np
+import jax
 from omegaconf import DictConfig
 import sonnet as snt
 import tensorflow as tf
@@ -183,6 +184,8 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/offline/discrete_bc.py

@@ -18,6 +18,7 @@
 import numpy as np
 import sonnet as snt
 import tensorflow as tf
+import jax
 import tensorflow_probability as tfp
 import tree
 import hydra
@@ -202,6 +203,8 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/offline/iddpg_bc.py

@@ -18,6 +18,7 @@
 import copy
 import hydra
 import numpy as np
+import jax
 from omegaconf import DictConfig
 import tensorflow as tf
 from tensorflow import Tensor
@@ -292,6 +293,8 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/offline/iddpg_cql.py

@@ -18,6 +18,7 @@
 import copy
 import hydra
 import numpy as np
+import jax
 from omegaconf import DictConfig
 import tensorflow as tf
 from tensorflow import Tensor
@@ -414,6 +415,8 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/offline/iql_bcq.py

@@ -17,6 +17,7 @@
 
 import copy
 import hydra
+import jax
 from omegaconf import DictConfig
 import sonnet as snt
 import tensorflow as tf
@@ -296,6 +297,8 @@ def _tf_train_step(self, train_step: int, experience: Dict[str, Any]) -> Dict[st
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/offline/iql_cql.py

@@ -16,6 +16,7 @@
 
 import copy
 import numpy as np
+import jax
 import tensorflow as tf
 import sonnet as snt
 import tree
@@ -260,6 +261,8 @@ def _tf_train_step(self, train_step: int, experience: Dict[str, Any]) -> Dict[st
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/offline/maddpg_cql.py

@@ -18,6 +18,7 @@
 
 import hydra
 import numpy as np
+import jax
 import sonnet as snt
 from omegaconf import DictConfig
 import tensorflow as tf
@@ -408,6 +409,8 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/offline/maicq.py

@@ -18,6 +18,7 @@
 import copy
 import hydra
 import numpy as np
+import jax
 from omegaconf import DictConfig
 import sonnet as snt
 import tensorflow as tf
@@ -318,6 +319,8 @@ def _tf_train_step(
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/offline/omar.py

@@ -19,6 +19,7 @@
 import hydra
 import tree
 import numpy as np
+import jax
 from omegaconf import DictConfig
 import tensorflow as tf
 import sonnet as snt
@@ -490,6 +491,8 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/offline/qmix_bcq.py

@@ -19,6 +19,7 @@
 import hydra
 import tree
 import numpy as np
+import jax
 from omegaconf import DictConfig
 import sonnet as snt
 import tensorflow as tf
@@ -333,6 +334,8 @@ def mixing(
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/offline/qmix_cql.py

@@ -19,6 +19,7 @@
 from omegaconf import DictConfig
 import tree
 import numpy as np
+import jax
 import tensorflow as tf
 from tensorflow import Tensor
 import sonnet as snt
@@ -331,6 +332,8 @@ def mixing(
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"])
 
     buffer = FlashbaxReplayBuffer(

og_marl/tf2_systems/online/iddpg.py

@@ -18,6 +18,7 @@
 import copy
 import hydra
 import numpy as np
+import jax
 from omegaconf import DictConfig
 import tensorflow as tf
 from tensorflow import Tensor
@@ -293,6 +294,8 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(
         cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"]
     )

og_marl/tf2_systems/online/maddpg.py

@@ -18,6 +18,7 @@
 import copy
 import hydra
 import numpy as np
+import jax
 from omegaconf import DictConfig
 import tensorflow as tf
 from tensorflow import Tensor
@@ -294,6 +295,8 @@ def _tf_train_step(self, experience: Dict[str, Any]) -> Dict[str, Numeric]:
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(
         cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"]
     )

og_marl/tf2_systems/online/qmix.py

@@ -18,6 +18,7 @@
 import copy
 import hydra
 import numpy as np
+import jax
 from omegaconf import DictConfig
 import tensorflow as tf
 from tensorflow import Tensor
@@ -298,6 +299,8 @@ def mixing(
 def run_experiment(cfg: DictConfig) -> None:
     print(cfg)
 
+    jax.config.update('jax_platform_name', 'cpu')
+
     env = get_environment(
         cfg["task"]["source"], cfg["task"]["env"], cfg["task"]["scenario"], seed=cfg["seed"]
     )