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(next_tensordict: TensorDictBase) TensorDictBase ¶
读取输入 tensordict,并对选定的键应用转换。
默认情况下,此方法
直接调用
_apply_transform()
。不调用
_step()
或_call()
。
此方法在任何时候都不会在 env.step 中调用。但是,它会在
sample()
中调用。注意
forward
还可以使用dispatch
与常规关键字参数一起使用,将参数名称转换为键。示例
>>> class TransformThatMeasuresBytes(Transform): ... '''Measures the number of bytes in the tensordict, and writes it under `"bytes"`.''' ... def __init__(self): ... super().__init__(in_keys=[], out_keys=["bytes"]) ... ... def forward(self, tensordict: TensorDictBase) -> TensorDictBase: ... bytes_in_td = tensordict.bytes() ... tensordict["bytes"] = bytes ... return tensordict >>> t = TransformThatMeasuresBytes() >>> env = env.append_transform(t) # works within envs >>> t(TensorDict(a=0)) # Works offline too.
- transform_input_spec(input_spec: Composite) Composite [source]¶
转换输入规范,使结果规范与转换映射匹配。
- 参数:
input_spec (TensorSpec) – 转换前的规范
- 返回:
转换后的预期规范