[Feature] D4rl direct download (#1430)

MateuszGuzek · vmoens · web-flow · commit 301881091be6 · 2023-10-04T17:50:26.000+01:00
Co-authored-by: Mateusz Guzek &lt;matguzek@meta.com&gt;
Co-authored-by: vmoens &lt;vincentmoens@gmail.com&gt;
diff --git a/.github/unittest/linux_libs/scripts_d4rl/run_test.sh b/.github/unittest/linux_libs/scripts_d4rl/run_test.sh
@@ -40,3 +40,22 @@ python -c "import gym, d4rl"
 python .github/unittest/helpers/coverage_run_parallel.py -m pytest test/test_libs.py --instafail -v --durations 200 --capture no -k TestD4RL --error-for-skips
 coverage combine
 coverage xml -i
+
+## check what happens if we update gym
+#pip install gym -U
+#python -c """
+#from torchrl.data.datasets import D4RLExperienceReplay
+#data = D4RLExperienceReplay('halfcheetah-medium-v2', batch_size=10, from_env=False, direct_download=True)
+#for batch in data:
+#    print(batch)
+#    break
+#
+#data = D4RLExperienceReplay('halfcheetah-medium-v2', batch_size=10, from_env=False, direct_download=False)
+#for batch in data:
+#    print(batch)
+#    break
+#
+#import d4rl
+#import gym
+#gym.make('halfcheetah-medium-v2')
+#"""
diff --git a/test/test_libs.py b/test/test_libs.py
@@ -1775,7 +1775,7 @@ class TestD4RL:
     def test_terminate_on_end(self, task, use_truncated_as_done, split_trajs):
 
         with pytest.warns(
-            UserWarning, match="Using terminate_on_end=True with from_env=False"
+            UserWarning, match="Using use_truncated_as_done=True"
         ) if use_truncated_as_done else nullcontext():
             data_true = D4RLExperienceReplay(
                 task,
@@ -1823,6 +1823,37 @@ def test_terminate_on_end(self, task, use_truncated_as_done, split_trajs):
             ]
             assert "truncated" not in leaf_names
 
+    @pytest.mark.parametrize("task", ["walker2d-medium-replay-v2"])
+    def test_direct_download(self, task):
+        data_direct = D4RLExperienceReplay(
+            task,
+            split_trajs=False,
+            from_env=False,
+            batch_size=2,
+            use_truncated_as_done=True,
+            direct_download=True,
+        )
+        data_d4rl = D4RLExperienceReplay(
+            task,
+            split_trajs=False,
+            from_env=True,
+            batch_size=2,
+            use_truncated_as_done=True,
+            direct_download=False,
+            terminate_on_end=True,  # keep the last time step
+        )
+        keys = set(data_direct._storage._storage.keys(True, True))
+        keys = keys.intersection(data_d4rl._storage._storage.keys(True, True))
+        assert len(keys)
+        assert_allclose_td(
+            data_direct._storage._storage.select(*keys).apply(
+                lambda t: t.as_tensor().float()
+            ),
+            data_d4rl._storage._storage.select(*keys).apply(
+                lambda t: t.as_tensor().float()
+            ),
+        )
+
     @pytest.mark.parametrize(
         "task",
         [
diff --git a/torchrl/data/datasets/d4rl.py b/torchrl/data/datasets/d4rl.py
@@ -2,15 +2,22 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
+from __future__ import annotations
+
+import os
+import urllib
 import warnings
-from typing import Callable, Optional
+from typing import Callable
 
 import numpy as np
 
 import torch
+
+from tensordict import PersistentTensorDict
 from tensordict.tensordict import make_tensordict
 
 from torchrl.collectors.utils import split_trajectories
+from torchrl.data.datasets.d4rl_infos import D4RL_DATASETS
 from torchrl.data.replay_buffers import TensorDictReplayBuffer
 from torchrl.data.replay_buffers.samplers import Sampler
 from torchrl.data.replay_buffers.storages import LazyMemmapStorage
@@ -75,18 +82,25 @@ class D4RLExperienceReplay(TensorDictReplayBuffer):
               differ. In particular, the ``"truncated"`` key (used to determine the
               end of an episode) may be absent when ``from_env=False`` but present
               otherwise, leading to a different slicing when ``traj_splits`` is enabled.
-
+        direct_download (bool): if ``True``, the data will be downloaded without
+            requiring D4RL. If ``None``, if ``d4rl`` is present in the env it will
+            be used to download the dataset, otherwise the download will fall back
+            on ``direct_download=True``.
+            This is not compatible with ``from_env=True``.
+            Defaults to ``None``.
         use_truncated_as_done (bool, optional): if ``True``, ``done = terminated | truncated``.
             Otherwise, only the ``terminated`` key is used. Defaults to ``True``.
+        terminate_on_end (bool, optional): Set ``done=True`` on the last timestep
+            in a trajectory. Default is ``False``, and will discard the
+            last timestep in each trajectory.
         **env_kwargs (key-value pairs): additional kwargs for
-            :func:`d4rl.qlearning_dataset`. Supports ``terminate_on_end``
-            (``False`` by default) or other kwargs if defined by D4RL library.
+            :func:`d4rl.qlearning_dataset`.
 
 
     Examples:
         >>> from torchrl.data.datasets.d4rl import D4RLExperienceReplay
         >>> from torchrl.envs import ObservationNorm
-        >>> data = D4RLExperienceReplay("maze2d-umaze-v1")
+        >>> data = D4RLExperienceReplay("maze2d-umaze-v1", 128)
         >>> # we can append transforms to the dataset
         >>> data.append_transform(ObservationNorm(loc=-1, scale=1.0))
         >>> data.sample(128)
@@ -109,34 +123,63 @@ def __init__(
         self,
         name,
         batch_size: int,
-        sampler: Optional[Sampler] = None,
-        writer: Optional[Writer] = None,
-        collate_fn: Optional[Callable] = None,
+        sampler: Sampler | None = None,
+        writer: Writer | None = None,
+        collate_fn: Callable | None = None,
         pin_memory: bool = False,
-        prefetch: Optional[int] = None,
-        transform: Optional["Transform"] = None,  # noqa-F821
+        prefetch: int | None = None,
+        transform: "torchrl.envs.Transform" | None = None,  # noqa-F821
         split_trajs: bool = False,
-        from_env: bool = True,
+        from_env: bool = None,
         use_truncated_as_done: bool = True,
+        direct_download: bool = None,
+        terminate_on_end: bool = None,
         **env_kwargs,
     ):
-
-        type(self)._import_d4rl()
-
-        if not self._has_d4rl:
-            raise ImportError("Could not import d4rl") from self.D4RL_ERR
+        if from_env is None:
+            warnings.warn(
+                "from_env will soon default to ``False``, ie the data will be "
+                "downloaded without relying on d4rl by default. "
+                "For now, ``True`` will still be the default. "
+                "To disable this warning, explicitly pass the ``from_env`` argument "
+                "during construction of the dataset.",
+                category=DeprecationWarning,
+            )
+            from_env = True
         self.from_env = from_env
         self.use_truncated_as_done = use_truncated_as_done
-        if from_env:
-            dataset = self._get_dataset_from_env(name, env_kwargs)
+
+        if not from_env and direct_download is None:
+            self._import_d4rl()
+            direct_download = not self._has_d4rl
+
+        if not direct_download:
+            if terminate_on_end is None:
+                # we use the default of d4rl
+                terminate_on_end = False
+            self._import_d4rl()
+
+            if not self._has_d4rl:
+                raise ImportError("Could not import d4rl") from self.D4RL_ERR
+
+            if from_env:
+                dataset = self._get_dataset_from_env(name, env_kwargs)
+            else:
+                if self.use_truncated_as_done:
+                    warnings.warn(
+                        "Using use_truncated_as_done=True + terminate_on_end=True "
+                        "with from_env=False may not have the intended effect "
+                        "as the timeouts (truncation) "
+                        "can be absent from the static dataset."
+                    )
+                env_kwargs.update({"terminate_on_end": terminate_on_end})
+                dataset = self._get_dataset_direct(name, env_kwargs)
         else:
-            if self.use_truncated_as_done:
-                warnings.warn(
-                    "Using terminate_on_end=True with from_env=False "
-                    "may not have the intended effect as the timeouts (truncation) "
-                    "can be absent from the static dataset."
+            if terminate_on_end is False:
+                raise ValueError(
+                    "Using terminate_on_end=False is not compatible with direct_download=True."
                 )
-            dataset = self._get_dataset_direct(name, env_kwargs)
+            dataset = self._get_dataset_direct_download(name, env_kwargs)
         # Fill unknown next states with 0
         dataset["next", "observation"][dataset["next", "done"].squeeze()] = 0
 
@@ -157,6 +200,23 @@ def __init__(
         )
         self.extend(dataset)
 
+    def _get_dataset_direct_download(self, name, env_kwargs):
+        """Directly download and use a D4RL dataset."""
+        if env_kwargs:
+            raise RuntimeError(
+                f"Cannot pass env_kwargs when `direct_download=True`. Got env_kwargs keys: {env_kwargs.keys()}"
+            )
+        url = D4RL_DATASETS.get(name, None)
+        if url is None:
+            raise KeyError(f"Env {name} not found.")
+        h5path = _download_dataset_from_url(url)
+        # h5path_parent = Path(h5path).parent
+        dataset = PersistentTensorDict.from_h5(h5path)
+        dataset = dataset.to_tensordict()
+        with dataset.unlock_():
+            dataset = self._process_data_from_env(dataset)
+        return dataset
+
     def _get_dataset_direct(self, name, env_kwargs):
         from torchrl.envs.libs.gym import GymWrapper
 
@@ -247,6 +307,10 @@ def _get_dataset_from_env(self, name, env_kwargs):
             }
         )
         dataset = dataset.unflatten_keys("/")
+        dataset = self._process_data_from_env(dataset, env)
+        return dataset
+
+    def _process_data_from_env(self, dataset, env=None):
         if "metadata" in dataset.keys():
             metadata = dataset.get("metadata")
             dataset = dataset.exclude("metadata")
@@ -277,10 +341,11 @@ def _get_dataset_from_env(self, name, env_kwargs):
             pass
 
         # let's make sure that the dtypes match what's expected
-        for key, spec in env.observation_spec.items(True, True):
-            dataset[key] = dataset[key].to(spec.dtype)
-        dataset["action"] = dataset["action"].to(env.action_spec.dtype)
-        dataset["reward"] = dataset["reward"].to(env.reward_spec.dtype)
+        if env is not None:
+            for key, spec in env.observation_spec.items(True, True):
+                dataset[key] = dataset[key].to(spec.dtype)
+            dataset["action"] = dataset["action"].to(env.action_spec.dtype)
+            dataset["reward"] = dataset["reward"].to(env.reward_spec.dtype)
 
         # format done
         dataset["done"] = dataset["done"].bool().unsqueeze(-1)
@@ -300,7 +365,10 @@ def _get_dataset_from_env(self, name, env_kwargs):
             dataset.clone()
         )  # make sure that all tensors have a different data_ptr
         self._shift_reward_done(dataset)
-        self.specs = env.specs.clone()
+        if env is not None:
+            self.specs = env.specs.clone()
+        else:
+            self.specs = None
         return dataset
 
     def _shift_reward_done(self, dataset):
@@ -313,3 +381,39 @@ def _shift_reward_done(self, dataset):
             dataset[key] = dataset[key].clone()
             dataset[key][1:] = dataset[key][:-1].clone()
             dataset[key][0] = 0
+
+
+def _download_dataset_from_url(dataset_url):
+    dataset_filepath = _filepath_from_url(dataset_url)
+    if not os.path.exists(dataset_filepath):
+        print("Downloading dataset:", dataset_url, "to", dataset_filepath)
+        urllib.request.urlretrieve(dataset_url, dataset_filepath)
+    if not os.path.exists(dataset_filepath):
+        raise IOError("Failed to download dataset from %s" % dataset_url)
+    return dataset_filepath
+
+
+def _filepath_from_url(dataset_url):
+    _, dataset_name = os.path.split(dataset_url)
+    dataset_filepath = os.path.join(DATASET_PATH, dataset_name)
+    return dataset_filepath
+
+
+def _set_dataset_path(path):
+    global DATASET_PATH
+    DATASET_PATH = path
+    os.makedirs(path, exist_ok=True)
+
+
+_set_dataset_path(
+    os.environ.get(
+        "D4RL_DATASET_DIR", os.path.expanduser("~/.cache/torchrl/data/d4rl/datasets")
+    )
+)
+
+if __name__ == "__main__":
+    data = D4RLExperienceReplay("kitchen-partial-v0", batch_size=128)
+    print(data)
+    for sample in data:
+        print(sample)
+        break
diff --git a/torchrl/data/datasets/d4rl_infos.py b/torchrl/data/datasets/d4rl_infos.py
