CatTensors

class torchrl.envs.transforms.CatTensors(in_keys: Sequence[NestedKey] | None = None, out_key: NestedKey = 'observation_vector', dim: int = - 1, *, del_keys: bool = True, unsqueeze_if_oor: bool = False, sort: bool = True)

将多个键连接到单个张量中。

如果多个键描述单个状态（例如，“observation_position”和“observation_velocity”），这将特别有用。

参数：

• in_keys (嵌套键序列) – 要连接的键。如果为 None（或未提供），则在第一次使用转换时从父环境中检索键。此行为仅在设置了父级时才有效。

• out_key (嵌套键) – 结果张量的键。

• dim (int, 可选) – 将发生连接的维度。默认为 -1。

关键字参数：

• del_keys (bool, 可选) – 如果为 True，则连接后将删除输入值。默认为 True。

• unsqueeze_if_oor (bool, 可选) – 如果为 True，则 CatTensor 将检查要连接的张量是否存在指示的维度。如果不是，则将在该维度上对张量进行压缩。默认为 False。

• sort (bool, 可选) – 如果为 True，则将在转换中对键进行排序。否则，将使用用户提供的顺序。默认为 True。

示例

>>> transform = CatTensors(in_keys=["key1", "key2"])
>>> td = TensorDict({"key1": torch.zeros(1, 1),
...     "key2": torch.ones(1, 1)}, [1])
>>> _ = transform(td)
>>> print(td.get("observation_vector"))
tensor([[0., 1.]])
>>> transform = CatTensors(in_keys=["key1", "key2"], dim=-2, unsqueeze_if_oor=True)
>>> td = TensorDict({"key1": torch.zeros(1),
...     "key2": torch.ones(1)}, [])
>>> _ = transform(td)
>>> print(td.get("observation_vector").shape)
torch.Size([2, 1])

forward(tensordict: TensorDictBase) → TensorDictBase

读取输入 tensordict，并针对所选键应用转换。

transform_observation_spec(observation_spec: TensorSpec) → TensorSpec

转换观察规范，以便生成的规范与转换映射匹配。

参数：

observation_spec (TensorSpec) – 转换前的规范

返回值：

转换后的预期规范