[Feature Request] Optimizing an EnsembleModule

[btx0424]

Motivation

Ensembled modules should also support optimization with an API consistent with ordinary modules. However, passing EnsembledModule.parameters() to an optimizer, as usual, does not yield correct behavior.

import torch
import torch.nn as nn

from tensordict import TensorDict
from tensordict.nn import EnsembleModule, TensorDictModule

m = TensorDictModule(nn.Linear(128, 1), ["a"], ["a_out"])

m = EnsembleModule(m, 3, expand_input=True)
x = TensorDict({"a": torch.randn(32, 128)}, [32])

# this does not work
params = m.parameters()

# this works
params = list(m.params_td.values(True, True))
for param in params:
    param.retain_grad() # cannot optimize non-leaf tensors

opt = torch.optim.Adam(params)
for i in range(10):
    y = m(x)
    loss = y["a_out"].sum()
    opt.zero_grad()
    loss.backward()
    opt.step()
    print(loss.item())

Solution

A direct solution would be to override EnsembledModule.named_parameters(), which however may lead to other inconsistencies. It seems the issue is with nn.Parameter, which shares data with the tensordict from which it is created but does not receive gradients.

• [ x ] I have checked that there is no similar issue in the repo (required)

[vmoens]

Thanks for reporting!
We need to make the params apparent within tensordict modules in some way, even if they're contained in the tensordict.
One option is to put them in parameter list, but I'm not fan of that solution...
I'm currently facing a similar problem with the loss modules in torchrl where we hack our way to register the params contained in a tensordict. I'll try to come up with an elegant solution for both and ping you once i'm done!