[Major] major refactoring of collectors

benchmarks/ecosystem/gym_env_throughput.py

@@ -27,7 +27,7 @@
 )
 from torchrl.envs import EnvCreator, GymEnv, ParallelEnv
 from torchrl.envs.libs.gym import gym_backend as gym_bc, set_gym_backend
-from torchrl.envs.utils import RandomPolicy
+from torchrl.modules import RandomPolicy
 
 if __name__ == "__main__":
     avail_devices = ("cpu",)

benchmarks/storage/benchmark_sample_latency_over_rpc.py

@@ -144,7 +144,7 @@ def __init__(self, capacity: int):
     rank = args.rank
     storage_type = args.storage
 
-    torchrl_logger.info(f"Rank: {rank}; Storage: {storage_type}")
+    torchrl_logger.debug(f"RANK: {rank}; Storage: {storage_type}")
 
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "29500"

benchmarks/test_collectors_benchmark.py

@@ -18,7 +18,7 @@
 from torchrl.data.utils import CloudpickleWrapper
 from torchrl.envs import EnvCreator, GymEnv, ParallelEnv, StepCounter, TransformedEnv
 from torchrl.envs.libs.dm_control import DMControlEnv
-from torchrl.envs.utils import RandomPolicy
+from torchrl.modules import RandomPolicy
 
 
 def single_collector_setup():

docs/source/reference/envs_api.rst

@@ -198,7 +198,6 @@ Helpers
     :toctree: generated/
     :template: rl_template_fun.rst
 
-    RandomPolicy
     check_env_specs
     exploration_type
     get_available_libraries

docs/source/reference/modules_actors.rst

@@ -20,6 +20,7 @@ TensorDictModules and SafeModules
     SafeModule
     SafeSequential
     TanhModule
+    RandomPolicy
 
 Probabilistic actors
 --------------------

examples/collectors/multi_weight_updates.py

@@ -25,7 +25,7 @@
 from torchrl.data import LazyTensorStorage, ReplayBuffer
 from torchrl.envs.libs.gym import GymEnv
 from torchrl.envs.transforms.module import ModuleTransform
-from torchrl.weight_update.weight_sync_schemes import MultiProcessWeightSyncScheme
+from torchrl.weight_update import MultiProcessWeightSyncScheme
 
 
 def make_module():

examples/collectors/weight_sync_collectors.py

@@ -90,7 +90,7 @@ def example_multi_collector_shared_memory():
     env.close()
 
     # Shared memory is more efficient for frequent updates
-    scheme = SharedMemWeightSyncScheme(strategy="tensordict", auto_register=True)
+    scheme = SharedMemWeightSyncScheme(strategy="tensordict")
 
     print("Creating multi-collector with shared memory...")
     collector = MultiSyncDataCollector(

examples/distributed/collectors/multi_nodes/delayed_dist.py

@@ -116,7 +116,7 @@ def main():
     from torchrl.collectors import MultiSyncDataCollector, SyncDataCollector
     from torchrl.data import Bounded
     from torchrl.envs.libs.gym import GymEnv, set_gym_backend
-    from torchrl.envs.utils import RandomPolicy
+    from torchrl.modules import RandomPolicy
 
     collector_class = SyncDataCollector if num_workers == 1 else MultiSyncDataCollector
     device_str = "device" if num_workers == 1 else "devices"

examples/distributed/collectors/multi_nodes/delayed_rpc.py

@@ -115,7 +115,7 @@ def main():
     from torchrl.collectors import MultiSyncDataCollector, SyncDataCollector
     from torchrl.data import Bounded
     from torchrl.envs.libs.gym import GymEnv, set_gym_backend
-    from torchrl.envs.utils import RandomPolicy
+    from torchrl.modules import RandomPolicy
 
     collector_class = SyncDataCollector if num_workers == 1 else MultiSyncDataCollector
     device_str = "device" if num_workers == 1 else "devices"

examples/distributed/collectors/multi_nodes/generic.py

@@ -14,7 +14,7 @@
 from torchrl.collectors.distributed import DistributedDataCollector
 from torchrl.envs import EnvCreator
 from torchrl.envs.libs.gym import GymEnv, set_gym_backend
-from torchrl.envs.utils import RandomPolicy
+from torchrl.modules import RandomPolicy
 
 parser = ArgumentParser()
 parser.add_argument(

examples/distributed/collectors/multi_nodes/rpc.py

@@ -15,7 +15,7 @@
 from torchrl.collectors.distributed import RPCDataCollector
 from torchrl.envs import EnvCreator
 from torchrl.envs.libs.gym import GymEnv, set_gym_backend
-from torchrl.envs.utils import RandomPolicy
+from torchrl.modules import RandomPolicy
 
 parser = ArgumentParser()
 parser.add_argument(

examples/distributed/collectors/multi_nodes/sync.py

@@ -14,7 +14,7 @@
 from torchrl.collectors.distributed import DistributedSyncDataCollector
 from torchrl.envs import EnvCreator
 from torchrl.envs.libs.gym import GymEnv, set_gym_backend
-from torchrl.envs.utils import RandomPolicy
+from torchrl.modules import RandomPolicy
 
 parser = ArgumentParser()
 parser.add_argument(

examples/distributed/collectors/single_machine/generic.py

@@ -34,7 +34,7 @@
 from torchrl.collectors.distributed import DistributedDataCollector
 from torchrl.envs import EnvCreator, ParallelEnv
 from torchrl.envs.libs.gym import GymEnv, set_gym_backend
-from torchrl.envs.utils import RandomPolicy
+from torchrl.modules import RandomPolicy
 
 parser = ArgumentParser()
 parser.add_argument(

examples/distributed/collectors/single_machine/rpc.py

@@ -30,7 +30,7 @@
 from torchrl.collectors.distributed import RPCDataCollector
 from torchrl.envs import EnvCreator, ParallelEnv
 from torchrl.envs.libs.gym import GymEnv, set_gym_backend
-from torchrl.envs.utils import RandomPolicy
+from torchrl.modules import RandomPolicy
 
 parser = ArgumentParser()
 parser.add_argument(

examples/distributed/collectors/single_machine/sync.py

@@ -31,7 +31,7 @@
 from torchrl.collectors.distributed import DistributedSyncDataCollector
 from torchrl.envs import EnvCreator, ParallelEnv
 from torchrl.envs.libs.gym import GymEnv, set_gym_backend
-from torchrl.envs.utils import RandomPolicy
+from torchrl.modules import RandomPolicy
 
 parser = ArgumentParser()
 parser.add_argument(

examples/distributed/replay_buffers/distributed_replay_buffer.py

@@ -172,7 +172,7 @@ def __init__(self, capacity: int):
 if __name__ == "__main__":
     args = parser.parse_args()
     rank = args.rank
-    torchrl_logger.info(f"Rank: {rank}")
+    torchrl_logger.debug(f"RANK: {rank}")
 
     os.environ["MASTER_ADDR"] = "localhost"
     os.environ["MASTER_PORT"] = "29500"

sota-implementations/expert-iteration/ei_utils.py

@@ -5,7 +5,6 @@
 from __future__ import annotations
 
 import time
-
 from typing import Any, Literal
 
 import torch
@@ -612,7 +611,6 @@ def get_wandb_run_id(wandb_logger):
     """
     try:
         # Wait a bit for wandb to initialize
-        import time
 
         max_attempts = 10
         for attempt in range(max_attempts):

test/llm/test_wrapper.py

@@ -7,8 +7,10 @@
 import argparse
 import gc
 import importlib.util
+import threading
 
 import time
+from concurrent.futures import ThreadPoolExecutor, wait
 from functools import partial
 
 import pytest
@@ -412,8 +414,6 @@ def slow_forward(self, td_input, **kwargs):
 @pytest.fixture
 def monkey_patch_forward_for_instrumentation():
     """Fixture to monkey patch the forward method to add detailed processing event tracking."""
-    import threading
-    import time
 
     # Track processing events
     processing_events = []
@@ -2706,8 +2706,6 @@ def test_batching_min_batch_size_one_immediate_processing(
         monkey_patch_forward_for_timing,
     ):
         """Test that with min_batch_size=1, first request is processed immediately and subsequent ones are grouped."""
-        import time
-        from concurrent.futures import ThreadPoolExecutor, wait
 
         # Create wrapper using helper function
         wrapper = create_batching_test_wrapper(