Romain Guillaumé

Hi, I have two very specific questions (maybe too precise and irrelevant for performance) about the losses of TD-MPC2 and TD-MPC: 1) In the papers TD-MPC, FOWL, TD-MPC2, it looks...

The Q networks NormedLinear layers are not initialized.

1

Since the merge enabling `torch.compile`, the line `self.apply(init.weight_init)` in `WorldModel.__init__()` does not initialize the Q networks layers since `self._Qs.params` is not a `nn.ParameterList` anymore, but a `TensorDictParams` (see `init.weight_init`). They...

Request: more permissible shape checking for regression losses (equal -> broadcastable)

4

Hi, Despite what is written in the documentation [optax.losses.squared_error](https://optax.readthedocs.io/en/latest/api/losses.html#optax.losses.squared_error) i.e. "targets: a vector with shape broadcastable to that of predictions", I noticed that the `squared_error` loss in [_regression.py](https://github.com/google-deepmind/optax/blob/main/optax/losses/_regression.py) and all...