Table of Contents
Shortcuts

Binary

class torchrl.data.Binary(n: int | None = None, shape: torch.Size | None = None, device: DEVICE_TYPING | None = None, dtype: str | torch.dtype = torch.int8)[source]

A binary discrete tensor spec.

A binary tensor spec encodes tensors of arbitrary size where the values are either 0 or 1 (or True or False if the dtype it torch.bool).

Unlike OneHot, Binary can have more than one non-null element along the last dimension.

Parameters:

  • n (int) –

    length of the binary vector. If provided along with shape, shape[-1] must match n. If not provided, shape must be passed.

    Warning

    the n argument from Binary must not be confused with the n argument from Categorical or OneHot which denotes the maximum number of elements that can be sampled. For clarity, use shape instead.

  • shape (torch.Size, optional) – total shape of the sampled tensors. If provided, the last dimension must match n.

  • device (str, int or torch.device, optional) – device of the tensors.

  • dtype (str or torch.dtype, optional) – dtype of the tensors. Defaults to torch.int8.

Examples

>>> torch.manual_seed(0)
>>> spec = Binary(n=4, shape=(2, 4))
>>> print(spec.rand())
tensor([[0, 1, 1, 0],
        [1, 1, 1, 1]], dtype=torch.int8)
>>> spec = Binary(shape=(2, 4))
>>> print(spec.rand())
tensor([[1, 1, 1, 0],
        [0, 1, 0, 0]], dtype=torch.int8)
>>> spec = Binary(n=4)
>>> print(spec.rand())
tensor([0, 0, 0, 1], dtype=torch.int8)
assert_is_in(value: Tensor) → None

Asserts whether a tensor belongs to the box, and raises an exception otherwise.

Parameters:

value (torch.Tensor) – value to be checked.

cardinality() → int

The cardinality of the spec.

This refers to the number of possible outcomes in a spec. It is assumed that the cardinality of a composite spec is the cartesian product of all possible outcomes.

clear_device_() → T

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

For Composite specs, this method will erase the device.

clone() → Binary[source]

Creates a copy of the TensorSpec.

contains(item: torch.Tensor | TensorDictBase) → bool

If the value val could have been generated by the TensorSpec, returns True, otherwise False.

See is_in() for more information.

cpu()

Casts the TensorSpec to ‘cpu’ device.

cuda(device=None)

Casts the TensorSpec to ‘cuda’ device.

device: torch.device | None = None
encode(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.

This method is to be used in environments that return a value (eg, a numpy array) that can be easily mapped to the TorchRL required domain. If the value is already a tensor, the spec will not change its value and return it as-is.

Parameters:

val (np.ndarray or torch.Tensor) – value to be encoded as tensor.

Keyword Arguments:

ignore_device (bool, optional) – if True, the spec device will be ignored. This is used to group tensor casting within a call to TensorDict(..., device="cuda") which is faster.

Returns:

torch.Tensor matching the required tensor specs.

enumerate(use_mask: bool = False) → Tensor

Returns all the samples that can be obtained from the TensorSpec.

The samples will be stacked along the first dimension.

This method is only implemented for discrete specs.

Parameters:

use_mask (bool, optional) – If True and the spec has a mask, samples that are masked are excluded. Default is False.

erase_memoize_cache() → None

Clears the memoized cache for cached encode execution.

See also

memoize_encode().

expand(*shape)[source]

Returns a new Spec with the expanded shape.

Parameters:

*shape (tuple or iterable of int) – the new shape of the Spec. Must be broadcastable with the current shape: its length must be at least as long as the current shape length, and its last values must be compliant too; ie they can only differ from it if the current dimension is a singleton.

flatten(start_dim: int, end_dim: int) → T

Flattens a TensorSpec.

Check flatten() for more information on this method.

classmethod implements_for_spec(torch_function: Callable) → Callable

Register a torch function override for TensorSpec.

index(index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase) → torch.Tensor | TensorDictBase

Indexes the input tensor.

This method is to be used with specs that encode one or more categorical variables (e.g., OneHot or Categorical), such that indexing of a tensor with a sample can be done without caring about the actual representation of the index.

Parameters:

  • index (int, torch.Tensor, slice or list) – index of the tensor

  • tensor_to_index – tensor to be indexed

Returns:

indexed tensor

Exanples:
>>> from torchrl.data import OneHot
>>> import torch
>>>
>>> one_hot = OneHot(n=100)
>>> categ = one_hot.to_categorical_spec()
>>> idx_one_hot = torch.zeros((100,), dtype=torch.bool)
>>> idx_one_hot[50] = 1
>>> print(one_hot.index(idx_one_hot, torch.arange(100)))
tensor(50)
>>> idx_categ = one_hot.to_categorical(idx_one_hot)
>>> print(categ.index(idx_categ, torch.arange(100)))
tensor(50)
is_in(val: Tensor) → bool

If the value val could have been generated by the TensorSpec, returns True, otherwise False.

More precisely, the is_in methods checks that the value val is within the limits defined by the space attribute (the box), and that the dtype, device, shape potentially other metadata match those of the spec. If any of these checks fails, the is_in method will return False.

Parameters:

val (torch.Tensor) – value to be checked.

Returns:

boolean indicating if values belongs to the TensorSpec box.

make_neg_dim(dim: int) → T

Converts a specific dimension to -1.

memoize_encode(mode: bool = True) → None

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

This should only be used whenever the input type, shape etc. are expected to be consistent across calls for a given spec.

Parameters:

mode (bool, optional) – Whether the cache should be used. Defaults to True.

See also

the cache can be erased via erase_memoize_cache().

property ndim: int

Number of dimensions of the spec shape.

Shortcut for len(spec.shape).

ndimension() → int

Number of dimensions of the spec shape.

Shortcut for len(spec.shape).

one(shape: torch.Size = None) → torch.Tensor | TensorDictBase

Returns a one-filled tensor in the box.

Note

Even though there is no guarantee that 1 belongs to the spec domain, this method will not raise an exception when this condition is violated. The primary use case of one is to generate empty data buffers, not meaningful data.

Parameters:

shape (torch.Size) – shape of the one-tensor

Returns:

a one-filled tensor sampled in the TensorSpec box.

ones(shape: torch.Size = None) → torch.Tensor | TensorDictBase

Proxy to one().

project(val: torch.Tensor | TensorDictBase) → torch.Tensor | TensorDictBase

If the input tensor is not in the TensorSpec box, it maps it back to it given some defined heuristic.

Parameters:

val (torch.Tensor) – tensor to be mapped to the box.

Returns:

a torch.Tensor belonging to the TensorSpec box.

rand(shape: Optional[Size] = None) → Tensor

Returns a random tensor in the space defined by the spec.

The sampling will be done uniformly over the space, unless the box is unbounded in which case normal values will be drawn.

Parameters:

shape (torch.Size) – shape of the random tensor

Returns:

a random tensor sampled in the TensorSpec box.

reshape(*shape) → T

Reshapes a TensorSpec.

Check reshape() for more information on this method.

sample(shape: torch.Size = None) → torch.Tensor | TensorDictBase

Returns a random tensor in the space defined by the spec.

See rand() for details.

set_provisional_n(n: int)

Set the cardinality of the Categorical spec temporarily.

This method is required to be called before sampling from the spec when n is -1.

Parameters:

n (int) – The cardinality of the Categorical spec.

squeeze(dim=None)[source]

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

When dim is given, a squeeze operation is done only in that dimension.

Parameters:

dim (int or None) – the dimension to apply the squeeze operation to

to(dest: torch.dtype | DEVICE_TYPING) → Binary[source]

Casts a TensorSpec to a device or a dtype.

Returns the same spec if no change is made.

to_categorical(val: torch.Tensor, safe: bool | None = None) → torch.Tensor

No-op for categorical.

to_categorical_spec() → Categorical

No-op for categorical.

to_numpy(val: torch.Tensor, safe: bool | None = None) → dict

Returns the np.ndarray correspondent of an input tensor.

This is intended to be the inverse operation of encode().

Parameters:

  • val (torch.Tensor) – tensor to be transformed_in to numpy.

  • safe (bool) – boolean value indicating whether a check should be performed on the value against the domain of the spec. Defaults to the value of the CHECK_SPEC_ENCODE environment variable.

Returns:

a np.ndarray.

to_one_hot(val: torch.Tensor, safe: bool | None = None) → torch.Tensor

Encodes a discrete tensor from the spec domain into its one-hot correspondent.

Parameters:

  • val (torch.Tensor, optional) – Tensor to one-hot encode.

  • safe (bool) – boolean value indicating whether a check should be performed on the value against the domain of the spec. Defaults to the value of the CHECK_SPEC_ENCODE environment variable.

Returns:

The one-hot encoded tensor.

Examples

>>> categ = Categorical(3)
>>> categ_sample = categ.zero()
>>> categ_sample
tensor(0)
>>> onehot_sample = categ.to_one_hot(categ_sample)
>>> onehot_sample
tensor([ True, False, False])
to_one_hot_spec() → OneHot

Converts the spec to the equivalent one-hot spec.

Examples

>>> categ = Categorical(3)
>>> categ.to_one_hot_spec()
OneHot(
    shape=torch.Size([3]),
    space=CategoricalBox(n=3),
    device=cpu,
    dtype=torch.bool,
    domain=discrete)
type_check(value: Tensor, key: Optional[NestedKey] = None) → None

Checks the input value dtype against the TensorSpec dtype and raises an exception if they don’t match.

Parameters:

  • value (torch.Tensor) – tensor whose dtype has to be checked.

  • key (str, optional) – if the TensorSpec has keys, the value dtype will be checked against the spec pointed by the indicated key.

unflatten(dim: int, sizes: tuple[int]) → T

Unflattens a TensorSpec.

Check unflatten() for more information on this method.

unsqueeze(dim: int)[source]

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

Parameters:

dim (int or None) – the dimension to apply the unsqueeze operation to.

update_mask(mask)

Sets a mask to prevent some of the possible outcomes when a sample is taken.

The mask can also be set during initialization of the spec.

Parameters:

mask (torch.Tensor or None) – boolean mask. If None, the mask is disabled. Otherwise, the shape of the mask must be expandable to the shape of the equivalent one-hot spec. False masks an outcome and True leaves the outcome unmasked. If all of the possible outcomes are masked, then an error is raised when a sample is taken.

Examples

>>> mask = torch.tensor([True, False, True])
>>> ts = Categorical(3, (10,), dtype=torch.int64, mask=mask)
>>> # One of the three possible outcomes is masked
>>> ts.rand()
tensor([0, 2, 2, 0, 2, 0, 2, 2, 0, 2])
view(*shape) → T

Reshapes a TensorSpec.

Check reshape() for more information on this method.

zero(shape: torch.Size = None) → torch.Tensor | TensorDictBase

Returns a zero-filled tensor in the box.

Note

Even though there is no guarantee that 0 belongs to the spec domain, this method will not raise an exception when this condition is violated. The primary use case of zero is to generate empty data buffers, not meaningful data.

Parameters:

shape (torch.Size) – shape of the zero-tensor

Returns:

a zero-filled tensor sampled in the TensorSpec box.

zeros(shape: torch.Size = None) → torch.Tensor | TensorDictBase

Proxy to zero().

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

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.