Revert None filtering logic and fix bugs

Author: Luca Carminati

torchrl/collectors/collectors.py

@@ -3843,15 +3843,17 @@ def iterator(self) -> Iterator[TensorDictBase]:
                     # mask buffers if cat, and create a mask if stack
                     if cat_results != "stack":
                         buffers = [None] * self.num_workers
-                        for idx, buffer in enumerate(filter(None.__ne__, self.buffers)):
+                        for worker_idx, buffer in enumerate(
+                            filter(lambda x: x is not None, self.buffers)
+                        ):
                             valid = buffer.get(("collector", "traj_ids")) != -1
                             if valid.ndim > 2:
                                 valid = valid.flatten(0, -2)
                             if valid.ndim == 2:
                                 valid = valid.any(0)
-                            buffers[idx] = buffer[..., valid]
+                            buffers[worker_idx] = buffer[..., valid]
                     else:
-                        for buffer in filter(None.__ne__, self.buffers):
+                        for buffer in filter(lambda x: x is not None, self.buffers):
                             with buffer.unlock_():
                                 buffer.set(
                                     ("collector", "mask"),
@@ -3861,6 +3863,11 @@ def iterator(self) -> Iterator[TensorDictBase]:
                 else:
                     buffers = self.buffers
 
+                # Skip frame counting if this worker didn't send data this iteration
+                # (happens when reusing buffers or on first iteration with some workers)
+                if self.buffers[idx] is None:
+                    continue
+
                 workers_frames[idx] = workers_frames[idx] + buffers[idx].numel()
 
                 if workers_frames[idx] >= self.total_frames:
@@ -3892,7 +3899,7 @@ def iterator(self) -> Iterator[TensorDictBase]:
             if same_device is None:
                 prev_device = None
                 same_device = True
-                for item in filter(None.__ne__, self.buffers):
+                for item in filter(lambda x: x is not None, self.buffers):
                     if prev_device is None:
                         prev_device = item.device
                     else:
@@ -3912,9 +3919,13 @@ def iterator(self) -> Iterator[TensorDictBase]:
                     )
             else:
                 if self._use_buffers is None:
-                    torchrl_logger.warning("use_buffer not specified and not yet inferred from data, assuming `True`.")
+                    torchrl_logger.warning(
+                        "use_buffer not specified and not yet inferred from data, assuming `True`."
+                    )
                 elif not self._use_buffers:
-                    raise RuntimeError("Cannot concatenate results with use_buffers=False")
+                    raise RuntimeError(
+                        "Cannot concatenate results with use_buffers=False"
+                    )
                 try:
                     if same_device:
                         self.out_buffer = torch.cat(
@@ -3926,7 +3937,11 @@ def iterator(self) -> Iterator[TensorDictBase]:
                             cat_results,
                         )
                 except RuntimeError as err:
-                    if preempt and cat_results != -1 and "Sizes of tensors must match" in str(err):
+                    if (
+                        preempt
+                        and cat_results != -1
+                        and "Sizes of tensors must match" in str(err)
+                    ):
                         raise RuntimeError(
                             "The value provided to cat_results isn't compatible with the collectors outputs. "
                             "Consider using `cat_results=-1`."
@@ -3944,7 +3959,11 @@ def iterator(self) -> Iterator[TensorDictBase]:
             self._frames += n_collected
 
             if self.postprocs:
-                self.postprocs = self.postprocs.to(out.device) if hasattr(self.postprocs, "to") else self.postprocs
+                self.postprocs = (
+                    self.postprocs.to(out.device)
+                    if hasattr(self.postprocs, "to")
+                    else self.postprocs
+                )
                 out = self.postprocs(out)
             if self._exclude_private_keys:
                 excluded_keys = [key for key in out.keys() if key.startswith("_")]