Feature Request: Consistent Dropout Implementation

J_Johnson: <p>I would like to put in a feature request to implement Consistent Dropout, as described in this paper:</p>
<aside class="onebox allowlistedgeneric" data-onebox-src="https://arxiv.org/abs/2202.11818">
  <header class="source">
      <img src="https://static.arxiv.org/static/browse/0.3.4/images/icons/favicon-32x32.png" class="site-icon" width="" height="">

      <a href="https://arxiv.org/abs/2202.11818" target="_blank" rel="noopener nofollow ugc">arXiv.org</a>
  </header>

  <article class="onebox-body">
    <img src="https://static.arxiv.org/static/browse/0.3.4/images/arxiv-logo-fb.png" class="thumbnail" width="" height="">

<h3><a href="https://arxiv.org/abs/2202.11818" target="_blank" rel="noopener nofollow ugc">Consistent Dropout for Policy Gradient Reinforcement Learning</a></h3>

  <p>Dropout has long been a staple of supervised learning, but is rarely used in
reinforcement learning. We analyze why naive application of dropout is
problematic for policy-gradient learning algorithms and introduce consistent
dropout, a simple...</p>


  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>

<p><strong>Background</strong></p>
<p>Reinforcement Learning typically does not use dropout layers as they cause instability.</p>
<p>What the above paper found, though, is fixing the dropout masks to be unchanged within each episode provided stable training and improved overall performance for RL networks vs. no dropout layers.</p>
<p><strong>Implementation</strong></p>
<p>One approach could be to supply an argument in <code>nn.Dropout</code> when instantiating called <code>freeze</code> which defaults to <code>False</code>.</p>
<p>If <code>freeze = True</code>, the masks stay fixed on those dropout layers until a manual <code>model.dropout_reset()</code> method is called.</p>
<p>Thank you.</p>

vmoens: <p>Thanks for raising this.<br>
Indeed it would be a pretty neat feature to have.<br>
I think this would nicely fit in TorchRL (rather than torch core) as my understanding is that it is quite RL specific.<br>
Wdyt?</p>

J_Johnson: <p><a class="mention" href="/u/vmoens">@vmoens</a> While I think the paper demonstrates it’s usefulness in RL, I don’t see why it wouldn’t be a potential improvement in other time sequence prediction networks, such as RNNs, for the same reasons.</p>

vmoens: <p>The reason I was suggesting TorchRL is that we make sure that this sort of traj-consistent modules are easy to code and well standardized / tested / robust and that the proper errors / warnings are raised whenever you use them.</p>
<p>If you think this should be a pytorch core feature under torch.nn, feel free to open an issue on pytorch github. Given the scope of the feature, my 2 cents is that the answer will be that this is more suited for a domain library.</p>

vmoens: <p>We merged this in torchrl.<br>
You can use the layer independently or within a torchrl script, here’s an example:</p>
<pre data-code-wrap="python"><code class="lang-python">            &gt;&gt;&gt; from torchrl.modules import ConsistentDropoutModule
            &gt;&gt;&gt; from tensordict.nn import TensorDictSequential as Seq, TensorDictModule as Mod
            &gt;&gt;&gt; from torchrl.envs import GymEnv, StepCounter, SerialEnv
            &gt;&gt;&gt; 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"]),
            ... )
            &gt;&gt;&gt; primer = get_primers_from_module(m)
            &gt;&gt;&gt; env0 = GymEnv("Pendulum-v1").append_transform(StepCounter(5))
            &gt;&gt;&gt; env1 = GymEnv("Pendulum-v1").append_transform(StepCounter(6))
            &gt;&gt;&gt; env = SerialEnv(2, [lambda env=env0: env, lambda env=env1: env])
            &gt;&gt;&gt; env = env.append_transform(primer)
            &gt;&gt;&gt; r = env.rollout(10, m, break_when_any_done=False)
            &gt;&gt;&gt; mask = [k for k in r.keys() if k.startswith("mask")][0]
            &gt;&gt;&gt; assert (r[mask][0, :5] != r[mask][0, 5:6]).any()
            &gt;&gt;&gt; assert (r[mask][0, :4] == r[mask][0, 4:5]).all()
</code></pre>
<aside class="onebox githubpullrequest" data-onebox-src="https://github.com/pytorch/rl/pull/2399">
  <header class="source">

      <a href="https://github.com/pytorch/rl/pull/2399" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/rl</a>
  </header>

  <article class="onebox-body">
    <div class="github-row" data-github-private-repo="false">



    <div class="github-icon-container" title="Pull Request">
      <svg width="60" height="60" class="github-icon" viewBox="0 0 12 16" aria-hidden="true"><path fill-rule="evenodd" d="M11 11.28V5c-.03-.78-.34-1.47-.94-2.06C9.46 2.35 8.78 2.03 8 2H7V0L4 3l3 3V4h1c.27.02.48.11.69.31.21.2.3.42.31.69v6.28A1.993 1.993 0 0 0 10 15a1.993 1.993 0 0 0 1-3.72zm-1 2.92c-.66 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2zM4 3c0-1.11-.89-2-2-2a1.993 1.993 0 0 0-1 3.72v6.56A1.993 1.993 0 0 0 2 15a1.993 1.993 0 0 0 1-3.72V4.72c.59-.34 1-.98 1-1.72zm-.8 10c0 .66-.55 1.2-1.2 1.2-.65 0-1.2-.55-1.2-1.2 0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2zM2 4.2C1.34 4.2.8 3.65.8 3c0-.65.55-1.2 1.2-1.2.65 0 1.2.55 1.2 1.2 0 .65-.55 1.2-1.2 1.2z"></path></svg>
    </div>

  <div class="github-info-container">



      <h4>
        <a href="https://github.com/pytorch/rl/pull/2399" target="_blank" rel="noopener nofollow ugc">[Feature] Consistent Dropout</a>
      </h4>

    <div class="branches">
      <code>pytorch:main</code> ← <code>N00bcak:consistent-dropout</code>
    </div>

      <div class="github-info">
        <div class="date">
          opened <span class="discourse-local-date" data-format="ll" data-date="2024-08-15" data-time="15:30:11" data-timezone="UTC">03:30PM - 15 Aug 24 UTC</span>
        </div>

        <div class="user">
          <a href="https://github.com/N00bcak" target="_blank" rel="noopener nofollow ugc">
            <img alt="N00bcak" src="https://avatars.githubusercontent.com/u/10582246?v=4" class="onebox-avatar-inline" width="20" height="20">
            N00bcak
          </a>
        </div>

        <div class="lines" title="12 commits changed 7 files with 401 additions and 16 deletions">
          <a href="https://github.com/pytorch/rl/pull/2399/files" target="_blank" rel="noopener nofollow ugc">
            <span class="added">+401</span>
            <span class="removed">-16</span>
          </a>
        </div>
      </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## Description

Introduces [Consistent Dropout (Hausknecht &amp; Wagener, 2022)](h<span class="show-more-container"><a href="https://github.com/pytorch/rl/pull/2399" target="_blank" rel="noopener nofollow ugc" class="show-more">…</a></span><span class="excerpt hidden">ttps://arxiv.org/abs/2202.11818) to TorchRL.

Consists of the following changes: 
- `ConsistentDropout` PyTorch module 
- Companion `TensorDictModule`: `ConsistentDropoutModule`
- Tests for `ConsistentDropoutModule`

## Motivation and Context

Addresses a [feature request](https://discuss.pytorch.org/t/feature-request-consistent-dropout-implementation/183778).
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)
- [x] New feature (non-breaking change which adds core functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to change)
- [x] 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!

- [x] I have read the [CONTRIBUTION](https://github.com/pytorch/rl/blob/main/CONTRIBUTING.md) guide (**required**)
- [x] My change requires a change to the documentation.
- [x] I have updated the tests accordingly (*required for a bug fix or a new feature*).
- [x] I have updated the documentation accordingly.</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>