Repository: pytorch/rl

Issue:
[Feature Request] Replay buffer support for storing and sampling whole trajectories
## Motivation

I was trying to work with recurrent models (for example implementing https://arxiv.org/pdf/2110.05038.pdf) and got a problem trying to store and sample whole trajectories in the replay buffer instead of tuples. When training on-policy algorithm (without replay buffer) it is possible to do so using BatchSubSampler but I couldn't find a proper way to do it with off-policy algorithms.

## Solution

The main obstacle is the following lines in the ReplayBufferTrainer class:
![image](https://user-images.githubusercontent.com/49375140/190733842-5f545dd2-c233-4467-963f-3d921def3ba0.png)
It flattens the batch of trajectories into a batch of tuples. Maybe add a flag that controls this behavior?

## Checklist

- [x] I have checked that there is no similar issue in the repo (**required**)



Discussion:
Good point this helper class could have a flag to store trajectories.
Feel free to make a PR for this. 

Otherwise I'll take care of it on Tuesday when I'm back to work!

Resolving patch:
diff --git a/torchrl/trainers/trainers.py b/torchrl/trainers/trainers.py
--- a/torchrl/trainers/trainers.py
+++ b/torchrl/trainers/trainers.py
@@ -30,7 +30,7 @@
     TensorDictPrioritizedReplayBuffer,
     TensorDictReplayBuffer,
 )
-from torchrl.data.tensordict.tensordict import TensorDictBase
+from torchrl.data.tensordict.tensordict import TensorDictBase, pad
 from torchrl.data.utils import expand_right, DEVICE_TYPING
 from torchrl.envs.common import EnvBase
 from torchrl.envs.transforms import TransformedEnv
@@ -532,6 +532,17 @@ class ReplayBufferTrainer:
             Default is False.
         device (device, optional): device where the samples must be placed.
             Default is cpu.
+        flatten_tensordicts (bool, optional): if True, the tensordicts will be
+            flattened (or equivalently masked with the valid mask obtained from
+            the collector) before being passed to the replay buffer. Otherwise,
+            no transform will be achieved other than padding (see `max_dims` arg below).
+            Defaults to True
+        max_dims (sequence of int, optional): if `flatten_tensordicts` is set to False,
+            this will be a list of the length of the batch_size of the provided
+            tensordicts that represent the maximum size of each. If provided,
+            this list of sizes will be used to pad the tensordict and make their shape
+            match before they are passed to the replay buffer. If there is no
+            maximum value, a -1 value should be provided.
 
     Examples:
         >>> rb_trainer = ReplayBufferTrainer(replay_buffer=replay_buffer, batch_size=N)
@@ -547,17 +558,33 @@ def __init__(
         batch_size: int,
         memmap: bool = False,
         device: DEVICE_TYPING = "cpu",
+        flatten_tensordicts: bool = True,
+        max_dims: Optional[Sequence[int]] = None,
     ) -> None:
         self.replay_buffer = replay_buffer
         self.batch_size = batch_size
         self.memmap = memmap
         self.device = device
+        self.flatten_tensordicts = flatten_tensordicts
+        self.max_dims = max_dims
 
     def extend(self, batch: TensorDictBase) -> TensorDictBase:
-        if "mask" in batch.keys():
-            batch = batch[batch.get("mask").squeeze(-1)]
+        if self.flatten_tensordicts:
+            if "mask" in batch.keys():
+                batch = batch[batch.get("mask").squeeze(-1)]
+            else:
+                batch = batch.reshape(-1)
         else:
-            batch = batch.reshape(-1)
+            if self.max_dims is not None:
+                pads = []
+                for d in range(batch.ndimension()):
+                    pad_value = (
+                        0
+                        if self.max_dims[d] == -1
+                        else self.max_dims[d] - batch.batch_size[d]
+                    )
+                    pads += [0, pad_value]
+                batch = pad(batch, pads)
         # reward_training = batch.get("reward").mean().item()
         batch = batch.cpu()
         if self.memmap:


Test patch: