ConsistentDropoutModule

class torchrl.modules.ConsistentDropoutModule(*args, **kwargs)

用於 ConsistentDropout 的 TensorDictModule 包裝器。

參數:

• p (float, optional) – Dropout 概率。預設值：0.5。

• in_keys (NestedKey 或 NestedKeys 的清單) – 要從輸入 tensordict 讀取並傳遞給此模組的鍵。

• out_keys (NestedKey 或 NestedKeys 的可迭代物件) – 要寫入輸入 tensordict 的鍵。預設為 in_keys 值。

關鍵字引數:

• input_shape (tuple, optional) – 輸入的形狀（非批次化的），用於使用 make_tensordict_primer() 產生 tensordict primers。

• input_dtype (torch.dtype, optional) – primer 輸入的 dtype。如果未傳遞，則假設為 torch.get_default_dtype。

注意

若要在策略中使用此類別，需要在重設時重設遮罩。這可以透過 TensorDictPrimer 轉換來實現，該轉換可以使用 make_tensordict_primer() 獲得。有關更多資訊，請參閱此方法。

範例

>>> from tensordict import TensorDict
>>> module = ConsistentDropoutModule(p = 0.1)
>>> td = TensorDict({"x": torch.randn(3, 4)}, [3])
>>> module(td)
TensorDict(
    fields={
        mask_6127171760: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.bool, is_shared=False),
        x: Tensor(shape=torch.Size([3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},
    batch_size=torch.Size([3]),
    device=None,
    is_shared=False)

forward(tensordict)

定義每次呼叫時執行的計算。

應由所有子類別覆寫。

注意

雖然需要在這個函式中定義正向傳遞的配方，但之後應該呼叫 Module 實例，而不是這個，因為前者會負責執行註冊的掛鉤，而後者會靜默地忽略它們。

make_tensordict_primer()

為環境建立一個 tensordict primer，以便在重置呼叫期間產生隨機遮罩。

另請參閱

torchrl.modules.utils.get_primers_from_module() 關於為給定模組產生所有 primer 的方法

模組。

範例

>>> from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod
>>> from torchrl.envs import GymEnv, StepCounter, SerialEnv
>>> m = Seq(
...     Mod(torch.nn.Linear(7, 4), in_keys=["observation"], out_keys=["intermediate"]),
...     ConsistentDropoutModule(
...         p=0.5,
...         input_shape=(2, 4),
...         in_keys="intermediate",
...     ),
...     Mod(torch.nn.Linear(4, 7), in_keys=["intermediate"], out_keys=["action"]),
... )
>>> primer = get_primers_from_module(m)
>>> env0 = GymEnv("Pendulum-v1").append_transform(StepCounter(5))
>>> env1 = GymEnv("Pendulum-v1").append_transform(StepCounter(6))
>>> env = SerialEnv(2, [lambda env=env0: env, lambda env=env1: env])
>>> env = env.append_transform(primer)
>>> r = env.rollout(10, m, break_when_any_done=False)
>>> mask = [k for k in r.keys() if k.startswith("mask")][0]
>>> assert (r[mask][0, :5] != r[mask][0, 5:6]).any()
>>> assert (r[mask][0, :4] == r[mask][0, 4:5]).all()