[Feature] reset_parameters for multiagent nets

[matteobettini]

Fixes #1967

Todo:

• Tests

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 6 New Failures, 21 Unrelated Failures

As of commit 48757d0 with merge base 3d65083 ():

NEW FAILURES - The following jobs have failed:

• Continuous Benchmark (PR) / CPU Pytest benchmark (gh)
  Workflow failed! Resource not accessible by integration
• Continuous Benchmark (PR) / GPU Pytest benchmark (gh)
  Workflow failed! Resource not accessible by integration
• Unit-tests on MacOS CPU / tests (3.11) / macos-job (gh)
  The process '/usr/local/bin/git' failed with exit code 128
• Unit-tests on MacOS CPU / tests (3.8) / macos-job (gh)
  The process '/usr/local/bin/git' failed with exit code 128
• Unit-tests on Windows / unittests-cpu / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128
• Unit-tests on Windows / unittests-gpu / windows-job (gh)
  The process 'C:\Program Files\Git\cmd\git.exe' failed with exit code 128

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

• Examples Tests on Linux / tests (3.9, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Generate documentation / build-docs (3.9, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Habitat Tests on Linux / tests (3.9, 11.6) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Libs Tests on Linux / unittests-brax (3.9, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Libs Tests on Linux / unittests-gym (3.9, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Libs Tests on Linux / unittests-jumanji (3.9, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Libs Tests on Linux / unittests-pettingzoo / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Libs Tests on Linux / unittests-robohive (3.9, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Libs Tests on Linux / unittests-sklearn (3.9, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Libs Tests on Linux / unittests-vmas (3.9, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Lint / c-source / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Lint / python-source-and-configs / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• RLHF Tests on Linux / unittests (3.9, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Unit-tests on Linux / tests-cpu (3.8) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Unit-tests on Linux / tests-cpu (3.9) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Unit-tests on Linux / tests-gpu (3.8, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Unit-tests on Linux / tests-olddeps (3.8, 11.6) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Unit-tests on Linux / tests-optdeps (3.9, 12.1) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128
• Unit-tests on Linux / tests-stable-gpu (3.8, 11.8) / linux-job (gh)
  The process '/usr/bin/git' failed with exit code 128

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

[matteobettini]

I would like to discuss the vmap randomness of this class

    @property
    def _vmap_randomness(self):
        if self.initialized:
            return "error"
        return "same"

Why would this be a class property and why are we having those values?

For me this should be

    @property
    def _vmap_randomness(self):
        return "different"

[matteobettini]

In this case it needs to be different to have a different reset values for each agent.

But also in the forward pass I feel like it should be "different"

[vmoens]

IMO it should be

    @property
    def _vmap_randomness(self):
        if self.initialized:
            return self.vmap_randomness
        return "different"

and users are in charge of telling the module what randomness they want.

[matteobettini]

Could you explain this a bit? Why do we have a switch on initialization?

[vmoens]

For init you need "different" because you must have different weights for each net.
But in other settings you can't tell, and the best is to let the user choose.
They may as well want the same random number for each element of the batch

[matteobettini]

but what I do not understand is why before we had

    @property
    def _vmap_randomness(self):
        if self.initialized:
            return "error"
        return "same"

[matteobettini]

if I try to change this to

    @property
    def _vmap_randomness(self):
        if self.initialized:
            return self.vmap_randomness
        return "different"

the lazy layers will crash

[vmoens]

the lazy layers will crash

this is a statement that is hard to reproduce, can you share more?

For instance this works fine on my end:

from torchrl.modules import MLP
from tensordict import TensorDict
import torch
from functorch import dim

d0 = dim.dims(1)
modules = [torch.nn.Linear(2, 3) for _ in range(3)]

td = TensorDict.from_modules(*modules, as_module=True)

def reset(td):
    with td.to_module(modules[0]):
        modules[0].reset_parameters()
    return td

td = torch.vmap(reset, randomness="same")(td)
print(td["weight"])
td = torch.vmap(reset, randomness="different")(td)
print(td["weight"])

the first produces a stack of identical tensors, the second different

[vmoens]

Thanks!