[BugFix] RB.add unsqueezes tds when applying the transform (#3047)

Author: vmoens

test/test_transforms.py

@@ -47,6 +47,7 @@
     Composite,
     LazyTensorStorage,
     NonTensor,
+    RandomSampler,
     ReplayBuffer,
     TensorDictReplayBuffer,
     TensorSpec,
@@ -13271,6 +13272,52 @@ def test_multistep_transform_changes(self):
             assert rb[:]["next", "steps"][-1] == data["steps"][-1]
             assert t._buffer["steps"][-1] == data["steps"][-1]
 
+    @pytest.mark.parametrize("add_or_extend", ["add", "extend"])
+    def test_multisteptransform_single_item(self, add_or_extend):
+        # Configuration
+        buffer_size = 1000
+        n_step = 3
+        gamma = 0.99
+        device = "cpu"
+
+        rb = ReplayBuffer(
+            storage=LazyTensorStorage(max_size=buffer_size, device=device, ndim=1),
+            sampler=RandomSampler(),
+            transform=MultiStepTransform(n_steps=n_step, gamma=gamma),
+        )
+        obs_dict = lambda i: {"observation": torch.full((4,), i)}  # 4-dim observation
+        next_obs_dict = lambda i: {"observation": torch.full((4,), i)}
+
+        for i in range(10):
+            # Create transition with batch_size=[] (no batch dimension)
+            transition = TensorDict(
+                {
+                    "obs": TensorDict(obs_dict(i), batch_size=[]),
+                    "action": torch.full((2,), i),  # 2-dim action
+                    "next": TensorDict(
+                        {
+                            "obs": TensorDict(next_obs_dict(i), batch_size=[]),
+                            "done": torch.tensor(False, dtype=torch.bool),
+                            "reward": torch.tensor(float(i), dtype=torch.float32),
+                        },
+                        batch_size=[],
+                    ),
+                },
+                batch_size=[],
+            )
+
+            if add_or_extend == "add":
+                rb.add(transition)
+            else:
+                rb.extend(transition.unsqueeze(0))
+        rbcontent = rb[:]
+        assert (rbcontent["steps_to_next_obs"] == 3).all()
+        assert rbcontent.shape == (7,)
+        assert (rbcontent["next", "original_reward"] == torch.arange(7)).all()
+        assert (
+            rbcontent["next", "reward"] > rbcontent["next", "original_reward"]
+        ).all()
+
 
 class TestBatchSizeTransform(TransformBase):
     class MyEnv(EnvBase):

torchrl/data/replay_buffers/replay_buffers.py

@@ -702,7 +702,18 @@ def add(self, data: Any) -> int:
         """
         if self._transform is not None and len(self._transform):
             with _set_dispatch_td_nn_modules(is_tensor_collection(data)):
-                data = self._transform.inv(data)
+                make_none = False
+                # Transforms usually expect a time batch dimension when called within a RB, so we unsqueeze the data temporarily
+                is_tc = is_tensor_collection(data)
+                with data.unsqueeze(-1) if is_tc else contextlib.nullcontext(data) as data_unsq:
+                    data_unsq_r = self._transform.inv(data_unsq)
+                    if is_tc and data_unsq_r is not None:
+                        # this is a no-op whenever the result matches the input
+                        data_unsq.update(data_unsq_r)
+                    else:
+                        make_none = data_unsq_r is None
+                if make_none:
+                    data = None
         if data is None:
             return torch.zeros((0, self._storage.ndim), dtype=torch.long)
         return self._add(data)