duo_ai.models.impala

Classes

Impala	IMPALA convolutional neural network for feature extraction from image observations.
ImpalaBlock	A convolutional block used in the IMPALA architecture, consisting of a convolution, max pooling, and two residual blocks.
ResidualBlock	Residual block with two convolutional layers and ReLU activations.
Flatten	Module to flatten a tensor except for the batch dimension.

Module Contents

class duo_ai.models.impala.Impala(input_size: tuple, scale: int = 1)

Bases: torch.nn.Module

IMPALA convolutional neural network for feature extraction from image observations.

Examples

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

block1

block2

block3

fc

output_dim = 256

_get_fc_input_size(input_size: tuple) → int

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

Parameters:

input_size (tuple of int) – Shape of the input observation (C, H, W).

Returns:

Flattened feature size after convolutional blocks.

Return type:

int

Examples

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

forward(x: torch.Tensor) → torch.Tensor

Forward pass of the IMPALA model.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, C, H, W).

Returns:

Output feature tensor of shape (batch_size, output_dim).

Return type:

torch.Tensor

Examples

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

class duo_ai.models.impala.ImpalaBlock(in_channels: int, out_channels: int)

Bases: torch.nn.Module

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

Examples

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

conv

res1

res2

forward(x: torch.Tensor) → torch.Tensor

Forward pass of the ImpalaBlock.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, in_channels, H, W).

Returns:

Output tensor after convolution, pooling, and residual blocks.

Return type:

torch.Tensor

Examples

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

class duo_ai.models.impala.ResidualBlock(in_channels: int)

Bases: torch.nn.Module

Residual block with two convolutional layers and ReLU activations.

Examples

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

conv1

conv2

forward(x: torch.Tensor) → torch.Tensor

Forward pass of the residual block.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, in_channels, H, W).

Returns:

Output tensor after residual connection.

Return type:

torch.Tensor

Examples

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

class duo_ai.models.impala.Flatten(*args, **kwargs)

Bases: torch.nn.Module

Module to flatten a tensor except for the batch dimension.

Examples

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

forward(x: torch.Tensor) → torch.Tensor

Flatten the input tensor except for the batch dimension.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, …).

Returns:

Flattened tensor of shape (batch_size, -1).

Return type:

torch.Tensor

Examples

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