Add LSTM layer after Conv2D layers and add some other inputs

Question

I'm working on a racing game that uses reinforcement learning. To train the model I'm facing an issue when implementing the neural network. I found some examples that use CNN. But it seems like adding extra LSTM layer will increase the model efficiency. I found the following example.

https://team.inria.fr/rits/files/2018/02/ICRA18_EndToEndDriving_CameraReady.pdf

The network I need to implement

The problem is I'm not sure how can I implement the LSTM layer here. How can I give following inputs to LSTM layer

1. Processed image output
2. current speed
3. last action

Here is the code I'm currently using. I want to add the LSTM layer after Conv2D.

    self.__nb_actions = 28
    self.__gamma = 0.99

    #Define the model
    activation = 'relu'
    pic_input = Input(shape=(59,255,3))

    img_stack = Conv2D(16, (3, 3), name='convolution0', padding='same', activation=activation, trainable=train_conv_layers)(pic_input)
    img_stack = MaxPooling2D(pool_size=(2,2))(img_stack)
    img_stack = Conv2D(32, (3, 3), activation=activation, padding='same', name='convolution1', trainable=train_conv_layers)(img_stack)
    img_stack = MaxPooling2D(pool_size=(2, 2))(img_stack)
    img_stack = Conv2D(32, (3, 3), activation=activation, padding='same', name='convolution2', trainable=train_conv_layers)(img_stack)
    img_stack = MaxPooling2D(pool_size=(2, 2))(img_stack)
    img_stack = Flatten()(img_stack)
    img_stack = Dropout(0.2)(img_stack)

    img_stack = Dense(128, name='rl_dense', kernel_initializer=random_normal(stddev=0.01))(img_stack)

    img_stack=Dropout(0.2)(img_stack)
    output = Dense(self.__nb_actions, name='rl_output', kernel_initializer=random_normal(stddev=0.01))(img_stack)

    opt = Adam()
    self.__action_model = Model(inputs=[pic_input], outputs=output)

    self.__action_model.compile(optimizer=opt, loss='mean_squared_error')
    self.__action_model.summary()

Thanks

Answer 1

There are various methods to do that, First, reshape the output of conv output and feed it to lstm layer. Here is an explained example with various method Shaping data for LSTM, and feeding output of dense layers to LSTM