ConditionalSkip

class torchrl.envs.transforms.ConditionalSkip(cond: Callable[[TensorDict], bool | torch.Tensor])

当满足特定条件时，跳过环境中步骤的转换。

此转换将 `cond(tensordict)` 的结果写入传递给 `TransformedEnv.base_env._step` 方法的 tensordict 的 “_step” 条目中。如果 `base_env` 不是批处理锁定的（一般来说，它是无状态的），则 tensordict 将被缩减到需要通过步骤的元素。如果它是批处理锁定的（一般来说，它是状态化的），如果 “_step” 中的任何值不是 `True`，则完全跳过该步骤。否则，可以信任环境会相应地处理 “_step” 信号。

注意

跳过操作也会影响修改环境输出的转换，即，如果满足条件，任何将在 `step()` 返回的 tensordict 上执行的转换都将被跳过。如果此效果不理想，可以通过将转换后的环境包装在另一个转换后的环境中来缓解，因为跳过操作只会影响 `ConditionalSkip` 转换的第一级父级。请参阅下面的示例。

参数:

cond (Callable[[TensorDictBase], bool | torch.Tensor]) – 一个可调用对象，用于输入 tensordict，它检查下一个环境步骤是否必须被跳过（`True` = 跳过，`False` = 执行 `env.step`）。

示例

>>> import torch
>>>
>>> from torchrl.envs import GymEnv
>>> from torchrl.envs.transforms.transforms import ConditionalSkip, StepCounter, TransformedEnv, Compose
>>>
>>> torch.manual_seed(0)
>>>
>>> base_env = TransformedEnv(
...     GymEnv("Pendulum-v1"),
...     StepCounter(step_count_key="inner_count"),
... )
>>> middle_env = TransformedEnv(
...     base_env,
...     Compose(
...         StepCounter(step_count_key="middle_count"),
...         ConditionalSkip(cond=lambda td: td["step_count"] % 2 == 1),
...     ),
...     auto_unwrap=False)  # makes sure that transformed envs are properly wrapped
>>> env = TransformedEnv(
...     middle_env,
...     StepCounter(step_count_key="step_count"),
...     auto_unwrap=False)
>>> env.set_seed(0)
>>>
>>> r = env.rollout(10)
>>> print(r["observation"])
tensor([[-0.9670, -0.2546, -0.9669],
        [-0.9802, -0.1981, -1.1601],
        [-0.9802, -0.1981, -1.1601],
        [-0.9926, -0.1214, -1.5556],
        [-0.9926, -0.1214, -1.5556],
        [-0.9994, -0.0335, -1.7622],
        [-0.9994, -0.0335, -1.7622],
        [-0.9984,  0.0561, -1.7933],
        [-0.9984,  0.0561, -1.7933],
        [-0.9895,  0.1445, -1.7779]])
>>> print(r["inner_count"])
tensor([[0],
        [1],
        [1],
        [2],
        [2],
        [3],
        [3],
        [4],
        [4],
        [5]])
>>> print(r["middle_count"])
tensor([[0],
        [1],
        [1],
        [2],
        [2],
        [3],
        [3],
        [4],
        [4],
        [5]])
>>> print(r["step_count"])
tensor([[0],
        [1],
        [2],
        [3],
        [4],
        [5],
        [6],
        [7],
        [8],
        [9]])

forward(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.