[Feature Request] Q Ensembles

Motivation

Twin Q/ensemble Q functions are used in many RL algorithms and mitigate Q overestimation. My understanding is that TorchRL only deals with ensembles in the loss functions. This is fine for actor/critic methods since we only use the critics to compute actor loss. But for critic-only methods (e.g. DQN), we need the Q ensemble at sample collection time. Doing so would also simplify the loss functions for DDPG/SAC/REDQ/etc.

Solution

I would like to add ensemble Q function support to TorchRL, but I'm not sure on the best way to do this. I was thinking of creating a TensorDictModuleEnsemble in tensordict_module.py that could be used for more than just Q functions. The issue is that we essentially need two forward functions: one at sample-time to compute some reduce operation like min over the ensemble outputs, and one at training time that does not reduce, but rather does something like

self.ensemble(tensordict.expand(tensordict.shape, ensemble_size))

so we can compute the loss for all Q functions. I'm not sure if there is a good way to tell a TensorDictModule whether it is in "sampling" or "training" mode.

I think it also makes sense to provide an option to keep separate datasets for each model/Q function, e.g.

tensordict.set("ensemble_idx", tensordict.get("ensemble_idx", torch.randint(ensemble_size, 1)))
for i in range(ensemble_size):
  sub_td = tensordict[tensordict["ensemble_idx"] == i]
  self.models[i](sub_td)
...

I'd like to avoid the for-loop here if possible, but I'm not sure how.

Additional context

Related to https://github.com/pytorch/rl/issues/876

Checklist

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