BinaryDiscreteTensorSpec¶
- class torchrl.data.BinaryDiscreteTensorSpec(n: int, shape: Optional[Size] = None, device: Optional[Union[device, str, int]] = None, dtype: Union[str, dtype] = torch.int8)[source]¶
A binary discrete tensor spec.
- Parameters:
n (int) – length of the binary vector.
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.long.
Examples
>>> spec = BinaryDiscreteTensorSpec(n=4, shape=(5, 4), device="cpu", dtype=torch.bool) >>> print(spec.zero())
- 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], *, ignore_device=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 toTensorDict(..., device="cuda")
which is faster.- Returns:
torch.Tensor matching the required tensor specs.
- expand(*shape)[source]¶
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.
- abstract 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.
- 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_numpy(val: Tensor, safe: Optional[bool] = None) dict ¶
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
- to_one_hot(val: Tensor, safe: Optional[bool] = None) 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.
- to_one_hot_spec() OneHotDiscreteTensorSpec ¶
Converts the spec to the equivalent one-hot spec.
- 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 equivalent one-hot spec.
False
masks an outcome andTrue
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 = DiscreteTensorSpec(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])
- 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.