[BugFix] Replay Buffer prefetch & SliceSampler

Fix pickling for `ReplayBuffer` with prefetch by removing non-picklable
prefetch objects and recreating them on unpickle.
Fix logging typo in `SliceSampler`.


ghstack-source-id: 7a8b94a
Pull-Request: #3322

Author: vmoens

torchrl/data/replay_buffers/replay_buffers.py

@@ -1158,6 +1158,13 @@ def __getstate__(self) -> dict[str, Any]:
             state["_replay_lock_placeholder"] = None
         if _futures_lock is not None:
             state["_futures_lock_placeholder"] = None
+        # Remove non-picklable prefetch objects - they will be recreated on unpickle
+        _prefetch_queue = state.pop("_prefetch_queue", None)
+        _prefetch_executor = state.pop("_prefetch_executor", None)
+        if _prefetch_queue is not None:
+            state["_prefetch_queue_placeholder"] = None
+        if _prefetch_executor is not None:
+            state["_prefetch_executor_placeholder"] = None
         return state
 
     def __setstate__(self, state: dict[str, Any]):
@@ -1176,6 +1183,15 @@ def __setstate__(self, state: dict[str, Any]):
             state.pop("_futures_lock_placeholder")
             _futures_lock = threading.RLock()
             state["_futures_lock"] = _futures_lock
+        # Recreate prefetch objects after unpickling if they were present
+        if "_prefetch_queue_placeholder" in state:
+            state.pop("_prefetch_queue_placeholder")
+            state["_prefetch_queue"] = collections.deque()
+        if "_prefetch_executor_placeholder" in state:
+            state.pop("_prefetch_executor_placeholder")
+            state["_prefetch_executor"] = ThreadPoolExecutor(
+                max_workers=state["_prefetch_cap"]
+            )
         self.__dict__.update(state)
         if rngstate is not None:
             self.set_rng(rng)

torchrl/data/replay_buffers/samplers.py

@@ -1154,10 +1154,10 @@ def __init__(
     def __getstate__(self):
         if get_spawning_popen() is not None and self.cache_values:
             logger.warning(
-                f"It seems you are sharing a {type(self).__name__} across processes with"
+                f"It seems you are sharing a {type(self).__name__} across processes with "
                 f"cache_values=True. "
                 f"While this isn't forbidden and could perfectly work if your dataset "
-                f"is unaltered on both processes, remember that calling extend/add on"
+                f"is unaltered on both processes, remember that calling extend/add on "
                 f"one process will NOT erase the cache on another process's sampler, "
                 f"which will cause synchronization issues."
             )