[Feature] Reset parameters of multiagent networks

[matteobettini]

Hello!

So I usually used a function like this to reset parameters of the multiagent networks

def reset_child_params(module):
    for layer in module.children():
        if hasattr(layer, "reset_parameters"):
            layer.reset_parameters()
        reset_child_params(layer)

After #1921 this seems to have no effect.

Is there a suggested way to reset the parameters?

Thanks!

[matteobettini]

Even when resetting through the dedicated tensordict function it does not work

from tensordict.nn import TensorDictModule
from torch import nn

from torchrl.modules.models.multiagent import MultiAgentMLP

if __name__ == "__main__":
    actor_net = MultiAgentMLP(
        n_agent_inputs=4,
        n_agent_outputs=6,
        n_agents=2,
        centralised=False,
        share_params=False,
        device="cpu",
        depth=2,
        num_cells=256,
        activation_class=nn.Tanh,
    )

    policy_module = TensorDictModule(
        actor_net,
        in_keys=[("agents", "observation")],
        out_keys=[("agents", "action")],
    )
    params_before = list(policy_module.parameters())
    policy_module.reset_parameters_recursive()
    params_after = list(policy_module.parameters())
    for p1, p2 in zip(params_before, params_after):
        assert (p1 != p2).all()

[vmoens]

That sounds like something we should support! I have a limited bandwidth and that doesn't seem very complex so feel free to submit a PR if you need this (semi-)urgently

[matteobettini]

Do we have any insights of why policy_module.reset_parameters_recursive() runs without error but does not apply the reset?

[matteobettini]

From what I investigated it seems like when reset_parameters_recursive is called it looks for an nn.Module with the reset_parameters() function. It will not find anything with TensorDictParams

I made a PR to warn when this is a no-op

Will make another PR to implement reset_parameters() for the multiagent nets