[Feature] Enabling worker level control on frames_per_batch (#3020)

Author: Ghelfi

test/test_collector.py

@@ -1553,6 +1553,94 @@ def create_env():
             collector.shutdown()
             del collector
 
+    @pytest.mark.parametrize("num_env", [1, 2])
+    @pytest.mark.parametrize("env_name", ["vec"])
+    @pytest.mark.parametrize("frames_per_batch_worker", [[10, 10], [15, 5]])
+    def test_collector_frames_per_batch_worker(
+        self,
+        num_env,
+        env_name,
+        frames_per_batch_worker,
+        seed=100,
+        num_workers=2,
+    ):
+        """Tests that there are 'sum(frames_per_batch_worker)' frames in each batch of a collection."""
+        if num_env == 1:
+
+            def env_fn():
+                env = make_make_env(env_name)()
+                return env
+
+        else:
+
+            def env_fn():
+                # 1226: For efficiency, we don't use Parallel but Serial
+                # env = ParallelEnv(
+                env = SerialEnv(
+                    num_workers=num_env, create_env_fn=make_make_env(env_name)
+                )
+                return env
+
+        policy = make_policy(env_name)
+
+        torch.manual_seed(0)
+        np.random.seed(0)
+
+        frames_per_batch = sum(frames_per_batch_worker)
+
+        collector = MultiaSyncDataCollector(
+            create_env_fn=[env_fn for _ in range(num_workers)],
+            policy=policy,
+            frames_per_batch=frames_per_batch_worker,
+            max_frames_per_traj=1000,
+            total_frames=frames_per_batch * 100,
+        )
+        try:
+            collector.set_seed(seed)
+            for i, b in enumerate(collector):
+                assert b.numel() == -(-frames_per_batch // num_env) * num_env
+                if i == 5:
+                    break
+            assert b.names[-1] == "time"
+        finally:
+            collector.shutdown()
+
+        collector = MultiSyncDataCollector(
+            create_env_fn=[env_fn for _ in range(num_workers)],
+            policy=policy,
+            frames_per_batch=frames_per_batch,
+            max_frames_per_traj=1000,
+            total_frames=frames_per_batch * 100,
+            cat_results="stack",
+        )
+        try:
+            collector.set_seed(seed)
+            for i, b in enumerate(collector):
+                assert (
+                    b.numel()
+                    == -(-frames_per_batch // num_env // num_workers)
+                    * num_env
+                    * num_workers
+                )
+                if i == 5:
+                    break
+            assert b.names[-1] == "time"
+        finally:
+            collector.shutdown()
+            del collector
+
+        with pytest.raises(
+            ValueError,
+            match="If `frames_per_batch` is provided as a sequence, it should contain exactly one value per worker.",
+        ):
+            collector = MultiSyncDataCollector(
+                create_env_fn=[env_fn for _ in range(num_workers)],
+                policy=policy,
+                frames_per_batch=frames_per_batch_worker[:-1],
+                max_frames_per_traj=1000,
+                total_frames=frames_per_batch * 100,
+            )
+
 
 class TestCollectorDevices:
     class DeviceLessEnv(EnvBase):

torchrl/collectors/collectors.py

@@ -1765,8 +1765,9 @@ class _MultiDataCollector(DataCollectorBase):
             .. warning:: `policy_factory` is currently not compatible with multiprocessed data
                 collectors.
 
-        frames_per_batch (int): A keyword-only argument representing the
-            total number of elements in a batch.
+        frames_per_batch (int, Sequence[int]): A keyword-only argument representing the
+            total number of elements in a batch. If a sequence is provided, represents the number of elements in a
+            batch per worker. Total number of elements in a batch is then the sum over the sequence.
         total_frames (int, optional): A keyword-only argument representing the
             total number of frames returned by the collector
             during its lifespan. If the ``total_frames`` is not divisible by
@@ -1923,7 +1924,7 @@ def __init__(
         policy_factory: Callable[[], Callable]
         | list[Callable[[], Callable]]
         | None = None,
-        frames_per_batch: int,
+        frames_per_batch: int | Sequence[int],
         total_frames: int | None = -1,
         device: DEVICE_TYPING | Sequence[DEVICE_TYPING] | None = None,
         storing_device: DEVICE_TYPING | Sequence[DEVICE_TYPING] | None = None,
@@ -1959,6 +1960,22 @@ def __init__(
         self.closed = True
         self.num_workers = len(create_env_fn)
 
+        if (
+            isinstance(frames_per_batch, Sequence)
+            and len(frames_per_batch) != self.num_workers
+        ):
+            raise ValueError(
+                "If `frames_per_batch` is provided as a sequence, it should contain exactly one value per worker."
+                f"Got {len(frames_per_batch)} values for {self.num_workers} workers."
+            )
+
+        self._frames_per_batch = frames_per_batch
+        total_frames_per_batch = (
+            sum(frames_per_batch)
+            if isinstance(frames_per_batch, Sequence)
+            else frames_per_batch
+        )
+
         self.set_truncated = set_truncated
         self.num_sub_threads = num_sub_threads
         self.num_threads = num_threads
@@ -2076,11 +2093,11 @@ def __init__(
         if total_frames is None or total_frames < 0:
             total_frames = float("inf")
         else:
-            remainder = total_frames % frames_per_batch
+            remainder = total_frames % total_frames_per_batch
             if remainder != 0 and RL_WARNINGS:
                 warnings.warn(
-                    f"total_frames ({total_frames}) is not exactly divisible by frames_per_batch ({frames_per_batch}). "
-                    f"This means {frames_per_batch - remainder} additional frames will be collected. "
+                    f"total_frames ({total_frames}) is not exactly divisible by frames_per_batch ({total_frames_per_batch}). "
+                    f"This means {total_frames_per_batch - remainder} additional frames will be collected. "
                     "To silence this message, set the environment variable RL_WARNINGS to False."
                 )
         self.total_frames = (
@@ -2091,7 +2108,8 @@ def __init__(
         self.max_frames_per_traj = (
             int(max_frames_per_traj) if max_frames_per_traj is not None else 0
         )
-        self.requested_frames_per_batch = int(frames_per_batch)
+
+        self.requested_frames_per_batch = total_frames_per_batch
         self.reset_when_done = reset_when_done
         if split_trajs is None:
             split_trajs = False
@@ -2221,8 +2239,7 @@ def _get_devices(
         )
         return storing_device, policy_device, env_device
 
-    @property
-    def frames_per_batch_worker(self):
+    def frames_per_batch_worker(self, worker_idx: int | None = None) -> int:
         raise NotImplementedError
 
     @property
@@ -2281,7 +2298,7 @@ def _run_processes(self) -> None:
                     "create_env_kwargs": env_fun_kwargs,
                     "policy": policy,
                     "max_frames_per_traj": self.max_frames_per_traj,
-                    "frames_per_batch": self.frames_per_batch_worker,
+                    "frames_per_batch": self.frames_per_batch_worker(worker_idx=i),
                     "reset_at_each_iter": self.reset_at_each_iter,
                     "policy_device": policy_device,
                     "storing_device": storing_device,
@@ -2773,8 +2790,9 @@ def update_policy_weights_(
             policy_or_weights=policy_or_weights, worker_ids=worker_ids, **kwargs
         )
 
-    @property
-    def frames_per_batch_worker(self):
+    def frames_per_batch_worker(self, worker_idx: int | None) -> int:
+        if worker_idx is not None and isinstance(self._frames_per_batch, Sequence):
+            return self._frames_per_batch[worker_idx]
         if self.requested_frames_per_batch % self.num_workers != 0 and RL_WARNINGS:
             warnings.warn(
                 f"frames_per_batch {self.requested_frames_per_batch} is not exactly divisible by the number of collector workers {self.num_workers},"
@@ -2855,9 +2873,9 @@ def iterator(self) -> Iterator[TensorDictBase]:
                 use_buffers = self._use_buffers
                 if self.replay_buffer is not None:
                     idx = new_data
-                    workers_frames[idx] = (
-                        workers_frames[idx] + self.frames_per_batch_worker
-                    )
+                    workers_frames[idx] = workers_frames[
+                        idx
+                    ] + self.frames_per_batch_worker(worker_idx=idx)
                     continue
                 elif j == 0 or not use_buffers:
                     try:
@@ -2903,7 +2921,12 @@ def iterator(self) -> Iterator[TensorDictBase]:
 
             if self.replay_buffer is not None:
                 yield
-                self._frames += self.frames_per_batch_worker * self.num_workers
+                self._frames += sum(
+                    [
+                        self.frames_per_batch_worker(worker_idx)
+                        for worker_idx in range(self.num_workers)
+                    ]
+                )
                 continue
 
             # we have to correct the traj_ids to make sure that they don't overlap
@@ -3156,8 +3179,7 @@ def update_policy_weights_(
             policy_or_weights=policy_or_weights, worker_ids=worker_ids, **kwargs
         )
 
-    @property
-    def frames_per_batch_worker(self):
+    def frames_per_batch_worker(self, worker_idx: int | None = None) -> int:
         return self.requested_frames_per_batch
 
     def _get_from_queue(self, timeout=None) -> tuple[int, int, TensorDictBase]:
@@ -3221,7 +3243,7 @@ def iterator(self) -> Iterator[TensorDictBase]:
                 if self.split_trajs:
                     out = split_trajectories(out, prefix="collector")
             else:
-                worker_frames = self.frames_per_batch_worker
+                worker_frames = self.frames_per_batch_worker()
             self._frames += worker_frames
             workers_frames[idx] = workers_frames[idx] + worker_frames
             if self.postprocs:

torchrl/data/llm/chat.py

@@ -11,7 +11,6 @@
 
 import torch
 
-
 from tensordict import lazy_stack, LazyStackedTensorDict, list_to_stack, TensorClass
 from tensordict.utils import _maybe_correct_neg_dim