Composite¶
- class torchrl.data.Composite(*args, **kwargs)[source]¶
A composition of TensorSpecs.
If a
TensorSpec
is the set-description of Tensor category, theComposite
class is akin to theTensorDict
class. LikeTensorDict
, it has ashape
(akin to theTensorDict
’sbatch_size
) and an optionaldevice
.- Parameters:
*args – if an unnamed argument is passed, it must be a dictionary with keys matching the expected keys to be found in the
Composite
object. This is useful to build nested CompositeSpecs with tuple indices.**kwargs (key (str) – value (TensorSpec)): dictionary of tensorspecs to be stored. Values can be None, in which case is_in will be assumed to be
True
for the corresponding tensors, andproject()
will have no effect. spec.encode cannot be used with missing values.
- Variables:
device (torch.device or None) – if not specified, the device of the composite spec is
None
(as it is the case for TensorDicts). A non-none device constraints all leaves to be of the same device. On the other hand, aNone
device allows leaves to have different devices. Defaults toNone
.shape (torch.Size) – the leading shape of all the leaves. Equivalent to the batch-size of the corresponding tensordicts.
Examples
>>> pixels_spec = Bounded( ... low=torch.zeros(4, 3, 32, 32), ... high=torch.ones(4, 3, 32, 32), ... dtype=torch.uint8 ... ) >>> observation_vector_spec = Bounded( ... low=torch.zeros(4, 33), ... high=torch.ones(4, 33), ... dtype=torch.float) >>> composite_spec = Composite( ... pixels=pixels_spec, ... observation_vector=observation_vector_spec, ... shape=(4,) ... ) >>> composite_spec Composite( pixels: BoundedDiscrete( shape=torch.Size([4, 3, 32, 32]), space=ContinuousBox( low=Tensor(shape=torch.Size([4, 3, 32, 32]), device=cpu, dtype=torch.uint8, contiguous=True), high=Tensor(shape=torch.Size([4, 3, 32, 32]), device=cpu, dtype=torch.uint8, contiguous=True)), device=cpu, dtype=torch.uint8, domain=discrete), observation_vector: BoundedContinuous( shape=torch.Size([4, 33]), space=ContinuousBox( low=Tensor(shape=torch.Size([4, 33]), device=cpu, dtype=torch.float32, contiguous=True), high=Tensor(shape=torch.Size([4, 33]), device=cpu, dtype=torch.float32, contiguous=True)), device=cpu, dtype=torch.float32, domain=continuous), device=None, shape=torch.Size([4])) >>> td = composite_spec.rand() >>> td TensorDict( fields={ observation_vector: Tensor(shape=torch.Size([4, 33]), device=cpu, dtype=torch.float32, is_shared=False), pixels: Tensor(shape=torch.Size([4, 3, 32, 32]), device=cpu, dtype=torch.uint8, is_shared=False)}, batch_size=torch.Size([4]), device=None, is_shared=False) >>> # we can build a nested composite spec using unnamed arguments >>> print(Composite({("a", "b"): None, ("a", "c"): None})) Composite( a: Composite( b: None, c: None, device=None, shape=torch.Size([])), device=None, shape=torch.Size([]))
- 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.
- contains(item: torch.Tensor | tensordict.base.TensorDictBase) bool ¶
If the value
val
could have been generated by theTensorSpec
, returnsTrue
, otherwiseFalse
.See
is_in()
for more information.
- cpu()¶
Casts the TensorSpec to ‘cpu’ device.
- cuda(device=None)¶
Casts the TensorSpec to ‘cuda’ device.
- property device: Union[device, str, int]¶
The device of the spec.
Only
Composite
specs can have aNone
device. All leaves must have a non-null device.
- encode(vals: Dict[str, Any], *, ignore_device: bool = False) Dict[str, Tensor] [source]¶
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 toTensorDict(..., device="cuda")
which is faster.- Returns:
torch.Tensor matching the required tensor specs.
- 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.
- get(item, default=_NoDefault.ZERO)[source]¶
Gets an item from the Composite.
If the item is absent, a default value can be passed.
- 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: torch.Tensor | tensordict.base.TensorDictBase) torch.Tensor | tensordict.base.TensorDictBase ¶
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: Union[dict, TensorDictBase]) bool [source]¶
If the value
val
could have been generated by theTensorSpec
, returnsTrue
, otherwiseFalse
.More precisely, the
is_in
methods checks that the valueval
is within the limits defined by thespace
attribute (the box), and that thedtype
,device
,shape
potentially other metadata match those of the spec. If any of these checks fails, theis_in
method will returnFalse
.- Parameters:
val (torch.Tensor) – value to be checked.
- Returns:
boolean indicating if values belongs to the TensorSpec box.
- items(include_nested: bool = False, leaves_only: bool = False, *, is_leaf: Optional[Callable[[type], bool]] = None) _CompositeSpecItemsView [source]¶
Items of the Composite.
- Parameters:
include_nested (bool, optional) – if
False
, the returned keys will not be nested. They will represent only the immediate children of the root, and not the whole nested sequence, i.e. aComposite(next=Composite(obs=None))
will lead to the keys["next"]. Default is ``False`
, i.e. nested keys will not be returned.leaves_only (bool, optional) – if
False
, the values returned will contain every level of nesting, i.e. aComposite(next=Composite(obs=None))
will lead to the keys["next", ("next", "obs")]
. Default isFalse
.
- Keyword Arguments:
is_leaf (callable, optional) – reads a type and returns a boolean indicating if that type should be seen as a leaf. By default, all non-Composite nodes are considered as leaves.
- keys(include_nested: bool = False, leaves_only: bool = False, *, is_leaf: Optional[Callable[[type], bool]] = None) _CompositeSpecKeysView [source]¶
Keys of the Composite.
The keys argument reflect those of
tensordict.TensorDict
.- Parameters:
include_nested (bool, optional) – if
False
, the returned keys will not be nested. They will represent only the immediate children of the root, and not the whole nested sequence, i.e. aComposite(next=Composite(obs=None))
will lead to the keys["next"]. Default is ``False`
, i.e. nested keys will not be returned.leaves_only (bool, optional) – if
False
, the values returned will contain every level of nesting, i.e. aComposite(next=Composite(obs=None))
will lead to the keys["next", ("next", "obs")]
. Default isFalse
.
- Keyword Arguments:
is_leaf (callable, optional) – reads a type and returns a boolean indicating if that type should be seen as a leaf. By default, all non-Composite nodes are considered as leaves.
- lock_(recurse=False)[source]¶
Locks the Composite and prevents modification of its content.
This is only a first-level lock, unless specified otherwise through the
recurse
arg.Leaf specs can always be modified in place, but they cannot be replaced in their Composite parent.
Examples
>>> shape = [3, 4, 5] >>> spec = Composite( ... a=Composite( ... b=Composite(shape=shape[:3], device="cpu"), shape=shape[:2] ... ), ... shape=shape[:1], ... ) >>> spec["a"] = spec["a"].clone() >>> recurse = False >>> spec.lock_(recurse=recurse) >>> try: ... spec["a"] = spec["a"].clone() ... except RuntimeError: ... print("failed!") failed! >>> try: ... spec["a", "b"] = spec["a", "b"].clone() ... print("succeeded!") ... except RuntimeError: ... print("failed!") succeeded! >>> recurse = True >>> spec.lock_(recurse=recurse) >>> try: ... spec["a", "b"] = spec["a", "b"].clone() ... print("succeeded!") ... except RuntimeError: ... print("failed!") failed!
- make_neg_dim(dim: int) T ¶
Converts a specific dimension to
-1
.
- property ndim¶
Number of dimensions of the spec shape.
Shortcut for
len(spec.shape)
.
- one(shape: Optional[Size] = None) torch.Tensor | tensordict.base.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 ofone
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: Optional[Size] = None) torch.Tensor | tensordict.base.TensorDictBase ¶
Proxy to
one()
.
- project(val: TensorDictBase) TensorDictBase [source]¶
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) TensorDictBase [source]¶
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.
- sample(shape: Optional[Size] = None) torch.Tensor | tensordict.base.TensorDictBase ¶
Returns a random tensor in the space defined by the spec.
See
rand()
for details.
- squeeze(dim: Optional[int] = 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: Union[dtype, device, str, int]) Composite [source]¶
Casts a TensorSpec to a device or a dtype.
Returns the same spec if no change is made.
- to_numpy(val: TensorDict, safe: Optional[bool] = None) dict [source]¶
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.
- type_check(value: Union[Tensor, TensorDictBase], selected_keys: Optional[Union[str, Sequence[str]]] = None)[source]¶
Checks the input value
dtype
against theTensorSpec
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.
- unlock_(recurse=False)[source]¶
Unlocks the Composite and allows modification of its content.
This is only a first-level lock modification, unless specified otherwise through the
recurse
arg.
- 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.
- values(include_nested: bool = False, leaves_only: bool = False, *, is_leaf: Optional[Callable[[type], bool]] = None) _CompositeSpecValuesView [source]¶
Values of the Composite.
- Parameters:
include_nested (bool, optional) – if
False
, the returned keys will not be nested. They will represent only the immediate children of the root, and not the whole nested sequence, i.e. aComposite(next=Composite(obs=None))
will lead to the keys["next"]. Default is ``False`
, i.e. nested keys will not be returned.leaves_only (bool, optional) – if
False
, the values returned will contain every level of nesting, i.e. aComposite(next=Composite(obs=None))
will lead to the keys["next", ("next", "obs")]
. Default isFalse
.
- Keyword Arguments:
is_leaf (callable, optional) – reads a type and returns a boolean indicating if that type should be seen as a leaf. By default, all non-Composite nodes are considered as leaves.
- zero(shape: Optional[Size] = None) TensorDictBase [source]¶
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 ofzero
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: Optional[Size] = None) torch.Tensor | tensordict.base.TensorDictBase ¶
Proxy to
zero()
.