FlattenObservation

class torchrl.envs.transforms.FlattenObservation(first_dim: int, last_dim: int, in_keys: Sequence[NestedKey] | None = None, out_keys: Sequence[NestedKey] | None = None, allow_positive_dim: bool = False)

展平张量的相邻维度。

参数:

• first_dim (int) – 要展平的第一个维度。

• last_dim (int) – 要展平的最后一个维度。

• in_keys (sequence of NestedKey, optional) – 要展平的条目。如果未提供，则假定为 ["pixels"]。

• out_keys (sequence of NestedKey, optional) – 展平后的观察键。如果未提供，则假定为 in_keys。

• allow_positive_dim (bool, optional) – 如果为 True，则接受正维度。 FlattenObservation 将这些维度映射到输入张量的第 n 个特征维度（即父环境批次大小之后的第 n 个维度）。默认为 False，即不允许正维度。

forward(next_tensordict: TensorDictBase) → TensorDictBase

读取输入 tensordict，并对选定的键应用转换。

_call 可以在需要独立于上一步收集的数据（包括动作和状态）修改 env.step 的输出时进行重写。

对于任何仅与父环境相关的操作（例如 FrameSkip），请改用 _step() 方法。仅当需要修改输入 tensordict 时，才应重写 _call()。

_call() 将由 step() 和 reset() 调用，但在 forward() 期间不调用。

transform_observation_spec(observation_spec: TensorSpec) → TensorSpec

转换观察规范，使结果规范与转换映射匹配。

参数:

observation_spec (TensorSpec) – 转换前的规范

返回:

转换后的预期规范