Update

[ghstack-poisoned]

Author: vmoens

examples/collectors/weight_sync_standalone.py

@@ -129,7 +129,7 @@ def example_multiprocess_sync():
         print(
             f"Main: Worker completed. Worker's weight sum: {model_state['weight_sum']:.4f}"
         )
-        print(f"✓ Weight synchronization successful!")
+        print("Weight synchronization successful!")
 
 
 def example_shared_memory_sync():
@@ -163,7 +163,7 @@ def example_shared_memory_sync():
         weights_td["weight"].fill_(2.0)
         weights_td["bias"].fill_(1.0)
 
-    print(f"Main: Sending weights via shared memory...")
+    print("Main: Sending weights via shared memory...")
     sender.update_weights(weights_td)
 
     # Workers automatically see updates via shared memory!
@@ -179,7 +179,7 @@ def example_shared_memory_sync():
         print(
             f"Main: Worker completed. Worker's weight sum: {model_state['weight_sum']:.4f}"
         )
-        print(f"✓ Shared memory synchronization successful!")
+        print("Shared memory synchronization successful!")
 
 
 def main():

test/test_weightsync.py

@@ -10,10 +10,7 @@
 import torch
 import torch.nn as nn
 from tensordict import TensorDict
-from tensordict.nn import TensorDictModule
 from torch import multiprocessing as mp
-from torchrl.collectors import SyncDataCollector
-from torchrl.envs import GymEnv
 
 from torchrl.weight_update.weight_sync_schemes import (
     _resolve_model,