OneHotDiscreteTensorSpec

class torchrl.data.OneHotDiscreteTensorSpec(n: int, shape: Optional[torch.Size] = None, device: Optional[DEVICE_TYPING] = None, dtype: Optional[Union[str, torch.dtype]] = torch.bool, use_register: bool = False, mask: torch.Tensor | None = None)

A unidimensional, one-hot discrete tensor spec.

By default, TorchRL assumes that categorical variables are encoded as one-hot encodings of the variable. This allows for simple indexing of tensors, e.g.

>>> batch, size = 3, 4
>>> action_value = torch.arange(batch*size)
>>> action_value = action_value.view(batch, size).to(torch.float)
>>> action = (action_value == action_value.max(-1,
...    keepdim=True)[0]).to(torch.long)
>>> chosen_action_value = (action * action_value).sum(-1)
>>> print(chosen_action_value)
tensor([ 3.,  7., 11.])

The last dimension of the shape (variable domain) cannot be indexed.

Parameters:

• n (int) – number of possible outcomes.

• 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.

• user_register (bool) – experimental feature. If True, every integer will be mapped onto a binary vector in the order in which they appear. This feature is designed for environment with no a-priori definition of the number of possible outcomes (e.g. discrete outcomes are sampled from an arbitrary set, whose elements will be mapped in a register to a series of unique one-hot binary vectors).

• mask (torch.Tensor or None) – mask some of the possible outcomes when a sample is taken. See update_mask() for more information.

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.

clear_device_()

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

contains(item)

Returns whether a sample is contained within the space defined by the TensorSpec.

See is_in() for more information.

encode(val: Union[ndarray, Tensor], space: Optional[DiscreteBox] = None, *, ignore_device: bool = False) → Tensor

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

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.

expand(*shape)

Returns a new Spec with the extended shape.

Parameters:

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

flatten(start_dim, end_dim)

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: Union[int, Tensor, ndarray, slice, List], tensor_to_index: Tensor) → Tensor

Indexes the input tensor.

Parameters:

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

• tensor_to_index – tensor to be indexed

Returns:

indexed tensor

is_in(val: Tensor) → bool

If the value val is in the box defined by the TensorSpec, returns True, otherwise False.

Parameters:

val (torch.Tensor) – value to be checked

Returns:

boolean indicating if values belongs to the TensorSpec box

project(val: Tensor) → Tensor

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

Parameters:

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

Returns:

a torch.Tensor belonging to the TensorSpec box.

rand(shape=None) → Tensor

Returns a random tensor in the space defined by the spec. The sampling will be uniform unless the box is unbounded.

Parameters:

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

Returns:

a random tensor sampled in the TensorSpec box.

reshape(*shape)

Reshapes a tensorspec.

Check reshape() for more information on this method.

property sample

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

See rand() for details.

squeeze(dim=None)

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_categorical(val: Tensor, safe: Optional[bool] = None) → Tensor

Converts a given one-hot tensor in categorical format.

Parameters:

• val (torch.Tensor, optional) – One-hot tensor to convert in categorical format.

• 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 categorical tensor.

to_categorical_spec() → DiscreteTensorSpec

Converts the spec to the equivalent categorical spec.

to_numpy(val: Tensor, safe: Optional[bool] = None) → ndarray

Returns the np.ndarray correspondent of an input tensor.

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

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, sizes)

Unflattens a tensorspec.

Check unflatten() for more information on this method.

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 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, False])
>>> ts = OneHotDiscreteTensorSpec(3, (2, 3,), dtype=torch.int64, mask=mask)
>>> # All but one of the three possible outcomes are masked
>>> ts.rand()
tensor([[1, 0, 0],
        [1, 0, 0]])

view(*shape)

Reshapes a tensorspec.

Check reshape() for more information on this method.

zero(shape=None) → Tensor

Returns a zero-filled tensor in the box.

Parameters:

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

Returns:

a zero-filled tensor sampled in the TensorSpec box.