[Feature] Storage Shared Initialization for Multiprocessing

ghstack-source-id: fbbc221
Pull-Request: #3183

Author: vmoens

examples/distributed/collectors/multi_nodes/delayed_dist.py

@@ -150,7 +150,7 @@ def make_env():
         if i == 10:
             t0 = time.time()
     t1 = time.time()
-    torchrl_logger.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    torchrl_logger.info(f"time elapsed: {t1 - t0}s, rate: {counter / (t1 - t0)} fps")
     collector.shutdown()
     exit()
 

examples/distributed/collectors/multi_nodes/delayed_rpc.py

@@ -148,7 +148,7 @@ def make_env():
         if i == 10:
             t0 = time.time()
     t1 = time.time()
-    torchrl_logger.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    torchrl_logger.info(f"time elapsed: {t1 - t0}s, rate: {counter / (t1 - t0)} fps")
     collector.shutdown()
     exit()
 

examples/distributed/collectors/multi_nodes/generic.py

@@ -125,5 +125,5 @@ def gym_make():
             t0 = time.time()
     collector.shutdown()
     t1 = time.time()
-    torchrl_logger.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    torchrl_logger.info(f"time elapsed: {t1 - t0}s, rate: {counter / (t1 - t0)} fps")
     exit()

examples/distributed/collectors/multi_nodes/rpc.py

@@ -113,5 +113,5 @@ def gym_make():
             t0 = time.time()
     collector.shutdown()
     t1 = time.time()
-    torchrl_logger.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    torchrl_logger.info(f"time elapsed: {t1 - t0}s, rate: {counter / (t1 - t0)} fps")
     exit()

examples/distributed/collectors/multi_nodes/sync.py

@@ -119,5 +119,5 @@ def gym_make():
             t0 = time.time()
     collector.shutdown()
     t1 = time.time()
-    torchrl_logger.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    torchrl_logger.info(f"time elapsed: {t1 - t0}s, rate: {counter / (t1 - t0)} fps")
     exit()

examples/distributed/collectors/single_machine/generic.py

@@ -160,5 +160,5 @@ def gym_make():
             t0 = time.time()
     collector.shutdown()
     t1 = time.time()
-    torchrl_logger.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    torchrl_logger.info(f"time elapsed: {t1 - t0}s, rate: {counter / (t1 - t0)} fps")
     exit()

examples/distributed/collectors/single_machine/rpc.py

@@ -129,5 +129,5 @@ def gym_make():
             t0 = time.time()
     collector.shutdown()
     t1 = time.time()
-    torchrl_logger.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    torchrl_logger.info(f"time elapsed: {t1 - t0}s, rate: {counter / (t1 - t0)} fps")
     exit()

examples/distributed/collectors/single_machine/sync.py

@@ -149,5 +149,5 @@ def gym_make():
             t0 = time.time()
     collector.shutdown()
     t1 = time.time()
-    torchrl_logger.info(f"time elapsed: {t1-t0}s, rate: {counter/(t1-t0)} fps")
+    torchrl_logger.info(f"time elapsed: {t1 - t0}s, rate: {counter / (t1 - t0)} fps")
     exit()

examples/rlhf/train.py

@@ -147,7 +147,7 @@ def main(cfg):
         elif it % log_interval == 0:
             # loss as float. note: this is a CPU-GPU sync point
             loss = batch.loss.item()
-            msg = f"TRAIN: {it=}: {loss=:.4f}, time {dt*1000:.2f}ms"
+            msg = f"TRAIN: {it=}: {loss=:.4f}, time {dt * 1000:.2f}ms"
             torchrl_logger.info(msg)
             loss_logger.info(msg)
 

examples/rlhf/train_reward.py

@@ -155,7 +155,7 @@ def main(cfg):
             acc = _accuracy(
                 batch.chosen_data.end_scores, batch.rejected_data.end_scores
             )
-            msg = f"TRAIN: {it=}: {loss=:.4f}, {acc=:.4f} time={dt*1000:.2f}ms"
+            msg = f"TRAIN: {it=}: {loss=:.4f}, {acc=:.4f} time={dt * 1000:.2f}ms"
             torchrl_logger.info(msg)
             loss_logger.info(msg)