What is the target output for updating a Deep Q Network I'm trying to implement Deep Q-Learning for a pet problem having a continuous state space and discretized action space.
The algorithm for table-based Q-Learning updates a single entry of the Q table - i.e. a single $Q(s, a)$. However, a neural network outputs an entire row of the table - i.e. the Q-values for every possible action for a given state. So, what should the target output vector be for the network?
I've been trying to get it to work with something like the following:
q_values = model(state)
action = argmax(q_values)
next_state = env.step(state, action)
next_q_values = model(next_state)
max_next_q = max(next_q_values)

target_q_values = q_values
target_q_values[action] = reward(next_state) + gamma * max_next_q

The result is that my model tends to converge on some set of fixed values for every possible action - in other words, I get the same Q-values no matter what the input state is. (My guess is that this is because, since only 1 Q-value is updated, the training is teaching my model that most of its output is already fine.)
What should I be using for the target output vector for training? Should I calculate the target Q value for every action, instead of just one?
As you say, the output of a $Q$ network is typically a value for all actions of the given state. Let us call this output $\mathbf{x} \in \mathbb{R}^{|\mathcal{A}|}$. To train your network using the squared bellman error you need first calculate the scalar target $y = r(s, a) + \max_a Q(s', a)$. Then, to train the network we take a vector $\mathbf{x'} = \mathbf{x}$ and change the $a$th element of it to be equal to $y$, where $a$ is the action you took in state $s$; call this modified vector $\mathbf{x'}_a$. We calculate the loss $\mathcal{L}(\mathbf{x}, \mathbf{x'}_a)$ and back propagate through this to update the parameters of our network.
Note that when we use $Q$ to calculate $y$ we typically use some form of target network; this can be a copy of $Q$ where the parameters are only updated every $i$th update or a network whose weights are updated using a polyak average with the main networks weights after every update.
Judging by your code it looks as though your action selection is what might be causing you some problems. As far as I can tell you're always acting greedily with respect to your $Q$-function. You should be looking to act $\epsilon$-greedily, i.e. with probability $\epsilon$ take a random action and act greedily otherwise. Typically you start with $\epsilon=1$ and decay it each time a random action is taken down to some small value such as 0.05.