It would also be very cool to utilize the EnsembleModule here instead of the previous approach of stacking network params. This will ensure the ensemble members are initialized correctly and give a bit more power to the user.

I am not stacking params, I am using vmap like in the normal SAC so I guess it behaves exactly like the ensemble module?
lemme see if I can swap it in

@smorad at a quick look it seems that the losses have their way of making ensembles in https://github.com/pytorch/rl/blob/main/torchrl/objectives/common.py#L174 (which also creates target params).
so I think that if we want to use EnsembleModule the best place would be to integrate it there in the future.
for now i think I am just going to stick to normal vmap as it is used in the other losses (like sac).

Hey @matteobettini, thanks for fixing this for me! It seems to work on my side as well

@smorad at a quick look it seems that the losses have their way of making ensembles in https://github.com/pytorch/rl/blob/main/torchrl/objectives/common.py#L174 (which also creates target params).

Yes they do, as the losses were implemented before the addition of EnsembleModule. @vmoens and I briefly discussed this in this issue as a more generic form of what is currently done in the losses. It is also vmapped so it should be very fast.

Adding it would be fairly straightforward and simplify the loss function:

loss_fn = SACLoss(
  q_network=EnsembleModule(MyQNetwork(), num_copies=5)
  ...
)

class SACLoss:
...
def forward(...):
  ...
  q_values = self.q_network(td)
  if isinstance(self.q_network, EnsembleModule):
    # Min reduce over q values
    q_values['action_value'] = q_values['action_value'].min(0).values
    # EnsembleModule will expand at dim 0, so this just squeezes it back
    q_values = q_values[0]
 ...

It would open up a lot of interesting options. For example, one could consider using an Ensemble of actors too, and taking the combined mode action in the discrete SAC case.

Yea ideally it would be cool to pass ensembles in place of each nn module passed to the losses

I am not super convinced by the isinstance you called in the loss.

Wouldn't there be a seamless way of passing ensembles in place of current modules and manage the dimensionality outside of the loss?

smth like what you said here

policy = TensorDictSequential(
  Ensemble(in_keys=['observation'], out_keys=['ensemble_state_action_value']),
  Reduce(in_keys=['ensemble_state_action_value'], out_keys=['state_value_action'] reduce_fn=lambda x, dim: x.min(dim=dim),
  QValueActor(env.action_spec)
)

i feel like losses should explicitly consider ensembles only when they are required by default (as SAC)
if you want to use them for other losses (like combining actor ensembles) you can do it in your module like in the snippet above

Actually, my above code is indeed wrong, we don't want to reduce at all in the q function loss! Usually, we want to update all ensemble members at once:

You could instead just do

loss_fn = SACLoss(
  q_network=EnsembleModule(MyQFn(), num_copies=5),
  q_reduce=MinReduce(reduce_key='action_value')
  ...
)

class SACLoss:
  def forward_q_loss(...):
    ...
    # Of shape [ensemble_size, batch, time, agent, ...]
    q_values = self.q_network(td)
    # Unsqueeze so this is [1, batch, time, agent, ...]
    if self.q_reduce is not None:
      target = td_estimator.estimate(td.unsqueeze(0))
    error = (q_values - target).pow(2).mean()
    return error

That code using Reduce would be used when training the actor and in the agent passed to the collector. There, we just want the "best" Q value estimate instead of all of them.

def forward_actor_loss(...):
  # TD is originally of shape [ensemble, batch, ...] from the compute_q_loss
  # Now let's reduce it to [batch, ...]
  if self.q_reduce is not None:
    td = self.q_reduce(td)
  # Continue to compute actor loss as normal
  ...

You could easily do min/softmin/mean/median/whatever reduce you want. For the case you don't want an ensemble, just do this:

loss_fn = SACLoss(
  q_network=MyQFn()
  q_reduce=None
  ...
)

I have run the discrete sac example available on main to compare main vs this PR.

The reported metrics seems the same

It would also be very cool to utilize the EnsembleModule here instead of the previous approach of stacking network params. This will ensure the ensemble members are initialized correctly and give a bit more power to the user.

@smorad happy to consider this in the future. We should open an issue about it