Method __init__

utils/spatial_att_github.py:152–172

(self)

class VisionNetwork(nn.Module):
    def __init__(self):
        super(VisionNetwork, self).__init__()
        self.vision_cnn = nn.Sequential(
            nn.Conv2d(
                in_channels=3,
                out_channels=32,
                kernel_size=(8, 8),
                stride=4,
                padding=1,  # Padding s.t. the output shapes match the paper.
            ),
            nn.Conv2d(
                in_channels=32,
                out_channels=64,
                kernel_size=(4, 4),
                stride=2,
                padding=2,  # Padding s.t. the output shapes match the paper.
            ),
        )
        self.vision_lstm = ConvLSTMCell(
            input_channels=64, hidden_channels=128, kernel_size=3
        )

    def reset(self):
        self.vision_lstm.reset()