# 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.
# import tree
from __future__ import annotations
import contextlib
import itertools
import math
import operator
import os
import threading
import typing
import weakref
from collections.abc import Callable
from pathlib import Path
from typing import Any, Union
import numpy as np
import torch
from tensordict import (
lazy_stack,
MemoryMappedTensor,
NonTensorData,
TensorDict,
TensorDictBase,
unravel_key,
)
from torch import Tensor
from torch.nn import functional as F
from torch.utils._pytree import LeafSpec, tree_flatten, tree_unflatten
# DEFAULT_DONE_KEYS lives in torchrl._utils (dependency-light so envs and
# collectors can share it) and is re-exported here and as
# torchrl.data.DEFAULT_DONE_KEYS, the public path.
from torchrl._utils import DEFAULT_DONE_KEYS, implement_for, logger as torchrl_logger
SINGLE_TENSOR_BUFFER_NAME = os.environ.get(
"SINGLE_TENSOR_BUFFER_NAME", "_-single-tensor-_"
)
INT_CLASSES_TYPING = Union[int, np.integer]
if hasattr(typing, "get_args"):
INT_CLASSES = typing.get_args(INT_CLASSES_TYPING)
else:
# python 3.7
INT_CLASSES = (int, np.integer)
_REPLAY_BOUNDARY_CACHE: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
_REPLAY_BOUNDARY_CACHE_LOCK = threading.RLock()
def _freeze_boundary_cache_key(value: Any) -> Any:
if isinstance(value, (list, tuple)):
return tuple(_freeze_boundary_cache_key(item) for item in value)
if isinstance(value, torch.device):
return str(value)
if isinstance(value, torch.Tensor):
value = value.reshape(-1)
return tuple(value.detach().cpu().tolist())
if isinstance(value, range):
return (value.start, value.stop, value.step)
return value
def _boundary_cursor_cache_key(cursor: Any) -> int | None:
if isinstance(cursor, torch.Tensor):
cursor = cursor.reshape(-1)
return int(cursor[-1].item()) if cursor.numel() else None
if isinstance(cursor, range):
return int(cursor[-1]) if len(cursor) else None
if cursor is None:
return None
return int(cursor)
def _boundary_distances_1d(
anchor: torch.Tensor,
start: torch.Tensor,
stop: torch.Tensor,
length: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Pure searchsorted boundary query used by sequence sample units."""
trajectory_index = torch.searchsorted(stop, anchor).remainder(stop.shape[0])
anchor_start = start[trajectory_index]
anchor_stop = stop[trajectory_index]
distance_from_start = torch.remainder(anchor - anchor_start, length)
distance_to_stop = torch.remainder(anchor_stop - anchor, length)
return distance_from_start.to(torch.long), distance_to_stop.to(torch.long)
def _boundary_distances_nd(
anchor: torch.Tensor,
start: torch.Tensor,
stop: torch.Tensor,
stop_key: torch.Tensor,
boundary_index: torch.Tensor,
lane_stride: torch.Tensor,
length: int,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Pure lane-aware searchsorted query for multidimensional storage."""
lane = (anchor[:, 1:] * lane_stride).sum(-1)
anchor_key = lane * (2 * length) + anchor[:, 0]
search_index = torch.searchsorted(stop_key, anchor_key)
trajectory_index = boundary_index[search_index]
anchor_start = start[trajectory_index, 0]
anchor_stop = stop[trajectory_index, 0]
distance_from_start = torch.remainder(anchor[:, 0] - anchor_start, length)
distance_to_stop = torch.remainder(anchor_stop - anchor[:, 0], length)
return distance_from_start.to(torch.long), distance_to_stop.to(torch.long)
def _make_boundary_search_nd(
start: torch.Tensor,
stop: torch.Tensor,
storage_shape: tuple[int, ...],
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
lane_shape = storage_shape[1:]
lane_stride = torch.tensor(
[math.prod(lane_shape[index + 1 :]) for index in range(len(lane_shape))],
dtype=torch.long,
device=stop.device,
)
lane = (stop[:, 1:] * lane_stride).sum(-1)
regular_key = lane * (2 * storage_shape[0]) + stop[:, 0]
order = torch.argsort(regular_key)
regular_key = regular_key[order]
ordered_lane = lane[order]
first = torch.ones_like(ordered_lane, dtype=torch.bool)
first[1:] = ordered_lane[1:] != ordered_lane[:-1]
first_index = torch.arange(order.shape[0], device=order.device)[first]
sentinel_key = regular_key[first_index] + storage_shape[0]
stop_key = torch.cat([regular_key, sentinel_key])
boundary_index = torch.cat([order, order[first_index]])
sort_order = torch.argsort(stop_key)
return stop_key[sort_order], boundary_index[sort_order], lane_stride
def _to_numpy(data: Tensor) -> np.ndarray:
return data.detach().cpu().numpy() if isinstance(data, torch.Tensor) else data
def _to_torch(
data: Tensor, device, pin_memory: bool = False, non_blocking: bool = False
) -> torch.Tensor:
if isinstance(data, np.generic):
return torch.as_tensor(data, device=device)
elif isinstance(data, np.ndarray):
data = torch.from_numpy(data)
elif not isinstance(data, Tensor):
data = torch.as_tensor(data, device=device)
if pin_memory:
data = data.pin_memory()
if device is not None:
data = data.to(device, non_blocking=non_blocking)
return data
def pin_memory_output(fun) -> Callable:
"""Calls pin_memory on outputs of decorated function if they have such method."""
def decorated_fun(self, *args, **kwargs):
output = fun(self, *args, **kwargs)
if self._pin_memory:
_tuple_out = True
if not isinstance(output, tuple):
_tuple_out = False
output = (output,)
output = tuple(_pin_memory(_output) for _output in output)
if _tuple_out:
return output
return output[0]
return output
return decorated_fun
def _pin_memory(output: Any) -> Any:
if hasattr(output, "pin_memory") and output.device == torch.device("cpu"):
return output.pin_memory()
else:
return output
def _reduce(
tensor: torch.Tensor, reduction: str, dim: int | None = None
) -> float | torch.Tensor:
"""Reduces a tensor given the reduction method."""
if reduction == "max":
result = tensor.max(dim=dim)
elif reduction == "min":
result = tensor.min(dim=dim)
elif reduction == "mean":
result = tensor.mean(dim=dim)
elif reduction == "median":
result = tensor.median(dim=dim)
elif reduction == "sum":
result = tensor.sum(dim=dim)
else:
raise NotImplementedError(f"Unknown reduction method {reduction}")
if isinstance(result, tuple):
result = result[0]
return result.item() if dim is None else result
def _is_int(index):
if isinstance(index, INT_CLASSES):
return True
if isinstance(index, (np.ndarray, torch.Tensor)):
return index.ndim == 0
return False
[docs]
def find_start_stop_traj(
*,
trajectory: torch.Tensor | None = None,
end: torch.Tensor | None = None,
at_capacity: bool,
cursor: int | torch.Tensor | range | None = None,
device: torch.device | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Recover trajectory boundaries from trajectory ids or end-of-trajectory flags.
This is the canonical trajectory-boundary recovery routine used by
:class:`~torchrl.data.replay_buffers.SliceSampler` and its subclasses. It
understands the storage layout of TorchRL's replay buffers: circular
(ring-buffer) storages where a trajectory may span the wrap point, and
partially-filled storages where the write cursor acts as an implicit
truncation. See :ref:`the trajectory-boundary documentation
<ref_traj_boundaries>` for the full contract.
Keyword Args:
trajectory (torch.Tensor, optional): a tensor of trajectory ids laid
out in storage order, with time along dim ``0`` and any extra
batch dims after it (shape ``[T, *B]``). A boundary is detected
wherever the id changes between two consecutive steps. Exclusive
with ``end``.
end (torch.Tensor, optional): a boolean tensor of end-of-trajectory
flags laid out in storage order (shape ``[T, *B]``, time along
dim ``0``). ``True`` marks the last step of a trajectory.
Exclusive with ``trajectory``.
at_capacity (bool): whether the storage is full and behaves as a
circular buffer. If ``True``, a trajectory that has no end flag
after the last row is assumed to continue at row ``0`` (it spans
the wrap point). If ``False``, the last valid row is always
treated as an end.
cursor (int, torch.Tensor or range, optional): the index of the last
written row (e.g. ``storage._last_cursor``). Only used when
``at_capacity=True``: the row under the cursor is forced to be an
end, since the data that followed it has been overwritten and the
stored trajectory is implicitly truncated there.
device (torch.device, optional): a device on which to run the
boundary computation (the underlying ``nonzero()`` call can
benefit from an accelerator for large storages). Results are
returned on the device of the input tensor. Defaults to ``None``
(compute where the input lives).
Returns:
A ``(start, stop, lengths)`` tuple where ``start`` and ``stop`` are
``[N, 1 + len(B)]`` integer tensors holding, for each of the ``N``
recovered trajectories, the time index of its first (resp. last,
inclusive) step in column ``0`` and the batch coordinates in the
remaining columns; ``lengths`` is a ``[N]`` tensor of trajectory
lengths. For a trajectory spanning the wrap point of a full circular
storage, ``start[i, 0] > stop[i, 0]`` and the length accounts for the
wrap.
Examples:
>>> import torch
>>> from torchrl.data import find_start_stop_traj
>>> # A full circular storage with 10 rows and end flags at rows 2 and 7.
>>> # The write cursor sits at row 4: row 4 is an implicit truncation.
>>> end = torch.zeros(10, dtype=torch.bool)
>>> end[2] = end[7] = True
>>> start, stop, lengths = find_start_stop_traj(end=end, at_capacity=True, cursor=4)
>>> start.squeeze(-1) # the trajectory starting at row 8 wraps around to row 2
tensor([8, 3, 5])
>>> stop.squeeze(-1)
tensor([2, 4, 7])
>>> lengths
tensor([5, 2, 3])
>>> # The same boundaries recovered from trajectory ids
>>> trajectory = torch.tensor([5, 5, 5, 0, 0, 1, 1, 1, 5, 5])
>>> start, stop, lengths = find_start_stop_traj(trajectory=trajectory, at_capacity=True)
>>> start.squeeze(-1), stop.squeeze(-1), lengths
(tensor([8, 3, 5]), tensor([2, 4, 7]), tensor([5, 2, 3]))
.. seealso::
:func:`~torchrl.collectors.utils.split_trajectories` splits a
contiguous *rollout batch* (fresh collector output) into a padded
``[B, T]`` layout. This function instead operates on *storage-order*
data and returns indices, leaving the data untouched.
"""
end, length = _derive_end_flags(
trajectory=trajectory, end=end, at_capacity=at_capacity, cursor=cursor
)
return _end_to_start_stop(end=end, length=length, device=device)
class _ReplayBoundaryIndex:
"""Shared trajectory boundaries for replay components."""
def __init__(
self,
*,
trajectory: torch.Tensor | None = None,
end: torch.Tensor | None = None,
at_capacity: bool,
cursor: int | torch.Tensor | range | None = None,
device: torch.device | None = None,
end_to_start_stop: (
Callable[
[torch.Tensor, int], tuple[torch.Tensor, torch.Tensor, torch.Tensor]
]
| None
) = None,
storage: Any | None = None,
source: Any | None = None,
cache_values: bool = False,
) -> None:
self.device = device
self._end_to_start_stop_fn = end_to_start_stop
self._boundaries = None
self._search = None
self._cache_storage = None
self._cache_key = None
shape = getattr(storage, "shape", None) if storage is not None else None
self.storage_shape = tuple(shape) if shape is not None else None
if cache_values and storage is not None and source is not None:
revision = storage._mutation_revision
cache_key = (
_freeze_boundary_cache_key(source),
tuple(shape) if shape is not None else None,
len(storage),
bool(storage._is_full),
_boundary_cursor_cache_key(
getattr(storage, "_last_cursor_index", None)
),
str(device) if device is not None else None,
)
with _REPLAY_BOUNDARY_CACHE_LOCK:
storage_cache = _REPLAY_BOUNDARY_CACHE.get(storage)
if storage_cache is None or storage_cache[0] != revision:
storage_cache = (revision, {})
_REPLAY_BOUNDARY_CACHE[storage] = storage_cache
cached = storage_cache[1].get(cache_key)
self._cache_storage = storage
self._cache_key = cache_key
if cached is not None:
self.end = None
self.length, self._boundaries, self._search = cached
return
end, length = _derive_end_flags(
trajectory=trajectory,
end=end,
at_capacity=at_capacity,
cursor=cursor,
)
self.end = end
self.length = length
def boundaries(self) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Returns canonical inclusive starts, stops, and trajectory lengths."""
boundaries = self._boundaries
if boundaries is None:
if self._end_to_start_stop_fn is None:
boundaries = _end_to_start_stop(
end=self.end, length=self.length, device=self.device
)
else:
boundaries = self._end_to_start_stop_fn(self.end, self.length)
if self.device is not None and self._end_to_start_stop_fn is None:
boundaries = tuple(value.to(self.device) for value in boundaries)
self._boundaries = boundaries
if self._cache_storage is not None:
with _REPLAY_BOUNDARY_CACHE_LOCK:
revision, storage_cache = _REPLAY_BOUNDARY_CACHE[
self._cache_storage
]
if revision == self._cache_storage._mutation_revision:
storage_cache[self._cache_key] = (
self.length,
boundaries,
self._search,
)
return boundaries
def _update_cached_search(self) -> None:
if self._cache_storage is None:
return
with _REPLAY_BOUNDARY_CACHE_LOCK:
revision, storage_cache = _REPLAY_BOUNDARY_CACHE[self._cache_storage]
if revision == self._cache_storage._mutation_revision:
storage_cache[self._cache_key] = (
self.length,
self._boundaries,
self._search,
)
def distances(self, anchor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
"""Returns distances from starts and to stops for each anchor."""
start, stop, _ = self.boundaries()
start = start.to(anchor.device)
stop = stop.to(anchor.device)
if start.shape[-1] == 1:
return _boundary_distances_1d(anchor, start[:, 0], stop[:, 0], self.length)
if anchor.ndim != 2 or anchor.shape[-1] != start.shape[-1]:
raise RuntimeError(
"Multidimensional anchors must have shape [batch, storage.ndim]."
)
if self.storage_shape is None:
raise RuntimeError(
"Multidimensional boundary queries require the storage shape."
)
search = self._search
if search is None:
search = _make_boundary_search_nd(start, stop, self.storage_shape)
self._search = search
self._update_cached_search()
else:
search = tuple(value.to(anchor.device) for value in search)
return _boundary_distances_nd(
anchor,
start,
stop,
*search,
self.length,
)
def _derive_end_flags(
*,
trajectory: torch.Tensor | None = None,
end: torch.Tensor | None = None,
at_capacity: bool,
cursor: int | torch.Tensor | range | None = None,
) -> tuple[torch.Tensor, int]:
"""Preprocessing stage of :func:`find_start_stop_traj`.
Turns trajectory ids or raw end flags into a fully-marked ``[T, *B]``
boolean end tensor (every trajectory has at least one end flag, the
cursor and partial-storage implicit truncations are applied) and returns
it together with the time length ``T``.
"""
if (trajectory is None) == (end is None):
raise TypeError(
"Exactly one of `trajectory` and `end` must be provided, got "
f"trajectory={trajectory} and end={end}."
)
signal = trajectory if trajectory is not None else end
if signal.ndim == 0:
raise RuntimeError(
"Expected the end-of-trajectory signal to be at least 1-dimensional."
)
if signal.shape[0] == 0:
raise RuntimeError(
"Cannot recover trajectory boundaries from an empty (T=0) input."
)
if trajectory is not None:
# slower
# _, stop_idx = torch.unique_consecutive(trajectory, return_counts=True)
# stop_idx = stop_idx.cumsum(0) - 1
# even slower
# t = trajectory.unsqueeze(0)
# w = torch.tensor([1, -1], dtype=torch.int).view(1, 1, 2)
# stop_idx = torch.conv1d(t, w).nonzero()
# faster
end = trajectory[:-1] != trajectory[1:]
if not at_capacity:
# ones_like(trajectory[:1]) rather than ones_like(end[:1]): for a
# single-step (T=1) input the diff above is empty, and the last
# valid row must still be marked as an end.
end = torch.cat([end, torch.ones_like(trajectory[:1], dtype=torch.bool)], 0)
else:
end = torch.cat([end, trajectory[-1:] != trajectory[:1]], 0)
length = trajectory.shape[0]
else:
# We presume that not done at the end means that the traj spans across end and beginning of storage
length = end.shape[0]
if not at_capacity:
end = end.clone()
end[length - 1] = True
if at_capacity:
# we must have at least one end by traj to individuate trajectories
# so if no end can be found we set it manually
if cursor is not None:
if isinstance(cursor, torch.Tensor):
cursor = cursor.reshape(-1)
if not cursor.numel():
raise RuntimeError("cursor must not be an empty tensor.")
cursor = cursor[-1].item()
elif isinstance(cursor, range):
if not len(cursor):
raise RuntimeError("cursor must not be an empty range.")
cursor = cursor[-1]
if not _is_int(cursor):
raise RuntimeError(
"cursor should be an integer or a 1d tensor or a range."
)
end = torch.index_fill(
end,
index=torch.tensor(cursor, device=end.device, dtype=torch.long),
dim=0,
value=1,
)
if not end.any(0).all():
mask = ~end.any(0, True)
mask = torch.cat([torch.zeros_like(end[:-1]), mask])
end = torch.masked_fill(mask, end, 1)
return end, length
def _end_to_start_stop(
end: torch.Tensor, length: int, device: torch.device | None = None
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""Convert a time-first end-flag tensor into ``(start, stop, lengths)`` indices.
Low-level companion of :func:`find_start_stop_traj`: expects a boolean
``[T, *B]`` tensor where every trajectory has at least one end flag, and
returns the boundary indices in the format documented there.
"""
orig_device = None
if device is not None and end.device != device:
orig_device = end.device
end = end.to(device)
# Using transpose ensures the start and stop are sorted the same way
stop_idx = end.transpose(0, -1).nonzero()
stop_idx[:, [0, -1]] = stop_idx[:, [-1, 0]].clone()
# First build the start indices as the stop + 1, we'll shift it later
start_idx = stop_idx.clone()
start_idx[:, 0] += 1
start_idx[:, 0] %= end.shape[0]
# shift start: to do this, we check when the non-first dim indices are identical
# and get a mask like [False, True, True, False, True, ...] where False means
# that there's a switch from one dim to another (ie, a switch from one element of the batch
# to another). We roll this one step along the time dimension and these two
# masks provide us with the indices of the permutation matrix we need
# to apply to start_idx.
if start_idx.shape[0] > 1:
start_idx_mask = (start_idx[1:, 1:] == start_idx[:-1, 1:]).all(-1)
m1 = torch.cat([torch.zeros_like(start_idx_mask[:1]), start_idx_mask])
m2 = torch.cat([start_idx_mask, torch.zeros_like(start_idx_mask[:1])])
start_idx_replace = torch.empty_like(start_idx)
start_idx_replace[m1] = start_idx[m2]
start_idx_replace[~m1] = start_idx[~m2]
start_idx = start_idx_replace
else:
# In this case we have only one start and stop has already been set
pass
lengths = stop_idx[:, 0] - start_idx[:, 0] + 1
lengths[lengths <= 0] = lengths[lengths <= 0] + length
if orig_device is not None:
return (
start_idx.to(orig_device),
stop_idx.to(orig_device),
lengths.to(orig_device),
)
return start_idx, stop_idx, lengths
[docs]
class TED2Flat:
"""A storage saving hook to serialize TED data in a compact format.
Args:
done_key (NestedKey, optional): the key where the done states should be read.
Defaults to ``("next", "done")``.
shift_key (NestedKey, optional): the key where the shift will be written.
Defaults to "shift".
is_full_key (NestedKey, optional): the key where the is_full attribute will be written.
Defaults to "is_full".
done_keys (Tuple[NestedKey], optional): a tuple of nested keys indicating the done entries.
Defaults to :data:`~torchrl.data.DEFAULT_DONE_KEYS`, i.e.
``("done", "truncated", "terminated")``.
reward_keys (Tuple[NestedKey], optional): a tuple of nested keys indicating the reward entries.
Defaults to ("reward",)
Examples:
>>> import tempfile
>>>
>>> from tensordict import TensorDict
>>>
>>> from torchrl.collectors import Collector
>>> from torchrl.data import ReplayBuffer, TED2Flat, LazyMemmapStorage
>>> from torchrl.envs import GymEnv
>>> import torch
>>>
>>> env = GymEnv("CartPole-v1")
>>> env.set_seed(0)
>>> torch.manual_seed(0)
>>> collector = Collector(env, policy=env.rand_step, total_frames=200, frames_per_batch=200)
>>> rb = ReplayBuffer(storage=LazyMemmapStorage(200))
>>> rb.register_save_hook(TED2Flat())
>>> with tempfile.TemporaryDirectory() as tmpdir:
... for i, data in enumerate(collector):
... rb.extend(data)
... rb.dumps(tmpdir)
... # load the data to represent it
... td = TensorDict.load(tmpdir + "/storage/")
... print(td)
TensorDict(
fields={
action: MemoryMappedTensor(shape=torch.Size([200, 2]), device=cpu, dtype=torch.int64, is_shared=True),
collector: TensorDict(
fields={
traj_ids: MemoryMappedTensor(shape=torch.Size([200]), device=cpu, dtype=torch.int64, is_shared=True)},
batch_size=torch.Size([]),
device=cpu,
is_shared=False),
done: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.bool, is_shared=True),
observation: MemoryMappedTensor(shape=torch.Size([220, 4]), device=cpu, dtype=torch.float32, is_shared=True),
reward: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.float32, is_shared=True),
terminated: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.bool, is_shared=True),
truncated: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.bool, is_shared=True)},
batch_size=torch.Size([]),
device=cpu,
is_shared=False)
"""
_shift: int | None = None
_is_full: bool | None = None
def __init__(
self,
done_key=("next", "done"),
shift_key="shift",
is_full_key="is_full",
done_keys=DEFAULT_DONE_KEYS,
reward_keys=("reward",),
):
self.done_key = done_key
self.shift_key = shift_key
self.is_full_key = is_full_key
self.done_keys = {unravel_key(key) for key in done_keys}
self.reward_keys = {unravel_key(key) for key in reward_keys}
@property
def shift(self):
return self._shift
@shift.setter
def shift(self, value: int):
self._shift = value
@property
def is_full(self):
return self._is_full
@is_full.setter
def is_full(self, value: int):
self._is_full = value
def __call__(self, data: TensorDictBase, path: Path = None):
# Get the done state
shift = self.shift
is_full = self.is_full
# Create an output storage
output = TensorDict()
output.set_non_tensor(self.is_full_key, is_full)
output.set_non_tensor(self.shift_key, shift)
output.set_non_tensor("_storage_shape", tuple(data.shape))
output.memmap_(path)
# Preallocate the output
done = data.get(self.done_key).squeeze(-1).clone()
if not is_full:
# shift is the cursor place
done[shift - 1] = True
else:
done = done.roll(-shift, dims=0)
done[-1] = True
ntraj = done.sum()
# Get the keys that require extra storage
keys_to_expand = set(data.get("next").keys(True, True)) - (
self.done_keys.union(self.reward_keys)
)
total_keys = data.exclude("next").keys(True, True)
total_keys = set(total_keys).union(set(data.get("next").keys(True, True)))
len_with_offset = data.numel() + ntraj # + done[0].numel()
for key in total_keys:
if key in (self.done_keys.union(self.reward_keys)):
entry = data.get(("next", key))
else:
entry = data.get(key)
if key in keys_to_expand:
shape = torch.Size([len_with_offset, *entry.shape[data.ndim :]])
dtype = entry.dtype
output.make_memmap(key, shape=shape, dtype=dtype)
else:
shape = torch.Size([data.numel(), *entry.shape[data.ndim :]])
output.make_memmap(key, shape=shape, dtype=entry.dtype)
if data.ndim == 1:
return self._call(
data=data,
output=output,
is_full=is_full,
shift=shift,
done=done,
total_keys=total_keys,
keys_to_expand=keys_to_expand,
)
with data.flatten(1, -1) if data.ndim > 2 else contextlib.nullcontext(
data
) as data_flat:
if data.ndim > 2:
done = done.flatten(1, -1)
traj_per_dim = done.sum(0)
nsteps = data_flat.shape[0]
start = 0
start_with_offset = start
stop_with_offset = 0
stop = 0
for data_slice, done_slice, traj_for_dim in zip(
data_flat.unbind(1), done.unbind(1), traj_per_dim
):
stop_with_offset = stop_with_offset + nsteps + traj_for_dim
cur_slice_offset = slice(start_with_offset, stop_with_offset)
start_with_offset = stop_with_offset
stop = stop + data.shape[0]
cur_slice = slice(start, stop)
start = stop
def _index(
key,
val,
keys_to_expand=keys_to_expand,
cur_slice=cur_slice,
cur_slice_offset=cur_slice_offset,
):
if key in keys_to_expand:
return val[cur_slice_offset]
return val[cur_slice]
out_slice = output.named_apply(_index, nested_keys=True)
self._call(
data=data_slice,
output=out_slice,
is_full=is_full,
shift=shift,
done=done_slice,
total_keys=total_keys,
keys_to_expand=keys_to_expand,
)
return output
def _call(self, *, data, output, is_full, shift, done, total_keys, keys_to_expand):
# capture for each item in data where the observation should be written
idx = torch.arange(data.shape[0])
idx_done = (idx + done.cumsum(0))[done]
idx += torch.nn.functional.pad(done, [1, 0])[:-1].cumsum(0)
for key in total_keys:
if key in (self.done_keys.union(self.reward_keys)):
entry = data.get(("next", key))
else:
entry = data.get(key)
if key in keys_to_expand:
mmap = output.get(key)
shifted_next = data.get(("next", key))
if is_full:
_roll_inplace(entry, shift=-shift, out=mmap, index_dest=idx)
_roll_inplace(
shifted_next,
shift=-shift,
out=mmap,
index_dest=idx_done,
index_source=done,
)
else:
mmap[idx] = entry
mmap[idx_done] = shifted_next[done]
elif is_full:
mmap = output.get(key)
_roll_inplace(entry, shift=-shift, out=mmap)
else:
mmap = output.get(key)
mmap.copy_(entry)
return output
[docs]
class Flat2TED:
"""A storage loading hook to deserialize flattened TED data to TED format.
Args:
done_key (NestedKey, optional): the key where the done states should be read.
Defaults to ``("next", "done")``.
shift_key (NestedKey, optional): the key where the shift will be written.
Defaults to "shift".
is_full_key (NestedKey, optional): the key where the is_full attribute will be written.
Defaults to "is_full".
done_keys (Tuple[NestedKey], optional): a tuple of nested keys indicating the done entries.
Defaults to :data:`~torchrl.data.DEFAULT_DONE_KEYS`, i.e.
``("done", "truncated", "terminated")``.
reward_keys (Tuple[NestedKey], optional): a tuple of nested keys indicating the reward entries.
Defaults to ("reward",)
Examples:
>>> import tempfile
>>>
>>> from tensordict import TensorDict
>>>
>>> from torchrl.collectors import Collector
>>> from torchrl.data import ReplayBuffer, TED2Flat, LazyMemmapStorage, Flat2TED
>>> from torchrl.envs import GymEnv
>>> import torch
>>>
>>> env = GymEnv("CartPole-v1")
>>> env.set_seed(0)
>>> torch.manual_seed(0)
>>> collector = Collector(env, policy=env.rand_step, total_frames=200, frames_per_batch=200)
>>> rb = ReplayBuffer(storage=LazyMemmapStorage(200))
>>> rb.register_save_hook(TED2Flat())
>>> with tempfile.TemporaryDirectory() as tmpdir:
... for i, data in enumerate(collector):
... rb.extend(data)
... rb.dumps(tmpdir)
... # load the data to represent it
... td = TensorDict.load(tmpdir + "/storage/")
...
... rb_load = ReplayBuffer(storage=LazyMemmapStorage(200))
... rb_load.register_load_hook(Flat2TED())
... rb_load.load(tmpdir)
... print("storage after loading", rb_load[:])
... assert (rb[:] == rb_load[:]).all()
storage after loading TensorDict(
fields={
action: MemoryMappedTensor(shape=torch.Size([200, 2]), device=cpu, dtype=torch.int64, is_shared=False),
collector: TensorDict(
fields={
traj_ids: MemoryMappedTensor(shape=torch.Size([200]), device=cpu, dtype=torch.int64, is_shared=False)},
batch_size=torch.Size([200]),
device=cpu,
is_shared=False),
done: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.bool, is_shared=False),
next: TensorDict(
fields={
done: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.bool, is_shared=False),
observation: MemoryMappedTensor(shape=torch.Size([200, 4]), device=cpu, dtype=torch.float32, is_shared=False),
reward: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.float32, is_shared=False),
terminated: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.bool, is_shared=False),
truncated: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
batch_size=torch.Size([200]),
device=cpu,
is_shared=False),
observation: MemoryMappedTensor(shape=torch.Size([200, 4]), device=cpu, dtype=torch.float32, is_shared=False),
terminated: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.bool, is_shared=False),
truncated: MemoryMappedTensor(shape=torch.Size([200, 1]), device=cpu, dtype=torch.bool, is_shared=False)},
batch_size=torch.Size([200]),
device=cpu,
is_shared=False)
"""
def __init__(
self,
done_key="done",
shift_key="shift",
is_full_key="is_full",
done_keys=DEFAULT_DONE_KEYS,
reward_keys=("reward",),
):
self.done_key = done_key
self.shift_key = shift_key
self.is_full_key = is_full_key
self.done_keys = {unravel_key(key) for key in done_keys}
self.reward_keys = {unravel_key(key) for key in reward_keys}
def __call__(self, data: TensorDictBase, out: TensorDictBase = None):
_storage_shape = data.get_non_tensor("_storage_shape", default=None)
if isinstance(_storage_shape, int):
_storage_shape = torch.Size([_storage_shape])
shift = data.get_non_tensor(self.shift_key, default=None)
is_full = data.get_non_tensor(self.is_full_key, default=None)
done = (
data.get("done")
.reshape((*_storage_shape[1:], -1))
.contiguous()
.permute(-1, *range(0, len(_storage_shape) - 1))
.clone()
)
if not is_full:
# shift is the cursor place
done[shift - 1] = True
else:
# done = done.roll(-shift, dims=0)
done[-1] = True
if _storage_shape is not None and len(_storage_shape) > 1:
# iterate over data and allocate
if out is None:
# out = TensorDict(batch_size=_storage_shape)
# for i in range(out.ndim):
# if i >= 2:
# # FLattening the lazy stack will make the data unavailable - we need to find a way to make this
# # possible.
# raise RuntimeError(
# "Checkpointing an uninitialized buffer with more than 2 dimensions is currently not supported. "
# "Please file an issue on GitHub to ask for this feature!"
# )
# out = LazyStackedTensorDict(*out.unbind(i), stack_dim=i)
out = TensorDict(batch_size=_storage_shape)
for i in range(1, out.ndim):
if i >= 2:
# FLattening the lazy stack will make the data unavailable - we need to find a way to make this
# possible.
raise RuntimeError(
"Checkpointing an uninitialized buffer with more than 2 dimensions is currently not supported. "
"Please file an issue on GitHub to ask for this feature!"
)
out_list = [
out._get_sub_tensordict((slice(None),) * i + (j,))
for j in range(out.shape[i])
]
out = lazy_stack(out_list, i)
# Create a function that reads slices of the input data
with out.flatten(1, -1) if out.ndim > 2 else contextlib.nullcontext(
out
) as out_flat:
nsteps = done.shape[0]
n_elt_batch = done.shape[1:].numel()
traj_per_dim = done.sum(0)
start = 0
start_with_offset = start
stop_with_offset = 0
stop = 0
for out_unbound, traj_for_dim in zip(out_flat.unbind(-1), traj_per_dim):
stop_with_offset = stop_with_offset + nsteps + traj_for_dim
cur_slice_offset = slice(start_with_offset, stop_with_offset)
start_with_offset = stop_with_offset
stop = stop + nsteps
cur_slice = slice(start, stop)
start = stop
def _index(
key,
val,
cur_slice=cur_slice,
nsteps=nsteps,
n_elt_batch=n_elt_batch,
cur_slice_offset=cur_slice_offset,
):
if val.shape[0] != (nsteps * n_elt_batch):
return val[cur_slice_offset]
return val[cur_slice]
data_slice = data.named_apply(
_index, nested_keys=True, batch_size=[]
)
self._call(
data=data_slice,
out=out_unbound,
is_full=is_full,
shift=shift,
_storage_shape=_storage_shape,
)
return out
return self._call(
data=data,
out=out,
is_full=is_full,
shift=shift,
_storage_shape=_storage_shape,
)
def _call(self, *, data, out, _storage_shape, shift, is_full):
done = data.get(self.done_key)
done = done.clone()
nsteps = done.shape[0]
# capture for each item in data where the observation should be written
idx = torch.arange(done.shape[0])
padded_done = F.pad(done.squeeze(-1), [1, 0])
root_idx = idx + padded_done[:-1].cumsum(0)
next_idx = root_idx + 1
if out is None:
out = TensorDict(batch_size=[nsteps])
def maybe_roll(entry, out=None):
if is_full and shift is not None:
if out is not None:
_roll_inplace(entry, shift=shift, out=out)
return
else:
return entry.roll(shift, dims=0)
if out is not None:
out.copy_(entry)
return
return entry
root_idx = maybe_roll(root_idx)
next_idx = maybe_roll(next_idx)
if not is_full:
next_idx = next_idx[:-1]
for key, entry in data.items(True, True):
if entry.shape[0] == nsteps:
if key in (self.done_keys.union(self.reward_keys)):
if key != "reward" and key not in out.keys(True, True):
# Create a done state at the root full of 0s
out.set(key, torch.zeros_like(entry), inplace=True)
entry = maybe_roll(entry, out=out.get(("next", key), None))
if entry is not None:
out.set(("next", key), entry, inplace=True)
else:
# action and similar
entry = maybe_roll(entry, out=out.get(key, default=None))
if entry is not None:
# then out is not locked
out.set(key, entry, inplace=True)
else:
dest_next = out.get(("next", key), None)
if dest_next is not None:
if not is_full:
dest_next = dest_next[:-1]
dest_next.copy_(entry[next_idx])
else:
if not is_full:
val = entry[next_idx]
val = torch.cat([val, torch.zeros_like(val[:1])])
out.set(("next", key), val, inplace=True)
else:
out.set(("next", key), entry[next_idx], inplace=True)
dest = out.get(key, None)
if dest is not None:
dest.copy_(entry[root_idx])
else:
out.set(key, entry[root_idx], inplace=True)
return out
class TED2Nested(TED2Flat):
"""Converts a TED-formatted dataset into a tensordict populated with nested tensors where each row is a trajectory.
.. seealso::
Trajectory lengths are recovered with
:func:`~torchrl.data.find_start_stop_traj`; see :ref:`the
trajectory-boundary documentation <ref_traj_boundaries>` for the
conventions this class relies on.
"""
_shift: int | None = None
_is_full: bool | None = None
def __init__(self, *args, **kwargs):
if not hasattr(torch, "_nested_compute_contiguous_strides_offsets"):
raise ValueError(
f"Unsupported torch version {torch.__version__}. "
f"torch>=2.4 is required for {type(self).__name__} to be used."
)
return super().__init__(*args, **kwargs)
def __call__(self, data: TensorDictBase, path: Path = None):
data = super().__call__(data, path=path)
shift = self.shift
is_full = self.is_full
storage_shape = data.get_non_tensor("_storage_shape", (-1,))
# place time at the end
storage_shape = (*storage_shape[1:], storage_shape[0])
done = data.get("done")
done = done.squeeze(-1).clone()
if not is_full:
done.view(storage_shape)[..., shift - 1] = True
# else:
done.view(storage_shape)[..., -1] = True
# The ring has been linearized by TED2Flat and the final boundary of
# every row was forced above, so the flat end flags fully delimit the
# trajectories.
_, _, traj_lengths = find_start_stop_traj(end=done, at_capacity=False)
ntraj = traj_lengths.numel()
keys_to_expand, keys_to_keep = zip(
*[
(key, None) if val.shape[0] != done.shape[0] else (None, key)
for key, val in data.items(True, True)
]
)
keys_to_expand = [key for key in keys_to_expand if key is not None]
keys_to_keep = [key for key in keys_to_keep if key is not None]
out = TensorDict(batch_size=[ntraj])
out.update(dict(data.non_tensor_items()))
out.memmap_(path)
traj_lengths = traj_lengths.unsqueeze(-1)
if not is_full:
# Increment by one only the trajectories that are not terminal
traj_lengths_expand = traj_lengths + (
traj_lengths.cumsum(0) % storage_shape[-1] != 0
)
else:
traj_lengths_expand = traj_lengths + 1
for key in keys_to_expand:
val = data.get(key)
shape = torch.cat(
[
traj_lengths_expand,
torch.tensor(val.shape[1:], dtype=torch.long).repeat(
traj_lengths.numel(), 1
),
],
-1,
)
# This works because the storage location is the same as the previous one - no copy is done
# but a new shape is written
out.make_memmap_from_storage(
key, val.untyped_storage(), dtype=val.dtype, shape=shape
)
for key in keys_to_keep:
val = data.get(key)
shape = torch.cat(
[
traj_lengths,
torch.tensor(val.shape[1:], dtype=torch.long).repeat(
traj_lengths.numel(), 1
),
],
-1,
)
out.make_memmap_from_storage(
key, val.untyped_storage(), dtype=val.dtype, shape=shape
)
return out
class Nested2TED(Flat2TED):
"""Converts a nested tensordict where each row is a trajectory into the TED format."""
def __call__(self, data, out: TensorDictBase = None):
# Get a flat representation of data
def flatten_het_dim(tensor):
shape = [tensor.size(i) for i in range(2, tensor.ndim)]
tensor = torch.tensor(tensor.untyped_storage(), dtype=tensor.dtype).view(
-1, *shape
)
return tensor
data = data.apply(flatten_het_dim, batch_size=[])
data.auto_batch_size_()
return super().__call__(data, out=out)
class H5Split(TED2Flat):
"""Splits a dataset prepared with TED2Nested into a TensorDict where each trajectory is stored as views on their parent nested tensors."""
_shift: int | None = None
_is_full: bool | None = None
def __call__(self, data):
nzeros = int(math.ceil(math.log10(data.shape[0])))
result = TensorDict(
{
f"traj_{str(i).zfill(nzeros)}": _data
for i, _data in enumerate(data.filter_non_tensor_data().unbind(0))
}
).update(dict(data.non_tensor_items()))
return result
class H5Combine:
"""Combines trajectories in a persistent tensordict into a single standing tensordict stored in filesystem."""
def __call__(self, data, out=None):
# TODO: this load the entire H5 in memory, which can be problematic
# Ideally we would want to load it on a memmap tensordict
# We currently ignore out in this call but we should leverage that
values = [val for key, val in data.items() if key.startswith("traj")]
metadata_keys = [key for key in data.keys() if not key.startswith("traj")]
result = TensorDict({key: NonTensorData(data[key]) for key in metadata_keys})
# Create a memmap in file system (no files associated)
result.memmap_()
# Create each entry
def initialize(key, *x):
result.make_memmap(
key,
shape=torch.stack([torch.tensor(_x.shape) for _x in x]),
dtype=x[0].dtype,
)
return
values[0].named_apply(
initialize,
*values[1:],
nested_keys=True,
batch_size=[],
filter_empty=True,
)
# Populate the entries
def populate(key, *x):
dest = result.get(key)
for i, _x in enumerate(x):
dest[i].copy_(_x)
values[0].named_apply(
populate,
*values[1:],
nested_keys=True,
batch_size=[],
filter_empty=True,
)
return result
@implement_for("torch", "2.3", None)
def _path2str(path, default_name=None):
# Uses the Keys defined in pytree to build a path
from torch.utils._pytree import MappingKey, SequenceKey
if default_name is None:
default_name = SINGLE_TENSOR_BUFFER_NAME
if not path:
return default_name
if isinstance(path, tuple):
return "/".join([_path2str(_sub, default_name=default_name) for _sub in path])
if isinstance(path, MappingKey):
if not isinstance(path.key, (int, str, bytes)):
raise ValueError("Values must be of type int, str or bytes in PyTree maps.")
result = str(path.key)
if result == default_name:
raise RuntimeError(
"A tensor had the same identifier as the default name used when the buffer contains "
f"a single tensor (name={default_name}). This behavior is not allowed. Please rename your "
f"tensor in the map/dict or set a new default name with the environment variable SINGLE_TENSOR_BUFFER_NAME."
)
return result
if isinstance(path, SequenceKey):
return str(path.idx)
@implement_for("torch", None, "2.3")
def _path2str(path, default_name=None): # noqa: F811
raise RuntimeError
def _save_pytree_common(tensor_path, path, tensor, metadata):
if "." in tensor_path:
tensor_path.replace(".", "_<dot>_")
total_tensor_path = path / (tensor_path + ".memmap")
if os.path.exists(total_tensor_path):
MemoryMappedTensor.from_filename(
shape=tensor.shape,
filename=total_tensor_path,
dtype=tensor.dtype,
).copy_(tensor)
else:
os.makedirs(total_tensor_path.parent, exist_ok=True)
MemoryMappedTensor.from_tensor(
tensor,
filename=total_tensor_path,
copy_existing=True,
copy_data=True,
)
key = tensor_path.replace("/", ".")
if key in metadata:
raise KeyError(
"At least two values have conflicting representations in "
f"the data structure to be serialized: {key}."
)
metadata[key] = {
"dtype": str(tensor.dtype),
"shape": list(tensor.shape),
}
@implement_for("torch", "2.3", None)
def _save_pytree(_storage, metadata, path):
from torch.utils._pytree import tree_map_with_path
def save_tensor(
tensor_path: tuple, tensor: torch.Tensor, metadata=metadata, path=path
):
tensor_path = _path2str(tensor_path)
_save_pytree_common(tensor_path, path, tensor, metadata)
tree_map_with_path(save_tensor, _storage)
@implement_for("torch", None, "2.3")
def _save_pytree(_storage, metadata, path): # noqa: F811
flat_storage, storage_specs = tree_flatten(_storage)
storage_paths = _get_paths(storage_specs)
def save_tensor(
tensor_path: str, tensor: torch.Tensor, metadata=metadata, path=path
):
_save_pytree_common(tensor_path, path, tensor, metadata)
for tensor, tensor_path in zip(flat_storage, storage_paths):
save_tensor(tensor_path, tensor)
def _get_paths(spec, cumulpath=""):
# alternative way to build a path without the keys
if isinstance(spec, LeafSpec):
yield cumulpath if cumulpath else SINGLE_TENSOR_BUFFER_NAME
contexts = spec.context
children_specs = spec.children_specs
if contexts is None:
contexts = range(len(children_specs))
for context, spec in zip(contexts, children_specs):
cpath = "/".join((cumulpath, str(context))) if cumulpath else str(context)
yield from _get_paths(spec, cpath)
def _init_pytree_common(tensor_path, scratch_dir, max_size_fn, tensor):
if "." in tensor_path:
tensor_path.replace(".", "_<dot>_")
if scratch_dir is not None:
total_tensor_path = Path(scratch_dir) / (tensor_path + ".memmap")
if os.path.exists(total_tensor_path):
raise RuntimeError(
f"The storage of tensor {total_tensor_path} already exists. "
f"To load an existing replay buffer, use storage.loads. "
f"Choose a different path to store your buffer or delete the existing files."
)
os.makedirs(total_tensor_path.parent, exist_ok=True)
else:
total_tensor_path = None
out = MemoryMappedTensor.empty(
shape=max_size_fn(tensor.shape),
filename=total_tensor_path,
dtype=tensor.dtype,
)
try:
filesize = os.path.getsize(tensor.filename) / 1024 / 1024
torchrl_logger.debug(
f"The storage was created in {out.filename} and occupies {filesize} Mb of storage."
)
except (RuntimeError, AttributeError):
pass
return out
@implement_for("torch", "2.3", None)
def _init_pytree(scratch_dir, max_size_fn, data):
from torch.utils._pytree import tree_map_with_path
# If not a tensorclass/tensordict, it must be a tensor(-like) or a PyTree
# if Tensor, we just create a MemoryMappedTensor of the desired shape, device and dtype
def save_tensor(tensor_path: tuple, tensor: torch.Tensor):
tensor_path = _path2str(tensor_path)
return _init_pytree_common(tensor_path, scratch_dir, max_size_fn, tensor)
out = tree_map_with_path(save_tensor, data)
return out
@implement_for("torch", None, "2.3")
def _init_pytree(scratch_dir, max_size, data): # noqa: F811
flat_data, data_specs = tree_flatten(data)
data_paths = _get_paths(data_specs)
data_paths = list(data_paths)
# If not a tensorclass/tensordict, it must be a tensor(-like) or a PyTree
# if Tensor, we just create a MemoryMappedTensor of the desired shape, device and dtype
def save_tensor(tensor_path: str, tensor: torch.Tensor):
return _init_pytree_common(tensor_path, scratch_dir, max_size, tensor)
out = []
for tensor, tensor_path in zip(flat_data, data_paths):
out.append(save_tensor(tensor_path, tensor))
return tree_unflatten(out, data_specs)
def _roll_inplace(tensor, shift, out, index_dest=None, index_source=None):
# slice 0
source0 = tensor[:-shift]
if index_source is not None:
source0 = source0[index_source[shift:]]
slice0_shift = source0.shape[0]
if index_dest is not None:
out[index_dest[-slice0_shift:]] = source0
else:
slice0 = out[-slice0_shift:]
slice0.copy_(source0)
# slice 1
source1 = tensor[-shift:]
if index_source is not None:
source1 = source1[index_source[:shift]]
if index_dest is not None:
out[index_dest[:-slice0_shift]] = source1
else:
slice1 = out[:-slice0_shift]
slice1.copy_(source1)
return out
# Copy-paste of unravel-index for PT 2.0
def _unravel_index(
indices: Tensor, shape: int | typing.Sequence[int] | torch.Size
) -> tuple[Tensor, ...]:
res_tensor = _unravel_index_impl(indices, shape)
return res_tensor.unbind(-1)
def _unravel_index_impl(indices: Tensor, shape: int | typing.Sequence[int]) -> Tensor:
if isinstance(shape, (int, torch.SymInt)):
shape = torch.Size([shape])
else:
shape = torch.Size(shape)
coefs = list(
reversed(
list(
itertools.accumulate(
reversed(shape[1:] + torch.Size([1])), func=operator.mul
)
)
)
)
return indices.unsqueeze(-1).floor_divide(
torch.tensor(coefs, device=indices.device, dtype=torch.int64)
) % torch.tensor(shape, device=indices.device, dtype=torch.int64)
@implement_for("torch", None, "2.2")
def unravel_index(indices, shape):
"""A version-compatible wrapper around torch.unravel_index."""
return _unravel_index(indices, shape)
@implement_for("torch", "2.2")
def unravel_index(indices, shape): # noqa: F811
"""A version-compatible wrapper around torch.unravel_index."""
return torch.unravel_index(indices, shape)
@implement_for("torch", None, "2.3")
def tree_iter(pytree):
"""A version-compatible wrapper around tree_iter."""
flat_tree, _ = torch.utils._pytree.tree_flatten(pytree)
yield from flat_tree
@implement_for("torch", "2.3", "2.4")
def tree_iter(pytree): # noqa: F811
"""A version-compatible wrapper around tree_iter."""
yield from torch.utils._pytree.tree_leaves(pytree)
@implement_for("torch", "2.4")
def tree_iter(pytree): # noqa: F811
"""A version-compatible wrapper around tree_iter."""
yield from torch.utils._pytree.tree_iter(pytree)
def _auto_device() -> torch.device:
if torch.cuda.is_available():
return torch.device("cuda:0")
elif torch.mps.is_available():
return torch.device("mps:0")
return torch.device("cpu")
