torchrl 0.11.0__cp314-cp314-win_amd64.whl

torchrl/data/datasets/minari_data.py

@@ -0,0 +1,643 @@
+# 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.util
+import json
+import os.path
+import shutil
+import tempfile
+from collections import defaultdict
+from collections.abc import Callable
+from contextlib import nullcontext
+from dataclasses import asdict
+from pathlib import Path
+
+import torch
+from tensordict import (
+    is_non_tensor,
+    is_tensor_collection,
+    NonTensorData,
+    NonTensorStack,
+    PersistentTensorDict,
+    set_list_to_stack,
+    TensorDict,
+    TensorDictBase,
+)
+
+from torchrl._utils import KeyDependentDefaultDict, logger as torchrl_logger
+from torchrl.data.datasets.common import BaseDatasetExperienceReplay
+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
+from torchrl.data.tensor_specs import Bounded, Categorical, Composite, Unbounded
+from torchrl.envs.utils import _classproperty
+
+_has_tqdm = importlib.util.find_spec("tqdm", None) is not None
+_has_minari = importlib.util.find_spec("minari", None) is not None
+
+_NAME_MATCH = KeyDependentDefaultDict(lambda key: key)
+_NAME_MATCH["observations"] = "observation"
+_NAME_MATCH["rewards"] = "reward"
+_NAME_MATCH["truncations"] = "truncated"
+_NAME_MATCH["terminations"] = "terminated"
+_NAME_MATCH["actions"] = "action"
+_NAME_MATCH["infos"] = "info"
+
+
+_DTYPE_DIR = {
+    "float16": torch.float16,
+    "float32": torch.float32,
+    "float64": torch.float64,
+    "int64": torch.int64,
+    "int32": torch.int32,
+    "uint8": torch.uint8,
+}
+
+
+class MinariExperienceReplay(BaseDatasetExperienceReplay):
+    """Minari Experience replay dataset.
+
+    Learn more about Minari on their website: https://minari.farama.org/
+
+    The data format follows the :ref:`TED convention <TED-format>`.
+
+    Args:
+        dataset_id (str): The dataset to be downloaded. Must be part of MinariExperienceReplay.available_datasets
+        batch_size (int): Batch-size used during sampling. Can be overridden by `data.sample(batch_size)` if
+            necessary.
+
+    Keyword Args:
+        root (Path or str, optional): The Minari 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/atari`.minari`.
+        download (bool or str, optional): Whether the dataset should be downloaded if
+            not found. Defaults to ``True``. Download can also be passed as ``"force"``,
+            in which case the downloaded data will be overwritten.
+        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.
+            Defaults to ``False``.
+        load_from_local_minari (bool, optional): if ``True``, the dataset will be loaded directly
+            from the local Minari cache (typically located at ``~/.minari/datasets``),
+            bypassing any remote download. This is useful when working with custom
+            Minari datasets previously generated and stored locally, or when network
+            access should be avoided. If the dataset is not found in the expected
+            cache directory, a ``FileNotFoundError`` will be raised.
+            Defaults to ``False``.
+
+
+    Attributes:
+        available_datasets: a list of accepted entries to be downloaded.
+
+    .. note::
+      Text data is currenrtly discarded from the wrapped dataset, as there is not
+      PyTorch native way of representing text data.
+      If this feature is required, please post an issue on TorchRL's GitHub
+      repository.
+
+    Examples:
+        >>> from torchrl.data.datasets.minari_data import MinariExperienceReplay
+        >>> data = MinariExperienceReplay("door-human-v1", batch_size=32, download="force")
+        >>> for sample in data:
+        ...     torchrl_logger.info(sample)
+        ...     break
+        TensorDict(
+            fields={
+                action: Tensor(shape=torch.Size([32, 28]), device=cpu, dtype=torch.float32, is_shared=False),
+                index: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.int64, is_shared=False),
+                info: TensorDict(
+                    fields={
+                        success: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([32]),
+                    device=cpu,
+                    is_shared=False),
+                next: TensorDict(
+                    fields={
+                        done: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        info: TensorDict(
+                            fields={
+                                success: Tensor(shape=torch.Size([32]), device=cpu, dtype=torch.bool, is_shared=False)},
+                            batch_size=torch.Size([32]),
+                            device=cpu,
+                            is_shared=False),
+                        observation: Tensor(shape=torch.Size([32, 39]), device=cpu, dtype=torch.float64, is_shared=False),
+                        reward: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.float64, is_shared=False),
+                        state: TensorDict(
+                            fields={
+                                door_body_pos: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
+                                qpos: Tensor(shape=torch.Size([32, 30]), device=cpu, dtype=torch.float64, is_shared=False),
+                                qvel: Tensor(shape=torch.Size([32, 30]), device=cpu, dtype=torch.float64, is_shared=False)},
+                            batch_size=torch.Size([32]),
+                            device=cpu,
+                            is_shared=False),
+                        terminated: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.bool, is_shared=False),
+                        truncated: Tensor(shape=torch.Size([32, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
+                    batch_size=torch.Size([32]),
+                    device=cpu,
+                    is_shared=False),
+                observation: Tensor(shape=torch.Size([32, 39]), device=cpu, dtype=torch.float64, is_shared=False),
+                state: TensorDict(
+                    fields={
+                        door_body_pos: Tensor(shape=torch.Size([32, 3]), device=cpu, dtype=torch.float64, is_shared=False),
+                        qpos: Tensor(shape=torch.Size([32, 30]), device=cpu, dtype=torch.float64, is_shared=False),
+                        qvel: Tensor(shape=torch.Size([32, 30]), device=cpu, dtype=torch.float64, is_shared=False)},
+                    batch_size=torch.Size([32]),
+                    device=cpu,
+                    is_shared=False)},
+            batch_size=torch.Size([32]),
+            device=cpu,
+            is_shared=False)
+
+    """
+
+    def __init__(
+        self,
+        dataset_id,
+        batch_size: int,
+        *,
+        root: str | Path | None = None,
+        download: bool = True,
+        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,
+        load_from_local_minari: bool = False,
+    ):
+        self.dataset_id = dataset_id
+        if root is None:
+            root = _get_root_dir("minari")
+            os.makedirs(root, exist_ok=True)
+        self.root = root
+        self.split_trajs = split_trajs
+        self.download = download
+        self.load_from_local_minari = load_from_local_minari
+
+        if (
+            self.download == "force"
+            or (self.download and not self._is_downloaded())
+            or self.load_from_local_minari
+        ):
+            if self.download == "force":
+                try:
+                    if os.path.exists(self.data_path_root):
+                        shutil.rmtree(self.data_path_root)
+                    if self.data_path != self.data_path_root:
+                        shutil.rmtree(self.data_path)
+                except FileNotFoundError:
+                    pass
+            storage = self._download_and_preproc()
+        elif self.split_trajs and not os.path.exists(self.data_path):
+            storage = self._make_split()
+        else:
+            storage = self._load()
+        storage = TensorStorage(storage)
+
+        if writer is None:
+            writer = ImmutableDatasetWriter()
+
+        super().__init__(
+            storage=storage,
+            sampler=sampler,
+            writer=writer,
+            collate_fn=collate_fn,
+            pin_memory=pin_memory,
+            prefetch=prefetch,
+            batch_size=batch_size,
+            transform=transform,
+        )
+
+    @_classproperty
+    def available_datasets(self):
+        if not _has_minari:
+            raise ImportError("minari library not found.")
+        import minari
+
+        return minari.list_remote_datasets().keys()
+
+    def _is_downloaded(self):
+        return os.path.exists(self.data_path_root)
+
+    @property
+    def data_path(self) -> Path:
+        if self.split_trajs:
+            return Path(self.root) / (self.dataset_id + "_split")
+        return self.data_path_root
+
+    @property
+    def data_path_root(self) -> Path:
+        return Path(self.root) / self.dataset_id
+
+    @property
+    def metadata_path(self) -> Path:
+        return Path(self.root) / self.dataset_id / "env_metadata.json"
+
+    def _download_and_preproc(self):
+        if not _has_minari:
+            raise ImportError("minari library not found.")
+        import minari
+
+        if _has_tqdm:
+            from tqdm import tqdm
+
+        prev_minari_datasets_path_save = prev_minari_datasets_path = os.environ.get(
+            "MINARI_DATASETS_PATH"
+        )
+        try:
+            if prev_minari_datasets_path is None:
+                prev_minari_datasets_path = os.path.expanduser("~/.minari/datasets")
+            with tempfile.TemporaryDirectory() as tmpdir:
+
+                total_steps = 0
+                td_data = TensorDict()
+
+                if self.load_from_local_minari:
+                    # Load minari dataset from user's local Minari cache
+
+                    parent_dir = (
+                        Path(prev_minari_datasets_path) / self.dataset_id / "data"
+                    )
+                    h5_path = parent_dir / "main_data.hdf5"
+
+                    if not h5_path.exists():
+                        raise FileNotFoundError(
+                            f"{h5_path} does not exist in local Minari cache!"
+                        )
+
+                    torchrl_logger.info(
+                        f"loading dataset from local Minari cache at {h5_path}"
+                    )
+                    h5_data = PersistentTensorDict.from_h5(h5_path)
+                    h5_data = h5_data.to_tensordict()
+
+                else:
+                    # temporarily change the minari cache path
+                    prev_minari_datasets_path_save2 = os.environ.get(
+                        "MINARI_DATASETS_PATH"
+                    )
+                    os.environ["MINARI_DATASETS_PATH"] = tmpdir
+                    try:
+                        minari.download_dataset(dataset_id=self.dataset_id)
+                    finally:
+                        if prev_minari_datasets_path_save2 is not None:
+                            os.environ[
+                                "MINARI_DATASETS_PATH"
+                            ] = prev_minari_datasets_path_save2
+
+                    parent_dir = Path(tmpdir) / self.dataset_id / "data"
+
+                    torchrl_logger.info(
+                        "first read through data to create data structure..."
+                    )
+                    h5_data = PersistentTensorDict.from_h5(
+                        parent_dir / "main_data.hdf5"
+                    )
+                    h5_data = h5_data.to_tensordict()
+
+                # populate the tensordict
+                episode_dict = {}
+                dataset_has_nontensor = False
+                for i, (episode_key, episode) in enumerate(h5_data.items()):
+                    episode_num = int(episode_key[len("episode_") :])
+                    episode_len = episode["actions"].shape[0]
+                    episode_dict[episode_num] = (episode_key, episode_len)
+                    # Get the total number of steps for the dataset
+                    total_steps += episode_len
+                    if i == 0:
+                        td_data.set("episode", 0)
+                        seen = set()
+                        for key, val in episode.items():
+                            match = _NAME_MATCH[key]
+                            if match in seen:
+                                continue
+                            seen.add(match)
+                            if key in ("observations", "state", "infos"):
+                                val = episode[key]
+                                if is_tensor_collection(val) and any(
+                                    isinstance(
+                                        val.get(k), (NonTensorData, NonTensorStack)
+                                    )
+                                    for k in val.keys()
+                                ):
+                                    non_tensor_probe = val.clone()
+                                    _extract_nontensor_fields(
+                                        non_tensor_probe, recursive=True
+                                    )
+                                    dataset_has_nontensor = True
+                                if (
+                                    not val.shape
+                                ):  # no need for this, we don't need the proper length: or steps != val.shape[0] - 1:
+                                    if val.is_empty():
+                                        continue
+                                    if is_non_tensor(val):
+                                        continue
+                                    val = _patch_info(val)
+                                td_data.set(("next", match), torch.zeros_like(val[0]))
+                                td_data.set(match, torch.zeros_like(val[0]))
+                            elif key not in ("terminations", "truncations", "rewards"):
+                                td_data.set(match, torch.zeros_like(val[0]))
+                            else:
+                                td_data.set(
+                                    ("next", match),
+                                    torch.zeros_like(val[0].unsqueeze(-1)),
+                                )
+
+                # give it the proper size
+                td_data["next", "done"] = (
+                    td_data["next", "truncated"] | td_data["next", "terminated"]
+                )
+                if "terminated" in td_data.keys():
+                    td_data["done"] = td_data["truncated"] | td_data["terminated"]
+                td_data = td_data.expand(total_steps).contiguous()
+                # save to designated location
+                torchrl_logger.info(
+                    f"creating tensordict data in {self.data_path_root}: "
+                )
+                if dataset_has_nontensor:
+                    _preallocate_nontensor_fields(
+                        td_data, episode, total_steps, name_map=_NAME_MATCH
+                    )
+                torchrl_logger.info(f"tensordict structure: {td_data}")
+
+                torchrl_logger.info(
+                    f"Reading data from {max(*episode_dict) + 1} episodes"
+                )
+                index = 0
+                with tqdm(total=total_steps) if _has_tqdm else nullcontext() as pbar:
+                    # iterate over episodes and populate the tensordict
+                    for episode_num in sorted(episode_dict):
+                        episode_key, steps = episode_dict[episode_num]
+                        episode = _patch_nontensor_data_to_stack(
+                            h5_data.get(episode_key)
+                        )
+                        idx = slice(index, (index + steps))
+                        data_view = td_data[idx]
+                        data_view.fill_("episode", episode_num)
+                        for key, val in episode.items():
+                            match = _NAME_MATCH[key]
+                            if key in (
+                                "observations",
+                                "state",
+                                "infos",
+                            ):
+                                if not val.shape or steps != val.shape[0] - 1:
+                                    if val.is_empty():
+                                        continue
+                                    if is_non_tensor(val):
+                                        continue
+                                    val = _patch_info(val)
+                                if steps != val.shape[0] - 1:
+                                    raise RuntimeError(
+                                        f"Mismatching number of steps for key {key}: was {steps} but got {val.shape[0] - 1}."
+                                    )
+                                val_next = val[1:].clone()
+                                val_copy = val[:-1].clone()
+
+                                data_view["next", match].copy_(val_next)
+                                data_view[match].copy_(val_copy)
+
+                                if is_tensor_collection(val_next):
+                                    non_tensors_next = _extract_nontensor_fields(
+                                        val_next
+                                    )
+                                    non_tensors_now = _extract_nontensor_fields(
+                                        val_copy
+                                    )
+                                    data_view["next", match].update_(non_tensors_next)
+                                    data_view[match].update_(non_tensors_now)
+
+                            elif key not in ("terminations", "truncations", "rewards"):
+                                if steps is None:
+                                    steps = val.shape[0]
+                                else:
+                                    if steps != val.shape[0]:
+                                        raise RuntimeError(
+                                            f"Mismatching number of steps for key {key}: was {steps} but got {val.shape[0]}."
+                                        )
+                                data_view[match].copy_(val)
+                            else:
+                                if steps is None:
+                                    steps = val.shape[0]
+                                else:
+                                    if steps != val.shape[0]:
+                                        raise RuntimeError(
+                                            f"Mismatching number of steps for key {key}: was {steps} but got {val.shape[0]}."
+                                        )
+                                data_view[("next", match)].copy_(val.unsqueeze(-1))
+                        data_view["next", "done"].copy_(
+                            data_view["next", "terminated"]
+                            | data_view["next", "truncated"]
+                        )
+                        if "done" in data_view.keys():
+                            data_view["done"].copy_(
+                                data_view["terminated"] | data_view["truncated"]
+                            )
+                        if pbar is not None:
+                            pbar.update(steps)
+                            pbar.set_description(
+                                f"index={index} - episode num {episode_num}"
+                            )
+                        index += steps
+
+                td_data = td_data.memmap_like(self.data_path_root)
+                # Add a "done" entry
+                if self.split_trajs:
+                    with td_data.unlock_():
+                        from torchrl.collectors.utils import split_trajectories
+
+                        td_data = split_trajectories(td_data).memmap_(self.data_path)
+                with open(self.metadata_path, "w") as metadata_file:
+                    dataset = minari.load_dataset(self.dataset_id)
+                    self.metadata = asdict(dataset.spec)
+                    self.metadata["observation_space"] = _spec_to_dict(
+                        self.metadata["observation_space"]
+                    )
+                    self.metadata["action_space"] = _spec_to_dict(
+                        self.metadata["action_space"]
+                    )
+                    json.dump(self.metadata, metadata_file)
+                self._load_and_proc_metadata()
+                return td_data
+        finally:
+            if prev_minari_datasets_path_save is not None:
+                os.environ["MINARI_DATASETS_PATH"] = prev_minari_datasets_path_save
+
+    def _make_split(self):
+        from torchrl.collectors.utils import split_trajectories
+
+        self._load_and_proc_metadata()
+        td_data = TensorDict.load_memmap(self.data_path_root)
+        td_data = split_trajectories(td_data).memmap_(self.data_path)
+        return td_data
+
+    def _load(self):
+        self._load_and_proc_metadata()
+        return TensorDict.load_memmap(self.data_path)
+
+    def _load_and_proc_metadata(self):
+        with open(self.metadata_path) as file:
+            self.metadata = json.load(file)
+        self.metadata["observation_space"] = _proc_spec(
+            self.metadata["observation_space"]
+        )
+        self.metadata["action_space"] = _proc_spec(self.metadata["action_space"])
+
+
+def _proc_spec(spec):
+    if spec is None:
+        return
+    if spec["type"] == "Dict":
+        return Composite(
+            {key: _proc_spec(subspec) for key, subspec in spec["subspaces"].items()}
+        )
+    elif spec["type"] == "Box":
+        if all(item == -float("inf") for item in spec["low"]) and all(
+            item == float("inf") for item in spec["high"]
+        ):
+            return Unbounded(spec["shape"], dtype=_DTYPE_DIR[spec["dtype"]])
+        return Bounded(
+            shape=spec["shape"],
+            low=torch.as_tensor(spec["low"]),
+            high=torch.as_tensor(spec["high"]),
+            dtype=_DTYPE_DIR[spec["dtype"]],
+        )
+    elif spec["type"] == "Discrete":
+        return Categorical(
+            spec["n"], shape=spec["shape"], dtype=_DTYPE_DIR[spec["dtype"]]
+        )
+    else:
+        raise NotImplementedError(f"{type(spec)}")
+
+
+def _spec_to_dict(spec):
+    from torchrl.envs.libs.gym import gym_backend
+
+    if isinstance(spec, gym_backend("spaces").Dict):
+        return {
+            "type": "Dict",
+            "subspaces": {key: _spec_to_dict(val) for key, val in spec.items()},
+        }
+    if isinstance(spec, gym_backend("spaces").Box):
+        return {
+            "type": "Box",
+            "low": spec.low.tolist(),
+            "high": spec.high.tolist(),
+            "dtype": str(spec.dtype),
+            "shape": tuple(spec.shape),
+        }
+    if isinstance(spec, gym_backend("spaces").Discrete):
+        return {
+            "type": "Discrete",
+            "dtype": str(spec.dtype),
+            "n": int(spec.n),
+            "shape": tuple(spec.shape),
+        }
+    if isinstance(spec, gym_backend("spaces").Text):
+        return
+    raise NotImplementedError(f"{type(spec)}, {str(spec)}")
+
+
+def _patch_info(info_td):
+    # Some info dicts have tensors with one less element than others
+    # We explicitly assume that the missing item is in the first position because
+    # it wasn't given at reset time.
+    # An alternative explanation could be that the last element is missing because
+    # deemed useless for training...
+    unique_shapes = defaultdict(list)
+    for subkey, subval in info_td.items():
+        unique_shapes[subval.shape[0]].append(subkey)
+    if len(unique_shapes) == 1:
+        unique_shapes[subval.shape[0] + 1] = []
+    if len(unique_shapes) != 2:
+        raise RuntimeError(
+            f"Unique shapes in a sub-tensordict can only be of length 2, got shapes {unique_shapes}."
+        )
+    val_td = info_td.to_tensordict()
+    min_shape = min(*unique_shapes)  # can only be found at root
+    max_shape = min_shape + 1
+    val_td_sel = val_td.select(*unique_shapes[min_shape])
+    val_td_sel = val_td_sel.apply(
+        lambda x: torch.cat([torch.zeros_like(x[:1]), x], 0), batch_size=[min_shape + 1]
+    )
+    source = val_td.select(*unique_shapes[max_shape])
+    # make sure source has no batch size
+    source.batch_size = ()
+    if not source.is_empty():
+        val_td_sel.update(source, update_batch_size=True)
+    return val_td_sel
+
+
+def _patch_nontensor_data_to_stack(tensordict: TensorDictBase):
+    """Recursively replaces all NonTensorData fields in the TensorDict with NonTensorStack."""
+    for key, val in tensordict.items():
+        if isinstance(val, TensorDictBase):
+            _patch_nontensor_data_to_stack(val)  # in-place recursive
+        elif isinstance(val, NonTensorData):
+            data_list = list(val.data)
+            with set_list_to_stack(True):
+                tensordict[key] = data_list
+    return tensordict
+
+
+def _extract_nontensor_fields(
+    tensordict: TensorDictBase, recursive: bool = False
+) -> TensorDict:
+    """Deletes the NonTensor fields from tensordict and returns the deleted tensordict."""
+    extracted = {}
+    for key in list(tensordict.keys()):
+        val = tensordict.get(key)
+        if is_non_tensor(val):
+            extracted[key] = val
+            del tensordict[key]
+        elif recursive and is_tensor_collection(val):
+            nested = _extract_nontensor_fields(val, recursive=True)
+            if len(nested) > 0:
+                extracted[key] = nested
+    return TensorDict(extracted, batch_size=tensordict.batch_size)
+
+
+def _preallocate_nontensor_fields(
+    td_data: TensorDictBase, example: TensorDictBase, total_steps: int, name_map: dict
+):
+    """Preallocates NonTensorStack fields in td_data based on an example TensorDict, applying key remapping."""
+    with set_list_to_stack(True):
+
+        def _recurse(src_td: TensorDictBase, dst_td: TensorDictBase, prefix=()):
+            for key, val in src_td.items():
+                mapped_key = name_map.get(key, key)
+                full_dst_key = prefix + (mapped_key,)
+
+                if is_non_tensor(val):
+                    dummy_stack = NonTensorStack(
+                        *[total_steps for _ in range(total_steps)]
+                    )
+                    dst_td.set(full_dst_key, dummy_stack)
+                    dst_td.set(("next",) + full_dst_key, dummy_stack)
+
+                elif is_tensor_collection(val):
+                    _recurse(val, dst_td, full_dst_key)
+
+        _recurse(example, td_data)