amend

pytorch · vmoens · May 31, 2024 · Apr 30, 2024 · May 26, 2024 · May 26, 2024
commit 3e579c1a91630025645801ffbdddaac0579c6f79
diff --git a/test/test_collector.py b/test/test_collector.py
@@ -798,12 +798,8 @@ def make_env(seed):
 
 
 @pytest.mark.parametrize("num_env", [1, 2])
-@pytest.mark.parametrize(
-    "env_name",
-    [
-        "vec",
-    ],
-)  # 1226: for efficiency, we just test vec, not "conv"
+# 1226: for efficiency, we just test vec, not "conv"
+@pytest.mark.parametrize("env_name", ["vec"])
 def test_collector_batch_size(
     num_env, env_name, seed=100, num_workers=2, frames_per_batch=20
 ):
@@ -1428,6 +1424,7 @@ def env_fn(seed):
         device=device,
         storing_device=storing_device,
     )
+    assert collector._use_buffers
     batch = next(collector.iterator())
     assert batch.device == torch.device(storing_device)
     collector.shutdown()
@@ -2581,6 +2578,7 @@ def test_unique_traj_sync(self, cat_results):
         try:
             for d in c:
                 buffer.extend(d)
+            assert c._use_buffers
             traj_ids = buffer[:].get(("collector", "traj_ids"))
             # check that we have as many trajs as expected (no skip)
             assert traj_ids.unique().numel() == traj_ids.max() + 1
@@ -2611,6 +2609,7 @@ def test_dynamic_sync_collector(self):
         )
         for data in collector:
             assert isinstance(data, LazyStackedTensorDict)
+            assert data.names[-1] == "time"
 
     def test_dynamic_multisync_collector(self):
         env = EnvWithDynamicSpec
@@ -2625,6 +2624,7 @@ def test_dynamic_multisync_collector(self):
         )
         for data in collector:
             assert isinstance(data, LazyStackedTensorDict)
+            assert data.names[-1] == "time"
 
     def test_dynamic_multiasync_collector(self):
         env = EnvWithDynamicSpec
@@ -2638,6 +2638,7 @@ def test_dynamic_multiasync_collector(self):
         )
         for data in collector:
             assert isinstance(data, LazyStackedTensorDict)
+            assert data.names[-1] == "time"
 
 
 if __name__ == "__main__":

diff --git a/torchrl/collectors/collectors.py b/torchrl/collectors/collectors.py
@@ -200,7 +200,7 @@ def update_policy_weights_(
             self.policy_weights.data.update_(self.get_weights_fn())
 
     def __iter__(self) -> Iterator[TensorDictBase]:
-        return self.iterator()
+        yield from self.iterator()
 
     def next(self):
         try:
@@ -796,6 +796,8 @@ def filter_policy(value_output, value_input, value_input_clone):
         )
         self._final_rollout.refine_names(..., "time")
 
+        assert self._final_rollout.names[-1] == "time"
+
     def _set_truncated_keys(self):
         self._truncated_keys = []
         if self.set_truncated:
@@ -1080,25 +1082,30 @@ def rollout(self) -> TensorDictBase:
                                 )
                     else:
                         result = TensorDict.maybe_dense_stack(tensordicts, dim=-1)
+                        assert result.names[-1] == "time"
                     break
             else:
                 if self._use_buffers:
                     result = self._final_rollout
                     try:
-                        self._final_rollout = torch.stack(
+                        result = torch.stack(
                             tensordicts,
                             self._final_rollout.ndim - 1,
                             out=self._final_rollout,
                         )
+                        assert result.names[-1] == "time"
+
                     except RuntimeError:
                         with self._final_rollout.unlock_():
-                            self._final_rollout = torch.stack(
+                            result = torch.stack(
                                 tensordicts,
                                 self._final_rollout.ndim - 1,
                                 out=self._final_rollout,
                             )
+                            assert result.names[-1] == "time"
                 else:
                     result = TensorDict.maybe_dense_stack(tensordicts, dim=-1)
+                    result.refine_names(..., "time")
 
         return self._maybe_set_truncated(result)
 
@@ -2213,7 +2220,11 @@ def iterator(self) -> Iterator[TensorDictBase]:
                     ("collector", "traj_ids"), torch.stack(traj_ids_list), inplace=True
                 )
             else:
-                if not self._use_buffers:
+                if self._use_buffers is None:
+                    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"
                     )
@@ -2455,7 +2466,6 @@ def iterator(self) -> Iterator[TensorDictBase]:
             _check_for_faulty_process(self.procs)
             self._iter += 1
             idx, j, out = self._get_from_queue()
-
             worker_frames = out.numel()
             if self.split_trajs:
                 out = split_trajectories(out, prefix="collector")
@@ -2854,6 +2864,7 @@ def _main_async_collector(
                             else x
                         )
                 data = (collected_tensordict, idx)
+                assert collected_tensordict.names[-1] == "time"
             else:
                 if next_data is not collected_tensordict:
                     raise RuntimeError(