[Feature] Collectors for async envs

ghstack-source-id: 764c21d
Pull Request resolved: #2893

Author: Vincent Moens

test/test_collector.py

@@ -67,6 +67,7 @@
 from torchrl.data.llm.dataset import _has_transformers
 from torchrl.data.utils import CloudpickleWrapper
 from torchrl.envs import (
+    AsyncEnvPool,
     EnvBase,
     EnvCreator,
     InitTracker,
@@ -3737,10 +3738,12 @@ async def test_llm_collector_start(self, vllm_instance):
     def test_llm_collector_completed(
         self, vllm_instance_opt, rb, yield_only_last_steps
     ):
+        torch.manual_seed(0)
         policy = vLLMWrapper(vllm_instance_opt)
         tokenizer = vllm_instance_opt.get_tokenizer()
         bsz = 4
         total_steps = 20
+        max_steps = 20
         dataloader = DummyStrDataLoader(bsz)
 
         env = LLMEnv.from_dataloader(
@@ -3751,7 +3754,7 @@ def test_llm_collector_completed(
             eos_token_id=tokenizer.eos_token_id,
         )
         # To make sure the env breaks at some point
-        env = env.append_transform(StepCounter(max_steps=100))
+        env = env.append_transform(StepCounter(max_steps=max_steps))
 
         if rb:
             rb = ReplayBuffer(storage=LazyStackStorage(max_size=total_steps * 2))
@@ -3774,11 +3777,27 @@ def test_llm_collector_completed(
         for data in collector:
             if rb is None:
                 assert data.ndim == 1
-                assert (data["next", "step_count"] < 99).all()
+                # assert (data["next", "step_count"] < max_steps-1).all()
                 cur_total_steps += data.numel()
                 for i in range(data.numel()):
-                    # Check that there are more chars in the next step
-                    assert len(data["text"][i]) < len(data["next", "text"][i])
+                    if data[i]["next", "step_count"] == max_steps:
+                        continue
+                    if data[i]["text_response"]:
+                        # Check that there are more chars in the next step
+                        assert len(data["text"][i]) < len(data["next", "text"][i]), (
+                            i,
+                            data[i]["next", "step_count"],
+                            data[i]["next", "done"],
+                            data[i]["text_response"],
+                        )
+                    else:
+                        assert len(data["text"][i]) == len(data["next", "text"][i]), (
+                            i,
+                            data[i]["next", "step_count"],
+                            data[i]["next", "done"],
+                            data[i]["text_response"],
+                        )
+
                 if yield_only_last_steps:
                     assert data.shape == (1,)
                 else:
@@ -3787,8 +3806,137 @@ def test_llm_collector_completed(
                 assert data is None
                 sample = rb.sample(5)
                 for i in range(sample.numel()):
-                    # Check that there are more chars in the next step
-                    assert len(sample["text"][i]) < len(sample["next", "text"][i])
+                    if sample[i]["next", "step_count"] == max_steps:
+                        continue
+                    if sample[i]["text_response"]:
+                        # Check that there are more chars in the next step
+                        assert len(sample["text"][i]) < len(
+                            sample["next", "text"][i]
+                        ), (
+                            i,
+                            sample[i]["next", "step_count"],
+                            sample[i]["next", "done"],
+                            sample[i]["text_response"],
+                        )
+                    else:
+                        assert len(sample["text"][i]) == len(
+                            sample["next", "text"][i]
+                        ), (
+                            i,
+                            sample[i]["next", "step_count"],
+                            sample[i]["next", "done"],
+                            sample[i]["text_response"],
+                        )
+
+                assert sample.ndim == 1
+                assert sample.shape == (5,)
+                assert (sample["next", "step_count"] < 99).all()
+                cur_total_steps += 1
+            assert collector._frames >= cur_total_steps
+        if rb is None and not yield_only_last_steps:
+            assert has_found_one_with_more_steps
+        assert collector._frames >= total_steps
+
+    @pytest.mark.slow
+    @pytest.mark.parametrize("rb", [False, True])
+    @pytest.mark.parametrize("yield_only_last_steps", [False, True])
+    def test_llm_collector_completed_async(
+        self, vllm_instance_opt, rb, yield_only_last_steps
+    ):
+        torch.manual_seed(0)
+        policy = vLLMWrapper(vllm_instance_opt)
+        tokenizer = vllm_instance_opt.get_tokenizer()
+        bsz = 4
+        total_steps = 20
+        max_steps = 20
+        dataloader = DummyStrDataLoader(bsz)
+
+        def env_maker():
+            env = LLMEnv.from_dataloader(
+                dataloader=dataloader,
+                str2str=True,
+                batch_size=(),
+                group_repeats=True,
+                eos_token_id=tokenizer.eos_token_id,
+            )
+            # To make sure the env breaks at some point
+            env = env.append_transform(StepCounter(max_steps=max_steps))
+            return env
+
+        env = AsyncEnvPool([env_maker] * bsz, backend="threading", stack="lazy")
+
+        if rb:
+            rb = ReplayBuffer(storage=LazyStackStorage(max_size=total_steps * 2))
+        else:
+            rb = None
+        collector = LLMCollector(
+            env=env,
+            policy_factory=lambda: policy,
+            steps_per_batch=env.batch_size[0],
+            replay_buffer=rb,
+            total_steps=total_steps,
+            yield_completed_trajectories=True,
+            yield_only_last_steps=yield_only_last_steps,
+        )
+        assert collector.yield_completed_trajectories
+        assert collector.yield_only_last_steps is yield_only_last_steps
+
+        cur_total_steps = 0
+        has_found_one_with_more_steps = False
+        for data in collector:
+            if rb is None:
+                assert data.ndim == 1
+                # assert (data["next", "step_count"] < max_steps-1).all()
+                cur_total_steps += data.numel()
+                for i in range(data.numel()):
+                    if data[i]["next", "step_count"] == max_steps:
+                        continue
+                    if data[i]["text_response"]:
+                        # Check that there are more chars in the next step
+                        assert len(data["text"][i]) < len(data["next", "text"][i]), (
+                            i,
+                            data[i]["next", "step_count"],
+                            data[i]["next", "done"],
+                            data[i]["text_response"],
+                        )
+                    else:
+                        assert len(data["text"][i]) == len(data["next", "text"][i]), (
+                            i,
+                            data[i]["next", "step_count"],
+                            data[i]["next", "done"],
+                            data[i]["text_response"],
+                        )
+
+                if yield_only_last_steps:
+                    assert data.shape == (1,)
+                else:
+                    has_found_one_with_more_steps |= data.numel() > 1
+            else:
+                assert data is None
+                sample = rb.sample(5)
+                for i in range(sample.numel()):
+                    if sample[i]["next", "step_count"] == max_steps:
+                        continue
+                    if sample[i]["text_response"]:
+                        # Check that there are more chars in the next step
+                        assert len(sample["text"][i]) < len(
+                            sample["next", "text"][i]
+                        ), (
+                            i,
+                            sample[i]["next", "step_count"],
+                            sample[i]["next", "done"],
+                            sample[i]["text_response"],
+                        )
+                    else:
+                        assert len(sample["text"][i]) == len(
+                            sample["next", "text"][i]
+                        ), (
+                            i,
+                            sample[i]["next", "step_count"],
+                            sample[i]["next", "done"],
+                            sample[i]["text_response"],
+                        )
+
                 assert sample.ndim == 1
                 assert sample.shape == (5,)
                 assert (sample["next", "step_count"] < 99).all()

torchrl/collectors/llm.py

@@ -17,6 +17,7 @@
     WeightUpdateSenderBase,
 )
 from torchrl.data.replay_buffers.replay_buffers import ReplayBuffer
+from torchrl.envs import AsyncEnvPool
 from torchrl.envs.common import EnvBase
 
 
@@ -57,7 +58,8 @@ class LLMCollector(SyncDataCollector):
             a :class:`~torchrl.envs.Transform` or a :class:`~torchrl.data.postprocs.MultiStep`
             instance.
             Defaults to ``None``.
-        async_envs (bool, optional): if ``True``, the environment will be run synchronously.
+        async_envs (bool, optional): if ``True``, the environment will be run asynchronously. Defaults to `True` if the
+            environment is a :class:`~torchrl.envs.AsyncEnvPool` instance.
         replay_buffer (ReplayBuffer, optional): if provided, the collector will not yield tensordicts
             but populate the buffer instead. Defaults to ``None``.
         reset_at_each_iter (bool, optional): if ``True``, the environment will be reset at each iteration.
@@ -149,7 +151,7 @@ def __init__(
         yield_completed_trajectories: bool | None = None,
         postproc: Callable[[TensorDictBase], TensorDictBase] | None = None,
         total_steps: int = -1,
-        async_envs: bool = False,
+        async_envs: bool | None = None,
         replay_buffer: ReplayBuffer | None = None,
         reset_at_each_iter: bool = False,
         flatten_data: bool | None = None,
@@ -160,8 +162,6 @@ def __init__(
         | Callable[[], WeightUpdateSenderBase]
         | None = None,
     ):
-        if async_envs:
-            raise NotImplementedError
         super().__init__(
             create_env_fn=env,
             policy=policy,
@@ -209,6 +209,13 @@ def __init__(
                 )
             self._yield_queues = [deque() for _ in range(self.env.batch_size[0])]
             self._trajectory_queue = deque()
+        self.async_envs = bool(async_envs) | isinstance(self.env, AsyncEnvPool)
+        if self.async_envs and not isinstance(self.env, AsyncEnvPool):
+            # This basically means that `async_envs` is automatically set and passing is it useless as of today,
+            #  except for the following error.
+            raise RuntimeError(
+                "async_envs requires the environment to be an AsyncEnvPool instance."
+            )
 
     @property
     def steps_per_batch(self) -> int:
@@ -218,7 +225,10 @@ def steps_per_batch(self) -> int:
     @property
     def rollout(self) -> Callable[[], TensorDictBase]:
         if self.yield_completed_trajectories:
-            return self._rollout_yield_trajs
+            if self.async_envs:
+                return self._rollout_yield_trajs_async
+            else:
+                return self._rollout_yield_trajs
         else:
             return self._rollout_all
 
@@ -250,27 +260,33 @@ def _rollout_all(self) -> TensorDictBase:  # A simplified version of rollout
 
     def _rollout_yield_trajs(self) -> TensorDictBase:  # A simplified version of rollout
         if self._shuttle is None:
-            data = self.env.reset()
+            raise RuntimeError("Data shuttle not found")
+            # next_output = self.env.reset()
         else:
-            data = self._shuttle
+            next_output = self._shuttle
 
         collected_steps = 0
         dones = torch.zeros(self.env.batch_size, dtype=torch.bool)
         while True:
             if self._trajectory_queue:
                 break
-            policy_input = data
-            env_input = self.policy(policy_input)
-            env_output, env_next_output = self.env.step_and_maybe_reset(env_input)
+            env_input = self.policy(next_output)
+            cur_output, next_output = self.env.step_and_maybe_reset(env_input)
+            # for i in range(cur_output.numel()):
+            #     print(len(cur_output[i]["text"]) < len(cur_output[i]["next", "text"]))
 
             # carry over collector data without messing up devices
-            collector_data = env_output.get("collector").copy()
-            env_next_output.set("collector", collector_data)
-            self._shuttle = env_next_output
-            self._update_traj_ids(env_output)
-            data = env_output
-            collected_steps += data.numel()
-            for i, (_data, queue) in enumerate(zip(data.unbind(0), self._yield_queues)):
+            self._update_traj_ids(cur_output)
+
+            collector_data = cur_output.get("collector").copy()
+            next_output.set("collector", collector_data)
+
+            # if the loop is interrupted
+            self._shuttle = next_output
+            collected_steps += next_output.numel()
+            for i, (_data, queue) in enumerate(
+                zip(cur_output.unbind(0), self._yield_queues)
+            ):
                 queue.append(_data)
                 dones[i] = _data["next", "done"].any()
             if dones.any():
@@ -290,3 +306,61 @@ def _rollout_yield_trajs(self) -> TensorDictBase:  # A simplified version of rol
 
         result = self._trajectory_queue.popleft()
         return result
+
+    started = False
+
+    def _rollout_yield_trajs_async(
+        self,
+    ) -> TensorDictBase:  # A simplified version of rollout
+        if not self.started:
+            next_output = self._shuttle
+            env_input = self.policy(next_output)
+            self.env.async_step_and_maybe_reset_send(env_input)
+        self.started = True
+
+        collected_steps = 0
+        dones = torch.zeros(self.env.batch_size, dtype=torch.bool)
+        while True:
+            if self._trajectory_queue:
+                break
+
+            cur_output, next_output = self.env.async_step_and_maybe_reset_recv()
+
+            # Get the env ids
+            env_ids = cur_output.get(self.env._env_idx_key).tolist()
+
+            # carry over collector data without messing up devices
+            self._update_traj_ids(cur_output)
+
+            collector_data = cur_output.get("collector").copy()
+            next_output.set("collector", collector_data)
+
+            collected_steps += next_output.numel()
+            dones.fill_(False)
+            for i, _data in zip(env_ids, cur_output.unbind(0)):
+                queue = self._yield_queues[i]
+                queue.append(_data)
+                dones[i] = _data["next", "done"].any()
+            if dones.any():
+                for idx in dones.nonzero()[0].tolist():
+                    if not self.yield_only_last_steps:
+                        self._trajectory_queue.append(
+                            lazy_stack(self._yield_queues[idx], -1)
+                        )
+                    else:
+                        # FIXME: We need to increment the step count here because iterator() won't
+                        #  see the extra steps
+                        # We use lazy-stack because unsqueeze doesn't nest the strings in lists
+                        self._trajectory_queue.append(
+                            lazy_stack([self._yield_queues[idx][-1]])
+                        )
+                    self._yield_queues[idx].clear()
+
+            # Launch the next batch:
+            # FIXME: Add a condition RE number of frames here
+            if True:
+                env_input = self.policy(next_output)
+                self.env.async_step_and_maybe_reset_send(env_input)
+
+        result = self._trajectory_queue.popleft()
+        return result

torchrl/envs/async_envs.py

@@ -219,7 +219,7 @@ def __init__(
 
         output_spec, input_spec = self._setup()
         input_spec["full_state_spec"].set(
-            self._env_idx_key, NonTensor(example_data=0, shape=self.batch_size)
+            self._env_idx_key, NonTensor(example_data=0, shape=input_spec.shape)
         )
         self.__dict__["_output_spec"] = output_spec
         self.__dict__["_input_spec"] = input_spec