Add NPU Support for Single Agent

lowdy1 · lowdy1 · commit 94b95e2ce92d · 2025-11-04T11:57:41.000Z
diff --git a/sota-implementations/a2c/a2c_atari.py b/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)
 
diff --git a/sota-implementations/a2c/a2c_mujoco.py b/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)
 
diff --git a/sota-implementations/cql/cql_offline.py b/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)
diff --git a/sota-implementations/cql/cql_online.py b/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)
diff --git a/sota-implementations/cql/discrete_cql_offline.py b/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
diff --git a/sota-implementations/cql/discrete_cql_online.py b/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)
diff --git a/sota-implementations/crossq/crossq.py b/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)
diff --git a/sota-implementations/ddpg/ddpg.py b/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)
diff --git a/sota-implementations/decision_transformer/dt.py b/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)
diff --git a/sota-implementations/decision_transformer/online_dt.py b/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)
diff --git a/sota-implementations/dqn/dqn_atari.py b/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)
diff --git a/sota-implementations/dqn/dqn_cartpole.py b/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)
diff --git a/sota-implementations/gail/gail.py b/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)
diff --git a/sota-implementations/iql/discrete_iql.py b/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)
diff --git a/sota-implementations/iql/iql_offline.py b/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,
diff --git a/sota-implementations/iql/iql_online.py b/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)
diff --git a/sota-implementations/ppo/ppo_atari.py b/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)
diff --git a/sota-implementations/ppo/ppo_mujoco.py b/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)
diff --git a/sota-implementations/sac/sac-async.py b/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)
diff --git a/sota-implementations/sac/sac.py b/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)
diff --git a/sota-implementations/td3/td3.py b/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:
diff --git a/sota-implementations/td3_bc/td3_bc.py b/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)
