Update

[ghstack-poisoned]

Author: Vincent Moens

torchrl/data/replay_buffers/storages.py

@@ -1099,12 +1099,19 @@ def max_size_along_dim0(data_shape):
 
         if is_tensor_collection(data):
             out = data.to(self.device)
-            if self.empty_lazy and shape is None:
-                raise RuntimeError(
-                    "Make sure you have called `extend` and not `add` first when setting `empty_lazy=True`."
+            if self.empty_lazy:
+                if shape is None:
+                    # shape is None in add
+                    raise RuntimeError(
+                        "Make sure you have called `extend` and not `add` first when setting `empty_lazy=True`."
+                    )
+                out: TensorDictBase = torch.empty_like(
+                    out.expand(max_size_along_dim0(data.shape))
                 )
             elif shape is None:
                 shape = data.shape
+            else:
+                out = out[0]
             out: TensorDictBase = out.new_empty(
                 max_size_along_dim0(shape), empty_lazy=self.empty_lazy
             )

torchrl/envs/transforms/transforms.py

@@ -7308,12 +7308,13 @@ def transform_reward_spec(self, reward_spec: TensorSpec) -> TensorSpec:
         return reward_spec
 
     def forward(self, tensordict: TensorDictBase) -> TensorDictBase:
-        time_dim = [i for i, name in enumerate(tensordict.names) if name == "time"]
-        if not time_dim:
+        try:
+            time_dim = list(tensordict.names).index("time")
+        except ValueError:
             raise ValueError(
                 "At least one dimension of the tensordict must be named 'time' in offline mode"
             )
-        time_dim = time_dim[0] - 1
+        time_dim = time_dim - 1
         for in_key, out_key in _zip_strict(self.in_keys, self.out_keys):
             reward = tensordict[in_key]
             cumsum = reward.cumsum(time_dim)