int,

range,

list[int],

np.ndarray,

slice,

None,

torch.Tensor,

type(...),

tuple[

int,

range,

list[int],

np.ndarray,

slice,

None,

torch.Tensor,

type(...),

tuple[Any],

],

"method is not currently implemented."

" If you are interested in this feature please submit"

" an issue at https://github.com/pytorch/rl/issues"

dtype: None | torch.dtype,

device: None | str | int | torch.device,

allow_none_device: bool = False,

) -> tuple[torch.dtype, torch.device | None]:

if dtype is None:

dtype = torch.get_default_dtype()

if device is not None:

device = _make_ordinal_device(torch.device(device))

elif not allow_none_device:

device = torch.zeros(()).device

if device is not None and device.type == "mps" and dtype == torch.float64:

warnings.warn(

"MPS device does not support float64. Downcasting dtype from float64 to float32.",

UserWarning,

stacklevel=2,

)

dtype = torch.float32

"""Raise an IndexError if idx is out of bounds for shape[axis].

if shape[axis] >= 0 and (idx >= shape[axis] or idx < 0 and -idx > shape[axis]):

raise IndexError(

f"index {idx} is out of bounds for axis {axis} with size {shape[axis]}"

idx: Iterable[Any], expected_type: type, iterable_classname: str

):

"""Raise an IndexError if the iterable contains a type different from the expected type or Iterable.

for item in idx:

if isinstance(item, Iterable):

_validate_iterable(item, expected_type, iterable_classname)

else:

if not isinstance(item, expected_type):

raise IndexError(

f"{iterable_classname} indexing expects {expected_type} indices"

"""Given an input shape and a slice index, returns the new indexed shape.

if idx.step == 0:

raise ValueError("slice step cannot be zero")

# Slicing an empty shape returns the shape

if len(shape) == 0:

return shape

if idx.start is None:

start = 0

else:

start = idx.start if idx.start >= 0 else max(shape[0] + idx.start, 0)

if idx.stop is None:

stop = shape[0]

else:

stop = idx.stop if idx.stop >= 0 else max(shape[0] + idx.stop, 0)

step = 1 if idx.step is None else idx.step

if step > 0:

if start >= stop:

n_items = 0

else:

stop = min(stop, shape[0])

n_items = math.ceil((stop - start) / step)

else:

if start <= stop:

n_items = 0

else:

start = min(start, shape[0] - 1)

n_items = math.ceil((stop - start) / step)

shape: list[int] | torch.Size | tuple[int], idx: SHAPE_INDEX_TYPING

) -> list[int]:

"""Given an input shape and an index, returns the size of the resulting indexed spec.

if not isinstance(shape, list):

shape = list(shape)

if idx is Ellipsis or (

isinstance(idx, slice) and (idx.step is idx.start is idx.stop is None)

):

return shape

if idx is None:

return [1] + shape

if len(shape) == 0 and (

isinstance(idx, int)

or isinstance(idx, range)

or isinstance(idx, list)

and len(idx) > 0

):

raise IndexError(

f"cannot use integer indices on 0-dimensional shape. `{idx}` received"

)

if isinstance(idx, int):

_validate_idx(shape, idx)

return shape[1:]

if isinstance(idx, range):

if len(idx) > 0 and (idx.start >= shape[0] or idx.stop > shape[0]):

raise IndexError(f"index out of bounds for axis 0 with size {shape[0]}")

return [len(idx)] + shape[1:]

if isinstance(idx, slice):

return _slice_indexing(shape, idx)

if isinstance(idx, tuple):

# Supports int, None, slice and ellipsis indices

# Index on the current shape dimension

shape_idx = 0

none_dims = 0

ellipsis = False

prev_is_list = False

shape_len = len(shape)

for item_idx, item in enumerate(idx):

if item is None:

shape = shape[:shape_idx] + [1] + shape[shape_idx:]

shape_idx += 1

none_dims += 1

elif isinstance(item, int):

_validate_idx(shape, item, shape_idx)

del shape[shape_idx]

elif isinstance(item, slice):

shape[shape_idx] = _slice_indexing([shape[shape_idx]], item)[0]

shape_idx += 1

elif item is Ellipsis:

if ellipsis:

raise IndexError("an index can only have a single ellipsis (`...`)")

# Move to the end of the shape, subtracted by the number of future indices impacting the dimensions (i.e. all except None and ...)

shape_idx = len(shape) - len(

[i for i in idx[item_idx + 1 :] if not (i is None or i is Ellipsis)]

)

ellipsis = True

elif any(

isinstance(item, _type)

for _type in [list, tuple, range, np.ndarray, torch.Tensor]

):

while isinstance(idx, tuple) and len(idx) == 1:

idx = idx[0]

# Nested tuples are handled as a list. Numpy behavior

if isinstance(item, tuple):

item = list(item)

if prev_is_list and isinstance(item, list):

del shape[shape_idx]

continue

if isinstance(item, list):

prev_is_list = True

if shape_idx >= len(shape):

raise IndexError("Raise IndexError: too many indices for array")

res = _shape_indexing([shape[shape_idx]], item)

shape = shape[:shape_idx] + res + shape[shape_idx + 1 :]

shape_idx += len(res)

else:

raise IndexError(

f"tuple indexing only supports integers, ranges, slices (`:`), ellipsis (`...`), new axis (`None`), tuples, list, tensor and ndarray indices. {str(type(idx))} received"

)

if len(idx) - none_dims - int(ellipsis) > shape_len:

raise IndexError(

f"shape is {shape_len}-dimensional, but {len(idx) - none_dims - int(ellipsis)} dimensions were indexed"

)

return shape

if isinstance(idx, list):

# int indexing only

_validate_iterable(idx, int, "list")

for item in np.array(idx).reshape(-1):

_validate_idx(shape, item, 0)

return list(np.array(idx).shape) + shape[1:]

if isinstance(idx, np.ndarray) or isinstance(idx, torch.Tensor):

# Out of bounds check

for item in idx.reshape(-1):

_validate_idx(shape, item)

"""An invertible dictionary.

def __init__(self, *args, inv_dict=None, **kwargs):

if inv_dict is None:

inv_dict = {}

super().__init__(*args, **kwargs)

self.inv_dict = inv_dict

def __setitem__(self, k, v):

if v in self.inv_dict or k in self:

raise Exception("overwriting in invertible_dict is not permitted")

self.inv_dict[v] = k

return super().__setitem__(k, v)

def update(self, d):

raise NotImplementedError

def invert(self):

d = invertible_dict()

for k, value in self.items():

d[value] = k

return d

def inverse(self):

"""A box of values."""

def __iter__(self):

raise NotImplementedError

def to(self, dest: torch.dtype | DEVICE_TYPING) -> ContinuousBox:

raise NotImplementedError

def __repr__(self):

return f"{self.__class__.__name__}()"

def clone(self) -> CategoricalBox:

"""A continuous box of values, in between a minimum (self.low) and a maximum (self.high)."""

_low: torch.Tensor

_high: torch.Tensor

device: torch.device | None = None

_batch_size: torch.Size | None = None

@property

def batch_size(self):

return self._batch_size

@batch_size.setter

def batch_size(self, value: torch.Size | tuple):

# Check batch size is compatible with low and high

value = _remove_neg_shapes(value)

if self._batch_size is None:

if value != self._low.shape[: len(value)]:

raise ValueError(

f"Batch size {value} is not compatible with low and high {self._low.shape}"

)

if value is None:

self._batch_size = None

self._low = self.low.clone()

self._high = self.high.clone()

return

# Remove batch size from low and high

if value:

# Check that low and high have a single value

td_low_high = TensorDict(

low=self.low, high=self.high, batch_size=value

).flatten()

td_low_high0 = td_low_high[0]

if torch.allclose(

td_low_high0["low"], td_low_high["low"]

) and torch.allclose(td_low_high0["high"], td_low_high["high"]):

self._low = td_low_high0["low"].clone()

self._high = td_low_high0["high"].clone()

self._batch_size = torch.Size(value)

else:

self._low = self.low.clone()

self._high = self.high.clone()

self._batch_size = torch.Size(value)

# We store the tensors on CPU to avoid overloading CUDA with tensors that are rarely used.

@property

def low(self):

low = self._low

if self.device is not None and low.device != self.device:

low = low.to(self.device)

if self._batch_size:

low = low.expand((*self._batch_size, *low.shape)).clone()

return low

@property

def high(self):

high = self._high

if self.device is not None and high.device != self.device:

high = high.to(self.device)

if self._batch_size:

high = high.expand((*self._batch_size, *high.shape)).clone()

return high

def unbind(self, dim: int = 0):

return tuple(

type(self)(low, high, self.device)

for (low, high) in zip(self.low.unbind(dim), self.high.unbind(dim))

)

@low.setter

def low(self, value):

self.device = value.device

self._low = value

if self._batch_size is not None:

if value.shape[: len(self._batch_size)] != self._batch_size:

raise ValueError(

f"Batch size {value.shape[: len(self._batch_size)]} is not compatible with low and high {self._batch_size}"

)

if self._batch_size:

self._low = self._low.flatten(0, len(self._batch_size) - 1)[0].clone()

@high.setter

def high(self, value):

self.device = value.device

self._high = value

if self._batch_size is not None:

if value.shape[: len(self._batch_size)] != self._batch_size:

raise ValueError(

f"Batch size {value.shape[: len(self._batch_size)]} is not compatible with low and high {self._batch_size}"

)

if self._batch_size:

self._high = self._high.flatten(0, len(self._batch_size) - 1)[0].clone()

def __post_init__(self):

self.low = self.low.clone()

self.high = self.high.clone()

self._batch_size = None

def __iter__(self):

yield self.low

yield self.high

def to(self, dest: torch.dtype | DEVICE_TYPING) -> ContinuousBox:

return self.__class__(self.low.to(dest), self.high.to(dest))

def clone(self) -> ContinuousBox:

return self.__class__(self.low.clone(), self.high.clone())

def __repr__(self):

min_str = indent(

f"\nlow=Tensor(shape={self.low.shape}, device={self.low.device}, dtype={self.low.dtype}, contiguous={self.high.is_contiguous()})",

" " * 4,

)

max_str = indent(

f"\nhigh=Tensor(shape={self.high.shape}, device={self.high.device}, dtype={self.high.dtype}, contiguous={self.high.is_contiguous()})",

" " * 4,

)

return f"{self.__class__.__name__}({min_str},{max_str})"

def __eq__(self, other):

if other is None:

minval, maxval = _minmax_dtype(self.low.dtype)

minval = torch.as_tensor(minval).to(self.low.device, self.low.dtype)

maxval = torch.as_tensor(maxval).to(self.low.device, self.low.dtype)

if (

torch.isclose(self.low, minval).all()

and torch.isclose(self.high, maxval).all()

):

return True

if (

not torch.isfinite(self.low).any()

and not torch.isfinite(self.high).any()

):

return True

return False

return (

type(self) is type(other)

and self.low.dtype == other.low.dtype

and self.high.dtype == other.high.dtype

and self.device == other.device

and torch.isclose(self.low, other.low).all()

and torch.isclose(self.high, other.high).all()

"""A box of discrete, categorical values."""

n: int

register = invertible_dict()

def __post_init__(self):

# n could be a numpy array or a tensor, making compile go a bit crazy

# We want to make sure we're working with a regular integer

self.__dict__["n"] = int(self.n)

def to(self, dest: torch.dtype | DEVICE_TYPING) -> CategoricalBox:

return deepcopy(self)

def __repr__(self):

"""A box of discrete values."""

boxes: list

def to(self, dest: torch.dtype | DEVICE_TYPING) -> BoxList:

return BoxList([box.to(dest) for box in self.boxes])

def __iter__(self):

yield from self.boxes

def __repr__(self):

return f"{self.__class__.__name__}(boxes={self.boxes})"

def __len__(self):

return len(self.boxes)

@staticmethod

def from_nvec(nvec: torch.Tensor):

if nvec.ndim == 0:

return CategoricalBox(nvec.item())

else:

"""A box of n binary values."""

n: int

def to(self, dest: torch.dtype | DEVICE_TYPING) -> ContinuousBox:

return deepcopy(self)

def __repr__(self):

"""Parent class of the tensor meta-data containers.

shape: torch.Size

space: None | Box

device: torch.device | None = None # noqa # type: ignore

dtype: torch.dtype = torch.get_default_dtype()

domain: str = ""

_encode_memo_dict: dict[Any, Callable[[Any], Any]] = field(

default_factory=dict,

)

SPEC_HANDLED_FUNCTIONS = {}

def memoize_encode(self, mode: bool = True) -> None:

"""Creates a cached sequence of callables for the `encode` method that speeds up its execution.

warnings.warn(

"memoized encoding is an experimental feature. Use at your own risks."

)

if mode:

self.encode = self._encode_memo

else:

self.encode = self._encode_eager

def erase_memoize_cache(self) -> None:

"""Clears the memoized cache for cached encode execution.

self._encode_memo_dict.clear()

def __getstate__(self):

state = dict(self.__dict__)

state["_encode"] = {}

# Clear device-specific bounds cache to avoid serializing CUDA tensors

state.pop("_bounds_cache", None)

return state

@classmethod

def implements_for_spec(cls, torch_function: Callable) -> Callable:

"""Register a torch function override for TensorSpec."""

@wraps(torch_function)

def decorator(func):

cls.SPEC_HANDLED_FUNCTIONS[torch_function] = func

return func

return decorator

@property

def device(self) -> torch.device: # noqa # type: ignore

"""The device of the spec.

return self._device

@device.setter

def device(self, device: torch.device | None) -> None:

self._device = _make_ordinal_device(device)

def clear_device_(self) -> T:

"""A no-op for all leaf specs (which must have a device).

return self

@abc.abstractmethod

def cardinality(self) -> int:

"""The cardinality of the spec.

...

def encode(

self,

val: np.ndarray | list | torch.Tensor | TensorDictBase,

*,

ignore_device: bool = False,

) -> torch.Tensor | TensorDictBase:

"""Encodes a value given the specified spec, and return the corresponding tensor.

raise NotImplementedError(

"This is a placeholder that needs to be set during construction"

)

def _encode_eager(

self,

val: np.ndarray | list | torch.Tensor | TensorDictBase,

*,

ignore_device: bool = False,

) -> torch.Tensor | TensorDictBase:

if not isinstance(val, torch.Tensor):

if isinstance(val, list):

if len(val) == 1:

# gym used to return lists of images since 0.26.0

# We convert these lists in np.array or take the first element

# if there is just one.

# See https://github.com/pytorch/rl/pull/403/commits/73d77d033152c61d96126ccd10a2817fecd285a1

val = val[0]

else:

val = np.array(val)

if isinstance(val, np.ndarray) and not all(

stride > 0 for stride in val.strides

):

val = val.copy()

if not ignore_device:

val = torch.as_tensor(val, device=self.device, dtype=self.dtype)

else:

val = torch.as_tensor(val, dtype=self.dtype)

if val.shape != self.shape:

# if val.shape[-len(self.shape) :] != self.shape:

# option 1: add a singleton dim at the end

if val.shape == self.shape and self.shape[-1] == 1:

val = val.unsqueeze(-1)

else:

try:

val = val.reshape(self.shape)

except Exception as err:

raise RuntimeError(

f"Shape mismatch: the value has shape {val.shape} which "

f"is incompatible with the spec shape {self.shape}."

) from err

if _CHECK_SPEC_ENCODE:

self.assert_is_in(val)

return val

def _encode_memo(

self,

val: np.ndarray | list | torch.Tensor | TensorDictBase,

*,

ignore_device: bool = False,

) -> torch.Tensor | TensorDictBase:

funcs = self._encode_memo_dict.get(ignore_device)

if funcs is not None:

return funcs(val)

funcs = []

val_orig = val

if not isinstance(val, torch.Tensor):

if isinstance(val, list):

if len(val) == 1:

# gym used to return lists of images since 0.26.0

# We convert these lists in np.array or take the first element

# if there is just one.

# See https://github.com/pytorch/rl/pull/403/commits/73d77d033152c61d96126ccd10a2817fecd285a1

funcs.append(lambda val: val[0])

else:

funcs.append(lambda val: np.array(val))

val = _reduce_funcs(funcs)(val_orig)

if isinstance(val, np.ndarray) and not all(

stride > 0 for stride in val.strides

):

funcs.append(lambda val: val.copy())

val = _reduce_funcs(funcs)(val_orig)

if not ignore_device:

funcs.append(

lambda val: torch.as_tensor(

val, device=self.device, dtype=self.dtype

)

)

else:

funcs.append(lambda val: torch.as_tensor(val, dtype=self.dtype))

val = _reduce_funcs(funcs)(val_orig)

if val.shape != self.shape:

# if val.shape[-len(self.shape) :] != self.shape:

# option 1: add a singleton dim at the end

if val.shape == self.shape and self.shape[-1] == 1:

funcs.append(lambda val: val.unsqueeze(-1))

else:

def reshape(val):

try:

return val.reshape(self.shape)

except Exception as err:

raise RuntimeError(

f"Shape mismatch: the value has shape {val.shape} which "

f"is incompatible with the spec shape {self.shape}."

) from err

funcs.append(reshape)

val = _reduce_funcs(funcs)(val_orig)

if _CHECK_SPEC_ENCODE:

def check(val):

self.assert_is_in(val)

return val

funcs.append(check)

if len(funcs) == 0:

self._encode_memo_dict[ignore_device] = lambda x: x

elif len(funcs) == 1:

self._encode_memo_dict[ignore_device] = funcs[0]

else:

self._encode_memo_dict[ignore_device] = _reduce_funcs(funcs)

return self._encode_memo_dict[ignore_device](val_orig)

@abc.abstractmethod

def __eq__(self, other: Any) -> bool:

# Implement minimal version if super() is called

return type(self) is type(other)

def __ne__(self, other):

return not (self == other)

def __setattr__(self, key, value):

if key == "shape":

value = _size(value)

super().__setattr__(key, value)

def to_numpy(

self, val: torch.Tensor | TensorDictBase, safe: bool | None = None

) -> np.ndarray | dict:

"""Returns the ``np.ndarray`` correspondent of an input tensor.

if safe is None:

safe = _CHECK_SPEC_ENCODE

if safe:

self.assert_is_in(val)

return val.detach().cpu().numpy()

@property

def ndim(self) -> int:

"""Number of dimensions of the spec shape.

return self.ndimension()

def ndimension(self) -> int:

"""Number of dimensions of the spec shape.

return len(self.shape)

@property

def _safe_shape(self) -> torch.Size:

"""Returns a shape where all heterogeneous values are replaced by one (to be expandable)."""

return _size([int(v) if v >= 0 else 1 for v in self.shape])

@abc.abstractmethod

def index(

self, index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase

) -> torch.Tensor | TensorDictBase:

"""Indexes the input tensor.

...

@overload

def expand(self, shape: torch.Size):

...

@abc.abstractmethod

def expand(self, *shape: int) -> T:

"""Returns a new Spec with the expanded shape.

...

def squeeze(self, dim: int | None = None) -> T:

"""Returns a new Spec with all the dimensions of size ``1`` removed.

shape = _squeezed_shape(self.shape, dim)

if shape is None:

return self

return self.__class__(shape=shape, device=self.device, dtype=self.dtype)

def unsqueeze(self, dim: int) -> T:

"""Returns a new Spec with one more singleton dimension (at the position indicated by ``dim``).

shape = _unsqueezed_shape(self.shape, dim)

return self.__class__(shape=shape, device=self.device, dtype=self.dtype)

def make_neg_dim(self, dim: int) -> T:

"""Converts a specific dimension to ``-1``."""

if dim < 0:

dim = self.ndim + dim

if dim < 0 or dim > self.ndim - 1:

raise ValueError(f"dim={dim} is out of bound for ndim={self.ndim}")

self.shape = _size([s if i != dim else -1 for i, s in enumerate(self.shape)])

@overload

def reshape(self, shape) -> T:

...

def reshape(self, *shape) -> T:

"""Reshapes a ``TensorSpec``.

if len(shape) == 1 and not isinstance(shape[0], int):