Class VisionNetwork

utils/spatial_att_github.py:151–180

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()

    def forward(self, X):
        X = X.transpose(1, 3)
        O, _ = self.vision_lstm(self.vision_cnn(X))
        return O.transpose(1, 3)


class QueryNetwork(nn.Module):