embodichain.lab.sim.types

embodichain.lab.sim.types#

Overview#

Shared tensor/type aliases used across simulation, environment, and policy interfaces.

Type Aliases

Array

alias of Tensor | ndarray | Sequence

Device

alias of str | device

EnvObs

alias of TensorDict[str, Tensor | TensorDict[str, Tensor]]

EnvAction

alias of Tensor | TensorDict[str, Tensor]
embodichain.lab.sim.types.Array(*args, **kwargs)#

alias of Tensor | ndarray | Sequence

embodichain.lab.sim.types.Device(*args, **kwargs)#

alias of str | device

embodichain.lab.sim.types.EnvObs(source=None, batch_size=None, device=None, names=None, non_blocking=None, lock=False, **kwargs)#

A batched dictionary of tensors.

TensorDict is a tensor container where all tensors are stored in a key-value pair fashion and where each element shares the same first N leading dimensions shape, where is an arbitrary number with N >= 0.

Additionally, if the tensordict has a specified device, then each element must share that device.

TensorDict instances support many regular tensor operations with the notable exception of algebraic operations:

  • operations on shape: when a shape operation is called (indexing, reshape, view, expand, transpose, permute, unsqueeze, squeeze, masking etc), the operations is done as if it was executed on a tensor of the same shape as the batch size then expended to the right, e.g.:

    >>> td = TensorDict({'a': torch.zeros(3, 4, 5)}, batch_size=[3, 4])
>>> # returns a TensorDict of batch size [3, 4, 1]:
>>> td_unsqueeze = td.unsqueeze(-1)
>>> # returns a TensorDict of batch size [12]
>>> td_view = td.view(-1)
>>> # returns a tensor of batch size [12, 4]
>>> a_view = td.view(-1).get("a")

  • casting operations: a TensorDict can be cast on a different device using

    >>> td_cpu = td.to("cpu")
>>> dictionary = td.to_dict()

    A call of the .to() method with a dtype will return an error.

  • Cloning (clone()), contiguous (contiguous());

  • Reading: td.get(key), td.get_at(key, index)

  • Content modification: td.set(key, value), td.set_(key, value), td.update(td_or_dict), td.update_(td_or_dict), td.fill_(key, value), td.rename_key_(old_name, new_name), etc.

  • Operations on multiple tensordicts: torch.cat(tensordict_list, dim), torch.stack(tensordict_list, dim), td1 == td2, td.apply(lambda x+y, other_td) etc.

Parameters:

  • source (TensorDict or Dict[NestedKey, Union[Tensor, TensorDictBase]]) – a data source. If empty, the tensordict can be populated subsequently. A TensorDict can also be built via a sequence of keyword arguments, as it is the case for dict(...).

  • batch_size (iterable of int, optional) – a batch size for the tensordict. The batch size can be modified subsequently as long as it is compatible with its content. If not batch-size is provided, an empty batch-size is assumed (it is not inferred automatically from the data). To automatically set the batch-size, refer to auto_batch_size_().

  • device (torch.device or compatible type, optional) – a device for the TensorDict. If provided, all tensors will be stored on that device. If not, tensors on different devices are allowed.

  • names (lsit of str, optional) – the names of the dimensions of the tensordict. If provided, its length must match the one of the batch_size. Defaults to None (no dimension name, or None for every dimension).

  • non_blocking (bool, optional) – if True and a device is passed, the tensordict is delivered without synchronization. This is the fastest option but is only safe when casting from cpu to cuda (otherwise a synchronization call must be implemented by the user). If False is passed, every tensor movement will be done synchronously. If None (default), the device casting will be done asynchronously but a synchronization will be executed after creation if required. This option should generally be faster than False and potentially slower than True.

  • lock (bool, optional) – if True, the resulting tensordict will be locked.

Examples

>>> import torch
>>> from tensordict import TensorDict
>>> source = {'random': torch.randn(3, 4),
...     'zeros': torch.zeros(3, 4, 5)}
>>> batch_size = [3]
>>> td = TensorDict(source, batch_size=batch_size)
>>> print(td.shape)  # equivalent to td.batch_size
torch.Size([3])
>>> td_unqueeze = td.unsqueeze(-1)
>>> print(td_unqueeze.get("zeros").shape)
torch.Size([3, 1, 4, 5])
>>> print(td_unqueeze[0].shape)
torch.Size([1])
>>> print(td_unqueeze.view(-1).shape)
torch.Size([3])
>>> print((td.clone()==td).all())
True

alias of TensorDict[str, Tensor | TensorDict[str, Tensor]]

embodichain.lab.sim.types.EnvAction(*args, **kwargs)#

alias of Tensor | TensorDict[str, Tensor]

Contents