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)[source]¶
展平张量的相邻维度。
- 参数:
first_dim (int) – 要展平维度的起始维度。
last_dim (int) – 要展平维度的结束维度。
in_keys (NestedKey 序列, 可选) – 要展平的条目。如果未提供,则假定为
["pixels"]
。out_keys (NestedKey 序列, 可选) – 展平后的观察键。如果未提供,则假定为
in_keys
。allow_positive_dim (bool, 可选) – 如果为
True
,则接受正数维度。FlattenObservation
会将这些维度映射到输入张量的第 n 个特征维度(即父环境批次大小后的第 n 个维度)。默认为 False,即不允许非负维度。
- forward(tensordict: TensorDictBase) TensorDictBase ¶
读取输入的 tensordict,并对选定的键应用转换。
对于仅与父环境相关的任何操作(例如 FrameSkip),应修改 _step 方法。
_call()
仅在需要修改输入的 tensordict 时才应被覆盖。_call()
将由TransformedEnv.step()
和TransformedEnv.reset()
调用。
- transform_observation_spec(observation_spec: TensorSpec) TensorSpec [source]¶
转换观察 Spec,使结果 Spec 与转换映射匹配。
- 参数:
observation_spec (TensorSpec) – 转换前的 Spec
- 返回:
转换后的预期 Spec