torchrl 0.11.0__cp314-cp314-manylinux_2_28_aarch64.whl

torchrl/envs/llm/transforms/dataloading.py

@@ -0,0 +1,859 @@
+# 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 warnings
+from collections import deque
+from collections.abc import Callable, Iterable, Mapping
+
+from typing import Any, Literal, TypeVar
+
+import torch
+from tensordict import is_tensor_collection, lazy_stack, TensorDict, TensorDictBase
+
+from torchrl.data.tensor_specs import Composite, DEVICE_TYPING, TensorSpec
+from torchrl.envs.common import EnvBase
+from torchrl.envs.transforms import TensorDictPrimer, Transform
+
+# Import ray service components
+from torchrl.envs.transforms.ray_service import (
+    _map_input_output_device,
+    _RayServiceMetaClass,
+    RayTransform,
+)
+from torchrl.envs.utils import make_composite_from_td
+
+T = TypeVar("T")
+
+
+def as_nested_tensor(list_of_tensordicts: list[TensorDictBase]) -> TensorDictBase:
+    """Stacks a list of tensordicts into a single tensordict with nested tensors.
+
+    Args:
+        list_of_tensordicts (list[TensorDictBase]): A list of tensordicts to stack.
+
+    Returns:
+        TensorDictBase: A tensordict with nested tensors.
+
+    """
+
+    def _as_nested_tensor(*list_of_tensors):
+        return torch.nested.as_nested_tensor(list_of_tensors, layout=torch.jagged)
+
+    batch_size = list(list_of_tensordicts[0].shape)
+    batch_size.insert(0, len(list_of_tensordicts))
+    result: TensorDictBase = list_of_tensordicts[0].apply(  # type: ignore[assignment]
+        _as_nested_tensor, *list_of_tensordicts[1:], batch_size=batch_size
+    )
+    return result
+
+
+def as_padded_tensor(
+    list_of_tensordicts: list[TensorDictBase], dim=0, stack_dim: int = 0
+) -> TensorDictBase:
+    """Stacks a list of tensordicts into a single tensordict with padded tensors.
+
+    Args:
+        list_of_tensordicts (list[[TensorDictBase]]): A list of tensordicts to stack.
+        dim (int, optional): The dimension along which to pad. Defaults to 0.
+        stack_dim (int, optional): The dimension along which to stack. Defaults to 0.
+
+    Returns:
+        TensorDictBase: A tensordict with padded tensors.
+    """
+
+    def _stack_tensors(*list_of_tensors):
+        if dim < 0:
+            raise ValueError("dim must be >= 0")
+        max_length = max([t.size(dim) for t in list_of_tensors])
+
+        def pad_tensor(tensor):
+            padding_length = max_length - tensor.size(dim)
+            shape = [
+                s if i != dim else padding_length for i, s in enumerate(tensor.shape)
+            ]
+            return torch.cat((tensor.new_zeros(shape), tensor), dim=dim)
+
+        return torch.stack([pad_tensor(t) for t in list_of_tensors], dim=stack_dim)
+
+    batch_size = list(list_of_tensordicts[0].shape)
+    batch_size.insert(dim, len(list_of_tensordicts))
+    result: TensorDictBase = list_of_tensordicts[0].apply(  # type: ignore[assignment]
+        _stack_tensors, *list_of_tensordicts[1:], batch_size=batch_size
+    )
+    return result
+
+
+class RayDataLoadingPrimer(RayTransform):
+    """A :class:`~torchrl.envs.llm.transforms.dataloading.DataLoadingPrimer` that creates a single actor that can be shared by multiple environments.
+
+    This class creates a Ray remote actor from DataLoadingPrimer that can be shared across multiple workers.
+    All method calls are delegated to the remote actor, ensuring that multiple environments iterate over
+    the same shared dataloader.
+
+    Keyword Args:
+        dataloader: A dataloader object to be used directly. Ray will handle serialization.
+        dataloader_factory: A callable that returns a dataloader. This allows for explicit
+            resource control and avoids serialization issues.
+        num_cpus (int, optional): Number of CPUs to allocate to the Ray actor.
+            Defaults to the dataloader's num_workers if available, otherwise 1.
+        num_gpus (int, optional): Number of GPUs to allocate to the Ray actor. Defaults to 0.
+        actor_name (str, optional): Name of the Ray actor to use. If provided, the actor will be reused if it already exists.
+        **kwargs: Additional keyword arguments to pass to DataLoadingPrimer.
+
+    Note:
+        Exactly one of `dataloader` or `dataloader_factory` must be provided.
+
+    Examples:
+        >>> # Option 1: Using a dataloader factory for explicit resource control
+        >>> def create_dataloader():
+        ...     return torch.utils.data.DataLoader(dataset, batch_size=32, num_workers=4)
+        >>> primer1 = RayDataLoadingPrimer(dataloader_factory=create_dataloader, num_cpus=4)
+        >>> primer2 = RayDataLoadingPrimer(dataloader_factory=create_dataloader, num_cpus=4)  # Same shared actor
+
+        >>> # Option 2: Pass dataloader directly (Ray handles serialization)
+        >>> dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, num_workers=4)
+        >>> primer1 = RayDataLoadingPrimer(dataloader=dataloader)  # num_cpus=4 inferred from num_workers
+        >>> primer2 = RayDataLoadingPrimer(dataloader=dataloader)  # Same shared actor
+    """
+
+    def __init__(
+        self,
+        *,
+        dataloader=None,
+        dataloader_factory=None,
+        num_cpus=None,
+        num_gpus=0,
+        device: DEVICE_TYPING | None = None,
+        actor_name: str | None = None,
+        **kwargs,
+    ):
+        # Validate arguments: exactly one of dataloader or dataloader_factory must be provided
+        if dataloader is not None and dataloader_factory is not None:
+            raise ValueError(
+                "Cannot provide both 'dataloader' and 'dataloader_factory'. Choose one."
+            )
+        if dataloader is None and dataloader_factory is None:
+            raise ValueError(
+                "Must provide exactly one of 'dataloader' or 'dataloader_factory'."
+            )
+
+        # Infer num_cpus from dataloader if not specified
+        if num_cpus is None:
+            if dataloader is not None:
+                num_cpus = getattr(dataloader, "num_workers", 1)
+            elif dataloader_factory is not None:
+                temp_dataloader = dataloader_factory()
+                num_cpus = getattr(temp_dataloader, "num_workers", 1)
+                del temp_dataloader
+            else:
+                num_cpus = 1
+
+            if num_cpus == 0:
+                num_cpus = 1
+
+        # Handle device setup for primers
+        primers = kwargs.get("primers", None)
+        if hasattr(primers, "device") and primers.device is not None:
+            if device is not None and device != primers.device:
+                raise ValueError(
+                    "Device mismatch between primers and device. "
+                    "Use the device argument to set the device."
+                )
+            device = primers.device
+        if hasattr(primers, "cpu"):
+            primers = primers.cpu()
+        elif hasattr(primers, "to"):
+            primers = primers.to("cpu")
+        if primers is not None:
+            kwargs["primers"] = primers
+
+        # Store creation parameters for actor creation
+        self._dataloader = dataloader
+        self._dataloader_factory = dataloader_factory
+        self._creation_kwargs = kwargs
+
+        # Call parent constructor
+        super().__init__(
+            num_cpus=num_cpus,
+            num_gpus=num_gpus,
+            device=device,
+            actor_name=actor_name,
+            **kwargs,
+        )
+
+    # Actor initialization is handled by the parent RayTransform class
+
+    def _create_actor(self, **kwargs):
+        """Create the remote DataLoadingPrimer actor."""
+        # Create the remote DataLoadingPrimer with resource specifications
+        RemoteDataLoadingPrimer = self._ray.remote(
+            num_cpus=self._num_cpus, num_gpus=self._num_gpus
+        )(DataLoadingPrimer)
+
+        if self._actor_name is not None:
+            RemoteDataLoadingPrimer = RemoteDataLoadingPrimer.options(
+                name=self._actor_name
+            )
+
+        # Create the shared actor, passing factory or dataloader as appropriate
+        if self._dataloader_factory is not None:
+            actor = RemoteDataLoadingPrimer.remote(
+                dataloader_factory=self._dataloader_factory, **self._creation_kwargs
+            )
+        else:
+            actor = RemoteDataLoadingPrimer.remote(
+                dataloader=self._dataloader, **self._creation_kwargs
+            )
+
+        return actor
+
+    @property
+    def dataloader(self):
+        """Get dataloader property."""
+        raise NotImplementedError(
+            "dataloader is not implemented for RayDataLoadingPrimer"
+        )
+
+    @property
+    def endless_dataloader(self):
+        """Get endless_dataloader property."""
+        raise NotImplementedError(
+            "endless_dataloader is not implemented for RayDataLoadingPrimer"
+        )
+
+    def __repr__(self):
+        """String representation."""
+        try:
+            if hasattr(self, "_actor") and self._actor is not None:
+                return self._ray.get(self._actor.__repr__.remote())
+            else:
+                return "RayDataLoadingPrimer(actor=None)"
+        except Exception:
+            return f"RayDataLoadingPrimer(actor={getattr(self, '_actor', 'None')})"
+
+    @property
+    def stack_method(self):
+        """Get stack_method property."""
+        raise NotImplementedError(
+            "stack_method is not implemented for RayDataLoadingPrimer"
+        )
+
+    @property
+    def repeats(self):
+        """Get repeats property."""
+        return self._ray.get(self._actor.__getattribute__.remote("repeats"))
+
+    @property
+    def data_keys(self):
+        """Get data_keys property."""
+        return self._ray.get(self._actor.__getattribute__.remote("data_keys"))
+
+    @property
+    def primers(self):
+        """Get primers property."""
+        return self._ray.get(self._actor.__getattribute__.remote("primers"))
+
+    @primers.setter
+    def primers(self, value: TensorSpec):
+        """Set primers property."""
+        self._ray.get(self._actor._set_attr.remote("primers", value))
+
+    # TensorDictPrimer methods
+    def init(self, tensordict: TensorDictBase | None):
+        """Initialize."""
+        return self._ray.get(self._actor.init.remote(tensordict))
+
+    def reset_dataloader(self):
+        """Reset the dataloader."""
+        return self._delegate_method_call("reset_dataloader")
+
+    @_map_input_output_device
+    def _load_from_dataloader(
+        self, reset: torch.Tensor | None = None
+    ) -> TensorDictBase:
+        """Load data from the dataloader."""
+        result = self._delegate_method_call("_load_from_dataloader", reset)
+
+        # Ensure proper batch dimensions: if result is scalar (ndim=0), unsqueeze to add batch dim
+        # This matches the behavior in the original DataLoadingPrimer._load_from_dataloader
+        if hasattr(result, "ndim") and not result.ndim:
+            result = result.unsqueeze(0)
+
+        return result
+
+    @property
+    def primers(self):
+        """Get primers property."""
+        return self._delegate_property_get("primers")
+
+    @primers.setter
+    def primers(self, value: TensorSpec):
+        """Set primers property."""
+        self._delegate_property_set("primers", value)
+
+    @_map_input_output_device
+    def _reset_func(
+        self, tensordict: TensorDictBase | None, tensordict_reset: TensorDictBase | None
+    ) -> TensorDictBase | None:
+        """Reset function."""
+        result = super()._reset_func(tensordict, tensordict_reset)
+
+        # Handle batch size expansion locally since remote actor lacks parent context
+        # This mimics the batch expansion logic from TensorDictPrimer._reset_func
+        if (
+            self.parent
+            and self.parent.batch_locked
+            and hasattr(result, "apply")  # Check if it's a TensorDict-like object
+        ):
+            # Ensure result has proper batch dimensions to match parent
+            expected_batch_size = self.parent.batch_size
+            if result.batch_size != expected_batch_size:
+                # Expand result to match expected batch size
+                result = result.expand(expected_batch_size)
+
+        return result
+
+    # Additional methods and properties are handled by the parent RayTransform class
+
+
+class DataLoadingPrimer(TensorDictPrimer, metaclass=_RayServiceMetaClass):
+    """A primer that loads data from a dataloader and converts it into a tensordict using ``stack_method``.
+
+    Args:
+        dataloader (Iterable[Dict[str, Any]]): The dataloader to load data from.
+            During collection, we will attempt to convert it into a tensordict using :func:`~tensordict.from_dict` or a
+            similar function.
+            It is assumed that the elements retrieved from the dataloader come in batches along the first dimension
+            of every tensor, unless `dataloader.batch_size=0`.
+            The dataloader must yield mappable data structures (e.g., dictionaries).
+            If a dataloader_factory is provided, it will be used to create a fresh dataloader and this argument can be
+            omitted.
+
+    Keyword Args:
+        primers (Composite | None, optional): The primers to use for each key in the dataloader. Defaults to None.
+        stack_method (Callable[[Any], Any] | Literal["as_nested_tensor", "as_padded_tensor"], optional): The method to
+            use for stacking the data. Defaults to ``maybe_dense_stack``.
+        repeats (int, optional): How many times the same sample needs to appear successively. This can be useful in
+            situations like GRPO where a single prompt is used multiple times to estimate the advantage using Monte-Carlo
+            samples (rather than an advantage module).
+        batch_size (int, torch.Size or None): the batch-size of the data delivered by the transform.
+            This is somewhat unrelated to the batch-size of the dataloader, in the sense that this number may or may
+            not match the DL's batch size.
+            If left empty, the batch-size is inferred from `dataloader.batch_size` if that attribute exists. If not,
+            an empty batch-size will be used (`torch.Size([])`).
+
+            .. note:: The batch-size of the Primer must match the batch-size of the parent environment (typically a
+                wrapper around :class:`~torchrl.envs.LLMEnv`).
+
+        group_repeats (bool, optional): if ``True``, the batch-size is multiplied by the number of repeats such that
+            all repeats are grouped in a single batch collected from the buffer. Defaults to ``False``.
+        dataloader_factory (Callable[[], Iterable[dict[str, Any]]], optional): A callable that returns a dataloader.
+            This allows for explicit resource control and avoids serialization issues.
+        use_ray_service (bool, optional): if ``True``, returns a :class:`RayDataLoadingPrimer` instance instead,
+            which creates a Ray actor for shared dataloader access across multiple environments.
+            Defaults to ``False``.
+
+    Attributes:
+        dataloader (Iterable[Any]): The dataloader to load data from.
+        endless_dataloader (Iterable[Any]): An endless iterator over the dataloader.
+        stack_method (Callable[[Any], Any]): The method to use for stacking the data.
+
+    .. seealso:: :class:`~torchrl.envs.LLMEnv` and :class:`~torchrl.envs.LLMEnv.from_dataloader`.
+
+    Examples:
+        Using the regular DataLoadingPrimer:
+
+        >>> dataloader = torch.utils.data.DataLoader(dataset, batch_size=32)
+        >>> primer = DataLoadingPrimer(dataloader=dataloader)  # Regular implementation
+
+        Using the Ray-based implementation for shared dataloader access:
+
+        >>> primer = DataLoadingPrimer(dataloader=dataloader, use_ray_service=True)  # Returns RayDataLoadingPrimer
+        >>> # Multiple environments can now share the same dataloader through the Ray actor
+
+    Example of a dataloader yielding strings:
+        >>> import random
+        >>> import string
+        >>> import tensordict as td
+        >>> import torch
+        >>> from tensordict import TensorDict
+        >>> from torchrl.data import Unbounded
+        >>> from torchrl.envs import DataLoadingPrimer, LLMEnv
+        >>> td.set_capture_non_tensor_stack(False).set()
+        >>> class DummyDataLoader:
+        ...     '''A dummy dataloader that generates random strings.'''
+        ...     def __init__(self, batch_size: int = 0):
+        ...         self.batch_size = batch_size
+        ...     def generate_random_string(self, length: int = 10) -. str:
+        ...         '''Generate a random string of a given length.'''
+        ...         return ''.join(random.choice(string.ascii_lowercase) for _ in range(length))
+        ...     def __iter__(self):
+        ...         return self
+        ...     def __next__(self):
+        ...         if self.batch_size == 0:
+        ...             return self.generate_random_string()
+        ...         else:
+        ...             return [self.generate_random_string() for _ in range(self.batch_size)]
+        >>> # Create an LLM environment with string-to-string input/output.
+        >>> env = LLMEnv(from_text=True)
+        >>> # Append a DataLoadingPrimer to the environment.
+        >>> env = env.append_transform(
+        >>>     DataLoadingPrimer(
+        >>>         dataloader=DummyDataLoader(),
+        >>>         example_data="a string!",
+        >>>     )
+        >>> )
+        >>> # Test the environment.
+        >>> print(env.rand_action(TensorDict()))
+        TensorDict(
+            fields={
+                action: NonTensorData(data=a string, batch_size=torch.Size([]), device=None)},
+            batch_size=torch.Size([]),
+            device=None,
+            is_shared=False)
+        >>> print(env.rollout(3))
+        TensorDict(
+            fields={
+                action: NonTensorStack(
+                    ['a string', 'a string', 'a string'],
+                    batch_size=torch.Size([3]),
+                    device=None),
+                done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                next: TensorDict(
+                    fields={
+                        done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        observation: NonTensorStack(
+                            ['zxwvupirska string', 'zxwvupirska stringa string...,
+                            batch_size=torch.Size([3]),
+                            device=None),
+                        terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([3]),
+                    device=None,
+                    is_shared=False),
+                observation: NonTensorStack(
+                    ['zxwvupirsk', 'zxwvupirska string', 'zxwvupirska ...,
+                    batch_size=torch.Size([3]),
+                    device=None),
+                terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+            batch_size=torch.Size([3]),
+            device=None,
+            is_shared=False)
+        >>> # Roll out the environment with a specific initial state.
+        >>> init_state = env.reset(TensorDict(batch_size=[3]))
+        >>> print(env.rollout(3, auto_reset=False, tensordict=init_state))
+        TensorDict(
+            fields={
+                action: NonTensorStack(
+                    [['a string', 'a string', 'a string'], ['a string'...,
+                    batch_size=torch.Size([3, 3]),
+                    device=None),
+                done: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                next: TensorDict(
+                    fields={
+                        done: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        observation: NonTensorStack(
+                            [[array(['nngcmflsana string', 'vrrbnhzpmga string...,
+                            batch_size=torch.Size([3, 3]),
+                            device=None),
+                        terminated: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        truncated: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([3, 3]),
+                    device=None,
+                    is_shared=False),
+                observation: NonTensorStack(
+                    [['nngcmflsan', array(['nngcmflsana string', 'vrrb...,
+                    batch_size=torch.Size([3, 3]),
+                    device=None),
+                terminated: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                truncated: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+            batch_size=torch.Size([3, 3]),
+            device=None,
+            is_shared=False)
+
+    Example of dataloader yielding tensors:
+        >>> import random
+        >>> import string
+        >>>
+        >>> import tensordict as td
+        >>> import torch
+        >>> from tensordict import TensorDict
+        >>> from torchrl.data import Unbounded
+        >>> from torchrl.envs import DataLoadingPrimer, LLMEnv
+        >>>
+        >>> td.set_capture_non_tensor_stack(False).set()
+        >>>
+        >>>
+        >>> class DummyTensorDataLoader:
+        ...     '''A dummy dataloader that generates tensors of random int64 values.'''
+        ...
+        ...     def __init__(self, batch_size: int = 0, max_length: int = 10, padding: bool = False):
+        ...         '''
+        ...         Args:
+        ...             batch_size (int, optional): The batch size of the generated tensors. Defaults to 0.
+        ...             max_length (int, optional): The maximum length of the generated tensors. Defaults to 10.
+        ...             padding (bool, optional): Whether to pad the tensors to the maximum length. Defaults to `False`.
+        ...         '''
+        ...         self.batch_size = batch_size
+        ...         self.max_length = max_length
+        ...         self.padding = padding
+        ...
+        ...     def generate_random_tensor(self) -. torch.Tensor:
+        ...         '''Generate a tensor of random int64 values.'''
+        ...         length = random.randint(1, self.max_length)
+        ...         return torch.tensor([random.randint(0, 100) for _ in range(length)], dtype=torch.int64)
+        ...
+        ...     def pad_tensor(self, tensor: torch.Tensor) -. torch.Tensor:
+        ...         '''Pad a tensor to the maximum length.'''
+        ...         padding_length = self.max_length - len(tensor)
+        ...         return torch.cat((torch.zeros(padding_length, dtype=torch.int64), tensor))
+        ...
+        ...     def __iter__(self):
+        ...         return self
+        ...
+        ...     def __next__(self):
+        ...         if self.batch_size == 0:
+        ...             tensor = self.generate_random_tensor()
+        ...             return self.pad_tensor(tensor) if self.padding else tensor
+        ...         else:
+        ...             tensors = [self.generate_random_tensor() for _ in range(self.batch_size)]
+        ...             if self.padding:
+        ...                 tensors = [self.pad_tensor(tensor) for tensor in tensors]
+        ...                 return torch.stack(tensors)
+        ...             else:
+        ...                 return tensors
+        >>>
+        >>> # Create an LLM environment with non-string input/output and append a DataLoadingPrimer.
+        >>> env = LLMEnv(from_text=False)
+        >>> env = env.append_transform(
+        >>>     DataLoadingPrimer(
+        >>>         dataloader=DummyTensorDataLoader(),
+        >>>         data_specs=[Unbounded(shape=(-1,), dtype=torch.int64)],
+        >>>     )
+        >>> )
+        >>> print(env.rand_action(TensorDict()))
+        TensorDict(
+            fields={
+                action: Tensor(shape=torch.Size([1]), device=cpu, dtype=torch.int64, is_shared=False)},
+            batch_size=torch.Size([]),
+            device=None,
+            is_shared=False)
+        >>> print(env.rollout(3))
+        LazyStackedTensorDict(
+            fields={
+                action: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+                done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                next: LazyStackedTensorDict(
+                    fields={
+                        done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        observation: Tensor(shape=torch.Size([3, -1]), device=cpu, dtype=torch.int64, is_shared=False),
+                        terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    exclusive_fields={
+                    },
+                    batch_size=torch.Size([3]),
+                    device=None,
+                    is_shared=False,
+                    stack_dim=0),
+                observation: Tensor(shape=torch.Size([3, -1]), device=cpu, dtype=torch.int64, is_shared=False),
+                terminated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                truncated: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+            exclusive_fields={
+            },
+            batch_size=torch.Size([3]),
+            device=None,
+            is_shared=False,
+            stack_dim=0)
+        >>> # Create an LLM environment with padded tensor input/output and append a DataLoadingPrimer.
+        >>> env = LLMEnv(from_text=False)
+        >>> env = env.append_transform(
+        >>>     DataLoadingPrimer(
+        >>>         dataloader=DummyTensorDataLoader(padding=True),
+        >>>         data_specs=[Unbounded(shape=(-1,), dtype=torch.int64)],
+        >>>         stack_method="as_padded_tensor",
+        >>>     )
+        >>> )
+        >>> print(env.rollout(3, auto_reset=False, tensordict=env.reset(TensorDict(batch_size=[3]))))
+        LazyStackedTensorDict(
+            fields={
+                action: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.int64, is_shared=False),
+                done: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                next: LazyStackedTensorDict(
+                    fields={
+                        done: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        observation: Tensor(shape=torch.Size([3, 3, -1]), device=cpu, dtype=torch.int64, is_shared=False),
+                        terminated: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        truncated: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    exclusive_fields={
+                    },
+                    batch_size=torch.Size([3, 3]),
+                    device=None,
+                    is_shared=False,
+                    stack_dim=1),
+                observation: Tensor(shape=torch.Size([3, 3, -1]), device=cpu, dtype=torch.int64, is_shared=False),
+                terminated: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                truncated: Tensor(shape=torch.Size([3, 3, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+            exclusive_fields={
+            },
+            batch_size=torch.Size([3, 3]),
+            device=None,
+            is_shared=False,
+            stack_dim=1)
+
+    """
+
+    _RayServiceClass = RayDataLoadingPrimer
+
+    def __init__(
+        self,
+        dataloader: Iterable[dict[str, Any]] | None = None,
+        *,
+        dataloader_factory: Callable[[], Iterable[dict[str, Any]]] | None = None,
+        primers: Composite | None = None,
+        stack_method: Callable[[Any], Any]
+        | Literal["as_nested_tensor", "as_padded_tensor"]
+        | None = None,
+        batch_size: int | torch.Size | None = None,
+        repeats: int | None = None,
+        device: torch.device | None = None,
+        group_repeats: bool = False,
+        use_ray_service: bool = False,
+    ):
+        # Validate arguments: exactly one of dataloader or dataloader_factory must be provided
+        if dataloader is not None and dataloader_factory is not None:
+            raise ValueError(
+                "Cannot provide both 'dataloader' and 'dataloader_factory'. Choose one."
+            )
+        if dataloader is None and dataloader_factory is None:
+            raise ValueError(
+                "Must provide exactly one of 'dataloader' or 'dataloader_factory'."
+            )
+
+        # Initialize dataloader from factory if provided
+        if dataloader_factory is not None:
+            self.dataloader = dataloader_factory()
+            self.dataloader_factory = dataloader_factory
+        else:
+            self.dataloader = dataloader
+            self.dataloader_factory = None
+
+        if repeats is None:
+            repeats = 0
+        self.repeats = repeats
+
+        # Determine batch-size
+        #  We must distinguish the batch-size of the DL and the batch size of the transform.
+        #  We may want more or less elements than the DL and the logic is slightly different so we
+        #  allow to recompose batches on the fly. If the DL has a batch-size, every element will be
+        #  unbound and stored in a queue. Otherwise, we get as many elements from the DL to fulfill
+        #  the required batch-size.
+        #
+        #  If the batch-size is passed, we will stack as many elements as necessary to fulfill this.
+        #  If not, we try to get it from the dataloader. Contrary to the dataloader, we will always
+        #  deliver the same batch-size (we create an infinite dataloader and reset when it's done),
+        #  whereas DLs with drop_last=False may return batches of different sizes.
+        #
+        # If the batch size passed to the transform is empty (torch.Size(())) or 0, we will consider that
+        #  the batch-size is determined on-the-fly.
+        #
+        # A batch-size of 0 in the dataloader means no batch-size.
+        #
+        # If needed, the various repeats can be grouped in a single batch through group_repeats.
+        #
+        # If auto_batch_size is on, we call auto_batch_size=True when doing TensorDict.from_dict:
+        #  That way we get a tensordict of the right batch-size.
+        # If the dataloader has no batch-size, we're not sure that we can determine the batch-size
+        #  automatically so we will consider that each element in the DL has a batch-size of 0 (ie,
+        #  a single non-batched element is returned at a time).
+
+        if batch_size is None:
+            batch_size = getattr(dataloader, "batch_size", torch.Size([]))
+        if batch_size == 0:
+            batch_size = torch.Size(())
+        if not isinstance(batch_size, (list, tuple)):
+            if isinstance(batch_size, int):
+                batch_size_tuple = (batch_size,)
+            elif isinstance(batch_size, torch.Size):
+                batch_size_tuple = tuple(batch_size)
+            else:
+                batch_size_tuple = (batch_size,)
+        else:
+            batch_size_tuple = batch_size
+        batch_size = torch.Size(batch_size_tuple)
+        auto_batch_size = getattr(dataloader, "batch_size", 1) != 0
+
+        if len(batch_size) > 1:
+            raise ValueError(
+                f"batch_size can only be 0 or 1D, got batch_size={batch_size}."
+            )
+
+        # We deliver all the repeats in the same batch
+        if repeats and group_repeats:
+            if batch_size == torch.Size([]):
+                batch_size = torch.Size((repeats,))
+            else:
+                batch_size = torch.Size([batch_size[0] * repeats])
+
+        self._queue = deque()
+        self.auto_batch_size = auto_batch_size
+        self.batch_size = batch_size
+        self.endless_dataloader = self._endless_iter(self.dataloader)
+
+        if stack_method is None:
+            stack_method = lazy_stack
+        elif stack_method == "as_nested_tensor":
+            stack_method = as_nested_tensor
+        elif stack_method == "as_padded_tensor":
+            stack_method = as_padded_tensor
+        elif not callable(stack_method):
+            raise ValueError(f"Unknown stack_method={stack_method}")
+        self.stack_method = stack_method
+
+        if primers is None:
+            # We can get the primer from the dataloader itself
+            data = self._load_from_dataloader()
+            primers = make_composite_from_td(
+                data, dynamic_shape=True, unsqueeze_null_shapes=False
+            )
+            if batch_size:
+                primers = primers.expand(batch_size)
+            self._queue.insert(0, data)
+            self.data_keys = list(primers.keys(True, True))
+        else:
+            self.data_keys = list(primers.keys(True, True))
+
+        super().__init__(
+            primers=primers,
+            default_value=self._load_from_dataloader,
+            reset_key=None,
+            expand_specs=None,
+            single_default_value=True,
+            call_before_env_reset=True,
+            device=device,
+        )
+        self._reset_key = "_reset"
+
+    def reset_dataloader(self):
+        """Reset the dataloader.
+
+        This is useful when the dataloader is not infinite and we want to reset it.
+        If a dataloader_factory was provided, it will be used to create a fresh dataloader.
+
+        Returns:
+            self: The transform itself.
+        """
+        self._queue.clear()
+        if self.dataloader_factory is not None:
+            # Create a fresh dataloader from the factory
+            self.dataloader = self.dataloader_factory()
+        self.endless_dataloader = self._endless_iter(self.dataloader)
+        return self
+
+    @classmethod
+    def _endless_iter(self, obj):
+        while True:
+            yield from obj
+
+    _device: torch.device | None = None
+
+    @property
+    def device(self) -> torch.device | None:
+        if self._device is None:
+            primers = getattr(self, "primers", None)
+            if primers is not None:
+                device = self.primers.device
+            else:
+                parent = getattr(self, "parent", None)
+                if parent is not None:
+                    device = getattr(parent, "device", None)
+                else:
+                    device = None
+            self._device = device
+        return self._device
+
+    @device.setter
+    def device(self, device: torch.device | None):
+        self._device = device
+
+    def _load_from_dataloader(
+        self, reset: torch.Tensor | None = None
+    ) -> TensorDictBase:
+        """Loads a single element from the dataloader, or alternatively from the buffer.
+
+        If `reset` is passed, then one element per reset will be loaded.
+        """
+        device = self.device
+
+        if reset is not None:
+            if not reset.any():
+                raise RuntimeError("reset must have at least one True value.")
+            if reset.ndim > 0:
+                loaded = [
+                    self._load_from_dataloader().to(device) for _ in range(reset.sum())
+                ]
+                return self.stack_method(loaded)
+
+        if len(self._queue) > 0:
+            result = self._queue.popleft()
+            if result.device != device:
+                result = result.to(device)
+            return result
+
+        data = next(self.endless_dataloader)
+        # Some heuristic here:
+        # if data is a map, assume its keys match the keys in spec
+        # TODO: one could rename the keys too
+        if is_tensor_collection(data):
+            out = data
+        elif isinstance(data, Mapping):
+            out = TensorDict.from_dict(
+                data,
+                auto_batch_size=self.auto_batch_size,
+                batch_dims=int(bool(self.auto_batch_size or self.batch_size)),
+                device=device,
+            )
+        else:
+            raise TypeError(
+                "Data loader must return a mapping that can be automatically cast to a tensordict. Check that you have "
+                "the appropriate collate_fn in your dataloader to do so."
+            )
+        if not out.ndim:
+            out = out.unsqueeze(0)
+        self._queue.extend(
+            [d for d in out.unbind(0) for _ in range(max(1, self.repeats))]
+        )
+        out = self._queue.popleft()
+        return out
+
+    def set_container(self, container: Transform | EnvBase) -> None:
+        result = super().set_container(container)
+        # Check batch size
+        parent = getattr(self, "parent", None)
+        if (
+            self.batch_size is not None
+            and parent is not None
+            and parent.batch_size != self.batch_size
+        ):
+            warnings.warn(
+                f"The parent env has a different batch size than the {type(self).__name__} transform."
+            )
+        return result
+
+    def _update_primers_batch_size(self, parent_batch_size):
+        """Update the primers to match the parent's batch size.
+
+        This method is called remotely to ensure the remote actor's primers
+        have the correct batch dimensions.
+        """
+        if hasattr(self.primers, "expand"):
+            # Expand primers to match the parent batch size
+            if self.primers.shape != parent_batch_size:
+                self.primers = self.primers.expand(parent_batch_size)
+
+    def __repr__(self) -> str:
+        class_name = self.__class__.__name__
+        return f"{class_name}(primers={self.primers}, dataloader={self.dataloader})"