[Feature] Add NPU Support for Single Agent

sota-implementations/a2c/a2c_atari.py

@@ -35,7 +35,12 @@ def main(cfg: DictConfig):  # noqa: F821
 
     device = cfg.loss.device
     if not device:
-        device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")
+        if torch.cuda.is_available():
+            device = torch.device("cuda:0")
+        elif torch.npu.is_available():
+            device = torch.device("npu:0")
+        else:
+            device = torch.device("cpu")
     else:
         device = torch.device(device)
 

sota-implementations/a2c/a2c_mujoco.py

@@ -38,7 +38,12 @@ def main(cfg: DictConfig):  # noqa: F821
 
     device = cfg.loss.device
     if not device:
-        device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")
+        if torch.cuda.is_available():
+            device = torch.device("cuda:0")
+        elif torch.npu.is_available():
+            device = torch.device("npu:0")
+        else:
+            device = torch.device("cpu")
     else:
         device = torch.device(device)
 

sota-implementations/cql/cql_offline.py

@@ -59,6 +59,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)

sota-implementations/cql/cql_online.py

@@ -64,6 +64,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)

sota-implementations/cql/discrete_cql_offline.py

@@ -38,7 +38,12 @@
 def main(cfg):  # noqa: F821
     device = cfg.optim.device
     if device in ("", None):
-        device = "cuda:0" if torch.cuda.is_available() else "cpu"
+        if torch.cuda.is_available():
+            device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
+        else:
+            device = "cpu"
     device = torch.device(device)
 
     # Create logger

sota-implementations/cql/discrete_cql_online.py

@@ -42,6 +42,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)

sota-implementations/crossq/crossq.py

@@ -44,6 +44,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = torch.device("cuda:0")
+        elif torch.npu.is_available():
+            device = torch.device("npu:0")
         else:
             device = torch.device("cpu")
     device = torch.device(device)

sota-implementations/ddpg/ddpg.py

@@ -43,6 +43,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)
@@ -51,6 +53,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if collector_device in ("", None):
         if torch.cuda.is_available():
             collector_device = "cuda:0"
+        elif torch.npu.is_available():
+            collector_device = "npu:0"
         else:
             collector_device = "cpu"
     collector_device = torch.device(collector_device)

sota-implementations/decision_transformer/dt.py

@@ -42,6 +42,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if model_device in ("", None):
         if torch.cuda.is_available():
             model_device = "cuda:0"
+        elif torch.npu.is_available():
+            model_device = "npu:0"
         else:
             model_device = "cpu"
     model_device = torch.device(model_device)

sota-implementations/decision_transformer/online_dt.py

@@ -40,6 +40,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if model_device in ("", None):
         if torch.cuda.is_available():
             model_device = "cuda:0"
+        elif torch.npu.is_available():
+            model_device = "npu:0"
         else:
             model_device = "cpu"
     model_device = torch.device(model_device)

sota-implementations/dqn/dqn_atari.py

@@ -37,6 +37,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)

sota-implementations/dqn/dqn_cartpole.py

@@ -32,6 +32,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)

sota-implementations/gail/gail.py

@@ -42,6 +42,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)

sota-implementations/iql/discrete_iql.py

@@ -67,6 +67,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)

sota-implementations/iql/iql_offline.py

@@ -63,11 +63,13 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)
 
-    # Creante env
+    # Create env
     train_env, eval_env = make_environment(
         cfg,
         cfg.logger.eval_envs,

sota-implementations/iql/iql_online.py

@@ -66,6 +66,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)

sota-implementations/ppo/ppo_atari.py

@@ -41,6 +41,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)

sota-implementations/ppo/ppo_mujoco.py

@@ -41,6 +41,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)

sota-implementations/sac/sac-async.py

@@ -56,6 +56,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = torch.device("cuda:0")
+        elif torch.npu.is_available():
+            device = torch.device("npu:0")
         else:
             device = torch.device("cpu")
     device = torch.device(device)

sota-implementations/sac/sac.py

@@ -45,6 +45,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = torch.device("cuda:0")
+        elif torch.npu.is_available():
+            device = torch.device("npu:0")
         else:
             device = torch.device("cpu")
     device = torch.device(device)

sota-implementations/td3/td3.py

@@ -43,6 +43,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = torch.device("cuda:0")
+        elif torch.npu.is_available():
+            device = torch.device("npu:0")
         else:
             device = torch.device("cpu")
     else:

sota-implementations/td3_bc/td3_bc.py

@@ -61,6 +61,8 @@ def main(cfg: DictConfig):  # noqa: F821
     if device in ("", None):
         if torch.cuda.is_available():
             device = "cuda:0"
+        elif torch.npu.is_available():
+            device = "npu:0"
         else:
             device = "cpu"
     device = torch.device(device)