Keras LSTM with stateful in Reinforcement Learning - tensorflow

Question

I'm doing a simple DQN RL algorithm with Keras, but using an LSTM in the network. The idea is that a stateful LSTM will remember the relevant information from all prior states and thus predict rewards for different actions better. This problem is more of a keras problem than RL. I think the stateful LSTM is not being handled by me correctly.
MODEL CODE - functional api used:
state_input = Input( batch_shape=(batch_size,look_back,1,resolution[0], resolution[1]))
conv1 = TimeDistributed(Conv2D(8, 6, strides=3, activation='relu', data_format="channels_first"))(
state_input) # filters, kernal_size, stride
conv2 = TimeDistributed(Conv2D(8, 3, strides=2, activation='relu', data_format="channels_first"))(
conv1) # filters, kernal_size, stride
flatten = TimeDistributed(Flatten())(conv2)
fc1 = TimeDistributed(Dense(128,activation='relu'))(flatten)
fc2 = TimeDistributed(Dense(64, activation='relu'))(fc1)
lstm_layer = LSTM(4, stateful=True)(fc2)
fc3 = Dense(128, activation='relu')(lstm_layer)
fc4 = Dense(available_actions_count)(fc3)
model = keras.models.Model(input=state_input, output=fc4)
adam = RMSprop(lr= learning_rate)#0.001
model.compile(loss="mse", optimizer=adam)
print(model.summary())
This is the model summary:
Layer (type) Output Shape Param
=================================================================
input_1 (InputLayer) (1, 1, 1, 30, 45) 0
_________________________________________________________________
time_distributed_1 (TimeDist (1, 1, 8, 9, 14) 296
_________________________________________________________________
time_distributed_2 (TimeDist (1, 1, 8, 4, 6) 584
_________________________________________________________________
time_distributed_3 (TimeDist (1, 1, 192) 0
_________________________________________________________________
time_distributed_4 (TimeDist (1, 1, 128) 24704
_________________________________________________________________
time_distributed_5 (TimeDist (1, 1, 64) 8256
_________________________________________________________________
lstm_1 (LSTM) (1, 4) 1104
_________________________________________________________________
dense_3 (Dense) (1, 128) 640
_________________________________________________________________
dense_4 (Dense) (1, 8) 1032
=================================================================
Total params: 36,616
Trainable params: 36,616
Non-trainable params: 0
================================================================
I feed in one frame at a time to fit the model. When I need to predict to act, I make sure I save the model state and restore it as follows.
CODE TO FIT/TRAIN THE MODEL:
#save the state (lstm memory) for recovery before fitting.
prev_state = get_model_states(model)
target_q = model.predict(s1, batch_size=batch_size)
#lstm predict updates the state of the lstm modules
q_next = model.predict(s2, batch_size=batch_size)
max_q_next = np.max(q_next, axis=1)
target_q[np.arange(target_q.shape[0]), a] = r + discount_factor * (1 - isterminal) * max_q_next
#now recover states for fitting the model correctly
set_model_states(model,prev_state)#to before s1 prediction
model.fit(s1, target_q,batch_size=batch_size, verbose=0)
#after fitting, the state and weights both get updated !!
#so lstm has already moved forward in the sequence
The model does not seem to be working at all. The variance remains very high across different epochs. I am resetting the model after every episode , as one would expect. So stateful does not affect the training between episodes. Each episode is fed in one frame at a time, that is why I need stateful.
I've tried different discount factors and learning rates.In theory, this should be a superior model to the vanilla dqn (CNN with 4 frames )
What am I doing wrong ? Any help would be appreciated.