Optimal implementation of vanilla DQN loss in Keras

Question

I've implemented vanilla DQN for continuous/non-images (no CNN) states in keras. But, I'm not sure if my implementation of the loss computation is optimal.

For reminder the loss is defined as : $loss=(r+\gamma max_{a'} Q(s',a')-Q(s,a))^2 \frac{1}{2}$

Here is my implementation of the network + loss function :

self.network = Sequential()
self.network.add(Dense(256, activation='relu', input_dim=input_dim))
self.network.add(Dense(32, activation='relu'))
self.network.add(Dense(output_dim))

def q_loss(data, y_pred):
    # Extract the concatenated data tensor
    action, reward, next_state, done = K.cast(data[:, 0], 'int32'), data[:, 1], data[:, 2:-1], K.cast(data[:, -1], 'bool')
    # Compute Q(s,a)
    mask = tf.one_hot(action, depth=y_pred.shape[1], dtype=tf.bool, on_value=True, off_value=False)
    q_action = tf.boolean_mask(y_pred, mask)
    # Compute the max of values at next state except if done=True
    max_q_next = K.max(self.network(next_state), axis=1) * K.cast(tf.logical_not(done), 'float32')
    # Compute the TD-error, do not propagate the gradient into the next state value
    td_error = reward + 0.95 * K.stop_gradient(max_q_next) - q_action
    # Compute the MSE
    loss = K.square(td_error) / 2
    return loss

self.network.compile(loss=q_loss, optimizer=RMSprop(lr=self.learning_rate))

Here is my train function :

def train(self):
    # Sample a batch (a tuple of narray) from the replay buffer
    # States (B*S), actions (B), rewards (B), next_states(B*S), dones (B)
    # B=batch_size and S=state_size
    states, actions, rewards, next_states, dones = self.memory.sample(self.batch_size)
    # Concatenate actions, rewards, next_states, dones together because keras loss only accept one tensor
    data = np.concatenate([np.expand_dims(actions, axis=1), np.expand_dims(rewards, axis=1), next_states, np.expand_dims(dones, axis=1)], axis=1)
    # Train on a batch
    self.network.train_on_batch(states, data)

I find the way I compute the DQN td-error and loss being ugly and probably sub-optimal. Do you have a better workaround (maybe with a mix of Keras and Tensorflow combined) ?

I've checked multiple existing Keras implementations (like link), but sadly they mostly compute the td-error outside of keras in full python/numpy, witch is IMO sub-optimal.