- class torchrl.data.TensorSpec(shape: Size, space: Union[None, Box], device: device = device(type='cpu'), dtype: dtype = torch.float32, domain: str = '')¶
Parent class of the tensor meta-data containers for observation, actions and rewards.
- assert_is_in(value: Tensor) None ¶
Asserts whether a tensor belongs to the box, and raises an exception otherwise.
value (torch.Tensor) – value to be checked.
- encode(val: Union[ndarray, Tensor], *, ignore_device=False) Tensor ¶
Encodes a value given the specified spec, and return the corresponding tensor.
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.
torch.Tensor matching the required tensor specs.
- abstract expand(*shape)¶
Returns a new Spec with the extended shape.
*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.
- 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.
index (int, torch.Tensor, slice or list) – index of the tensor
tensor_to_index – tensor to be indexed
- abstract is_in(val: Tensor) bool ¶
If the value
valis in the box defined by the TensorSpec, returns True, otherwise False.
val (torch.Tensor) – value to be checked
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.
val (torch.Tensor) – tensor to be mapped to the box.
a torch.Tensor belonging to the TensorSpec box.
- abstract rand(shape=None) Tensor ¶
Returns a random tensor in the box. The sampling will be uniform unless the box is unbounded.
shape (torch.Size) – shape of the random tensor
a random tensor sampled in the TensorSpec box.
- squeeze(dim: int | None = None)¶
Returns a new Spec with all the dimensions of size
dimis given, a squeeze operation is done only in that dimension.
dim (int or None) – the dimension to apply the squeeze operation to
- to_numpy(val: Tensor, safe: Optional[bool] = None) ndarray ¶
Returns the np.ndarray correspondent of an input tensor.
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
- type_check(value: Tensor, key: Optional[str] = None) None ¶
Checks the input value dtype against the TensorSpec dtype and raises an exception if they don’t match.
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.