[Feature] Add NPU Support for Single Agent (#3229)

Co-authored-by: vmoens <[email protected]>

Author: lowdy1

docs/source/reference/utils.rst

@@ -10,6 +10,7 @@ Set of utility methods that are used internally by the library.
     :toctree: generated/
     :template: rl_template.rst
 
+    get_available_device
     implement_for
     set_auto_unwrap_transformed_env
     auto_unwrap_transformed_env

sota-implementations/a2c/a2c_atari.py

@@ -22,7 +22,7 @@ def main(cfg: DictConfig):  # noqa: F821
     from tensordict import from_module
     from tensordict.nn import CudaGraphModule
 
-    from torchrl._utils import timeit
+    from torchrl._utils import get_available_device, timeit
     from torchrl.collectors import SyncDataCollector
     from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
     from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
@@ -33,11 +33,9 @@ def main(cfg: DictConfig):  # noqa: F821
     from torchrl.record.loggers import generate_exp_name, get_logger
     from utils_atari import eval_model, make_parallel_env, make_ppo_models
 
-    device = cfg.loss.device
-    if not device:
-        device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")
-    else:
-        device = torch.device(device)
+    device = (
+        torch.device(cfg.loss.device) if cfg.loss.device else get_available_device()
+    )
 
     # Correct for frame_skip
     frame_skip = 4

sota-implementations/a2c/a2c_mujoco.py

@@ -23,7 +23,7 @@ def main(cfg: DictConfig):  # noqa: F821
     from tensordict import from_module
     from tensordict.nn import CudaGraphModule
 
-    from torchrl._utils import timeit
+    from torchrl._utils import get_available_device, timeit
     from torchrl.collectors import SyncDataCollector
     from torchrl.data import LazyTensorStorage, TensorDictReplayBuffer
     from torchrl.data.replay_buffers.samplers import SamplerWithoutReplacement
@@ -36,11 +36,9 @@ def main(cfg: DictConfig):  # noqa: F821
 
     # Define paper hyperparameters
 
-    device = cfg.loss.device
-    if not device:
-        device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")
-    else:
-        device = torch.device(device)
+    device = (
+        torch.device(cfg.loss.device) if cfg.loss.device else get_available_device()
+    )
 
     num_mini_batches = cfg.collector.frames_per_batch // cfg.loss.mini_batch_size
     total_network_updates = (

sota-implementations/cql/cql_offline.py

@@ -18,7 +18,7 @@
 import torch
 import tqdm
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
@@ -55,13 +55,9 @@ def main(cfg: DictConfig):  # noqa: F821
     # Set seeds
     torch.manual_seed(cfg.env.seed)
     np.random.seed(cfg.env.seed)
-    device = cfg.optim.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = "cuda:0"
-        else:
-            device = "cpu"
-    device = torch.device(device)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Create replay buffer
     replay_buffer = make_offline_replay_buffer(cfg.replay_buffer)

sota-implementations/cql/cql_online.py

@@ -21,7 +21,7 @@
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
@@ -60,13 +60,9 @@ def main(cfg: DictConfig):  # noqa: F821
     # Set seeds
     torch.manual_seed(cfg.env.seed)
     np.random.seed(cfg.env.seed)
-    device = cfg.optim.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = "cuda:0"
-        else:
-            device = "cpu"
-    device = torch.device(device)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Create env
     train_env, eval_env = make_environment(

sota-implementations/cql/discrete_cql_offline.py

@@ -18,7 +18,7 @@
 import torch
 import tqdm
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.record.loggers import generate_exp_name, get_logger
 from utils import (
@@ -36,10 +36,9 @@
 
 @hydra.main(version_base="1.1", config_path="", config_name="discrete_offline_config")
 def main(cfg):  # noqa: F821
-    device = cfg.optim.device
-    if device in ("", None):
-        device = "cuda:0" if torch.cuda.is_available() else "cpu"
-    device = torch.device(device)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Create logger
     exp_name = generate_exp_name("DiscreteCQL", cfg.logger.exp_name)

sota-implementations/cql/discrete_cql_online.py

@@ -20,7 +20,7 @@
 import torch.cuda
 import tqdm
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.record.loggers import generate_exp_name, get_logger
 from utils import (
@@ -38,13 +38,9 @@
 
 @hydra.main(version_base="1.1", config_path="", config_name="discrete_online_config")
 def main(cfg: DictConfig):  # noqa: F821
-    device = cfg.optim.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = "cuda:0"
-        else:
-            device = "cpu"
-    device = torch.device(device)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Create logger
     exp_name = generate_exp_name("DiscreteCQL", cfg.logger.exp_name)

sota-implementations/crossq/crossq.py

@@ -21,7 +21,7 @@
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
@@ -40,13 +40,11 @@
 
 @hydra.main(version_base="1.1", config_path=".", config_name="config")
 def main(cfg: DictConfig):  # noqa: F821
-    device = cfg.network.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = torch.device("cuda:0")
-        else:
-            device = torch.device("cpu")
-    device = torch.device(device)
+    device = (
+        torch.device(cfg.network.device)
+        if cfg.network.device
+        else get_available_device()
+    )
 
     # Create logger
     exp_name = generate_exp_name("CrossQ", cfg.logger.exp_name)

sota-implementations/ddpg/ddpg.py

@@ -21,7 +21,7 @@
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import timeit
+from torchrl._utils import get_available_device, timeit
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.objectives import group_optimizers
 from torchrl.record.loggers import generate_exp_name, get_logger
@@ -39,21 +39,14 @@
 
 @hydra.main(version_base="1.1", config_path="", config_name="config")
 def main(cfg: DictConfig):  # noqa: F821
-    device = cfg.optim.device
-    if device in ("", None):
-        if torch.cuda.is_available():
-            device = "cuda:0"
-        else:
-            device = "cpu"
-    device = torch.device(device)
-
-    collector_device = cfg.collector.device
-    if collector_device in ("", None):
-        if torch.cuda.is_available():
-            collector_device = "cuda:0"
-        else:
-            collector_device = "cpu"
-    collector_device = torch.device(collector_device)
+    device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
+    collector_device = (
+        torch.device(cfg.collector.device)
+        if cfg.collector.device
+        else get_available_device()
+    )
 
     # Create logger
     exp_name = generate_exp_name("DDPG", cfg.logger.exp_name)

sota-implementations/decision_transformer/dt.py

@@ -17,7 +17,7 @@
 import tqdm
 from tensordict import TensorDict
 from tensordict.nn import CudaGraphModule
-from torchrl._utils import logger as torchrl_logger, timeit
+from torchrl._utils import get_available_device, logger as torchrl_logger, timeit
 from torchrl.envs.libs.gym import set_gym_backend
 from torchrl.envs.utils import ExplorationType, set_exploration_type
 from torchrl.modules.tensordict_module import DecisionTransformerInferenceWrapper
@@ -38,13 +38,9 @@
 def main(cfg: DictConfig):  # noqa: F821
     set_gym_backend(cfg.env.backend).set()
 
-    model_device = cfg.optim.device
-    if model_device in ("", None):
-        if torch.cuda.is_available():
-            model_device = "cuda:0"
-        else:
-            model_device = "cpu"
-    model_device = torch.device(model_device)
+    model_device = (
+        torch.device(cfg.optim.device) if cfg.optim.device else get_available_device()
+    )
 
     # Set seeds
     torch.manual_seed(cfg.env.seed)