models.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from __future__ import annotations

import dataclasses
from copy import deepcopy
from numbers import Number
from typing import Callable, Sequence

import torch
from torch import nn

from torchrl._utils import prod
from torchrl.data.utils import DEVICE_TYPING
from torchrl.modules.models.decision_transformer import DecisionTransformer
from torchrl.modules.models.utils import (
    _find_depth,
    create_on_device,
    LazyMapping,
    SquashDims,
    Squeeze2dLayer,
    SqueezeLayer,
)
from torchrl.modules.tensordict_module.common import DistributionalDQNnet  # noqa


class MLP(nn.Sequential):
    """A multi-layer perceptron.

    If MLP receives more than one input, it concatenates them all along the last dimension before passing the
    resulting tensor through the network. This is aimed at allowing for a seamless interface with calls of the type of

        >>> model(state, action)  # compute state-action value

    In the future, this feature may be moved to the ProbabilisticTDModule, though it would require it to handle
    different cases (vectors, images, ...)

    Args:
        in_features (int, optional): number of input features;
        out_features (int, torch.Size or equivalent): number of output
            features. If iterable of integers, the output is reshaped to the
            desired shape.
        depth (int, optional): depth of the network. A depth of 0 will produce
            a single linear layer network with the desired input and output size.
            A length of 1 will create 2 linear layers etc. If no depth is indicated,
            the depth information should be contained in the ``num_cells``
            argument (see below). If ``num_cells`` is an iterable and depth is
            indicated, both should match: ``len(num_cells)`` must be equal to
            ``depth``.
            Defaults to ``0`` (no depth - the network contains a single linear layer).
        num_cells (int or sequence of int, optional): number of cells of every
            layer in between the input and output. If an integer is provided,
            every layer will have the same number of cells. If an iterable is provided,
            the linear layers ``out_features`` will match the content of
            ``num_cells``. Defaults to ``32``;
        activation_class (Type[nn.Module] or callable, optional): activation
            class or constructor to be used.
            Defaults to :class:`~torch.nn.Tanh`.
        activation_kwargs (dict or list of dicts, optional): kwargs to be used
            with the activation class. Also accepts a list of kwargs of length
            ``depth + int(activate_last_layer)``.
        norm_class (Type or callable, optional): normalization class or
            constructor, if any.
        norm_kwargs (dict or list of dicts, optional): kwargs to be used with
            the normalization layers. Also accepts a list of kwargs of length
            ``depth + int(activate_last_layer)``.
        dropout (:obj:`float`, optional): dropout probability. Defaults to ``None`` (no
            dropout);
        bias_last_layer (bool): if ``True``, the last Linear layer will have a bias parameter.
            default: True;
        single_bias_last_layer (bool): if ``True``, the last dimension of the bias of the last layer will be a singleton
            dimension.
            default: True;
        layer_class (Type[nn.Module] or callable, optional): class to be used
            for the linear layers;
        layer_kwargs (dict or list of dicts, optional): kwargs for the linear
            layers. Also accepts a list of kwargs of length ``depth + 1``.
        activate_last_layer (bool): whether the MLP output should be activated. This is useful when the MLP output
            is used as the input for another module.
            default: False.
        device (torch.device, optional): device to create the module on.

    Examples:
        >>> # All of the following examples provide valid, working MLPs
        >>> mlp = MLP(in_features=3, out_features=6, depth=0) # MLP consisting of a single 3 x 6 linear layer
        >>> print(mlp)
        MLP(
          (0): Linear(in_features=3, out_features=6, bias=True)
        )
        >>> mlp = MLP(in_features=3, out_features=6, depth=4, num_cells=32)
        >>> print(mlp)
        MLP(
          (0): Linear(in_features=3, out_features=32, bias=True)
          (1): Tanh()
          (2): Linear(in_features=32, out_features=32, bias=True)
          (3): Tanh()
          (4): Linear(in_features=32, out_features=32, bias=True)
          (5): Tanh()
          (6): Linear(in_features=32, out_features=32, bias=True)
          (7): Tanh()
          (8): Linear(in_features=32, out_features=6, bias=True)
        )
        >>> mlp = MLP(out_features=6, depth=4, num_cells=32)  # LazyLinear for the first layer
        >>> print(mlp)
        MLP(
          (0): LazyLinear(in_features=0, out_features=32, bias=True)
          (1): Tanh()
          (2): Linear(in_features=32, out_features=32, bias=True)
          (3): Tanh()
          (4): Linear(in_features=32, out_features=32, bias=True)
          (5): Tanh()
          (6): Linear(in_features=32, out_features=32, bias=True)
          (7): Tanh()
          (8): Linear(in_features=32, out_features=6, bias=True)
        )
        >>> mlp = MLP(out_features=6, num_cells=[32, 33, 34, 35])  # defines the depth by the num_cells arg
        >>> print(mlp)
        MLP(
          (0): LazyLinear(in_features=0, out_features=32, bias=True)
          (1): Tanh()
          (2): Linear(in_features=32, out_features=33, bias=True)
          (3): Tanh()
          (4): Linear(in_features=33, out_features=34, bias=True)
          (5): Tanh()
          (6): Linear(in_features=34, out_features=35, bias=True)
          (7): Tanh()
          (8): Linear(in_features=35, out_features=6, bias=True)
        )
        >>> mlp = MLP(out_features=(6, 7), num_cells=[32, 33, 34, 35])  # returns a view of the output tensor with shape [*, 6, 7]
        >>> print(mlp)
        MLP(
          (0): LazyLinear(in_features=0, out_features=32, bias=True)
          (1): Tanh()
          (2): Linear(in_features=32, out_features=33, bias=True)
          (3): Tanh()
          (4): Linear(in_features=33, out_features=34, bias=True)
          (5): Tanh()
          (6): Linear(in_features=34, out_features=35, bias=True)
          (7): Tanh()
          (8): Linear(in_features=35, out_features=42, bias=True)
        )
        >>> from torchrl.modules import NoisyLinear
        >>> mlp = MLP(out_features=(6, 7), num_cells=[32, 33, 34, 35], layer_class=NoisyLinear)  # uses NoisyLinear layers
        >>> print(mlp)
        MLP(
          (0): NoisyLazyLinear(in_features=0, out_features=32, bias=False)
          (1): Tanh()
          (2): NoisyLinear(in_features=32, out_features=33, bias=True)
          (3): Tanh()
          (4): NoisyLinear(in_features=33, out_features=34, bias=True)
          (5): Tanh()
          (6): NoisyLinear(in_features=34, out_features=35, bias=True)
          (7): Tanh()
          (8): NoisyLinear(in_features=35, out_features=42, bias=True)
        )

    """

    def __init__(
        self,
        in_features: int | None = None,
        out_features: int | torch.Size = None,
        depth: int | None = None,
        num_cells: Sequence[int] | int | None = None,
        activation_class: type[nn.Module] | Callable = nn.Tanh,
        activation_kwargs: dict | list[dict] | None = None,
        norm_class: type[nn.Module] | Callable | None = None,
        norm_kwargs: dict | list[dict] | None = None,
        dropout: float | None = None,
        bias_last_layer: bool = True,
        single_bias_last_layer: bool = False,
        layer_class: type[nn.Module] | Callable = nn.Linear,
        layer_kwargs: dict | None = None,
        activate_last_layer: bool = False,
        device: DEVICE_TYPING | None = None,
    ):
        if out_features is None:
            raise ValueError("out_features must be specified for MLP.")

        if num_cells is None:
            default_num_cells = 32
            if depth is None:
                num_cells = []
                depth = 0
            else:
                num_cells = [default_num_cells] * depth

        self.in_features = in_features

        _out_features_num = out_features
        if not isinstance(out_features, Number):
            _out_features_num = prod(out_features)
        self.out_features = out_features
        self._reshape_out = not isinstance(
            self.out_features, (int, torch.SymInt, Number)
        )
        self._out_features_num = _out_features_num
        self.activation_class = activation_class
        self.norm_class = norm_class
        self.dropout = dropout
        self.bias_last_layer = bias_last_layer
        self.single_bias_last_layer = single_bias_last_layer
        self.layer_class = layer_class

        self.activation_kwargs = activation_kwargs
        self.norm_kwargs = norm_kwargs
        self.layer_kwargs = layer_kwargs

        self.activate_last_layer = activate_last_layer
        if single_bias_last_layer:
            raise NotImplementedError

        if not (isinstance(num_cells, Sequence) or depth is not None):
            raise RuntimeError(
                "If num_cells is provided as an integer, \
            depth must be provided too."
            )
        self.num_cells = (
            list(num_cells) if isinstance(num_cells, Sequence) else [num_cells] * depth
        )
        self.depth = depth if depth is not None else len(self.num_cells)
        if not (len(self.num_cells) == depth or depth is None):
            raise RuntimeError(
                "depth and num_cells length conflict, \
            consider matching or specifying a constant num_cells argument together with a a desired depth"
            )

        self._activation_kwargs_iter = _iter_maybe_over_single(
            activation_kwargs, n=self.depth + self.activate_last_layer
        )
        self._norm_kwargs_iter = _iter_maybe_over_single(
            norm_kwargs, n=self.depth + self.activate_last_layer
        )
        self._layer_kwargs_iter = _iter_maybe_over_single(
            layer_kwargs, n=self.depth + 1
        )
        layers = self._make_net(device)
        layers = [
            layer if isinstance(layer, nn.Module) else _ExecutableLayer(layer)
            for layer in layers
        ]
        super().__init__(*layers)

    def _make_net(self, device: DEVICE_TYPING | None) -> list[nn.Module]:
        layers = []
        in_features = [self.in_features] + self.num_cells
        out_features = self.num_cells + [self._out_features_num]
        for i, (_in, _out) in enumerate(zip(in_features, out_features)):
            layer_kwargs = next(self._layer_kwargs_iter)
            _bias = layer_kwargs.pop(
                "bias", self.bias_last_layer if i == self.depth else True
            )
            if _in is not None:
                layers.append(
                    create_on_device(
                        self.layer_class,
                        device,
                        _in,
                        _out,
                        bias=_bias,
                        **layer_kwargs,
                    )
                )
            else:
                try:
                    lazy_version = LazyMapping[self.layer_class]
                except KeyError:
                    raise KeyError(
                        f"The lazy version of {self.layer_class.__name__} is not implemented yet. "
                        "Consider providing the input feature dimensions explicitly when creating an MLP module"
                    )
                layers.append(
                    create_on_device(
                        lazy_version, device, _out, bias=_bias, **layer_kwargs
                    )
                )

            if i < self.depth or self.activate_last_layer:
                norm_kwargs = next(self._norm_kwargs_iter)
                activation_kwargs = next(self._activation_kwargs_iter)
                if self.dropout is not None:
                    layers.append(create_on_device(nn.Dropout, device, p=self.dropout))
                if self.norm_class is not None:
                    layers.append(
                        create_on_device(self.norm_class, device, **norm_kwargs)
                    )
                layers.append(
                    create_on_device(self.activation_class, device, **activation_kwargs)
                )

        return layers

    def forward(self, *inputs: tuple[torch.Tensor]) -> torch.Tensor:
        if len(inputs) > 1:
            inputs = (torch.cat([*inputs], -1),)

        out = super().forward(*inputs)
        if self._reshape_out:
            out = out.view(*out.shape[:-1], *self.out_features)
        return out


class ConvNet(nn.Sequential):
    """A convolutional neural network.

    Args:
        in_features (int, optional): number of input features. If ``None``, a
            :class:`~torch.nn.LazyConv2d` module is used for the first layer.;
        depth (int, optional): depth of the network. A depth of 1 will produce
            a single linear layer network with the desired input size, and
            with an output size equal to the last element of the num_cells
            argument.
            If no depth is indicated, the depth information should be contained
            in the ``num_cells`` argument (see below).
            If ``num_cells`` is an iterable and ``depth`` is indicated, both
            should match: ``len(num_cells)`` must be equal to the ``depth``.
        num_cells (int or Sequence of int, optional): number of cells of
            every layer in between the input and output. If an integer is
            provided, every layer will have the same number of cells. If an
            iterable is provided, the linear layers ``out_features`` will match
            the content of num_cells. Defaults to ``[32, 32, 32]``.
        kernel_sizes (int, sequence of int, optional): Kernel size(s) of the
            conv network. If iterable, the length must match the depth,
            defined by the ``num_cells`` or depth arguments.
            Defaults to ``3``.
        strides (int or sequence of int, optional): Stride(s) of the conv network. If
            iterable, the length must match the depth, defined by the
            ``num_cells`` or depth arguments. Defaults to ``1``.
        activation_class (Type[nn.Module] or callable, optional): activation
            class or constructor to be used.
            Defaults to :class:`~torch.nn.Tanh`.
        activation_kwargs (dict or list of dicts, optional): kwargs to be used
            with the activation class. A list of kwargs of length ``depth``
            can also be passed, with one element per layer.
        norm_class (Type or callable, optional): normalization class or
            constructor, if any.
        norm_kwargs (dict or list of dicts, optional): kwargs to be used with
            the normalization layers. A list of kwargs of length ``depth`` can
            also be passed, with one element per layer.
        bias_last_layer (bool): if ``True``, the last Linear layer will have a
            bias parameter. Defaults to ``True``.
        aggregator_class (Type[nn.Module] or callable): aggregator class or
            constructor to use at the end of the chain.
            Defaults to :class:`torchrl.modules.utils.models.SquashDims`;
        aggregator_kwargs (dict, optional): kwargs for the
            ``aggregator_class``.
        squeeze_output (bool): whether the output should be squeezed of its
            singleton dimensions.
            Defaults to ``False``.
        device (torch.device, optional): device to create the module on.

    Examples:
        >>> # All of the following examples provide valid, working MLPs
        >>> cnet = ConvNet(in_features=3, depth=1, num_cells=[32,]) # MLP consisting of a single 3 x 6 linear layer
        >>> print(cnet)
        ConvNet(
          (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1))
          (1): ELU(alpha=1.0)
          (2): SquashDims()
        )
        >>> cnet = ConvNet(in_features=3, depth=4, num_cells=32)
        >>> print(cnet)
        ConvNet(
          (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1))
          (1): ELU(alpha=1.0)
          (2): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
          (3): ELU(alpha=1.0)
          (4): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
          (5): ELU(alpha=1.0)
          (6): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1))
          (7): ELU(alpha=1.0)
          (8): SquashDims()
        )
        >>> cnet = ConvNet(in_features=3, num_cells=[32, 33, 34, 35])  # defines the depth by the num_cells arg
        >>> print(cnet)
        ConvNet(
          (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1))
          (1): ELU(alpha=1.0)
          (2): Conv2d(32, 33, kernel_size=(3, 3), stride=(1, 1))
          (3): ELU(alpha=1.0)
          (4): Conv2d(33, 34, kernel_size=(3, 3), stride=(1, 1))
          (5): ELU(alpha=1.0)
          (6): Conv2d(34, 35, kernel_size=(3, 3), stride=(1, 1))
          (7): ELU(alpha=1.0)
          (8): SquashDims()
        )
        >>> cnet = ConvNet(in_features=3, num_cells=[32, 33, 34, 35], kernel_sizes=[3, 4, 5, (2, 3)])  # defines kernels, possibly rectangular
        >>> print(cnet)
        ConvNet(
          (0): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1))
          (1): ELU(alpha=1.0)
          (2): Conv2d(32, 33, kernel_size=(4, 4), stride=(1, 1))
          (3): ELU(alpha=1.0)
          (4): Conv2d(33, 34, kernel_size=(5, 5), stride=(1, 1))
          (5): ELU(alpha=1.0)
          (6): Conv2d(34, 35, kernel_size=(2, 3), stride=(1, 1))
          (7): ELU(alpha=1.0)
          (8): SquashDims()
        )

    """

    def __init__(
        self,
        in_features: int | None = None,
        depth: int | None = None,
        num_cells: Sequence[int] | int = None,
        kernel_sizes: Sequence[int] | int = 3,
        strides: Sequence[int] | int = 1,
        paddings: Sequence[int] | int = 0,
        activation_class: type[nn.Module] | Callable = nn.ELU,
        activation_kwargs: dict | list[dict] | None = None,
        norm_class: type[nn.Module] | Callable | None = None,
        norm_kwargs: dict | list[dict] | None = None,
        bias_last_layer: bool = True,
        aggregator_class: type[nn.Module] | Callable | None = SquashDims,
        aggregator_kwargs: dict | None = None,
        squeeze_output: bool = False,
        device: DEVICE_TYPING | None = None,
    ):
        if num_cells is None:
            num_cells = [32, 32, 32]

        self.in_features = in_features
        self.activation_class = activation_class
        self.norm_class = norm_class
        self.bias_last_layer = bias_last_layer
        self.aggregator_class = aggregator_class
        self.aggregator_kwargs = (
            aggregator_kwargs if aggregator_kwargs is not None else {"ndims_in": 3}
        )
        self.squeeze_output = squeeze_output
        # self.single_bias_last_layer = single_bias_last_layer

        self.activation_kwargs = (
            activation_kwargs if activation_kwargs is not None else {}
        )
        self.norm_kwargs = norm_kwargs if norm_kwargs is not None else {}

        depth = _find_depth(depth, num_cells, kernel_sizes, strides, paddings)
        self.depth = depth
        if depth == 0:
            raise ValueError("Null depth is not permitted with ConvNet.")

        for _field, _value in zip(
            ["num_cells", "kernel_sizes", "strides", "paddings"],
            [num_cells, kernel_sizes, strides, paddings],
        ):
            _depth = depth
            setattr(
                self,
                _field,
                (_value if isinstance(_value, Sequence) else [_value] * _depth),
            )
            if not (isinstance(_value, Sequence) or _depth is not None):
                raise RuntimeError(
                    f"If {_field} is provided as an integer, "
                    "depth must be provided too."
                )
            if not (len(getattr(self, _field)) == _depth or _depth is None):
                raise RuntimeError(
                    f"depth={depth} and {_field}={len(getattr(self, _field))} length conflict, "
                    + f"consider matching or specifying a constant {_field} argument together with a a desired depth"
                )

        self.out_features = self.num_cells[-1]

        self.depth = len(self.kernel_sizes)

        self._activation_kwargs_iter = _iter_maybe_over_single(
            activation_kwargs, n=self.depth
        )
        self._norm_kwargs_iter = _iter_maybe_over_single(norm_kwargs, n=self.depth)

        layers = self._make_net(device)
        layers = [
            layer if isinstance(layer, nn.Module) else _ExecutableLayer(layer)
            for layer in layers
        ]
        super().__init__(*layers)

    def _make_net(self, device: DEVICE_TYPING | None) -> nn.Module:
        layers = []
        in_features = [self.in_features] + list(self.num_cells[: self.depth])
        out_features = list(self.num_cells) + [self.out_features]
        kernel_sizes = self.kernel_sizes
        strides = self.strides
        paddings = self.paddings
        for i, (_in, _out, _kernel, _stride, _padding) in enumerate(
            zip(in_features, out_features, kernel_sizes, strides, paddings)
        ):
            _bias = (i < len(in_features) - 1) or self.bias_last_layer
            if _in is not None:
                layers.append(
                    nn.Conv2d(
                        _in,
                        _out,
                        kernel_size=_kernel,
                        stride=_stride,
                        bias=_bias,
                        padding=_padding,
                        device=device,
                    )
                )
            else:
                layers.append(
                    nn.LazyConv2d(
                        _out,
                        kernel_size=_kernel,
                        stride=_stride,
                        bias=_bias,
                        padding=_padding,
                        device=device,
                    )
                )

            activation_kwargs = next(self._activation_kwargs_iter)
            layers.append(
                create_on_device(self.activation_class, device, **activation_kwargs)
            )
            if self.norm_class is not None:
                norm_kwargs = next(self._norm_kwargs_iter)
                layers.append(create_on_device(self.norm_class, device, **norm_kwargs))

        if self.aggregator_class is not None:
            layers.append(
                create_on_device(
                    self.aggregator_class, device, **self.aggregator_kwargs
                )
            )

        if self.squeeze_output:
            layers.append(Squeeze2dLayer())
        return layers

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        *batch, C, L, W = inputs.shape
        if len(batch) > 1:
            inputs = inputs.flatten(0, len(batch) - 1)
        out = super().forward(inputs)
        if len(batch) > 1:
            out = out.unflatten(0, batch)
        return out

    @classmethod
    def default_atari_dqn(cls, num_actions: int):
        """Returns the default DQN as presented in the seminal DQN paper.

        Args:
            num_actions (int): the action space of the atari game.

        """
        cnn = ConvNet(
            activation_class=torch.nn.ReLU,
            num_cells=[32, 64, 64],
            kernel_sizes=[8, 4, 3],
            strides=[4, 2, 1],
        )
        mlp = MLP(
            activation_class=torch.nn.ReLU,
            out_features=num_actions,
            num_cells=[512],
        )
        return nn.Sequential(cnn, mlp)


Conv2dNet = ConvNet


class Conv3dNet(nn.Sequential):
    """A 3D-convolutional neural network.

    Args:
        in_features (int, optional): number of input features. A lazy
            implementation that automatically retrieves the input size will be
            used if none is provided.
        depth (int, optional): depth of the network. A depth of ``1`` will
            produce a single linear layer network with the desired input size,
            and with an output size equal to the last element of the
            ``num_cells`` argument. If no ``depth`` is indicated, the ``depth``
            information should be contained in the ``num_cells`` argument
            (see below).
            If ``num_cells`` is an iterable and ``depth`` is indicated,
            both should match: ``len(num_cells)`` must be equal to
            the ``depth``.
        num_cells (int or sequence of int, optional): number of cells of every
            layer in between the input and output. If an integer is provided,
            every layer will have the same number of cells and the depth will
            be retrieved from ``depth``. If an iterable is
            provided, the linear layers ``out_features`` will match the content
            of num_cells. Defaults to ``[32, 32, 32]`` or ``[32] * depth` is
            depth is not ``None``.
        kernel_sizes (int, sequence of int, optional): Kernel size(s) of the
            conv network. If iterable, the length must match the depth,
            defined by the ``num_cells`` or depth arguments. Defaults to ``3``.
        strides (int or sequence of int): Stride(s) of the conv network.
            If iterable, the length must match the depth, defined by the
            ``num_cells`` or depth arguments. Defaults to ``1``.
        activation_class (Type[nn.Module] or callable): activation class or
            constructor to be used. Defaults to :class:`~torch.nn.Tanh`.
        activation_kwargs (dict or list of dicts, optional): kwargs to be used
            with the activation class. A list of kwargs of length ``depth``
            with one element per layer can also be provided.
        norm_class (Type or callable, optional): normalization class, if any.
        norm_kwargs (dict or list of dicts, optional): kwargs to be used with
            the normalization layers. A list of kwargs of length ``depth``
            with one element per layer can also be provided.
        bias_last_layer (bool): if ``True``, the last Linear layer will have a
            bias parameter. Defaults to ``True``.
        aggregator_class (Type[nn.Module] or callable): aggregator class or
            constructor to use at the end of the chain. Defaults to
            :class:`~torchrl.modules.models.utils.SquashDims`.
        aggregator_kwargs (dict, optional): kwargs for the ``aggregator_class``
            constructor.
        squeeze_output (bool): whether the output should be squeezed of its
            singleton dimensions. Defaults to ``False``.
        device (torch.device, optional): device to create the module on.

    Examples:
        >>> # All of the following examples provide valid, working MLPs
        >>> cnet = Conv3dNet(in_features=3, depth=1, num_cells=[32,])
        >>> print(cnet)
        Conv3dNet(
            (0): Conv3d(3, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1))
            (1): ELU(alpha=1.0)
            (2): SquashDims()
        )
        >>> cnet = Conv3dNet(in_features=3, depth=4, num_cells=32)
        >>> print(cnet)
        Conv3dNet(
            (0): Conv3d(3, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1))
            (1): ELU(alpha=1.0)
            (2): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1))
            (3): ELU(alpha=1.0)
            (4): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1))
            (5): ELU(alpha=1.0)
            (6): Conv3d(32, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1))
            (7): ELU(alpha=1.0)
            (8): SquashDims()
        )
        >>> cnet = Conv3dNet(in_features=3, num_cells=[32, 33, 34, 35])  # defines the depth by the num_cells arg
        >>> print(cnet)
        Conv3dNet(
            (0): Conv3d(3, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1))
            (1): ELU(alpha=1.0)
            (2): Conv3d(32, 33, kernel_size=(3, 3, 3), stride=(1, 1, 1))
            (3): ELU(alpha=1.0)
            (4): Conv3d(33, 34, kernel_size=(3, 3, 3), stride=(1, 1, 1))
            (5): ELU(alpha=1.0)
            (6): Conv3d(34, 35, kernel_size=(3, 3, 3), stride=(1, 1, 1))
            (7): ELU(alpha=1.0)
            (8): SquashDims()
        )
        >>> cnet = Conv3dNet(in_features=3, num_cells=[32, 33, 34, 35], kernel_sizes=[3, 4, 5, (2, 3, 4)])  # defines kernels, possibly rectangular
        >>> print(cnet)
        Conv3dNet(
            (0): Conv3d(3, 32, kernel_size=(3, 3, 3), stride=(1, 1, 1))
            (1): ELU(alpha=1.0)
            (2): Conv3d(32, 33, kernel_size=(4, 4, 4), stride=(1, 1, 1))
            (3): ELU(alpha=1.0)
            (4): Conv3d(33, 34, kernel_size=(5, 5, 5), stride=(1, 1, 1))
            (5): ELU(alpha=1.0)
            (6): Conv3d(34, 35, kernel_size=(2, 3, 4), stride=(1, 1, 1))
            (7): ELU(alpha=1.0)
            (8): SquashDims()
        )

    """

    def __init__(
        self,
        in_features: int | None = None,
        depth: int | None = None,
        num_cells: Sequence[int] | int = None,
        kernel_sizes: Sequence[int] | int = 3,
        strides: Sequence[int] | int = 1,
        paddings: Sequence[int] | int = 0,
        activation_class: type[nn.Module] | Callable = nn.ELU,
        activation_kwargs: dict | list[dict] | None = None,
        norm_class: type[nn.Module] | Callable | None = None,
        norm_kwargs: dict | list[dict] | None = None,
        bias_last_layer: bool = True,
        aggregator_class: type[nn.Module] | Callable | None = SquashDims,
        aggregator_kwargs: dict | None = None,
        squeeze_output: bool = False,
        device: DEVICE_TYPING | None = None,
    ):
        if num_cells is None:
            if depth is None:
                num_cells = [32, 32, 32]
            else:
                num_cells = [32] * depth

        self.in_features = in_features
        self.activation_class = activation_class
        self.norm_class = norm_class

        self.activation_kwargs = (
            activation_kwargs if activation_kwargs is not None else {}
        )
        self.norm_kwargs = norm_kwargs if norm_kwargs is not None else {}

        self.bias_last_layer = bias_last_layer
        self.aggregator_class = aggregator_class
        self.aggregator_kwargs = (
            aggregator_kwargs if aggregator_kwargs is not None else {"ndims_in": 4}
        )
        self.squeeze_output = squeeze_output
        # self.single_bias_last_layer = single_bias_last_layer

        depth = _find_depth(depth, num_cells, kernel_sizes, strides, paddings)
        self.depth = depth
        if depth == 0:
            raise ValueError("Null depth is not permitted with Conv3dNet.")

        for _field, _value in zip(
            ["num_cells", "kernel_sizes", "strides", "paddings"],
            [num_cells, kernel_sizes, strides, paddings],
        ):
            _depth = depth
            setattr(
                self,
                _field,
                (_value if isinstance(_value, Sequence) else [_value] * _depth),
            )
            if not (len(getattr(self, _field)) == _depth or _depth is None):
                raise ValueError(
                    f"depth={depth} and {_field}={len(getattr(self, _field))} length conflict, "
                    + f"consider matching or specifying a constant {_field} argument together with a a desired depth"
                )

        self.out_features = self.num_cells[-1]

        self.depth = len(self.kernel_sizes)

        self._activation_kwargs_iter = _iter_maybe_over_single(
            activation_kwargs, n=self.depth
        )
        self._norm_kwargs_iter = _iter_maybe_over_single(norm_kwargs, n=self.depth)

        layers = self._make_net(device)
        layers = [
            layer if isinstance(layer, nn.Module) else _ExecutableLayer(layer)
            for layer in layers
        ]
        super().__init__(*layers)

    def _make_net(self, device: DEVICE_TYPING | None) -> nn.Module:
        layers = []
        in_features = [self.in_features] + self.num_cells[: self.depth]
        out_features = self.num_cells + [self.out_features]
        kernel_sizes = self.kernel_sizes
        strides = self.strides
        paddings = self.paddings
        for i, (_in, _out, _kernel, _stride, _padding) in enumerate(
            zip(in_features, out_features, kernel_sizes, strides, paddings)
        ):
            _bias = (i < len(in_features) - 1) or self.bias_last_layer
            if _in is not None:
                layers.append(
                    nn.Conv3d(
                        _in,
                        _out,
                        kernel_size=_kernel,
                        stride=_stride,
                        bias=_bias,
                        padding=_padding,
                        device=device,
                    )
                )
            else:
                layers.append(
                    nn.LazyConv3d(
                        _out,
                        kernel_size=_kernel,
                        stride=_stride,
                        bias=_bias,
                        padding=_padding,
                        device=device,
                    )
                )

            activation_kwargs = next(self._activation_kwargs_iter)
            layers.append(
                create_on_device(self.activation_class, device, **activation_kwargs)
            )
            if self.norm_class is not None:
                norm_kwargs = next(self._norm_kwargs_iter)
                layers.append(create_on_device(self.norm_class, device, **norm_kwargs))

        if self.aggregator_class is not None:
            layers.append(
                create_on_device(
                    self.aggregator_class, device, **self.aggregator_kwargs
                )
            )

        if self.squeeze_output:
            layers.append(SqueezeLayer((-3, -2, -1)))
        return layers

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        try:
            *batch, C, D, L, W = inputs.shape
        except ValueError as err:
            raise ValueError(
                f"The input value of {self.__class__.__name__} must have at least 4 dimensions, got {inputs.ndim} instead."
            ) from err
        if len(batch) > 1:
            inputs = inputs.flatten(0, len(batch) - 1)
        out = super().forward(inputs)
        if len(batch) > 1:
            out = out.unflatten(0, batch)
        return out


class DuelingMlpDQNet(nn.Module):
    """Creates a Dueling MLP Q-network.

    Presented in https://arxiv.org/abs/1511.06581

    Args:
        out_features (int, torch.Size or equivalent): number of features for the advantage network
        out_features_value (int): number of features for the value network.
            Defaults to ``1``.
        mlp_kwargs_feature (dict, optional): kwargs for the feature network.
            Default is

            >>> mlp_kwargs_feature = {
            ...     'num_cells': [256, 256],
            ...     'activation_class': nn.ELU,
            ...     'out_features': 256,
            ...     'activate_last_layer': True,
            ... }

        mlp_kwargs_output (dict, optional): kwargs for the advantage and
            value networks. Default is

            >>> mlp_kwargs_output = {
            ...     "depth": 1,
            ...     "activation_class": nn.ELU,
            ...     "num_cells": 512,
            ...     "bias_last_layer": True,
            ... }

        device (torch.device, optional): device to create the module on.

    Examples:
        >>> import torch
        >>> from torchrl.modules import DuelingMlpDQNet
        >>> # we can ask for a specific output shape
        >>> net = DuelingMlpDQNet(out_features=(3, 2))
        >>> print(net)
        DuelingMlpDQNet(
          (features): MLP(
            (0): LazyLinear(in_features=0, out_features=256, bias=True)
            (1): ELU(alpha=1.0)
            (2): Linear(in_features=256, out_features=256, bias=True)
            (3): ELU(alpha=1.0)
            (4): Linear(in_features=256, out_features=256, bias=True)
            (5): ELU(alpha=1.0)
          )
          (advantage): MLP(
            (0): LazyLinear(in_features=0, out_features=512, bias=True)
            (1): ELU(alpha=1.0)
            (2): Linear(in_features=512, out_features=6, bias=True)
          )
          (value): MLP(
            (0): LazyLinear(in_features=0, out_features=512, bias=True)
            (1): ELU(alpha=1.0)
            (2): Linear(in_features=512, out_features=1, bias=True)
          )
        )
        >>> x = torch.zeros(1, 5)
        >>> y = net(x)
        >>> print(y)
        tensor([[[ 0.0232, -0.0477],
                 [-0.0226, -0.0019],
                 [-0.0314,  0.0069]]], grad_fn=<SubBackward0>)

    """

    def __init__(
        self,
        out_features: int | torch.Size,
        out_features_value: int = 1,
        mlp_kwargs_feature: dict | None = None,
        mlp_kwargs_output: dict | None = None,
        device: DEVICE_TYPING | None = None,
    ):
        super().__init__()

        mlp_kwargs_feature = (
            mlp_kwargs_feature if mlp_kwargs_feature is not None else {}
        )
        _mlp_kwargs_feature = {
            "num_cells": [256, 256],
            "out_features": 256,
            "activation_class": nn.ELU,
            "activate_last_layer": True,
        }
        _mlp_kwargs_feature.update(mlp_kwargs_feature)
        self.features = MLP(device=device, **_mlp_kwargs_feature)

        _mlp_kwargs_output = {
            "depth": 1,
            "activation_class": nn.ELU,
            "num_cells": 512,
            "bias_last_layer": True,
        }
        mlp_kwargs_output = mlp_kwargs_output if mlp_kwargs_output is not None else {}
        _mlp_kwargs_output.update(mlp_kwargs_output)
        self.out_features = out_features
        self.out_features_value = out_features_value
        self.advantage = MLP(
            out_features=out_features, device=device, **_mlp_kwargs_output
        )
        self.value = MLP(
            out_features=out_features_value, device=device, **_mlp_kwargs_output
        )
        for layer in self.modules():
            if isinstance(layer, (nn.Conv2d, nn.Linear)) and isinstance(
                layer.bias, torch.Tensor
            ):
                layer.bias.data.zero_()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.features(x)
        advantage = self.advantage(x)
        value = self.value(x)
        return value + advantage - advantage.mean(dim=-1, keepdim=True)


class DuelingCnnDQNet(nn.Module):
    """Dueling CNN Q-network.

    Presented in https://arxiv.org/abs/1511.06581

    Args:
        out_features (int): number of features for the advantage network.
        out_features_value (int): number of features for the value network.
        cnn_kwargs (dict or list of dicts, optional): kwargs for the feature
            network. Default is

            >>> cnn_kwargs = {
            ...     'num_cells': [32, 64, 64],
            ...     'strides': [4, 2, 1],
            ...     'kernel_sizes': [8, 4, 3],
            ... }

        mlp_kwargs (dict or list of dicts, optional): kwargs for the advantage
            and value network. Default is

            >>> mlp_kwargs = {
            ...     "depth": 1,
            ...     "activation_class": nn.ELU,
            ...     "num_cells": 512,
            ...     "bias_last_layer": True,
            ... }

        device (torch.device, optional): device to create the module on.

    Examples:
        >>> import torch
        >>> from torchrl.modules import DuelingCnnDQNet
        >>> net = DuelingCnnDQNet(out_features=20)
        >>> print(net)
        DuelingCnnDQNet(
          (features): ConvNet(
            (0): LazyConv2d(0, 32, kernel_size=(8, 8), stride=(4, 4))
            (1): ELU(alpha=1.0)
            (2): Conv2d(32, 64, kernel_size=(4, 4), stride=(2, 2))
            (3): ELU(alpha=1.0)
            (4): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1))
            (5): ELU(alpha=1.0)
            (6): SquashDims()
          )
          (advantage): MLP(
            (0): LazyLinear(in_features=0, out_features=512, bias=True)
            (1): ELU(alpha=1.0)
            (2): Linear(in_features=512, out_features=20, bias=True)
          )
          (value): MLP(
            (0): LazyLinear(in_features=0, out_features=512, bias=True)
            (1): ELU(alpha=1.0)
            (2): Linear(in_features=512, out_features=1, bias=True)
          )
        )
        >>> x = torch.zeros(1, 3, 64, 64)
        >>> y = net(x)
        >>> print(y.shape)
        torch.Size([1, 20])

    """

    def __init__(
        self,
        out_features: int,