amend d4rl

pytorch · Nov 29, 2023 · 2de0247 · 2de0247
1 parent ae10b94
commit 2de0247
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 3 deletions.
diff --git a/torchrl/data/datasets/d4rl.py b/torchrl/data/datasets/d4rl.py
@@ -152,6 +152,10 @@ def __init__(
         self.name = name
         dataset = None
 
+        if not from_env and direct_download is None:
+            self._import_d4rl()
+            direct_download = not self._has_d4rl
+
         if not direct_download:
             if from_env is None:
                 warnings.warn(
@@ -169,9 +173,6 @@ def __init__(
                 from_env = False
             self.from_env = from_env
 
-        if not from_env and direct_download is None:
-            self._import_d4rl()
-            direct_download = not self._has_d4rl
         if download and not self._is_downloaded():
             if not direct_download:
                 if terminate_on_end is None:

diff --git a/torchrl/data/datasets/minari.py b/torchrl/data/datasets/minari.py
@@ -28,6 +28,16 @@
 
 
 class MinariExperienceReplay(TensorDictReplayBuffer):
+    """Minari Experience replay dataset.
+
+    Args:
+        dataset_id (str):
+        batch_size (int):
+
+    Keyword Args:
+        root (Path or str, optional):
+        download (bool, optional):
+    """
     def __init__(
         self,
         dataset_id,