[Feature] Consistent Dropout

[N00bcak]

Description

Introduces Consistent Dropout (Hausknecht & Wagener, 2022) to TorchRL.

Consists of the following changes:

• ConsistentDropout PyTorch module
• Companion TensorDictModule: ConsistentDropoutModule
• Tests for ConsistentDropoutModule

Motivation and Context

Addresses a feature request.
Fleshes out a PR draft #1587.

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/2399

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 2 New Failures, 8 Unrelated Failures

As of commit 3231780 with merge base 6aa4b53 ():

NEW FAILURES - The following jobs have failed:

• Habitat Tests on Linux / tests (3.9, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 8f8dedc71dd868db7724a1bf63aa1a72fd40d11e79650e254e813a1efcee4aa3 /exec failed with exit code 139
• Libs Tests on Linux / unittests-gym (3.9, 12.1) / linux-job (gh)
  RuntimeError: Command docker exec -t 9e960e3181dff173511480976674b5f0c1eca2d4942947a2c6395cb4b63388ca /exec failed with exit code 1

FLAKY - The following jobs failed but were likely due to flakiness present on trunk:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh) (detected as infra flaky with no log or failing log classifier)
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh) (detected as infra flaky with no log or failing log classifier)
• Unit-tests on Windows / unittests-cpu / windows-job (gh) (detected as infra flaky with no log or failing log classifier)

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Build Windows Wheels / pytorch/rl (pytorch/rl, .github/scripts/td_script.sh, .github/scripts/version_script.bat, python ... / upload / wheel-py3_9-cpu (gh) (trunk failure)
• Build Windows Wheels / pytorch/rl (pytorch/rl, .github/scripts/td_script.sh, .github/scripts/version_script.bat, python ... / upload / wheel-py3_9-cuda11_8 (gh) (trunk failure)
• Build Windows Wheels / pytorch/rl (pytorch/rl, .github/scripts/td_script.sh, .github/scripts/version_script.bat, python ... / upload / wheel-py3_9-cuda12_1 (gh) (trunk failure)
• Build Windows Wheels / pytorch/rl (pytorch/rl, .github/scripts/td_script.sh, .github/scripts/version_script.bat, python ... / upload / wheel-py3_9-cuda12_4 (gh) (trunk failure)
• Unit-tests on Linux / tests-olddeps (3.8, 11.6) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[vmoens]

I need to make a second pass over this but here's what I got running:

from torchrl.modules import ConsistentDropoutModule, get_primers_from_module

from tensordict.nn import TensorDictModule as Mod, TensorDictSequential as Seq
import torch
from torchrl.envs import GymEnv
from tensordict import TensorDict

m = Seq(
    Mod(torch.nn.Linear(3, 4), in_keys=["observation"], out_keys=["intermediate"]),
    ConsistentDropoutModule(p=0.5, input_shape=(4,), in_keys="intermediate"),
    Mod(torch.nn.Linear(4, 1), in_keys=["intermediate"], out_keys=["action"]),
)
primer = get_primers_from_module(m)
env = GymEnv("Pendulum-v1").append_transform(primer)
r = env.reset()
env.rollout(32, m)["next", "mask_13266468624"]

With this you generate your mask during reset, such that it stays the same during execution.

This works, but as you can see we need to change the size of the primer, which means that we won't be able to use the same primer with different envs:

from torchrl.modules import ConsistentDropoutModule, get_primers_from_module

from tensordict.nn import TensorDictModule as Mod, TensorDictSequential as Seq
import torch
from torchrl.envs import GymEnv, SerialEnv, StepCounter
from tensordict import TensorDict

m = Seq(
    Mod(torch.nn.Linear(3, 4), in_keys=["observation"], out_keys=["intermediate"]),
    ConsistentDropoutModule(p=0.5, input_shape=(2, 4,), in_keys="intermediate"),
    Mod(torch.nn.Linear(4, 1), 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.reset()

There should be a way to solve this though!