SelectTransform

类 torchrl.envs.transforms.SelectTransform(*selected_keys: NestedKey, keep_rewards: bool = True, keep_dones: bool = True)

从输入 tensordict 中选择键。

通常，应首选 ExcludeTransform：此变换还会

选择“action”（或 input_spec 中的其他键）、“done”和“reward”键，但可能还需要其他键。

参数：

*selected_keys (NestedKey 的可迭代对象) – 要选择的键的名称。如果键不存在，则会被忽略。

关键字参数：

• keep_rewards (bool, 可选的) – 如果为 False，如果要保留奖励键，则必须提供。默认为 True。

• keep_dones (bool, 可选的) – 如果为 False，如果要保留 done 键，则必须必须提供。默认为 True。

示例

>>> import gymnasium
>>> from torchrl.envs import GymWrapper
>>> env = TransformedEnv(
...     GymWrapper(gymnasium.make("Pendulum-v1")),
...     SelectTransform("observation", "reward", "done", keep_dones=False), # we leave done behind
... )
>>> env.rollout(3)  # the truncated key is now absent
TensorDict(
    fields={
        action: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False),
        done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
        next: TensorDict(
            fields={
                done: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.bool, is_shared=False),
                observation: Tensor(shape=torch.Size([3, 3]), device=cpu, dtype=torch.float32, is_shared=False),
                reward: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False)},
            batch_size=torch.Size([3]),
            device=cpu,
            is_shared=False),
        observation: Tensor(shape=torch.Size([3, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=cpu,
    is_shared=False)

forward(tensordict: TensorDictBase) → TensorDictBase

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

transform_output_spec(output_spec: Composite) → Composite

变换输出规范，使结果规范与变换映射匹配。

通常应保持此方法不变。更改应使用 transform_observation_spec()、transform_reward_spec() 和 transform_full_done_spec() 来实现。 :param output_spec: 变换前的规范 :type output_spec: TensorSpec

返回：

变换后的预期规范