RenameTransform¶
- class torchrl.envs.transforms.RenameTransform(in_keys, out_keys, in_keys_inv=None, out_keys_inv=None, create_copy=False)[source]¶
一个用于重命名输出 tensordict (或通过逆向键重命名输入 tensordict) 中条目的变换。
- 参数:
in_keys (sequence of NestedKey) – 待重命名的条目。
out_keys (sequence of NestedKey) – 重命名后的条目名称。
in_keys_inv (sequence of NestedKey, optional) – 在输入 tensordict 中待重命名的条目,这些条目将传递给
EnvBase._step()
。out_keys_inv (sequence of NestedKey, optional) – 在输入 tensordict 中重命名后的条目名称。
create_copy (bool, optional) – 如果为
True
,则将条目复制一份并使用不同的名称,而不是直接重命名。这允许重命名不可变条目,如"reward"
和"done"
。
示例
>>> from torchrl.envs.libs.gym import GymEnv >>> env = TransformedEnv( ... GymEnv("Pendulum-v1"), ... RenameTransform(["observation", ], ["stuff",], create_copy=False), ... ) >>> tensordict = env.rollout(3) >>> print(tensordict) 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), reward: Tensor(shape=torch.Size([3, 1]), device=cpu, dtype=torch.float32, is_shared=False), stuff: Tensor(shape=torch.Size([3, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3]), device=cpu, is_shared=False), stuff: Tensor(shape=torch.Size([3, 3]), device=cpu, dtype=torch.float32, is_shared=False)}, batch_size=torch.Size([3]), device=cpu, is_shared=False) >>> # if the output is also an input, we need to rename if both ways: >>> from torchrl.envs.libs.brax import BraxEnv >>> env = TransformedEnv( ... BraxEnv("fast"), ... RenameTransform(["state"], ["newname"], ["state"], ["newname"]) ... ) >>> _ = env.set_seed(1) >>> tensordict = env.rollout(3) >>> assert "newname" in tensordict.keys() >>> assert "state" not in tensordict.keys()
- forward(tensordict: TensorDictBase) TensorDictBase ¶
读取输入 tensordict,并对选定的键应用变换。
- transform_input_spec(input_spec: Composite) Composite [source]¶
变换输入规格,使得到的规格与变换映射匹配。
- 参数:
input_spec (TensorSpec) – 变换前的规格
- 返回值:
变换后预期的规格