predictive_autoencoder.py 文件源码

def __call__(self, x, action):
        h1 = F.relu(self.conv1(x))
        index = F.expand_dims(np.array(action, dtype=np.int32), axis=0)
        h2 = F.relu(self.embed_conv2d(index, x))
        h = F.concat((h1, h2), axis=1)  # Glue together the action convolutions
        h = F.relu(self.conv2(h))
        h = F.relu(self.conv3(h))
        h = F.relu(self.conv_gru1(h))

        h_img = F.relu(self.deconv1(h))
        h_img = self.deconv2(h_img)

        h_action = F.relu(self.linear1(h))
        h_action = self.linear2(h_action)

        return h_img, h_action