duo_ai.models.impala
====================

.. py:module:: duo_ai.models.impala


Classes
-------

.. autoapisummary::

   duo_ai.models.impala.Impala
   duo_ai.models.impala.ImpalaBlock
   duo_ai.models.impala.ResidualBlock
   duo_ai.models.impala.Flatten


Module Contents
---------------

.. py:class:: Impala(input_size: tuple, scale: int = 1)

   Bases: :py:obj:`torch.nn.Module`


   IMPALA convolutional neural network for feature extraction from image observations.

   .. rubric:: Examples

   >>> model = Impala((3, 64, 64))
   >>> x = torch.randn(8, 3, 64, 64)
   >>> out = model(x)
   >>> print(out.shape)


   .. py:attribute:: block1


   .. py:attribute:: block2


   .. py:attribute:: block3


   .. py:attribute:: fc


   .. py:attribute:: output_dim
      :value: 256



   .. py:method:: _get_fc_input_size(input_size: tuple) -> int

      Compute the input size for the fully connected layer after convolutions.

      :param input_size: Shape of the input observation (C, H, W).
      :type input_size: tuple of int

      :returns: Flattened feature size after convolutional blocks.
      :rtype: int

      .. rubric:: Examples

      >>> model = Impala((3, 64, 64))
      >>> size = model._get_fc_input_size((3, 64, 64))



   .. py:method:: forward(x: torch.Tensor) -> torch.Tensor

      Forward pass of the IMPALA model.

      :param x: Input tensor of shape (batch_size, C, H, W).
      :type x: torch.Tensor

      :returns: Output feature tensor of shape (batch_size, output_dim).
      :rtype: torch.Tensor

      .. rubric:: Examples

      >>> model = Impala((3, 64, 64))
      >>> x = torch.randn(8, 3, 64, 64)
      >>> out = model(x)



.. py:class:: ImpalaBlock(in_channels: int, out_channels: int)

   Bases: :py:obj:`torch.nn.Module`


   A convolutional block used in the IMPALA architecture, consisting of a convolution, max pooling, and two residual blocks.

   .. rubric:: Examples

   >>> block = ImpalaBlock(3, 16)
   >>> x = torch.randn(8, 3, 64, 64)
   >>> out = block(x)
   >>> print(out.shape)


   .. py:attribute:: conv


   .. py:attribute:: res1


   .. py:attribute:: res2


   .. py:method:: forward(x: torch.Tensor) -> torch.Tensor

      Forward pass of the ImpalaBlock.

      :param x: Input tensor of shape (batch_size, in_channels, H, W).
      :type x: torch.Tensor

      :returns: Output tensor after convolution, pooling, and residual blocks.
      :rtype: torch.Tensor

      .. rubric:: Examples

      >>> block = ImpalaBlock(3, 16)
      >>> x = torch.randn(8, 3, 64, 64)
      >>> out = block(x)



.. py:class:: ResidualBlock(in_channels: int)

   Bases: :py:obj:`torch.nn.Module`


   Residual block with two convolutional layers and ReLU activations.

   .. rubric:: Examples

   >>> block = ResidualBlock(16)
   >>> x = torch.randn(8, 16, 32, 32)
   >>> out = block(x)
   >>> print(out.shape)


   .. py:attribute:: conv1


   .. py:attribute:: conv2


   .. py:method:: forward(x: torch.Tensor) -> torch.Tensor

      Forward pass of the residual block.

      :param x: Input tensor of shape (batch_size, in_channels, H, W).
      :type x: torch.Tensor

      :returns: Output tensor after residual connection.
      :rtype: torch.Tensor

      .. rubric:: Examples

      >>> block = ResidualBlock(16)
      >>> x = torch.randn(8, 16, 32, 32)
      >>> out = block(x)



.. py:class:: Flatten(*args, **kwargs)

   Bases: :py:obj:`torch.nn.Module`


   Module to flatten a tensor except for the batch dimension.

   .. rubric:: Examples

   >>> flatten = Flatten()
   >>> x = torch.randn(8, 16, 4, 4)
   >>> out = flatten(x)
   >>> print(out.shape)


   .. py:method:: forward(x: torch.Tensor) -> torch.Tensor

      Flatten the input tensor except for the batch dimension.

      :param x: Input tensor of shape (batch_size, ...).
      :type x: torch.Tensor

      :returns: Flattened tensor of shape (batch_size, -1).
      :rtype: torch.Tensor

      .. rubric:: Examples

      >>> flatten = Flatten()
      >>> x = torch.randn(8, 16, 4, 4)
      >>> out = flatten(x)