How to use Deep Q-Network with two-dimensional input?

Hello, I'm studying with the book and i'm trying to implement my own environment inheriting from gym, and with a one-dimensional input works fine, similar to the CartPole-v1 example

But in this case the input is not a current state, but several previous temperature states, and what I want is for the system to learn to predict when the temperature is going to be maintained, when it is going to rise or when it is going to fall.

The way I see that the robot learns is that it manages to predict the trend based on the previous states, but I need to be able to enter more input parameters, such as amount of fuel, revolutions, voltage, speed and also temperature for the prediction.

The problem in particular is that since I am studying until now, I lack understanding and experience, I get several errors that I have not been able to repair, (if I manage to implement it using an RNN correctly) and when trying to calculate the Q_values ​​in the training_step function I get several errors just when I get to step 50 and the function is launched in training.

The key parts that I changed were the following

Model
```python
input_shape = [49, 5] # == env.observation_space.shape
n_outputs = 3 # == env.action_space.n

model = keras.models.Sequential([
keras.layers.Dense(32, activation="elu", input_shape=input_shape),
keras.layers.Dense(32, activation="elu"),
keras.layers.Dense(n_outputs)
])
```

Action
```python
def epsilon_greedy_policy(state, epsilon=0):
if np.random.rand() < epsilon:
return np.random.randint(3)
else:
Q_values = model.predict(state[np.newaxis])[:, -1][..., np.newaxis]
return np.argmax(Q_values[0])
```

training_step
```python
def training_step(batch_size):
experiences = sample_experiences(batch_size)
states, actions, rewards, next_states, dones = experiences
next_Q_values = model.predict(next_states)[:, -1][..., np.newaxis]
best_next_actions = np.argmax(next_Q_values, axis=1)
next_mask = tf.one_hot(best_next_actions, n_outputs).numpy()
next_best_Q_values = (target.predict(next_states)[:, -1][..., np.newaxis] * next_mask).sum(axis=1)
target_Q_values = (rewards + (1- dones) * discount_rate * next_best_Q_values)
target_Q_values = target_Q_values.reshape(-1, 1)
mask = tf.one_hot(actions, n_outputs)
with tf.GradientTape() as tape:
all_Q_values = model(states)
Q_values = tf.reduce_sum(all_Q_values * mask, axis=1, keepdims=True)
loss = tf.reduce_mean(loss_fn(target_Q_values, Q_values))
grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
```

Error
```stacktrace
in training_step(batch_size)
6 next_mask = tf.one_hot(best_next_actions, n_outputs).numpy()
7 next_best_Q_values = (target.predict(next_states)[:, -1][..., np.newaxis] * next_mask).sum(axis=1)
----> 8 target_Q_values = (rewards + (1- dones) * discount_rate * next_best_Q_values)
9 target_Q_values = target_Q_values.reshape(-1, 1)
10 mask = tf.one_hot(actions, n_outputs)

ValueError: operands could not be broadcast together with shapes (32,) (32,3)
```

I just tried to adjust the prediction based on what I learned in RNN (which is how I verify the prediction data) but it is evident that I am not understanding well how the Q values work, or how to adjust everything to use multiple inputs, I appreciate if you can help me.