Transform¶
- class torchrl.envs.transforms.Transform(in_keys: Sequence[NestedKey] = None, out_keys: Sequence[NestedKey] | None = None, in_keys_inv: Sequence[NestedKey] | None = None, out_keys_inv: Sequence[NestedKey] | None = None)[source]¶
环境变换的基类,用于修改或创建 tensordict 中的新数据。
变换用于操作环境的输入和输出数据。它们可用于预处理观测值、修改奖励或变换动作。变换可以组合在一起以创建更复杂的变换。
变换接收一个 tensordict 作为输入,并返回(相同或另一个)tensordict 作为输出,其中一系列值已被修改或以新键创建。
- 变量:
parent – 变换的父环境。
container – 包含该变换的容器。
in_keys – 变换将从中读取的输入 tensordict 的键。
out_keys – 变换将写入的输出 tensordict 的键。
另请参阅
子类化 Transform
有多种子类化变换的方法。需要考虑的事情是:
变换对于每个被变换的张量/项目都是相同的吗?使用
_apply_transform()
和_inv_apply_transform()
。变换是否需要访问 env.step 的输入数据以及输出?重写
_step()
。否则,重写_call()
(或_inv_call()
)。变换是否要在回放缓冲区中使用?重写
forward()
、inv()
、_apply_transform()
或_inv_apply_transform()
。在变换中,您可以通过
parent
(基础环境 + 直到此处的所有变换)或container()
(封装该变换的对象)来访问(并调用)父环境。别忘了根据需要编辑规范:顶层:
transform_output_spec()
、transform_input_spec()
。叶级:transform_observation_spec()
、transform_action_spec()
、transform_state_spec()
、transform_reward_spec()
和transform_reward_spec()
。
有关实际示例,请参阅上面列出的方法。
- property collector: DataCollectorBase | None¶
返回与容器关联的收集器(如果存在)。
每当变换需要了解收集器或与之关联的策略时,都可以使用此属性。请确保仅在未嵌套在子进程中的变换上调用此属性。收集器引用不会传递给
ParallelEnv
或类似的批处理环境的 worker。确保仅在未嵌套在子进程中的变换上调用此属性。收集器引用将不会传递给
ParallelEnv
或类似的批处理环境的 worker。
- property container: EnvBase | None¶
返回包含该变换的环境。
示例
>>> from torchrl.envs import TransformedEnv, Compose, RewardSum, StepCounter >>> from torchrl.envs.libs.gym import GymEnv >>> env = TransformedEnv(GymEnv("Pendulum-v1"), Compose(RewardSum(), StepCounter())) >>> env.transform[0].container is env True
- forward(tensordict: TensorDictBase = None) TensorDictBase [source]¶
读取输入 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.
- inv(tensordict: TensorDictBase = None) TensorDictBase [source]¶
读取输入 tensordict,并对选定的键应用逆变换。
默认情况下,此方法
直接调用
_inv_apply_transform()
。不调用
_inv_call()
。
注意
inv
还使用dispatch
通过将参数名称映射到键来处理常规关键字参数。注意
inv
由extend()
调用。
- property parent: TransformedEnv | None¶
返回变换的父环境。
父环境是包含直到当前变换的所有变换的环境。
示例
>>> from torchrl.envs import TransformedEnv, Compose, RewardSum, StepCounter >>> from torchrl.envs.libs.gym import GymEnv >>> env = TransformedEnv(GymEnv("Pendulum-v1"), Compose(RewardSum(), StepCounter())) >>> env.transform[1].parent TransformedEnv( env=GymEnv(env=Pendulum-v1, batch_size=torch.Size([]), device=cpu), transform=Compose( RewardSum(keys=['reward'])))
- to(*args, **kwargs)[source]¶
移动和/或转换参数和缓冲区。
这可以这样调用
- to(device=None, dtype=None, non_blocking=False)[source]
- to(dtype, non_blocking=False)[source]
- to(tensor, non_blocking=False)[source]
- to(memory_format=torch.channels_last)[source]
其签名类似于
torch.Tensor.to()
,但只接受浮点或复数dtype
。此外,此方法将仅将浮点或复数参数和缓冲区转换为dtype
(如果已给出)。整数参数和缓冲区将移动到device
(如果已给出),但 dtype 不变。当non_blocking
设置为 True 时,它会尝试异步(相对于主机)转换/移动,如果可能,例如将具有固定内存的 CPU 张量移动到 CUDA 设备。有关示例,请参阅下文。
注意
此方法就地修改模块。
- 参数:
device (
torch.device
) – 此模块中的参数和缓冲区的目标设备dtype (
torch.dtype
) – 此模块中的参数和缓冲区的目标浮点或复数 dtypetensor (torch.Tensor) – 其 dtype 和设备是此模块中所有参数和缓冲区的目标 dtype 和设备
memory_format (
torch.memory_format
) – 此模块中 4D 参数和缓冲区的目标内存格式(仅关键字参数)
- 返回:
self
- 返回类型:
模块
示例
>>> # xdoctest: +IGNORE_WANT("non-deterministic") >>> linear = nn.Linear(2, 2) >>> linear.weight Parameter containing: tensor([[ 0.1913, -0.3420], [-0.5113, -0.2325]]) >>> linear.to(torch.double) Linear(in_features=2, out_features=2, bias=True) >>> linear.weight Parameter containing: tensor([[ 0.1913, -0.3420], [-0.5113, -0.2325]], dtype=torch.float64) >>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_CUDA1) >>> gpu1 = torch.device("cuda:1") >>> linear.to(gpu1, dtype=torch.half, non_blocking=True) Linear(in_features=2, out_features=2, bias=True) >>> linear.weight Parameter containing: tensor([[ 0.1914, -0.3420], [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1') >>> cpu = torch.device("cpu") >>> linear.to(cpu) Linear(in_features=2, out_features=2, bias=True) >>> linear.weight Parameter containing: tensor([[ 0.1914, -0.3420], [-0.5112, -0.2324]], dtype=torch.float16) >>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble) >>> linear.weight Parameter containing: tensor([[ 0.3741+0.j, 0.2382+0.j], [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128) >>> linear(torch.ones(3, 2, dtype=torch.cdouble)) tensor([[0.6122+0.j, 0.1150+0.j], [0.6122+0.j, 0.1150+0.j], [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)
- transform_action_spec(action_spec: TensorSpec) TensorSpec [source]¶
转换动作规范,使结果规范与变换映射匹配。
- 参数:
action_spec (TensorSpec) – 变换前的规范
- 返回:
转换后的预期规范
- transform_done_spec(done_spec: TensorSpec) TensorSpec [source]¶
变换 done spec,使结果 spec 与变换映射匹配。
- 参数:
done_spec (TensorSpec) – 变换前的 spec
- 返回:
转换后的预期规范
- transform_input_spec(input_spec: TensorSpec) TensorSpec [source]¶
转换输入规范,使结果规范与转换映射匹配。
- 参数:
input_spec (TensorSpec) – 转换前的规范
- 返回:
转换后的预期规范
- transform_observation_spec(observation_spec: TensorSpec) TensorSpec [source]¶
转换观察规范,使结果规范与转换映射匹配。
- 参数:
observation_spec (TensorSpec) – 转换前的规范
- 返回:
转换后的预期规范
- transform_output_spec(output_spec: Composite) Composite [source]¶
转换输出规范,使结果规范与转换映射匹配。
此方法通常应保持不变。应通过
transform_observation_spec()
、transform_reward_spec()
和transform_full_done_spec()
来实现更改。:param output_spec: 变换前的 spec :type output_spec: TensorSpec- 返回:
转换后的预期规范
- transform_reward_spec(reward_spec: TensorSpec) TensorSpec [source]¶
转换奖励的 spec,使其与变换映射匹配。
- 参数:
reward_spec (TensorSpec) – 变换前的 spec
- 返回:
转换后的预期规范
- transform_state_spec(state_spec: TensorSpec) TensorSpec [source]¶
转换状态规范,使结果规范与变换映射匹配。
- 参数:
state_spec (TensorSpec) – 变换前的规范
- 返回:
转换后的预期规范