torchrl/data/datasets/d4rl.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# 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 importlib
import os
import shutil
import tempfile
import urllib
import warnings

from pathlib import Path
from typing import Callable

import numpy as np

import torch

from tensordict import make_tensordict, PersistentTensorDict, TensorDict

from torchrl._utils import logger as torchrl_logger

from torchrl.collectors.utils import split_trajectories
from torchrl.data.datasets.common import BaseDatasetExperienceReplay
from torchrl.data.datasets.d4rl_infos import D4RL_DATASETS

from torchrl.data.datasets.utils import _get_root_dir
from torchrl.data.replay_buffers.samplers import Sampler
from torchrl.data.replay_buffers.storages import TensorStorage
from torchrl.data.replay_buffers.writers import ImmutableDatasetWriter, Writer


class D4RLExperienceReplay(BaseDatasetExperienceReplay):
    """An Experience replay class for D4RL.

    To install D4RL, follow the instructions on the
    `official repo <https://github.com/Farama-Foundation/D4RL>`__.

    The data format follows the :ref:`TED convention <TED-format>`.
    The replay buffer contains the env specs under D4RLExperienceReplay.specs.

    If present, metadata will be written in ``D4RLExperienceReplay.metadata``
    and excluded from the dataset.

    The transitions are reconstructed using ``done = terminated | truncated`` and
    the ``("next", "observation")`` of ``"done"`` states are zeroed.

    Args:
        dataset_id (str): the dataset_id of the D4RL env to get the data from.
        batch_size (int): the batch size to use during sampling.
        sampler (Sampler, optional): the sampler to be used. If none is provided
            a default RandomSampler() will be used.
        writer (Writer, optional): the writer to be used. If none is provided
            a default :class:`~torchrl.data.replay_buffers.writers.ImmutableDatasetWriter` will be used.
        collate_fn (callable, optional): merges a list of samples to form a
            mini-batch of Tensor(s)/outputs.  Used when using batched
            loading from a map-style dataset.
        pin_memory (bool): whether pin_memory() should be called on the rb
            samples.
        prefetch (int, optional): number of next batches to be prefetched
            using multithreading.
        transform (Transform, optional): Transform to be executed when sample() is called.
            To chain transforms use the :class:`~torchrl.envs.transforms.transforms.Compose` class.
        split_trajs (bool, optional): if ``True``, the trajectories will be split
            along the first dimension and padded to have a matching shape.
            To split the trajectories, the ``"done"`` signal will be used, which
            is recovered via ``done = truncated | terminated``. In other words,
            it is assumed that any ``truncated`` or ``terminated`` signal is
            equivalent to the end of a trajectory. For some datasets from
            ``D4RL``, this may not be true. It is up to the user to make
            accurate choices regarding this usage of ``split_trajs``.
            Defaults to ``False``.
        from_env (bool, optional): if ``True``, :meth:`env.get_dataset` will
            be used to retrieve the dataset. Otherwise :func:`d4rl.qlearning_dataset`
            will be used. Defaults to ``True``.

            .. note::

              Using ``from_env=False`` will provide fewer data than ``from_env=True``.
              For instance, the info keys will be left out.
              Usually, ``from_env=False`` with ``terminate_on_end=True`` will
              lead to the same result as ``from_env=True``, with the latter
              containing meta-data and info entries that the former does
              not possess.

            .. note::

              The keys in ``from_env=True`` and ``from_env=False`` *may* unexpectedly
              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. This is to be used only with
            ``direct_download=False``.
        root (Path or str, optional): The D4RL dataset root directory.
            The actual dataset memory-mapped files will be saved under
            `<root>/<dataset_id>`. If none is provided, it defaults to
            ``~/.cache/torchrl/d4rl`.
        download (bool, optional): Whether the dataset should be downloaded if
            not found. Defaults to ``True``.
        **env_kwargs (key-value pairs): additional kwargs for
            :func:`d4rl.qlearning_dataset`.


    Examples:
        >>> from torchrl.data.datasets.d4rl import D4RLExperienceReplay
        >>> from torchrl.envs import ObservationNorm
        >>> data = D4RLExperienceReplay("maze2d-umaze-v1", 128)
        >>> # we can append transforms to the dataset
        >>> data.append_transform(ObservationNorm(loc=-1, scale=1.0, in_keys=["observation"]))
        >>> data.sample(128)

    """

    D4RL_ERR = None

    @classmethod
    def _import_d4rl(cls):
        cls._has_d4rl = importlib.util.find_spec("d4rl") is not None
        try:
            import d4rl  # noqa

        except ModuleNotFoundError as err:
            cls.D4RL_ERR = err
        except Exception:
            pass

    def __init__(
        self,
        dataset_id,
        batch_size: int,
        sampler: Sampler | None = None,
        writer: Writer | None = None,
        collate_fn: Callable | None = None,
        pin_memory: bool = False,
        prefetch: int | None = None,
        transform: "torchrl.envs.Transform" | None = None,  # noqa-F821
        split_trajs: bool = False,
        from_env: bool = False,
        use_truncated_as_done: bool = True,
        direct_download: bool = None,
        terminate_on_end: bool = None,
        download: bool = True,
        root: str | Path | None = None,
        **env_kwargs,
    ):
        self.use_truncated_as_done = use_truncated_as_done
        if root is None:
            root = _get_root_dir("d4rl")
        self.root = Path(root)
        self.dataset_id = dataset_id

        if not from_env and direct_download is None:
            self._import_d4rl()
            direct_download = not self._has_d4rl

        if not direct_download:
            warnings.warn(
                "You are using the D4RL library for collecting data. "
                "We advise against this use, as D4RL formatting can be "
                "inconsistent. "
                "To download the D4RL data without the D4RL library, use "
                "direct_download=True in the dataset constructor. "
                "Recurring to `direct_download=False` will soon be deprecated."
            )
            self.from_env = from_env
        else:
            self.from_env = from_env

        if (download == "force") or (download and not self._is_downloaded()):
            if download == "force" and os.path.exists(self.data_path_root):
                shutil.rmtree(self.data_path_root)

            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(dataset_id, 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(dataset_id, env_kwargs)
            else:
                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_download(dataset_id, env_kwargs)
            # Fill unknown next states with 0
            dataset["next", "observation"][dataset["next", "done"].squeeze()] = 0

            if split_trajs:
                dataset = split_trajectories(dataset)
                dataset["next", "done"][:, -1] = True

            storage = TensorStorage(dataset.memmap(self._dataset_path))
        elif self._is_downloaded():
            storage = TensorStorage(TensorDict.load_memmap(self._dataset_path))
        else:
            raise RuntimeError(
                f"The dataset could not be found in {self._dataset_path}."
            )

        if writer is None:
            writer = ImmutableDatasetWriter()
        super().__init__(
            batch_size=batch_size,
            storage=storage,
            sampler=sampler,
            writer=writer,
            collate_fn=collate_fn,
            pin_memory=pin_memory,
            prefetch=prefetch,
            transform=transform,
        )

    @property
    def data_path(self) -> Path:
        return self._dataset_path

    @property
    def data_path_root(self) -> Path:
        return self._dataset_path

    @property
    def _dataset_path(self):
        return Path(self.root) / self.dataset_id

    def _is_downloaded(self):
        return os.path.exists(self._dataset_path)

    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.")
        with tempfile.TemporaryDirectory() as tmpdir:
            os.environ["D4RL_DATASET_DIR"] = tmpdir
            h5path = _download_dataset_from_url(url, tmpdir)
            # 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

        type(self)._import_d4rl()

        if not self._has_d4rl:
            raise ImportError("Could not import d4rl") from self.D4RL_ERR
        import d4rl
        import gym

        env = GymWrapper(gym.make(name))
        with tempfile.TemporaryDirectory() as tmpdir:
            os.environ["D4RL_DATASET_DIR"] = tmpdir
            dataset = d4rl.qlearning_dataset(env._env, **env_kwargs)

            dataset = make_tensordict(
                {
                    k: torch.from_numpy(item)
                    for k, item in dataset.items()
                    if isinstance(item, np.ndarray)
                }
            )
        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 = dataset.unflatten_keys("/")
        else:
            self.metadata = {}
        dataset.rename_key_("observations", "observation")
        dataset.create_nested("next")
        dataset.rename_key_("next_observations", ("next", "observation"))
        dataset.rename_key_("terminals", "terminated")
        if "timeouts" in dataset.keys():
            dataset.rename_key_("timeouts", "truncated")
        if self.use_truncated_as_done:
            done = dataset.get("terminated") | dataset.get("truncated", False)
            dataset.set("done", done)
        else:
            dataset.set("done", dataset.get("terminated"))
        dataset.rename_key_("rewards", "reward")
        dataset.rename_key_("actions", "action")

        # 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["next", key] = dataset["next", key].to(spec.dtype)
        dataset["action"] = dataset["action"].to(env.action_spec.dtype)
        dataset["reward"] = dataset["reward"].to(env.reward_spec.dtype)

        # format done etc
        dataset["done"] = dataset["done"].bool().unsqueeze(-1)
        dataset["terminated"] = dataset["terminated"].bool().unsqueeze(-1)
        if "truncated" in dataset.keys():
            dataset["truncated"] = dataset["truncated"].bool().unsqueeze(-1)
        dataset["reward"] = dataset["reward"].unsqueeze(-1)
        dataset["next"].update(
            dataset.select("reward", "done", "terminated", "truncated", strict=False)
        )
        dataset = (
            dataset.clone()
        )  # make sure that all tensors have a different data_ptr
        self._shift_reward_done(dataset)
        self.specs = env.specs.clone()
        return dataset

    def _get_dataset_from_env(self, name, env_kwargs):
        """Creates an environment and retrieves the dataset using env.get_dataset().

        This method does not accept extra arguments.

        """
        if env_kwargs:
            raise RuntimeError("env_kwargs cannot be passed with using from_env=True")
        import gym

        # we do a local import to avoid circular import issues
        from torchrl.envs.libs.gym import GymWrapper

        with tempfile.TemporaryDirectory() as tmpdir:
            os.environ["D4RL_DATASET_DIR"] = tmpdir
            env = GymWrapper(gym.make(name))
            dataset = make_tensordict(
                {
                    k: torch.from_numpy(item)
                    for k, item in env.get_dataset().items()
                    if isinstance(item, np.ndarray)
                }
            )
        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")
            self.metadata = metadata
            # find batch size
            dataset = make_tensordict(dataset.flatten_keys("/").to_dict())
            dataset = dataset.unflatten_keys("/")
        else:
            self.metadata = {}

        dataset.rename_key_("observations", "observation")
        dataset.rename_key_("terminals", "terminated")
        if "timeouts" in dataset.keys():
            dataset.rename_key_("timeouts", "truncated")
        if self.use_truncated_as_done:
            dataset.set(
                "done",
                dataset.get("terminated") | dataset.get("truncated", False),
            )
        else:
            dataset.set("done", dataset.get("terminated"))

        dataset.rename_key_("rewards", "reward")
        dataset.rename_key_("actions", "action")
        try:
            dataset.rename_key_("infos", "info")
        except KeyError:
            pass

        # let's make sure that the dtypes match what's expected
        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)
        dataset["terminated"] = dataset["terminated"].bool().unsqueeze(-1)
        if "truncated" in dataset.keys():
            dataset["truncated"] = dataset["truncated"].bool().unsqueeze(-1)

        dataset["reward"] = dataset["reward"].unsqueeze(-1)
        if "next_observations" in dataset.keys():
            dataset = dataset[:-1].set(
                "next",
                dataset.select("info", strict=False)[1:],
            )
            dataset.rename_key_("next_observations", ("next", "observation"))
        else:
            dataset = dataset[:-1].set(
                "next",
                dataset.select("observation", "info", strict=False)[1:],
            )
        dataset["next"].update(
            dataset.select("reward", "done", "terminated", "truncated", strict=False)
        )
        dataset = (
            dataset.clone()
        )  # make sure that all tensors have a different data_ptr
        self._shift_reward_done(dataset)
        if env is not None:
            self.specs = env.specs.clone()
        else:
            self.specs = None
        return dataset

    def _shift_reward_done(self, dataset):
        dataset["reward"] = dataset["reward"].clone()
        dataset["reward"][1:] = dataset["reward"][:-1].clone()
        dataset["reward"][0] = 0
        for key in ("done", "terminated", "truncated"):
            if key not in dataset.keys():
                continue
            dataset[key] = dataset[key].clone()
            dataset[key][1:] = dataset[key][:-1].clone()
            dataset[key][0] = 0


def _download_dataset_from_url(dataset_url, dataset_path):
    dataset_filepath = _filepath_from_url(dataset_url, dataset_path)
    if not os.path.exists(dataset_filepath):
        torchrl_logger.info(f"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_path):
    _, 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(_get_root_dir("d4rl")))

if __name__ == "__main__":
    data = D4RLExperienceReplay("kitchen-partial-v0", batch_size=128)
    torchrl_logger.info(data)
    for sample in data:
        torchrl_logger.info(sample)
        break