torchrl-nightly 2025.6.19__cp313-cp313-macosx_10_13_universal2.whl → 2025.6.21__cp313-cp313-macosx_10_13_universal2.whl

torchrl/collectors/collectors.py

@@ -352,8 +352,8 @@ class DataCollectorBase(IterableDataset, metaclass=abc.ABCMeta):
                 self._iterator = iter(self)
             out = next(self._iterator)
             # if any, we don't want the device ref to be passed in distributed settings
-            if out is not None:
-                out.clear_device_()
+            if out is not None and (out.device != "cpu"):
+                out = out.copy().clear_device_()
             return out
         except StopIteration:
             return None
@@ -892,7 +892,10 @@ class SyncDataCollector(DataCollectorBase):
             and hasattr(self.postproc, "to")
             and self.storing_device
         ):
-            self.postproc.to(self.storing_device)
+            postproc = self.postproc.to(self.storing_device)
+            if postproc is not self.postproc and postproc is not None:
+                self.postproc = postproc
+
         if frames_per_batch % self.n_env != 0 and RL_WARNINGS:
             warnings.warn(
                 f"frames_per_batch ({frames_per_batch}) is not exactly divisible by the number of batched environments ({self.n_env}), "
@@ -1253,9 +1256,9 @@ class SyncDataCollector(DataCollectorBase):
                     yield
                     continue
                 self._increment_frames(tensordict_out.numel())
-                if self.verbose:
-                    torchrl_logger.info("Collector: postproc.")
                 tensordict_out = self._postproc(tensordict_out)
+                if self.verbose:
+                    torchrl_logger.info("Collector: postproc done.")
                 if self.return_same_td:
                     # This is used with multiprocessed collectors to use the buffers
                     # stored in the tensordict.
@@ -1765,8 +1768,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 +1927,7 @@ class _MultiDataCollector(DataCollectorBase):
         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 +1963,22 @@ class _MultiDataCollector(DataCollectorBase):
         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 +2096,11 @@ class _MultiDataCollector(DataCollectorBase):
         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 +2111,8 @@ class _MultiDataCollector(DataCollectorBase):
         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 +2242,7 @@ class _MultiDataCollector(DataCollectorBase):
         )
         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 +2301,7 @@ class _MultiDataCollector(DataCollectorBase):
                     "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 +2793,9 @@ class MultiSyncDataCollector(_MultiDataCollector):
             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 +2876,9 @@ class MultiSyncDataCollector(_MultiDataCollector):
                 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 +2924,12 @@ class MultiSyncDataCollector(_MultiDataCollector):
 
             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 +3182,7 @@ class MultiaSyncDataCollector(_MultiDataCollector):
             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 +3246,7 @@ class MultiaSyncDataCollector(_MultiDataCollector):
                 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/collectors/llm/base.py

@@ -242,6 +242,11 @@ class LLMCollector(SyncDataCollector):
         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:
@@ -361,9 +366,10 @@ class LLMCollector(SyncDataCollector):
                         )
                     self._yield_queues[idx].clear()
         result = self._trajectory_queue.popleft()
-        torchrl_logger.info(
-            f"LLMCollector: Yielding completed trajectory with shape {result.shape}."
-        )
+        if self.verbose:
+            torchrl_logger.info(
+                f"LLMCollector: Yielding completed trajectory with shape {result.shape}."
+            )
         return result
 
     started = False
@@ -422,9 +428,10 @@ class LLMCollector(SyncDataCollector):
                 self.env.async_step_and_maybe_reset_send(env_input)
 
         result = self._trajectory_queue.popleft()
-        torchrl_logger.info(
-            f"LLMCollector: Yielding completed trajectory with shape {result.shape}."
-        )
+        if self.verbose:
+            torchrl_logger.info(
+                f"LLMCollector: Yielding completed trajectory with shape {result.shape}."
+            )
         return result
 
     as_remote = as_remote

torchrl/collectors/llm/ray_collector.py

@@ -134,6 +134,9 @@ class RayLLMCollector(LLMCollector):
             verbose=verbose,
         )
 
+    def set_postproc(self, postproc: Callable[[TensorDictBase], TensorDictBase]):
+        return ray.get(self._collector.set_postproc.remote(postproc))
+
     def _next_remote(self) -> None:
         return self._collector.next.remote()
 

torchrl/data/__init__.py

@@ -17,6 +17,7 @@ from .llm import (
     RolloutFromModel,
     TensorDictTokenizer,
     TokenizedDatasetLoader,
+    TopKRewardSelector,
 )
 from .map import (
     BinaryToDecimal,
@@ -116,6 +117,7 @@ __all__ = [
     "Categorical",
     "Choice",
     "ContentBase",
+    "TopKRewardSelector",
     "Composite",
     "CompositeSpec",
     "ConstantKLController",

torchrl/data/datasets/minari_data.py

@@ -350,7 +350,7 @@ class MinariExperienceReplay(BaseDatasetExperienceReplay):
             # Add a "done" entry
             if self.split_trajs:
                 with td_data.unlock_():
-                    from torchrl.objectives.utils import split_trajectories
+                    from torchrl.collectors.utils import split_trajectories
 
                     td_data = split_trajectories(td_data).memmap_(self.data_path)
             with open(self.metadata_path, "w") as metadata_file:

torchrl/data/llm/__init__.py

@@ -13,6 +13,7 @@ from .dataset import (
 )
 from .prompt import PromptData, PromptTensorDictTokenizer
 from .reward import PairwiseDataset, RewardData
+from .topk import TopKRewardSelector
 from .utils import AdaptiveKLController, ConstantKLController, RolloutFromModel
 
 __all__ = [
@@ -30,4 +31,5 @@ __all__ = [
     "TokenizedDatasetLoader",
     "create_infinite_iterator",
     "get_dataloader",
+    "TopKRewardSelector",
 ]

torchrl/data/llm/chat.py

@@ -11,19 +11,27 @@ from typing import Literal
 
 import torch
 
-
-from tensordict import lazy_stack, LazyStackedTensorDict, list_to_stack, TensorClass
+from tensordict import (
+    lazy_stack,
+    LazyStackedTensorDict,
+    list_to_stack,
+    TensorClass,
+    TensorDict,
+)
 from tensordict.utils import _maybe_correct_neg_dim
-
 from torchrl._utils import logger as torchrl_logger
 
 
 _CHAT_TEMPLATES = {
     "chatml_format": """{% for message in messages %}
+    {%- if message['role'] == 'assistant' %}
+    {% generation %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endgeneration %}
+    {%- else %}
     {{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}
+    {%- endif %}
 {% endfor %}
 {%- if add_generation_prompt %}
-    {{- '<|im_start|>assistant\n' }}
+    {% generation %}{{- '<|im_start|>assistant\n' }}{% endgeneration %}
 {%- endif %}
 """,
     "qwen": """
@@ -283,7 +291,7 @@ class History(TensorClass["nocast"]):
 
         Keyword Args:
             tokenizer (transformers.PreTrainedTokenizer | transformers.AutoProcessor): The tokenizer to use.
-            add_generation_prompt (bool, optional): Whether to add a generation prompt. Defaults to `True`.
+            add_generation_prompt (bool, optional): Whether to add a generation prompt (e.g. `"<|im_start|>assistant"`). Defaults to `True`.
             chat_template (str, optional): The chat template to use. Defaults to the tokenizer's default template.
             chat_template_name (Literal["chatml_format", "qwen"], optional): The name of the chat template to use.
                 Prevalent over `tokenizer.chat_template`. Defaults to `None`.
@@ -294,6 +302,7 @@ class History(TensorClass["nocast"]):
             return_tensors (str | None, optional): The type of tensors to return. Defaults to "pt".
             return_dict (bool, optional): Whether to return a dictionary. Defaults to `False`.
             return_assistant_tokens_mask (bool, optional): Whether to return a mask of the assistant generated tokens.
+                If `True`, the mask will be written to the `assistant_masks` key.
                 For tokens generated by the assistant, the mask will contain `1`.
                 For user and system tokens, the mask will contain `0`.
                 This functionality is only available for chat templates that support it via the `{% generation %}` keyword.
@@ -316,6 +325,11 @@ class History(TensorClass["nocast"]):
                 raise RuntimeError(
                     "You must specify a tokenizer to use when chat_template is not specified."
                 )
+            elif "qwen" in getattr(tokenizer, "name_or_path", "").lower():
+                # We prefer our implementation of the Qwen template,
+                #  since it accounts for the assistant's masking.
+                chat_template = _CHAT_TEMPLATES["qwen"]
+                chat_template_name = None
             else:
                 chat_template = tokenizer.chat_template
         if chat_template is None:
@@ -334,7 +348,7 @@ class History(TensorClass["nocast"]):
             return_dict = False
 
         if self.ndim > 1:
-            return [
+            result = [
                 self[i].apply_chat_template(
                     tokenizer=tokenizer,
                     add_generation_prompt=add_generation_prompt,
@@ -351,12 +365,16 @@ class History(TensorClass["nocast"]):
                 )
                 for i in range(self.batch_size[0])
             ]
+            if return_dict:
+                return lazy_stack(result)
+            else:
+                return result
         self_flat = self.view(-1)
         # tolist_first=True is needed to avoid having a list of dict of dicts, but a list of dicts of lists of dicts
         self_flat = self_flat.tolist(tolist_first=True)
         # Remove the "<none>" role
         self_flat = [item for item in self_flat if item["role"] != "<none>"]
-        return tokenizer.apply_chat_template(
+        result = tokenizer.apply_chat_template(
             conversation=self_flat,
             add_generation_prompt=add_generation_prompt,
             chat_template=chat_template,
@@ -369,6 +387,16 @@ class History(TensorClass["nocast"]):
             return_assistant_tokens_mask=return_assistant_tokens_mask,
             **kwargs,
         )
+        if not isinstance(result, (torch.Tensor, list, str)):
+            result = TensorDict.from_dict(result, auto_batch_size=True, batch_dims=1)
+            # If self has a batch_dims of 1, we have just the time dimension, so we need to remove the batch dim from the result
+            if self.batch_dims == 1:
+                if result.batch_size[0] != 1:
+                    raise RuntimeError(
+                        f"Expected a batch size of 1, got {result.batch_size[0]}."
+                    )
+                result = result.squeeze(0)
+        return result
 
     @classmethod
     def from_text(
@@ -376,10 +404,20 @@ class History(TensorClass["nocast"]):
         text: str | list[str],
         chat_template_name: Literal["chatml_format", "qwen"] | None = None,
         chat_template: str | None = None,
+        tokenizer: transformers.AutoTokenizer  # noqa: F821
+        | transformers.AutoProcessor  # noqa: F821
+        | None = None,
     ) -> History:
-        if chat_template_name in ("chatml_format", None):
+        if chat_template_name is None and chat_template is None:
+            if "qwen" in getattr(tokenizer, "name_or_path", "").lower():
+                # We can automatically detect the template name from the tokenizer
+                #  and use the precoded parser.
+                chat_template_name = "qwen"
+            else:
+                chat_template_name = "chatml_format"
+        elif chat_template_name in ("chatml_format",):
             func = cls._inv_chatml
-        elif chat_template_name == "qwen":
+        elif chat_template_name in ("qwen",):
             func = cls._inv_qwen
         else:
             raise NotImplementedError(
@@ -736,3 +774,15 @@ class History(TensorClass["nocast"]):
         }
 
         return Composite(defaults, shape=shape[:-1], data_cls=cls)
+
+    @classmethod
+    def from_chats(cls, chats: list[list[dict]]) -> History:
+        """Create a History object from a list of chats.
+
+        Args:
+            chats (list[list[dict]]): A list of chats, where each chat is a list of dictionaries.
+        """
+        if isinstance(chats[0], dict):
+            return lazy_stack([cls(**chat) for chat in chats])
+        else:
+            return lazy_stack([cls.from_chats(chat) for chat in chats])

torchrl/data/llm/topk.py

@@ -0,0 +1,186 @@
+# 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 defaultdict, deque
+from typing import Any
+
+import torch
+from tensordict import NestedKey, TensorDictBase
+from torchrl._utils import logger as torchrl_logger
+from torchrl.envs.transforms import Transform
+
+
+class TopKRewardSelector(Transform):
+    """A replay-buffer transform that selects the top-k rewards for each prompt.
+
+    Args:
+        total_dialog_turns (int): Number of dialog turns to keep in memory for the top-k selection.
+        topk_size (int): Number of top-k rewards to select. Must be smaller than or equal to total_dialog_turns.
+        prompt_key (NestedKey): Key to the prompt in the tensordict. Defaults to "text".
+        rewards_key (NestedKey): Key to the rewards in the tensordict. Defaults to ("next", "reward").
+        done_key (NestedKey): Key to the done state in the tensordict. Defaults to ("next", "done").
+        verbose (bool): Whether to print verbose information. Defaults to `False`.
+
+    Example:
+        >>> from torchrl.data import ReplayBuffer, LazyStackStorage, SamplerWithoutReplacement
+        >>> from tensordict import TensorDict, lazy_stack
+        >>> import torch
+        >>> from torchrl.data.llm.topk import TopKRewardSelector
+        >>> # Create a replay buffer with 50 items, a sampler that samples without replacement, and a batch size of 5
+        >>> rb = ReplayBuffer(
+        ...     storage=LazyStackStorage(50),
+        ...     sampler=SamplerWithoutReplacement,
+        ...     batch_size=5,
+        ... )
+        >>> # Create a tensordict with 50 items, each with 10 dialog turns
+        >>> td = lazy_stack(
+        ...     [
+        ...         TensorDict(
+        ...             {
+        ...                 ("next", "done"): torch.full((1, 1), True),
+        ...                 # Reward for i+5 tokens
+        ...                 ("next", "reward"): torch.full((i + 5, 1), i),
+        ...                 # total of 10 dialogs per prompt
+        ...                 "text": f"Prompt {i // 5}",
+        ...             }
+        ...         )
+        ...         for i in range(50)
+        ...     ]
+        ... )
+        >>> # Create a top-k reward selector with 5 dialog turns and a top-k size of 3
+        >>> topk = TopKRewardSelector(total_dialog_turns=5, topk_size=3)
+        >>> rb.append_transform(topk)
+        >>> for _td in td.chunk(25):
+        ...     rb.extend(_td)
+        >>> # Only wrote top3 of 50 items in 10 groups of 5
+        >>>  assert rb.write_count == 30
+        >>> assert len(rb) == 30
+        >>> r3 = rb[:3].get(("next", "reward"), as_padded_tensor=True).squeeze()
+        >>> # 0 and 1 are missing because they're not part of the top-k
+        >>> assert (
+        ...     r3 == torch.tensor(
+        ...         [
+        ...             [4, 4, 4, 4, 4, 4, 4, 4, 4],
+        ...             [3, 3, 3, 3, 3, 3, 3, 3, 0],
+        ...             [2, 2, 2, 2, 2, 2, 2, 0, 0],
+        ...         ]
+        ...     )
+        ... ).all()
+    """
+
+    def __init__(
+        self,
+        total_dialog_turns: int,
+        topk_size: int,
+        prompt_key: NestedKey = "text",
+        rewards_key: NestedKey = ("next", "reward"),
+        done_key: NestedKey = ("next", "done"),
+        verbose: bool = True,
+    ):
+        super().__init__()
+        self.in_keys = [prompt_key, rewards_key, done_key]
+        self.prompt_key = prompt_key
+        self.rewards_key = rewards_key
+        self.done_key = done_key
+        self.queues = defaultdict(lambda: deque(maxlen=total_dialog_turns))
+        self.total_dialog_turns = total_dialog_turns
+        self.topk_size = topk_size
+        if topk_size > total_dialog_turns:
+            raise ValueError(
+                f"topk_size must be smaller than or equal to total_dialog_turns, got {topk_size=} and {total_dialog_turns=}"
+            )
+        self.verbose = verbose
+
+    def forward(self, tensordict: TensorDictBase) -> Any:
+        return tensordict
+
+    def _inv_call(self, tensordict: TensorDictBase) -> TensorDictBase:
+        # Tensordict can be any number of dims, but it must contain entire trajectories
+        if tensordict.ndim == 1:
+            # Check how many done states we have
+            num_done = tensordict[self.done_key].sum()
+            if num_done > 1:
+                done_idx = tensordict[self.done_key].nonzero(as_tuple=True)[0] + 1
+                splits = torch.cat([done_idx.new_zeros((1,)), done_idx], dim=0).diff()
+                tensordicts = tensordict.split(splits)
+                tensordicts = [self._inv_call(td) for td in tensordicts]
+                tensordicts = [td for td in tensordicts if td is not None]
+                return torch.cat(tensordicts) if tensordicts else None
+            # Then we have a single trajectory
+            if not tensordict[-1][self.done_key].all():
+                raise RuntimeError("Expected the trajectory to be done.")
+            prompt = tensordict[0][self.prompt_key]
+            if not isinstance(prompt, str):
+                raise TypeError(f"Expected a string as prompt, got {type(prompt)=}")
+            self.queues[prompt].append(tensordict)
+            if len(self.queues[prompt]) == self.total_dialog_turns:
+                if self.verbose:
+                    torchrl_logger.info(f"Getting top-k rewards for {prompt=}")
+                # Cat is the most robust way to combine the trajs
+                tds = torch.cat(list(self.queues[prompt]), -1)
+                # Collect rewards
+                reward = tds.get(self.rewards_key, as_nested_tensor=True)
+                reward = self._aggregate_rewards(reward)
+                # Check if all rewards are equal
+                if (reward == reward[0]).all():
+                    # If all rewards are equal, we can't select top-k
+                    if self.verbose:
+                        torchrl_logger.warning(
+                            f"All rewards are equal ({reward.unique()=})"
+                        )
+                    return
+                # Filter out rewards below median
+                median_reward = reward.median(dim=-1, keepdim=True)[0]
+                mask = reward > median_reward
+                filtered_reward = reward[mask]
+                filtered_indices = mask.nonzero(as_tuple=True)[0]
+                # Get top-k from filtered rewards
+                topk_reward = filtered_reward.topk(
+                    k=min(self.topk_size, len(filtered_indices)), dim=-1
+                )
+                if not topk_reward.indices.numel():
+                    if self.verbose:
+                        torchrl_logger.warning(
+                            f"No top-{self.topk_size} rewards found ({reward=})"
+                        )
+                    return
+                # Map back to original indices
+                selected_indices = filtered_indices[topk_reward.indices]
+                tds = tds[selected_indices]
+                if self.verbose:
+                    torchrl_logger.info(
+                        f"Selected top-{self.topk_size} rewards, with reward {topk_reward.values=}"
+                    )
+                return tds
+            return
+        elif tensordict.ndim > 2:
+            # keep the time dim at the end
+            tensordict = tensordict.flatten(0, -2)
+        trajs = tensordict.unbind(-1)
+        # Iterate over the trajectories
+        result = []
+        for traj in trajs:
+            td_out = self._inv_call(traj)
+            if td_out is None:
+                continue
+            result.append(td_out)
+        if result:
+            return torch.cat(result, -1)
+        return
+
+    def _aggregate_rewards(self, reward: torch.Tensor) -> torch.Tensor:
+        """Aggregate the rewards across the dialog turns.
+
+        `reward` is expected to be a nested tensor.
+
+        The default implementation is to take the mean of the rewards across the dialog turns.
+        """
+        # reward = reward.to_padded_tensor(padding=0.0)
+        if reward.ndim < 2 or reward.ndim > 3:
+            raise ValueError(
+                f"Expected reward to be a 2D or 3D tensor, got {reward.ndim}D tensor"
+            )
+        return reward.mean(dim=-2).squeeze(-1)

torchrl/data/replay_buffers/ray_buffer.py

@@ -54,9 +54,12 @@ class RayReplayBuffer(ReplayBuffer):
     """A Ray implementation of the Replay Buffer that can be extended and sampled remotely.
 
     Keyword Args:
+        replay_buffer_cls (type[ReplayBuffer], optional): the class to use for the replay buffer.
+            Defaults to :class:`~torchrl.data.ReplayBuffer`.
         ray_init_config (dict[str, Any], optiona): keyword arguments to pass to `ray.init()`.
         remote_config (dict[str, Any], optiona): keyword arguments to pass to `cls.as_remote()`.
             Defaults to `torchrl.collectors.distributed.ray.DEFAULT_REMOTE_CLASS_CONFIG`.
+        **kwargs: keyword arguments to pass to the replay buffer class.
 
     .. seealso:: :class:`~torchrl.data.ReplayBuffer` for a list of other keyword arguments.
 
@@ -119,6 +122,7 @@ class RayReplayBuffer(ReplayBuffer):
     def __init__(
         self,
         *args,
+        replay_buffer_cls: type[ReplayBuffer] | None = ReplayBuffer,
         ray_init_config: dict[str, Any] | None = None,
         remote_config: dict[str, Any] | None = None,
         **kwargs,
@@ -134,7 +138,13 @@ class RayReplayBuffer(ReplayBuffer):
                 ray_init_config = DEFAULT_RAY_INIT_CONFIG
             ray.init(**ray_init_config)
 
-        remote_cls = ReplayBuffer.as_remote(remote_config).remote
+        remote_cls = replay_buffer_cls.as_remote(remote_config).remote
+        # We can detect if the buffer has a GPU allocated, if not
+        #  we'll make sure that the data is sent to CPU when needed.
+        if remote_config is not None:
+            self.has_gpu = remote_config.get("num_gpus", 0) > 0
+        else:
+            self.has_gpu = False
         self._rb = remote_cls(*args, **kwargs)
 
     def close(self):
@@ -158,6 +168,10 @@ class RayReplayBuffer(ReplayBuffer):
         return ray.get(pending_task)
 
     def extend(self, *args, **kwargs):
+        if not self.has_gpu:
+            # Move the data to GPU
+            args = [arg.to("cpu") for arg in args if hasattr(arg, "to")]
+            kwargs = {k: v.to("cpu") for k, v in kwargs.items() if hasattr(v, "to")}
         pending_task = self._rb.extend.remote(*args, **kwargs)
         return ray.get(pending_task)