Shortcuts

Source code for torchrl.data.datasets.minari_data

# 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 contextlib import nullcontext
from dataclasses import asdict
from pathlib import Path
from typing import Callable

import torch

from tensordict import PersistentTensorDict, TensorDict
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,
}


[docs]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/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``. 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, ): 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 if self.download == "force" or (self.download and not self._is_downloaded()): 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 with tempfile.TemporaryDirectory() as tmpdir: os.environ["MINARI_DATASETS_PATH"] = tmpdir minari.download_dataset(dataset_id=self.dataset_id) parent_dir = Path(tmpdir) / self.dataset_id / "data" td_data = TensorDict({}, []) total_steps = 0 torchrl_logger.info("first read through data to create data structure...") h5_data = PersistentTensorDict.from_h5(parent_dir / "main_data.hdf5") # populate the tensordict episode_dict = {} 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) for key, val in episode.items(): match = _NAME_MATCH[key] if key in ("observations", "state", "infos"): 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 val = _patch_info(val) td_data.set(("next", match), torch.zeros_like(val[0])) td_data.set(match, torch.zeros_like(val[0])) if 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) # save to designated location torchrl_logger.info(f"creating tensordict data in {self.data_path_root}: ") td_data = td_data.memmap_like(self.data_path_root) 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 = 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 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}." ) data_view["next", match].copy_(val[1:]) data_view[match].copy_(val[:-1]) 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 h5_data.close() # Add a "done" entry if self.split_trajs: with td_data.unlock_(): from torchrl.objectives.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 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, "r") 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 explicitely 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] ) val_td_sel.update(val_td.select(*unique_shapes[max_shape])) return val_td_sel

Docs

Access comprehensive developer documentation for PyTorch

View Docs

Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

Resources

Find development resources and get your questions answered

View Resources