torchrl 0.11.0__cp314-cp314t-win_amd64.whl

torchrl/collectors/distributed/utils.py

@@ -0,0 +1,160 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import subprocess
+import time
+
+from torchrl._utils import logger as torchrl_logger, VERBOSE
+from torchrl.collectors.distributed.default_configs import (
+    DEFAULT_SLURM_CONF,
+    DEFAULT_SLURM_CONF_MAIN,
+    TCP_PORT,
+)
+from torchrl.collectors.distributed.generic import _distributed_init_delayed
+from torchrl.collectors.distributed.rpc import _rpc_init_collection_node
+
+try:
+    import submitit
+
+    _has_submitit = True
+except ModuleNotFoundError as err:
+    _has_submitit = False
+    SUBMITIT_ERR = err
+
+
+class submitit_delayed_launcher:
+    """Delayed launcher for submitit.
+
+    In some cases, launched jobs cannot spawn other jobs on their own and this
+    can only be done at the jump-host level.
+
+    In these cases, the :func:`submitit_delayed_launcher` can be used to
+    pre-launch collector nodes that will wait for the main worker to provide
+    the launching instruction.
+
+    Args:
+        num_jobs (int): the number of collection jobs to be launched.
+        framework (str, optional): the framework to use. Can be either ``"distributed"``
+            or ``"rpc"``. ``"distributed"`` requires a :class:`~.DistributedDataCollector`
+            collector whereas ``"rpc"`` requires a :class:`RPCDataCollector`.
+            Defaults to ``"distributed"``.
+        backend (str, optional): torch.distributed backend in case ``framework``
+            points to ``"distributed"``. This value must match the one passed to
+            the collector, otherwise main and satellite nodes will fail to
+            reach the rendezvous and hang forever (ie no exception will be raised!)
+            Defaults to ``'gloo'``.
+        tcpport (int or str, optional): the TCP port to use.
+            Defaults to :obj:`torchrl.collectors.distributed.default_configs.TCP_PORT`
+        submitit_main_conf (dict, optional): the main node configuration to be passed to submitit.
+            Defaults to :obj:`torchrl.collectors.distributed.default_configs.DEFAULT_SLURM_CONF_MAIN`
+        submitit_collection_conf (dict, optional): the configuration to be passed to submitit.
+            Defaults to :obj:`torchrl.collectors.distributed.default_configs.DEFAULT_SLURM_CONF`
+
+    Examples:
+        >>> num_jobs=2
+        >>> @submitit_delayed_launcher(num_jobs=num_jobs)
+        ... def main():
+        ...     from torchrl.modules.utils.utils import RandomPolicyfrom torchrl.envs.libs.gym import GymEnv
+        ...     from torchrl.data import BoundedContinuous
+        ...     collector = DistributedDataCollector(
+        ...         [EnvCreator(lambda: GymEnv("Pendulum-v1"))] * num_jobs,
+        ...         policy=RandomPolicy(BoundedContinuous(-1, 1, shape=(1,))),
+        ...         launcher="submitit_delayed",
+        ...     )
+        ...     for data in collector:
+        ...         print(data)
+        ...
+        >>> if __name__ == "__main__":
+        ...     main()
+        ...
+    """
+
+    _VERBOSE = VERBOSE  # for debugging
+
+    def __init__(
+        self,
+        num_jobs,
+        framework="distributed",
+        backend="gloo",
+        tcpport=TCP_PORT,
+        submitit_main_conf: dict = DEFAULT_SLURM_CONF_MAIN,
+        submitit_collection_conf: dict = DEFAULT_SLURM_CONF,
+    ):
+        self.num_jobs = num_jobs
+        self.backend = backend
+        self.framework = framework
+        self.submitit_collection_conf = submitit_collection_conf
+        self.submitit_main_conf = submitit_main_conf
+        self.tcpport = tcpport
+
+    def __call__(self, main_func):
+        def exec_fun():
+            if not _has_submitit:
+                raise ModuleNotFoundError(
+                    "Failed to import submitit. Check installation of the library."
+                ) from SUBMITIT_ERR
+            # submit main
+            executor = submitit.AutoExecutor(folder="log_test")
+            executor.update_parameters(**self.submitit_main_conf)
+            main_job = executor.submit(main_func)
+            # listen to output file looking for IP address
+            torchrl_logger.debug(f"job id: {main_job.job_id}")
+            time.sleep(2.0)
+            node = None
+            while not node:
+                cmd = f"squeue -j {main_job.job_id} -o %N | tail -1"
+                node = subprocess.check_output(cmd, shell=True, text=True).strip()
+                try:
+                    node = int(node)
+                except ValueError:
+                    time.sleep(0.5)
+                    continue
+            torchrl_logger.debug(f"node: {node}")
+            # by default, sinfo will truncate the node name at char 20, we increase this to 200
+            cmd = f"sinfo -n {node} -O nodeaddr:200 | tail -1"
+            rank0_ip = subprocess.check_output(cmd, shell=True, text=True).strip()
+            torchrl_logger.debug(f"IP: {rank0_ip}")
+            world_size = self.num_jobs + 1
+
+            # submit jobs
+            executor = submitit.AutoExecutor(folder="log_test")
+            executor.update_parameters(**self.submitit_collection_conf)
+            jobs = []
+            if self.framework == "rpc":
+                from .rpc import DEFAULT_TENSORPIPE_OPTIONS
+
+                tensorpipe_options = DEFAULT_TENSORPIPE_OPTIONS
+            for i in range(self.num_jobs):
+                rank = i + 1
+                if self.framework == "distributed":
+                    job = executor.submit(
+                        _distributed_init_delayed,
+                        rank,
+                        self.backend,
+                        rank0_ip,
+                        self.tcpport,
+                        world_size,
+                        self._VERBOSE,
+                    )
+                elif self.framework == "rpc":
+                    job = executor.submit(
+                        _rpc_init_collection_node,
+                        rank,
+                        rank0_ip,
+                        self.tcpport,
+                        world_size,
+                        None,
+                        tensorpipe_options,
+                    )
+                else:
+                    raise NotImplementedError(f"Unknown framework {self.framework}.")
+                jobs.append(job)
+            for job in jobs:
+                job.result()
+            main_job.result()
+
+        return exec_fun

torchrl/collectors/llm/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .base import LLMCollector
+from .ray_collector import RayLLMCollector
+from .weight_update import vLLMUpdater, vLLMUpdaterV2
+
+__all__ = ["vLLMUpdater", "vLLMUpdaterV2", "LLMCollector", "RayLLMCollector"]

torchrl/collectors/llm/base.py

@@ -0,0 +1,494 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+from collections import deque
+from collections.abc import Callable
+from typing import Any
+
+import torch
+
+from tensordict import lazy_stack, TensorDictBase
+
+from torchrl._utils import as_remote, logger as torchrl_logger
+
+from torchrl.collectors._single import Collector
+from torchrl.collectors.llm.utils import _QueueAsRB
+from torchrl.collectors.weight_update import WeightUpdaterBase
+from torchrl.data.replay_buffers.replay_buffers import ReplayBuffer
+from torchrl.envs import AsyncEnvPool
+from torchrl.envs.common import EnvBase
+from torchrl.envs.llm.transforms.policy_version import PolicyVersion
+
+
+class LLMCollector(Collector):
+    """A simplified version of Collector for LLM inference.
+
+    Args:
+        env (EnvBase or EnvBase constructor): the environment to be used for data collection.
+
+    Keyword Args:
+        policy (Callable[[TensorDictBase], TensorDictBase]): the policy to be used for data collection.
+        policy_factory (Callable[[], Callable], optional): a callable that returns
+            a policy instance. This is exclusive with the `policy` argument.
+
+            .. note:: `policy_factory` comes in handy whenever the policy cannot be serialized.
+
+        dialog_turns_per_batch (int, optional): A keyword-only argument representing the total
+            number of elements in a batch. It is always required except when `yield_completed_trajectories=True`.
+        total_dialog_turns (int): A keyword-only argument representing the total
+            number of steps returned by the collector during its lifespan. -1 is never ending (until shutdown).
+            Defaults to -1.
+        yield_completed_trajectories (bool, optional): whether to yield batches of rollouts with a given number of steps
+            (`yield_completed_trajectories=False`, default) or single, completed trajectories
+            (`yield_completed_trajectories=True`).
+            Defaults to `False` unless `yield_only_last_steps=True`, where it cannot be `False`.
+
+            .. warning:: If the `done` state of the environment is not properly set, this may lead to a collector
+                that never leads any data.
+
+        yield_only_last_steps (bool, optional): whether to yield every step of a trajectory, or only the
+            last (done) steps.
+            If `True`, a single trajectory is yielded (or written in the buffer) at a time.
+
+            .. warning:: If the `done` state of the environment is not properly set, this may lead to a collector
+                that never leads any data.
+
+        postproc (Callable, optional): A post-processing transform, such as
+            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 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.
+        flatten_data (bool, optional): if ``True``, the collector will flatten the collected data
+            before returning it. In practice, this means that if an environment of batch-size `(B,)` is used
+            and run for `T` steps, `flatten_data=True` will present data of shape `(B*T,)`, whereas
+            `flatten_data=False` will not present data of shape `(B, T)`.
+            Defaults to `True` when `replay_buffer` is provided, `False` otherwise.
+        weight_updater (WeightUpdaterBase or constructor, optional): An instance of :class:`~torchrl.collectors.WeightUpdaterBase`
+            or its subclass, responsible for updating the policy weights on remote inference workers.
+            This is typically not used in :class:`~torchrl.collectors.Collector` as it operates in a single-process environment.
+            Consider using a constructor if the updater needs to be serialized.
+        track_policy_version (bool or PolicyVersion, optional): if ``True``, the collector will track the version of the policy.
+            This will be mediated by the :class:`~torchrl.envs.llm.transforms.policy_version.PolicyVersion` transform, which will be added to the environment.
+            Alternatively, a :class:`~torchrl.envs.llm.transforms.policy_version.PolicyVersion` instance can be passed, which will be used to track
+            the policy version.
+            Defaults to `False`.
+        verbose (bool, optional): if ``True``, the collector will print progress information.
+            Defaults to `False`.
+
+    Examples:
+        >>> import vllm
+        >>> from torchrl.modules import vLLMWrapper
+        >>> from torchrl.testing.mocking_classes import DummyStrDataLoader
+        >>> from torchrl.envs import LLMEnv
+        >>> llm_model = vllm.LLM("gpt2")
+        >>> tokenizer = llm_model.get_tokenizer()
+        >>> tokenizer.pad_token = tokenizer.eos_token
+        >>> policy = vLLMWrapper(llm_model)
+        >>> dataloader = DummyStrDataLoader(1)
+        >>> env = LLMEnv.from_dataloader(
+        ...    dataloader=dataloader,
+        ...    tokenizer=tokenizer,
+        ...    from_text=True,
+        ...    batch_size=1,
+        ...    group_repeats=True,
+        ... )
+        >>> collector = LLMCollector(
+        ...    env=env,
+        ...    policy_factory=lambda: policy,
+        ...    dialog_turns_per_batch=env.batch_size[0],
+        ...    total_dialog_turns=3,
+        ... )
+        >>> for i, data in enumerate(collector):
+        ...     if i == 2:
+        ...         print(data)
+        ...         break
+        LazyStackedTensorDict(
+        fields={
+            attention_mask: Tensor(shape=torch.Size([1, 1, 22]), device=cpu, dtype=torch.int64, is_shared=False),
+            collector: LazyStackedTensorDict(
+                fields={
+                    traj_ids: Tensor(shape=torch.Size([1, 1]), device=cpu, dtype=torch.int64, is_shared=False)},
+                exclusive_fields={
+                },
+                batch_size=torch.Size([1, 1]),
+                device=None,
+                is_shared=False,
+                stack_dim=1),
+            done: Tensor(shape=torch.Size([1, 1, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+            terminated: Tensor(shape=torch.Size([1, 1, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+            text: NonTensorStack(
+                [['plsgqejeyd']],
+                batch_size=torch.Size([1, 1]),
+                device=None),
+            text_response: NonTensorStack(
+                [['ec.n.n.n.tjbjz3perwhz']],
+                batch_size=torch.Size([1, 1]),
+                device=None),
+            tokens: Tensor(shape=torch.Size([1, 1, 22]), device=cpu, dtype=torch.int64, is_shared=False),
+            tokens_response: Tensor(shape=torch.Size([1, 1, 16]), device=cpu, dtype=torch.int64, is_shared=False)},
+        exclusive_fields={
+        },
+        batch_size=torch.Size([1, 1]),
+        device=None,
+        is_shared=False,
+        stack_dim=1)
+        >>> del collector
+
+    """
+
+    def __init__(
+        self,
+        env: EnvBase | Callable[[], EnvBase],
+        *,
+        policy: Callable[[TensorDictBase], TensorDictBase] | None = None,
+        policy_factory: Callable[[], Callable[[TensorDictBase], TensorDictBase]]
+        | None = None,
+        dialog_turns_per_batch: int | None = None,
+        yield_only_last_steps: bool | None = None,
+        yield_completed_trajectories: bool | None = None,
+        postproc: Callable[[TensorDictBase], TensorDictBase] | None = None,
+        total_dialog_turns: int = -1,
+        async_envs: bool | None = None,
+        replay_buffer: ReplayBuffer | None = None,
+        reset_at_each_iter: bool = False,
+        flatten_data: bool | None = None,
+        weight_updater: WeightUpdaterBase
+        | Callable[[], WeightUpdaterBase]
+        | None = None,
+        queue: Any | None = None,
+        track_policy_version: bool | PolicyVersion = False,
+        verbose: bool = False,
+    ):
+        if queue is not None and replay_buffer is not None:
+            raise RuntimeError(
+                "Handling both a buffer and a queue is not possible at the moment."
+            )
+        elif queue is not None:
+            # disguise the queue as a replay buffer
+            replay_buffer = _QueueAsRB(queue)
+        if dialog_turns_per_batch is None and yield_completed_trajectories:
+            dialog_turns_per_batch = 1
+        super().__init__(
+            create_env_fn=env,
+            policy=policy,
+            policy_factory=policy_factory,
+            frames_per_batch=dialog_turns_per_batch,
+            replay_buffer=replay_buffer,
+            total_frames=total_dialog_turns,
+            weight_updater=weight_updater,
+            reset_at_each_iter=reset_at_each_iter,
+            trust_policy=True,
+            use_buffers=False,
+            no_cuda_sync=True,
+            extend_buffer=True,
+            postproc=postproc,
+        )
+        if hasattr(self.policy, "register_collector"):
+            self.policy.register_collector(self)
+
+        if yield_only_last_steps is None:
+            yield_only_last_steps = False
+
+        if yield_completed_trajectories is None:
+            yield_completed_trajectories = yield_only_last_steps
+        elif yield_only_last_steps and not yield_completed_trajectories:
+            raise TypeError(
+                "yield_only_last_steps=True requires yield_completed_trajectories=True (or None)"
+            )
+
+        if yield_only_last_steps:
+            if flatten_data is not None:
+                raise TypeError(
+                    "`yield_only_last_steps` cannot be `True` when `flatten_data` is passed."
+                )
+            if self.reset_at_each_iter:
+                raise TypeError(
+                    "`yield_only_last_steps` cannot be `True` when `reset_at_each_iter=True`."
+                )
+        if flatten_data is None:
+            flatten_data = replay_buffer is not None
+        self.flatten_data = flatten_data
+        self.yield_completed_trajectories = yield_completed_trajectories
+        self.yield_only_last_steps = yield_only_last_steps
+        self.verbose = verbose
+        self._shuttle = None  # Initialize shuttle for rollout
+        if self.yield_completed_trajectories:
+            # For async envs, we route by env_id so we only care about batch_size[0].
+            # For non-async envs, we need exactly one batch dimension.
+            if not isinstance(self.env, AsyncEnvPool) and len(self.env.batch_size) != 1:
+                raise ValueError(
+                    "`yield_completed_trajectories` only works with envs that have a single batch dimension. Got "
+                    f"env.batch_size={self.env.batch_size}."
+                )
+            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."
+            )
+        self.policy_version_tracker = track_policy_version
+        if isinstance(track_policy_version, bool) and track_policy_version:
+            if isinstance(self.env, AsyncEnvPool):
+                raise RuntimeError(
+                    "AsyncEnvPool is not supported for policy version tracking. Please add the PolicyVersion transform to the environment manually, "
+                    "and pass that transform to the collector."
+                )
+            self.policy_version_tracker = PolicyVersion()
+            self.env = self.env.append_transform(self.policy_version_tracker)  # type: ignore
+        elif isinstance(track_policy_version, PolicyVersion):
+            self.policy_version_tracker = track_policy_version
+            self.env = self.env.append_transform(self.policy_version_tracker)  # type: ignore
+        else:
+            self.policy_version_tracker = None
+
+    def set_postproc(self, postproc: Callable[[TensorDictBase], TensorDictBase]):
+        if self.postproc is not None:
+            raise RuntimeError("Postproc already set")
+        self.postproc = postproc
+
+    def increment_version(self):
+        """Increment the policy version."""
+        if self.policy_version_tracker is not None:
+            if not isinstance(self.policy_version_tracker, PolicyVersion):
+                raise RuntimeError(
+                    "Policy version tracker is not a PolicyVersion instance. Please pass a PolicyVersion instance to the collector."
+                )
+            self.policy_version_tracker.increment_version()
+
+    @property
+    def policy_version(self) -> str | int | None:
+        """The current policy version."""
+        if not isinstance(self.policy_version_tracker, PolicyVersion):
+            return None
+        return self.policy_version_tracker.version
+
+    def get_policy_version(self) -> str | int | None:
+        """Get the current policy version.
+
+        This method exists to support remote calls in Ray actors, since properties
+        cannot be accessed directly through Ray's RPC mechanism.
+
+        Returns:
+            The current version number (int) or UUID (str), or None if version tracking is disabled.
+        """
+        return self.policy_version
+
+    @property
+    def total_dialog_turns(self):
+        return self.total_frames
+
+    @property
+    def dialog_turns_per_batch(self) -> int:
+        """Alias to `frames_per_batch`."""
+        return self.requested_frames_per_batch
+
+    @property
+    def rollout(self) -> Callable[[], TensorDictBase]:
+        if self.yield_completed_trajectories:
+            if self.async_envs:
+                return self._rollout_yield_trajs_async
+            else:
+                return self._rollout_yield_trajs
+        else:
+            return self._rollout_all
+
+    def _rollout_all(self) -> TensorDictBase:  # A simplified version of rollout
+        if self.reset_at_each_iter or self._shuttle is None:
+            self._shuttle = self.env.reset()
+
+        trajectory = []
+        collected_steps = 0
+        policy_input = self._shuttle
+        while collected_steps < self.dialog_turns_per_batch:
+            if self.verbose:
+                torchrl_logger.debug(
+                    f"LLMCollector: Collected {collected_steps} steps over {self.dialog_turns_per_batch} requested."
+                )
+            env_input = self.policy(policy_input)
+            env_output, env_next_output = self.env.step_and_maybe_reset(env_input)
+
+            # carry over collector data without messing up devices
+            collector_data = env_output.get("collector", default=None)
+            if collector_data is not None:
+                env_next_output.set("collector", collector_data.copy())
+                self._update_traj_ids(env_output)
+            trajectory.append(env_output.clone())
+            collected_steps += env_output.numel()
+            policy_input = self._shuttle = env_next_output
+        trajectory = lazy_stack(trajectory, -1)
+        if self.flatten_data:
+            return trajectory.view(-1)
+        return trajectory
+
+    _result_numel = 0
+
+    def _rollout_yield_trajs(self) -> TensorDictBase:  # A simplified version of rollout
+        if self._shuttle is None:
+            self._shuttle = self.env.reset()
+        next_output = self._shuttle
+
+        collected_steps = 0
+        dones = torch.zeros(self.env.batch_size, dtype=torch.bool)
+        while True:
+            if self._result_numel >= self.dialog_turns_per_batch:
+                break
+            elif self.verbose:
+                torchrl_logger.debug(
+                    f"LLMCollector: Collected {collected_steps} steps with {self._result_numel} elements in the resulting batch, over {self.dialog_turns_per_batch} requested."
+                )
+            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 = cur_output.get("collector", default=None)
+            if collector_data is not None:
+                self._update_traj_ids(cur_output)
+                next_output.set("collector", collector_data.copy())
+
+            # 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():
+                for idx in dones.nonzero(as_tuple=True)[0].tolist():
+                    if not self.yield_only_last_steps:
+                        _result = lazy_stack(self._yield_queues[idx], -1)
+                        self._trajectory_queue.append(_result)
+                    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
+                        _result = lazy_stack([self._yield_queues[idx][-1]])
+                        self._trajectory_queue.append(_result)
+                    self._result_numel += _result.numel()
+                    self._yield_queues[idx].clear()
+        result = [self._trajectory_queue.popleft()]
+        elt = result[0].numel()
+        self._result_numel -= result[0].numel()
+        while elt < self.dialog_turns_per_batch:
+            result.append(self._trajectory_queue.popleft())
+            elt += result[-1].numel()
+            self._result_numel -= result[-1].numel()
+        result = torch.cat(result, -1)
+        if self.verbose:
+            torchrl_logger.debug(
+                f"LLMCollector: Yielding completed trajectory with shape {result.shape}."
+            )
+        return result
+
+    started = False
+
+    def _rollout_yield_trajs_async(
+        self,
+    ) -> TensorDictBase:  # A simplified version of rollout
+        if not self.started:
+            if self._shuttle is None:
+                self._shuttle = self.env.reset()
+            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
+        # Use only the first dimension (num_envs) for done tracking, since we route by env_id
+        dones = torch.zeros(self.env.batch_size[0], 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 - flatten to handle multi-dimensional batch sizes
+            # (e.g., AsyncEnvPool with batch_size=[4, 1] gives [[0], [1], [2], [3]])
+            env_ids_raw = cur_output.get(self.env._env_idx_key).tolist()
+            # Flatten nested lists to get scalar env indices
+            env_ids = []
+            for eid in env_ids_raw:
+                while isinstance(eid, list) and len(eid) == 1:
+                    eid = eid[0]
+                env_ids.append(eid)
+
+            # carry over collector data without messing up devices
+            collector_data = cur_output.get("collector", default=None)
+            if collector_data is not None:
+                self._update_traj_ids(cur_output)
+                next_output.set("collector", collector_data.copy())
+
+            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(as_tuple=True)[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()
+        # Flatten the result - AsyncEnvPool child envs with batch_size=(1,) produce
+        # trajectories with shape [1, T] but we want [T] for consistency
+        result = result.view(-1)
+        if self.verbose:
+            torchrl_logger.debug(
+                f"LLMCollector: Yielding completed trajectory with shape {result.shape}."
+            )
+        return result
+
+    as_remote = as_remote
+
+    def get_policy_model(self):
+        """Get the policy model.
+
+        This method is used by RayLLMCollector to get the remote LLM instance
+        for weight updates.
+
+        Returns:
+            The policy model instance
+        """
+        return self.policy.model
+
+    def is_initialized(self) -> bool:
+        """Check if the collector is initialized and ready.
+
+        Returns:
+            bool: True if the collector is initialized and ready to collect data.
+        """
+        # The collector is initialized if it has a valid environment and policy
+        return hasattr(self, "_env") and hasattr(self, "_policy")
+
+    def set_weight_updater(self, weight_updater: WeightUpdaterBase):
+        self.weight_updater = weight_updater
+        return True