# Single-value loss/training in a CNN with a tensor output

I am playing around with an idea of using using Q-learning with a DQN (Deep Q-Network), to determine the optimal position of a number of 'units' on a grid of allowed locations, according to some reward-metric that is only calculated after placing all of the units. I am following much of Deep Convolutional Q-Learning with Python and TensorFlow 2.0, but I diverge with the grid output.

To do this, I use a CNN that takes in a one-hot encoded grid of the units (0=no units here, 1=one unit here), and outputs a grid with the same shape as the input with expected rewards of placing a unit at this location.

I can use the following Tensorflow/Keras code to get the expected rewards, where units can be placed in 3 dimensions and channels determining the different unit-styles:

from tensorflow.keras import layers, models

model = models.Sequential(name="DQN")
layers.Conv3D(
filters=10,
kernel_size=3,
activation="relu",
input_shape=input_shape,
bias_initializer="random_normal"
)
)



Currently I am using a very simple training scheme, where Q-values are first generated from the current state. At the position where the agent earlier placed a unit, the calculated true reward is given and trained against. If the following state was a terminated state, the calculated reward is used directly, while the discounted reward is used in non-terminated states.

for state, action, next_state, reward, terminated in minibatch:
# state: one-hot grid
# action: index in the state where a unit is placed by the agent
# terminated: True/False whether the 'next_state' is the terminated state

q_target = model.predict(state)

if terminated:
q_target[0][action] = reward
else:
following_target = model.predict(next_state)
q_target[0][action] = reward + gamma * np.amax(following_target)

model.fit(state, q_target, epochs=1, verbose=0)


This means, that only a single value in the entire training tensor is the true reward - all other are approximated by the CNN. However, all of the expected rewards are used in training, instead of this singular value. So I was considering whether it would be possible to train the CNN towards this single value only, and whether it would make any sense at all?

I thought of creating a custom loss function that would calculate the loss function for this single action, so training is done against this. However, I can't really figure out how I would go about doing this. I've looked at something like Custom training with tf.distribute.Strategy, but I wasn't successful at it..