Revert None filtering logic and fix bugs

Author: Luca Carminati

torchrl/collectors/collectors.py

@@ -3843,15 +3843,15 @@ 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 +3861,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 +3897,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: