[BugFix] Load collector frames and iter (#1557)

Signed-off-by: Matteo Bettini <matbet@meta.com>

Author: matteobettini

test/test_collector.py

@@ -1740,6 +1740,49 @@ def test_param_sync(self, give_weights, collector, policy_device, env_device):
             col.shutdown()
 
 
+@pytest.mark.parametrize(
+    "collector_class",
+    [MultiSyncDataCollector, MultiaSyncDataCollector, SyncDataCollector],
+)
+def test_collector_reloading(collector_class):
+    def make_env():
+        return ContinuousActionVecMockEnv()
+
+    dummy_env = make_env()
+    obs_spec = dummy_env.observation_spec["observation"]
+    policy_module = nn.Linear(obs_spec.shape[-1], dummy_env.action_spec.shape[-1])
+    policy = Actor(policy_module, spec=dummy_env.action_spec)
+    policy_explore = OrnsteinUhlenbeckProcessWrapper(policy)
+
+    collector_kwargs = {
+        "create_env_fn": make_env,
+        "policy": policy_explore,
+        "frames_per_batch": 30,
+        "total_frames": 90,
+    }
+    if collector_class is not SyncDataCollector:
+        collector_kwargs["create_env_fn"] = [
+            collector_kwargs["create_env_fn"] for _ in range(3)
+        ]
+
+    collector = collector_class(**collector_kwargs)
+    for i, _ in enumerate(collector):
+        if i == 3:
+            break
+    collector_frames = collector._frames
+    collector_iter = collector._iter
+    collector_state_dict = collector.state_dict()
+    collector.shutdown()
+
+    collector = collector_class(**collector_kwargs)
+    collector.load_state_dict(collector_state_dict)
+    assert collector._frames == collector_frames
+    assert collector._iter == collector_iter
+    for _ in enumerate(collector):
+        raise AssertionError
+    collector.shutdown()
+
+
 if __name__ == "__main__":
     args, unknown = argparse.ArgumentParser().parse_known_args()
     pytest.main([__file__, "--capture", "no", "--exitfirst"] + unknown)

torchrl/collectors/collectors.py

@@ -695,6 +695,8 @@ def __init__(
         self.split_trajs = split_trajs
         self._exclude_private_keys = True
         self.interruptor = interruptor
+        self._frames = 0
+        self._iter = -1
 
     # for RPC
     def next(self):
@@ -745,11 +747,9 @@ def iterator(self) -> Iterator[TensorDictBase]:
             stream = None
         with torch.cuda.stream(stream):
             total_frames = self.total_frames
-            i = -1
-            self._frames = 0
-            while True:
-                i += 1
-                self._iter = i
+
+            while self._frames < self.total_frames:
+                self._iter += 1
                 tensordict_out = self.rollout()
                 self._frames += tensordict_out.numel()
                 if self._frames >= total_frames:
@@ -788,9 +788,6 @@ def iterator(self) -> Iterator[TensorDictBase]:
                     # >>> assert data0["done"] is not data1["done"]
                     yield tensordict_out.clone()
 
-                if self._frames >= self.total_frames:
-                    break
-
     def _step_and_maybe_reset(self) -> None:
 
         any_done = False
@@ -985,6 +982,8 @@ def state_dict(self) -> OrderedDict:
         else:
             state_dict = OrderedDict(env_state_dict=env_state_dict)
 
+        state_dict.update({"frames": self._frames, "iter": self._iter})
+
         return state_dict
 
     def load_state_dict(self, state_dict: OrderedDict, **kwargs) -> None:
@@ -1000,6 +999,8 @@ def load_state_dict(self, state_dict: OrderedDict, **kwargs) -> None:
             self.env.load_state_dict(state_dict["env_state_dict"], **kwargs)
         if strict or "policy_state_dict" in state_dict:
             self.policy.load_state_dict(state_dict["policy_state_dict"], **kwargs)
+        self._frames = state_dict["frames"]
+        self._iter = state_dict["iter"]
 
     def __repr__(self) -> str:
         env_str = indent(f"env={self.env}", 4 * " ")
@@ -1284,6 +1285,8 @@ def device_err_msg(device_name, devices_list):
             self.interruptor = None
         self._run_processes()
         self._exclude_private_keys = True
+        self._frames = 0
+        self._iter = -1
 
     @property
     def frames_per_batch_worker(self):
@@ -1471,6 +1474,7 @@ def state_dict(self) -> OrderedDict:
             if msg != "state_dict":
                 raise RuntimeError(f"Expected msg='state_dict', got {msg}")
             state_dict[f"worker{idx}"] = _state_dict
+        state_dict.update({"frames": self._frames, "iter": self._iter})
 
         return state_dict
 
@@ -1488,6 +1492,8 @@ def load_state_dict(self, state_dict: OrderedDict) -> None:
             _, msg = self.pipes[idx].recv()
             if msg != "loaded":
                 raise RuntimeError(f"Expected msg='loaded', got {msg}")
+        self._frames = state_dict["frames"]
+        self._iter = state_dict["iter"]
 
 
 @accept_remote_rref_udf_invocation
@@ -1639,27 +1645,26 @@ def _queue_len(self) -> int:
         return self.num_workers
 
     def iterator(self) -> Iterator[TensorDictBase]:
-        i = -1
-        frames = 0
+
         self.buffers = {}
         dones = [False for _ in range(self.num_workers)]
         workers_frames = [0 for _ in range(self.num_workers)]
         same_device = None
         self.out_buffer = None
 
-        while not all(dones) and frames < self.total_frames:
+        while not all(dones) and self._frames < self.total_frames:
             _check_for_faulty_process(self.procs)
             if self.update_at_each_batch:
                 self.update_policy_weights_()
 
             for idx in range(self.num_workers):
-                if frames < self.init_random_frames:
+                if self._frames < self.init_random_frames:
                     msg = "continue_random"
                 else:
                     msg = "continue"
                 self.pipes[idx].send((None, msg))
 
-            i += 1
+            self._iter += 1
             max_traj_idx = None
 
             if self.interruptor is not None and self.preemptive_threshold < 1.0:
@@ -1714,10 +1719,10 @@ def iterator(self) -> Iterator[TensorDictBase]:
 
             if self.split_trajs:
                 out = split_trajectories(self.out_buffer, prefix="collector")
-                frames += out.get(("collector", "mask")).sum().item()
+                self._frames += out.get(("collector", "mask")).sum().item()
             else:
                 out = self.out_buffer.clone()
-                frames += prod(out.shape)
+                self._frames += prod(out.shape)
             if self.postprocs:
                 self.postprocs = self.postprocs.to(out.device)
                 out = self.postprocs(out)
@@ -1894,13 +1899,11 @@ def iterator(self) -> Iterator[TensorDictBase]:
             else:
                 self.pipes[i].send((None, "continue"))
         self.running = True
-        i = -1
-        self._frames = 0
 
         workers_frames = [0 for _ in range(self.num_workers)]
         while self._frames < self.total_frames:
             _check_for_faulty_process(self.procs)
-            i += 1
+            self._iter += 1
             idx, j, out = self._get_from_queue()
 
             worker_frames = out.numel()