[Refactor] Put all buffers on CPU in examples (#1645)

Author: Vincent Moens

examples/cql/cql_online.py

@@ -51,7 +51,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         batch_size=cfg.optim.batch_size,
         prb=cfg.replay_buffer.prb,
         buffer_size=cfg.replay_buffer.size,
-        device=device,
+        device="cpu",
     )
 
     # Make Model
@@ -104,7 +104,13 @@ def main(cfg: "DictConfig"):  # noqa: F821
             (actor_losses, q_losses, alpha_losses, alpha_primes) = ([], [], [], [])
             for _ in range(num_updates):
                 # sample from replay buffer
-                sampled_tensordict = replay_buffer.sample().clone()
+                sampled_tensordict = replay_buffer.sample()
+                if sampled_tensordict.device != device:
+                    sampled_tensordict = sampled_tensordict.to(
+                        device, non_blocking=True
+                    )
+                else:
+                    sampled_tensordict = sampled_tensordict.clone()
 
                 loss_td = loss_module(sampled_tensordict)
 

examples/ddpg/ddpg.py

@@ -70,7 +70,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         prb=cfg.replay_buffer.prb,
         buffer_size=cfg.replay_buffer.size,
         buffer_scratch_dir=cfg.replay_buffer.scratch_dir,
-        device=device,
+        device="cpu",
     )
 
     # Create optimizers
@@ -118,7 +118,13 @@ def main(cfg: "DictConfig"):  # noqa: F821
             ) = ([], [])
             for _ in range(num_updates):
                 # Sample from replay buffer
-                sampled_tensordict = replay_buffer.sample().clone()
+                sampled_tensordict = replay_buffer.sample()
+                if sampled_tensordict.device != device:
+                    sampled_tensordict = sampled_tensordict.to(
+                        device, non_blocking=True
+                    )
+                else:
+                    sampled_tensordict = sampled_tensordict.clone()
 
                 # Update critic
                 q_loss, *_ = loss_module.loss_value(sampled_tensordict)

examples/discrete_sac/discrete_sac.py

@@ -201,7 +201,7 @@ def env_factory(num_workers):
         prb=cfg.prb,
         buffer_size=cfg.buffer_size,
         batch_size=cfg.batch_size,
-        device=device,
+        device="cpu",
     )
 
     # Optimizers
@@ -255,7 +255,13 @@ def env_factory(num_workers):
             ) = ([], [], [], [], [], [])
             for _ in range(cfg.frames_per_batch * int(cfg.utd_ratio)):
                 # sample from replay buffer
-                sampled_tensordict = replay_buffer.sample().clone()
+                sampled_tensordict = replay_buffer.sample()
+                if sampled_tensordict.device != device:
+                    sampled_tensordict = sampled_tensordict.to(
+                        device, non_blocking=True
+                    )
+                else:
+                    sampled_tensordict = sampled_tensordict.clone()
 
                 loss_td = loss_module(sampled_tensordict)
 

examples/dqn/dqn.py

@@ -115,7 +115,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         cfg=cfg,
     )
 
-    replay_buffer = make_replay_buffer(device, cfg)
+    replay_buffer = make_replay_buffer("cpu", cfg)
 
     recorder = transformed_env_constructor(
         cfg,

examples/dreamer/dreamer.py

@@ -186,7 +186,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
     )
     print("collector:", collector)
 
-    replay_buffer = make_replay_buffer(device, cfg)
+    replay_buffer = make_replay_buffer("cpu", cfg)
 
     record = Recorder(
         record_frames=cfg.record_frames,

examples/iql/iql_online.py

@@ -218,7 +218,7 @@ def env_factory(num_workers):
 
     # Make Replay Buffer
     replay_buffer = make_replay_buffer(
-        buffer_size=cfg.buffer_size, device=device, batch_size=cfg.batch_size
+        buffer_size=cfg.buffer_size, device="cpu", batch_size=cfg.batch_size
     )
 
     # Optimizers

examples/redq/redq.py

@@ -161,7 +161,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         # ],
     )
 
-    replay_buffer = make_replay_buffer(device, cfg)
+    replay_buffer = make_replay_buffer("cpu", cfg)
 
     recorder = transformed_env_constructor(
         cfg,

examples/sac/sac.py

@@ -70,7 +70,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         prb=cfg.replay_buffer.prb,
         buffer_size=cfg.replay_buffer.size,
         buffer_scratch_dir=cfg.replay_buffer.scratch_dir,
-        device=device,
+        device="cpu",
     )
 
     # Create optimizers
@@ -122,7 +122,13 @@ def main(cfg: "DictConfig"):  # noqa: F821
             )
             for i in range(num_updates):
                 # Sample from replay buffer
-                sampled_tensordict = replay_buffer.sample().clone()
+                sampled_tensordict = replay_buffer.sample()
+                if sampled_tensordict.device != device:
+                    sampled_tensordict = sampled_tensordict.to(
+                        device, non_blocking=True
+                    )
+                else:
+                    sampled_tensordict = sampled_tensordict.clone()
 
                 # Compute loss
                 loss_td = loss_module(sampled_tensordict)

examples/td3/td3.py

@@ -70,7 +70,7 @@ def main(cfg: "DictConfig"):  # noqa: F821
         prb=cfg.replay_buffer.prb,
         buffer_size=cfg.replay_buffer.size,
         buffer_scratch_dir=cfg.replay_buffer.scratch_dir,
-        device=device,
+        device="cpu",
     )
 
     # Create optimizers
@@ -124,7 +124,13 @@ def main(cfg: "DictConfig"):  # noqa: F821
                 update_actor = update_counter % delayed_updates == 0
 
                 # Sample from replay buffer
-                sampled_tensordict = replay_buffer.sample().clone()
+                sampled_tensordict = replay_buffer.sample()
+                if sampled_tensordict.device != device:
+                    sampled_tensordict = sampled_tensordict.to(
+                        device, non_blocking=True
+                    )
+                else:
+                    sampled_tensordict = sampled_tensordict.clone()
 
                 # Compute loss
                 q_loss, *_ = loss_module.value_loss(sampled_tensordict)

torchrl/data/replay_buffers/storages.py

@@ -740,7 +740,7 @@ def _collate_contiguous(x):
 
 
 def _collate_as_tensor(x):
-    return x.contiguous()
+    return x.as_tensor()
 
 
 def _get_default_collate(storage, _is_tensordict=False):