# 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 abc
import gc
import json
import os
import shutil
import tempfile
import warnings
from pathlib import Path
import numpy as np
import torch
from tensordict import (
is_tensor_collection,
lazy_stack,
NonTensorData,
PersistentTensorDict,
TensorDict,
)
from tensordict.memmap import MemoryMappedTensor
from tensordict.utils import _zip_strict
from torchrl._utils import _STRDTYPE2DTYPE
from torchrl.data.replay_buffers.utils import (
_save_pytree,
Flat2TED,
H5Combine,
H5Split,
Nested2TED,
TED2Flat,
TED2Nested,
)
def _map_metadata(function, metadata, *, is_leaf):
if is_leaf(metadata):
return function(metadata)
if isinstance(metadata, dict):
return {
key: _map_metadata(function, value, is_leaf=is_leaf)
for key, value in metadata.items()
}
if isinstance(metadata, list):
return [_map_metadata(function, value, is_leaf=is_leaf) for value in metadata]
if isinstance(metadata, tuple):
return tuple(
_map_metadata(function, value, is_leaf=is_leaf) for value in metadata
)
return function(metadata)
def _write_compressed_data_to_memmap(compressed_data, path):
processed_data = []
for item in compressed_data:
if item is None:
processed_data.append(None)
elif isinstance(item, torch.Tensor):
processed_data.append(TensorDict({"data": item}, batch_size=[]))
elif isinstance(item, dict):
processed_data.append(TensorDict(item, batch_size=[]))
else:
processed_data.append(TensorDict({"data": item}, batch_size=[]))
non_none_data = [item for item in processed_data if item is not None]
if not non_none_data:
return []
stacked_data = lazy_stack(non_none_data)
stacked_data.memmap_(path / "compressed_data")
data_indices = []
current_idx = 0
for item in processed_data:
if item is None:
data_indices.append(None)
else:
data_indices.append(current_idx)
current_idx += 1
return data_indices
[docs]
class StorageCheckpointerBase:
"""Public base class for storage checkpointers.
Each storage checkpointer must implement a `save` and `load` method that take as input a storage and a
path.
"""
def __init__(self):
self._save_hooks = []
self._load_hooks = []
[docs]
def register_save_hook(self, hook):
"""Registers a save hook for this checkpointer."""
self._save_hooks.append(hook)
[docs]
def register_load_hook(self, hook):
"""Registers a load hook for this checkpointer."""
self._load_hooks.append(hook)
def _detach_from_load_path(self, storage):
del storage
def _get_shift_from_last_cursor(self, last_cursor):
"""Computes shift from the last cursor position."""
if isinstance(last_cursor, slice):
return last_cursor.stop + 1
if isinstance(last_cursor, int):
return last_cursor + 1
if isinstance(last_cursor, range):
return last_cursor[-1] + 1
if isinstance(last_cursor, torch.Tensor):
return last_cursor.reshape(-1)[-1].item() + 1
if isinstance(last_cursor, np.ndarray):
return last_cursor.reshape(-1)[-1].item() + 1
raise ValueError(f"Unrecognised last_cursor type {type(last_cursor)}.")
def _set_hooks_shift_is_full(self, storage):
"""Sets shift and is_full attributes on save hooks that have them."""
is_full = storage._is_full
last_cursor = storage._last_cursor
for hook in self._save_hooks:
if hasattr(hook, "is_full"):
hook.is_full = is_full
if last_cursor is None:
warnings.warn(
"last_cursor is None. The replay buffer "
"may not be saved properly in this setting. To solve this issue, make "
"sure the storage updates the _last_cursor value during calls to `set`."
)
shift = 0
else:
shift = self._get_shift_from_last_cursor(last_cursor)
for hook in self._save_hooks:
if hasattr(hook, "shift"):
hook.shift = shift
@abc.abstractmethod
def dumps(self, storage, path):
...
@abc.abstractmethod
def loads(self, storage, path):
...
[docs]
class ListStorageCheckpointer(StorageCheckpointerBase):
"""A storage checkpointer for ListStoage.
Currently not implemented.
"""
@staticmethod
def dumps(storage, path):
raise NotImplementedError(
"ListStorage doesn't support serialization via `dumps` - `loads` API."
)
@staticmethod
def loads(storage, path):
raise NotImplementedError(
"ListStorage doesn't support serialization via `dumps` - `loads` API."
)
[docs]
class CompressedListStorageCheckpointer(StorageCheckpointerBase):
"""A storage checkpointer for CompressedListStorage.
This checkpointer saves compressed data and metadata using memory-mapped storage
for efficient disk I/O and memory usage.
"""
[docs]
def dumps(self, storage, path):
"""Save compressed storage to disk using memory-mapped storage.
Args:
storage: The CompressedListStorage instance to save
path: Directory path where to save the storage
"""
path = Path(path)
path.mkdir(exist_ok=True)
if not hasattr(storage, "_storage") or len(storage._storage) == 0:
raise RuntimeError(
"Cannot save an empty or non-initialized CompressedListStorage."
)
# Get state dict from storage
state_dict = storage.state_dict()
compressed_data = state_dict["_storage"]
metadata = state_dict["_metadata"]
# Create a temporary directory for processing
with tempfile.TemporaryDirectory() as tmp_dir:
tmp_path = Path(tmp_dir)
try:
data_indices = _write_compressed_data_to_memmap(
compressed_data, tmp_path
)
finally:
# Windows cannot remove a file while a mapping is alive. The helper
# owns all temporary TensorDict references, so collecting after it
# returns releases any reference cycles before TemporaryDirectory
# removes the staging tree.
if os.name == "nt":
gc.collect()
# Process metadata for JSON serialization
def is_leaf(item):
return isinstance(
item,
(
torch.Size,
torch.dtype,
torch.device,
str,
int,
float,
bool,
torch.Tensor,
NonTensorData,
),
)
def map_to_json_serializable(item):
if isinstance(item, torch.Size):
return {"__type__": "torch.Size", "value": list(item)}
elif isinstance(item, torch.dtype):
return {"__type__": "torch.dtype", "value": str(item)}
elif isinstance(item, torch.device):
return {"__type__": "torch.device", "value": str(item)}
elif isinstance(item, torch.Tensor):
return {"__type__": "torch.Tensor", "value": item.tolist()}
elif isinstance(item, NonTensorData):
return {"__type__": "NonTensorData", "value": item.data}
return item
serializable_metadata = _map_metadata(
map_to_json_serializable, metadata, is_leaf=is_leaf
)
# Save metadata and indices
metadata_file = tmp_path / "metadata.json"
with open(metadata_file, "w") as f:
json.dump(serializable_metadata, f, indent=2)
indices_file = tmp_path / "data_indices.json"
with open(indices_file, "w") as f:
json.dump(data_indices, f, indent=2)
# Copy all files from temp directory to final destination
for item in tmp_path.iterdir():
if item.is_file():
shutil.copy2(item, path / item.name)
elif item.is_dir():
shutil.copytree(item, path / item.name, dirs_exist_ok=True)
[docs]
def loads(self, storage, path):
"""Load compressed storage from disk.
Args:
storage: The CompressedListStorage instance to load into
path: Directory path where the storage was saved
"""
path = Path(path)
# Load metadata
metadata_file = path / "metadata.json"
if not metadata_file.exists():
raise FileNotFoundError(f"Metadata file not found at {metadata_file}")
with open(metadata_file) as f:
serializable_metadata = json.load(f)
# Load data indices
indices_file = path / "data_indices.json"
if not indices_file.exists():
raise FileNotFoundError(f"Data indices file not found at {indices_file}")
with open(indices_file) as f:
data_indices = json.load(f)
# Convert serializable metadata back to original format
def is_leaf(item):
return isinstance(item, dict) and "__type__" in item
def map_from_json_serializable(item):
if isinstance(item, dict) and "__type__" in item:
if item["__type__"] == "torch.Size":
return torch.Size(item["value"])
elif item["__type__"] == "torch.dtype":
# Handle torch.dtype conversion
dtype_str = item["value"]
if hasattr(torch, dtype_str.replace("torch.", "")):
return getattr(torch, dtype_str.replace("torch.", ""))
else:
# Handle cases like 'torch.float32' -> torch.float32
return eval(dtype_str)
elif item["__type__"] == "torch.device":
return torch.device(item["value"])
elif item["__type__"] == "torch.Tensor":
return torch.tensor(item["value"])
elif item["__type__"] == "NonTensorData":
return NonTensorData(item["value"])
return item
metadata = _map_metadata(
map_from_json_serializable, serializable_metadata, is_leaf=is_leaf
)
# Load compressed data from memmap
compressed_data = []
memmap_path = path / "compressed_data"
if memmap_path.exists():
# Load the memmapped data
stacked_data = TensorDict.load_memmap(memmap_path)
compressed_data = stacked_data.tolist()
if len(compressed_data) != len(data_indices):
raise ValueError(
f"Length of compressed data ({len(compressed_data)}) does not match length of data indices ({len(data_indices)})"
)
for i, (data, mtdt) in enumerate(_zip_strict(compressed_data, metadata)):
if mtdt["type"] == "tensor":
compressed_data[i] = data["data"]
else:
compressed_data[i] = data
else:
# No data to load
compressed_data = [None] * len(data_indices)
# Load into storage
storage._storage = compressed_data
storage._metadata = metadata
def _detach_from_load_path(self, storage):
storage._storage = _map_metadata(
lambda item: item.clone() if isinstance(item, MemoryMappedTensor) else item,
storage._storage,
is_leaf=lambda item: isinstance(item, torch.Tensor),
)
if os.name == "nt":
gc.collect()
[docs]
class TensorStorageCheckpointer(StorageCheckpointerBase):
"""A storage checkpointer for TensorStorages.
This class supports TensorDict-based storages as well as pytrees.
This class will call save and load hooks if provided. These hooks should take as input the
data being transformed as well as the path where the data should be saved.
"""
def dumps(self, storage, path):
path = Path(path)
path.mkdir(exist_ok=True)
if not storage.initialized:
raise RuntimeError("Cannot save a non-initialized storage.")
metadata = {}
_storage = storage._storage
self._set_hooks_shift_is_full(storage)
for hook in self._save_hooks:
_storage = hook(_storage, path=path)
if is_tensor_collection(_storage):
if (
_storage.is_memmap()
and _storage.saved_path
and Path(_storage.saved_path).absolute() == Path(path).absolute()
):
_storage.memmap_refresh_()
else:
# try to load the path and overwrite.
_storage.memmap(
path,
copy_existing=True, # num_threads=torch.get_num_threads()
)
is_pytree = False
else:
_save_pytree(_storage, metadata, path)
is_pytree = True
with open(path / "storage_metadata.json", "w") as file:
json.dump(
{
"metadata": metadata,
"is_pytree": is_pytree,
"len": storage._len,
},
file,
)
def loads(self, storage, path):
with open(path / "storage_metadata.json") as file:
metadata = json.load(file)
is_pytree = metadata["is_pytree"]
_len = metadata["len"]
if is_pytree:
if self._load_hooks:
raise RuntimeError(
"Loading hooks are not compatible with PyTree storages."
)
path = Path(path)
for local_path, md in metadata["metadata"].items():
# load tensor
local_path_dot = local_path.replace(".", "/")
total_tensor_path = path / (local_path_dot + ".memmap")
shape = torch.Size(md["shape"])
dtype = _STRDTYPE2DTYPE[md["dtype"]]
tensor = MemoryMappedTensor.from_filename(
filename=total_tensor_path, shape=shape, dtype=dtype
)
# split path
local_path = local_path.split(".")
# replace potential dots
local_path = [_path.replace("_<dot>_", ".") for _path in local_path]
if storage.initialized:
# copy in-place
_storage_tensor = storage._storage
# in this case there is a single tensor, so we skip
if local_path != ["_-single-tensor-_"]:
for _path in local_path:
if _path.isdigit():
_path_attempt = int(_path)
try:
_storage_tensor = _storage_tensor[_path_attempt]
continue
except IndexError:
pass
_storage_tensor = _storage_tensor[_path]
_storage_tensor.copy_(tensor)
else:
raise RuntimeError(
"Cannot fill a non-initialized pytree-based TensorStorage."
)
else:
_storage = TensorDict.load_memmap(path)
if storage.initialized:
dest = storage._storage
else:
# TODO: This could load the RAM a lot, maybe try to catch this within the hook and use memmap instead
dest = None
for hook in self._load_hooks:
_storage = hook(_storage, out=dest)
if not storage.initialized:
from torchrl.data.replay_buffers.storages import LazyMemmapStorage
if (
isinstance(storage, LazyMemmapStorage)
and storage.scratch_dir
and Path(storage.scratch_dir).absolute() == Path(path).absolute()
):
storage._storage = TensorDict.load_memmap(path)
storage.initialized = True
else:
# this should not be reached if is_pytree=True
storage._init(_storage[0])
storage._storage.update_(_storage)
elif (
storage._storage.is_memmap()
and storage._storage.saved_path
and Path(storage._storage.saved_path).absolute()
== Path(path).absolute()
):
# If the storage is already where it should be, we don't need to load anything.
storage._storage.memmap_refresh_()
else:
storage._storage.copy_(_storage)
storage._len = _len
[docs]
class FlatStorageCheckpointer(TensorStorageCheckpointer):
"""Saves the storage in a compact form, saving space on the TED format.
This class explicitly assumes and does NOT check that:
- done states (including terminated and truncated) at the root are always False;
- observations in the "next" tensordict are shifted by one step in the future (this
is not the case when a multi-step transform is used for instance) unless `done` is `True`
in which case the observation in `("next", key)` at time `t` and the one in `key` at time
`t+1` should not match.
.. seealso: The full list of arguments can be found in :class:`~torchrl.data.TED2Flat`.
"""
def __init__(self, done_keys=None, reward_keys=None):
super().__init__()
kwargs = {}
if done_keys is not None:
kwargs["done_keys"] = done_keys
if reward_keys is not None:
kwargs["reward_keys"] = reward_keys
self._save_hooks = [TED2Flat(**kwargs)]
self._load_hooks = [Flat2TED(**kwargs)]
[docs]
class NestedStorageCheckpointer(FlatStorageCheckpointer):
"""Saves the storage in a compact form, saving space on the TED format and using memory-mapped nested tensors.
This class explicitly assumes and does NOT check that:
- done states (including terminated and truncated) at the root are always False;
- observations in the "next" tensordict are shifted by one step in the future (this
is not the case when a multi-step transform is used for instance).
.. seealso: The full list of arguments can be found in :class:`~torchrl.data.TED2Flat`.
"""
def __init__(self, done_keys=None, reward_keys=None):
super().__init__()
kwargs = {}
if done_keys is not None:
kwargs["done_keys"] = done_keys
if reward_keys is not None:
kwargs["reward_keys"] = reward_keys
self._save_hooks = [TED2Nested(**kwargs)]
self._load_hooks = [Nested2TED(**kwargs)]
[docs]
class H5StorageCheckpointer(NestedStorageCheckpointer):
"""Saves the storage in a compact form, saving space on the TED format and using H5 format to save the data.
This class explicitly assumes and does NOT check that:
- done states (including terminated and truncated) at the root are always False;
- observations in the "next" tensordict are shifted by one step in the future (this
is not the case when a multi-step transform is used for instance).
Keyword Args:
checkpoint_file: the filename where to save the checkpointed data.
This will be ignored iff the path passed to dumps / loads ends with the ``.h5``
suffix. Defaults to ``"checkpoint.h5"``.
h5_kwargs (Dict[str, Any] or Tuple[Tuple[str, Any], ...]): kwargs to be
passed to :meth:`h5py.File.create_dataset`.
.. note:: To prevent out-of-memory issues, the data of the H5 file will be temporarily written
on memory-mapped tensors stored in shared file system. The physical memory usage may increase
during loading as a consequence.
.. seealso: The full list of arguments can be found in :class:`~torchrl.data.TED2Flat`. Note that this class only
supports keyword arguments.
"""
def __init__(
self,
*,
checkpoint_file: str = "checkpoint.h5",
done_keys=None,
reward_keys=None,
h5_kwargs=None,
**kwargs,
):
StorageCheckpointerBase.__init__(self)
ted2_kwargs = kwargs
if done_keys is not None:
ted2_kwargs["done_keys"] = done_keys
if reward_keys is not None:
ted2_kwargs["reward_keys"] = reward_keys
self._save_hooks = [TED2Nested(**ted2_kwargs), H5Split()]
self._load_hooks = [H5Combine(), Nested2TED(**ted2_kwargs)]
self.kwargs = {} if h5_kwargs is None else dict(h5_kwargs)
self.checkpoint_file = checkpoint_file
def dumps(self, storage, path):
path = self._get_path(path)
self._set_hooks_shift_is_full(storage)
if not storage.initialized:
raise RuntimeError("Cannot save a non-initialized storage.")
_storage = storage._storage
length = storage._len
for hook in self._save_hooks:
# we don't pass a path here since we're not reusing the tensordict
_storage = hook(_storage)
if is_tensor_collection(_storage):
# try to load the path and overwrite.
data = PersistentTensorDict.from_dict(_storage, path, **self.kwargs)
data["_len"] = NonTensorData(data=length)
else:
raise ValueError("Only tensor collections are supported.")
def loads(self, storage, path):
path = self._get_path(path)
data = PersistentTensorDict.from_h5(path)
if storage.initialized:
dest = storage._storage
else:
# TODO: This could load the RAM a lot, maybe try to catch this within the hook and use memmap instead
dest = None
_len = data["_len"]
for hook in self._load_hooks:
data = hook(data, out=dest)
if not storage.initialized:
# this should not be reached if is_pytree=True
storage._init(data[0])
storage._storage.update_(data)
else:
storage._storage.copy_(data)
storage._len = _len
def _get_path(self, path):
path = Path(path)
if path.suffix == ".h5":
return str(path.absolute())
try:
path.mkdir(exist_ok=True)
except Exception:
raise RuntimeError(f"Failed to create the checkpoint directory {path}.")
path = path / self.checkpoint_file
return str(path.absolute())
[docs]
class StorageEnsembleCheckpointer(StorageCheckpointerBase):
"""Checkpointer for ensemble storages."""
@staticmethod
def dumps(storage, path: Path):
path = Path(path).absolute()
path.mkdir(parents=True, exist_ok=True)
empty = [i for i, member in enumerate(storage._storages) if not len(member)]
with open(path / "metadata.json", "w") as file:
json.dump({"empty_members": empty}, file)
for i, member in enumerate(storage._storages):
if i not in empty:
member.dumps(path / str(i))
if storage._transforms is not None:
for i, transform in enumerate(storage._transforms):
if transform is not None:
torch.save(transform.state_dict(), path / f"{i}_transform.pt")
@staticmethod
def loads(storage, path: Path):
path = Path(path).absolute()
metadata_path = path / "metadata.json"
if metadata_path.exists():
with open(metadata_path) as file:
empty = set(json.load(file)["empty_members"])
else:
empty = set()
for i, member in enumerate(storage._storages):
if i in empty:
member._empty()
else:
member.loads(path / str(i))
if storage._transforms is not None:
for i, transform in enumerate(storage._transforms):
if transform is not None:
transform.load_state_dict(torch.load(path / f"{i}_transform.pt"))
