[BugFix] Fix auto-batch-size

Author: Vincent Moens

torchrl/data/datasets/d4rl.py

@@ -291,15 +291,18 @@ def _get_dataset_direct(self, name, env_kwargs):
                     k: torch.from_numpy(item)
                     for k, item in dataset.items()
                     if isinstance(item, np.ndarray)
-                }
+                },
+                auto_batch_size=True,
             )
         dataset = dataset.unflatten_keys("/")
         if "metadata" in dataset.keys():
             metadata = dataset.get("metadata")
             dataset = dataset.exclude("metadata")
             self.metadata = metadata
             # find batch size
-            dataset = make_tensordict(dataset.flatten_keys("/").to_dict())
+            dataset = make_tensordict(
+                dataset.flatten_keys("/").to_dict(), auto_batch_size=True
+            )
             dataset = dataset.unflatten_keys("/")
         else:
             self.metadata = {}
@@ -361,7 +364,8 @@ def _get_dataset_from_env(self, name, env_kwargs):
                     k: torch.from_numpy(item)
                     for k, item in env.get_dataset().items()
                     if isinstance(item, np.ndarray)
-                }
+                },
+                auto_batch_size=True,
             )
         dataset = dataset.unflatten_keys("/")
         dataset = self._process_data_from_env(dataset, env)
@@ -373,7 +377,9 @@ def _process_data_from_env(self, dataset, env=None):
             dataset = dataset.exclude("metadata")
             self.metadata = metadata
             # find batch size
-            dataset = make_tensordict(dataset.flatten_keys("/").to_dict())
+            dataset = make_tensordict(
+                dataset.flatten_keys("/").to_dict(), auto_batch_size=True
+            )
             dataset = dataset.unflatten_keys("/")
         else:
             self.metadata = {}

torchrl/data/datasets/gen_dgrl.py

@@ -272,7 +272,7 @@ def _unpack_category_file(
                     npybuffer = tar.extractfile(member=member)
                     # npyfile = Path(download_folder) / member.name
                     npfile = np.load(npybuffer, allow_pickle=True)
-                    td = TensorDict.from_dict(npfile.tolist())
+                    td = TensorDict.from_dict(npfile.tolist(), auto_batch_size=True)
                     td.set("observations", td.get("observations").to(torch.uint8))
                     td.set(("next", "observation"), td.get("observations")[1:])
                     td.set("observations", td.get("observations")[:-1])

torchrl/data/datasets/openx.py

@@ -787,4 +787,4 @@ def _make_tensordict_image_conv(data):
         data["observation"]["image"] = tensor
     except KeyError:
         pass
-    return make_tensordict(data)
+    return make_tensordict(data, auto_batch_size=True)

torchrl/data/datasets/vd4rl.py

@@ -418,7 +418,7 @@ def _is_downloaded(self):
 def _from_npz(npz_path):
     npz = np.load(npz_path)
     npz_dict = {file: npz[file] for file in npz.files}
-    return TensorDict.from_dict(npz_dict)
+    return TensorDict.from_dict(npz_dict, auto_batch_size=True)
 
 
 _NAME_MATCH = KeyDependentDefaultDict(lambda x: x)

torchrl/data/rlhf/dataset.py

@@ -148,9 +148,10 @@ def load(self):
         dataset = self._load_dataset()
         dataset = self._tokenize(dataset)
         prefix = (split, str(max_length))
-        return self.dataset_to_tensordict(
+        result = self.dataset_to_tensordict(
             dataset, data_dir=data_dir, prefix=prefix, valid_mask_key="valid_sample"
-        )[prefix]
+        )
+        return result[prefix]
 
     def _load_dataset(self):
         """Loads a text dataset from ``datasets``.
@@ -213,7 +214,9 @@ def _tokenize(
                     for key, value in dataset_dict.items()
                     if key not in excluded_features
                 }
-            dataset = TensorDict.from_dict(dataset_dict)
+            dataset = TensorDict.from_dict(
+                dataset_dict, auto_batch_size=True, batch_dims=1
+            )
         elif excluded_features:
             dataset = dataset.exclude(*excluded_features)
         # keep non empty rows (i.e. where at least one token is not eos)
@@ -294,14 +297,16 @@ def dataset_to_tensordict(
             if prefix is None:
                 prefix = ()
             data_dict = {key: torch.as_tensor(dataset[key]) for key in features}
-            out = TensorDict.from_dict(data_dict, batch_dims=batch_dims)
+            out = TensorDict.from_dict(
+                data_dict, batch_dims=batch_dims, auto_batch_size=True
+            )
         else:
             out = dataset
         if valid_mask_key is not None and valid_mask_key in out.keys(
             include_nested=True
         ):
             out = out[out.get(valid_mask_key)]
-        out = TensorDict({prefix: out}, [])
+        out = TensorDict({prefix: out})
         out.memmap_(prefix=data_dir)
         return out
 
@@ -481,6 +486,9 @@ def __call__(self, sample):
         batch_size = [] if isinstance(input, str) else [len(input)]
         if self.return_tensordict:
             return TensorDict.from_dict(
-                dict(tokenized_sample), batch_size=batch_size, device=self.device
+                dict(tokenized_sample),
+                batch_size=batch_size,
+                device=self.device,
+                auto_batch_size=True,
             )
         return tokenized_sample