[Quality] IMPALA auto-device

ghstack-source-id: 05e59ba466e3db83a824dc9e7c215581c1677ca5
Pull Request resolved: #2654

sota-implementations/impala/config_multi_node_ray.yaml

@@ -24,7 +24,7 @@ ray_init_config:
   storage: null
 
 # Device for the forward and backward passes
-local_device: "cuda:0"
+local_device:
 
 # Resources assigned to each IMPALA rollout collection worker
 remote_worker_resources:

sota-implementations/impala/config_multi_node_submitit.yaml

@@ -3,7 +3,7 @@ env:
   env_name: PongNoFrameskip-v4
 
 # Device for the forward and backward passes
-local_device: "cuda:0"
+local_device:
 
 # SLURM config
 slurm_config:

sota-implementations/impala/config_single_node.yaml

@@ -3,7 +3,7 @@ env:
   env_name: PongNoFrameskip-v4
 
 # Device for the forward and backward passes
-device: "cuda:0"
+device:
 
 # collector
 collector:

sota-implementations/impala/impala_multi_node_ray.py

@@ -32,7 +32,11 @@ def main(cfg: "DictConfig"):  # noqa: F821
     from torchrl.record.loggers import generate_exp_name, get_logger
     from utils import eval_model, make_env, make_ppo_models
 
-    device = torch.device(cfg.local_device)
+    device = cfg.local_device
+    if not device:
+        device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")
+    else:
+        device = torch.device(device)
 
     # Correct for frame_skip
     frame_skip = 4

sota-implementations/impala/impala_multi_node_submitit.py

@@ -34,7 +34,11 @@ def main(cfg: "DictConfig"):  # noqa: F821
     from torchrl.record.loggers import generate_exp_name, get_logger
     from utils import eval_model, make_env, make_ppo_models
 
-    device = torch.device(cfg.local_device)
+    device = cfg.local_device
+    if not device:
+        device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")
+    else:
+        device = torch.device(device)
 
     # Correct for frame_skip
     frame_skip = 4

sota-implementations/impala/impala_single_node.py

@@ -31,7 +31,11 @@ def main(cfg: "DictConfig"):  # noqa: F821
     from torchrl.record.loggers import generate_exp_name, get_logger
     from utils import eval_model, make_env, make_ppo_models
 
-    device = torch.device(cfg.device)
+    device = cfg.device
+    if not device:
+        device = torch.device("cpu" if not torch.cuda.is_available() else "cuda:0")
+    else:
+        device = torch.device(device)
 
     # Correct for frame_skip
     frame_skip = 4
@@ -55,7 +59,6 @@ def main(cfg: "DictConfig"):  # noqa: F821
 
     # Create models (check utils.py)
     actor, critic = make_ppo_models(cfg.env.env_name)
-    actor, critic = actor.to(device), critic.to(device)
 
     # Create collector
     collector = MultiaSyncDataCollector(

sota-implementations/ppo/ppo_atari.py

@@ -178,6 +178,7 @@ def update(batch, num_network_updates):
         # Update the networks
         optim.step()
         return loss.detach().set("alpha", alpha), num_network_updates.clone()
+
     if cfg.compile.compile:
         update = compile_with_warmup(update, mode=compile_mode, warmup=1)
         adv_module = compile_with_warmup(adv_module, mode=compile_mode, warmup=1)

sota-implementations/ppo/ppo_mujoco.py

@@ -163,6 +163,7 @@ def update(batch, num_network_updates):
         # Update the networks
         optim.step()
         return loss.detach().set("alpha", alpha), num_network_updates.clone()
+
     if cfg.compile.compile:
         update = compile_with_warmup(update, mode=compile_mode, warmup=1)
         adv_module = compile_with_warmup(adv_module, mode=compile_mode, warmup=1)